diff --git a/studio/backend/core/inference/diffusion_engine_router.py b/studio/backend/core/inference/diffusion_engine_router.py index e8b7a7981d..f89ceae089 100644 --- a/studio/backend/core/inference/diffusion_engine_router.py +++ b/studio/backend/core/inference/diffusion_engine_router.py @@ -227,6 +227,48 @@ def select_and_activate_engine( return _activate(ENGINE_DIFFUSERS, reason) +def predict_engine(fam: DiffusionFamily, *, model_kind: Optional[str] = None) -> str: + """The engine a load of ``fam`` would select on this host, WITHOUT any side effect. + + Same policy as ``select_and_activate_engine`` -- and it has to be, because the download plan + is built from it: the two engines need different files (sharded diffusers components vs + sd-cli's single-file VAE + text encoders), so a plan built for the wrong one stages GB the + load never opens and then fetches the right files inline, outside the manager. + + It differs from selection in exactly two ways, both deliberate. Nothing is activated (staging + a download must not unload the resident model), and the binary is only LOCATED, never + installed -- but an absent binary on a host where install is allowed still counts as + available, because that is what the load will do. Getting that wrong the other way (planning + diffusers for the very first native load, when no binary is on disk yet) would mispredict the + common case: a fresh CPU host. + """ + if model_kind and model_kind != "gguf": + return ENGINE_DIFFUSERS + + forced, sd_cpp_pref, mps_enabled = _engine_config() + if forced == ENGINE_DIFFUSERS: + return ENGINE_DIFFUSERS + prefer_native = forced == ENGINE_SD_CPP + if sd_cpp_pref in _DISABLE_TOKENS and not prefer_native: + return ENGINE_DIFFUSERS + + backend = resolve_diffusion_device_target().backend + policy_eligible = backend == "cpu" or (backend == "mps" and mps_enabled) or prefer_native + if not (policy_eligible and family_sd_cpp_supported(fam)): + return ENGINE_DIFFUSERS + + server_binary = ensure_sd_server_binary(allow_install = False) + if server_binary and not _server_binary_runnable(server_binary): + server_binary = None + binary = ensure_sd_cpp_binary(allow_install = False) + if binary and SdCppEngine(binary = binary).version() is None: + binary = None + native_available = bool(binary or server_binary) or _install_allowed() + return select_diffusion_engine( + backend, native_available = native_available, prefer_native = prefer_native + ) + + def annotate_status(status: dict[str, Any]) -> dict[str, Any]: """Tag a backend status dict with the active engine + any fallback reason.""" out = dict(status) diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py index fdf07a0de8..6a52fc9ec9 100644 --- a/studio/backend/core/inference/sd_cpp_backend.py +++ b/studio/backend/core/inference/sd_cpp_backend.py @@ -633,6 +633,92 @@ class SdCppDiffusionBackend: if self._load_token == _load_token and self._loading is not None: self._loading.error = redact_native_paths(str(exc)) + def download_plan( + self, + repo_id: str, + *, + gguf_filename: Optional[str] = None, + base_repo: Optional[str] = None, + family_override: Optional[str] = None, + model_kind: Optional[str] = None, + hf_token: Optional[str] = None, + **_load_kwargs: Any, + ) -> dict[str, Any]: + """The repos + exact files a NATIVE load of this pick needs, in the same envelope the + diffusers backend returns, so the Hub download manager stages what sd-cli will actually + open. + + The two engines want different files: diffusers builds a pipeline around the base repo's + sharded components, while sd-cli reads the single-file VAE + text encoders declared in + ``diffusion_families``. Planning with the wrong engine stages tens of GB the load never + opens and then pulls the native assets inline, outside the manager's progress and disk + preflight -- so the route asks whichever engine it predicts the load will select. + + The diffusers-only kwargs (quant / memory / LoRA) are accepted and ignored, exactly as + ``begin_load`` accepts them: nothing sd-cli fetches depends on them.""" + if not gguf_filename: + raise ValueError( + "gguf_filename is required: the native engine loads single-file GGUF checkpoints only." + ) + fam = detect_family_for_pick(repo_id, gguf_filename, family_override) + if fam is None or not family_sd_cpp_supported(fam): + # Unreachable through the route (it only plans natively for a family the router would + # send to sd.cpp), but a direct caller gets the same message begin_load would raise. + raise ValueError(f"'{repo_id}' has no native sd.cpp asset mapping.") + + specs = self._asset_specs(repo_id, gguf_filename, fam) + # Group by repo so a family whose VAE and encoder live in one repo yields one entry, and + # keep first-seen order (transformer first) so the manager fetches the big file first. + by_repo: dict[str, list[str]] = {} + for repo, filename, kind in specs: + # A local GGUF directory is already on disk; nothing to stage for it. + if kind == "diffusion_model" and Path(repo).expanduser().exists(): + continue + names = by_repo.setdefault(repo, []) + if filename not in names: + names.append(filename) + + sizes = self._plan_file_sizes(by_repo, hf_token) + entries: list[dict[str, Any]] = [] + total = 0 + for repo, names in by_repo.items(): + repo_bytes = int(sum(sizes.get((repo, n), 0) for n in names)) + total += repo_bytes + entries.append( + { + "repo_id": repo, + "files": names, + "bytes": repo_bytes, + # Only the transformer entry carries the GGUF filename; the VAE / encoder + # entries are plain single files. + "gguf_filename": gguf_filename if repo == repo_id else None, + } + ) + return {"entries": entries, "total_bytes": total} + + @staticmethod + def _plan_file_sizes( + by_repo: dict[str, list[str]], hf_token: Optional[str] + ) -> dict[tuple[str, str], int]: + """(repo, filename) -> size in bytes, best-effort (0 for anything the Hub won't answer). + + A missing size only understates the manager's progress total; it must not fail the plan, + which is the cheap pre-flight for a load that would otherwise download inline.""" + out: dict[tuple[str, str], int] = {} + try: + from huggingface_hub import HfApi + + api = HfApi(token = hf_token) + except Exception: # noqa: BLE001 -- sizes are best-effort + return out + for repo, names in by_repo.items(): + try: + for info in api.get_paths_info(repo, paths = names, expand = False): + out[(repo, getattr(info, "path", ""))] = int(getattr(info, "size", 0) or 0) + except Exception: # noqa: BLE001 -- one unreadable repo is non-fatal + continue + return out + def _asset_specs( self, repo_id: str, gguf_filename: str, fam: DiffusionFamily ) -> list[tuple[str, str, str]]: diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index d53cb38ba4..4701610bf3 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -16066,6 +16066,8 @@ async def diffusion_download_plan( resolve_local_single_file, resolve_model_kind, ) + from core.inference.diffusion_engine_router import predict_engine + from core.inference.sd_cpp_engine import ENGINE_SD_CPP from utils.native_path_leases import redact_native_paths backend = get_diffusion_backend() @@ -16078,7 +16080,7 @@ async def diffusion_download_plan( if sole is not None: request.gguf_filename = sole kind = resolve_model_kind(sole) - await asyncio.to_thread( + fam = await asyncio.to_thread( backend.validate_load_request, request.model_path, gguf_filename = request.gguf_filename, @@ -16086,8 +16088,19 @@ async def diffusion_download_plan( model_kind = kind, base_repo = request.base_repo, ) + # Plan for the engine /images/load will pick, not for diffusers unconditionally. On a host + # with no usable GPU a GGUF pick routes to native sd.cpp, which reads single-file VAE + + # text encoders and never opens the base repo's sharded components: planning with the wrong + # engine stages GB the load discards and then leaves the assets it does need to be fetched + # inline, outside the manager's progress and disk preflight. predict_engine applies the + # selection policy without activating anything (staging must not unload a resident model). + planner = backend + if fam is not None and predict_engine(fam, model_kind = kind) == ENGINE_SD_CPP: + from core.inference.sd_cpp_backend import get_sd_cpp_backend + + planner = get_sd_cpp_backend() plan = await asyncio.to_thread( - backend.download_plan, + planner.download_plan, request.model_path, gguf_filename = request.gguf_filename, base_repo = request.base_repo, diff --git a/studio/backend/tests/test_diffusion_engine_router.py b/studio/backend/tests/test_diffusion_engine_router.py index c0f9ececbf..09ac9e8398 100644 --- a/studio/backend/tests/test_diffusion_engine_router.py +++ b/studio/backend/tests/test_diffusion_engine_router.py @@ -362,3 +362,67 @@ def test_switch_aborts_when_the_old_engine_fails_to_unload(monkeypatch): r._activate(ENGINE_DIFFUSERS, "switch test") # Still the old engine, so the resident model remains reachable and reclaimable. assert r.active_engine_name() == ENGINE_SD_CPP + + +# ── predict_engine: the download plan's read-only twin of the selection ─────── + + +def test_predict_engine_matches_the_selection_without_activating(monkeypatch): + # The download plan is built from this prediction, so it has to agree with the selection the + # load will make -- but staging a download must not unload a resident model, so it activates + # nothing. + _set_device(monkeypatch, "cpu") + _set_binary(monkeypatch, "/usr/bin/sd-cli") + _set_runnable(monkeypatch) + r._active_engine_name = ENGINE_DIFFUSERS + assert r.predict_engine(detect_family("z-image"), model_kind = "gguf") == ENGINE_SD_CPP + assert r.active_engine_name() == ENGINE_DIFFUSERS + + +def test_predict_engine_counts_an_installable_binary_as_available(monkeypatch): + # The first native load on a fresh CPU host installs the binary, so predicting diffusers just + # because nothing is on disk yet would mispredict the common case and stage the diffusers + # components for a load that opens none of them. It must not install anything itself. + _set_device(monkeypatch, "cpu") + installs: list[dict] = [] + + def _ensure(**kwargs): + installs.append(kwargs) + return None + + monkeypatch.setattr(r, "ensure_sd_cpp_binary", _ensure) + monkeypatch.setattr(r, "ensure_sd_server_binary", _ensure) + assert r.predict_engine(detect_family("z-image"), model_kind = "gguf") == ENGINE_SD_CPP + assert installs and all(k.get("allow_install") is False for k in installs) + + +def test_predict_engine_falls_back_when_install_is_disabled_and_nothing_is_installed(monkeypatch): + # With installs off and no binary present, the load really will fall back to diffusers. + monkeypatch.setenv("UNSLOTH_DIFFUSION_SD_CPP_INSTALL", "0") + _set_device(monkeypatch, "cpu") + _set_binary(monkeypatch, None) + monkeypatch.setattr(r, "ensure_sd_server_binary", lambda **_: None) + assert r.predict_engine(detect_family("z-image"), model_kind = "gguf") == ENGINE_DIFFUSERS + + +@pytest.mark.parametrize( + "kwargs, device", + [ + ({"model_kind": "pipeline"}, "cpu"), # native is GGUF-only + ({"model_kind": "single_file"}, "cpu"), + ({"model_kind": "gguf"}, "cuda"), # a usable GPU always means diffusers + ], +) +def test_predict_engine_returns_diffusers_where_the_load_would(monkeypatch, kwargs, device): + _set_device(monkeypatch, device) + _set_binary(monkeypatch, "/usr/bin/sd-cli") + _set_runnable(monkeypatch) + assert r.predict_engine(detect_family("z-image"), **kwargs) == ENGINE_DIFFUSERS + + +def test_predict_engine_returns_diffusers_for_a_family_without_native_assets(monkeypatch): + # sdxl has no single-file VAE / text-encoder mapping, so sd-cli cannot serve it. + _set_device(monkeypatch, "cpu") + _set_binary(monkeypatch, "/usr/bin/sd-cli") + _set_runnable(monkeypatch) + assert r.predict_engine(detect_family("sdxl"), model_kind = "gguf") == ENGINE_DIFFUSERS diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index 855a42496d..477e772be6 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -1071,6 +1071,88 @@ def test_download_plan_forwards_the_load_time_controls(client, monkeypatch): assert len(seen["loras"] or []) == 1 +def test_download_plan_uses_the_engine_the_load_will_pick(client, monkeypatch): + # On a host with no usable GPU a GGUF pick routes to native sd.cpp, which reads the single-file + # VAE + text encoders and never opens the base repo's sharded components. Planning with + # diffusers there staged GB the load discards and left the assets sd-cli actually needs to be + # fetched inline, outside the manager's progress and disk preflight. + from core.inference import diffusion_engine_router as router + from core.inference import sd_cpp_backend as sd_cpp + from core.inference.sd_cpp_engine import ENGINE_SD_CPP + + monkeypatch.setattr(router, "predict_engine", lambda fam, **_: ENGINE_SD_CPP) + native_plan = { + "entries": [{"repo_id": "Comfy-Org/z_image_turbo", "files": ["ae.safetensors"], + "bytes": 7, "gguf_filename": None}], + "total_bytes": 7, + } + seen: dict = {} + + class _Native: + def download_plan(self, model_path, **kwargs): + seen["model_path"] = model_path + seen.update(kwargs) + return native_plan + + monkeypatch.setattr(sd_cpp, "get_sd_cpp_backend", lambda: _Native()) + diffusers_backend = diffusion_module.get_diffusion_backend() + monkeypatch.setattr( + diffusers_backend, + "download_plan", + lambda *a, **k: pytest.fail("planned with diffusers for a native-routed load"), + raising = False, + ) + + resp = client.post( + "/api/inference/images/download-plan", + json = { + "model_path": "unsloth/Z-Image-Turbo-GGUF", + "gguf_filename": "z-image-turbo-Q4_K_M.gguf", + "model_kind": "gguf", + "hf_token": "hf_secret", + }, + ) + + assert resp.status_code == 200 + assert resp.json()["total_bytes"] == 7 + # The native planner gets the same identity + token the load would use. + assert seen["model_path"] == "unsloth/Z-Image-Turbo-GGUF" + assert seen["gguf_filename"] == "z-image-turbo-Q4_K_M.gguf" + assert seen["hf_token"] == "hf_secret" + + +def test_download_plan_stays_on_diffusers_when_the_load_will(client, monkeypatch): + # The mirror of the above: a GPU host (or any non-GGUF kind) loads through diffusers, so the + # plan must keep describing the diffusers component set. + from core.inference import diffusion_engine_router as router + from core.inference import sd_cpp_backend as sd_cpp + from core.inference.sd_cpp_engine import ENGINE_DIFFUSERS + + monkeypatch.setattr(router, "predict_engine", lambda fam, **_: ENGINE_DIFFUSERS) + monkeypatch.setattr( + sd_cpp, + "get_sd_cpp_backend", + lambda: pytest.fail("planned natively for a diffusers load"), + ) + backend = diffusion_module.get_diffusion_backend() + monkeypatch.setattr( + backend, + "download_plan", + lambda *a, **k: {"entries": [], "total_bytes": 11}, + raising = False, + ) + + resp = client.post( + "/api/inference/images/download-plan", + json = { + "model_path": "unsloth/Z-Image-Turbo-GGUF", + "gguf_filename": "z-image-turbo-Q4_K_M.gguf", + "model_kind": "gguf", + }, + ) + assert resp.status_code == 200 and resp.json()["total_bytes"] == 11 + + def test_load_refused_when_only_the_diffusion_probe_can_be_read(client, monkeypatch): # The two training probes are independent. An LLM backend that raises used to short-circuit the # guard entirely, so an image load sailed past a KNOWN-active diffusion trainer and contended diff --git a/studio/backend/tests/test_sd_cpp_backend.py b/studio/backend/tests/test_sd_cpp_backend.py index 5fa14c1e0d..2460986aec 100644 --- a/studio/backend/tests/test_sd_cpp_backend.py +++ b/studio/backend/tests/test_sd_cpp_backend.py @@ -214,6 +214,73 @@ def test_asset_specs_cover_required_files(fam_name, expect_kinds): assert tr[0] == "unsloth/x-GGUF" and tr[1] == "x-Q4_K_M.gguf" +def test_download_plan_stages_exactly_what_sd_cli_opens(monkeypatch): + # The plan feeds the Hub download manager. Native reads the single-file VAE + text encoder and + # never opens the base repo's sharded components, so a native-routed pick must be staged from + # the asset specs, not from the diffusers plan. + b = SdCppDiffusionBackend(engine = _FakeEngine()) + sizes = { + ("unsloth/Z-Image-Turbo-GGUF", "z-image-turbo-Q4_K_M.gguf"): 4_000, + ("Comfy-Org/z_image_turbo", "split_files/vae/ae.safetensors"): 300, + ("Comfy-Org/z_image_turbo", "split_files/text_encoders/qwen_3_4b.safetensors"): 8_000, + } + monkeypatch.setattr( + SdCppDiffusionBackend, + "_plan_file_sizes", + staticmethod(lambda by_repo, token: sizes), + ) + + plan = b.download_plan( + "unsloth/Z-Image-Turbo-GGUF", + gguf_filename = "z-image-turbo-Q4_K_M.gguf", + model_kind = "gguf", + # diffusers-only knobs are accepted and ignored, exactly as begin_load accepts them. + transformer_quant = "int8", + memory_mode = "low_vram", + ) + + fam = detect_family("z-image") + expected = {(r, f) for r, f, _k in b._asset_specs( + "unsloth/Z-Image-Turbo-GGUF", "z-image-turbo-Q4_K_M.gguf", fam + )} + listed = {(e["repo_id"], f) for e in plan["entries"] for f in e["files"]} + assert listed == expected + assert plan["total_bytes"] == 12_300 + # The transformer entry is the only one that carries the GGUF filename (the manager treats a + # GGUF pick differently from a plain file), and the VAE + encoder share one repo entry. + tr = [e for e in plan["entries"] if e["gguf_filename"]] + assert len(tr) == 1 and tr[0]["repo_id"] == "unsloth/Z-Image-Turbo-GGUF" + assert len([e for e in plan["entries"] if e["repo_id"] == "Comfy-Org/z_image_turbo"]) == 1 + + +def test_download_plan_skips_a_local_transformer_but_still_stages_the_assets(monkeypatch, tmp_path): + # A local GGUF folder is already on disk; its VAE + encoder still have to come from the Hub. + b = SdCppDiffusionBackend(engine = _FakeEngine()) + (tmp_path / "z-image-turbo-Q4_K_M.gguf").write_bytes(b"gguf") + monkeypatch.setattr( + SdCppDiffusionBackend, "_plan_file_sizes", staticmethod(lambda by_repo, token: {}) + ) + + plan = b.download_plan( + str(tmp_path), gguf_filename = "z-image-turbo-Q4_K_M.gguf", model_kind = "gguf" + ) + + assert str(tmp_path) not in {e["repo_id"] for e in plan["entries"]} + assert {e["repo_id"] for e in plan["entries"]} == {"Comfy-Org/z_image_turbo"} + # An unreadable size understates the total; it must never fail the plan. + assert plan["total_bytes"] == 0 + + +def test_download_plan_refuses_a_pick_native_cannot_serve(): + b = SdCppDiffusionBackend(engine = _FakeEngine()) + with pytest.raises(ValueError, match = "gguf_filename is required"): + b.download_plan("unsloth/Z-Image-Turbo-GGUF", model_kind = "gguf") + with pytest.raises(ValueError, match = "native sd.cpp asset mapping"): + b.download_plan( + "stabilityai/sdxl-turbo", gguf_filename = "sdxl-Q4_K_M.gguf", model_kind = "gguf" + ) + + def test_asset_specs_flux2_klein_selects_encoder_by_variant(): # FLUX.2-klein 4B pairs with Qwen3-4B, 9B with Qwen3-8B, so the encoder must be chosen from the # load identity, not the family default (a mismatched encoder fails deep in sd-cli).