Emil_Shanaty/open-webui-ai4me
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: comfyui image edit support

Browse Source
This commit is contained in:
Timothy Jaeryang Baek
2025-11-06 03:43:59 -05:00
parent 74db2b9f36
commit 63e8ab7a05
5 changed files with 296 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/open_webui/config.py
+23
View File
@@ -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
####################################
backend/open_webui/main.py
+8
View File
@@ -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
########################################
backend/open_webui/routers/images.py
+95 -39
View File
@@ -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
backend/open_webui/utils/images/comfyui.py
+120
View File
@@ -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