refac/enh: dedicated enable image edit toggle

This commit is contained in:
Timothy Jaeryang Baek 2025-11-20 19:52:31 -05:00
parent 0c430629e5
commit 557170c0b6
5 changed files with 72 additions and 47 deletions

View file

@ -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",

View file

@ -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

View file

@ -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,

View file

@ -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:
if len(input_images) > 0 and request.app.state.config.ENABLE_IMAGE_EDIT:
# Edit image(s)
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_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"<context>Image generation was attempted but failed. The system is currently unable to generate the image. Tell the user that an error occurred: {error_message}</context>"
else:
# Edit image(s)
# Create image(s)
if request.app.state.config.ENABLE_IMAGE_PROMPT_GENERATION:
try:
images = await image_edits(
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_generations(
request=request,
form_data=EditImageForm(**{"prompt": prompt, "image": input_images}),
form_data=CreateImageForm(**{"prompt": prompt}),
user=user,
)

View file

@ -888,23 +888,15 @@
<div class="flex w-full justify-between items-center">
<div class="text-xs pr-2">
<div class="">
{$i18n.t('Image Edit Engine')}
{$i18n.t('Image Edit')}
</div>
</div>
<select
class=" dark:bg-gray-900 w-fit pr-8 cursor-pointer rounded-sm px-2 text-xs bg-transparent outline-hidden text-right"
bind:value={config.IMAGE_EDIT_ENGINE}
placeholder={$i18n.t('Select Engine')}
>
<option value="openai">{$i18n.t('Default (Open AI)')}</option>
<option value="comfyui">{$i18n.t('ComfyUI')}</option>
<option value="gemini">{$i18n.t('Gemini')}</option>
</select>
<Switch bind:state={config.ENABLE_IMAGE_EDIT} />
</div>
</div>
{#if config.ENABLE_IMAGE_GENERATION}
{#if config?.ENABLE_IMAGE_GENERATION && config?.ENABLE_IMAGE_EDIT}
<div class="mb-2.5">
<div class="flex w-full justify-between items-center">
<div class="text-xs pr-2">
@ -949,6 +941,26 @@
</div>
{/if}
<div class="mb-2.5">
<div class="flex w-full justify-between items-center">
<div class="text-xs pr-2">
<div class="">
{$i18n.t('Image Edit Engine')}
</div>
</div>
<select
class=" dark:bg-gray-900 w-fit pr-8 cursor-pointer rounded-sm px-2 text-xs bg-transparent outline-hidden text-right"
bind:value={config.IMAGE_EDIT_ENGINE}
placeholder={$i18n.t('Select Engine')}
>
<option value="openai">{$i18n.t('Default (Open AI)')}</option>
<option value="comfyui">{$i18n.t('ComfyUI')}</option>
<option value="gemini">{$i18n.t('Gemini')}</option>
</select>
</div>
</div>
{#if config?.IMAGE_EDIT_ENGINE === 'openai'}
<div class="mb-2.5">
<div class="flex w-full justify-between items-center">