ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Resolve pre-quantized checkpoints per base variant

Browse source 
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.
This commit is contained in:
Daniel Han 2026-07-17 11:07:17 +00:00
commit 570cef6c79
6 changed files with 105 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

9
studio/backend/core/inference/diffusion.py
View file

@ -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,

4
studio/backend/core/inference/diffusion_auto_policy.py
View file

@ -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

31
studio/backend/core/inference/diffusion_families.py
View file

@ -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

11
studio/backend/core/inference/diffusion_prequant.py
View file

@ -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

6
studio/backend/tests/test_diffusion_auto_policy.py
View file

@ -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.

54
studio/backend/tests/test_diffusion_prequant.py
View file

@ -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