#!/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.")