diff --git a/studio/backend/core/inference/diffusion_precision.py b/studio/backend/core/inference/diffusion_precision.py index d744c7a864..460824ed4b 100644 --- a/studio/backend/core/inference/diffusion_precision.py +++ b/studio/backend/core/inference/diffusion_precision.py @@ -118,7 +118,10 @@ def te_quant_supported(target: Any, mode: str) -> bool: # 4-bit is a steeper quality cost), never an auto pick, consistent with the DiT ladder. _TE_AUTO_LADDER: tuple[tuple[tuple[int, int], tuple[str, ...]], ...] = ( ((8, 9), (TE_QUANT_FP8_DYNAMIC, TE_QUANT_INT8, TE_QUANT_FP8)), # Ada sm_89 / Hopper / Blackwell - ((8, 0), (TE_QUANT_INT8, TE_QUANT_FP8)), # Ampere sm_80/86: no fp8 GEMM -> int8 or layerwise fp8 + ( + (8, 0), + (TE_QUANT_INT8, TE_QUANT_FP8), + ), # Ampere sm_80/86: no fp8 GEMM -> int8 or layerwise fp8 ) # Text encoders whose activation ranges break a scheme at the MODEL level (measured hidden-state @@ -145,7 +148,6 @@ def _te_scheme_probe(scheme: str, device: str) -> bool: return True try: from .diffusion_transformer_quant import _smoke_probe - return _smoke_probe(tq, device) except Exception: return False @@ -179,7 +181,9 @@ def select_te_quant_scheme( # Consumer GDDR parts run int8 full-rate but halve fp8 FP32-accumulate: prefer int8. ordered = ( (TE_QUANT_INT8,) + tuple(s for s in schemes if s != TE_QUANT_INT8) - if TE_QUANT_INT8 in schemes and schemes[0] != TE_QUANT_INT8 and _is_consumer_gpu(device) + if TE_QUANT_INT8 in schemes + and schemes[0] != TE_QUANT_INT8 + and _is_consumer_gpu(device) else schemes ) for scheme in ordered: diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index e342617e23..0be18d2e93 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1745,7 +1745,9 @@ class DiffusionLoadRequest(BaseModel): "default (also regional torch.compile where eligible), " "max (also TF32 + fused QKV).", ) - text_encoder_quant: Optional[Literal["auto", "none", "off", "fp8", "fp8_dynamic", "int8", "nvfp4"]] = Field( + text_encoder_quant: Optional[ + Literal["auto", "none", "off", "fp8", "fp8_dynamic", "int8", "nvfp4"] + ] = Field( None, description = "Quantise the companion text encoder(s). auto (the default when unset) picks " "the fastest accurate scheme for this GPU + model family (fp8_dynamic on fp8-GEMM silicon, " @@ -2348,7 +2350,9 @@ class VideoLoadRequest(BaseModel): "backend's transformer_quant field.", ) ) - text_encoder_quant: Optional[Literal["auto", "none", "off", "fp8", "fp8_dynamic", "int8", "nvfp4"]] = Field( + text_encoder_quant: Optional[ + Literal["auto", "none", "off", "fp8", "fp8_dynamic", "int8", "nvfp4"] + ] = Field( None, description = "Quantise the dense companion text encoder (Gemma3 / UMT5 / Qwen2.5-VL), " "which loads bf16 from the base repo regardless of how the DiT was sourced and is often " diff --git a/studio/backend/tests/test_diffusion_precision.py b/studio/backend/tests/test_diffusion_precision.py index 133ab58173..2bba13fea8 100644 --- a/studio/backend/tests/test_diffusion_precision.py +++ b/studio/backend/tests/test_diffusion_precision.py @@ -404,7 +404,13 @@ def test_nvfp4_filter_keeps_vision_tower_dense(monkeypatch): # ── auto ladder (select_te_quant_scheme) ──────────────────────────────────────── -def _stub_tq_select(monkeypatch, *, cc, consumer = False, smoke = True): +def _stub_tq_select( + monkeypatch, + *, + cc, + consumer = False, + smoke = True, +): """Stub the transformer module's shared helpers that select_te_quant_scheme imports: capability, GPU class, and the kernel smoke probe (bool or a (tq, dev) predicate).""" dtq = types.ModuleType("core.inference.diffusion_transformer_quant")