aboutsummaryrefslogtreecommitdiff
path: root/harness.py
blob: 061ec02881ccabda320c26d1fbe037fbe62d1201 (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
123
124
125
126
127
128
129
130
131
132
133
134
#!/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=None,
        help="Tool argument: sampling steps (default: omitted, use --default-steps)",
    )
    parser.add_argument(
        "--default-steps",
        type=int,
        default=30,
        help="Open WebUI's global IMAGE_STEPS, used when --steps is omitted",
    )
    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.default_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,
            steps=args.steps,
            __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.")