""" title: Comfier author: Felix Morgner author_url: https://felixmorgner.com description: A ComfyUI image generation tool extending the built-in tool by allowing for more parameters. required_open_webui_version: 0.12.0 version: 0.1.0 license: 3-clause-BSD """ import asyncio import copy import inspect import json import mimetypes import random import time import uuid from typing import Annotated, Any, Awaitable, Callable, Optional import aiohttp from fastapi import Request from open_webui.models.config import Config from open_webui.models.users import Users from open_webui.routers.images import get_image_config, upload_image from pydantic import BaseModel, Field EventEmitter = Callable[[dict[str, Any]], Awaitable[None]] EventCall = Callable[[dict[str, Any]], Awaitable[Any]] class Tools: class Valves(BaseModel): negative_prompt: str = Field( default="watermarks", description="Tool level value for the negative prompt.", ) negative_prompt_key: str = Field( default="negative_prompt", description="Tool-level input key name for the negative prompt (e.g. 'negative_prompt' or 'text').", ) negative_prompt_node_ids: str = Field( default="", description="Tool-level node IDs for negative prompt targets, comma-separated (e.g. '459:452'). Leave empty to disable negative prompting.", ) def __init__(self): self.valves = self.Valves() async def comfier_generate_seed( self, bits: Annotated[int, Field(ge=8, le=63)] = 63 ) -> str: """ Generate a new seed for image generation. This seed can then be used for later calls of comfier_generate_image. :param bits: The maximum bit-width of the generated seed (between 8 and 63). :return: The generated seed. """ clamped = max(8, min(int(bits), 63)) return f"{random.getrandbits(clamped)}" async def comfier_generate_image( self, positive_prompt: str, negative_prompt: Optional[str] = None, size: Optional[Annotated[str, Field(pattern=r"^\d+[xX]\d+$")]] = None, seed: Optional[int | str] = None, __request__: Optional[Request] = None, __user__: Optional[dict[str, Any]] = None, __event_emitter__: Optional[EventEmitter] = None, __event_call__: Optional[EventCall] = None, ) -> dict[str, Any]: """ Generate an image from a text prompt using an optional reproducible seed. Additionally, this tool allows for the specification of a 'negative prompt'. :param positive_prompt: The descriptive text prompt for the image to generate. :param negative_prompt: Optional descriptive text prompt of what to avoid generating in the image. :param size: Optional image dimensions in 'WIDTHxHEIGHT' format (e.g. '1024x768'). Default to the global Open WebUI size. :param seed: Optional seed for reproducible generation. :return: Generation status. """ if __request__ is None: raise RuntimeError( "FastAPI Request context (__request__) was not injected." ) if not __user__ or "id" not in __user__: raise RuntimeError("User context (__user__) was not injected.") normalized_seed = self._normalize_seed(seed) workflow = await self._load_workflow() global_mappings = await self._load_global_node_mappings() effective_mappings = self._resolve_node_mappings(global_mappings) effective_negative_prompt = ( negative_prompt if negative_prompt is not None else self.valves.negative_prompt ) width, height = await self._resolve_dimensions(size) steps = await self._resolve_steps() parameter_values: dict[str, Any] = { "positive_prompt": positive_prompt, "negative_prompt": effective_negative_prompt, "seed": normalized_seed, "width": width, "height": height, "steps": steps, } self._apply_workflow_parameters( workflow=workflow, mappings=effective_mappings, values=parameter_values, ) user = Users.get_user_by_id(__user__["id"]) if inspect.isawaitable(user): user = await user while True: try: image_urls = await self._execute( workflow, request=__request__, user=user, metadata={ "prompt": positive_prompt, "negative_prompt": effective_negative_prompt, "seed": str(normalized_seed), "width": width, "height": height, "steps": steps, }, event_emitter=__event_emitter__, ) break except (ConnectionError, aiohttp.ClientError) as err: if __event_call__ is None: raise await self._emit_status( __event_emitter__, f"ComfyUI connection lost: {err}. Waiting for retry confirmation...", done=False, ) should_retry = await __event_call__( { "type": "confirmation", "data": { "title": "ComfyUI Backend Disconnected", "message": ( f"Lost connection to ComfyUI during generation ({err}).\n\n" f"Would you like to retry with the same parameters (seed {normalized_seed}, {width}x{height})?" ), }, } ) if not should_retry: await self._emit_status( __event_emitter__, f"Generation failed: {err}", done=True, ) raise if __event_emitter__ is not None and image_urls: await __event_emitter__( { "type": "files", "data": { "files": [{"type": "image", "url": url} for url in image_urls] }, } ) return { "status": "success", "seed": str(normalized_seed), "message": "The image has been successfully generated and is already visible to the user in the chat. You do not need to display or embed the image again - just acknowledge that it has been created.", "images": image_urls, } def _apply_workflow_parameters( self, workflow: dict[str, Any], mappings: dict[str, dict[str, Any]], values: dict[str, Any], ) -> dict[str, list[str]]: """Inject the resolved parameter values into the workflow dictionary in-place.""" applied: dict[str, list[str]] = {} for param_name, value in values.items(): mapping = mappings.get(param_name) if not mapping or value is None: continue input_key = mapping.get("key") node_ids = mapping.get("node_ids", []) if not input_key or not node_ids: continue updated_targets: list[str] = [] for node_id in node_ids: node = workflow.get(node_id) if not isinstance(node, dict) or not isinstance( node.get("inputs"), dict ): raise KeyError( f"Mapped node {node_id!r} for {param_name!r} does not exist in COMFYUI_WORKFLOW." ) node["inputs"][input_key] = value updated_targets.append(f"{node_id}.{input_key}") if updated_targets: applied[param_name] = updated_targets return applied async def _cancel( self, http_base: str, headers: dict[str, str], prompt_id: Optional[str] ) -> None: """Try to cancel a ComfyUI request.""" if not prompt_id: return try: timeout = aiohttp.ClientTimeout(total=3.0) async with aiohttp.ClientSession( headers=headers, timeout=timeout ) as cancel_session: await cancel_session.post( f"{http_base}/queue", json={"delete": [prompt_id]} ) await cancel_session.post( f"{http_base}/interrupt", json={"prompt_id": prompt_id} ) except Exception: pass async def _verify_prompt_alive( self, session: aiohttp.ClientSession, http_base: str, prompt_id: str, ) -> bool: """ Verify with ComfyUI that the backend is reachable and still tracking prompt_id in either /queue (running/pending) or /history. """ timeout = aiohttp.ClientTimeout(total=3.0) try: async with session.get(f"{http_base}/queue", timeout=timeout) as q_resp: q_resp.raise_for_status() q_data = await q_resp.json() running = q_data.get("queue_running") or [] pending = q_data.get("queue_pending") or [] in_queue = any( isinstance(item, list) and len(item) > 1 and item[1] == prompt_id for item in (running + pending) ) if in_queue: return True async with session.get( f"{http_base}/history/{prompt_id}", timeout=timeout ) as h_resp: h_resp.raise_for_status() h_data = await h_resp.json() or {} if prompt_id not in h_data: return False status_info = (h_data[prompt_id] or {}).get("status") or {} if status_info.get("status_str") == "error": raise RuntimeError( "The job was terminated with an error or interruption on the backend." ) return True except aiohttp.ClientError as err: raise ConnectionError( f"Lost connection to ComfyUI backend ({err})." ) from err async def _emit_status( self, event_emitter: Optional[EventEmitter], description: str, done: bool = False, ) -> None: """Send a live status update banner to the Open WebUI frontend.""" if event_emitter is None: return await event_emitter( { "type": "status", "data": { "description": description, "done": done, "hidden": False, }, } ) async def _execute( self, workflow: dict[str, Any], request: Request, user: Any, metadata: dict[str, Any], event_emitter: Optional[EventEmitter], ) -> list[str]: """Submit the workflow to ComfyUI, stream progress, and store the generated images.""" http_base, api_key = await self._load_connection_information() ws_base = http_base.replace("https://", "wss://").replace("http://", "ws://") headers = {"Authorization": f"Bearer {api_key}"} if api_key else {} client_id = str(uuid.uuid4()) output_images: list[dict[str, Any]] = [] stored_urls: list[str] = [] last_emit_time = 0.0 prompt_id: Optional[str] = None idle_ticks = 0 try: async with aiohttp.ClientSession(headers=headers) as session: async with session.ws_connect( f"{ws_base}/ws?clientId={client_id}", heartbeat=5.0, ) as ws: await self._emit_status( event_emitter, "Submitting workflow to ComfyUI...", done=False ) async with session.post( f"{http_base}/prompt", json={"prompt": workflow, "client_id": client_id}, ) as response: if response.status >= 400: error_body = await response.text() raise RuntimeError( f"ComfyUI rejected workflow (HTTP {response.status}): {error_body}" ) queued = await response.json() prompt_id = queued["prompt_id"] await self._emit_status( event_emitter, "Waiting for ComfyUI queue...", done=False ) while True: try: msg = await asyncio.wait_for(ws.receive(), timeout=1.0) idle_ticks = 0 except asyncio.TimeoutError: idle_ticks += 1 if idle_ticks >= 5 and prompt_id: idle_ticks = 0 is_alive = await self._verify_prompt_alive( session, http_base, prompt_id ) if not is_alive: raise ConnectionError( f"ComfyUI backend lost prompt {prompt_id} (backend may have restarted or crashed)." ) from None continue if msg.type != aiohttp.WSMsgType.TEXT: if msg.type in ( aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.CLOSING, aiohttp.WSMsgType.CLOSED, aiohttp.WSMsgType.ERROR, ): ws_err = ws.exception() detail = ( f": {ws_err}" if ws_err else f" (code {ws.close_code})" ) raise ConnectionError( f"ComfyUI WebSocket disconnected unexpectedly before completion{detail}." ) continue payload = json.loads(msg.data) or {} msg_type = payload.get("type") data = payload.get("data") or {} if data.get("prompt_id") and data.get("prompt_id") != prompt_id: continue if msg_type == "executing": node_id = data.get("node") if node_id is None: break node_info = workflow.get(str(node_id)) or {} node_meta = node_info.get("_meta") or {} node_title = ( node_meta.get("title") or node_info.get("class_type") or f"Node {node_id}" ) await self._emit_status( event_emitter, f"Running: {node_title}...", done=False ) elif msg_type == "progress": val = data.get("value", 0) maximum = data.get("max", 1) percentage = int((val / maximum) * 100) if maximum else 0 now = time.monotonic() if (now - last_emit_time) >= 0.25 or val == maximum: last_emit_time = now await self._emit_status( event_emitter, f"Sampling: step {val}/{maximum} ({percentage}%)", done=False, ) elif msg_type == "executed": output_data = data.get("output") or {} images = output_data.get("images") or [] if isinstance(images, list): output_images.extend( img for img in images if isinstance(img, dict) ) elif msg_type == "execution_interrupted": raise RuntimeError( "Generation was interrupted on the backend." ) elif msg_type == "execution_error": error_message = data.get( "exception_message", "Unknown ComfyUI error" ) raise RuntimeError( f"ComfyUI execution error: {error_message}" ) if not output_images: raise RuntimeError( "Generation finished without producing any output. The Job may have been cancelled or removed from the queue." ) await self._emit_status( event_emitter, "Saving generated images...", done=False ) for img in output_images: filename = img.get("filename") if not filename: continue params = { "filename": filename, "subfolder": img.get("subfolder", ""), "type": img.get("type", "output"), } async with session.get( f"{http_base}/view", params=params ) as img_response: img_response.raise_for_status() image_bytes = await img_response.read() content_type = ( img_response.headers.get("Content-Type") or mimetypes.guess_type(filename)[0] or "image/png" ) uploaded: Any = upload_image( request=request, image_data=image_bytes, content_type=content_type, metadata=metadata, user=user, ) if inspect.isawaitable(uploaded): uploaded = await uploaded if isinstance(uploaded, tuple): uploaded = uploaded[-1] if isinstance(uploaded, str): url = uploaded elif isinstance(uploaded, dict) and "id" in uploaded: url = f"/api/v1/files/{uploaded['id']}/content" elif hasattr(uploaded, "id"): url = f"/api/v1/files/{uploaded.id}/content" # pyright: ignore[reportAttributeAccessIssue] else: raise RuntimeError( f"Unexpected return type from upload_image: {type(uploaded)!r}" ) stored_urls.append(url) except asyncio.CancelledError: await asyncio.shield(self._cancel(http_base, headers, prompt_id)) await asyncio.shield( self._emit_status( event_emitter, "Generation aborted by user", done=True ) ) raise except Exception as err: await self._emit_status( event_emitter, f"Generation failed: {err}", done=True ) raise await self._emit_status(event_emitter, "Generation complete", done=True) return stored_urls async def _load_connection_information(self) -> tuple[str, str]: """Load COMFYUI_BASE_URL and COMFYUI_API_KEY from Open WebUI config.""" image_cfg = await get_image_config() base_url = getattr(image_cfg, "COMFYUI_BASE_URL", None) api_key = getattr(image_cfg, "COMFYUI_API_KEY", None) if not base_url: base_url = await Config.get("image_generation.comfyui.base_url", "") if api_key is None: api_key = await Config.get("image_generation.comfyui.api_key", "") base_url = str(base_url or "").strip().rstrip("/") if not base_url: raise RuntimeError("COMFYUI_BASE_URL is not configured in Open WebUI.") return base_url, str(api_key or "").strip() async def _load_global_node_mappings(self) -> list[dict[str, Any]]: """Load the globally configured node mappings (COMFYUI_WORKFLOW_NODES) list.""" image_cfg = await get_image_config() nodes = getattr(image_cfg, "COMFYUI_WORKFLOW_NODES", None) if nodes is None: nodes = await Config.get("image_generation.comfyui.nodes", []) return nodes if isinstance(nodes, list) else [] async def _load_workflow(self) -> dict[str, Any]: """Load and parse the globally configured ComfyUI workflow JSON.""" image_cfg = await get_image_config() raw_workflow = getattr(image_cfg, "COMFYUI_WORKFLOW", None) if not raw_workflow: raw_workflow = await Config.get("image_generation.comfyui.workflow", "") if isinstance(raw_workflow, dict): return copy.deepcopy(raw_workflow) if not isinstance(raw_workflow, str) or not raw_workflow.strip(): raise RuntimeError( "No global COMFYUI_WORKFLOW is configured in Open WebUI." ) try: workflow = json.loads(raw_workflow) except json.JSONDecodeError as err: raise RuntimeError( "Configured COMFYUI_WORKFLOW is not valid JSON." ) from err if not isinstance(workflow, dict): raise RuntimeError( "Configured COMFYUI_WORKFLOW root must be a JSON object." ) return workflow def _normalize_seed(self, seed: Optional[int | str]) -> int: """Normalize an optional int or string seed into a valid 63-bit positive integer.""" if seed is None: return random.getrandbits(63) if isinstance(seed, str): stripped = seed.strip() if not stripped: return random.getrandbits(63) try: parsed = int(stripped, 0) except ValueError as err: raise ValueError(f"Invalid seed string: {seed!r}") from err else: parsed = int(seed) if parsed < 0: return random.getrandbits(63) return parsed & ((1 << 63) - 1) async def _resolve_dimensions(self, size: Optional[str]) -> tuple[int, int]: """Parse 'WIDTHxHEIGHT' from the size parameter or fall back to global IMAGE_SIZE.""" raw_size = (size or "").strip() if not raw_size: image_cfg = await get_image_config() raw_size = getattr(image_cfg, "IMAGE_SIZE", None) if not raw_size: raw_size = await Config.get("image_generation.size", "1024x1024") try: width_string, height_string = str(raw_size).lower().split("x", 1) width, height = int(width_string.strip()), int(height_string.strip()) except (ValueError, AttributeError) as err: raise ValueError( f"Invalid image size {raw_size!r}. Expected 'WIDTHxHEIGHT' (e.g. '1024x768')." ) from err if width <= 0 or height <= 0: raise ValueError(f"Image dimensions must be positive: {width}x{height}") return width, height def _resolve_node_mappings( self, global_nodes: list[dict[str, Any]] ) -> dict[str, dict[str, Any]]: """Combine global prompt/seed mappings with the tool's negative_prompt Valve.""" resolved: dict[str, dict[str, Any]] = { "positive_prompt": {"key": "text", "node_ids": []}, "negative_prompt": { "key": self.valves.negative_prompt_key.strip() or "negative_prompt", "node_ids": [ nid.strip() for nid in self.valves.negative_prompt_node_ids.split(",") if nid.strip() ], }, "seed": {"key": "seed", "node_ids": []}, "width": {"key": "width", "node_ids": []}, "height": {"key": "height", "node_ids": []}, "steps": {"key": "steps", "node_ids": []}, } type_map = { "prompt": "positive_prompt", "positive_prompt": "positive_prompt", "seed": "seed", "width": "width", "height": "height", "steps": "steps", } for entry in global_nodes: if not isinstance(entry, dict): continue target = type_map.get(entry.get("type", "")) if not target: continue if entry.get("key"): resolved[target]["key"] = str(entry["key"]).strip() if isinstance(entry.get("node_ids"), list): resolved[target]["node_ids"] = [ str(nid).strip() for nid in entry["node_ids"] if str(nid).strip() ] return resolved async def _resolve_steps(self) -> int: """Load the global IMAGE_STEPS setting from Open WebUI config.""" image_cfg = await get_image_config() raw_steps = getattr(image_cfg, "IMAGE_STEPS", None) if raw_steps is None: raw_steps = await Config.get("image_generation.steps", 25) return int(raw_steps)