diff options
Diffstat (limited to 'comfier.py')
| -rw-r--r-- | comfier.py | 702 |
1 files changed, 702 insertions, 0 deletions
diff --git a/comfier.py b/comfier.py new file mode 100644 index 0000000..a97fd0d --- /dev/null +++ b/comfier.py @@ -0,0 +1,702 @@ +""" +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 aiohttp +import asyncio +import copy +import inspect +import json +import mimetypes +import random +import time +import uuid + +from fastapi import Request +from pydantic import BaseModel, Field +from typing import Annotated, Any, Awaitable, Callable, Optional + +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 + +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 win 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)." + ) + 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 beend 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 = 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 hasattr(uploaded, "id"): + url = f"/api/v1/files/{uploaded.id}/content" + elif isinstance(uploaded, dict) and "id" in uploaded: + url = f"/api/v1/files/{uploaded['id']}/content" + 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) + |