diff --git a/backend/open_webui/config.py b/backend/open_webui/config.py index 9b562af5ab..4070e03376 100644 --- a/backend/open_webui/config.py +++ b/backend/open_webui/config.py @@ -3539,6 +3539,11 @@ IMAGES_GEMINI_ENDPOINT_METHOD = PersistentConfig( os.getenv("IMAGES_GEMINI_ENDPOINT_METHOD", ""), ) +ENABLE_IMAGE_EDIT = PersistentConfig( + "ENABLE_IMAGE_EDIT", + "images.edit.enable", + os.environ.get("ENABLE_IMAGE_EDIT", "").lower() == "true", +) IMAGE_EDIT_ENGINE = PersistentConfig( "IMAGE_EDIT_ENGINE", diff --git a/backend/open_webui/main.py b/backend/open_webui/main.py index 34da15eb67..0740482263 100644 --- a/backend/open_webui/main.py +++ b/backend/open_webui/main.py @@ -164,6 +164,7 @@ from open_webui.config import ( IMAGES_GEMINI_API_BASE_URL, IMAGES_GEMINI_API_KEY, IMAGES_GEMINI_ENDPOINT_METHOD, + ENABLE_IMAGE_EDIT, IMAGE_EDIT_ENGINE, IMAGE_EDIT_MODEL, IMAGE_EDIT_SIZE, @@ -1118,6 +1119,7 @@ app.state.config.COMFYUI_WORKFLOW = COMFYUI_WORKFLOW app.state.config.COMFYUI_WORKFLOW_NODES = COMFYUI_WORKFLOW_NODES +app.state.config.ENABLE_IMAGE_EDIT = ENABLE_IMAGE_EDIT app.state.config.IMAGE_EDIT_ENGINE = IMAGE_EDIT_ENGINE app.state.config.IMAGE_EDIT_MODEL = IMAGE_EDIT_MODEL app.state.config.IMAGE_EDIT_SIZE = IMAGE_EDIT_SIZE diff --git a/backend/open_webui/routers/images.py b/backend/open_webui/routers/images.py index 5f035695a9..c4e67ae9ea 100644 --- a/backend/open_webui/routers/images.py +++ b/backend/open_webui/routers/images.py @@ -126,6 +126,7 @@ class ImagesConfig(BaseModel): IMAGES_GEMINI_API_KEY: str IMAGES_GEMINI_ENDPOINT_METHOD: str + ENABLE_IMAGE_EDIT: bool IMAGE_EDIT_ENGINE: str IMAGE_EDIT_MODEL: str IMAGE_EDIT_SIZE: Optional[str] @@ -164,6 +165,7 @@ async def get_config(request: Request, user=Depends(get_admin_user)): "IMAGES_GEMINI_API_BASE_URL": request.app.state.config.IMAGES_GEMINI_API_BASE_URL, "IMAGES_GEMINI_API_KEY": request.app.state.config.IMAGES_GEMINI_API_KEY, "IMAGES_GEMINI_ENDPOINT_METHOD": request.app.state.config.IMAGES_GEMINI_ENDPOINT_METHOD, + "ENABLE_IMAGE_EDIT": request.app.state.config.ENABLE_IMAGE_EDIT, "IMAGE_EDIT_ENGINE": request.app.state.config.IMAGE_EDIT_ENGINE, "IMAGE_EDIT_MODEL": request.app.state.config.IMAGE_EDIT_MODEL, "IMAGE_EDIT_SIZE": request.app.state.config.IMAGE_EDIT_SIZE, @@ -253,6 +255,7 @@ async def update_config( ) # Edit Image + request.app.state.config.ENABLE_IMAGE_EDIT = form_data.ENABLE_IMAGE_EDIT request.app.state.config.IMAGE_EDIT_ENGINE = form_data.IMAGE_EDIT_ENGINE request.app.state.config.IMAGE_EDIT_MODEL = form_data.IMAGE_EDIT_MODEL request.app.state.config.IMAGE_EDIT_SIZE = form_data.IMAGE_EDIT_SIZE @@ -308,6 +311,7 @@ async def update_config( "IMAGES_GEMINI_API_BASE_URL": request.app.state.config.IMAGES_GEMINI_API_BASE_URL, "IMAGES_GEMINI_API_KEY": request.app.state.config.IMAGES_GEMINI_API_KEY, "IMAGES_GEMINI_ENDPOINT_METHOD": request.app.state.config.IMAGES_GEMINI_ENDPOINT_METHOD, + "ENABLE_IMAGE_EDIT": request.app.state.config.ENABLE_IMAGE_EDIT, "IMAGE_EDIT_ENGINE": request.app.state.config.IMAGE_EDIT_ENGINE, "IMAGE_EDIT_MODEL": request.app.state.config.IMAGE_EDIT_MODEL, "IMAGE_EDIT_SIZE": request.app.state.config.IMAGE_EDIT_SIZE, diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py index 7653536e34..5095bb418b 100644 --- a/backend/open_webui/utils/middleware.py +++ b/backend/open_webui/utils/middleware.py @@ -791,42 +791,13 @@ async def chat_image_generation_handler( input_images = get_last_images(message_list) system_message_content = "" - if len(input_images) == 0: - # Create image(s) - if request.app.state.config.ENABLE_IMAGE_PROMPT_GENERATION: - try: - res = await generate_image_prompt( - request, - { - "model": form_data["model"], - "messages": form_data["messages"], - }, - user, - ) - - response = res["choices"][0]["message"]["content"] - - try: - bracket_start = response.find("{") - bracket_end = response.rfind("}") + 1 - - if bracket_start == -1 or bracket_end == -1: - raise Exception("No JSON object found in the response") - - response = response[bracket_start:bracket_end] - response = json.loads(response) - prompt = response.get("prompt", []) - except Exception as e: - prompt = user_message - - except Exception as e: - log.exception(e) - prompt = user_message + if len(input_images) > 0 and request.app.state.config.ENABLE_IMAGE_EDIT: + # Edit image(s) try: - images = await image_generations( + images = await image_edits( request=request, - form_data=CreateImageForm(**{"prompt": prompt}), + form_data=EditImageForm(**{"prompt": prompt, "image": input_images}), user=user, ) @@ -874,12 +845,43 @@ async def chat_image_generation_handler( ) system_message_content = f"Image generation was attempted but failed. The system is currently unable to generate the image. Tell the user that an error occurred: {error_message}" + else: - # Edit image(s) + # Create image(s) + if request.app.state.config.ENABLE_IMAGE_PROMPT_GENERATION: + try: + res = await generate_image_prompt( + request, + { + "model": form_data["model"], + "messages": form_data["messages"], + }, + user, + ) + + response = res["choices"][0]["message"]["content"] + + try: + bracket_start = response.find("{") + bracket_end = response.rfind("}") + 1 + + if bracket_start == -1 or bracket_end == -1: + raise Exception("No JSON object found in the response") + + response = response[bracket_start:bracket_end] + response = json.loads(response) + prompt = response.get("prompt", []) + except Exception as e: + prompt = user_message + + except Exception as e: + log.exception(e) + prompt = user_message + try: - images = await image_edits( + images = await image_generations( request=request, - form_data=EditImageForm(**{"prompt": prompt, "image": input_images}), + form_data=CreateImageForm(**{"prompt": prompt}), user=user, ) diff --git a/src/lib/components/admin/Settings/Images.svelte b/src/lib/components/admin/Settings/Images.svelte index 2794f2bbaa..1c0c0b07f1 100644 --- a/src/lib/components/admin/Settings/Images.svelte +++ b/src/lib/components/admin/Settings/Images.svelte @@ -888,23 +888,15 @@
- {$i18n.t('Image Edit Engine')} + {$i18n.t('Image Edit')}
- +
- {#if config.ENABLE_IMAGE_GENERATION} + {#if config?.ENABLE_IMAGE_GENERATION && config?.ENABLE_IMAGE_EDIT}
@@ -949,6 +941,26 @@
{/if} +
+
+
+
+ {$i18n.t('Image Edit Engine')} +
+
+ + +
+
+ {#if config?.IMAGE_EDIT_ENGINE === 'openai'}