diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 34ae2d9406..f9bc5ea135 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -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 diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index 20ba774b1a..b09ef97d85 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -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", diff --git a/studio/backend/core/inference/diffusion_krea2.py b/studio/backend/core/inference/diffusion_krea2.py index 3a9cd4f864..1d9ac0ebff 100644 --- a/studio/backend/core/inference/diffusion_krea2.py +++ b/studio/backend/core/inference/diffusion_krea2.py @@ -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 ) diff --git a/studio/backend/core/inference/diffusion_te_prequant.py b/studio/backend/core/inference/diffusion_te_prequant.py index 4bd2a4797a..5e9d90df1b 100644 --- a/studio/backend/core/inference/diffusion_te_prequant.py +++ b/studio/backend/core/inference/diffusion_te_prequant.py @@ -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": }`` - 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: ``{: }`` + 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 diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 152620c909..c7ddaea868 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -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() diff --git a/studio/backend/tests/test_diffusion_te_prequant.py b/studio/backend/tests/test_diffusion_te_prequant.py index 52aaeed83b..fee18b69b9 100644 --- a/studio/backend/tests/test_diffusion_te_prequant.py +++ b/studio/backend/tests/test_diffusion_te_prequant.py @@ -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