diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index d5a80b667f..2c531eda1e 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -1380,6 +1380,7 @@ class DiffusionBackend: prequant_path = transformer_prequant_path, allow_dense_fallback = dense_fallback_allowed, lora_specs = loras, + text_encoder_quant = text_encoder_quant, ) except Exception as exc: # noqa: BLE001 — fall back to the GGUF build logger.warning( @@ -1812,6 +1813,7 @@ class DiffusionBackend: base_local_dir: Optional[str] = None, allow_dense_fallback: bool = True, lora_specs: Optional[list[tuple[str, float]]] = None, + text_encoder_quant: Optional[str] = None, ) -> tuple[Any, str]: """Build the opt-in fast pipeline and return ``(pipe, engaged_scheme)``. @@ -1867,7 +1869,7 @@ class DiffusionBackend: if transformer is not None: pipe = self._assemble_pipe( pipeline_cls, base, transformer, dtype, hf_token, device, base_local_dir, - fam = fam, + fam = fam, te_quant_mode = text_encoder_quant, target = target, ) return pipe, scheme @@ -1882,7 +1884,8 @@ class DiffusionBackend: base, subfolder = "transformer", torch_dtype = dtype, token = hf_token ) pipe = self._assemble_pipe( - pipeline_cls, base, transformer, dtype, hf_token, device, base_local_dir, fam = fam + pipeline_cls, base, transformer, dtype, hf_token, device, base_local_dir, + fam = fam, te_quant_mode = text_encoder_quant, target = target, ) if lora_specs: # Bake the adapters BEFORE quantize_: peft injects its wrappers on the dense @@ -1929,6 +1932,8 @@ class DiffusionBackend: device: str, base_local_dir: Optional[str] = None, fam: Optional[DiffusionFamily] = None, + te_quant_mode: Optional[str] = None, + target: Any = None, ) -> Any: """Assemble the diffusers pipeline around ``transformer`` and place it on ``device`` (a no-op for an already-placed pre-quantized transformer; it moves the companions).""" @@ -1948,6 +1953,20 @@ class DiffusionBackend: # The repo ships no Llama text_encoder_4; assemble it from the open mirror # (diffusion_hidream.py) exactly like the full-pipeline load branch. pipe_kwargs.update(hidream_te4_kwargs(dtype, hf_token)) + # Same pre-cast TE injection as the full-pipeline and GGUF branches: the dense + # fast path supplies only the transformer, so the companion TE is the big download. + if target is not None: + pipe_kwargs.update( + te_prequant_pipe_kwargs( + fam, + base, + te_quant_mode = te_quant_mode, + target = target, + dtype = dtype, + hf_token = hf_token, + logger = logger, + ) + ) pipe = pipeline_cls.from_pretrained(base_local_dir or base, **pipe_kwargs) pipe.to(device) return pipe diff --git a/studio/backend/core/inference/diffusion_precision.py b/studio/backend/core/inference/diffusion_precision.py index f0919a179e..7b19e9fa98 100644 --- a/studio/backend/core/inference/diffusion_precision.py +++ b/studio/backend/core/inference/diffusion_precision.py @@ -288,6 +288,22 @@ def _cast_fp8(encoder: Any, target: Any) -> None: skip_modules_classes = (torch.nn.Embedding,), ) + # Module.dtype reports the first floating parameter, which is now fp8 STORAGE; pipelines + # derive tensor dtypes from encoder.dtype (Flux2 casts prompt embeds to it and feeds the + # result to randn_tensor, which has no fp8 kernel; VLM pipelines cast pixel_values to it, + # racing the upcast hooks). The encoder computes in target.dtype, so report that. + compute_dtype = getattr(target, "dtype", None) + if compute_dtype is not None and not getattr(encoder, "_unsloth_te_dtype_override", False): + cls = type(encoder) + encoder.__class__ = type( + cls.__name__, + (cls,), + { + "dtype": property(lambda self, _d = compute_dtype: _d), + "_unsloth_te_dtype_override": True, + }, + ) + def _has_layerwise_hooks(encoder: Any) -> bool: """True when any submodule already carries the diffusers layerwise-casting hook.""" diff --git a/studio/backend/tests/test_diffusion_te_prequant.py b/studio/backend/tests/test_diffusion_te_prequant.py index 6a479a7703..4ddb7af59b 100644 --- a/studio/backend/tests/test_diffusion_te_prequant.py +++ b/studio/backend/tests/test_diffusion_te_prequant.py @@ -263,6 +263,39 @@ def test_hosted_te_prequant_entries(): ) == "LTX-2-text_encoder-FP8.pt" +def test_assemble_pipe_injects_precast_te(monkeypatch): + """The dense transformer_quant fast path assembles companions through _assemble_pipe, + which must inject the hosted pre-cast TE like the full-pipeline and GGUF branches.""" + import core.inference.diffusion as dif + + seen: dict = {} + + class FakePipe: + def to(self, device): + return self + + class FakePipelineCls: + @staticmethod + def from_pretrained(base, **kw): + seen.update(kw) + return FakePipe() + + monkeypatch.setattr( + dif, "te_prequant_pipe_kwargs", lambda *a, **k: {"text_encoder": "PRECAST"} + ) + dif.DiffusionBackend._assemble_pipe( + FakePipelineCls, "org/base", "TR", None, None, "cpu", None, + fam = None, te_quant_mode = "fp8", target = object(), + ) + assert seen["text_encoder"] == "PRECAST" + seen.clear() + # No target (defensive default) keeps the assembly unchanged. + dif.DiffusionBackend._assemble_pipe( + FakePipelineCls, "org/base", "TR", None, None, "cpu", None, fam = None, + ) + assert "text_encoder" not in seen + + def test_cast_fp8_is_idempotent_on_precast_encoder(): """A pre-cast encoder arrives with the layerwise hooks installed; the runtime re-apply in quantize_text_encoders must be a no-op (re-registering the hook name raises, which made @@ -275,8 +308,13 @@ def test_cast_fp8_is_idempotent_on_precast_encoder(): enc = torch.nn.Sequential(torch.nn.Linear(64, 64), torch.nn.LayerNorm(64)) _cast_fp8(enc, target) assert enc[0].weight.dtype == torch.float8_e4m3fn + # 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) _cast_fp8(enc, target) # must not raise assert enc[0].weight.dtype == torch.float8_e4m3fn + assert enc.dtype == torch.bfloat16 def test_builder_metadata_survives_weights_only_load(tmp_path):