aboutsummaryrefslogtreecommitdiff
path: root/harness.py
blob: 450df24ed7299a714f16d62931fdfd98ccfe70d6 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
#!/usr/bin/env python3
"""Standalone CLI harness: run the Comfier tool against a real ComfyUI."""

import argparse
import asyncio
import json
import os
import pathlib
from types import SimpleNamespace
from typing import Any, cast

from devkit import FakeEnv, install

FIXTURES = pathlib.Path(__file__).parent / "fixtures"


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:
    data = event.get("data", {})
    print(f"\n--- MODAL ({event.get('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")


async def main() -> None:
    parser = argparse.ArgumentParser(description="Standalone test harness for Comfier")
    parser.add_argument(
        "--url",
        default=os.environ.get("COMFYUI_URL", "http://127.0.0.1:8188"),
        help="ComfyUI base URL (default: $COMFYUI_URL or http://127.0.0.1:8188)",
    )
    parser.add_argument("--api-key", default=os.environ.get("COMFYUI_API_KEY", ""))
    parser.add_argument(
        "--workflow", type=pathlib.Path, default=FIXTURES / "workflow.json"
    )
    parser.add_argument(
        "--nodes",
        type=pathlib.Path,
        default=FIXTURES / "nodes.json",
        help="JSON file mirroring Open WebUI's COMFYUI_WORKFLOW_NODES",
    )
    parser.add_argument("--prompt", default="a tiny blue marble on a white surface")
    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")
    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("--out", type=pathlib.Path, default=pathlib.Path("test_out"))
    parser.add_argument(
        "--cancel-after",
        type=float,
        default=None,
        help="Simulate user clicking Stop after N seconds",
    )
    args = parser.parse_args()

    install(
        FakeEnv(
            comfyui_url=args.url,
            api_key=args.api_key,
            workflow=args.workflow.read_text(encoding="utf-8"),
            nodes=json.loads(args.nodes.read_text(encoding="utf-8")),
            image_size="1024x1024",
            image_steps=args.steps,
            output_dir=args.out,
        )
    )

    # Import the tool AFTER the fake Open WebUI modules 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__=cast(Any, 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.")