Build the image download plan for the engine that will load
/images/download-plan always asked the diffusers backend, while /images/load picks the engine per host: a GGUF pick on a machine with no usable GPU routes to native sd.cpp, which reads a single-file VAE plus text encoders and never opens the base repo's sharded components. Measured on unsloth/FLUX.2-klein-4B-GGUF (Q2_K): the plan staged 7.66 GB of FLUX.2-klein-4B components the native load discards, and the 7.80 GB sd-cli actually needs was then fetched inline by the loader, outside the download manager's progress and its disk preflight. Z-Image-Turbo is the same shape. The plan now asks whichever engine the load will select. predict_engine() applies the selection policy without any side effect: it activates nothing (staging a download must not unload the resident model) and only locates the binary rather than installing it, but still counts an installable binary as available, since that is what the load does on a fresh host. The native backend gains a download_plan built from the same _asset_specs the loader fetches, returning the same envelope, so the manager stages exactly the files sd-cli opens.
This commit is contained in:
Daniel Han
parent
94ad906f7f
commit
d4a17b2ca0
6 changed files with 356 additions and 2 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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]]:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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).
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue