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"