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

Host pre-cast fp8 text encoders for four more families

Browse source 
Round 2 of the hosted TE set, each bit-identical to dense-load-then-cast
and gated through the real backend (marker + status fp8 + same-seed LPIPS
vs dense TEs):

- FLUX.1 T5-XXL (text_encoder_2): 9.52 -> 5.90 GB, one artifact for
  schnell/dev/Krea-dev (T5 shards byte-identical across all three,
  verified sha256). 220 tensors, 144 fp8, LPIPS 0.109.
- Lumina Gemma2-2B: fp32 hub store 10.46 -> 3.20 GB (3.3x download cut).
  288 tensors, 182 fp8, LPIPS 0.041.
- Z-Image Qwen3-4B: 8.04 -> 4.41 GB. 399 tensors, 252 fp8, LPIPS 0.112.
  NOT shared with flux.2-klein-4B: klein retrained layer 35's MLP
  (verified tensor diff, maxdiff 0.86), so klein hosts no entry.
- Krea-2 Qwen3-VL-4B: 8.88 -> 4.83 GB. 713 tensors, 460 fp8, LPIPS 0.082.
  The constructor-assembled krea pipeline takes the encoder directly
  (load_krea2_pipeline text_encoder kwarg); the loader remaps 5.x
  rope_parameters and re-ties weights after assign so the rebuilt encoder
  matches the builder's structure.

HunyuanImage 2.1 reuses the Qwen-Image artifact outright: its Qwen2.5-VL
text encoder is byte-identical (every shard sha256, 16,584,414,544 bytes),
recorded in the new component-level base-equivalence table the checkpoint
validator consults. The injection loop now covers text_encoder.._3 so a
family can host several components. Live check: LPIPS 0.123 vs dense.
This commit is contained in:
Daniel Han 2026-07-18 10:24:51 +00:00
commit 4a1ab9ac36
6 changed files with 245 additions and 30 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -1425,8 +1425,22 @@ class DiffusionBackend:
# any embedded quantization_config (e.g. bnb-4bit).
if fam.name == KREA2_FAMILY_NAME:
# krea ships transformers-5.x configs the 4.x line can't parse; assemble
# per-component (see diffusion_krea2.py).
pipe = load_krea2_pipeline(repo_id, dtype, hf_token = hf_token)
# per-component (see diffusion_krea2.py). The constructor path never
# sees pipe_kwargs, so the pre-cast TE is handed in directly.
pipe = load_krea2_pipeline(
repo_id,
dtype,
hf_token = hf_token,
text_encoder = te_prequant_pipe_kwargs(
fam,
repo_id,
te_quant_mode = text_encoder_quant,
target = target,
dtype = dtype,
hf_token = hf_token,
logger = logger,
).get("text_encoder"),
)
elif fam.name == IDEOGRAM4_FAMILY_NAME:
# ideogram ships the same transformers-5.x Qwen stack as krea; assemble
# per-component too (see diffusion_ideogram4.py).
@ -1495,7 +1509,20 @@ class DiffusionBackend:
if fam.name == KREA2_FAMILY_NAME:
pipe = load_krea2_pipeline(
base, dtype, hf_token = hf_token, transformer = transformer
base,
dtype,
hf_token = hf_token,
transformer = transformer,
# Same pre-cast TE hand-in as the full-pipeline branch.
text_encoder = te_prequant_pipe_kwargs(
fam,
base,
te_quant_mode = text_encoder_quant,
target = target,
dtype = dtype,
hf_token = hf_token,
logger = logger,
).get("text_encoder"),
)
else:
pipe_kwargs = {"torch_dtype": dtype, "transformer": transformer}
@ -1957,8 +1984,23 @@ class DiffusionBackend:
# krea ships transformers-5.x configs and no top-level tokenizer files, so
# Pipeline.from_pretrained dies in the tokenizer (vocab_file = None); assemble
# per-component like every other krea load path (see diffusion_krea2.py).
krea_te = None
if target is not None:
krea_te = te_prequant_pipe_kwargs(
fam,
base,
te_quant_mode = te_quant_mode,
target = target,
dtype = dtype,
hf_token = hf_token,
logger = logger,
).get("text_encoder")
pipe = load_krea2_pipeline(
base_local_dir or base, dtype, hf_token = hf_token, transformer = transformer
base_local_dir or base,
dtype,
hf_token = hf_token,
transformer = transformer,
text_encoder = krea_te,
)
pipe.to(device)
return pipe

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

@ -141,6 +141,10 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
("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"),
),
# Pre-cast T5-XXL (9.52 -> 5.90 GB; CLIP-L stays dense, 0.25 GB). One artifact
# serves schnell/dev/Krea-dev: the T5 shards are byte-identical across all three
# (verified sha256, see diffusion_te_prequant._TE_EQUIVALENT_BASES).
te_prequant_repos = (("fp8", "text_encoder_2", "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,
@ -284,6 +288,9 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
("int8", "unsloth/Z-Image-Turbo-FP8"),
("fp8", "unsloth/Z-Image-Turbo-FP8"),
),
# Pre-cast Qwen3-4B TE (8.04 -> 4.41 GB). NOT shared with flux.2-klein-4B: klein's
# TE retrained layer 35's MLP (up/down_proj maxdiff 0.86 vs this checkpoint).
te_prequant_repos = (("fp8", "text_encoder", "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,
@ -309,6 +316,9 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
("int8", "unsloth/Krea-2-Turbo-FP8"),
("fp8", "unsloth/Krea-2-Turbo-FP8"),
),
# Pre-cast Qwen3-VL-4B TE (8.88 -> 4.83 GB); handed into load_krea2_pipeline
# directly (constructor assembly never sees pipe_kwargs).
te_prequant_repos = (("fp8", "text_encoder", "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
@ -336,6 +346,9 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
("int8", "unsloth/Lumina-Image-2.0-FP8"),
("fp8", "unsloth/Lumina-Image-2.0-FP8"),
),
# Pre-cast Gemma2-2B TE. The Hub stores it fp32 (10.46 GB), so the 3.20 GB
# artifact is a 3.3x download cut even though the model is small.
te_prequant_repos = (("fp8", "text_encoder", "unsloth/Lumina-Image-2.0-FP8"),),
aliases = ("lumina-image-2.0", "lumina-image-2", "lumina2"),
# Published and validated bf16-only upstream; keep the fp16 fallback off like z-image.
fp16_incompatible = True,
@ -358,6 +371,10 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
("int8", "unsloth/HunyuanImage-2.1-FP8"),
("fp8", "unsloth/HunyuanImage-2.1-FP8"),
),
# The Qwen2.5-VL TE is byte-identical to Qwen-Image's (verified sha256, see
# _TE_EQUIVALENT_BASES), so the family reuses the Qwen-Image artifact: zero new
# hosting, 16.58 -> 8.84 GB download. ByT5 (text_encoder_2) stays dense.
te_prequant_repos = (("fp8", "text_encoder", "unsloth/Qwen-Image-FP8"),),
pipeline_class = "HunyuanImagePipeline",
transformer_class = "HunyuanImageTransformer2DModel",
base_repo = "hunyuanvideo-community/HunyuanImage-2.1-Diffusers",

10
studio/backend/core/inference/diffusion_krea2.py
View file

@ -104,13 +104,16 @@ def load_krea2_pipeline(
hf_token: Optional[str] = None,
transformer = None,
with_transformer: bool = True,
text_encoder = None,
):
"""A ready ``Krea2Pipeline`` for ``repo_id`` (still on CPU; caller places it).
``transformer`` lets the single-file/quant paths hand in a prebuilt denoiser;
``with_transformer = False`` skips the (26 GB) denoiser entirely for a
conditioning-only pipeline (the trainer's phased load). The remaining components
(VAE, text encoder, tokenizer, scheduler) come from the repo.
conditioning-only pipeline (the trainer's phased load). ``text_encoder`` lets the
pre-cast TE path (diffusion_te_prequant) hand in an already-built encoder, skipping
the dense Qwen3-VL download. The remaining components (VAE, tokenizer, scheduler)
come from the repo.
"""
import diffusers
@ -125,7 +128,8 @@ def load_krea2_pipeline(
token = hf_token or None
tokenizer = load_krea2_tokenizer(repo_id, hf_token = token)
text_encoder = load_krea2_text_encoder(repo_id, dtype, hf_token = token)
if text_encoder is None:
text_encoder = load_krea2_text_encoder(repo_id, dtype, hf_token = token)
scheduler = diffusers.FlowMatchEulerDiscreteScheduler.from_pretrained(
repo_id, subfolder = "scheduler", token = token
)

96
studio/backend/core/inference/diffusion_te_prequant.py
View file

@ -41,6 +41,46 @@ TE_PREQUANT_FORMAT = "unsloth_prequant_text_encoder_state_dict_v1"
# The one scheme hosted in v1 (see module docstring).
TE_PREQUANT_SCHEMES = ("fp8",)
# Components the pipeline-assembly injection covers (the attrs quantize_text_encoders
# casts; text_encoder_4 is family-assembled separately, see diffusion_hidream.py).
TE_PREQUANT_COMPONENTS = ("text_encoder", "text_encoder_2", "text_encoder_3")
# Bases whose text-encoder weights are VERIFIED byte-identical, so one hosted artifact
# serves all of them. Verification: every safetensors shard's LFS sha256 compared across
# repos on 2026-07-18 (huggingface_hub list_repo_tree; no local download needed). The
# checkpoint validator accepts a base_model_id from the same group as the loading base;
# everything else keeps the strict refusal. Ids are lowercased.
_TE_EQUIVALENT_BASES: tuple[frozenset[str], ...] = (
# Qwen2.5-VL-7B text encoder: 4 shards, 16,584,414,544 bytes, identical sha256 set.
frozenset(
{
"qwen/qwen-image",
"hunyuanvideo-community/hunyuanimage-2.1-diffusers",
}
),
# T5-XXL (text_encoder_2): 2 shards, 9,524,648,584 bytes, identical sha256 set across
# every FLUX.1 release; HiDream-I1 ships the same bytes as text_encoder_3 (cross-
# component filename/metadata mapping is not wired yet, entry documents the identity).
frozenset(
{
"black-forest-labs/flux.1-schnell",
"black-forest-labs/flux.1-dev",
"black-forest-labs/flux.1-krea-dev",
"hidream-ai/hidream-i1-full",
}
),
)
def te_base_equivalent(ckpt_base: str, base: str) -> bool:
"""True when the checkpoint's baked base and the loading base carry byte-identical
weights for the component: the same repo (``_same_base_model``) or a verified
equivalence group above."""
if _same_base_model(ckpt_base, base):
return True
a, b = str(ckpt_base).strip().lower(), str(base).strip().lower()
return any(a in group and b in group for group in _TE_EQUIVALENT_BASES)
@dataclass(frozen = True)
class TePrequantSource:
@ -174,6 +214,13 @@ def load_prequant_text_encoder(
if subfolder:
config_kwargs["subfolder"] = subfolder
config = transformers.AutoConfig.from_pretrained(base, **config_kwargs)
# Krea-2 ships transformers-5.x configs whose rope lives under rope_parameters;
# the runtime component loader remaps it for a 4.x runtime, and the meta-init
# here must match or the rebuilt encoder forwards with a broken rope. No-op for
# every other family (and on a 5.x runtime).
from .diffusion_krea2 import remap_rope_parameters
remap_rope_parameters(getattr(config, "text_config", config))
for key, value in (config_overrides or {}).items():
setattr(config, key, value)
from accelerate import init_empty_weights
@ -188,6 +235,14 @@ def load_prequant_text_encoder(
# meta. Rebuild on CPU so they hold real values, then re-assign the cast weights.
encoder = encoder_cls(config)
encoder.load_state_dict(state_dict, strict = True, assign = True)
# assign=True swaps in SEPARATE tensors for tied weights (the saved dict carries a
# copy per key), untying e.g. Qwen3's lm_head from embed_tokens. An untied head
# defeats _cast_fp8's tied-projection skip below (the head would get cast while the
# builder's did not, breaking bit-identity and duplicating the embedding). Re-tie to
# the builder-identical structure; a no-op for untied configs.
tie = getattr(encoder, "tie_weights", None)
if callable(tie):
tie()
encoder.eval()
# Install the SAME upcast hooks the runtime cast applies. The weight cast inside is
@ -225,9 +280,10 @@ def te_prequant_pipe_kwargs(
hf_token: Optional[str] = None,
logger: Any = None,
) -> dict[str, Any]:
"""Component overrides for pipeline assembly: ``{"text_encoder": <pre-cast encoder>}``
when the requested TE quant is layerwise fp8 and this family hosts a pre-cast
checkpoint for its primary encoder; ``{}`` otherwise (assembly loads dense as today).
"""Component overrides for pipeline assembly: ``{<component>: <pre-cast encoder>}``
for every ``TE_PREQUANT_COMPONENTS`` attr the family hosts a pre-cast checkpoint for
(e.g. flux.1 hosts its T5-XXL as ``text_encoder_2``); ``{}`` when none resolve
(assembly loads dense as today).
Gated exactly like the runtime cast (mode normalized, device-supported, family not
denied), so injection can never engage where ``quantize_text_encoders`` would not.
@ -252,21 +308,23 @@ def te_prequant_pipe_kwargs(
return {}
if not te_quant_supported(target, mode):
return {}
source = resolve_te_prequant_source(fam, "text_encoder", mode)
if source is None:
return {}
encoder = load_prequant_text_encoder(
base,
"text_encoder",
source,
dtype = dtype,
hf_token = hf_token,
scheme = mode,
logger = logger,
)
if encoder is None:
return {}
return {"text_encoder": encoder}
injected: dict[str, Any] = {}
for component in TE_PREQUANT_COMPONENTS:
source = resolve_te_prequant_source(fam, component, mode)
if source is None:
continue
encoder = load_prequant_text_encoder(
base,
component,
source,
dtype = dtype,
hf_token = hf_token,
scheme = mode,
logger = logger,
)
if encoder is not None:
injected[component] = encoder
return injected
except Exception as exc: # noqa: BLE001 — injection is an optimisation, never a blocker
_warn(logger, "pipe_kwargs", exc)
return {}
@ -324,7 +382,7 @@ def _validate_checkpoint(ckpt: Any, scheme: str, component: str, base: str, logg
),
)
return False
if not _same_base_model(ckpt_base, base):
if not te_base_equivalent(ckpt_base, base):
_warn(logger, scheme, ValueError(f"checkpoint base {ckpt_base!r} != {base!r}"))
return False
return True

2
studio/backend/tests/test_diffusion_backend.py
View file

@ -3073,7 +3073,7 @@ def test_assemble_pipe_routes_krea2_per_component(monkeypatch):
calls["device"] = device
return self
def fake_loader(base, dtype, hf_token = None, transformer = None):
def fake_loader(base, dtype, hf_token = None, transformer = None, text_encoder = None):
calls["base"] = base
calls["transformer"] = transformer
return Pipe()

100
studio/backend/tests/test_diffusion_te_prequant.py
View file

@ -225,6 +225,70 @@ def test_pipe_kwargs_empty_when_load_fails(monkeypatch):
) == {}
def test_pipe_kwargs_injects_every_hosted_component(monkeypatch):
"""A family hosting several TE components (flux.1: T5 as text_encoder_2) gets each
one injected under its own attr; unhosted components stay dense."""
import core.inference.diffusion_precision as precision
fam = _fam(
te_prequant_repos = (
("fp8", "text_encoder", "org/hosted"),
("fp8", "text_encoder_2", "org/hosted-2"),
)
)
monkeypatch.setattr(precision, "te_quant_supported", lambda target, mode: True)
markers = {"text_encoder": object(), "text_encoder_2": object()}
monkeypatch.setattr(
tpq,
"load_prequant_text_encoder",
lambda base, component, source, **kw: markers[component],
)
out = te_prequant_pipe_kwargs(
fam, "some/base", te_quant_mode = "fp8", target = _target(), dtype = None
)
assert out == markers
# ── base equivalence ─────────────────────────────────────────────────────────
def test_te_base_equivalent_groups():
from core.inference.diffusion_te_prequant import te_base_equivalent
# Same repo (case-folded) always matches.
assert te_base_equivalent("Qwen/Qwen-Image", "qwen/qwen-image")
# Verified byte-identical groups match across repos, both directions.
assert te_base_equivalent(
"Qwen/Qwen-Image", "hunyuanvideo-community/HunyuanImage-2.1-Diffusers"
)
assert te_base_equivalent(
"black-forest-labs/FLUX.1-schnell", "black-forest-labs/FLUX.1-dev"
)
assert te_base_equivalent(
"black-forest-labs/FLUX.1-Krea-dev", "black-forest-labs/FLUX.1-schnell"
)
# Unrelated bases stay refused, including across groups.
assert not te_base_equivalent("Qwen/Qwen-Image", "black-forest-labs/FLUX.1-schnell")
assert not te_base_equivalent("Tongyi-MAI/Z-Image-Turbo", "black-forest-labs/FLUX.2-klein-4B")
def test_validate_accepts_equivalent_base():
ckpt = {
"format": TE_PREQUANT_FORMAT,
"state_dict": {},
"metadata": {
"scheme": "fp8",
"component": "text_encoder",
"base_model_id": "Qwen/Qwen-Image",
},
}
assert tpq._validate_checkpoint(
ckpt, "fp8", "text_encoder",
"hunyuanvideo-community/HunyuanImage-2.1-Diffusers", None,
)
assert not tpq._validate_checkpoint(
ckpt, "fp8", "text_encoder", "black-forest-labs/FLUX.1-schnell", None
)
# ── family field wiring ──────────────────────────────────────────────────────
def test_family_dataclasses_declare_te_prequant_field():
from core.inference.diffusion_families import DiffusionFamily, detect_family
@ -232,8 +296,9 @@ def test_family_dataclasses_declare_te_prequant_field():
assert DiffusionFamily.__dataclass_fields__["te_prequant_repos"].default_factory is tuple
assert VideoFamily.__dataclass_fields__["te_prequant_repos"].default_factory is tuple
# Families without a hosted TE checkpoint keep the empty default.
fam = detect_family("unsloth/FLUX.1-schnell-GGUF")
# Families without a hosted TE checkpoint keep the empty default (sdxl's CLIPs
# stay dense; flux.1 now hosts its T5 and is asserted below).
fam = detect_family("stabilityai/stable-diffusion-xl-base-1.0")
assert fam.te_prequant_repos == ()
@ -269,6 +334,29 @@ def test_hosted_te_prequant_entries():
assert te_prequant_repo_filename(
"unsloth/HiDream-I1-Full-FP8", "text_encoder_4", "fp8"
) == "HiDream-I1-Full-text_encoder_4-FP8.pt"
# Round 2: T5-XXL for every flux.1 base (byte-identical weights, one artifact),
# Gemma2-2B, Qwen3-4B, Qwen3-VL-4B, and hunyuanimage reusing the Qwen-Image artifact.
assert detect_family("black-forest-labs/FLUX.1-schnell").te_prequant_repos == (
("fp8", "text_encoder_2", "unsloth/FLUX.1-schnell-FP8"),
)
assert te_prequant_repo_filename(
"unsloth/FLUX.1-schnell-FP8", "text_encoder_2", "fp8"
) == "FLUX.1-schnell-text_encoder_2-FP8.pt"
assert detect_family("Alpha-VLLM/Lumina-Image-2.0").te_prequant_repos == (
("fp8", "text_encoder", "unsloth/Lumina-Image-2.0-FP8"),
)
assert detect_family("Tongyi-MAI/Z-Image-Turbo").te_prequant_repos == (
("fp8", "text_encoder", "unsloth/Z-Image-Turbo-FP8"),
)
assert detect_family("krea/Krea-2-Turbo").te_prequant_repos == (
("fp8", "text_encoder", "unsloth/Krea-2-Turbo-FP8"),
)
assert detect_family(
"hunyuanvideo-community/HunyuanImage-2.1-Diffusers"
).te_prequant_repos == (("fp8", "text_encoder", "unsloth/Qwen-Image-FP8"),)
# flux.2-klein-4B hosts NO TE entry: its Qwen3-4B retrained layer 35's MLP, so the
# z-image artifact must not serve it (verified tensor diff, maxdiff 0.86).
assert detect_family("black-forest-labs/FLUX.2-klein-4B").te_prequant_repos == ()
def _hidream_transformers_stub(monkeypatch, recorder):
@ -441,7 +529,13 @@ def test_cast_fp8_is_idempotent_on_precast_encoder():
# Module.dtype must report the COMPUTE dtype: pipelines derive tensor dtypes from it
# (Flux2 feeds it to randn_tensor, which has no fp8 kernel).
assert enc.dtype == torch.bfloat16
assert isinstance(enc, torch.nn.Sequential)
# EXACT class identity: a dynamic-subclass swap here broke transformers' kwargs-based
# output recording (Qwen3VLModel returned hidden_states=None; krea-2 crashed at encode).
assert type(enc) is torch.nn.Sequential
# An uncast sibling of the same (now property-patched) class keeps original behaviour.
sibling = torch.nn.Sequential(torch.nn.Linear(8, 8))
with pytest.raises(AttributeError):
sibling.dtype
_cast_fp8(enc, target) # must not raise
assert enc[0].weight.dtype == torch.float8_e4m3fn
assert enc.dtype == torch.bfloat16