From 570cef6c79c608a4dcf45419637604398dc564fd Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 17 Jul 2026 11:07:17 +0000 Subject: [PATCH] Resolve pre-quantized checkpoints per base variant One family entry covers several published variants whose weights differ (flux.1: schnell, dev, Krea-dev), but prequant resolution was keyed on (family, scheme) alone, so only the default base could ever be served: the loader's baked base_model_id validation correctly refused the schnell checkpoint for dev and Krea-dev bases and every such load paid the dense download plus on-the-fly quantise. Add an optional prequant_variant_repos table on DiffusionFamily as (base_repo, scheme, repo_id) triples and thread the resolved base repo through resolve_prequant_source / usable_prequant_source and their three call sites (load fast path, memory-plan probe, auto-policy candidate). A base without its own entry keeps returning the family default, preserving the existing refuse-then-dense behavior exactly. Wire the flux.1 variants: the gate-validated unsloth/FLUX.1-dev-FP8 checkpoints (built in the earlier campaign but never reachable) and the new unsloth/FLUX.1-Krea-dev-FP8. --- studio/backend/core/inference/diffusion.py | 9 +++- .../core/inference/diffusion_auto_policy.py | 4 +- .../core/inference/diffusion_families.py | 31 ++++++++++- .../core/inference/diffusion_prequant.py | 11 ++-- .../tests/test_diffusion_auto_policy.py | 6 ++- .../backend/tests/test_diffusion_prequant.py | 54 ++++++++++++++++++- 6 files changed, 105 insertions(+), 10 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index b3a303106b..305d143540 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -1292,7 +1292,10 @@ class DiffusionBackend: None if loras else usable_prequant_source( - fam, scheme, path_override = transformer_prequant_path + fam, + scheme, + path_override = transformer_prequant_path, + base_repo = base, ) if scheme is not None else None @@ -1755,7 +1758,9 @@ class DiffusionBackend: if fam is not None and not lora_specs: # A LoRA bake needs the DENSE transformer (adapters attach before quantize_), so # the prequant shortcut is skipped when adapters were requested. - source = resolve_prequant_source(fam, scheme, path_override = prequant_path) + source = resolve_prequant_source( + fam, scheme, path_override = prequant_path, base_repo = base + ) if source is not None: transformer = load_prequantized_transformer( transformer_cls, diff --git a/studio/backend/core/inference/diffusion_auto_policy.py b/studio/backend/core/inference/diffusion_auto_policy.py index b7d805c8c6..74f629c257 100644 --- a/studio/backend/core/inference/diffusion_auto_policy.py +++ b/studio/backend/core/inference/diffusion_auto_policy.py @@ -183,7 +183,9 @@ def resolve_dense_quant_candidate( # accept it (allowlisted AND present), else load_prequantized_transformer refuses it # and rebuilds dense after the resident pipe is unloaded (the evict-then-OOM this # prefetch avoids). - src = usable_prequant_source(fam, scheme, path_override = prequant_path) + src = usable_prequant_source( + fam, scheme, path_override = prequant_path, base_repo = base_repo + ) prequant_available = src is not None except Exception: # noqa: BLE001 -- prequant probing must never sink the candidate prequant_available = False diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index 1efedd17bc..3865e5849b 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -78,6 +78,12 @@ class DiffusionFamily: # path resolves a scheme with a hosted checkpoint, the loader fetches the already-quantized # weights instead of the dense bf16 (lower load VRAM + smaller download). Empty -> unchanged. prequant_repos: tuple[tuple[str, str], ...] = field(default_factory = tuple) + # Hosted checkpoints for NON-DEFAULT bases of the family, as (base_repo, scheme, repo_id) + # triples with base_repo lowercased. One family entry covers several published variants + # (flux.1: schnell/dev/Krea-dev) whose weights differ, so each variant needs its own baked + # checkpoint; the loader's base_model_id validation correctly refuses the default entry for + # them. Resolution prefers an exact variant match, then falls back to ``prequant_repos``. + prequant_variant_repos: tuple[tuple[str, str, str], ...] = field(default_factory = tuple) # Native (sd.cpp) single-file assets, used only on the no-GPU sd.cpp engine. The transformer GGUF # is shared with diffusers; sd-cli also needs a single-file VAE + text encoder(s) (the base repo # ships those sharded). Each is a (repo_id, filename); ``sd_cpp_text_encoders`` carries a trailing @@ -120,6 +126,15 @@ _FAMILIES: tuple[DiffusionFamily, ...] = ( ("int8", "unsloth/FLUX.1-schnell-FP8"), ("fp8", "unsloth/FLUX.1-schnell-FP8"), ), + # Gate-validated checkpoints baked from the dev / Krea-dev weights (same arch, different + # weights): without these entries the default schnell checkpoint is refused for those + # bases and every int8/fp8 load pays the dense download + on-the-fly quantise. + prequant_variant_repos = ( + ("black-forest-labs/flux.1-dev", "int8", "unsloth/FLUX.1-dev-FP8"), + ("black-forest-labs/flux.1-dev", "fp8", "unsloth/FLUX.1-dev-FP8"), + ("black-forest-labs/flux.1-krea-dev", "int8", "unsloth/FLUX.1-Krea-dev-FP8"), + ("black-forest-labs/flux.1-krea-dev", "fp8", "unsloth/FLUX.1-Krea-dev-FP8"), + ), aliases = ("flux1", "flux-1"), # LoRA training targets FLUX.1-dev via the DiT trainer (QLoRA nf4); the dev repo is gated. trainable = True, @@ -493,8 +508,20 @@ def default_generation_params(*identifiers: Optional[str]) -> tuple[int, float]: return _GENERATION_DEFAULT_FALLBACK -def family_prequant_repo(fam: DiffusionFamily, scheme: str) -> Optional[str]: - """The hosted pre-quantized transformer repo for ``scheme`` in this family, or None.""" +def family_prequant_repo( + fam: DiffusionFamily, scheme: str, base_repo: Optional[str] = None +) -> Optional[str]: + """The hosted pre-quantized transformer repo for ``scheme`` in this family, or None. + + ``base_repo`` (when known) selects a variant-specific checkpoint first: a checkpoint is + baked from ONE base's weights and the loader refuses it for any other base, so a variant + without its own entry still returns the family default (harmless: the base_model_id + validation then falls back to dense-quantise, exactly as before this table existed).""" + base = (base_repo or "").strip().lower() + if base: + for entry_base, entry_scheme, repo_id in fam.prequant_variant_repos: + if entry_base == base and entry_scheme == scheme: + return repo_id for entry_scheme, repo_id in fam.prequant_repos: if entry_scheme == scheme: return repo_id diff --git a/studio/backend/core/inference/diffusion_prequant.py b/studio/backend/core/inference/diffusion_prequant.py index 3fc3f3531e..fc40450385 100644 --- a/studio/backend/core/inference/diffusion_prequant.py +++ b/studio/backend/core/inference/diffusion_prequant.py @@ -121,18 +121,20 @@ def resolve_prequant_source( scheme: str, *, path_override: Optional[str] = None, + base_repo: Optional[str] = None, ) -> Optional[PrequantSource]: """Resolve where the checkpoint for ``(fam, scheme)`` comes from. Priority: (1) explicit local ``path_override``; (2) the family's hosted repo for - ``scheme``; (3) None -> no pre-quant, caller quantises dense. Pure: no IO, no torch. + ``scheme`` (variant-specific when ``base_repo`` names a base with its own baked + checkpoint); (3) None -> no pre-quant, caller quantises dense. Pure: no IO, no torch. """ override = (path_override or "").strip() if override: return PrequantSource(kind = "path", location = override, filename = None) try: from .diffusion_families import family_prequant_repo - repo_id = family_prequant_repo(fam, scheme) + repo_id = family_prequant_repo(fam, scheme, base_repo = base_repo) except Exception: # noqa: BLE001 — a bad family object must not break the load repo_id = None if repo_id: @@ -150,13 +152,16 @@ def usable_prequant_source( scheme: str, *, path_override: Optional[str] = None, + base_repo: Optional[str] = None, ) -> Optional[PrequantSource]: """``resolve_prequant_source``, but a local path counts only when the loader would accept it: inside the allowlist AND present on disk. Otherwise resolves to None so memory planning falls back to dense-fit checks up front, instead of the loader refusing the path only after the resident pipeline was evicted and dense bf16 materialises under a plan that never budgeted for it (evict-then-OOM). Hosted-repo sources are unaffected.""" - src = resolve_prequant_source(fam, scheme, path_override = path_override) + src = resolve_prequant_source( + fam, scheme, path_override = path_override, base_repo = base_repo + ) if src is not None and src.kind == "path" and not local_prequant_path_ready(src.location): return None return src diff --git a/studio/backend/tests/test_diffusion_auto_policy.py b/studio/backend/tests/test_diffusion_auto_policy.py index 239360a82d..d4c0dfe85f 100644 --- a/studio/backend/tests/test_diffusion_auto_policy.py +++ b/studio/backend/tests/test_diffusion_auto_policy.py @@ -105,7 +105,11 @@ def _patch_selector( ) import core.inference.diffusion_prequant as pq - monkeypatch.setattr(pq, "resolve_prequant_source", lambda fam, s, path_override = None: prequant) + monkeypatch.setattr( + pq, + "resolve_prequant_source", + lambda fam, s, path_override = None, base_repo = None: prequant, + ) # Neutralize the cache-disk gate by default so resolution tests are independent of the # runner's free space (a small CI disk otherwise drops the candidate). The two disk-gate # tests re-patch this after calling the helper to exercise the gate explicitly. diff --git a/studio/backend/tests/test_diffusion_prequant.py b/studio/backend/tests/test_diffusion_prequant.py index e538ee5ae2..af4ef3ba16 100644 --- a/studio/backend/tests/test_diffusion_prequant.py +++ b/studio/backend/tests/test_diffusion_prequant.py @@ -26,13 +26,14 @@ from core.inference.diffusion_prequant import ( # ── resolve_prequant_source ────────────────────────────────────────────────────── -def _fam(prequant_repos = ()): +def _fam(prequant_repos = (), prequant_variant_repos = ()): return DiffusionFamily( name = "z-image", pipeline_class = "ZImagePipeline", transformer_class = "ZImageTransformer2DModel", base_repo = "Tongyi-MAI/Z-Image-Turbo", prequant_repos = prequant_repos, + prequant_variant_repos = prequant_variant_repos, ) @@ -60,6 +61,57 @@ def test_prequant_repo_filename_convention(): assert prequant_repo_filename("org/PlainRepo", "int8") == "PlainRepo-INT8.pt" +def test_resolve_variant_base_picks_variant_repo(): + # A base with its own baked checkpoint resolves to the variant repo; case-insensitive. + fam = _fam( + prequant_repos = (("int8", "org/default-fp8"),), + prequant_variant_repos = (("org/model-dev", "int8", "org/dev-fp8"),), + ) + src = resolve_prequant_source(fam, "int8", base_repo = "Org/Model-DEV") + assert src.kind == "repo" and src.location == "org/dev-fp8" + assert src.filename == "dev-INT8.pt" + + +def test_resolve_variant_base_falls_back_to_default(): + # An unknown variant base (or no base at all) keeps the family default entry: the + # loader's base_model_id validation then refuses it and dense-quantises, as before. + fam = _fam( + prequant_repos = (("int8", "org/default-fp8"),), + prequant_variant_repos = (("org/model-dev", "int8", "org/dev-fp8"),), + ) + assert resolve_prequant_source(fam, "int8").location == "org/default-fp8" + assert ( + resolve_prequant_source(fam, "int8", base_repo = "org/other-variant").location + == "org/default-fp8" + ) + # Scheme still has to match within the variant table. + assert ( + resolve_prequant_source(fam, "int8", base_repo = "org/model-dev").location + == "org/dev-fp8" + ) + + +def test_flux1_variant_prequant_wiring(): + # The real flux.1 entry serves schnell by default and dev / Krea-dev via variants. + from core.inference.diffusion_families import detect_family, family_prequant_repo + + fam = detect_family("black-forest-labs/FLUX.1-schnell") + for scheme in ("int8", "fp8"): + assert family_prequant_repo(fam, scheme) == "unsloth/FLUX.1-schnell-FP8" + assert ( + family_prequant_repo( + fam, scheme, base_repo = "black-forest-labs/FLUX.1-dev" + ) + == "unsloth/FLUX.1-dev-FP8" + ) + assert ( + family_prequant_repo( + fam, scheme, base_repo = "black-forest-labs/FLUX.1-Krea-dev" + ) + == "unsloth/FLUX.1-Krea-dev-FP8" + ) + + def test_resolve_wrong_scheme_is_none(): fam = _fam(prequant_repos = (("fp8", "org/hosted-fp8"),)) assert resolve_prequant_source(fam, "int8") is None