diff --git a/comfy_cli/cmdline.py b/comfy_cli/cmdline.py index 9a3594f9..0ef9c961 100644 --- a/comfy_cli/cmdline.py +++ b/comfy_cli/cmdline.py @@ -787,8 +787,9 @@ def run( help=( "Positive text prompt for the bundled default text2img workflow " "(used when --workflow is omitted). Cannot be combined with --workflow. " - "The bundled graph loads an SD1.5 checkpoint (v1-5-pruned-emaonly.ckpt) " - "that is NOT downloaded for you — install it, or point elsewhere with " + "The bundled graph prefers an SD1.5 checkpoint " + "(v1-5-pruned-emaonly-fp16.safetensors); if the target doesn't have it, " + "an installed checkpoint is substituted and reported. Pin your own with " "--set checkpoint=." ), ), @@ -944,7 +945,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( @@ -956,7 +957,7 @@ def run( from comfy_cli.cql.default_workflow import ( PromptInjectionError, build_default_workflow, - default_checkpoint, + overrides_set_checkpoint, ) try: @@ -964,29 +965,10 @@ def run( 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) - # The bundled graph pins an SD1.5 checkpoint that comfy-cli neither - # ships nor auto-downloads. Without it the run dies server-side on a - # bare validation error, so state the dependency up front. Pretty - # output only — the JSON dialects carry a fixed event contract. - ckpt = default_checkpoint(injected) - if ckpt and renderer.is_pretty(): - from rich.markup import escape as _escape - - # The checkpoint has to exist wherever the run is routed, so - # name that environment: pointing a `--where cloud` run at the - # local models/checkpoints sends the user to fix the wrong box. - if decision.target is where_module.WhereTarget.CLOUD: - where_ckpt = "in your cloud assets (`comfy models search --where cloud`)" - else: - where_ckpt = "in models/checkpoints" - # `--set checkpoint=…` puts a user string here; escape it so a - # value containing [brackets] can't be read as rich markup. - rprint( - f"[dim]Using the bundled default text2img workflow — it needs the[/dim] " - f"[bold]{_escape(ckpt)}[/bold] [dim]checkpoint {where_ckpt}. " - f"Override it with --set checkpoint=.[/dim]" - ) + # 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 ad299286..9e304b96 100644 --- a/comfy_cli/command/run/__init__.py +++ b/comfy_cli/command/run/__init__.py @@ -40,6 +40,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 @@ -165,7 +166,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, allow_spend: bool = False, ): # `0.0.0.0` is a wildcard bind, not a connect address. macOS / Windows @@ -185,10 +186,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: @@ -254,11 +258,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)) @@ -270,12 +293,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) partner_nodes = _detect_partner_nodes(workflow, object_info) # Spend gate (BE-4326): partner-API nodes spend Comfy credits. Require # explicit consent before resolving a credential or submitting. Fires @@ -715,7 +732,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, allow_spend: bool = False, ): """Run a workflow against Comfy Cloud via the stored OAuth session. @@ -731,14 +748,22 @@ 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: renderer.error(code=e.code, message=str(e), hint=e.hint) raise typer.Exit(code=1) from e + # The cloud object_info snapshot is used twice below (UI→API conversion and + # checkpoint resolution/preflight). `_load_from_target` is a live, uncached + # HTTPS fetch, so load it at most once and share it across both. `None` + # means "not fetched yet" — distinct from a fetched-but-empty snapshot, + # which must NOT trigger a second round-trip. + cloud_object_info: dict | None = None + if is_ui: # Frontend-format workflows (the `nodes`+`links` shape from the canvas # exporter and `comfy templates fetch`) have to be lowered to the API @@ -749,7 +774,7 @@ def execute_cloud( if renderer.is_pretty(): pprint("[yellow]Detected UI-format workflow, converting to API format…[/yellow]") try: - object_info = _load_from_target(mode="cloud") + object_info = cloud_object_info = _load_from_target(mode="cloud") except Exception as e: # noqa: BLE001 renderer.error( code="cql_no_graph", @@ -793,6 +818,25 @@ 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. + # Already fetched above when the workflow arrived in UI format. + if cloud_object_info is None: + 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. @@ -808,14 +852,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 = {} - _preflight_validate(renderer, parsed_workflow, cloud_object_info, target_label="cloud") # Spend gate (BE-4326): the cloud also bills partner-API nodes, so apply the diff --git a/comfy_cli/command/run/preflight.py b/comfy_cli/command/run/preflight.py index f22a8ba6..f27bee55 100644 --- a/comfy_cli/command/run/preflight.py +++ b/comfy_cli/command/run/preflight.py @@ -119,6 +119,67 @@ def _preflight_validate(renderer, workflow: dict, object_info: dict, *, target_l ) +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 → 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. + """ + 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) + + # 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: + # 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=( + 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(): + 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: """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 5f9b4a7b..aa015032 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" @@ -239,3 +245,140 @@ 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. + """ + 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 + 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 _basename(name: str) -> str: + """Last path segment of a ComfyUI-enumerated model name. + + ``folder_paths`` builds these names with ``os.path.relpath``, so the + separator is the SERVER's, not ours: a subfoldered checkpoint arrives as + ``SD1.5/name.safetensors`` from a POSIX host and ``SD1.5\\name.safetensors`` + from a Windows one. Normalize both before comparing basenames — splitting on + ``/`` alone would miss the Windows form and trigger a needless substitution. + """ + return name.replace("\\", "/").rsplit("/", 1)[-1] + + +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() + + # 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 = _basename(pinned) + for entry in enum: + if isinstance(entry, str) and _basename(entry) == pinned_base: + inputs["ckpt_name"] = entry + 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 1cd00e37..1d0e1764 100644 --- a/comfy_cli/error_codes.py +++ b/comfy_cli/error_codes.py @@ -201,6 +201,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 2733e54c..49992966 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") @@ -133,38 +154,88 @@ def test_bad_set_address_is_rejected(self): assert e.value.exit_code == 1 -class TestDefaultCheckpointNotice: - """The bundled graph's checkpoint is not downloaded for you, so `run` says so. +class TestRuntimeCheckpointResolutionLocal: + """Wiring: `execute` resolves the bundled default's checkpoint against the + server's object_info before submit (BE-2994).""" - The notice has to name the environment the run was ROUTED to: telling a - `--where cloud` user to drop a file in the local ``models/checkpoints`` - points them at the wrong machine (the README says cloud runs need it in - cloud assets). - """ + 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 _run(self, where: str, capsys): - renderer = MagicMock() - renderer.is_pretty.return_value = True - with ( - patch("comfy_cli.cmdline.tracking.track_event"), - patch("comfy_cli.cmdline.get_renderer", return_value=renderer), - patch("comfy_cli.cmdline.where_module.cloud_preflight", return_value=None), - patch("comfy_cli.command.run.execute"), - patch("comfy_cli.command.run.execute_cloud"), - ): - run_command(where=where, prompt="a red fox in snow") - return capsys.readouterr().out - - def test_local_run_points_at_the_local_model_dir(self, capsys): - out = self._run("local", capsys) - assert "models/checkpoints" in out - assert "cloud assets" not in out - - def test_cloud_run_points_at_cloud_assets(self, capsys): - out = self._run("cloud", capsys) - assert "cloud assets" in out - assert "models/checkpoints" not in out - - def test_notice_names_the_checkpoint(self, capsys): - out = self._run("local", capsys) - assert "v1-5-pruned-emaonly.ckpt" in out + 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" + + +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 7f1fa9b1..8694f612 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,127 @@ 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_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_pinned_matched_by_basename_windows_separator(self): + # A Windows server's folder_paths builds relative names with `\`, so the + # basename match must normalize separators — otherwise the intended + # checkpoint is missed and enum[0] is substituted needlessly. This is + # about the SERVER's separator, so it must hold on any client OS. + 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") + _, 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