From e55f20e2440ea13469a12f8a29617e4e056d521c Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Tue, 14 Jul 2026 13:23:41 -0700 Subject: [PATCH 1/2] fix(run): resolve the bundled default checkpoint at runtime (BE-2994) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `comfy run --prompt` (zero-config) hard-pinned one checkpoint and failed wherever it was absent — Comfy Cloud and fresh local installs. - Fix the stale pin: node "4" ckpt_name is now v1-5-pruned-emaonly-fp16.safetensors (the gallery/tutorial SD1.5 file). - Add a pure, offline-testable `resolve_default_checkpoint(workflow, object_info)` in cql/default_workflow.py: if the pinned checkpoint is in the target's CheckpointLoaderSimple enum → no change; else if the enum is non-empty → substitute the first available checkpoint + emit a note; else (positively-empty enum) → flag no_checkpoint; enum absent / object_info {} → fail open (unchanged behavior). - Wire it into both run paths (local + cloud) after object_info is fetched and before preflight, guarded to the bundled default graph and skipped when the user pinned the checkpoint via --set (honored verbatim). Thread a checkpoint_user_set flag through the `preloaded` tuple. - Emit a `checkpoint_substituted` event (--json) / pretty notice on substitution; raise a new `no_checkpoint_available` error with an actionable download/cloud hint when the target has zero checkpoints. --- comfy_cli/cmdline.py | 13 ++- comfy_cli/command/run/__init__.py | 30 ++++- comfy_cli/command/run/preflight.py | 51 ++++++++ comfy_cli/cql/data/default_text2img.json | 2 +- comfy_cli/cql/default_workflow.py | 115 +++++++++++++++++++ comfy_cli/error_codes.py | 10 ++ tests/comfy_cli/command/test_run_prompt.py | 92 ++++++++++++++- tests/comfy_cli/cql/test_default_workflow.py | 98 +++++++++++++++- 8 files changed, 396 insertions(+), 15 deletions(-) diff --git a/comfy_cli/cmdline.py b/comfy_cli/cmdline.py index 4662b788..5b6a8df0 100644 --- a/comfy_cli/cmdline.py +++ b/comfy_cli/cmdline.py @@ -822,7 +822,7 @@ def run( # against OUR pinned node ids, so mixing them with a user --workflow — # whose node ids are arbitrary — is rejected rather than silently # misapplied. `preloaded` is handed straight to run's execute path. - preloaded: tuple[dict, str, bool] | None = None + preloaded: tuple[dict, str, bool, bool] | None = None if prompt is not None or set_overrides: if workflow is not None: renderer.error( @@ -831,14 +831,21 @@ def run( hint="drop --workflow to use the bundled default, or edit the workflow file directly", ) raise typer.Exit(code=1) - from comfy_cli.cql.default_workflow import PromptInjectionError, build_default_workflow + from comfy_cli.cql.default_workflow import ( + PromptInjectionError, + build_default_workflow, + overrides_set_checkpoint, + ) try: injected = build_default_workflow(prompt=prompt, overrides=set_overrides) except PromptInjectionError as e: renderer.error(code=e.code, message=str(e), hint=e.hint) raise typer.Exit(code=1) from e - preloaded = (injected, "default_text2img", False) + # If the user pinned the checkpoint (--set checkpoint=… / 4.ckpt_name=…), + # honor it verbatim: runtime resolution is skipped downstream. + checkpoint_user_set = overrides_set_checkpoint(set_overrides, injected) + preloaded = (injected, "default_text2img", False, checkpoint_user_set) elif workflow is None: renderer.error( code="prompt_rejected", diff --git a/comfy_cli/command/run/__init__.py b/comfy_cli/command/run/__init__.py index fc230d00..13bfbc8b 100644 --- a/comfy_cli/command/run/__init__.py +++ b/comfy_cli/command/run/__init__.py @@ -38,6 +38,7 @@ from comfy_cli.command.run.preflight import _detect_partner_nodes as _detect_partner_nodes from comfy_cli.command.run.preflight import _fetch_object_info as _fetch_object_info from comfy_cli.command.run.preflight import _preflight_validate as _preflight_validate +from comfy_cli.command.run.preflight import _resolve_default_checkpoint_or_exit as _resolve_default_checkpoint_or_exit from comfy_cli.command.run.preflight import fetch_object_info as fetch_object_info from comfy_cli.command.run.watcher import _spawn_watcher as _spawn_watcher from comfy_cli.command.run.watcher import _tail_state_file as _tail_state_file @@ -104,7 +105,7 @@ def execute( notify: bool = False, api_key: str | None = None, print_prompt: bool = False, - preloaded: tuple[dict, str, bool] | None = None, + preloaded: tuple[dict, str, bool, bool] | None = None, ): # `0.0.0.0` is a wildcard bind, not a connect address. macOS / Windows # clients can't reach it; on Linux it happens to resolve to a loopback. @@ -123,10 +124,13 @@ def execute( # `preloaded` short-circuits file loading: an in-memory API-format graph # (e.g. the `comfy run --prompt` injected default) is handed straight in as - # (workflow_dict, display_name, is_ui). Everything downstream is unchanged. + # (workflow_dict, display_name, is_ui, checkpoint_user_set). Everything + # downstream is unchanged; `checkpoint_user_set` gates runtime checkpoint + # resolution for the bundled default (skip it when the user pinned one). if preloaded is not None: - raw_workflow, workflow_name, is_ui = preloaded + raw_workflow, workflow_name, is_ui, checkpoint_user_set = preloaded else: + checkpoint_user_set = False try: raw_workflow, workflow_name, is_ui = _load_workflow_file(workflow) except WorkflowLoadError as e: @@ -214,6 +218,14 @@ def execute( # extra_data so the partner node finds it server-side — same shape # the cloud submit path uses. object_info = _fetch_object_info(host, port) + + # Runtime checkpoint resolution for the bundled `--prompt` default: swap the + # pinned checkpoint for one the local server actually has (or hard-error if + # it has none). Guarded to the bundled default graph and skipped when the + # user pinned the checkpoint explicitly (honor it; let preflight reject it). + if preloaded is not None and workflow_name == "default_text2img" and not checkpoint_user_set: + _resolve_default_checkpoint_or_exit(renderer, workflow, object_info, where="local") + partner_nodes = _detect_partner_nodes(workflow, object_info) extra_data: dict | None = None if api_key: @@ -493,7 +505,7 @@ def execute_cloud( timeout: int = 600, notify: bool = False, print_prompt: bool = False, - preloaded: tuple[dict, str, bool] | None = None, + preloaded: tuple[dict, str, bool, bool] | None = None, ): """Run a workflow against Comfy Cloud via the stored OAuth session. @@ -508,8 +520,9 @@ def execute_cloud( renderer = get_renderer() if preloaded is not None: - raw_workflow, workflow_name, is_ui = preloaded + raw_workflow, workflow_name, is_ui, checkpoint_user_set = preloaded else: + checkpoint_user_set = False try: raw_workflow, workflow_name, is_ui = _load_workflow_file(workflow) except WorkflowLoadError as e: @@ -593,6 +606,13 @@ def execute_cloud( except Exception: # noqa: BLE001 cloud_object_info = {} + # Runtime checkpoint resolution for the bundled `--prompt` default (mirrors + # the local path): swap the pinned checkpoint for one Comfy Cloud actually + # has, or hard-error if it enumerates none. Guarded to the bundled default + # and skipped when the user pinned the checkpoint explicitly. + if preloaded is not None and workflow_name == "default_text2img" and not checkpoint_user_set: + _resolve_default_checkpoint_or_exit(renderer, parsed_workflow, cloud_object_info, where="cloud") + _preflight_validate(renderer, parsed_workflow, cloud_object_info, target_label="cloud") target = resolve_target(where="cloud") diff --git a/comfy_cli/command/run/preflight.py b/comfy_cli/command/run/preflight.py index 36fcf737..dfed8908 100644 --- a/comfy_cli/command/run/preflight.py +++ b/comfy_cli/command/run/preflight.py @@ -111,6 +111,57 @@ def _preflight_validate(renderer, workflow: dict, object_info: dict, *, target_l pprint(f"[yellow]⚠ {w.get('field', '?')}: {w.get('message', '')}[/yellow]") +def _resolve_default_checkpoint_or_exit(renderer, workflow: dict, object_info: dict, *, where: str) -> None: + """Runtime-resolve the bundled default's pinned checkpoint against the + target, in place, then report the outcome through the renderer. + + Call ONLY for the bundled default graph (``workflow_name == + "default_text2img"``) when the user did NOT explicitly ``--set`` the + checkpoint. Three outcomes: + + - pinned present / can't tell (object_info empty or not enumerated) → no-op + (fail open — preflight + the server decide); + - pinned absent but the target has ≥1 checkpoint → substitute the first + available one and emit a ``checkpoint_substituted`` note; + - target positively has zero checkpoints → hard ``no_checkpoint_available`` + error (exit 1) instead of a cryptic server-side reject. + + ``where`` is ``"local"`` or ``"cloud"`` and drives the target label + hint. + """ + from comfy_cli.cql.default_workflow import resolve_default_checkpoint + + target_label = "the local server" if where == "local" else "Comfy Cloud" + _, res = resolve_default_checkpoint(workflow, object_info, target=target_label) + + if res.no_checkpoint: + if where == "local": + hint = ( + "download a checkpoint, e.g. `comfy model download --url " + "https://huggingface.co/Comfy-Org/stable-diffusion-v1-5-archive/resolve/main/" + "v1-5-pruned-emaonly-fp16.safetensors`, then re-run — or `--set checkpoint=`" + ) + else: + hint = ( + "run a published gallery template (Comfy Cloud provisions its models), or " + "`--set checkpoint=` once a checkpoint is available on the target" + ) + renderer.error( + code="no_checkpoint_available", + message=( + f"the bundled default text2img workflow needs a checkpoint, but {target_label} has none installed" + ), + hint=hint, + details={"where": where}, + ) + raise typer.Exit(code=1) + + if res.note: + # Event fires in NDJSON/stream mode only; the pretty line covers humans. + renderer.event("checkpoint_substituted", message=res.note, checkpoint=res.substituted_to, where=where) + if renderer.is_pretty(): + pprint(f"[yellow]⚠ {res.note}[/yellow]") + + def _fetch_object_info(host: str, port: int) -> dict: """Fetch object_info for partner-node detection + validation. Fail open.""" try: diff --git a/comfy_cli/cql/data/default_text2img.json b/comfy_cli/cql/data/default_text2img.json index e42de994..6cc511b1 100644 --- a/comfy_cli/cql/data/default_text2img.json +++ b/comfy_cli/cql/data/default_text2img.json @@ -18,7 +18,7 @@ "4": { "class_type": "CheckpointLoaderSimple", "inputs": { - "ckpt_name": "v1-5-pruned-emaonly.ckpt" + "ckpt_name": "v1-5-pruned-emaonly-fp16.safetensors" }, "_meta": {"title": "Load Checkpoint"} }, diff --git a/comfy_cli/cql/default_workflow.py b/comfy_cli/cql/default_workflow.py index 173471ba..09a8fab8 100644 --- a/comfy_cli/cql/default_workflow.py +++ b/comfy_cli/cql/default_workflow.py @@ -16,8 +16,14 @@ import json import math +from dataclasses import dataclass from importlib import resources +# The name pinned in ``data/default_text2img.json`` node "4" ``ckpt_name``. +# Runtime resolution (``resolve_default_checkpoint``) swaps this for a checkpoint +# the target actually has when it's absent; keep the two in sync. +DEFAULT_CHECKPOINT_NAME = "v1-5-pruned-emaonly-fp16.safetensors" + # -- Pinned node ids (must match data/default_text2img.json) -- CHECKPOINT_LOADER_ID = "4" POSITIVE_PROMPT_ID = "6" @@ -219,3 +225,112 @@ def build_default_workflow(*, prompt: str | None = None, overrides: list[str] | _apply_set(workflow, node_id, field, value) return workflow + + +def overrides_set_checkpoint(overrides: list[str] | None, workflow: dict) -> bool: + """True if any ``--set`` override targets the checkpoint (node ``"4"`` + ``ckpt_name``), via an alias (``checkpoint``/``ckpt``) or the raw + ``4.ckpt_name`` form. + + Used by the caller to decide whether the user pinned the checkpoint + explicitly — if so, runtime resolution is skipped and the value is honored + verbatim. ``workflow`` must be a built default graph so addresses resolve; + malformed entries are ignored here (``build_default_workflow`` already + validated/rejected them upstream). + """ + for raw in overrides or []: + if "=" not in raw: + continue + address = raw.partition("=")[0].strip() + try: + node_id, field = _resolve_address(address, workflow) + except PromptInjectionError: + continue + if (node_id, field) == (CHECKPOINT_LOADER_ID, "ckpt_name"): + return True + return False + + +@dataclass(frozen=True) +class CheckpointResolution: + """Outcome of :func:`resolve_default_checkpoint`. + + - ``note``/``substituted_to`` are set only when the pinned checkpoint was + absent and a different one was substituted. + - ``no_checkpoint`` is True ONLY when ``object_info`` positively enumerated + an EMPTY checkpoint list (the target has zero checkpoints). It stays False + when we can't tell (``object_info`` absent/empty, or it didn't enumerate + ``CheckpointLoaderSimple.ckpt_name``) so callers fail open there. + """ + + note: str | None = None + substituted_to: str | None = None + no_checkpoint: bool = False + + +def _checkpoint_enum(object_info: dict) -> list | None: + """Return the ``CheckpointLoaderSimple.ckpt_name`` option list from + ``object_info``, or ``None`` when it isn't enumerated at all. + + The raw object_info shape is ``ckpt_name: [[, …], {opts}]`` — the + option list is element 0. ``None`` (not an empty list) means "can't tell" + so the caller can distinguish a positively-empty enum from an absent one. + """ + node = object_info.get("CheckpointLoaderSimple") + if not isinstance(node, dict): + return None + inp = node.get("input") + if not isinstance(inp, dict): + return None + req = inp.get("required") + if not isinstance(req, dict): + return None + spec = req.get("ckpt_name") + if isinstance(spec, list) and spec and isinstance(spec[0], list): + return spec[0] + return None + + +def resolve_default_checkpoint( + workflow: dict, object_info: dict, *, target: str = "the server" +) -> tuple[dict, CheckpointResolution]: + """Resolve the bundled default's pinned checkpoint against a live target. + + Pure and offline-testable. Given the bundled default graph and a target's + ``object_info``, decide what checkpoint node ``"4"`` should carry: + + - pinned name present in the target's enum → no change; + - pinned absent but the enum is non-empty → substitute the first available + checkpoint and return a ``note``; + - enum positively empty → leave unchanged, flag ``no_checkpoint`` (the + caller emits an actionable error); + - enum absent / ``object_info`` empty → leave unchanged, no flag (fail open). + + Mutates ``workflow`` in place on substitution and returns it alongside the + :class:`CheckpointResolution`. Callers must guard this to the bundled + default graph only (``workflow_name == "default_text2img"``). + """ + node = workflow.get(CHECKPOINT_LOADER_ID) + inputs = node.get("inputs") if isinstance(node, dict) else None + if not isinstance(inputs, dict): + return workflow, CheckpointResolution() + + enum = _checkpoint_enum(object_info) + if enum is None: + # Not enumerated (fresh/unfetched object_info) — fail open. + return workflow, CheckpointResolution() + if not enum: + # Positively empty: the target has zero checkpoints installed. + return workflow, CheckpointResolution(no_checkpoint=True) + + pinned = inputs.get("ckpt_name") + if pinned in enum: + return workflow, CheckpointResolution() + + replacement = enum[0] + inputs["ckpt_name"] = replacement + note = ( + f"default checkpoint {pinned} not found on {target}; using {replacement} " + f"instead (override with --set checkpoint=)" + ) + return workflow, CheckpointResolution(note=note, substituted_to=replacement) diff --git a/comfy_cli/error_codes.py b/comfy_cli/error_codes.py index 4cba21d3..5039babe 100644 --- a/comfy_cli/error_codes.py +++ b/comfy_cli/error_codes.py @@ -160,6 +160,16 @@ class ErrorCode: "(missing or corrupt package data). A packaging fault, not user input.", "reinstall comfy-cli", ), + ErrorCode( + "no_checkpoint_available", + "`comfy run --prompt`/`--set` (bundled default text2img) needs a checkpoint, but the " + "target positively enumerated ZERO installed checkpoints. Only raised when object_info " + "was fetched and its checkpoint list is empty — never when object_info couldn't be fetched " + "(that path fails open and submits).", + "install a checkpoint (local: `comfy model download --url `; cloud: run a " + "published gallery template, which provisions models), then re-run — or `--set " + "checkpoint=` once one is available", + ), ErrorCode( "conversion_error", "UI-format workflow could not be converted to API format.", diff --git a/tests/comfy_cli/command/test_run_prompt.py b/tests/comfy_cli/command/test_run_prompt.py index 37690081..d7d68e38 100644 --- a/tests/comfy_cli/command/test_run_prompt.py +++ b/tests/comfy_cli/command/test_run_prompt.py @@ -16,7 +16,16 @@ from comfy_cli.cmdline import run as run_command from comfy_cli.command.run import execute -from comfy_cli.cql.default_workflow import POSITIVE_PROMPT_ID +from comfy_cli.cql.default_workflow import ( + CHECKPOINT_LOADER_ID, + DEFAULT_CHECKPOINT_NAME, + POSITIVE_PROMPT_ID, + build_default_workflow, +) + + +def _object_info_with_checkpoints(names): + return {"CheckpointLoaderSimple": {"input": {"required": {"ckpt_name": [list(names), {}]}}}} class TestExecuteSubmitsPreloadedGraph: @@ -43,7 +52,7 @@ def test_preloaded_graph_is_submitted(self): port=8188, wait=True, timeout=30, - preloaded=(injected, "default_text2img", False), + preloaded=(injected, "default_text2img", False, False), ) # The exact injected graph is what got handed to the submit path. @@ -67,7 +76,7 @@ def test_preloaded_skips_file_loading(self): port=8188, wait=True, timeout=30, - preloaded=(injected, "default_text2img", False), + preloaded=(injected, "default_text2img", False, False), ) MockExec.assert_called_once() @@ -90,15 +99,27 @@ def test_prompt_builds_injected_graph_and_forwards(self): mock_exec.assert_called_once() preloaded = mock_exec.call_args.kwargs["preloaded"] assert preloaded is not None - graph, name, is_ui = preloaded + graph, name, is_ui, checkpoint_user_set = preloaded assert is_ui is False assert name == "default_text2img" + assert checkpoint_user_set is False assert graph[POSITIVE_PROMPT_ID]["inputs"]["text"] == "a red fox in snow" def test_set_checkpoint_override_forwarded(self): mock_exec, _ = self._call_run(prompt="fox", set_overrides=["checkpoint=sd_xl.safetensors"]) - graph = mock_exec.call_args.kwargs["preloaded"][0] + preloaded = mock_exec.call_args.kwargs["preloaded"] + graph = preloaded[0] assert graph["4"]["inputs"]["ckpt_name"] == "sd_xl.safetensors" + # A user-pinned checkpoint flips the flag so runtime resolution is skipped. + assert preloaded[3] is True + + def test_set_checkpoint_raw_form_flags_user_set(self): + mock_exec, _ = self._call_run(prompt="fox", set_overrides=["4.ckpt_name=sd_xl.safetensors"]) + assert mock_exec.call_args.kwargs["preloaded"][3] is True + + def test_non_checkpoint_set_leaves_flag_false(self): + mock_exec, _ = self._call_run(prompt="fox", set_overrides=["seed=42"]) + assert mock_exec.call_args.kwargs["preloaded"][3] is False def test_workflow_path_forwards_no_preloaded(self): mock_exec, _ = self._call_run(workflow="wf.json") @@ -131,3 +152,64 @@ def test_bad_set_address_is_rejected(self): with pytest.raises(typer.Exit) as e: self._call_run(prompt="fox", set_overrides=["bogus=1"]) assert e.value.exit_code == 1 + + +class TestRuntimeCheckpointResolutionLocal: + """Wiring: `execute` resolves the bundled default's checkpoint against the + server's object_info before submit (BE-2994).""" + + def _run_local(self, preloaded, object_info, patch_pprint=False): + stack = [ + patch("comfy_cli.command.run.check_comfy_server_running", return_value=True), + patch("comfy_cli.command.run._fetch_object_info", return_value=object_info), + # Isolate resolution from the unrelated class_type validation the + # bundled graph would otherwise trip against a stub object_info. + patch("comfy_cli.command.run._preflight_validate"), + patch("comfy_cli.command.run.ExecutionProgress"), + patch("comfy_cli.command.run.WorkflowExecution"), + ] + with stack[0], stack[1], stack[2], stack[3], stack[4] as MockExec: + MockExec.return_value = MagicMock(outputs=[]) + if patch_pprint: + with patch("comfy_cli.command.run.preflight.pprint") as mock_pprint: + execute(None, host="127.0.0.1", port=8188, wait=True, timeout=30, preloaded=preloaded) + return MockExec, mock_pprint + execute(None, host="127.0.0.1", port=8188, wait=True, timeout=30, preloaded=preloaded) + return MockExec, None + + def _default_preloaded(self, *, checkpoint_user_set=False): + return (build_default_workflow(prompt="fox"), "default_text2img", False, checkpoint_user_set) + + def test_absent_pinned_is_substituted(self): + oi = _object_info_with_checkpoints(["dreamshaper.safetensors", "sd_xl.safetensors"]) + MockExec, mock_pprint = self._run_local(self._default_preloaded(), oi, patch_pprint=True) + submitted = MockExec.call_args.args[0] + assert submitted[CHECKPOINT_LOADER_ID]["inputs"]["ckpt_name"] == "dreamshaper.safetensors" + # A human-facing substitution notice is printed. + assert any("dreamshaper.safetensors" in str(c.args[0]) for c in mock_pprint.call_args_list) + + def test_present_pinned_is_unchanged(self): + oi = _object_info_with_checkpoints(["other.safetensors", DEFAULT_CHECKPOINT_NAME]) + MockExec, _ = self._run_local(self._default_preloaded(), oi) + submitted = MockExec.call_args.args[0] + assert submitted[CHECKPOINT_LOADER_ID]["inputs"]["ckpt_name"] == DEFAULT_CHECKPOINT_NAME + + def test_empty_enum_errors_no_checkpoint_available(self): + oi = _object_info_with_checkpoints([]) + with pytest.raises(typer.Exit) as e: + self._run_local(self._default_preloaded(), oi) + assert e.value.exit_code == 1 + + def test_empty_object_info_fails_open_and_submits(self): + MockExec, _ = self._run_local(self._default_preloaded(), {}) + # No error; the pinned default is submitted as-is. + submitted = MockExec.call_args.args[0] + assert submitted[CHECKPOINT_LOADER_ID]["inputs"]["ckpt_name"] == DEFAULT_CHECKPOINT_NAME + + def test_user_pinned_checkpoint_is_never_substituted(self): + graph = build_default_workflow(prompt="fox", overrides=["checkpoint=userpick.safetensors"]) + preloaded = (graph, "default_text2img", False, True) # checkpoint_user_set=True + oi = _object_info_with_checkpoints(["dreamshaper.safetensors"]) + MockExec, _ = self._run_local(preloaded, oi) + submitted = MockExec.call_args.args[0] + assert submitted[CHECKPOINT_LOADER_ID]["inputs"]["ckpt_name"] == "userpick.safetensors" diff --git a/tests/comfy_cli/cql/test_default_workflow.py b/tests/comfy_cli/cql/test_default_workflow.py index 7f1fa9b1..cad89cb9 100644 --- a/tests/comfy_cli/cql/test_default_workflow.py +++ b/tests/comfy_cli/cql/test_default_workflow.py @@ -10,11 +10,14 @@ from comfy_cli.cql.default_workflow import ( CHECKPOINT_LOADER_ID, + DEFAULT_CHECKPOINT_NAME, NEGATIVE_PROMPT_ID, POSITIVE_PROMPT_ID, PromptInjectionError, build_default_workflow, load_default_workflow, + overrides_set_checkpoint, + resolve_default_checkpoint, ) @@ -64,7 +67,12 @@ def test_prompt_sets_positive_text_only(self): def test_no_prompt_no_overrides_is_default_graph(self): wf = build_default_workflow() assert wf[POSITIVE_PROMPT_ID]["inputs"]["text"] == "" - assert wf[CHECKPOINT_LOADER_ID]["inputs"]["ckpt_name"] == "v1-5-pruned-emaonly.ckpt" + assert wf[CHECKPOINT_LOADER_ID]["inputs"]["ckpt_name"] == DEFAULT_CHECKPOINT_NAME + + def test_default_checkpoint_constant_matches_bundle(self): + # The runtime-resolution constant must stay in sync with the JSON pin. + assert DEFAULT_CHECKPOINT_NAME == "v1-5-pruned-emaonly-fp16.safetensors" + assert load_default_workflow()[CHECKPOINT_LOADER_ID]["inputs"]["ckpt_name"] == DEFAULT_CHECKPOINT_NAME class TestSetOverrides: @@ -182,3 +190,91 @@ def test_corrupt_bundle_surfaces_controlled_error(self, monkeypatch): with pytest.raises(PromptInjectionError) as e: load_default_workflow() assert e.value.code == "default_workflow_unavailable" + + +def _object_info_with_checkpoints(names: list[str]) -> dict: + """Minimal object_info enumerating a CheckpointLoaderSimple ckpt_name enum, + in the raw ``[[], {opts}]`` server shape.""" + return {"CheckpointLoaderSimple": {"input": {"required": {"ckpt_name": [list(names), {}]}}}} + + +class TestResolveDefaultCheckpoint: + """Runtime checkpoint resolution (pure, offline) for the bundled default.""" + + def test_pinned_present_no_change(self): + wf = build_default_workflow(prompt="fox") + oi = _object_info_with_checkpoints(["other.safetensors", DEFAULT_CHECKPOINT_NAME]) + out, res = resolve_default_checkpoint(wf, oi, target="the local server") + assert out[CHECKPOINT_LOADER_ID]["inputs"]["ckpt_name"] == DEFAULT_CHECKPOINT_NAME + assert res.note is None + assert res.substituted_to is None + assert res.no_checkpoint is False + + def test_pinned_absent_substitutes_first_available(self): + wf = build_default_workflow(prompt="fox") + oi = _object_info_with_checkpoints(["dreamshaper.safetensors", "sd_xl.safetensors"]) + out, res = resolve_default_checkpoint(wf, oi, target="the local server") + assert out[CHECKPOINT_LOADER_ID]["inputs"]["ckpt_name"] == "dreamshaper.safetensors" + assert res.substituted_to == "dreamshaper.safetensors" + assert res.no_checkpoint is False + assert "not found on the local server" in res.note + assert "dreamshaper.safetensors" in res.note + assert "--set checkpoint=" in res.note + + def test_empty_enum_flags_no_checkpoint(self): + wf = build_default_workflow(prompt="fox") + oi = _object_info_with_checkpoints([]) + out, res = resolve_default_checkpoint(wf, oi) + # Unchanged; the caller emits the actionable error. + assert out[CHECKPOINT_LOADER_ID]["inputs"]["ckpt_name"] == DEFAULT_CHECKPOINT_NAME + assert res.no_checkpoint is True + assert res.note is None + + def test_empty_object_info_fails_open(self): + wf = build_default_workflow(prompt="fox") + out, res = resolve_default_checkpoint(wf, {}) + assert out[CHECKPOINT_LOADER_ID]["inputs"]["ckpt_name"] == DEFAULT_CHECKPOINT_NAME + assert res.no_checkpoint is False + assert res.note is None + + def test_object_info_without_checkpoint_node_fails_open(self): + # object_info present but no CheckpointLoaderSimple → can't tell → fail open. + wf = build_default_workflow(prompt="fox") + out, res = resolve_default_checkpoint(wf, {"KSampler": {"input": {"required": {}}}}) + assert res.no_checkpoint is False + assert res.note is None + + def test_missing_checkpoint_node_in_graph_is_noop(self): + # A caller misuse (non-default graph) must not crash. + _, res = resolve_default_checkpoint( + {"9": {"class_type": "SaveImage", "inputs": {}}}, _object_info_with_checkpoints(["a.safetensors"]) + ) + assert res.no_checkpoint is False + assert res.note is None + + +class TestOverridesSetCheckpoint: + def test_alias_detected(self): + wf = build_default_workflow() + assert overrides_set_checkpoint(["checkpoint=x.safetensors"], wf) is True + + def test_ckpt_alias_detected(self): + wf = build_default_workflow() + assert overrides_set_checkpoint(["ckpt=x.safetensors"], wf) is True + + def test_raw_form_detected(self): + wf = build_default_workflow() + assert overrides_set_checkpoint(["4.ckpt_name=x.safetensors"], wf) is True + + def test_non_checkpoint_override_not_detected(self): + wf = build_default_workflow() + assert overrides_set_checkpoint(["seed=42", "cfg=7.5"], wf) is False + + def test_none_and_empty(self): + wf = build_default_workflow() + assert overrides_set_checkpoint(None, wf) is False + assert overrides_set_checkpoint([], wf) is False + + def test_malformed_entry_ignored(self): + wf = build_default_workflow() + assert overrides_set_checkpoint(["bogus", "seed=42"], wf) is False From 93cfdfb0973dc654b2c18c53b566d7039d2166f4 Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Tue, 14 Jul 2026 14:08:34 -0700 Subject: [PATCH 2/2] fix(run): harden default-checkpoint resolution per cursor review (BE-2994) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address cursor-review findings on the runtime checkpoint resolver: - Cloud fails open on an empty checkpoint enum instead of hard-erroring no_checkpoint_available: Comfy Cloud provisions models per-job, so a zero-length cached enum can't prove the run would fail. Only the local path (enum reflects what's installed) treats empty as a hard stop. - Resolve the bundled default's checkpoint BEFORE emitting the prompt_preview event in the submit flow so the streamed audit trail advertises the graph we actually submit. --print-prompt stays a no-server-hit dry-run (documented) and prints the graph as-is. - Guard _checkpoint_enum against non-dict object_info (a hostile server returning `[]`) so it fails open rather than crashing with AttributeError. - Match a pinned bare filename against enum entries by basename so a checkpoint living in a subfolder (SD1.5/…safetensors) is found instead of triggering a needless substitution. - Escape the target-provided checkpoint name in the pretty substitution notice to prevent Rich markup injection. Co-Authored-By: Claude Opus 4.8 --- comfy_cli/command/run/__init__.py | 67 +++++++++++--------- comfy_cli/command/run/preflight.py | 38 +++++++---- comfy_cli/cql/default_workflow.py | 16 +++++ tests/comfy_cli/command/test_run_prompt.py | 26 ++++++++ tests/comfy_cli/cql/test_default_workflow.py | 22 +++++++ 5 files changed, 125 insertions(+), 44 deletions(-) diff --git a/comfy_cli/command/run/__init__.py b/comfy_cli/command/run/__init__.py index 13bfbc8b..6db1eb58 100644 --- a/comfy_cli/command/run/__init__.py +++ b/comfy_cli/command/run/__init__.py @@ -196,11 +196,30 @@ def execute( # foreach item map to stash on the job state at submit time. compose_meta = pop_compose_meta(workflow) + # Partner-API node preflight (below) and runtime checkpoint resolution both + # need the server's object_info. `--print-prompt` is a documented + # no-server-hit dry-run, so skip the fetch + resolution there and print the + # graph as-is; the real submit flow resolves BEFORE the prompt_preview event + # so the streamed audit trail advertises the graph we actually submit. + object_info: dict = {} + if not print_prompt: + object_info = _fetch_object_info(host, port) + + # Runtime checkpoint resolution for the bundled `--prompt` default: swap + # the pinned checkpoint for one the local server actually has (or + # hard-error if it has none). Guarded to the bundled default graph and + # skipped when the user pinned the checkpoint explicitly (honor it; let + # preflight reject it). + if preloaded is not None and workflow_name == "default_text2img" and not checkpoint_user_set: + _resolve_default_checkpoint_or_exit(renderer, workflow, object_info, where="local") + # Stream mode: emit the workflow graph so agents have a complete audit # trail of what the CLI is about to submit (no-op otherwise). renderer.event("prompt_preview", prompt=workflow) - # --print-prompt: emit/print the workflow and exit without submitting. + # --print-prompt: emit/print the workflow and exit without submitting. No + # server hit (documented) — the graph is shown as-is, before any + # server-dependent checkpoint resolution. if print_prompt: if renderer.is_pretty(): print(json.dumps(workflow, indent=2, ensure_ascii=False)) @@ -212,20 +231,6 @@ def execute( ) return - # Partner-API node preflight. Reject up-front when the workflow - # depends on a partner node (Veo/Kling/BFL/Gemini/…) and we have no - # credential to inject. If we DO have a credential, plumb it into - # extra_data so the partner node finds it server-side — same shape - # the cloud submit path uses. - object_info = _fetch_object_info(host, port) - - # Runtime checkpoint resolution for the bundled `--prompt` default: swap the - # pinned checkpoint for one the local server actually has (or hard-error if - # it has none). Guarded to the bundled default graph and skipped when the - # user pinned the checkpoint explicitly (honor it; let preflight reject it). - if preloaded is not None and workflow_name == "default_text2img" and not checkpoint_user_set: - _resolve_default_checkpoint_or_exit(renderer, workflow, object_info, where="local") - partner_nodes = _detect_partner_nodes(workflow, object_info) extra_data: dict | None = None if api_key: @@ -583,6 +588,23 @@ def execute_cloud( # its foreach item map to stash on the job state at submit time. compose_meta = pop_compose_meta(parsed_workflow) + # Cloud path uses cached/bundled object_info (no live server needed). Load + # it up front so checkpoint resolution can run BEFORE the preview/print + # below — the audit trail must advertise the graph we actually submit. + try: + from comfy_cli.cql.engine import _load_from_target + + cloud_object_info = _load_from_target(mode="cloud") + except Exception: # noqa: BLE001 + cloud_object_info = {} + + # Runtime checkpoint resolution for the bundled `--prompt` default (mirrors + # the local path): swap the pinned checkpoint for one Comfy Cloud actually + # has. Guarded to the bundled default and skipped when the user pinned the + # checkpoint explicitly. Cloud fails open on an empty enum (per-job models). + if preloaded is not None and workflow_name == "default_text2img" and not checkpoint_user_set: + _resolve_default_checkpoint_or_exit(renderer, parsed_workflow, cloud_object_info, where="cloud") + if print_prompt: # Documented dry-run: show the API-format graph that WOULD be sent and # exit WITHOUT POSTing. Mirrors local execute()'s print_prompt branch. @@ -598,21 +620,6 @@ def execute_cloud( raise typer.Exit(code=0) # Pre-submit validation via pure-Python CQL engine. - # Cloud path uses cached/bundled object_info (no live server needed). - try: - from comfy_cli.cql.engine import _load_from_target - - cloud_object_info = _load_from_target(mode="cloud") - except Exception: # noqa: BLE001 - cloud_object_info = {} - - # Runtime checkpoint resolution for the bundled `--prompt` default (mirrors - # the local path): swap the pinned checkpoint for one Comfy Cloud actually - # has, or hard-error if it enumerates none. Guarded to the bundled default - # and skipped when the user pinned the checkpoint explicitly. - if preloaded is not None and workflow_name == "default_text2img" and not checkpoint_user_set: - _resolve_default_checkpoint_or_exit(renderer, parsed_workflow, cloud_object_info, where="cloud") - _preflight_validate(renderer, parsed_workflow, cloud_object_info, target_label="cloud") target = resolve_target(where="cloud") diff --git a/comfy_cli/command/run/preflight.py b/comfy_cli/command/run/preflight.py index dfed8908..ec560bfa 100644 --- a/comfy_cli/command/run/preflight.py +++ b/comfy_cli/command/run/preflight.py @@ -123,8 +123,11 @@ def _resolve_default_checkpoint_or_exit(renderer, workflow: dict, object_info: d (fail open — preflight + the server decide); - pinned absent but the target has ≥1 checkpoint → substitute the first available one and emit a ``checkpoint_substituted`` note; - - target positively has zero checkpoints → hard ``no_checkpoint_available`` - error (exit 1) instead of a cryptic server-side reject. + - target positively has zero checkpoints → for ``where="local"`` a hard + ``no_checkpoint_available`` error (exit 1) instead of a cryptic + server-side reject; for ``where="cloud"`` a no-op (fail open), since Comfy + Cloud provisions its models per-job and the cached enum can't prove the + run would fail. ``where`` is ``"local"`` or ``"cloud"`` and drives the target label + hint. """ @@ -133,18 +136,21 @@ def _resolve_default_checkpoint_or_exit(renderer, workflow: dict, object_info: d target_label = "the local server" if where == "local" else "Comfy Cloud" _, res = resolve_default_checkpoint(workflow, object_info, target=target_label) + # Comfy Cloud provisions its models per-job at runtime, so an empty + # checkpoint enum in the cached/bundled cloud object_info does NOT mean the + # run would fail — hard-erroring there would wrongly block valid default + # cloud submits. Only the local path (where the enum reflects what's + # actually installed) treats a positively-empty enum as a hard stop. + if res.no_checkpoint and where == "cloud": + return + if res.no_checkpoint: - if where == "local": - hint = ( - "download a checkpoint, e.g. `comfy model download --url " - "https://huggingface.co/Comfy-Org/stable-diffusion-v1-5-archive/resolve/main/" - "v1-5-pruned-emaonly-fp16.safetensors`, then re-run — or `--set checkpoint=`" - ) - else: - hint = ( - "run a published gallery template (Comfy Cloud provisions its models), or " - "`--set checkpoint=` once a checkpoint is available on the target" - ) + # Only reachable for the local path (cloud returned above). + hint = ( + "download a checkpoint, e.g. `comfy model download --url " + "https://huggingface.co/Comfy-Org/stable-diffusion-v1-5-archive/resolve/main/" + "v1-5-pruned-emaonly-fp16.safetensors`, then re-run — or `--set checkpoint=`" + ) renderer.error( code="no_checkpoint_available", message=( @@ -159,7 +165,11 @@ def _resolve_default_checkpoint_or_exit(renderer, workflow: dict, object_info: d # Event fires in NDJSON/stream mode only; the pretty line covers humans. renderer.event("checkpoint_substituted", message=res.note, checkpoint=res.substituted_to, where=where) if renderer.is_pretty(): - pprint(f"[yellow]⚠ {res.note}[/yellow]") + from rich.markup import escape + + # res.note embeds a target-provided checkpoint name; escape it so a + # name containing Rich tags can't inject terminal markup. + pprint(f"[yellow]⚠ {escape(res.note)}[/yellow]") def _fetch_object_info(host: str, port: int) -> dict: diff --git a/comfy_cli/cql/default_workflow.py b/comfy_cli/cql/default_workflow.py index 09a8fab8..3d3eb4bb 100644 --- a/comfy_cli/cql/default_workflow.py +++ b/comfy_cli/cql/default_workflow.py @@ -276,6 +276,11 @@ def _checkpoint_enum(object_info: dict) -> list | None: option list is element 0. ``None`` (not an empty list) means "can't tell" so the caller can distinguish a positively-empty enum from an absent one. """ + if not isinstance(object_info, dict): + # A non-object /object_info payload (e.g. a hostile or misbehaving + # server returning ``[]``) means "can't tell" — fail open, mirroring + # Graph.from_object_info's isinstance guard rather than crashing. + return None node = object_info.get("CheckpointLoaderSimple") if not isinstance(node, dict): return None @@ -327,6 +332,17 @@ def resolve_default_checkpoint( if pinned in enum: return workflow, CheckpointResolution() + # ComfyUI enumerates checkpoints by their path relative to the models dir, + # so a pinned bare filename won't exact-match the same file living in a + # subfolder (e.g. ``SD1.5/v1-5-…safetensors``). Prefer a basename match to + # the *intended* checkpoint before falling back to an arbitrary substitute. + if isinstance(pinned, str): + pinned_base = pinned.rsplit("/", 1)[-1] + for entry in enum: + if isinstance(entry, str) and entry.rsplit("/", 1)[-1] == pinned_base: + inputs["ckpt_name"] = entry + return workflow, CheckpointResolution() + replacement = enum[0] inputs["ckpt_name"] = replacement note = ( diff --git a/tests/comfy_cli/command/test_run_prompt.py b/tests/comfy_cli/command/test_run_prompt.py index d7d68e38..49992966 100644 --- a/tests/comfy_cli/command/test_run_prompt.py +++ b/tests/comfy_cli/command/test_run_prompt.py @@ -213,3 +213,29 @@ def test_user_pinned_checkpoint_is_never_substituted(self): MockExec, _ = self._run_local(preloaded, oi) submitted = MockExec.call_args.args[0] assert submitted[CHECKPOINT_LOADER_ID]["inputs"]["ckpt_name"] == "userpick.safetensors" + + +class TestCheckpointResolutionEmptyEnumByTarget: + """`_resolve_default_checkpoint_or_exit` hard-errors on an empty enum only + for the local target; Comfy Cloud provisions models per-job so it fails + open (BE-2994).""" + + def test_local_empty_enum_hard_errors(self): + from comfy_cli.command.run.preflight import _resolve_default_checkpoint_or_exit + from comfy_cli.output import get_renderer + + wf = build_default_workflow(prompt="fox") + oi = _object_info_with_checkpoints([]) + with pytest.raises(typer.Exit) as e: + _resolve_default_checkpoint_or_exit(get_renderer(), wf, oi, where="local") + assert e.value.exit_code == 1 + + def test_cloud_empty_enum_fails_open(self): + from comfy_cli.command.run.preflight import _resolve_default_checkpoint_or_exit + from comfy_cli.output import get_renderer + + wf = build_default_workflow(prompt="fox") + oi = _object_info_with_checkpoints([]) + # No raise: the submit is allowed to proceed with the pinned default. + _resolve_default_checkpoint_or_exit(get_renderer(), wf, oi, where="cloud") + assert wf[CHECKPOINT_LOADER_ID]["inputs"]["ckpt_name"] == DEFAULT_CHECKPOINT_NAME diff --git a/tests/comfy_cli/cql/test_default_workflow.py b/tests/comfy_cli/cql/test_default_workflow.py index cad89cb9..40f27b45 100644 --- a/tests/comfy_cli/cql/test_default_workflow.py +++ b/tests/comfy_cli/cql/test_default_workflow.py @@ -237,6 +237,28 @@ def test_empty_object_info_fails_open(self): assert res.no_checkpoint is False assert res.note is None + def test_non_dict_object_info_fails_open(self): + # A hostile/misbehaving server returning non-object JSON (e.g. `[]`) + # must fail open, not crash with AttributeError. + wf = build_default_workflow(prompt="fox") + out, res = resolve_default_checkpoint(wf, []) + assert out[CHECKPOINT_LOADER_ID]["inputs"]["ckpt_name"] == DEFAULT_CHECKPOINT_NAME + assert res.no_checkpoint is False + assert res.note is None + + def test_pinned_matched_by_basename_in_subfolder(self): + # ComfyUI enumerates by path relative to the models dir; the pinned bare + # filename living in a subfolder should match (and be rewritten to the + # full path) rather than triggering an arbitrary substitution. + wf = build_default_workflow(prompt="fox") + subfoldered = f"SD1.5/{DEFAULT_CHECKPOINT_NAME}" + oi = _object_info_with_checkpoints(["other.safetensors", subfoldered]) + out, res = resolve_default_checkpoint(wf, oi) + assert out[CHECKPOINT_LOADER_ID]["inputs"]["ckpt_name"] == subfoldered + assert res.note is None + assert res.substituted_to is None + assert res.no_checkpoint is False + def test_object_info_without_checkpoint_node_fails_open(self): # object_info present but no CheckpointLoaderSimple → can't tell → fail open. wf = build_default_workflow(prompt="fox")