From 63e8ab7a05faf6d83ce056a5f11608537f1d1205 Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Thu, 6 Nov 2025 03:43:59 -0500 Subject: [PATCH] feat: comfyui image edit support --- backend/open_webui/config.py | 23 +++ backend/open_webui/main.py | 8 ++ backend/open_webui/routers/images.py | 134 +++++++++++++----- backend/open_webui/utils/images/comfyui.py | 120 ++++++++++++++++ .../components/admin/Settings/Images.svelte | 57 +++++++- 5 files changed, 296 insertions(+), 46 deletions(-) diff --git a/backend/open_webui/config.py b/backend/open_webui/config.py index 2e78754537..2576e8e995 100644 --- a/backend/open_webui/config.py +++ b/backend/open_webui/config.py @@ -3351,6 +3351,29 @@ IMAGES_EDIT_GEMINI_API_KEY = PersistentConfig( ) +IMAGES_EDIT_COMFYUI_BASE_URL = PersistentConfig( + "IMAGES_EDIT_COMFYUI_BASE_URL", + "images.edit.comfyui.base_url", + os.getenv("IMAGES_EDIT_COMFYUI_BASE_URL", ""), +) +IMAGES_EDIT_COMFYUI_API_KEY = PersistentConfig( + "IMAGES_EDIT_COMFYUI_API_KEY", + "images.edit.comfyui.api_key", + os.getenv("IMAGES_EDIT_COMFYUI_API_KEY", ""), +) + +IMAGES_EDIT_COMFYUI_WORKFLOW = PersistentConfig( + "IMAGES_EDIT_COMFYUI_WORKFLOW", + "images.edit.comfyui.workflow", + os.getenv("IMAGES_EDIT_COMFYUI_WORKFLOW", ""), +) + +IMAGES_EDIT_COMFYUI_WORKFLOW_NODES = PersistentConfig( + "IMAGES_EDIT_COMFYUI_WORKFLOW_NODES", + "images.edit.comfyui.nodes", + [], +) + #################################### # Audio #################################### diff --git a/backend/open_webui/main.py b/backend/open_webui/main.py index 7907e87505..f0aeeab02a 100644 --- a/backend/open_webui/main.py +++ b/backend/open_webui/main.py @@ -171,6 +171,10 @@ from open_webui.config import ( IMAGES_EDIT_OPENAI_API_VERSION, IMAGES_EDIT_GEMINI_API_BASE_URL, IMAGES_EDIT_GEMINI_API_KEY, + IMAGES_EDIT_COMFYUI_BASE_URL, + IMAGES_EDIT_COMFYUI_API_KEY, + IMAGES_EDIT_COMFYUI_WORKFLOW, + IMAGES_EDIT_COMFYUI_WORKFLOW_NODES, # Audio AUDIO_STT_ENGINE, AUDIO_STT_MODEL, @@ -1106,6 +1110,10 @@ app.state.config.IMAGES_EDIT_OPENAI_API_KEY = IMAGES_EDIT_OPENAI_API_KEY app.state.config.IMAGES_EDIT_OPENAI_API_VERSION = IMAGES_EDIT_OPENAI_API_VERSION app.state.config.IMAGES_EDIT_GEMINI_API_BASE_URL = IMAGES_EDIT_GEMINI_API_BASE_URL app.state.config.IMAGES_EDIT_GEMINI_API_KEY = IMAGES_EDIT_GEMINI_API_KEY +app.state.config.IMAGES_EDIT_COMFYUI_BASE_URL = IMAGES_EDIT_COMFYUI_BASE_URL +app.state.config.IMAGES_EDIT_COMFYUI_API_KEY = IMAGES_EDIT_COMFYUI_API_KEY +app.state.config.IMAGES_EDIT_COMFYUI_WORKFLOW = IMAGES_EDIT_COMFYUI_WORKFLOW +app.state.config.IMAGES_EDIT_COMFYUI_WORKFLOW_NODES = IMAGES_EDIT_COMFYUI_WORKFLOW_NODES ######################################## diff --git a/backend/open_webui/routers/images.py b/backend/open_webui/routers/images.py index 34973a292a..b1b3994968 100644 --- a/backend/open_webui/routers/images.py +++ b/backend/open_webui/routers/images.py @@ -22,8 +22,11 @@ from open_webui.utils.auth import get_admin_user, get_verified_user from open_webui.utils.headers import include_user_info_headers from open_webui.utils.images.comfyui import ( ComfyUICreateImageForm, + ComfyUIEditImageForm, ComfyUIWorkflow, + comfyui_upload_image, comfyui_create_image, + comfyui_edit_image, ) from pydantic import BaseModel @@ -126,6 +129,10 @@ class ImagesConfig(BaseModel): IMAGES_EDIT_OPENAI_API_VERSION: str IMAGES_EDIT_GEMINI_API_BASE_URL: str IMAGES_EDIT_GEMINI_API_KEY: str + IMAGES_EDIT_COMFYUI_BASE_URL: str + IMAGES_EDIT_COMFYUI_API_KEY: str + IMAGES_EDIT_COMFYUI_WORKFLOW: str + IMAGES_EDIT_COMFYUI_WORKFLOW_NODES: list[dict] @router.get("/config", response_model=ImagesConfig) @@ -158,6 +165,10 @@ async def get_config(request: Request, user=Depends(get_admin_user)): "IMAGES_EDIT_OPENAI_API_VERSION": request.app.state.config.IMAGES_EDIT_OPENAI_API_VERSION, "IMAGES_EDIT_GEMINI_API_BASE_URL": request.app.state.config.IMAGES_EDIT_GEMINI_API_BASE_URL, "IMAGES_EDIT_GEMINI_API_KEY": request.app.state.config.IMAGES_EDIT_GEMINI_API_KEY, + "IMAGES_EDIT_COMFYUI_BASE_URL": request.app.state.config.IMAGES_EDIT_COMFYUI_BASE_URL, + "IMAGES_EDIT_COMFYUI_API_KEY": request.app.state.config.IMAGES_EDIT_COMFYUI_API_KEY, + "IMAGES_EDIT_COMFYUI_WORKFLOW": request.app.state.config.IMAGES_EDIT_COMFYUI_WORKFLOW, + "IMAGES_EDIT_COMFYUI_WORKFLOW_NODES": request.app.state.config.IMAGES_EDIT_COMFYUI_WORKFLOW_NODES, } @@ -253,6 +264,19 @@ async def update_config( form_data.IMAGES_EDIT_GEMINI_API_KEY ) + request.app.state.config.IMAGES_EDIT_COMFYUI_BASE_URL = ( + form_data.IMAGES_EDIT_COMFYUI_BASE_URL.strip("/") + ) + request.app.state.config.IMAGES_EDIT_COMFYUI_API_KEY = ( + form_data.IMAGES_EDIT_COMFYUI_API_KEY + ) + request.app.state.config.IMAGES_EDIT_COMFYUI_WORKFLOW = ( + form_data.IMAGES_EDIT_COMFYUI_WORKFLOW + ) + request.app.state.config.IMAGES_EDIT_COMFYUI_WORKFLOW_NODES = ( + form_data.IMAGES_EDIT_COMFYUI_WORKFLOW_NODES + ) + return { "ENABLE_IMAGE_GENERATION": request.app.state.config.ENABLE_IMAGE_GENERATION, "ENABLE_IMAGE_PROMPT_GENERATION": request.app.state.config.ENABLE_IMAGE_PROMPT_GENERATION, @@ -281,6 +305,10 @@ async def update_config( "IMAGES_EDIT_OPENAI_API_VERSION": request.app.state.config.IMAGES_EDIT_OPENAI_API_VERSION, "IMAGES_EDIT_GEMINI_API_BASE_URL": request.app.state.config.IMAGES_EDIT_GEMINI_API_BASE_URL, "IMAGES_EDIT_GEMINI_API_KEY": request.app.state.config.IMAGES_EDIT_GEMINI_API_KEY, + "IMAGES_EDIT_COMFYUI_BASE_URL": request.app.state.config.IMAGES_EDIT_COMFYUI_BASE_URL, + "IMAGES_EDIT_COMFYUI_API_KEY": request.app.state.config.IMAGES_EDIT_COMFYUI_API_KEY, + "IMAGES_EDIT_COMFYUI_WORKFLOW": request.app.state.config.IMAGES_EDIT_COMFYUI_WORKFLOW, + "IMAGES_EDIT_COMFYUI_WORKFLOW_NODES": request.app.state.config.IMAGES_EDIT_COMFYUI_WORKFLOW_NODES, } @@ -790,6 +818,20 @@ async def image_edits( except Exception as e: raise HTTPException(status_code=400, detail=ERROR_MESSAGES.DEFAULT(e)) + def get_image_file_item(base64_string): + data = base64_string + header, encoded = data.split(",", 1) + mime_type = header.split(";")[0].lstrip("data:") + image_data = base64.b64decode(encoded) + return ( + "image", + ( + f"{uuid.uuid4()}.png", + io.BytesIO(image_data), + mime_type if mime_type else "image/png", + ), + ) + r = None try: if request.app.state.config.IMAGE_EDIT_ENGINE == "openai": @@ -807,25 +849,11 @@ async def image_edits( **({"size": size} if size else {}), **( {} - if "gpt-image-1" in request.app.state.config.IMAGE_GENERATION_MODEL + if "gpt-image-1" in request.app.state.config.IMAGE_EDIT_MODEL else {"response_format": "b64_json"} ), } - def get_image_file_item(base64_string): - data = base64_string - header, encoded = data.split(",", 1) - mime_type = header.split(";")[0].lstrip("data:") - image_data = base64.b64decode(encoded) - return ( - "image", - ( - f"{uuid.uuid4()}.png", - io.BytesIO(image_data), - mime_type if mime_type else "image/png", - ), - ) - files = [] if isinstance(form_data.image, str): files = [get_image_file_item(form_data.image)] @@ -840,7 +868,7 @@ async def image_edits( # Use asyncio.to_thread for the requests.post call r = await asyncio.to_thread( requests.post, - url=f"{request.app.state.config.IMAGES_OPENAI_API_BASE_URL}/images/edits{url_search_params}", + url=f"{request.app.state.config.IMAGES_EDIT_OPENAI_API_BASE_URL}/images/edits{url_search_params}", headers=headers, files=files, data=data, @@ -860,10 +888,10 @@ async def image_edits( images.append({"url": url}) return images - elif request.app.state.config.IMAGE_GENERATION_ENGINE == "gemini": + elif request.app.state.config.IMAGE_EDIT_ENGINE == "gemini": headers = { "Content-Type": "application/json", - "x-goog-api-key": request.app.state.config.IMAGES_GEMINI_API_KEY, + "x-goog-api-key": request.app.state.config.IMAGES_EDIT_GEMINI_API_KEY, } model = f"{model}:generateContent" @@ -894,7 +922,7 @@ async def image_edits( # Use asyncio.to_thread for the requests.post call r = await asyncio.to_thread( requests.post, - url=f"{request.app.state.config.IMAGES_GEMINI_API_BASE_URL}/models/{model}", + url=f"{request.app.state.config.IMAGES_EDIT_GEMINI_API_BASE_URL}/models/{model}", json=data, headers=headers, ) @@ -916,50 +944,77 @@ async def image_edits( return images - elif request.app.state.config.IMAGE_GENERATION_ENGINE == "comfyui": + elif request.app.state.config.IMAGE_EDIT_ENGINE == "comfyui": + try: + files = [] + if isinstance(form_data.image, str): + files = [get_image_file_item(form_data.image)] + elif isinstance(form_data.image, list): + for img in form_data.image: + files.append(get_image_file_item(img)) + + # Upload images to ComfyUI and get their names + comfyui_images = [] + for file_item in files: + res = await comfyui_upload_image( + file_item, + request.app.state.config.IMAGES_EDIT_COMFYUI_BASE_URL, + request.app.state.config.IMAGES_EDIT_COMFYUI_API_KEY, + ) + comfyui_images.append(res.get("name", file_item[1][0])) + except Exception as e: + log.debug(f"Error uploading images to ComfyUI: {e}") + raise Exception("Failed to upload images to ComfyUI.") + data = { + "image": comfyui_images, "prompt": form_data.prompt, - "width": width, - "height": height, - "n": form_data.n, + **({"width": width} if width is not None else {}), + **({"height": height} if height is not None else {}), + **({"n": form_data.n} if form_data.n else {}), } - if request.app.state.config.IMAGE_EDIT_STEPS is not None: - data["steps"] = request.app.state.config.IMAGE_EDIT_STEPS - - if form_data.negative_prompt is not None: - data["negative_prompt"] = form_data.negative_prompt - - form_data = ComfyUICreateImageForm( + form_data = ComfyUIEditImageForm( **{ "workflow": ComfyUIWorkflow( **{ - "workflow": request.app.state.config.COMFYUI_WORKFLOW, - "nodes": request.app.state.config.COMFYUI_WORKFLOW_NODES, + "workflow": request.app.state.config.IMAGES_EDIT_COMFYUI_WORKFLOW, + "nodes": request.app.state.config.IMAGES_EDIT_COMFYUI_WORKFLOW_NODES, } ), **data, } ) - res = await comfyui_create_image( + res = await comfyui_edit_image( model, form_data, user.id, - request.app.state.config.COMFYUI_BASE_URL, - request.app.state.config.COMFYUI_API_KEY, + request.app.state.config.IMAGES_EDIT_COMFYUI_BASE_URL, + request.app.state.config.IMAGES_EDIT_COMFYUI_API_KEY, ) log.debug(f"res: {res}") + image_urls = set() + for image in res["data"]: + image_urls.add(image["url"]) + image_urls = list(image_urls) + + # Prioritize output type URLs if available + output_type_urls = [url for url in image_urls if "type=output" in url] + if output_type_urls: + image_urls = output_type_urls + + log.debug(f"Image URLs: {image_urls}") images = [] - for image in res["data"]: + for image_url in image_urls: headers = None - if request.app.state.config.COMFYUI_API_KEY: + if request.app.state.config.IMAGES_EDIT_COMFYUI_API_KEY: headers = { - "Authorization": f"Bearer {request.app.state.config.COMFYUI_API_KEY}" + "Authorization": f"Bearer {request.app.state.config.IMAGES_EDIT_COMFYUI_API_KEY}" } - image_data, content_type = get_image_data(image["url"], headers) + image_data, content_type = get_image_data(image_url, headers) url = upload_image( request, image_data, @@ -968,6 +1023,7 @@ async def image_edits( user, ) images.append({"url": url}) + return images except Exception as e: error = e diff --git a/backend/open_webui/utils/images/comfyui.py b/backend/open_webui/utils/images/comfyui.py index 2bc2f584aa..506723bc92 100644 --- a/backend/open_webui/utils/images/comfyui.py +++ b/backend/open_webui/utils/images/comfyui.py @@ -2,6 +2,8 @@ import asyncio import json import logging import random +import requests +import aiohttp import urllib.parse import urllib.request from typing import Optional @@ -91,6 +93,25 @@ def get_images(ws, prompt, client_id, base_url, api_key): return {"data": output_images} +async def comfyui_upload_image(image_file_item, base_url, api_key): + url = f"{base_url}/api/upload/image" + headers = {} + + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + + _, (filename, file_bytes, mime_type) = image_file_item + + form = aiohttp.FormData() + form.add_field("image", file_bytes, filename=filename, content_type=mime_type) + form.add_field("type", "input") # required by ComfyUI + + async with aiohttp.ClientSession() as session: + async with session.post(url, data=form, headers=headers) as resp: + resp.raise_for_status() + return await resp.json() + + class ComfyUINodeInput(BaseModel): type: Optional[str] = None node_ids: list[str] = [] @@ -191,3 +212,102 @@ async def comfyui_create_image( ws.close() return images + + +class ComfyUIEditImageForm(BaseModel): + workflow: ComfyUIWorkflow + + image: str | list[str] + prompt: str + width: Optional[int] = None + height: Optional[int] = None + n: Optional[int] = None + + steps: Optional[int] = None + seed: Optional[int] = None + + +async def comfyui_edit_image( + model: str, payload: ComfyUIEditImageForm, client_id, base_url, api_key +): + ws_url = base_url.replace("http://", "ws://").replace("https://", "wss://") + workflow = json.loads(payload.workflow.workflow) + + for node in payload.workflow.nodes: + if node.type: + if node.type == "model": + for node_id in node.node_ids: + workflow[node_id]["inputs"][node.key] = model + elif node.type == "image": + if isinstance(payload.image, list): + # check if multiple images are provided + for idx, node_id in enumerate(node.node_ids): + if idx < len(payload.image): + workflow[node_id]["inputs"][node.key] = payload.image[idx] + else: + for node_id in node.node_ids: + workflow[node_id]["inputs"][node.key] = payload.image + elif node.type == "prompt": + for node_id in node.node_ids: + workflow[node_id]["inputs"][ + node.key if node.key else "text" + ] = payload.prompt + elif node.type == "negative_prompt": + for node_id in node.node_ids: + workflow[node_id]["inputs"][ + node.key if node.key else "text" + ] = payload.negative_prompt + elif node.type == "width": + for node_id in node.node_ids: + workflow[node_id]["inputs"][ + node.key if node.key else "width" + ] = payload.width + elif node.type == "height": + for node_id in node.node_ids: + workflow[node_id]["inputs"][ + node.key if node.key else "height" + ] = payload.height + elif node.type == "n": + for node_id in node.node_ids: + workflow[node_id]["inputs"][ + node.key if node.key else "batch_size" + ] = payload.n + elif node.type == "steps": + for node_id in node.node_ids: + workflow[node_id]["inputs"][ + node.key if node.key else "steps" + ] = payload.steps + elif node.type == "seed": + seed = ( + payload.seed + if payload.seed + else random.randint(0, 1125899906842624) + ) + for node_id in node.node_ids: + workflow[node_id]["inputs"][node.key] = seed + else: + for node_id in node.node_ids: + workflow[node_id]["inputs"][node.key] = node.value + + try: + ws = websocket.WebSocket() + headers = {"Authorization": f"Bearer {api_key}"} + ws.connect(f"{ws_url}/ws?clientId={client_id}", header=headers) + log.info("WebSocket connection established.") + except Exception as e: + log.exception(f"Failed to connect to WebSocket server: {e}") + return None + + try: + log.info("Sending workflow to WebSocket server.") + log.info(f"Workflow: {workflow}") + images = await asyncio.to_thread( + get_images, ws, workflow, client_id, base_url, api_key + ) + except Exception as e: + log.exception(f"Error while receiving images: {e}") + images = None + + ws.close() + + return images diff --git a/src/lib/components/admin/Settings/Images.svelte b/src/lib/components/admin/Settings/Images.svelte index a26de06513..f9be013adf 100644 --- a/src/lib/components/admin/Settings/Images.svelte +++ b/src/lib/components/admin/Settings/Images.svelte @@ -63,14 +63,19 @@ ]; let REQUIRED_EDIT_WORKFLOW_NODES = [ + { + type: 'image', + key: 'image', + node_ids: '' + }, { type: 'prompt', - key: 'text', + key: 'prompt', node_ids: '' }, { type: 'model', - key: 'ckpt_name', + key: 'unet_name', node_ids: '' }, { @@ -157,10 +162,25 @@ loading = false; return; } + + config.COMFYUI_WORKFLOW_NODES = REQUIRED_WORKFLOW_NODES.map((node) => { + return { + type: node.type, + key: node.key, + node_ids: + node.node_ids.trim() === '' ? [] : node.node_ids.split(',').map((id) => id.trim()) + }; + }); } - if (config?.COMFYUI_WORKFLOW) { - config.COMFYUI_WORKFLOW_NODES = REQUIRED_WORKFLOW_NODES.map((node) => { + if (config?.IMAGES_EDIT_COMFYUI_WORKFLOW) { + if (!validateJSON(config?.IMAGES_EDIT_COMFYUI_WORKFLOW)) { + toast.error($i18n.t('Invalid JSON format for ComfyUI Edit Workflow.')); + loading = false; + return; + } + + config.IMAGES_EDIT_COMFYUI_WORKFLOW_NODES = REQUIRED_EDIT_WORKFLOW_NODES.map((node) => { return { type: node.type, key: node.key, @@ -211,6 +231,30 @@ node_ids: typeof n.node_ids === 'string' ? n.node_ids : n.node_ids.join(',') }; }); + + if (config.IMAGES_EDIT_COMFYUI_WORKFLOW) { + try { + config.IMAGES_EDIT_COMFYUI_WORKFLOW = JSON.stringify( + JSON.parse(config.IMAGES_EDIT_COMFYUI_WORKFLOW), + null, + 2 + ); + } catch (e) { + console.error(e); + } + } + + REQUIRED_EDIT_WORKFLOW_NODES = REQUIRED_EDIT_WORKFLOW_NODES.map((node) => { + const n = + config.IMAGES_EDIT_COMFYUI_WORKFLOW_NODES.find((n) => n.type === node.type) ?? node; + console.debug(n); + + return { + type: n.type, + key: n.key, + node_ids: typeof n.node_ids === 'string' ? n.node_ids : n.node_ids.join(',') + }; + }); } }); @@ -814,7 +858,7 @@ placeholder={$i18n.t('Select Engine')} > - + @@ -835,7 +879,6 @@ class="text-right text-sm bg-transparent outline-hidden max-w-full w-52" bind:value={config.IMAGE_EDIT_MODEL} placeholder={$i18n.t('Select a model')} - required /> @@ -1086,7 +1129,7 @@
- {node.type}{node.type === 'prompt' ? '*' : ''} + {node.type}{['prompt', 'image'].includes(node.type) ? '*' : ''}