diff --git a/__init__.py b/__init__.py index 2a6478c..e195d39 100644 --- a/__init__.py +++ b/__init__.py @@ -14,6 +14,7 @@ from .config_builder_node import UltimateConfigBuilder from .json_text_node import SmartJSONTextNode from .metadata_packer import pack_metadata_into_image from .directory_scanner import scan_directory_for_images +from .workflow_sweep_collect import USCGSweepCollect import sys # Register a sys.modules alias for cross-companion imports. @@ -1120,6 +1121,79 @@ async def cancel_upscale(request): return web.json_response({"error": str(e)}, status=500) +# ============================================================================= +# API: WORKFLOW SWEEP (sweep any API-format workflow through ComfyUI's own executor) +# JSON bodies only, no path parameters: workflow JSON is request-body data, never a filesystem path. +# ============================================================================= + +@server.PromptServer.instance.routes.post("/config_tester/sweep/parse") +async def sweep_parse(request): + """Describe an API-format workflow for the Builder (nodes, outputs, image sources, suggested bindings).""" + try: + data = await request.json() + from .workflow_sweep import parse_for_ui + return web.json_response(parse_for_ui(data.get("api_json"))) + except ValueError as e: + return web.json_response({"ok": False, "errors": [str(e)], "error": str(e)}, status=400) + except Exception as e: + print(f"[ConfigTester] Error parsing workflow: {e}") + return web.json_response({"ok": False, "errors": [str(e)], "error": str(e)}, status=500) + +@server.PromptServer.instance.routes.post("/config_tester/sweep/dry_run") +async def sweep_dry_run(request): + """First N concrete prompts (with collector) + validation of job 1. Nothing is queued.""" + try: + data = await request.json() + from .workflow_sweep_runner import dry_run + return web.json_response(await dry_run(data)) + except ValueError as e: + return web.json_response({"error": str(e)}, status=400) + except Exception as e: + print(f"[ConfigTester] Error in sweep dry run: {e}") + return web.json_response({"error": str(e)}, status=500) + +@server.PromptServer.instance.routes.post("/config_tester/sweep/start") +async def sweep_start(request): + """Start a workflow sweep from the Builder (no Generator node needed; also the resume path).""" + try: + data = await request.json() + from .workflow_sweep_runner import api_start + status, body = await api_start(data) + return web.json_response(body, status=status) + except Exception as e: + print(f"[ConfigTester] Error starting sweep: {e}") + return web.json_response({"error": str(e)}, status=500) + +@server.PromptServer.instance.routes.get("/config_tester/sweep/status") +async def sweep_status(request): + from .workflow_sweep_runner import get_status + status = get_status(request.query.get("sweep_id") or None, request.query.get("session_name") or None) + if status is None: + return web.json_response({"error": "No such sweep"}, status=404) + return web.json_response(status) + +async def _sweep_control(request, action): + try: + data = await request.json() + from . import workflow_sweep_runner as runner + ok, state = getattr(runner, action)(data.get("sweep_id", "")) + return web.json_response({"ok": ok, "state": state}, status=200 if ok else 409) + except Exception as e: + return web.json_response({"error": str(e)}, status=500) + +@server.PromptServer.instance.routes.post("/config_tester/sweep/pause") +async def sweep_pause(request): + return await _sweep_control(request, "pause") + +@server.PromptServer.instance.routes.post("/config_tester/sweep/resume") +async def sweep_resume(request): + return await _sweep_control(request, "resume") + +@server.PromptServer.instance.routes.post("/config_tester/sweep/cancel") +async def sweep_cancel(request): + return await _sweep_control(request, "cancel") + + def _create_image_links(source_dir, link_dir, items): """ Create symlinks (or copies on Windows) from source images and videos into @@ -1173,7 +1247,8 @@ NODE_CLASS_MAPPINGS = { "UltimateSamplerGrid": SamplerGridTester, "UltimateGridDashboard": SamplerConfigDashboardViewer, "UltimateConfigBuilder": UltimateConfigBuilder, - "SmartJSONText": SmartJSONTextNode + "SmartJSONText": SmartJSONTextNode, + "USCGSweepCollect": USCGSweepCollect, } NODE_DISPLAY_NAME_MAPPINGS = { @@ -1181,6 +1256,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "UltimateGridDashboard": "Ultimate Grid Dashboard (Viewer)", "UltimateConfigBuilder": "Ultimate Config Builder", "SmartJSONText": "Smart JSON Text", + "USCGSweepCollect": "USCG Sweep Collect", } WEB_DIRECTORY = "./web" \ No newline at end of file diff --git a/manifest_utils.py b/manifest_utils.py index 23b27d8..8fedfdd 100644 --- a/manifest_utils.py +++ b/manifest_utils.py @@ -5,6 +5,7 @@ Handles saving and merging manifest.json with user modifications import json import os +import re import threading # Thread-safe lock for manifest file writes. @@ -123,9 +124,8 @@ def load_existing_manifest(manifest_path): def _item_file_path(manifest_path, item): """Local path of an item's image, derived from its /view URL, confined to /images.""" try: - from urllib.parse import urlparse, parse_qs - q = parse_qs(urlparse(item.get("file", "")).query) - filename = os.path.basename((q.get("filename") or [""])[0]) + m = re.search(r"[?&]filename=([^&]+)", str(item.get("file", ""))) + filename = os.path.basename(m.group(1)) if m else "" if not filename: return None images_dir = os.path.join(os.path.dirname(os.path.abspath(manifest_path)), "images") diff --git a/sampler_node.py b/sampler_node.py index 1f1bdcd..ff413c5 100644 --- a/sampler_node.py +++ b/sampler_node.py @@ -46,6 +46,20 @@ except ImportError: from .html_generator import get_html_template +def _workflow_configs(configs_json): + """The configs list when configs_json holds Workflow Sweep entries (model_type "workflow"), else None.""" + if not isinstance(configs_json, str) or "workflow" not in configs_json: + return None + try: + parsed = json.loads(configs_json) + except (TypeError, ValueError): + return None + configs = parsed.get("configs") if isinstance(parsed, dict) else parsed + if isinstance(configs, list) and any(isinstance(c, dict) and c.get("model_type") == "workflow" for c in configs): + return configs + return None + + class SamplerGridTester: @classmethod def INPUT_TYPES(s): @@ -96,6 +110,11 @@ class SamplerGridTester: "optional_positive", "optional_negative", "optional_latent" ] + # Workflow Sweep configs: a re-queue must restart/resume the sweep (finished jobs are skipped by + # fingerprint), never hit ComfyUI's node cache + if _workflow_configs(kwargs.get("configs_json")): + return float("NaN") + # Check if any optional input is connected (not None and present in kwargs) has_optional = any( kwargs.get(key) is not None for key in optional_keys @@ -327,6 +346,29 @@ class SamplerGridTester: except (TypeError, ValueError): pass + # Workflow Sweep: hand the template + bindings to the sweep runner and return the dashboard + # immediately. The sweep prompts go through ComfyUI's own queue AFTER this prompt finishes, so + # this node never waits on them (a node waiting on a prompt it queued would deadlock the executor). + wf_configs = _workflow_configs(configs_json) + if wf_configs is not None: + if any(x is not None for x in (optional_model, optional_clip, optional_vae, + optional_positive, optional_negative, optional_latent)): + print("[GridTester] ⚠️ Workflow Sweep ignores the optional_* inputs (the workflow has its own loaders)") + from .workflow_sweep_runner import start_sweep, SweepBusy + from .workflow_sweep import read_manifest_items, session_paths + settings = dict(session_settings or {}) + settings["session_name"] = session_name + settings["overwrite_existing"] = overwrite_existing + settings["add_random_seeds_to_gens"] = add_random_seeds_to_gens + try: + sweep = start_sweep(wf_configs, settings, origin_node=unique_id) + except SweepBusy as e: + raise RuntimeError(str(e)) + manifest = read_manifest_items(session_paths(sweep.session)["manifest"]) + manifest.setdefault("session_name", sweep.session) + manifest.setdefault("meta", {}) + return (get_html_template(sweep.session, manifest, unique_id),) + return run_generation_loop( self, ckpt_name, positive_text, negative_text, seed, denoise, vae_batch_size, diff --git a/tests/test_workflow_sweep_collector_node.py b/tests/test_workflow_sweep_collector_node.py new file mode 100644 index 0000000..8974113 --- /dev/null +++ b/tests/test_workflow_sweep_collector_node.py @@ -0,0 +1,191 @@ +"""USCGSweepCollect node: saves images + manifest items inside the benchmarks folder, notifies the runner.""" +import functools +import importlib.util +import json +import os +import sys +import time + +import numpy as np +import pytest +from PIL import Image + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +PKG = os.path.basename(ROOT) +sys.path.insert(0, ROOT) + +import workflow_sweep as wfs # noqa: E402 +import workflow_sweep_collect as collect_mod # noqa: E402 +import workflow_sweep_runner as runner # noqa: E402 + + +def _load_real(name, private): + spec = importlib.util.spec_from_file_location(f"{PKG}.{private}", os.path.join(ROOT, f"{name}.py")) + mod = importlib.util.module_from_spec(spec) + mod.__package__ = PKG + spec.loader.exec_module(mod) + return mod + + +class FakeTensor: + """Quacks like a torch IMAGE batch [B,H,W,C] float 0..1 (cpu().numpy()).""" + + def __init__(self, arr): + self.arr = arr + + def __len__(self): + return len(self.arr) + + def __getitem__(self, i): + return FakeTensor(self.arr[i]) + + def cpu(self): + return self + + def numpy(self): + return self.arr + + +@pytest.fixture +def env(tmp_path, monkeypatch): + import folder_paths + monkeypatch.setattr(folder_paths, "get_output_directory", lambda: str(tmp_path), raising=False) + mu = _load_real("manifest_utils", "manifest_utils_real2") + monkeypatch.setattr(collect_mod, "_append_items", mu.append_items) + torch = sys.modules.get("torch") + patched = torch is not None and not hasattr(torch, "inference_mode") + if patched: + torch.inference_mode = lambda *a, **k: (lambda f: f) + try: + cim = _load_real("image_generation", "image_generation_real2").create_image_metadata + finally: + if patched: + del torch.inference_mode + real_build = wfs.build_collect_item + monkeypatch.setattr(wfs, "build_collect_item", functools.partial(real_build, create_image_metadata=cim)) + runner.reset_registry() + yield tmp_path + runner.reset_registry() + + +def images(n=1, h=16, w=24): + return FakeTensor(np.random.rand(n, h, w, 3).astype(np.float32)) + + +CONF = {"model_type": "workflow", "model": "m", "sampler": "euler", "scheduler": "simple", "steps": 20, + "cfg": 4.0, "seed": 3, "denoise": 1.0, "positive": "a teapot", "negative": "", "wf_fingerprint": "ab" * 10} + + +def job_json(**extra): + d = {"conf": CONF, "sweep_id": "sw1", "job_index": 0, "gen_index": 7, "origin_node": "12", "image_format": "webp"} + d.update(extra) + return json.dumps(d) + + +def test_node_definition_matches_plan(): + cls = collect_mod.USCGSweepCollect + spec = cls.INPUT_TYPES() + assert spec["required"]["images"] == ("IMAGE",) and "session_name" in spec["required"] + assert spec["required"]["job_json"][1]["multiline"] is True + assert spec["hidden"] == {"prompt": "PROMPT", "unique_id": "UNIQUE_ID"} + assert cls.RETURN_TYPES == () and cls.OUTPUT_NODE is True and cls.FUNCTION == "collect" + assert cls.IS_CHANGED() != cls.IS_CHANGED() # NaN: never cached + assert collect_mod.NODE_CLASS_MAPPINGS["USCGSweepCollect"] is cls + + +def test_collect_saves_image_manifest_item_and_ui_output(env): + out = collect_mod.USCGSweepCollect().collect(images(2), "mysession", job_json(), prompt={"1": {}}, unique_id="40") + base = env / "benchmarks" / "mysession" + manifest = json.loads((base / "manifest.json").read_text()) + assert len(manifest["items"]) == 2 and len(out["ui"]["images"]) == 2 + item = manifest["items"][0] + assert item["model_type"] == "workflow" and item["width"] == 24 and item["height"] == 16 # from the tensor + assert item["gen_index"] == 7 and item["seed"] == 3 and item["denoise"] == 1 + assert item["wf_sweep_id"] == "sw1" and item["wf_job_index"] == 0 + assert item["favorited"] is False and item["rejected"] is False and item["note"] == "" + assert item["file"].startswith("/view?filename=img_") and item["file"].endswith( + "&type=output&subfolder=benchmarks/mysession/images") + assert {m["batch_idx"] for m in manifest["items"]} == {0, 1} + assert len({m["id"] for m in manifest["items"]}) == 2 + for f in out["ui"]["images"]: + assert f["subfolder"] == "benchmarks/mysession/images" and f["type"] == "output" + assert Image.open(base / "images" / f["filename"]).size == (24, 16) + + +def test_collect_png_embeds_prompt_chunk_and_formats(env): + prompt = {"3": {"class_type": "KSampler", "inputs": {}}} + out = collect_mod.USCGSweepCollect().collect(images(), "fmt", job_json(image_format="png"), prompt=prompt) + img = Image.open(env / "benchmarks" / "fmt" / "images" / out["ui"]["images"][0]["filename"]) + assert img.format == "PNG" and json.loads(img.text["prompt"]) == prompt + out = collect_mod.USCGSweepCollect().collect(images(), "fmt", job_json(image_format="jpg")) + assert out["ui"]["images"][0]["filename"].endswith(".jpg") + out = collect_mod.USCGSweepCollect().collect(images(), "fmt", job_json(image_format="weird")) + assert out["ui"]["images"][0]["filename"].endswith(".webp") + + +def test_collect_sanitises_session_name_and_stays_in_benchmarks(env): + collect_mod.USCGSweepCollect().collect(images(), "../../escape/..\\x", job_json()) + assert (env / "benchmarks" / "escapex" / "manifest.json").exists() + assert not (env.parent / "escape").exists() + written = [p for p in env.rglob("*") if p.is_file()] + assert all(str(p).startswith(str(env / "benchmarks")) for p in written) + + +def test_collect_works_without_a_sweep_manual_use(env): + out = collect_mod.USCGSweepCollect().collect(images(), "manual", "{}") + item = json.loads((env / "benchmarks" / "manual" / "manifest.json").read_text())["items"][0] + assert item["model_type"] == "workflow" and "wf_sweep_id" not in item and len(out["ui"]["images"]) == 1 + collect_mod.USCGSweepCollect().collect(images(), "manual", "not json at all") # garbage job_json is tolerated + + +def test_collect_appends_to_fresh_disk_manifest_preserving_favorites(env): + d = env / "benchmarks" / "fav" + d.mkdir(parents=True) + (d / "manifest.json").write_text(json.dumps({"session_name": "fav", "meta": {"k": 1}, + "items": [{"id": 1, "favorited": True}]})) + collect_mod.USCGSweepCollect().collect(images(), "fav", job_json()) + m = json.loads((d / "manifest.json").read_text()) + assert [i.get("id") for i in m["items"]][-1] == 1 and m["items"][-1]["favorited"] is True + assert m["meta"] == {"k": 1} and len(m["items"]) == 2 + + +def test_collect_replace_item_ids_removes_old_item_and_file(env): + d = env / "benchmarks" / "rep" + (d / "images").mkdir(parents=True) + (d / "images" / "img_old.webp").write_bytes(b"x") + (d / "manifest.json").write_text(json.dumps({"items": [ + {"id": 99, "file": "/view?filename=img_old.webp&type=output&subfolder=benchmarks/rep/images"}]})) + collect_mod.USCGSweepCollect().collect(images(), "rep", job_json(replace_item_ids=[99])) + m = json.loads((d / "manifest.json").read_text()) + assert [i["id"] for i in m["items"]] != [99] and all(i["id"] != 99 for i in m["items"]) and len(m["items"]) == 1 + assert not (d / "images" / "img_old.webp").exists() + + +def test_collect_notifies_runner_and_dashboard(env, monkeypatch): + events = [] + import server + ps = type("PS", (), {"send_sync": lambda self, e, d: events.append((e, d))})() + monkeypatch.setattr(server.PromptServer, "instance", ps, raising=False) + called = [] + monkeypatch.setattr(runner, "on_collected", lambda sid, idx, items: called.append((sid, idx, items)) or True) + collect_mod.USCGSweepCollect().collect(images(), "notify", job_json()) + assert called[0][0] == "sw1" and called[0][1] == 0 and len(called[0][2]) == 1 + (event, data), = events + assert event == "ultimate_grid.update" and data["node"] == "12" and data["session_name"] == "notify" + assert len(data["new_items"]) == 1 and isinstance(data["meta"], dict) + + +def test_collect_duration_uses_runner_start_time(env, monkeypatch): + monkeypatch.setattr(runner, "job_started_at", lambda sid, idx: time.time() - 10) + collect_mod.USCGSweepCollect().collect(images(), "dur", job_json()) + item = json.loads((env / "benchmarks" / "dur" / "manifest.json").read_text())["items"][0] + assert 9.5 < item["duration"] < 12 + + +def test_package_registers_collector_node_and_sweep_routes(): + pkg = sys.modules.get(PKG) + mappings = getattr(pkg, "NODE_CLASS_MAPPINGS", None) + assert mappings is not None and "USCGSweepCollect" in mappings + assert pkg.NODE_DISPLAY_NAME_MAPPINGS["USCGSweepCollect"] == "USCG Sweep Collect" + for fn in ("sweep_parse", "sweep_dry_run", "sweep_start", "sweep_status", "sweep_pause", "sweep_resume", "sweep_cancel"): + assert callable(getattr(pkg, fn)), fn diff --git a/tests/test_workflow_sweep_generator.py b/tests/test_workflow_sweep_generator.py new file mode 100644 index 0000000..c13f45b --- /dev/null +++ b/tests/test_workflow_sweep_generator.py @@ -0,0 +1,131 @@ +"""Generator node in workflow mode: IS_CHANGED, hand-off to the sweep runner, never blocks.""" +import importlib +import importlib.util +import json +import os +import sys + +import pytest + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +PKG = os.path.basename(ROOT) +sys.path.insert(0, ROOT) + +# the Generator imports the runner as .workflow_sweep_runner; use that same module instance here +runner = importlib.import_module(f"{PKG}.workflow_sweep_runner") + + +@pytest.fixture(scope="module") +def sampler_node(): + lora_utils = sys.modules[f"{PKG}.lora_utils"] + if not hasattr(lora_utils, "load_and_save_tags"): + lora_utils.load_and_save_tags = lambda *a, **k: None + sys.modules["comfy"].utils = sys.modules["comfy.utils"] + sys.modules["comfy"].sd = sys.modules["comfy.sd"] + # root conftest registers a stub for sampler_node; load the real file under a private name + spec = importlib.util.spec_from_file_location(f"{PKG}.sampler_node_real", os.path.join(ROOT, "sampler_node.py")) + mod = importlib.util.module_from_spec(spec) + mod.__package__ = PKG + spec.loader.exec_module(mod) + return mod + + +def wf_config(): + with open(os.path.join(ROOT, "tests", "fixtures", "image_moss_image_t2i.json"), encoding="utf-8") as f: + wf = json.load(f) + return {"model_type": "workflow", "name": "c", "workflow": wf, + "workflow_output": {"node_id": "38", "output_index": 0}, + "workflow_bindings": [{"id": "b", "kind": "literal", "targets": [{"node_id": "32", "input": "seed"}], + "values": [1, 2]}]} + + +def test_workflow_configs_detection(sampler_node): + detect = sampler_node._workflow_configs + assert detect(json.dumps({"configs": [wf_config()]})) and detect(json.dumps([wf_config()])) + assert detect(json.dumps({"configs": [{"model_type": "checkpoint", "model": "workflow.safetensors"}]})) is None + assert detect("not json workflow") is None and detect(None) is None and detect("{}") is None + + +def test_is_changed_is_nan_for_workflow_and_stable_for_standard(sampler_node): + cls = sampler_node.SamplerGridTester + wf = json.dumps({"configs": [wf_config()]}) + assert cls.IS_CHANGED(configs_json=wf) != cls.IS_CHANGED(configs_json=wf) # NaN + std = json.dumps({"configs": [{"model_type": "checkpoint", "model": "a"}]}) + assert cls.IS_CHANGED(configs_json=std) == cls.IS_CHANGED(configs_json=std) + + +class _FakeBridge: + def __init__(self): + self.enqueued = [] + + def run_on_loop(self, coro): + import asyncio + return asyncio.run(coro) + + async def validate(self, prompt, partial): + return True, None, {} + + async def enqueue(self, prompt, partial, client_id): + self.enqueued.append(prompt) + return f"p{len(self.enqueued)}", None, {} + + def queue_state(self): + return set(), {f"p{i + 1}" for i in range(len(self.enqueued))} + + def history_entry(self, pid): + return None + + def delete_queued(self, pid): + return True + + def interrupt(self, pid): + return False + + def send_sync(self, event, data): + pass + + +def test_run_tests_starts_sweep_and_returns_dashboard_without_waiting(sampler_node, tmp_path, monkeypatch): + import folder_paths + monkeypatch.setattr(folder_paths, "get_output_directory", lambda: str(tmp_path), raising=False) + monkeypatch.setattr(folder_paths, "get_filename_list", lambda t: ["a.safetensors"], raising=False) + runner.reset_registry() + bridge = _FakeBridge() + runner.set_bridge(bridge) + captured = {} + monkeypatch.setattr(sampler_node, "get_html_template", + lambda title, manifest, node_id: captured.update(title=title, manifest=manifest, node=node_id) or "") + monkeypatch.setattr(runner, "TICK", 0.05) + # class_info=None default arg of start_sweep would hit ComfyUI's registry; the stub `nodes` has none -> None + cfg = json.dumps({"configs": [wf_config()], "_session_settings": {"session_name": "gen_wf", "start_at_job": 0}}) + try: + out = sampler_node.SamplerGridTester().run_tests(cfg, unique_id="77") + assert out == ("",) and captured["title"] == "gen_wf" and captured["node"] == "77" + assert captured["manifest"]["session_name"] == "gen_wf" and captured["manifest"]["items"] == [] + sw = runner._active_sweep() + assert sw is not None and sw.origin_node == "77" and len(sw.jobs) == 2 # returned while jobs are pending + # a second Generator run while the sweep is active is refused with a readable error + with pytest.raises(RuntimeError, match="already active"): + sampler_node.SamplerGridTester().run_tests(cfg, unique_id="77") + finally: + sw = runner._active_sweep() + if sw is not None: + runner.cancel(sw.id) + sw.thread.join(5) + runner.set_bridge(None) + runner.reset_registry() + + +def test_mixed_workflow_and_standard_configs_raise(sampler_node, tmp_path, monkeypatch): + import folder_paths + monkeypatch.setattr(folder_paths, "get_output_directory", lambda: str(tmp_path), raising=False) + monkeypatch.setattr(folder_paths, "get_filename_list", lambda t: ["a.safetensors"], raising=False) + runner.reset_registry() + runner.set_bridge(_FakeBridge()) + cfg = json.dumps({"configs": [wf_config(), {"model_type": "checkpoint", "model": "x"}]}) + try: + with pytest.raises(ValueError, match="cannot be mixed"): + sampler_node.SamplerGridTester().run_tests(cfg, unique_id="1") + finally: + runner.set_bridge(None) + runner.reset_registry() diff --git a/tests/test_workflow_sweep_runner.py b/tests/test_workflow_sweep_runner.py new file mode 100644 index 0000000..2fecf68 --- /dev/null +++ b/tests/test_workflow_sweep_runner.py @@ -0,0 +1,628 @@ +"""Workflow Sweep runner: window, results, errors, breaker, interrupt/lost, cancel, resume, bridge mirror.""" +import asyncio +import json +import os +import sys +import threading +import time +import types + +import pytest + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, ROOT) + +import workflow_sweep as wfs # noqa: E402 +import workflow_sweep_runner as runner # noqa: E402 + + +def load(name): + with open(os.path.join(ROOT, "tests", "fixtures", name), "r", encoding="utf-8") as f: + return json.load(f) + + +def entry(values=(1, 2, 3, 4, 5, 6), **extra): + e = {"model_type": "workflow", "name": "c", "workflow": load("image_moss_image_t2i.json"), + "workflow_output": {"node_id": "38", "output_index": 0}, + "workflow_bindings": [{"id": "b", "kind": "literal", "targets": [{"node_id": "32", "input": "seed"}], + "values": list(values)}]} + e.update(extra) + return e + + +class FakeBridge: + """Synchronous stand-in for ComfyBridge with a hand-driven queue/history.""" + + def __init__(self): + self.queue, self.running, self.history, self.meta = [], [], {}, {} + self.enqueued, self.events, self.deleted, self.interrupted = [], [], [], [] + self.reject = None # callable(prompt) -> error string | None + self.n = 0 + + def run_on_loop(self, coro): + return asyncio.run(coro) + + async def validate(self, prompt, partial): + err = self.reject(prompt) if self.reject else None + if err: + return False, {"type": "prompt_outputs_failed_validation", "message": err}, \ + {"29": {"class_type": "UNETLoader", "errors": [{"message": err, "details": "unet_name"}]}} + return True, None, {} + + async def enqueue(self, prompt, partial, client_id): + ok, error, node_errors = await self.validate(prompt, partial) + if not ok: + return None, error, node_errors + self.n += 1 + pid = f"pid-{self.n}" + collector = [n for n in prompt.values() if n["class_type"] == "USCGSweepCollect"][0] + self.meta[pid] = json.loads(collector["inputs"]["job_json"]) + self.queue.append(pid) + self.enqueued.append((pid, prompt, partial, client_id)) + return pid, None, {} + + def queue_state(self): + return set(self.running), set(self.queue) + + def history_entry(self, pid): + return self.history.get(pid) + + def delete_queued(self, pid): + if pid in self.queue: + self.queue.remove(pid) + self.deleted.append(pid) + return True + return False + + def interrupt(self, pid): + self.interrupted.append(pid) + return pid in self.running + + def send_sync(self, event, data): + self.events.append((event, data)) + + # ---- test controls + def start(self, pid): + self.queue.remove(pid) + self.running.append(pid) + + def finish(self, pid, kind="success", collect=True, ts=(1000, 3500)): + if pid in self.queue: + self.queue.remove(pid) + if pid in self.running: + self.running.remove(pid) + msgs = [["execution_start", {"prompt_id": pid, "timestamp": ts[0]}]] + if kind == "success": + msgs.append(["execution_success", {"prompt_id": pid, "timestamp": ts[1]}]) + status = "success" + elif kind == "error": + msgs.append(["execution_error", {"prompt_id": pid, "node_id": "32", "node_type": "KSampler", + "exception_message": "CUDA out of memory", "timestamp": ts[1]}]) + status = "error" + else: + msgs.append(["execution_interrupted", {"prompt_id": pid, "timestamp": ts[1]}]) + status = "error" + self.history[pid] = {"status": {"status_str": status, "completed": kind == "success", "messages": msgs}, + "outputs": {}} + if kind == "success" and collect: + meta = self.meta[pid] + runner.on_collected(meta["sweep_id"], meta["job_index"], [{"id": 1}]) + + def progress_events(self): + return [d for e, d in self.events if e == "ultimate_grid.progress"] + + +@pytest.fixture +def env(tmp_path, monkeypatch): + import folder_paths + monkeypatch.setattr(folder_paths, "get_output_directory", lambda: str(tmp_path), raising=False) + runner.reset_registry() + bridge = FakeBridge() + yield types.SimpleNamespace(bridge=bridge, tmp=tmp_path) + runner.reset_registry() + + +def make(env, configs=None, settings=None, origin="node7", thread=False): + settings = {"session_name": "wf_test", **(settings or {})} + sw = runner.prepare_sweep(configs or [entry()], settings, origin, None, env.bridge, class_info=None) + runner.launch(sw, start_thread=thread) + return sw + + +def drive(sw, bridge, n=1): + for _ in range(n): + if not runner.tick(sw): + return False + return True + + +def complete_all(sw, bridge, limit=100): + for _ in range(limit): + for pid in list(bridge.queue): + bridge.start(pid) + for pid in list(bridge.running): + bridge.finish(pid) + if not runner.tick(sw): + return + raise AssertionError("sweep did not finish") + + +def test_window_respected_and_refilled(env): + sw = make(env) + drive(sw, env.bridge) + assert len(env.bridge.queue) == runner.WINDOW == 2 + drive(sw, env.bridge, 3) + assert len(env.bridge.enqueued) == 2 # nothing new while two are in flight + pid = env.bridge.queue[0] + env.bridge.start(pid) + env.bridge.finish(pid) + drive(sw, env.bridge) + assert len(env.bridge.enqueued) == 3 # the freed slot is refilled + assert len(sw.in_flight()) == 2 + + +def test_enqueued_prompt_shape_and_synthetic_client_id(env): + sw = make(env) + drive(sw, env.bridge) + pid, prompt, partial, client_id = env.bridge.enqueued[0] + collector_id = [k for k, v in prompt.items() if v["class_type"] == "USCGSweepCollect"][0] + assert partial == [collector_id] # the template's SaveImage / PreviewAny are not executed + assert client_id == f"uscg-sweep-{sw.id}" + job_json = json.loads(prompt[collector_id]["inputs"]["job_json"]) + assert job_json["sweep_id"] == sw.id and job_json["job_index"] == 0 + assert "class_type" not in json.dumps(job_json) # job conf only, never the template + assert prompt[collector_id]["inputs"]["session_name"] == "wf_test" + assert prompt["32"]["inputs"]["seed"] == 1 + + +def test_run_other_outputs_executes_everything(env): + sw = make(env, [entry(workflow_run_other_outputs=True)]) + drive(sw, env.bridge) + assert env.bridge.enqueued[0][2] is None + + +def test_success_path_progress_events_and_completion(env): + sw = make(env, [entry((1, 2, 3))]) + complete_all(sw, env.bridge) + assert sw.state == "completed" + assert [j["status"] for j in sw.jobs] == ["done"] * 3 + assert sw.jobs[0]["duration"] == pytest.approx(2.5) # from ComfyUI's execution_start/success timestamps + events = env.bridge.progress_events() + assert events[0]["current_job"] == 0 and events[0]["total_jobs"] == 3 and events[0]["phase"] == "sweep" + assert all(e["node"] == "node7" and e["session_name"] == "wf_test" for e in events) + last = events[-1] + assert last["complete"] is True and last["total_generated"] == 3 and last["eta_str"] == "Done" + assert any(e["current_job"] == 2 and not e.get("complete") for e in events) + status = runner.get_status(sw.id) + assert status["done"] == 3 and status["errors"] == 0 and status["state"] == "completed" + assert runner.get_status(session="wf_test")["sweep_id"] == sw.id + + +def test_spec_file_written_without_secrets(env): + sw = make(env, [entry((1, 2))]) + complete_all(sw, env.bridge) + path = env.tmp / "benchmarks" / "wf_test" / "workflow_sweeps" / f"{sw.id}.json" + spec = json.loads(path.read_text()) + assert spec["state"] == "completed" and len(spec["jobs"]) == 2 + assert spec["jobs"][0]["subs"] == [["32", "seed", 1]] and spec["jobs"][0]["prompt_id"] == "pid-1" + assert list(spec["templates"].values())[0]["output"] == ["38", 0] + + +def test_manifest_created_with_meta(env): + make(env) + data = json.loads((env.tmp / "benchmarks" / "wf_test" / "manifest.json").read_text()) + assert data["session_name"] == "wf_test" and data["items"] == [] and data["meta"]["image_format"] == "webp" + + +def test_job_error_recorded_and_sweep_continues(env): + sw = make(env, [entry((1, 2, 3))]) + drive(sw, env.bridge) + p1, p2 = env.bridge.queue[:2] + env.bridge.start(p1) + env.bridge.finish(p1, "error") + env.bridge.start(p2) + env.bridge.finish(p2) + complete_all(sw, env.bridge) + assert [j["status"] for j in sw.jobs] == ["error", "done", "done"] + assert "CUDA out of memory" in sw.jobs[0]["error"] and "KSampler" in sw.jobs[0]["error"] + assert sw.state == "completed" and "1 job" in sw.reason + status = runner.get_status(sw.id) + assert status["errors"] == 1 and status["last_errors"][0]["index"] == 0 + assert sw.consecutive_errors == 0 # reset by the later success + + +def test_three_consecutive_errors_pause(env): + sw = make(env, [entry(range(1, 9))]) + for _ in range(3): + drive(sw, env.bridge) + for pid in list(env.bridge.queue): + env.bridge.start(pid) + for pid in list(env.bridge.running): + env.bridge.finish(pid, "error") + break + drive(sw, env.bridge) + if sw.state == "paused": + break + drive(sw, env.bridge) + assert sw.state == "paused" and "in a row failed" in sw.reason and "CUDA" in sw.reason + + +def test_invalid_job_continues_and_first_invalid_fails_start(env): + env.bridge.reject = lambda prompt: "Value not in list: unet_name" if prompt["32"]["inputs"]["seed"] == 2 else None + sw = make(env, [entry((1, 2, 3))]) + complete_all(sw, env.bridge) + assert [j["status"] for j in sw.jobs] == ["done", "invalid", "done"] + assert "node 29 (UNETLoader" in sw.jobs[1]["error"] and "Value not in list" in sw.jobs[1]["error"] + # job 1 invalid at start -> the sweep is refused with the node errors (route answers 400) + env.bridge.reject = lambda prompt: "model missing" + sw2 = runner.prepare_sweep([entry((1, 2))], {"session_name": "other"}, None, None, env.bridge, class_info=None) + runner.reset_registry() + with pytest.raises(ValueError, match="first job is not valid.*model missing"): + asyncio.run(runner.validate_first(sw2)) + + +def test_interrupted_pauses_and_resume_requeues(env): + sw = make(env, [entry((1, 2, 3))]) + drive(sw, env.bridge) + p1 = env.bridge.queue[0] + env.bridge.start(p1) + env.bridge.finish(p1, "interrupted") + drive(sw, env.bridge) + assert sw.state == "paused" and "interrupted" in sw.reason + assert sw.jobs[0]["status"] == "interrupted" + # the second (still queued) prompt is pulled while paused + drive(sw, env.bridge) + assert env.bridge.deleted and sw.jobs[1]["status"] == "pending" + ok, state = runner.resume(sw.id) + assert ok and state == "running" and sw.jobs[0]["status"] == "pending" + complete_all(sw, env.bridge) + assert [j["status"] for j in sw.jobs] == ["done"] * 3 and sw.state == "completed" + + +def test_lost_prompt_pauses(env, monkeypatch): + monkeypatch.setattr(runner, "LOST_GRACE", -1) + sw = make(env, [entry((1, 2, 3))]) + drive(sw, env.bridge) + env.bridge.queue.clear() # user cleared ComfyUI's queue + drive(sw, env.bridge) + assert sw.state == "paused" and "removed from ComfyUI's queue" in sw.reason + assert {j["status"] for j in sw.jobs[:2]} == {"lost"} + runner.resume(sw.id) + assert [j["status"] for j in sw.jobs[:2]] == ["pending", "pending"] + + +def test_success_without_collected_item_is_an_error(env, monkeypatch): + monkeypatch.setattr(runner, "COLLECT_GRACE", -1) + sw = make(env, [entry((1,))]) + drive(sw, env.bridge) + pid = env.bridge.queue[0] + env.bridge.start(pid) + env.bridge.finish(pid, collect=False) + drive(sw, env.bridge, 2) + assert sw.jobs[0]["status"] == "error" and "collector" in sw.jobs[0]["error"] + + +def test_cancel_deletes_queued_interrupts_running_and_marks_rest(env): + sw = make(env, [entry(range(1, 6))]) + drive(sw, env.bridge) + p1, p2 = env.bridge.queue[:2] + env.bridge.start(p1) + drive(sw, env.bridge) + assert sw.jobs[0]["status"] == "running" + ok, state = runner.cancel(sw.id) + assert ok and state == "cancelled" + assert drive(sw, env.bridge) is False + assert env.bridge.deleted == [p2] and env.bridge.interrupted == [p1] + assert {j["status"] for j in sw.jobs} == {"cancelled"} + assert runner.cancel(sw.id)[0] is False # already finished + assert runner.get_status(sw.id)["state"] == "cancelled" + + +def test_pause_pulls_queued_but_lets_running_finish_then_resume(env): + sw = make(env, [entry((1, 2, 3))]) + drive(sw, env.bridge) + p1, p2 = env.bridge.queue[:2] + env.bridge.start(p1) + drive(sw, env.bridge) + assert runner.pause(sw.id) == (True, "paused") + drive(sw, env.bridge) + assert env.bridge.deleted == [p2] and sw.jobs[0]["status"] == "running" and sw.jobs[1]["status"] == "pending" + env.bridge.finish(p1) + drive(sw, env.bridge) + assert sw.jobs[0]["status"] == "done" and sw.state == "paused" and not env.bridge.queue + assert runner.pause(sw.id)[0] is False + runner.resume(sw.id) + complete_all(sw, env.bridge) + assert sw.state == "completed" + + +def _write_manifest(env, items): + d = env.tmp / "benchmarks" / "wf_test" + d.mkdir(parents=True, exist_ok=True) + (d / "manifest.json").write_text(json.dumps({"session_name": "wf_test", "items": items, "meta": {"x": 1}})) + + +def test_resume_skips_fingerprints_in_manifest(env): + plan = wfs.plan_workflow_configs([entry((1, 2, 3))]) + done_fps = [plan["jobs"][0]["wf_fingerprint"], plan["jobs"][2]["wf_fingerprint"]] + _write_manifest(env, [{"id": 11, "wf_fingerprint": done_fps[0], "gen_index": 4}, + {"id": 12, "wf_fingerprint": done_fps[1]}, {"id": 13}]) + sw = make(env, [entry((1, 2, 3))]) + assert [j["status"] for j in sw.jobs] == ["skipped", "pending", "skipped"] + assert sw.skipped_existing == 2 and sw.to_run == 1 + assert sw.jobs[1]["gen_index"] == 5 + 1 # continues after the highest existing gen_index + complete_all(sw, env.bridge) + assert len(env.bridge.enqueued) == 1 and env.bridge.enqueued[0][1]["32"]["inputs"]["seed"] == 2 + first_event = env.bridge.progress_events()[0] + assert first_event["current_job"] == 2 and first_event["total_jobs"] == 3 # "2/3 already done" + + +def test_overwrite_existing_reruns_and_carries_replace_ids(env): + plan = wfs.plan_workflow_configs([entry((1, 2))]) + _write_manifest(env, [{"id": 11, "wf_fingerprint": plan["jobs"][0]["wf_fingerprint"]}]) + sw = make(env, [entry((1, 2))], {"overwrite_existing": True}) + assert [j["status"] for j in sw.jobs] == ["pending", "pending"] and sw.jobs[0]["replace_ids"] == [11] + drive(sw, env.bridge) + jj = env.bridge.meta[env.bridge.enqueued[0][0]] + assert jj["replace_item_ids"] == [11] + + +def test_start_at_job_skips_first_n(env): + sw = make(env, [entry((1, 2, 3, 4))], {"start_at_job": 2}) + assert [j["status"] for j in sw.jobs] == ["skipped", "skipped", "pending", "pending"] + assert sw.skipped_existing == 0 + + +def test_all_done_sweep_completes_immediately(env): + plan = wfs.plan_workflow_configs([entry((1,))]) + _write_manifest(env, [{"id": 11, "wf_fingerprint": plan["jobs"][0]["wf_fingerprint"]}]) + sw = make(env, [entry((1,))]) + assert drive(sw, env.bridge) is False and sw.state == "completed" + + +def test_one_active_sweep_and_409(env): + sw = make(env) + with pytest.raises(runner.SweepBusy): + runner.prepare_sweep([entry()], {"session_name": "x"}, None, None, env.bridge, class_info=None) + status, body = asyncio.run(runner.api_start( + {"configs_json": json.dumps({"configs": [entry()], "_session_settings": {"session_name": "x"}})}, + bridge=env.bridge, class_info=None)) + assert status == 409 and body["sweep_id"] == sw.id + runner.cancel(sw.id) + drive(sw, env.bridge) + runner.prepare_sweep([entry()], {"session_name": "x"}, None, None, env.bridge, class_info=None) # free again + + +def test_api_start_ok_and_bad_input(env): + body = {"configs_json": json.dumps({"configs": [entry((1, 2))], "_session_settings": {"session_name": "api"}}), + "node_id": "5"} + status, out = asyncio.run(runner.api_start(body, bridge=env.bridge, class_info=None)) + assert status == 200 and out["total"] == 2 and out["to_run"] == 2 + runner.cancel(out["sweep_id"]) + time.sleep(runner.TICK * 3) + status, out = asyncio.run(runner.api_start({"configs_json": json.dumps({"configs": [{"model_type": "checkpoint"}]})}, + bridge=env.bridge, class_info=None)) + assert status == 400 and "cannot be mixed" in out["error"] + status, out = asyncio.run(runner.api_start({}, bridge=env.bridge, class_info=None)) + assert status == 400 + + +def test_start_sweep_returns_while_jobs_pending_and_never_blocks(env): + t0 = time.time() + sw = runner.start_sweep([entry(range(1, 50))], {"session_name": "nb"}, "n1", None, env.bridge, class_info=None) + assert time.time() - t0 < 2 and sw.state == "running" + time.sleep(runner.TICK * 3) # the feeder thread fills the window on its own + assert len(env.bridge.enqueued) == runner.WINDOW + assert runner.cancel(sw.id)[0] + sw.thread.join(5) + assert not sw.thread.is_alive() and sw.state == "cancelled" + + +def test_dry_run_shape_and_no_queueing(env): + body = {"configs_json": json.dumps({"configs": [entry((1, 2, 3))], "_session_settings": {"session_name": "dry"}}), + "limit": 2} + out = asyncio.run(runner.dry_run(body, bridge=env.bridge, class_info=None)) + assert out["total"] == 3 and out["to_run"] == 3 and len(out["jobs"]) == 2 + assert out["validation"]["ok"] is True and out["axes"][0]["count"] == 3 + j = out["jobs"][0] + assert j["prompt"]["32"]["inputs"]["seed"] == 1 and j["partial_execution_targets"] + assert any(n["class_type"] == "USCGSweepCollect" for n in j["prompt"].values()) + assert env.bridge.enqueued == [] and runner._active_sweep() is None + env.bridge.reject = lambda p: "bad model" + out = asyncio.run(runner.dry_run(body, bridge=env.bridge, class_info=None)) + assert out["validation"]["ok"] is False and "bad model" in out["validation"]["error"] + + +def test_dry_run_allowed_while_a_sweep_is_active(env): + make(env) + body = {"configs_json": json.dumps({"configs": [entry((1,))], "_session_settings": {"session_name": "d2"}})} + assert asyncio.run(runner.dry_run(body, bridge=env.bridge, class_info=None))["total"] == 1 + + +def test_configs_from_body_uses_builder_state_transformer(): + state = {"session_name": "st", "config_arrays": [{ + "name": "WF", "kind": "workflow", + "workflow": {"api_prompt": load("sdxl_default_api.json"), "output": {"node_id": "8", "output_index": 0}, + "bindings": []}}]} + configs, settings = runner.configs_from_body({"state": state}) + assert configs[0]["model_type"] == "workflow" and settings["session_name"] == "st" + with pytest.raises(ValueError): + runner.configs_from_body({}) + + +def test_classify_history_entries(): + ok = {"status": {"status_str": "success", "messages": [["execution_start", {"timestamp": 1}], + ["execution_success", {"timestamp": 3}]]}} + assert runner._classify(ok)[0] == "success" + err = {"status": {"status_str": "error", "messages": [["execution_error", {"node_id": "3", "node_type": "K", + "exception_message": "boom"}]]}} + kind, detail = runner._classify(err) + assert kind == "error" and "boom" in detail["error"] and "K #3" in detail["error"] + assert runner._classify({"status": {"status_str": "error", "messages": [["execution_interrupted", {}]]}})[0] == "interrupted" + assert runner._classify({"status": None})[0] == "unknown" + assert runner._classify(None)[0] == "unknown" + + +def test_summarize_errors(): + msg = runner.summarize_errors({"message": "Prompt outputs failed validation", "details": "x"}, + {"29": {"class_type": "UNETLoader", "errors": [{"message": "Value not in list", + "details": "unet_name: 'a'"}]}}, + load("image_moss_image_t2i.json")) + assert "Prompt outputs failed validation" in msg and "node 29 (UNETLoader / Load Diffusion Model)" in msg + + +# ------------------------------------------------------------------ ComfyBridge mirrors server.py post_prompt + +class _Queue: + def __init__(self): + self.items = [] + + def put(self, item): + self.items.append(item) + + +def _fake_comfy(monkeypatch, with_sensitive=True, validate_args=3, valid=True, handlers=None, replace=True): + calls = {"validate": [], "replace": [], "handlers": []} + + class PS: + number = 5.0 + prompt_queue = _Queue() + loop = None + + def trigger_on_prompt(self, json_data): + calls["handlers"].append(dict(json_data)) + if handlers: + json_data = handlers(json_data) + return json_data + + class node_replace_manager: # noqa: N801 + @staticmethod + def apply_replacements(prompt): + calls["replace"].append(prompt) + prompt["9999"] = {"class_type": "Replaced", "inputs": {}} + + ps = PS() + if not replace: + del PS.node_replace_manager + server_mod = sys.modules["server"] + monkeypatch.setattr(server_mod, "PromptServer", types.SimpleNamespace(instance=ps), raising=False) + monkeypatch.setattr(server_mod, "valid_workflow_metadata", lambda jd: jd.get("workflow_metadata"), raising=False) + monkeypatch.setattr(server_mod, "workflow_metadata_from_prompt", lambda ed: None, raising=False) + + ex = types.ModuleType("execution") + if validate_args == 3: + async def validate_prompt(prompt_id, prompt, partial_execution_list): + calls["validate"].append((prompt_id, partial_execution_list)) + return (valid, None if valid else {"message": "bad"}, ["38x"], {} if valid else {"3": {"errors": []}}) + elif validate_args == 2: + def validate_prompt(prompt_id, prompt): + calls["validate"].append((prompt_id, None)) + return (valid, None, ["38x"], {}) + else: + def validate_prompt(prompt): + calls["validate"].append((None, None)) + return (valid, None, ["38x"], {}) + ex.validate_prompt = validate_prompt + if with_sensitive: + ex.SENSITIVE_EXTRA_DATA_KEYS = ("auth_token_comfy_org", "api_key_comfy_org") + monkeypatch.setitem(sys.modules, "execution", ex) + return ps, calls + + +def test_bridge_enqueue_mirrors_post_prompt(monkeypatch): + ps, calls = _fake_comfy(monkeypatch) + prompt = {"1": {"class_type": "X", "inputs": {}}} + pid, err, node_errors = asyncio.run(runner.ComfyBridge().enqueue(prompt, ["1"], "cid")) + assert err is None and node_errors == {} + (item,) = ps.prompt_queue.items + assert len(item) == 6 # (number, prompt_id, prompt, extra_data, outputs, sensitive) + number, prompt_id, queued_prompt, extra, outputs, sensitive = item + assert number == 5.0 and ps.number == 6.0 # number taken, then incremented + assert prompt_id == pid and len(pid) == 36 + assert queued_prompt is prompt and "9999" in prompt # node replacements applied before validation + assert extra["client_id"] == "cid" and isinstance(extra["create_time"], int) + assert outputs == ["38x"] and sensitive == {} + assert calls["validate"] == [(pid, ["1"])] + assert calls["handlers"][0]["prompt"] is prompt and calls["handlers"][0]["partial_execution_targets"] == ["1"] + + +def test_bridge_enqueue_honours_on_prompt_handler_and_strips_sensitive(monkeypatch): + def handler(jd): + jd["extra_data"] = {"api_key_comfy_org": "SECRET", "other": 1} + jd["number"] = 2 + return jd + ps, _ = _fake_comfy(monkeypatch, handlers=handler) + asyncio.run(runner.ComfyBridge().enqueue({"1": {"class_type": "X", "inputs": {}}}, None, "c")) + number, _, _, extra, _, sensitive = ps.prompt_queue.items[0] + assert number == 2.0 and ps.number == 5.0 # an explicit number does not consume the counter + assert sensitive == {"api_key_comfy_org": "SECRET"} and "api_key_comfy_org" not in extra and extra["other"] == 1 + + +def test_bridge_enqueue_older_server_shapes(monkeypatch): + ps, calls = _fake_comfy(monkeypatch, with_sensitive=False, validate_args=2, replace=False) + pid, err, _ = asyncio.run(runner.ComfyBridge().enqueue({"1": {"class_type": "X", "inputs": {}}}, None, "c")) + assert err is None and len(ps.prompt_queue.items[0]) == 5 and calls["validate"][0][0] == pid + ps, calls = _fake_comfy(monkeypatch, validate_args=1) + asyncio.run(runner.ComfyBridge().enqueue({"1": {"class_type": "X", "inputs": {}}}, None, "c")) + assert len(ps.prompt_queue.items[0]) == 6 + + +def test_bridge_enqueue_invalid_prompt_is_not_queued(monkeypatch): + ps, _ = _fake_comfy(monkeypatch, valid=False) + pid, err, node_errors = asyncio.run(runner.ComfyBridge().enqueue({"1": {"class_type": "X", "inputs": {}}}, None, "c")) + assert pid is None and err == {"message": "bad"} and "3" in node_errors and ps.prompt_queue.items == [] + + +def test_bridge_validate_does_not_queue(monkeypatch): + ps, calls = _fake_comfy(monkeypatch) + ok, err, ne = asyncio.run(runner.ComfyBridge().validate({"1": {"class_type": "X", "inputs": {}}}, ["1"])) + assert ok and ps.prompt_queue.items == [] and len(calls["validate"]) == 1 + + +def test_run_on_loop_refuses_the_loop_thread_and_works_from_others(monkeypatch): + loop = asyncio.new_event_loop() + t = threading.Thread(target=loop.run_forever, daemon=True) + t.start() + ps, _ = _fake_comfy(monkeypatch) + ps.loop = loop + bridge = runner.ComfyBridge() + + async def double(x): + return x * 2 + assert bridge.run_on_loop(double(21)) == 42 + + async def from_loop(): + c = double(1) + with pytest.raises(RuntimeError, match="event loop thread"): + bridge.run_on_loop(c) + return "refused" + assert asyncio.run_coroutine_threadsafe(from_loop(), loop).result(5) == "refused" + loop.call_soon_threadsafe(loop.stop) + t.join(5) + + +def test_bridge_queue_helpers(monkeypatch): + removed = [] + + class Q: + def get_current_queue_volatile(self): + return [(0, "r1", {}, {}, [])], [(1, "q1", {}, {}, [], {}), (2, "q2", {}, {}, [], {})] + + def get_history(self, pid, map_function=None): + return {pid: map_function({"status": {"status_str": "success"}, "outputs": {}, "prompt": "big"})} \ + if pid == "h1" else {} + + def delete_queue_item(self, fn): + removed.append(fn((1, "q1"))) + return True + + def interrupt_if_running(self, pid): + return pid == "r1" + ps, _ = _fake_comfy(monkeypatch) + ps.prompt_queue = Q() + b = runner.ComfyBridge() + assert b.queue_state() == ({"r1"}, {"q1", "q2"}) + assert b.history_entry("h1") == {"status": {"status_str": "success"}, "outputs": {}} # no deep copy of prompt + assert b.history_entry("nope") is None + assert b.delete_queued("q1") is True and removed == [True] + assert b.interrupt("r1") is True and b.interrupt("q1") is False diff --git a/workflow_sweep.py b/workflow_sweep.py index 5db13c0..1f75c2e 100644 --- a/workflow_sweep.py +++ b/workflow_sweep.py @@ -950,3 +950,20 @@ def build_collect_item(conf, width, height, duration, batch_idx, gen_index, crea def new_item_id(): return int(time.time() * 100000) + random.randint(0, 1000) + + +def parse_for_ui(api_json, class_info=comfy_class_info): + """Response body of POST /config_tester/sweep/parse. Raises ValueError for unusable input.""" + prompt = parse_api_workflow(api_json) + sources = find_image_sources(prompt, class_info) + default = sources[0] if sources else None + suggestions = suggest_bindings(prompt, default, class_info) if default else [] + warnings = [] + nodes_ = describe_nodes(prompt, class_info) + missing = sorted({n["class_type"] for n in nodes_ if n["missing_class"]}) + if missing: + warnings.append("These node types are not installed in this ComfyUI: " + ", ".join(missing)) + if not sources: + warnings.append("No IMAGE output was found; pick the image socket to collect manually.") + return {"ok": True, "nodes": nodes_, "outputs": find_output_nodes(prompt, class_info), + "image_sources": sources, "suggestions": suggestions, "warnings": warnings, "errors": []} diff --git a/workflow_sweep_collect.py b/workflow_sweep_collect.py new file mode 100644 index 0000000..f206a86 --- /dev/null +++ b/workflow_sweep_collect.py @@ -0,0 +1,140 @@ +""" +USCGSweepCollect: the output node a Workflow Sweep splices onto each job's image socket. + +It saves the image(s) into the USCG session folder (benchmarks//images) and appends manifest +items with the standard fields, so the Dashboard, favourites, filters and export keep working. It never +waits on anything and never writes outside the benchmarks folder (the session name is sanitised with +the same rule as every other route). `job_json` carries the job's conf only, never the workflow template. +""" + +import json +import os +import time + +import numpy as np +from PIL import Image +from PIL.PngImagePlugin import PngInfo + +try: + from . import workflow_sweep as wfs +except ImportError: # top-level import (tests) + import workflow_sweep as wfs + + +def _append_items(*args, **kwargs): + try: + from .manifest_utils import append_items + except ImportError: + from manifest_utils import append_items + return append_items(*args, **kwargs) + + +def _runner(): + try: + from . import workflow_sweep_runner as runner + except ImportError: + import workflow_sweep_runner as runner + return runner + + +def _tensor_to_array(images, b): + """images[b] -> uint8 HxWxC numpy array (SaveImage's conversion), duck-typed so no torch import is needed.""" + arr = images[b] + if hasattr(arr, "cpu"): + arr = arr.cpu() + if hasattr(arr, "numpy"): + arr = arr.numpy() + return np.clip(255.0 * np.asarray(arr), 0, 255).astype(np.uint8) + + +class USCGSweepCollect: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images": ("IMAGE",), + "session_name": ("STRING", {"default": "workflow_sweep"}), + "job_json": ("STRING", {"multiline": True, "default": "{}"}), + }, + "hidden": {"prompt": "PROMPT", "unique_id": "UNIQUE_ID"}, + } + + RETURN_TYPES = () + OUTPUT_NODE = True + FUNCTION = "collect" + CATEGORY = "sampling/testing" + + @classmethod + def IS_CHANGED(cls, **kwargs): + # every job must run its collector, whatever the cache says about its ancestors + return float("NaN") + + def collect(self, images, session_name="workflow_sweep", job_json="{}", prompt=None, unique_id=None): + paths = wfs.session_paths(session_name) + session = paths["session"] + os.makedirs(paths["images"], exist_ok=True) + + try: + job = json.loads(job_json or "{}") + job = job if isinstance(job, dict) else {} + except (TypeError, ValueError): + job = {} + conf = job.get("conf") if isinstance(job.get("conf"), dict) else {} + sweep_id, job_index = job.get("sweep_id"), job.get("job_index") + + fmt = str(job.get("image_format") or "webp").lower().strip() + if fmt not in ("webp", "png", "jpg", "jpeg"): + fmt = "webp" + ext = "jpg" if fmt in ("jpg", "jpeg") else fmt + + started = None + if sweep_id is not None and job_index is not None: + started = _runner().job_started_at(sweep_id, job_index) + duration = round(time.time() - started, 2) if started else 0 + + gen_index = job.get("gen_index") + items, files = [], [] + for b in range(len(images)): + arr = _tensor_to_array(images, b) + img = Image.fromarray(arr) + item = wfs.build_collect_item(conf, arr.shape[1], arr.shape[0], duration, b, gen_index) + item["id"] = wfs.new_item_id() + filename = f"img_{item['id']}.{ext}" + filepath = os.path.join(paths["images"], filename) + if fmt == "png": + info = PngInfo() + if prompt is not None: # same text chunk SaveImage writes: the image can be dragged back into ComfyUI + info.add_text("prompt", json.dumps(prompt)) + img.save(filepath, format="PNG", pnginfo=info) + elif ext == "jpg": + img.save(filepath, format="JPEG", quality=95) + else: + img.save(filepath, quality=80) + item["file"] = f"/view?filename={filename}&type=output&subfolder=benchmarks/{session}/images" + if sweep_id is not None: + item["wf_sweep_id"] = sweep_id + if job_index is not None: + item["wf_job_index"] = job_index + items.append(item) + files.append({"filename": filename, "subfolder": f"benchmarks/{session}/images", "type": "output"}) + + manifest = _append_items(paths["manifest"], items, remove_ids=job.get("replace_item_ids") or (), + session_name=session) + try: # live dashboard update (same event flush_batch_with_vae sends) + from server import PromptServer + PromptServer.instance.send_sync("ultimate_grid.update", { + "node": job.get("origin_node"), "session_name": session, "new_items": items, + "meta": (manifest or {}).get("meta", {})}) + except Exception: + pass + if sweep_id is not None: + try: + _runner().on_collected(sweep_id, job_index, items) + except Exception: + pass + print(f"[GridTester] 🧪 Sweep collected {len(items)} image(s) into '{session}'") + return {"ui": {"images": files}} + + +NODE_CLASS_MAPPINGS = {"USCGSweepCollect": USCGSweepCollect} +NODE_DISPLAY_NAME_MAPPINGS = {"USCGSweepCollect": "USCG Sweep Collect"} diff --git a/workflow_sweep_runner.py b/workflow_sweep_runner.py new file mode 100644 index 0000000..0c60f63 --- /dev/null +++ b/workflow_sweep_runner.py @@ -0,0 +1,912 @@ +""" +Workflow Sweep runner: feeds a sweep's jobs to ComfyUI's own queue, one prompt per job. + +A sweep is prepared synchronously (expand, fingerprint, skip jobs already in the manifest), then a +daemon thread keeps at most WINDOW sweep prompts queued/running and follows each through ComfyUI's +history. Every enqueue repeats server.py's post_prompt in-process (on_prompt handlers, node +replacements, awaited validate_prompt, number, extra_data, queue tuple), so a queued sweep prompt is +exactly what ComfyUI would have queued from its UI. + +Hard rules honoured here: + * no network code (the enqueue is an in-process call on the server event loop), + * the executor is never blocked: nothing in this module (or the collector node) waits on a prompt + it queued; the Generator node starts a sweep and returns immediately, + * `run_on_loop` is never called from the loop thread, and no lock is held across it. + +All loop-bound ComfyUI access goes through `ComfyBridge`, so the runner is testable with a fake. +""" + +import asyncio +import inspect +import json +import os +import random +import threading +import time +import traceback +import uuid + +try: + from . import workflow_sweep as wfs +except ImportError: # imported as a top-level module (tests) + import workflow_sweep as wfs + + +WINDOW = 2 # max sweep prompts queued/running at once +TICK = 0.25 # seconds between runner polls +LOST_GRACE = 5.0 # a prompt missing from queue, running set and history this long is "lost" +COLLECT_GRACE = 5.0 # success without a collected item for this long is an error +BREAKER = 3 # consecutive error/invalid jobs that pause the sweep +SPEC_EVERY = 10 # spec file rewrite cadence (terminal jobs) + +ACTIVE_STATES = ("preparing", "running", "paused") +TERMINAL_JOB = ("done", "error", "invalid", "skipped", "cancelled") + + +class SweepBusy(Exception): + """Another workflow sweep is still active (route -> 409).""" + def __init__(self, sweep_id): + super().__init__(f"A workflow sweep is already active ({sweep_id}). Cancel it or wait for it to finish.") + self.sweep_id = sweep_id + + +def summarize_errors(error, node_errors, prompt=None): + """One readable line from ComfyUI's validate_prompt error + node_errors.""" + parts = [] + if isinstance(error, dict): + msg = error.get("message") or error.get("type") or "invalid prompt" + if error.get("details"): + msg += f" ({error['details']})" + parts.append(msg) + elif error: + parts.append(str(error)) + for nid, ne in (node_errors or {}).items(): + ne = ne if isinstance(ne, dict) else {} + ct = ne.get("class_type") or ((prompt or {}).get(str(nid)) or {}).get("class_type") or "?" + title = wfs.node_title(prompt, nid) if prompt else ct + label = ct if title == ct else f"{ct} / {title}" + msgs = "; ".join(f"{e.get('message', '')} {e.get('details', '')}".strip() for e in (ne.get("errors") or [])) + parts.append(f"node {nid} ({label}): {msgs or 'invalid'}") + return (" | ".join(parts) or "invalid prompt")[:1500] + + +# --------------------------------------------------------------------------- bridge to ComfyUI + +class ComfyBridge: + """Thin wrapper over PromptServer.instance / execution / nodes (lazy imports, version guards).""" + + def _ps(self): + import server + return server.PromptServer.instance + + def _execution(self): + import execution + return execution + + def run_on_loop(self, coro, timeout=300): + """Run a coroutine on the server event loop from a NON-loop thread and wait for it.""" + ps = self._ps() + loop = ps.loop + try: + running = asyncio.get_running_loop() + except RuntimeError: + running = None + if running is loop: + coro.close() + raise RuntimeError("run_on_loop was called from the server event loop thread (would deadlock)") + return asyncio.run_coroutine_threadsafe(coro, loop).result(timeout) + + async def _validate(self, prompt_id, prompt, partial): + execution = self._execution() + fn = execution.validate_prompt + n = len(inspect.signature(fn).parameters) + if n >= 3: + result = fn(prompt_id, prompt, partial) + elif n == 2: + result = fn(prompt_id, prompt) + else: + result = fn(prompt) + if inspect.isawaitable(result): + result = await result + return result # (valid, error, outputs, node_errors) + + async def validate(self, prompt, partial_targets): + """Validate without queueing. Returns (ok, error, node_errors).""" + valid = await self._validate(str(uuid.uuid4()), prompt, partial_targets) + return bool(valid[0]), valid[1], valid[3] if len(valid) > 3 else {} + + async def enqueue(self, prompt, partial_targets, client_id): + """Mirror of server.py post_prompt (ComfyUI 0.42.5). Returns (prompt_id|None, error, node_errors).""" + ps = self._ps() + execution = self._execution() + json_data = {"prompt": prompt, "client_id": client_id} + if partial_targets is not None: + json_data["partial_execution_targets"] = partial_targets + if hasattr(ps, "trigger_on_prompt"): + json_data = ps.trigger_on_prompt(json_data) + + if "number" in json_data: + number = float(json_data["number"]) + else: + number = ps.number + if "front" in json_data and json_data["front"]: + number = -number + ps.number += 1 + + if "prompt" not in json_data: + return None, {"type": "no_prompt", "message": "No prompt provided"}, {} + prompt = json_data["prompt"] + prompt_id = str(uuid.uuid4()) + partial = json_data.get("partial_execution_targets") + + manager = getattr(ps, "node_replace_manager", None) + if manager is not None and hasattr(manager, "apply_replacements"): + manager.apply_replacements(prompt) + + valid = await self._validate(prompt_id, prompt, partial) + extra_data = json_data.get("extra_data", {}) + if not isinstance(extra_data, dict): + extra_data = {} + if "client_id" in json_data: + extra_data["client_id"] = json_data["client_id"] + + # workflow metadata for websocket messages (0.42.x); absent on older servers + extra_data.pop("workflow_metadata", None) + try: + import server as _server + valid_meta = getattr(_server, "valid_workflow_metadata", None) + meta_from_prompt = getattr(_server, "workflow_metadata_from_prompt", None) + metadata = valid_meta(json_data) if valid_meta else None + if metadata is None and meta_from_prompt: + metadata = meta_from_prompt(extra_data) + if metadata is not None: + extra_data["workflow_metadata"] = metadata + except Exception: + pass + + if not valid[0]: + return None, valid[1], (valid[3] if len(valid) > 3 else {}) + + outputs_to_execute = valid[2] + sensitive_keys = getattr(execution, "SENSITIVE_EXTRA_DATA_KEYS", None) + extra_data["create_time"] = int(time.time() * 1000) + if sensitive_keys is not None: # 6-tuple queue items (0.42.x) + sensitive = {} + for key in sensitive_keys: + if key in extra_data: + sensitive[key] = extra_data.pop(key) + ps.prompt_queue.put((number, prompt_id, prompt, extra_data, outputs_to_execute, sensitive)) + else: + ps.prompt_queue.put((number, prompt_id, prompt, extra_data, outputs_to_execute)) + return prompt_id, None, (valid[3] if len(valid) > 3 else {}) + + def queue_state(self): + """(running prompt ids, queued prompt ids), one atomic snapshot.""" + running, queued = self._ps().prompt_queue.get_current_queue_volatile() + return {item[1] for item in running}, {item[1] for item in queued} + + def history_entry(self, prompt_id): + """{status:{status_str, messages}, outputs} or None when ComfyUI has no history for it.""" + q = self._ps().prompt_queue + + def slim(p): + return {"status": p.get("status"), "outputs": p.get("outputs")} + try: + got = q.get_history(prompt_id, map_function=slim) + except TypeError: # older servers: no map_function + got = q.get_history(prompt_id) + return (got or {}).get(prompt_id) + + def delete_queued(self, prompt_id): + return bool(self._ps().prompt_queue.delete_queue_item(lambda item: item[1] == prompt_id)) + + def interrupt(self, prompt_id): + q = self._ps().prompt_queue + if hasattr(q, "interrupt_if_running"): + return bool(q.interrupt_if_running(prompt_id)) + running, _ = self.queue_state() + if prompt_id in running: + import nodes + nodes.interrupt_processing() + return True + return False + + def send_sync(self, event, data): + try: + self._ps().send_sync(event, data) + except Exception: + pass + + +_bridge = None + + +def get_bridge(): + global _bridge + if _bridge is None: + _bridge = ComfyBridge() + return _bridge + + +def set_bridge(bridge): + """Tests inject a FakeBridge.""" + global _bridge + _bridge = bridge + + +# --------------------------------------------------------------------------- sweep model + +class Sweep: + def __init__(self, sweep_id, session, jobs, templates, origin_node, client_id, settings, warnings, axes, bridge): + self.id = sweep_id + self.session = session + self.jobs = jobs + self.templates = templates + self.origin_node = origin_node + self.client_id = client_id + self.settings = settings # {image_format, overwrite_existing, ...} + self.warnings = warnings + self.axes = axes + self.bridge = bridge + self.state = "preparing" + self.reason = "" + self.created_at = time.time() + self.started_at = None + self.finished_at = None + self.lock = threading.RLock() + self.wake = threading.Event() + self.thread = None + self.needs_cleanup = False + self.consecutive_errors = 0 + self.durations = [] + self.last_errors = [] + self.terminal_since_spec = 0 + self.eta_start = None + + # counters (computed from job states so they can never drift) + def count(self, *statuses): + return sum(1 for j in self.jobs if j["status"] in statuses) + + @property + def skipped_existing(self): + return sum(1 for j in self.jobs if j["status"] == "skipped" and j.get("skip_reason") == "existing") + + @property + def to_run(self): + return sum(1 for j in self.jobs if j["status"] != "skipped") + + def in_flight(self): + return [j for j in self.jobs if j["status"] in ("enqueued", "running")] + + +_sweeps = {} +_registry_lock = threading.RLock() + + +def reset_registry(): + """Tests only.""" + with _registry_lock: + _sweeps.clear() + + +def _active_sweep(): + with _registry_lock: + for sw in _sweeps.values(): + if sw.state in ACTIVE_STATES: + return sw + return None + + +def _find(sweep_id=None, session=None): + with _registry_lock: + if sweep_id: + return _sweeps.get(str(sweep_id)) + if session: + matches = [s for s in _sweeps.values() if s.session == session] + return max(matches, key=lambda s: s.created_at) if matches else None + return None + + +# --------------------------------------------------------------------------- prepare / validate / launch + +def _manifest_state(manifest_path): + """({fingerprint: [item ids]}, next gen_index) from the manifest on disk.""" + data = wfs.read_manifest_items(manifest_path) + by_fp, max_gen = {}, -1 + for it in data.get("items") or []: + if not isinstance(it, dict): + continue + fp = it.get("wf_fingerprint") + if fp and it.get("id") is not None: + by_fp.setdefault(fp, []).append(it["id"]) + g = it.get("gen_index") + if isinstance(g, int) and g > max_gen: + max_gen = g + return by_fp, max_gen + 1 + + +def prepare_sweep(configs, session_settings=None, origin_node=None, client_id=None, bridge=None, + class_info=wfs.comfy_class_info, check_busy=True): + """ + Expand + fingerprint + skip finished jobs. Raises ValueError (readable) or SweepBusy. + Pure apart from reading the session manifest. + """ + busy = _active_sweep() if check_busy else None + if busy is not None: + raise SweepBusy(busy.id) + settings = dict(session_settings or {}) + if not isinstance(configs, list) or not configs: + raise ValueError("There are no configs to run.") + extra_seeds = [] + try: + n_extra = int(settings.get("add_random_seeds_to_gens", 0) or 0) + except (TypeError, ValueError): + n_extra = 0 + if n_extra > 0: + extra_seeds = [random.randint(0, 2**32 - 1) for _ in range(n_extra)] + plan = wfs.plan_workflow_configs(configs, extra_seeds, class_info=class_info) + + paths = wfs.session_paths(settings.get("session_name") or "workflow_sweep") + overwrite = bool(settings.get("overwrite_existing", False)) + try: + start_at = max(0, int(settings.get("start_at_job", 0) or 0)) + except (TypeError, ValueError): + start_at = 0 + by_fp, gen_base = _manifest_state(paths["manifest"]) + + jobs = [] + for idx, conf in enumerate(plan["jobs"]): + job = {"index": idx, "fingerprint": conf["wf_fingerprint"], "conf": conf, "status": "pending", + "prompt_id": None, "error": None, "gen_index": gen_base + idx, "replace_ids": [], + "enqueued_at": None, "started_at": None, "finished_at": None, "duration": None, + "collected": False, "missing_since": None, "success_seen_at": None} + if idx < start_at: + job["status"], job["skip_reason"] = "skipped", "start_at_job" + elif conf["wf_fingerprint"] in by_fp: + if overwrite: + job["replace_ids"] = list(by_fp[conf["wf_fingerprint"]]) + else: + job["status"], job["skip_reason"] = "skipped", "existing" + jobs.append(job) + + sw = Sweep(uuid.uuid4().hex[:12], paths["session"], jobs, plan["templates"], origin_node, + client_id or None, {"image_format": str(settings.get("image_format", "webp")), + "overwrite_existing": overwrite, "start_at_job": start_at, + "total_jobs": len(jobs)}, + list(plan["warnings"]), plan["axes"], bridge or get_bridge()) + if sw.client_id is None: + sw.client_id = f"uscg-sweep-{sw.id}" + return sw + + +def build_job_request(sw, job): + """(prompt, partial_execution_targets) for one job.""" + conf = job["conf"] + tpl = sw.templates[conf["wf_template_hash"]] + job_json = {"conf": conf, "sweep_id": sw.id, "job_index": job["index"], "gen_index": job["gen_index"], + "origin_node": sw.origin_node, "replace_item_ids": job["replace_ids"], + "image_format": sw.settings.get("image_format", "webp")} + prompt, cid = wfs.build_job_prompt(tpl["prompt"], conf["wf_subs"], tpl["output"], + {"session_name": sw.session, + "job_json": json.dumps(job_json, ensure_ascii=False)}) + partial = None if tpl["run_other_outputs"] else [cid] + return prompt, partial + + +async def validate_first(sw): + """Validate the first job that will run (on the event loop). Raises ValueError with node errors.""" + job = next((j for j in sw.jobs if j["status"] == "pending"), None) + if job is None: + return + prompt, partial = build_job_request(sw, job) + ok, error, node_errors = await sw.bridge.validate(prompt, partial) + if not ok: + raise ValueError("The first job is not valid: " + summarize_errors(error, node_errors, prompt)) + + +def _new_meta(sw): + first = sw.jobs[0]["conf"] if sw.jobs else {} + return {"positive": first.get("positive", ""), "negative": first.get("negative", ""), + "model": first.get("model", ""), "seed": first.get("seed", 0) if isinstance(first.get("seed", 0), int) else 0, + "vae_batch_size": 1, "configs_json": "", "resolutions_json": "[]", + "image_format": sw.settings.get("image_format", "webp"), "workflow_sweep": True} + + +def ensure_manifest(sw): + """Create the session manifest when missing (existing manifests are left alone).""" + paths = wfs.session_paths(sw.session) + os.makedirs(paths["images"], exist_ok=True) + if not os.path.exists(paths["manifest"]): + data = {"session_name": sw.session, "meta": _new_meta(sw), "items": []} + with open(paths["manifest"], "w") as f: + json.dump(data, f, indent=4) + return data + data = wfs.read_manifest_items(paths["manifest"]) + return data + + +def launch(sw, start_thread=True): + """Register the sweep and start its feeder thread (tests drive tick() themselves with start_thread=False).""" + with _registry_lock: + busy = _active_sweep() + if busy is not None and busy is not sw: + raise SweepBusy(busy.id) + _sweeps[sw.id] = sw + ensure_manifest(sw) + with sw.lock: + sw.state = "running" + sw.started_at = sw.eta_start = time.time() + print(f"[GridTester] 🧪 Workflow sweep {sw.id}: {sw.to_run} job(s) to run, " + f"{sw.skipped_existing} already in the manifest (session '{sw.session}')") + for w in sw.warnings: + print(f"[GridTester] ⚠️ {w}") + _write_spec(sw) + _progress(sw) + if start_thread: + sw.thread = threading.Thread(target=_run, args=(sw,), name=f"uscg-sweep-{sw.id}", daemon=True) + sw.thread.start() + return sw + + +def start_sweep(configs, session_settings=None, origin_node=None, client_id=None, bridge=None, + class_info=wfs.comfy_class_info): + """Prepare + validate job 1 on the loop + launch. For a NON-loop thread (the Generator node). + Returns immediately after launch; never waits for any sweep prompt.""" + sw = prepare_sweep(configs, session_settings, origin_node, client_id, bridge, class_info) + sw.bridge.run_on_loop(validate_first(sw)) + return launch(sw) + + +# --------------------------------------------------------------------------- feeder thread + +def _run(sw): + try: + while True: + alive = tick(sw) + if not alive: + break + sw.wake.wait(TICK) + sw.wake.clear() + except Exception as e: # never let the feeder die silently + traceback.print_exc() + with sw.lock: + sw.state, sw.reason = "failed", f"internal error: {e}" + sw.finished_at = time.time() + _write_spec(sw) + _progress(sw) + + +def _fail_job(sw, job, status, error, params=None): + job["status"] = status + job["error"] = error + job["finished_at"] = time.time() + sw.last_errors.append({"index": job["index"], "status": status, "error": error, + "params": job["conf"].get("wf_params")}) + del sw.last_errors[:-20] + sw.consecutive_errors += 1 + sw.terminal_since_spec += 1 + print(f"[GridTester] ❌ Sweep job {job['index'] + 1}/{len(sw.jobs)} {status}: {error}") + if sw.consecutive_errors >= BREAKER and sw.state == "running": + sw.state = "paused" + sw.reason = f"{BREAKER} jobs in a row failed; last: {error}" + + +def _classify(history): + """-> (kind, detail) with kind in success|error|interrupted|unknown, from a ComfyUI history entry.""" + status = (history or {}).get("status") or {} + messages = status.get("messages") or [] + names = {} + for m in messages: + if isinstance(m, (list, tuple)) and len(m) == 2: + names[m[0]] = m[1] if isinstance(m[1], dict) else {} + if "execution_interrupted" in names: + return "interrupted", names["execution_interrupted"] + if "execution_error" in names or status.get("status_str") == "error": + d = names.get("execution_error", {}) + text = d.get("exception_message") or "execution error" + if d.get("node_type") or d.get("node_id"): + text = f"{d.get('node_type', '?')} #{d.get('node_id', '?')}: {text}" + return "error", {"error": text.strip()[:1500], "start": names.get("execution_start"), + "end": names.get("execution_error")} + if status.get("status_str") == "success" or "execution_success" in names: + return "success", {"start": names.get("execution_start"), "end": names.get("execution_success")} + return "unknown", {} + + +def _poll(sw, inflight): + b = sw.bridge + now = time.time() + running_ids, queued_ids = b.queue_state() # snapshot FIRST, then history (race-free ordering) + for job in inflight: + pid = job["prompt_id"] + if pid in running_ids: + job["missing_since"] = None + if job["status"] == "enqueued": + job["status"], job["started_at"] = "running", now + continue + if pid in queued_ids: + job["missing_since"] = None + continue + history = b.history_entry(pid) + if history is None: + job["missing_since"] = job["missing_since"] or now + if now - job["missing_since"] > LOST_GRACE: + with sw.lock: + job["status"] = "lost" + job["error"] = "the prompt disappeared from ComfyUI's queue and history (queue cleared?)" + if sw.state == "running": + sw.state, sw.reason = "paused", "a sweep prompt was removed from ComfyUI's queue; resume to re-run it" + print(f"[GridTester] ⚠️ Sweep job {job['index'] + 1} lost; sweep paused") + continue + kind, detail = _classify(history) + with sw.lock: + if kind == "success": + job["success_seen_at"] = job["success_seen_at"] or now + if job["collected"]: + _finish_done(sw, job, detail, now) + elif now - job["success_seen_at"] > COLLECT_GRACE: + _fail_job(sw, job, "error", "the job ran but the collector produced no image") + _after_terminal(sw) + elif kind == "error": + _fail_job(sw, job, "error", detail["error"]) + _after_terminal(sw) + elif kind == "interrupted": + job["status"], job["error"] = "interrupted", "interrupted in ComfyUI" + if sw.state == "running": + sw.state, sw.reason = "paused", "a job was interrupted in ComfyUI; resume to re-run it" + print(f"[GridTester] ⏸️ Sweep job {job['index'] + 1} interrupted; sweep paused") + else: + job["missing_since"] = job["missing_since"] or now + if now - job["missing_since"] > LOST_GRACE: + _fail_job(sw, job, "error", "ComfyUI finished the prompt without a recognisable status") + _after_terminal(sw) + + +def _duration_from(detail, fallback): + try: + start, end = detail.get("start") or {}, detail.get("end") or {} + d = (float(end["timestamp"]) - float(start["timestamp"])) / 1000.0 + if d >= 0: + return d + except (KeyError, TypeError, ValueError): + pass + return fallback + + +def _finish_done(sw, job, detail, now): + job["status"] = "done" + job["finished_at"] = now + fallback = (now - job["started_at"]) if job["started_at"] else 0.0 + job["duration"] = _duration_from(detail, fallback) + sw.durations.append(job["duration"]) + sw.consecutive_errors = 0 + sw.terminal_since_spec += 1 + _after_terminal(sw) + + +def _after_terminal(sw): + _progress(sw) + if sw.terminal_since_spec >= SPEC_EVERY: + sw.terminal_since_spec = 0 + _write_spec(sw) + + +def _enqueue_job(sw, job): + try: + prompt, partial = build_job_request(sw, job) + except Exception as e: + with sw.lock: + _fail_job(sw, job, "invalid", f"could not build the prompt: {e}") + _after_terminal(sw) + return + try: + pid, error, node_errors = sw.bridge.run_on_loop(sw.bridge.enqueue(prompt, partial, sw.client_id)) + except Exception as e: + with sw.lock: + _fail_job(sw, job, "invalid", f"could not queue the prompt: {e}") + _after_terminal(sw) + return + with sw.lock: + if pid is None: + _fail_job(sw, job, "invalid", summarize_errors(error, node_errors, prompt)) + _after_terminal(sw) + else: + job["status"], job["prompt_id"], job["enqueued_at"] = "enqueued", pid, time.time() + job["missing_since"] = None + + +def _do_cleanup(sw): + """Cancel: drop our queued prompts, interrupt the running one, mark the rest cancelled.""" + b = sw.bridge + for job in list(sw.in_flight()): + pid = job["prompt_id"] + try: + if job["status"] == "enqueued" and b.delete_queued(pid): + pass + else: + b.interrupt(pid) + except Exception as e: + print(f"[GridTester] ⚠️ Sweep cancel could not remove {pid}: {e}") + with sw.lock: + for job in sw.jobs: + if job["status"] in ("pending", "enqueued", "running", "interrupted", "lost"): + job["status"] = "cancelled" + sw.needs_cleanup = False + sw.finished_at = time.time() + _write_spec(sw) + _progress(sw) + + +def _finish(sw): + with sw.lock: + errors = sw.count("error", "invalid") + sw.state = "completed" + sw.reason = f"{errors} job(s) failed" if errors else "" + sw.finished_at = time.time() + print(f"[GridTester] 🎉 Workflow sweep {sw.id} complete: {sw.count('done')} image(s), {sw.count('error', 'invalid')} failed") + _write_spec(sw) + _progress(sw, complete=True) + + +def tick(sw): + """One feeder iteration. Returns False when the sweep is over (thread should exit).""" + b = sw.bridge + with sw.lock: + if sw.state in ("completed", "failed") or (sw.state == "cancelled" and not sw.needs_cleanup): + return False + cleanup = sw.state == "cancelled" and sw.needs_cleanup + inflight = sw.in_flight() + if cleanup: + _do_cleanup(sw) + return False + + if inflight: + _poll(sw, inflight) + + with sw.lock: + state = sw.state + queued = [j for j in sw.in_flight() if j["status"] == "enqueued"] + if state == "paused": + for job in queued: # pull what has not started; a started job finishes + try: + pulled = b.delete_queued(job["prompt_id"]) + except Exception: + pulled = False + if pulled: + with sw.lock: + job.update(status="pending", prompt_id=None, enqueued_at=None, missing_since=None) + return True + + while True: + with sw.lock: + if sw.state != "running": + break + nxt = next((j for j in sw.jobs if j["status"] == "pending"), None) + if nxt is None or len(sw.in_flight()) >= WINDOW: + break + _enqueue_job(sw, nxt) + + with sw.lock: + if sw.state != "running": + return True + if not sw.in_flight() and not any(j["status"] == "pending" for j in sw.jobs): + finished = True + else: + finished = False + if finished: + _finish(sw) + return False + return True + + +# --------------------------------------------------------------------------- public controls + +def pause(sweep_id): + sw = _find(sweep_id) + if sw is None: + return False, "unknown sweep" + with sw.lock: + if sw.state != "running": + return False, f"sweep is {sw.state}" + sw.state, sw.reason = "paused", "paused by user" + sw.wake.set() + return True, "paused" + + +def resume(sweep_id): + sw = _find(sweep_id) + if sw is None: + return False, "unknown sweep" + with sw.lock: + if sw.state != "paused": + return False, f"sweep is {sw.state}" + for job in sw.jobs: + if job["status"] in ("interrupted", "lost"): # failed jobs stay failed; these two never really ran + job.update(status="pending", prompt_id=None, error=None, missing_since=None, success_seen_at=None) + sw.consecutive_errors = 0 + sw.state, sw.reason = "running", "" + sw.wake.set() + return True, "running" + + +def cancel(sweep_id): + sw = _find(sweep_id) + if sw is None: + return False, "unknown sweep" + with sw.lock: + if sw.state not in ACTIVE_STATES: + return False, f"sweep is {sw.state}" + sw.state, sw.reason, sw.needs_cleanup = "cancelled", "cancelled by user", True + sw.wake.set() + return True, "cancelled" + + +def on_collected(sweep_id, job_index, items): + """Called by the collector node after it saved a job's image(s). False when the sweep is unknown.""" + sw = _find(sweep_id) + if sw is None: + return False + try: + job = sw.jobs[int(job_index)] + except (IndexError, TypeError, ValueError): + return False + with sw.lock: + job["collected"] = True + job["items"] = [it.get("id") for it in (items or []) if isinstance(it, dict)] + return True + + +def job_started_at(sweep_id, job_index): + sw = _find(sweep_id) + if sw is None: + return None + try: + return sw.jobs[int(job_index)]["started_at"] + except (IndexError, TypeError, ValueError): + return None + + +def get_status(sweep_id=None, session=None): + sw = _find(sweep_id, session) + if sw is None: + return None + with sw.lock: + done = sw.count("done") + errors = sw.count("error", "invalid") + recent = sw.durations[-10:] + avg = sum(recent) / len(recent) if recent else 0.0 + remaining = sw.count("pending", "enqueued", "running", "interrupted", "lost") + return {"sweep_id": sw.id, "session_name": sw.session, "state": sw.state, "reason": sw.reason, + "total": len(sw.jobs), "to_run": sw.to_run, "skipped_existing": sw.skipped_existing, + "done": done, "errors": errors, "in_flight": len(sw.in_flight()), "pending": sw.count("pending"), + "eta_str": _eta_str(avg * remaining) if recent and sw.state == "running" else "", + "avg_duration": round(avg, 1), "warnings": list(sw.warnings), + "last_errors": list(sw.last_errors)} + + +# --------------------------------------------------------------------------- progress + spec file + +def _eta_str(seconds): + seconds = int(max(0, seconds)) + h, m, s = seconds // 3600, (seconds % 3600) // 60, seconds % 60 + return f"{h}h {m}m {s}s" if h else (f"{m}m {s}s" if m else f"{s}s") + + +def _progress(sw, complete=False): + """ultimate_grid.progress for the Dashboard's ETA bar (phase 'sweep').""" + try: + with sw.lock: + total = len(sw.jobs) + done = sw.count("done") + errors = sw.count("error", "invalid") + current = sw.skipped_existing + done + remaining = sw.count("pending", "enqueued", "running", "interrupted", "lost") + recent = sw.durations[-10:] + avg = sum(recent) / len(recent) if recent else 0.0 + payload = { + "node": sw.origin_node, "session_name": sw.session, "phase": "sweep", "sweep_id": sw.id, + "state": sw.state, "errors": errors, "current_job": current, "total_jobs": total, + "progress_pct": int((current + errors) / total * 100) if total else 100, + "eta_str": _eta_str(avg * remaining) if recent and remaining else "—", + "finish_time": time.strftime("%H:%M:%S", time.localtime(time.time() + avg * remaining)) if recent and remaining else "", + "avg_duration": round(avg, 1), "last_duration": round(sw.durations[-1], 1) if sw.durations else 0, + } + if complete: + elapsed = time.time() - (sw.eta_start or sw.created_at) + payload.update({"complete": True, "eta_str": "Done", "progress_pct": 100, + "total_elapsed": f"{int(elapsed // 3600)}h {int((elapsed % 3600) // 60)}m {int(elapsed % 60)}s", + "total_generated": done}) + sw.bridge.send_sync("ultimate_grid.progress", payload) + except Exception: + pass + + +def _write_spec(sw): + """benchmarks//workflow_sweeps/.json: templates, bindings and the job list, so any + executed prompt can be reconstructed. Never served by any route.""" + try: + paths = wfs.session_paths(sw.session) + os.makedirs(paths["sweeps"], exist_ok=True) + with sw.lock: + spec = { + "sweep_id": sw.id, "session_name": sw.session, "state": sw.state, "reason": sw.reason, + "created_at": sw.created_at, "settings": sw.settings, "axes": [a for c in sw.axes for a in c["axes"]], "warnings": sw.warnings, + "templates": {h: {"prompt": t["prompt"], "output": t["output"], + "run_other_outputs": t["run_other_outputs"]} for h, t in sw.templates.items()}, + "jobs": [{"index": j["index"], "fingerprint": j["fingerprint"], "status": j["status"], + "prompt_id": j["prompt_id"], "error": j["error"], "duration": j["duration"], + "template": j["conf"].get("wf_template_hash"), "subs": j["conf"].get("wf_subs"), + "params": j["conf"].get("wf_params")} for j in sw.jobs], + } + path = os.path.join(paths["sweeps"], f"{sw.id}.json") + with open(path + ".tmp", "w") as f: + json.dump(spec, f, indent=1) + os.replace(path + ".tmp", path) + except Exception as e: + print(f"[GridTester] ⚠️ Could not write the sweep spec file: {e}") + + +# --------------------------------------------------------------------------- request-body helpers (routes) + +def configs_from_body(body): + """(configs list, session_settings) from {state} (Builder state, through the single source of truth) + or {configs_json}. Raises ValueError.""" + if not isinstance(body, dict): + raise ValueError("The request body must be a JSON object.") + if isinstance(body.get("state"), dict): + try: + from .config_builder_node import UltimateConfigBuilder + except ImportError: + from config_builder_node import UltimateConfigBuilder + raw = UltimateConfigBuilder.state_to_configs_json(body["state"]) + elif body.get("configs_json") is not None: + raw = body["configs_json"] + else: + raise ValueError("Send {state} (the Builder state) or {configs_json}.") + parsed = json.loads(raw) if isinstance(raw, (str, bytes)) else raw + if isinstance(parsed, dict): + return parsed.get("configs") or [], parsed.get("_session_settings") or {} + return parsed, {} + + +async def dry_run(body, bridge=None, class_info=wfs.comfy_class_info): + """Concrete first prompts (with collector) + validation of the first job, nothing is queued.""" + configs, settings = configs_from_body(body) + try: + limit = max(1, min(20, int(body.get("limit", 3)))) + except (TypeError, ValueError): + limit = 3 + sw = prepare_sweep(configs, settings, None, None, bridge, class_info, check_busy=False) + jobs_out = [] + for job in sw.jobs[:limit]: + prompt, partial = build_job_request(sw, job) + jobs_out.append({"index": job["index"], "fingerprint": job["fingerprint"], "skip": job["status"] == "skipped", + "params": job["conf"].get("wf_params"), "partial_execution_targets": partial, + "prompt": prompt}) + validation = {"ok": True, "error": None, "node_errors": {}} + first = next((j for j in sw.jobs if j["status"] == "pending"), None) + if first is not None: + prompt, partial = build_job_request(sw, first) + try: + ok, error, node_errors = await sw.bridge.validate(prompt, partial) + validation = {"ok": ok, "error": summarize_errors(error, node_errors, prompt) if not ok else None, + "node_errors": node_errors if not ok else {}} + except Exception as e: + validation = {"ok": False, "error": f"validation could not run: {e}", "node_errors": {}} + return {"total": len(sw.jobs), "to_run": sw.to_run, "skipped_existing": sw.skipped_existing, + "axes": [a for c in sw.axes for a in c["axes"]], "warnings": sw.warnings, "validation": validation, "jobs": jobs_out} + + +async def api_start(body, bridge=None, class_info=wfs.comfy_class_info): + """POST /config_tester/sweep/start. Returns (http_status, json_body). Runs on the event loop.""" + try: + configs, settings = configs_from_body(body) + sw = prepare_sweep(configs, settings, body.get("node_id"), body.get("client_id"), bridge, class_info) + await validate_first(sw) + launch(sw) + except SweepBusy as e: + return 409, {"error": str(e), "sweep_id": e.sweep_id} + except ValueError as e: + return 400, {"error": str(e)} + return 200, {"sweep_id": sw.id, "total": len(sw.jobs), "to_run": sw.to_run, + "skipped_existing": sw.skipped_existing, "warnings": sw.warnings}