From 182c49b41c477ac310e04b3532e0e999f06bbb8d Mon Sep 17 00:00:00 2001 From: Debanjum Date: Wed, 15 Jan 2025 17:55:21 +0700 Subject: [PATCH] Prefer explicitly configured OpenAI API url, key for image gen model Previously we'd use the general openai client, even if the image generation model has a different api key and base url set. This change uses the openai config of the image generation models when set. Otherwise it fallbacks to use the first openai api provider set --- src/khoj/processor/image/generate.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/src/khoj/processor/image/generate.py b/src/khoj/processor/image/generate.py index a4628115..daf0dbf9 100644 --- a/src/khoj/processor/image/generate.py +++ b/src/khoj/processor/image/generate.py @@ -119,25 +119,27 @@ async def text_to_image( def generate_image_with_openai( improved_image_prompt: str, text_to_image_config: TextToImageModelConfig, text2image_model: str ): - "Generate image using OpenAI API" + "Generate image using OpenAI (compatible) API" - # Get the API key from the user's configuration + # Get the API config from the user's configuration + api_key = None if text_to_image_config.api_key: api_key = text_to_image_config.api_key + openai_client = openai.OpenAI(api_key=api_key) elif text_to_image_config.ai_model_api: api_key = text_to_image_config.ai_model_api.api_key + api_base_url = text_to_image_config.ai_model_api.api_base_url + openai_client = openai.OpenAI(api_key=api_key, base_url=api_base_url) elif state.openai_client: - api_key = state.openai_client.api_key - auth_header = {"Authorization": f"Bearer {api_key}"} if api_key else {} + openai_client = state.openai_client # Generate image using OpenAI API OPENAI_IMAGE_GEN_STYLE = "vivid" - response = state.openai_client.images.generate( + response = openai_client.images.generate( prompt=improved_image_prompt, model=text2image_model, style=OPENAI_IMAGE_GEN_STYLE, response_format="b64_json", - extra_headers=auth_header, ) # Extract the base64 image from the response