From fbcf070fce55db23d295054ad85833fcdb090451 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 17 Jul 2026 05:47:45 +0000 Subject: [PATCH] Wire hosted pre-quantized DiT checkpoints into the image families Point prequant_repos for flux.1, flux.2-klein, flux.2-dev, qwen-image (int8 only there; fp8 is family-denied), z-image and krea-2 at the unsloth/-FP8 Hub repos carrying gate-validated int8 and fp8 transformer checkpoints, so the fast quant path loads the small pre-quantized file instead of materialising the dense bf16 transformer and quantising on device. Measured on FLUX.2-dev int8: build peak drops from 60.7 GB (dense + quantize) to 30.7 GB (hosted prequant), identical 30.7 GB resident after either path since loading a checkpoint is bit-identical to on-the-fly quantisation. The hosted repos name files -.pt, so resolve_prequant_source now derives that model-name filename from the repo id (scheme suffix stripped case-insensitively) and carries the legacy transformer_.pt as a fallback the resolver tries when the primary 404s, keeping older repos loadable. Wiring a repo also exposed a fallback hazard: with a prequant source present, the dense-fit preflight used to be skipped entirely, so a failed prequant download would fall through to the dense bf16 load the memory plan never budgeted, OOMing after eviction. The preflight now always runs and gates an allow_dense_fallback flag through _load_dense_quant_pipeline: a dense misfit still skips the fast path when no prequant exists, but with one it proceeds and a prequant failure raises to the GGUF build instead of loading dense. The same flag is set when the auto-policy replans an offloaded GGUF against a prequant-sized transient. Tests updated to the new filename convention plus new coverage for the derivation and the legacy-name fallback; the prequant-skips-refit test now asserts the re-check runs and forbids the dense fallback. Verified end to end on GPU: z-image int8 resolves the hosted repo, downloads the model-name file and renders (6.8s load, 5.9 GB peak). --- studio/backend/core/inference/diffusion.py | 54 ++++++++++++----- .../core/inference/diffusion_families.py | 26 ++++++++ .../core/inference/diffusion_prequant.py | 39 ++++++++++-- .../backend/tests/test_diffusion_backend.py | 17 +++--- .../backend/tests/test_diffusion_prequant.py | 60 ++++++++++++++++++- 5 files changed, 168 insertions(+), 28 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 8e99b16e39..296523bb6f 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -1210,6 +1210,10 @@ class DiffusionBackend: # The GGUF-size `plan` can mis-budget the fast path two ways, so preflight the real # footprint BEFORE eviction; both branches need the base repo + a resolved scheme. dense_declined = False + # False when the memory plan only holds a PREQUANT-sized build: if the prequant + # load then fails, the loader must raise to GGUF instead of materialising the + # dense bf16 transformer the plan never budgeted for. + dense_fallback_allowed = True if ( kind == "gguf" and normalize_transformer_quant(transformer_quant) is not None @@ -1246,6 +1250,10 @@ class DiffusionBackend: ) if replanned.offload_policy == OFFLOAD_NONE: quant_plan = replanned + # The GGUF plan already declined resident; a prequant-sized + # replan says nothing about the (larger) dense transformer. + if candidate.prequant: + dense_fallback_allowed = False else: # The GGUF fits resident, but this path first materialises the base's dense # bf16 transformer (bigger), so re-check the fit against THAT -- a card that @@ -1267,23 +1275,29 @@ class DiffusionBackend: if scheme is not None else None ) - if prequant is None: - dense_mib = int( - self._dense_transformer_resident_bytes(base) // (1024 * 1024) + dense_mib = int( + self._dense_transformer_resident_bytes(base) // (1024 * 1024) + ) + if dense_mib > 0: + dense_plan = self._plan_memory( + target, + single_file_path, + base, + fam, + memory_mode, + cpu_offload, + kind = kind, + repo_id = repo_id, + transformer_resident_override_mib = dense_mib, ) - if dense_mib > 0: - dense_plan = self._plan_memory( - target, - single_file_path, - base, - fam, - memory_mode, - cpu_offload, - kind = kind, - repo_id = repo_id, - transformer_resident_override_mib = dense_mib, - ) - dense_declined = dense_plan.offload_policy != OFFLOAD_NONE + if dense_plan.offload_policy != OFFLOAD_NONE: + dense_fallback_allowed = False + # Without a prequant source the dense build is the ONLY path, + # so a dense misfit skips the fast path entirely (as before); with + # one, the small prequant load proceeds and only the dense + # fallback is forbidden. + if prequant is None: + dense_declined = True if ( kind == "gguf" and normalize_transformer_quant(transformer_quant) is not None @@ -1305,6 +1319,7 @@ class DiffusionBackend: fam = fam, base_local_dir = _base_local_dir, prequant_path = transformer_prequant_path, + allow_dense_fallback = dense_fallback_allowed, ) except Exception as exc: # noqa: BLE001 — fall back to the GGUF build logger.warning( @@ -1713,6 +1728,7 @@ class DiffusionBackend: fam: Optional[DiffusionFamily] = None, prequant_path: Optional[str] = None, base_local_dir: Optional[str] = None, + allow_dense_fallback: bool = True, ) -> tuple[Any, str]: """Build the opt-in fast pipeline and return ``(pipe, engaged_scheme)``. @@ -1761,6 +1777,12 @@ class DiffusionBackend: return pipe, scheme # 2. Fallback: materialise the dense bf16 transformer and quantise it on-device. + if not allow_dense_fallback: + # The memory plan only budgeted the prequant-sized build; materialising the dense + # bf16 transformer here would exceed it after eviction. Raise to the GGUF build. + raise RuntimeError( + "prequant checkpoint unavailable and the dense transformer does not fit resident" + ) transformer = transformer_cls.from_pretrained( base, subfolder = "transformer", torch_dtype = dtype, token = hf_token ) diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index 78f42b4a4b..21de136cc5 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -112,6 +112,14 @@ _FAMILIES: tuple[DiffusionFamily, ...] = ( pipeline_class = "FluxPipeline", transformer_class = "FluxTransformer2DModel", base_repo = "black-forest-labs/FLUX.1-schnell", + # Hosted pre-quantized DiT checkpoints (gate-validated vs same-seed bf16). The loader + # verifies the checkpoint's baked base_model_id against the repo actually being loaded, + # so a non-default base (e.g. FLUX.1-dev under this family) safely falls back to the + # dense-quantize path instead of loading schnell weights. + prequant_repos = ( + ("int8", "unsloth/FLUX.1-schnell-FP8"), + ("fp8", "unsloth/FLUX.1-schnell-FP8"), + ), aliases = ("flux1", "flux-1"), # LoRA training targets FLUX.1-dev via the DiT trainer (QLoRA nf4); the dev repo is gated. trainable = True, @@ -133,6 +141,10 @@ _FAMILIES: tuple[DiffusionFamily, ...] = ( pipeline_class = "Flux2KleinPipeline", transformer_class = "Flux2Transformer2DModel", base_repo = "black-forest-labs/FLUX.2-klein-4B", + prequant_repos = ( + ("int8", "unsloth/FLUX.2-klein-4B-FP8"), + ("fp8", "unsloth/FLUX.2-klein-4B-FP8"), + ), aliases = ("flux2-klein",), # Flux2KleinPipeline takes reference image(s) via `image`, so it exposes a "reference" # workflow atop text-to-image. It has an inpaint pipeline (no img2img) -> inpaint + extend. @@ -156,6 +168,10 @@ _FAMILIES: tuple[DiffusionFamily, ...] = ( pipeline_class = "Flux2Pipeline", transformer_class = "Flux2Transformer2DModel", base_repo = "black-forest-labs/FLUX.2-dev", + prequant_repos = ( + ("int8", "unsloth/FLUX.2-dev-FP8"), + ("fp8", "unsloth/FLUX.2-dev-FP8"), + ), aliases = ("flux2-dev", "flux2dev"), sd_cpp_vae = ("Comfy-Org/flux2-dev", "split_files/vae/flux2-vae.safetensors"), sd_cpp_vae_format = "flux2", @@ -201,6 +217,8 @@ _FAMILIES: tuple[DiffusionFamily, ...] = ( pipeline_class = "QwenImagePipeline", transformer_class = "QwenImageTransformer2DModel", base_repo = "Qwen/Qwen-Image", + # int8 only: fp8 is family-denied (_FAMILY_SCHEME_DENY) so a repo entry would be dead. + prequant_repos = (("int8", "unsloth/Qwen-Image-FP8"),), cfg_kwarg = "true_cfg_scale", aliases = ("qwen_image", "qwenimage"), # LoRA training via the DiT trainer, defaulting to the prequant nf4 repo (QLoRA). @@ -229,6 +247,10 @@ _FAMILIES: tuple[DiffusionFamily, ...] = ( pipeline_class = "ZImagePipeline", transformer_class = "ZImageTransformer2DModel", base_repo = "Tongyi-MAI/Z-Image-Turbo", + prequant_repos = ( + ("int8", "unsloth/Z-Image-Turbo-FP8"), + ("fp8", "unsloth/Z-Image-Turbo-FP8"), + ), aliases = ("zimage", "z_image"), # LoRA training via the DiT trainer (bf16); defaults to the prequant nf4 repo for QLoRA. trainable = True, @@ -250,6 +272,10 @@ _FAMILIES: tuple[DiffusionFamily, ...] = ( pipeline_class = "Krea2Pipeline", transformer_class = "Krea2Transformer2DModel", base_repo = "krea/Krea-2-Turbo", + prequant_repos = ( + ("int8", "unsloth/Krea-2-Turbo-FP8"), + ("fp8", "unsloth/Krea-2-Turbo-FP8"), + ), aliases = ("krea2",), # LoRA training via the DiT trainer (no prequant repo yet, so nf4 quantizes on the fly). # Krea's guidance: train on the undistilled Raw, run adapters on Turbo, so Raw is the diff --git a/studio/backend/core/inference/diffusion_prequant.py b/studio/backend/core/inference/diffusion_prequant.py index 40d320a5e2..69a5779201 100644 --- a/studio/backend/core/inference/diffusion_prequant.py +++ b/studio/backend/core/inference/diffusion_prequant.py @@ -90,18 +90,32 @@ def local_prequant_path_ready(path: str) -> bool: @dataclass(frozen = True) class PrequantSource: """Where a pre-quantized checkpoint lives. ``kind`` is "path" (a local file) or "repo" - (Hub repo id in ``location`` + ``filename``).""" + (Hub repo id in ``location`` + ``filename``; ``fallback_filename`` is tried when the + primary name is absent, covering repos still on the legacy transformer_.pt).""" kind: str location: str filename: Optional[str] = None + fallback_filename: Optional[str] = None def prequant_filename(scheme: str) -> str: - """The conventional checkpoint filename for ``scheme`` inside a Hub repo.""" + """The legacy checkpoint filename for ``scheme`` inside a Hub repo.""" return f"transformer_{scheme}.pt" +def prequant_repo_filename(repo_id: str, scheme: str) -> str: + """The model-name checkpoint filename for ``scheme`` in ``repo_id``: the hosted repos are + named -FP8 (or -INT8 / -quantized) and carry -.pt files, e.g. + unsloth/Z-Image-Turbo-FP8 -> Z-Image-Turbo-INT8.pt / Z-Image-Turbo-FP8.pt.""" + model = repo_id.rsplit("/", 1)[-1] + for suffix in ("-fp8", "-int8", "-quantized"): + if model.lower().endswith(suffix): + model = model[: -len(suffix)] + break + return f"{model}-{scheme.upper()}.pt" + + def resolve_prequant_source( fam: Any, scheme: str, @@ -122,7 +136,12 @@ def resolve_prequant_source( except Exception: # noqa: BLE001 — a bad family object must not break the load repo_id = None if repo_id: - return PrequantSource(kind = "repo", location = repo_id, filename = prequant_filename(scheme)) + return PrequantSource( + kind = "repo", + location = repo_id, + filename = prequant_repo_filename(repo_id, scheme), + fallback_filename = prequant_filename(scheme), + ) return None @@ -243,7 +262,19 @@ def _resolve_checkpoint_path(source: PrequantSource, hf_token: Optional[str]) -> return expanded if os.path.isfile(expanded) else None if source.kind == "repo": from huggingface_hub import hf_hub_download - return hf_hub_download(repo_id = source.location, filename = source.filename, token = hf_token) + try: + from huggingface_hub.errors import EntryNotFoundError + except Exception: # noqa: BLE001 — older hub layouts; fall back to a private marker + class EntryNotFoundError(Exception): # type: ignore[no-redef] + pass + try: + return hf_hub_download(repo_id = source.location, filename = source.filename, token = hf_token) + except EntryNotFoundError: + if not source.fallback_filename or source.fallback_filename == source.filename: + raise + return hf_hub_download( + repo_id = source.location, filename = source.fallback_filename, token = hf_token + ) return None diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index d0c781cf05..d9f5652f80 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -2656,10 +2656,11 @@ def test_dense_quant_skipped_when_dense_transformer_does_not_fit( assert _FakeTransformer.last["path"] # GGUF path used -def test_dense_quant_prequant_skips_dense_refit(fake_runtime, tmp_path, monkeypatch): - # With a prequant checkpoint, the fast path loads the small quantized file, not the dense - # bf16 -- so the dense-transformer re-check must NOT run and must NOT decline the fast path, - # even when the base's dense shards are cached and large. +def test_dense_quant_prequant_proceeds_but_forbids_dense_fallback(fake_runtime, tmp_path, monkeypatch): + # With a prequant checkpoint, the fast path loads the small quantized file, so a dense + # misfit must NOT decline the fast path -- but the dense re-check still runs to gate the + # in-loader fallback: if the prequant later fails, the loader must raise to GGUF instead of + # materialising the dense bf16 the plan never budgeted (allow_dense_fallback=False). from core.inference import diffusion as dmod backend = DiffusionBackend() @@ -2688,13 +2689,15 @@ def test_dense_quant_prequant_skips_dense_refit(fake_runtime, tmp_path, monkeypa ): if transformer_resident_override_mib is not None: dense_refit_ran.append(True) + # GGUF budget fits (real plan -> none); the dense-transformer preflight does not. + return types.SimpleNamespace(offload_policy = "model") return orig_plan(self, *a, **k) monkeypatch.setattr(DiffusionBackend, "_plan_memory", spy_plan) attempted = [] def fake_dense_load(self, *a, **k): - attempted.append(True) + attempted.append(k.get("allow_dense_fallback")) return None, None # fall through to GGUF; we only assert the path was reached monkeypatch.setattr(DiffusionBackend, "_load_dense_quant_pipeline", fake_dense_load) @@ -2705,8 +2708,8 @@ def test_dense_quant_prequant_skips_dense_refit(fake_runtime, tmp_path, monkeypa family_override = "z-image", transformer_quant = "fp8", ) - assert dense_refit_ran == [] # prequant -> dense re-check skipped - assert attempted == [True] # fast path still attempted (with the prequant) + assert dense_refit_ran == [True] # the re-check runs (it gates the fallback)... + assert attempted == [False] # ...fast path still attempted, dense fallback forbidden def test_dense_quant_unusable_prequant_path_runs_dense_refit(fake_runtime, tmp_path, monkeypatch): diff --git a/studio/backend/tests/test_diffusion_prequant.py b/studio/backend/tests/test_diffusion_prequant.py index 3e77f58107..3e5d09b9c9 100644 --- a/studio/backend/tests/test_diffusion_prequant.py +++ b/studio/backend/tests/test_diffusion_prequant.py @@ -46,7 +46,18 @@ def test_resolve_family_repo_by_scheme(): fam = _fam(prequant_repos = (("fp8", "org/hosted-fp8"), ("int8", "org/hosted-int8"))) src = resolve_prequant_source(fam, "int8") assert src.kind == "repo" and src.location == "org/hosted-int8" - assert src.filename == "transformer_int8.pt" + # Model-name convention first (repo scheme suffix stripped), legacy name as fallback. + assert src.filename == "hosted-INT8.pt" + assert src.fallback_filename == "transformer_int8.pt" + + +def test_prequant_repo_filename_convention(): + from core.inference.diffusion_prequant import prequant_repo_filename + assert prequant_repo_filename("unsloth/Z-Image-Turbo-FP8", "int8") == "Z-Image-Turbo-INT8.pt" + assert prequant_repo_filename("unsloth/Z-Image-Turbo-FP8", "fp8") == "Z-Image-Turbo-FP8.pt" + assert prequant_repo_filename("unsloth/Qwen-Image-2512-INT8", "int8") == "Qwen-Image-2512-INT8.pt" + assert prequant_repo_filename("org/Some-Model-quantized", "fp8") == "Some-Model-FP8.pt" + assert prequant_repo_filename("org/PlainRepo", "int8") == "PlainRepo-INT8.pt" def test_resolve_wrong_scheme_is_none(): @@ -480,6 +491,53 @@ def test_load_repo_source_allowed_without_optin(monkeypatch, tmp_path): assert result is not None +def test_load_repo_source_falls_back_to_legacy_filename(monkeypatch, tmp_path): + # A repo still carrying the legacy transformer_.pt name serves the download after + # the model-name filename 404s; both names are requested in order. + _FakeTransformer.calls = {} + _stub_torch_accelerate(monkeypatch, _good_ckpt()) + monkeypatch.delenv(pq.ALLOW_LOCAL_PREQUANT_PATH_ENV, raising = False) + + downloaded = tmp_path / "transformer_fp8.pt" + downloaded.write_bytes(b"x") + + class _NotFound(Exception): + pass + + errors = types.ModuleType("huggingface_hub.errors") + errors.EntryNotFoundError = _NotFound + requested = [] + + def _dl(repo_id, filename, token = None): + requested.append(filename) + if filename != "transformer_fp8.pt": + raise _NotFound(filename) + return str(downloaded) + + hub = types.ModuleType("huggingface_hub") + hub.hf_hub_download = _dl + hub.errors = errors + monkeypatch.setitem(sys.modules, "huggingface_hub", hub) + monkeypatch.setitem(sys.modules, "huggingface_hub.errors", errors) + + source = PrequantSource( + kind = "repo", location = "org/Z-Image-Turbo-FP8", + filename = "Z-Image-Turbo-FP8.pt", fallback_filename = "transformer_fp8.pt", + ) + result = load_prequantized_transformer( + _FakeTransformer, + "Tongyi-MAI/Z-Image-Turbo", + source, + device = "cuda", + dtype = "bfloat16", + hf_token = None, + scheme = "fp8", + logger = None, + ) + assert result is not None + assert requested == ["Z-Image-Turbo-FP8.pt", "transformer_fp8.pt"] + + def test_load_local_path_outside_allowlist_refused(monkeypatch, tmp_path): # Even with the opt-in set, a path OUTSIDE every allowlisted directory must not be # unpickled: enabling one trusted dir is not a wildcard for arbitrary request paths.