diff options
| -rw-r--r-- | .gitignore | 1 | ||||
| -rw-r--r-- | comfier.py | 702 | ||||
| -rw-r--r-- | harness.py | 190 | ||||
| -rw-r--r-- | requirements.txt | 2 | ||||
| -rw-r--r-- | workflow.json | 379 |
5 files changed, 1274 insertions, 0 deletions
diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..ba0430d --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +__pycache__/
\ No newline at end of file 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) + diff --git a/harness.py b/harness.py new file mode 100644 index 0000000..f588328 --- /dev/null +++ b/harness.py @@ -0,0 +1,190 @@ +#!/usr/bin/env python3 +import argparse +import asyncio +import json +import pathlib +import sys +import types +from types import SimpleNamespace +from typing import Any + +# --------------------------------------------------------------------------- +# 1. Stub out `fastapi` and `open_webui.*` in sys.modules BEFORE importing tool +# --------------------------------------------------------------------------- +def install_open_webui_stubs( + comfyui_url: str, + api_key: str, + workflow_path: pathlib.Path, + nodes_config: list[dict[str, Any]], + image_size: str, + image_steps: int, + output_dir: pathlib.Path, +) -> None: + # Ensure fastapi.Request exists even in a minimal venv + if "fastapi" not in sys.modules: + fastapi_mod = types.ModuleType("fastapi") + fastapi_mod.Request = object + sys.modules["fastapi"] = fastapi_mod + + workflow_raw = workflow_path.read_text(encoding="utf-8") + + # Mock `open_webui.models.config.Config` + config_mod = types.ModuleType("open_webui.models.config") + + class MockConfig: + @staticmethod + async def get(key: str, default: Any = None) -> Any: + return default + + config_mod.Config = MockConfig + + # Mock `open_webui.models.users.Users` + users_mod = types.ModuleType("open_webui.models.users") + + class MockUsers: + @staticmethod + def get_user_by_id(user_id: str) -> Any: + return SimpleNamespace(id=user_id, name="Felix") + + users_mod.Users = MockUsers + + # Mock `open_webui.routers.images` + images_mod = types.ModuleType("open_webui.routers.images") + + async def mock_get_image_config() -> Any: + return SimpleNamespace( + COMFYUI_BASE_URL=comfyui_url, + COMFYUI_API_KEY=api_key, + COMFYUI_WORKFLOW=workflow_raw, + COMFYUI_WORKFLOW_NODES=nodes_config, + IMAGE_SIZE=image_size, + IMAGE_STEPS=image_steps, + ) + + def mock_upload_image( + request: Any, + image_data: bytes, + content_type: str, + metadata: dict[str, Any], + user: Any, + ) -> Any: + output_dir.mkdir(parents=True, exist_ok=True) + ext = ".png" if "png" in content_type else ".jpg" + file_id = f"seed_{metadata.get('seed', 'unknown')}_{int(asyncio.get_event_loop().time() * 1000)}" + out_file = output_dir / f"{file_id}{ext}" + out_file.write_bytes(image_data) + print(f" [Storage] Saved {len(image_data)} bytes -> {out_file}") + # Return a mock FileModel with an `.id` attribute just like Open WebUI + return SimpleNamespace(id=file_id) + + images_mod.get_image_config = mock_get_image_config + images_mod.upload_image = mock_upload_image + + # Register package hierarchy in sys.modules + sys.modules["open_webui"] = types.ModuleType("open_webui") + sys.modules["open_webui.models"] = types.ModuleType("open_webui.models") + sys.modules["open_webui.routers"] = types.ModuleType("open_webui.routers") + sys.modules["open_webui.models.config"] = config_mod + sys.modules["open_webui.models.users"] = users_mod + sys.modules["open_webui.routers.images"] = images_mod + + +# --------------------------------------------------------------------------- +# 2. Terminal Event Emitter & Interactive Event Call +# --------------------------------------------------------------------------- +async def console_event_emitter(event: dict[str, Any]) -> None: + ev_type = event.get("type") + data = event.get("data", {}) + if ev_type == "status": + done_marker = "✓" if data.get("done") else "…" + print(f" [{done_marker}] {data.get('description')}") + elif ev_type == "files": + print(f" [Files Emitted] {json.dumps(data.get('files', []))}") + else: + print(f" [Event:{ev_type}] {data}") + + +async def console_event_call(event: dict[str, Any]) -> bool: + ev_type = event.get("type") + data = event.get("data", {}) + print(f"\n--- MODAL ({ev_type}): {data.get('title')} ---") + print(data.get("message", "")) + loop = asyncio.get_running_loop() + answer = await loop.run_in_executor( + None, lambda: input("Confirm / Retry? [y/N]: ").strip().lower() + ) + return answer in ("y", "yes") + + +# --------------------------------------------------------------------------- +# 3. CLI Runner +# --------------------------------------------------------------------------- +async def main() -> None: + parser = argparse.ArgumentParser(description="Standalone test harness for Comfier") + parser.add_argument("--url", default="http://10.42.0.146:8188", help="ComfyUI base URL") + parser.add_argument("--workflow", type=pathlib.Path, required=True, help="Path to ComfyUI API workflow JSON") + parser.add_argument("--prompt", default="a tiny blue marble on a white surface", help="Positive prompt") + parser.add_argument("--neg", default=None, help="Negative prompt") + parser.add_argument("--size", default="256x256", help="Image size (WIDTHxHEIGHT)") + parser.add_argument("--seed", default="123", help="Seed") + parser.add_argument("--steps", type=int, default=30, help="Global IMAGE_STEPS") + parser.add_argument("--neg-nodes", default="459:452", help="Valve: negative_prompt_node_ids") + parser.add_argument("--cancel-after", type=float, default=None, help="Simulate user clicking Stop after N seconds") + args = parser.parse_args() + + # Match your Open WebUI global COMFYUI_WORKFLOW_NODES setup + global_nodes = [ + {"type": "prompt", "key": "value", "node_ids": ["520:506"]}, + {"type": "seed", "key": "seed", "node_ids": ["520:517"]}, + {"type": "width", "key": "value", "node_ids": ["520:518"]}, + {"type": "height", "key": "value", "node_ids": ["520:519"]}, + {"type": "steps", "key": "steps", "node_ids": ["459:458"]}, + ] + + install_open_webui_stubs( + comfyui_url=args.url, + api_key="", + workflow_path=args.workflow, + nodes_config=global_nodes, + image_size="1024x1024", + image_steps=args.steps, + output_dir=pathlib.Path("./test_out"), + ) + + # Import your tool AFTER stubs are installed + from comfier import Tools + + tools = Tools() + tools.valves.negative_prompt_node_ids = args.neg_nodes + + task = asyncio.create_task( + tools.comfier_generate_image( + positive_prompt=args.prompt, + negative_prompt=args.neg, + size=args.size, + seed=args.seed, + __request__=SimpleNamespace(), + __user__={"id": "local-dev"}, + __event_emitter__=console_event_emitter, + __event_call__=console_event_call, + ) + ) + + if args.cancel_after is not None: + await asyncio.sleep(args.cancel_after) + print(f"\n[Harness] Simulating user abort after {args.cancel_after}s...") + task.cancel() + + try: + result = await task + print("\n[Tool Return Payload]") + print(json.dumps(result, indent=2)) + except asyncio.CancelledError: + print("\n[Harness] Task cleanly cancelled.") + + +if __name__ == "__main__": + try: + asyncio.run(main()) + except KeyboardInterrupt: + print("\nInterrupted via Ctrl+C.") diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..7925b0b --- /dev/null +++ b/requirements.txt @@ -0,0 +1,2 @@ +aiohttp +pydantic
\ No newline at end of file diff --git a/workflow.json b/workflow.json new file mode 100644 index 0000000..40d24a8 --- /dev/null +++ b/workflow.json @@ -0,0 +1,379 @@ +{ + "513": { + "inputs": { + "images": [ + "459:525", + 0 + ] + }, + "class_type": "PreviewImage", + "_meta": { + "title": "Output" + } + }, + "520:506": { + "inputs": { + "value": "" + }, + "class_type": "PrimitiveStringMultiline", + "_meta": { + "title": "Prompt" + } + }, + "520:518": { + "inputs": { + "value": 1530 + }, + "class_type": "PrimitiveInt", + "_meta": { + "title": "Width" + } + }, + "520:519": { + "inputs": { + "value": 1024 + }, + "class_type": "PrimitiveInt", + "_meta": { + "title": "Height" + } + }, + "520:517": { + "inputs": { + "seed": 0 + }, + "class_type": "SeedNode", + "_meta": { + "title": "Generation Seed" + } + }, + "520:526": { + "inputs": { + "value": "" + }, + "class_type": "PrimitiveStringMultiline", + "_meta": { + "title": "Negative Prompt" + } + }, + "459:452": { + "inputs": { + "prompt": [ + "520:506", + 0 + ], + "negative_prompt": [ + "520:526", + 0 + ], + "resolution": 1024, + "clip": [ + "459:453", + 0 + ], + "vae": [ + "459:454", + 0 + ] + }, + "class_type": "TextEncodeQwenImage21", + "_meta": { + "title": "Text Encode Qwen Image 2.1" + } + }, + "459:453": { + "inputs": { + "clip_name": "qwen3vl_8b_bf16.safetensors", + "type": "qwen_image", + "device": "default" + }, + "class_type": "CLIPLoader", + "_meta": { + "title": "Load Conditioning Model" + } + }, + "459:454": { + "inputs": { + "vae_name": "qwen_image_2.1_vae_bf16.safetensors" + }, + "class_type": "VAELoader", + "_meta": { + "title": "Load VAE" + } + }, + "459:456": { + "inputs": { + "width": [ + "520:518", + 0 + ], + "height": [ + "520:519", + 0 + ], + "batch_size": 1 + }, + "class_type": "EmptyLatentImage", + "_meta": { + "title": "Empty Latent Image" + } + }, + "459:457": { + "inputs": { + "samples": [ + "459:458", + 0 + ], + "vae": [ + "459:454", + 0 + ] + }, + "class_type": "VAEDecode", + "_meta": { + "title": "VAE Decode" + } + }, + "459:458": { + "inputs": { + "seed": [ + "520:517", + 0 + ], + "steps": 25, + "cfg": 3, + "sampler_name": "euler", + "scheduler": "sgm_uniform", + "denoise": 1, + "model": [ + "459:468", + 0 + ], + "positive": [ + "459:452", + 0 + ], + "negative": [ + "459:452", + 1 + ], + "latent_image": [ + "459:456", + 0 + ] + }, + "class_type": "KSampler", + "_meta": { + "title": "KSampler" + } + }, + "459:468": { + "inputs": { + "unet_name": "qwen-image-2.1-UC-Q4_K_M.gguf" + }, + "class_type": "UnetLoaderGGUF", + "_meta": { + "title": "Load Inference Model" + } + }, + "459:472": { + "inputs": { + "model_name": "4xPurePhoto-RealPLSKR.pth" + }, + "class_type": "UpscaleModelLoader", + "_meta": { + "title": "Load Upscale Model" + } + }, + "459:473": { + "inputs": { + "upscale_model": [ + "459:472", + 0 + ], + "image": [ + "459:525", + 0 + ] + }, + "class_type": "ImageUpscaleWithModel", + "_meta": { + "title": "Upscale Image (using Model)" + } + }, + "459:488": { + "inputs": { + "filename_prefix": [ + "459:498", + 0 + ], + "format": "md", + "text": [ + "459:516", + 0 + ] + }, + "class_type": "SaveText", + "_meta": { + "title": "Save Prompt" + } + }, + "459:494": { + "inputs": { + "format_string": "%Y-%m-%d_%H-%M-%S", + "datetime": [ + "459:495", + 0 + ] + }, + "class_type": "Basic data handling: TimeFormat", + "_meta": { + "title": "Format Time String" + } + }, + "459:495": { + "inputs": { + "trigger": [ + "459:458", + 0 + ] + }, + "class_type": "Basic data handling: TimeNow", + "_meta": { + "title": "Time Now" + } + }, + "459:497": { + "inputs": { + "f_string": "ai.arknet.ch/generated-image_{a}", + "values.a": [ + "459:494", + 0 + ] + }, + "class_type": "StringFormat", + "_meta": { + "title": "Base Name" + } + }, + "459:498": { + "inputs": { + "f_string": "{a}_prompt", + "values.a": [ + "459:497", + 0 + ] + }, + "class_type": "StringFormat", + "_meta": { + "title": "Prompt Name" + } + }, + "459:514": { + "inputs": { + "sub_directory": "ai.arknet.ch", + "filename_text_1": "generated-image", + "filename_text_2": "base", + "filename_text_3": [ + "459:494", + 0 + ], + "filename_separator": "_", + "timestamp": "false", + "counter_type": "filename", + "filename_text_1_pos": 0, + "filename_text_2_pos": 2, + "filename_text_3_pos": 1, + "timestamp_pos": 4, + "timestamp_type": "job", + "counter_pos": 3, + "extra_metadata": "", + "images": [ + "459:525", + 0 + ] + }, + "class_type": "Save Images No Display", + "_meta": { + "title": "Save Base Image" + } + }, + "459:515": { + "inputs": { + "sub_directory": "ai.arknet.ch", + "filename_text_1": "generated-image", + "filename_text_2": "upscaled", + "filename_text_3": [ + "459:494", + 0 + ], + "filename_separator": "_", + "timestamp": "false", + "counter_type": "filename", + "filename_text_1_pos": 0, + "filename_text_2_pos": 2, + "filename_text_3_pos": 1, + "timestamp_pos": 4, + "timestamp_type": "job", + "counter_pos": 3, + "extra_metadata": "", + "images": [ + "459:523", + 0 + ] + }, + "class_type": "Save Images No Display", + "_meta": { + "title": "Save Upscaled Image" + } + }, + "459:516": { + "inputs": { + "f_string": "# Positive Prompt\n\n{a}\n\n---\n\n# Negative Prompt\n\n{b}\n\n---\n\nSEED: {c}", + "values.a": [ + "520:506", + 0 + ], + "values.b": [ + "520:526", + 0 + ], + "values.c": [ + "520:517", + 0 + ] + }, + "class_type": "StringFormat", + "_meta": { + "title": "Format Text" + } + }, + "459:523": { + "inputs": { + "value": [ + "459:473", + 0 + ], + "model": [ + "459:472", + 0 + ] + }, + "class_type": "UnloadModel", + "_meta": { + "title": "UnloadModel" + } + }, + "459:525": { + "inputs": { + "value": [ + "459:457", + 0 + ] + }, + "class_type": "UnloadAllModels", + "_meta": { + "title": "UnloadAllModels" + } + } +} |