path: root/comfier.py

"""
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)