Merge remote-tracking branch 'origin/image-generation' into fix/imggen-review-bugs

studio/backend/core/inference/diffusion.py

@@ -423,7 +423,9 @@ class DiffusionBackend:
             target = self._resolve_device_target(fam)
             if not dense_transformer_supported(target):
                 return False
-            scheme = select_transformer_quant_scheme(target, mode)
+            scheme = select_transformer_quant_scheme(
+                target, mode, family = getattr(fam, "name", None)
+            )
             if scheme is None:
                 return False
             source = resolve_prequant_source(
@@ -1329,7 +1331,7 @@ class DiffusionBackend:
         BEFORE the loader compiles the repeated block, so the order stays quantize ->
         compile -> placement."""
         # 1. Pre-quantized checkpoint, when one is configured for the resolved scheme.
-        scheme = select_transformer_quant_scheme(target, mode)
+        scheme = select_transformer_quant_scheme(target, mode, family = getattr(fam, "name", None))
         if scheme is None:
             # Bail BEFORE the (multi-GB) dense download: an explicit unsupported scheme
             # (e.g. fp8 on Ampere, nvfp4 off Blackwell) would otherwise materialise the
@@ -1368,7 +1370,14 @@ class DiffusionBackend:
             base, subfolder = "transformer", torch_dtype = dtype, token = hf_token
         )
         pipe = self._assemble_pipe(pipeline_cls, base, transformer, dtype, hf_token, device)
-        scheme = quantize_transformer(pipe, target, mode = mode, fast_accum = fast_accum, logger = logger)
+        scheme = quantize_transformer(
+            pipe,
+            target,
+            mode = mode,
+            family = getattr(fam, "name", None),
+            fast_accum = fast_accum,
+            logger = logger,
+        )
         if scheme is None:
             raise RuntimeError("transformer quant unsupported for this device/scheme")
         return pipe, scheme

studio/backend/core/inference/diffusion_eager_patches.py

@@ -181,6 +181,11 @@ def install_compile_safe_patches() -> int:
     for cls, new_fn in _specs():
         if cls is None:
             continue
+        # torch < 2.4 has no F.rms_norm: leave diffusers' original RMSNorm.forward in
+        # place rather than installing a patch whose fast path would AttributeError.
+        if cls is _RMSNorm and not hasattr(F, "rms_norm"):
+            logger.info("eager-patch: skipping RMSNorm (this torch has no F.rms_norm)")
+            continue
         # Capture the live original BEFORE patching so the RMSNorm fast path can fall back
         # to it for the uncommon (NPU / bias / fp32-weight / tuple-dim) cases.
         if cls is _RMSNorm:

studio/backend/core/inference/diffusion_transformer_quant.py

@@ -101,6 +101,30 @@ _AUTO_LADDER: tuple[tuple[tuple[int, int], tuple[str, ...]], ...] = (
     ((8, 0), (TQ_INT8,)),  # Ampere sm_80 / sm_86
 )
 
+# Families whose activation ranges break specific dense-quant schemes at the MODEL
+# level. The kernel smoke probe below cannot see this (it only proves the GEMM runs);
+# these were measured with the 28-pair prequant accuracy gate on a B200
+# (scripts/prequant_accuracy_gate.py) and reproduced with on-the-fly quantisation:
+#   qwen-image + fp8   -> every frame black (mean luma 0.0000, SSIM 0.016 vs bf16). The
+#                         same per-row fp8 that matches bf16 on Z-Image / FLUX: Qwen's
+#                         activation outliers exceed even per-row fp8's dynamic range.
+#   qwen-image + mxfp8 -> real semantic damage at 1024px (CLIP delta mean 0.0146, worst
+#                         cases 0.064 / 0.102 -- 2x the per-case bound).
+#   qwen-image + nvfp4 -> LPIPS mean 0.51 vs bf16: unusable.
+# int8 dynamic (per-token) is excellent on Qwen (LPIPS mean 0.069 / SSIM 0.958), so the
+# auto ladder falls through to it. The deny also applies to an EXPLICIT request: a
+# scheme that renders black frames has no legitimate use, and returning None gives the
+# caller the same fallback contract as an unsupported scheme (GGUF build).
+_FAMILY_SCHEME_DENY: dict[str, frozenset[str]] = {
+    "qwen-image": frozenset({TQ_FP8, TQ_MXFP8, TQ_NVFP4}),
+    "qwen-image-edit": frozenset({TQ_FP8, TQ_MXFP8, TQ_NVFP4}),  # same DiT + activations
+}
+
+
+def _family_denied(family, scheme: str) -> bool:
+    return scheme in _FAMILY_SCHEME_DENY.get(str(family or "").strip().lower(), ())
+
+
 # Cache of (scheme, device) -> bool so the quantise+matmul smoke test runs once.
 _SMOKE_CACHE: dict[tuple[str, str], bool] = {}
 
@@ -201,18 +225,27 @@ def dense_transformer_supported(target: Any) -> bool:
         return False
 
 
-def select_transformer_quant_scheme(target: Any, requested: Optional[str]) -> Optional[str]:
+def select_transformer_quant_scheme(
+    target: Any,
+    requested: Optional[str],
+    family: Optional[str] = None,
+) -> Optional[str]:
     """The concrete scheme to apply, or None to fall back to GGUF.
 
     ``auto`` walks the per-arch ladder and returns the first scheme that passes a real
     quantise+matmul smoke test, so on a box where the Blackwell fp4 / mx kernels are
     unavailable it lands on fp8 / int8 with no error. An explicit scheme is honored only
-    if supported (else None -> GGUF), never silently swapped for a different one."""
+    if supported (else None -> GGUF), never silently swapped for a different one.
+    ``family`` additionally applies the measured model-level deny list
+    (``_FAMILY_SCHEME_DENY``): schemes that produce black frames or out-of-bar drift on
+    that family are skipped by ``auto`` and refused when explicit."""
     requested = normalize_transformer_quant(requested)
     if requested is None or not dense_transformer_supported(target):
         return None
     device = str(getattr(target, "device", "cuda"))
     if requested != TQ_AUTO:
+        if _family_denied(family, requested):
+            return None
         return requested if _scheme_supported(requested, device) else None
     cap = _capability()
     if cap is None:
@@ -220,6 +253,8 @@ def select_transformer_quant_scheme(target: Any, requested: Optional[str]) -> Op
     for floor, schemes in _AUTO_LADDER:
         if cap >= floor:
             for scheme in _prefer_consumer_scheme(schemes, device):
+                if _family_denied(family, scheme):
+                    continue
                 if _scheme_supported(scheme, device):
                     return scheme
             return None
@@ -385,6 +420,7 @@ def quantize_transformer(
     target: Any,
     *,
     mode: Optional[str],
+    family: Optional[str] = None,
     min_features: int = DEFAULT_MIN_LINEAR_FEATURES,
     fast_accum: Optional[bool] = None,
     logger: Any = None,
@@ -396,7 +432,7 @@ def quantize_transformer(
 
     ``fast_accum`` (fp8 only) overrides the per-GPU-class accumulate choice: None
     auto-detects (fast on consumer, precise on data-center), True/False force it."""
-    scheme = select_transformer_quant_scheme(target, mode)
+    scheme = select_transformer_quant_scheme(target, mode, family = family)
     if scheme is None:
         return None
     transformer = getattr(pipe, "transformer", None)

studio/backend/core/training/diffusion_dit_trainer.py

@@ -486,6 +486,21 @@ def run_dit_lora_training(
     should_stop: Optional[StopCb] = None,
 ) -> str:
     """Train a flow-matching DiT LoRA (FLUX.1-dev / Qwen-Image / Z-Image) and export it."""
+    cfg = config.normalized()
+    spec = _SPECS.get(cfg.resolved_family)
+    if spec is None:
+        raise ValueError(f"No DiT trainer for family {cfg.resolved_family!r}")
+
+    # DiT families train in bf16 (Z-Image/Qwen require it; FLUX prefers it). A caller that
+    # explicitly asks for fp16 on a bf16-only family is refused rather than silently
+    # upgraded, so the choice is never misrepresented. Validation runs before the heavy
+    # imports so a host without diffusers still sees the real error.
+    if cfg.mixed_precision == "fp16" and spec.force_bf16:
+        raise ValueError(
+            f"{spec.family} LoRA training requires bf16: fp16 overflows its fp32 RoPE / "
+            f"embedder internals. Set mixed precision to bf16."
+        )
+
     import torch
     import torch.nn.functional as F
     from diffusers import FlowMatchEulerDiscreteScheduler
@@ -493,11 +508,6 @@ def run_dit_lora_training(
     from peft import LoraConfig
     from peft.utils import get_peft_model_state_dict
 
-    cfg = config.normalized()
-    spec = _SPECS.get(cfg.resolved_family)
-    if spec is None:
-        raise ValueError(f"No DiT trainer for family {cfg.resolved_family!r}")
-
     rng = random.Random(cfg.seed)
     torch.manual_seed(cfg.seed)
 
@@ -514,15 +524,6 @@ def run_dit_lora_training(
             save_on_stop = False
         return True
 
-    # DiT families train in bf16 (Z-Image/Qwen require it; FLUX prefers it). A caller that
-    # explicitly asks for fp16 on a bf16-only family is refused rather than silently
-    # upgraded, so the choice is never misrepresented.
-    if cfg.mixed_precision == "fp16" and spec.force_bf16:
-        raise ValueError(
-            f"{spec.family} LoRA training requires bf16: fp16 overflows its fp32 RoPE / "
-            f"embedder internals. Set mixed precision to bf16."
-        )
-
     device = "cuda" if torch.cuda.is_available() else "cpu"
     # The flow-matching + 4-bit path is bf16 throughout (fp32 on a CPU-only box, which is
     # unsupported for real runs but keeps import/unit tests architecture-agnostic).
@@ -639,8 +640,11 @@ def run_dit_lora_training(
             (loss / cfg.gradient_accumulation_steps).backward()
             step_loss += float(loss.detach()) / cfg.gradient_accumulation_steps
 
+        grad_norm: Optional[float] = None
         if cfg.max_grad_norm and cfg.max_grad_norm > 0:
-            torch.nn.utils.clip_grad_norm_(lora_params, cfg.max_grad_norm)
+            # clip_grad_norm_ returns the PRE-clip total norm: the signal the grad-norm
+            # chart wants (spikes stay visible even when clipping flattens the update).
+            grad_norm = float(torch.nn.utils.clip_grad_norm_(lora_params, cfg.max_grad_norm))
         optimizer.step()
 
         running_loss += step_loss
@@ -661,6 +665,7 @@ def run_dit_lora_training(
                 loss = round(step_loss, 5),
                 avg_loss = round(running_loss / done, 5),
                 learning_rate = cfg.learning_rate,
+                grad_norm = round(grad_norm, 5) if grad_norm is not None else None,
                 samples_per_second = sps,
                 peak_memory_gb = peak_gb or None,
             )

studio/backend/core/training/diffusion_lora_trainer.py

@@ -340,8 +340,10 @@ def run_diffusion_lora_training(
 
         # max_grad_norm <= 0 means "disable clipping" (the Studio payload sends 0.0 for that);
         # passing 0.0 to clip_grad_norm_ would scale every gradient to zero (no learning).
+        grad_norm: Optional[float] = None
         if cfg.max_grad_norm and cfg.max_grad_norm > 0:
-            torch.nn.utils.clip_grad_norm_(lora_params, cfg.max_grad_norm)
+            # clip_grad_norm_ returns the PRE-clip total norm (the grad-norm chart signal).
+            grad_norm = float(torch.nn.utils.clip_grad_norm_(lora_params, cfg.max_grad_norm))
         optimizer.step()
         lr_sched.step()
 
@@ -366,6 +368,7 @@ def run_diffusion_lora_training(
                 loss = round(step_loss, 5),
                 avg_loss = round(running_loss / done, 5),
                 learning_rate = lr_sched.get_last_lr()[0],
+                grad_norm = round(grad_norm, 5) if grad_norm is not None else None,
                 samples_per_second = samples_per_second,
                 peak_memory_gb = peak_gb or None,
             )

studio/backend/core/training/diffusion_train_common.py

@@ -384,9 +384,12 @@ def discover_image_caption_pairs(
             if sidecar.is_file():
                 caption = sidecar.read_text(encoding = "utf-8").strip()
                 break
-        # 2. metadata row keyed by file name (basename or the name as written).
+        # 2. metadata row keyed by file name (basename or the relative path; as_posix so a
+        #    Windows backslash path still matches the jsonl's forward-slash keys).
         if caption is None:
-            caption = meta_caption.get(img.name) or meta_caption.get(str(img.relative_to(root)))
+            caption = meta_caption.get(img.name) or meta_caption.get(
+                img.relative_to(root).as_posix()
+            )
         # 3. dreambooth instance prompt.
         if caption is None and instance_prompt:
             caption = instance_prompt

studio/backend/core/training/diffusion_training_service.py

@@ -67,6 +67,7 @@ def _idle_state() -> dict[str, Any]:
         "loss": None,
         "avg_loss": None,
         "learning_rate": None,
+        "grad_norm": None,
         "num_images": None,
         "in_model_load": False,
         "output_dir": None,
@@ -82,16 +83,25 @@ def _idle_state() -> dict[str, Any]:
         "metric_steps": [],
         "metric_loss": [],
         "metric_lr": [],
+        "metric_grad_norm": [],
     }
 
 
-def _append_metric(state: dict[str, Any], step: Any, loss: Any, lr: Any) -> None:
-    """Append one (step, loss, lr) point to the bounded history arrays on ``state``.
+def _append_metric(
+    state: dict[str, Any],
+    step: Any,
+    loss: Any,
+    lr: Any,
+    grad_norm: Any = None,
+) -> None:
+    """Append one (step, loss, lr, grad_norm) point to the bounded history arrays on
+    ``state``.
 
     Only records finite, positive-step points (mirrors the LLM trainer, which logs history
     only for step > 0 with a real loss). When the arrays hit ``_METRIC_CAP`` they are
     decimated in place (keep every other point) so appends stay bounded without losing the
-    curve's shape. lr may be None (kept as None so the LR series can be sparse)."""
+    curve's shape. lr / grad_norm may be None (kept as None so those series can be
+    sparse while staying index-aligned with ``steps``)."""
     try:
         istep = int(step)
     except (TypeError, ValueError):
@@ -104,22 +114,34 @@ def _append_metric(state: dict[str, Any], step: Any, loss: Any, lr: Any) -> None
         return
     if floss != floss:  # NaN guard
         return
-    flr: Optional[float]
-    try:
-        flr = float(lr) if lr is not None else None
-    except (TypeError, ValueError):
-        flr = None
+
+    def _opt_float(v: Any) -> Optional[float]:
+        try:
+            return float(v) if v is not None else None
+        except (TypeError, ValueError):
+            return None
+
+    flr = _opt_float(lr)
+    fgn = _opt_float(grad_norm)
     steps = state["metric_steps"]
     losses = state["metric_loss"]
     lrs = state["metric_lr"]
+    gns = state["metric_grad_norm"]
     if len(steps) >= _METRIC_CAP:
         state["metric_steps"] = steps[::2]
         state["metric_loss"] = losses[::2]
         state["metric_lr"] = lrs[::2]
-        steps, losses, lrs = state["metric_steps"], state["metric_loss"], state["metric_lr"]
+        state["metric_grad_norm"] = gns[::2]
+        steps, losses, lrs, gns = (
+            state["metric_steps"],
+            state["metric_loss"],
+            state["metric_lr"],
+            state["metric_grad_norm"],
+        )
     steps.append(istep)
     losses.append(floss)
     lrs.append(flr)
+    gns.append(fgn)
 
 
 class DiffusionTrainingService:
@@ -288,6 +310,7 @@ class DiffusionTrainingService:
                     loss = ev.get("loss", s["loss"]),
                     avg_loss = ev.get("avg_loss", s["avg_loss"]),
                     learning_rate = ev.get("learning_rate", s["learning_rate"]),
+                    grad_norm = ev.get("grad_norm", s["grad_norm"]),
                     message = "Training...",
                 )
                 # Fold optional perf fields (emitted by the trainers) so the UI can show
@@ -296,8 +319,14 @@ class DiffusionTrainingService:
                     s["samples_per_second"] = ev.get("samples_per_second")
                 if ev.get("peak_memory_gb") is not None:
                     s["peak_memory_gb"] = ev.get("peak_memory_gb")
-                # Retain a bounded (step, loss, lr) history for the live loss chart.
-                _append_metric(s, ev.get("step"), ev.get("loss"), ev.get("learning_rate"))
+                # Retain a bounded (step, loss, lr, grad_norm) history for the live charts.
+                _append_metric(
+                    s,
+                    ev.get("step"),
+                    ev.get("loss"),
+                    ev.get("learning_rate"),
+                    ev.get("grad_norm"),
+                )
             elif etype == "complete":
                 # Reset in_model_load: a stop during model load emits complete without a
                 # preceding model_load_completed, which would otherwise leave a stale

studio/backend/models/training.py

@@ -734,12 +734,14 @@ class DiffusionTrainingStartResponse(BaseModel):
 
 
 class DiffusionMetricHistory(BaseModel):
-    """Paired step-indexed history arrays for the live training charts. ``lr`` entries may
-    be null so a sparse learning-rate series still aligns with ``steps`` by index."""
+    """Paired step-indexed history arrays for the live training charts. ``lr`` and
+    ``grad_norm`` entries may be null so those sparse series still align with ``steps``
+    by index."""
 
     steps: List[int] = Field(default_factory = list)
     loss: List[float] = Field(default_factory = list)
     lr: List[Optional[float]] = Field(default_factory = list)
+    grad_norm: List[Optional[float]] = Field(default_factory = list)
 
 
 class DiffusionTrainingStatusResponse(BaseModel):
@@ -754,6 +756,9 @@ class DiffusionTrainingStatusResponse(BaseModel):
     loss: Optional[float] = None
     avg_loss: Optional[float] = None
     learning_rate: Optional[float] = None
+    # Pre-clip gradient norm from the trainer's progress events (None when clipping is
+    # disabled), feeding the grad-norm chart.
+    grad_norm: Optional[float] = None
     num_images: Optional[int] = None
     in_model_load: bool = False
     output_dir: Optional[str] = None

studio/backend/routes/training.py

@@ -1267,6 +1267,7 @@ async def diffusion_training_status(current_subject: str = Depends(get_current_s
         steps = snap.pop("metric_steps", []),
         loss = snap.pop("metric_loss", []),
         lr = snap.pop("metric_lr", []),
+        grad_norm = snap.pop("metric_grad_norm", []),
     )
     return DiffusionTrainingStatusResponse(**snap, metric_history = metric_history)
 
@@ -1515,7 +1516,15 @@ def _image_record(
                 caption = None
             break
     if caption is None:
+        # Basename first, then the relative path as written in the jsonl (as_posix so a
+        # Windows backslash path still matches forward-slash keys) -- the same lookup
+        # order discover_image_caption_pairs uses.
         meta = meta_captions.get(image_path.name)
+        if meta is None:
+            try:
+                meta = meta_captions.get(image_path.relative_to(folder).as_posix())
+            except ValueError:
+                meta = None
         if meta is not None:
             caption = meta
             source = "metadata"

studio/backend/tests/test_diffusion_backend.py

@@ -1764,7 +1764,9 @@ def _stub_dense_quant(monkeypatch, *, scheme = "fp8"):
     monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: True)
     # Resolve the scheme without the real GPU smoke probe, and configure no pre-quant
     # checkpoint so the dense materialise+quantise branch is the one exercised.
-    monkeypatch.setattr(dmod, "select_transformer_quant_scheme", lambda target, mode: scheme)
+    monkeypatch.setattr(
+        dmod, "select_transformer_quant_scheme", lambda target, mode, family = None: scheme
+    )
     monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: None)
 
     def _quantize(pipe, target, *, mode, **kw):
@@ -1828,7 +1830,9 @@ def test_transformer_quant_prequant_path_engaged(fake_runtime, tmp_path, monkeyp
     backend = DiffusionBackend()
     _force_cuda_target(backend, monkeypatch)
     monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: True)
-    monkeypatch.setattr(dmod, "select_transformer_quant_scheme", lambda target, mode: "fp8")
+    monkeypatch.setattr(
+        dmod, "select_transformer_quant_scheme", lambda target, mode, family = None: "fp8"
+    )
     monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: object())
     prequant_obj = object()
     loaded: dict = {"n": 0}
@@ -1958,7 +1962,9 @@ def test_transformer_quant_unsupported_scheme_skips_dense_download(
     backend = DiffusionBackend()
     _force_cuda_target(backend, monkeypatch)
     monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: True)
-    monkeypatch.setattr(dmod, "select_transformer_quant_scheme", lambda target, mode: None)
+    monkeypatch.setattr(
+        dmod, "select_transformer_quant_scheme", lambda target, mode, family = None: None
+    )
     monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: None)
 
     @classmethod
@@ -2005,7 +2011,9 @@ def test_dense_quant_prefetch_needed_gates(fake_runtime, monkeypatch):
     _force_cuda_target(backend, monkeypatch)
     fam = detect_family("unsloth/Z-Image-Turbo-GGUF")
     monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: True)
-    monkeypatch.setattr(dmod, "select_transformer_quant_scheme", lambda target, mode: "fp8")
+    monkeypatch.setattr(
+        dmod, "select_transformer_quant_scheme", lambda target, mode, family = None: "fp8"
+    )
     monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: None)
 
     assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8"}) is True
@@ -2016,10 +2024,14 @@ def test_dense_quant_prefetch_needed_gates(fake_runtime, monkeypatch):
     assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8"}) is False
     # Unsupported scheme bails before the dense path (and so must the prefetch).
     monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: None)
-    monkeypatch.setattr(dmod, "select_transformer_quant_scheme", lambda target, mode: None)
+    monkeypatch.setattr(
+        dmod, "select_transformer_quant_scheme", lambda target, mode, family = None: None
+    )
     assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8"}) is False
     # Device without dense support (e.g. non-CUDA) never widens.
-    monkeypatch.setattr(dmod, "select_transformer_quant_scheme", lambda target, mode: "fp8")
+    monkeypatch.setattr(
+        dmod, "select_transformer_quant_scheme", lambda target, mode, family = None: "fp8"
+    )
     monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: False)
     assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8"}) is False
 

studio/backend/tests/test_diffusion_routes.py

@@ -659,6 +659,17 @@ def test_in_progress_returns_409_after_validation_passes(client, monkeypatch):
     backend = _FakeBackend()
     backend.begin_load = _busy
     monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend)
+    # Pin the resolved device to cuda: the route only takes the arbiter for non-CPU
+    # loads, so on a CPU-only host the ownership assert below would never hold.
+    import types as _types
+
+    import core.inference.diffusion_device as devmod
+
+    monkeypatch.setattr(
+        devmod,
+        "resolve_diffusion_device_target",
+        lambda: _types.SimpleNamespace(device = "cuda"),
+    )
     resp = client.post(
         "/api/inference/images/load",
         json = {"model_path": "unsloth/Z-Image-Turbo-GGUF", "gguf_filename": "q.gguf"},

studio/backend/tests/test_diffusion_transformer_quant.py

@@ -433,7 +433,9 @@ def test_fp8_config_uses_per_row_granularity():
 
 
 def test_quantize_transformer_applies_and_marks(monkeypatch):
-    monkeypatch.setattr(tq, "select_transformer_quant_scheme", lambda target, mode: TQ_FP8)
+    monkeypatch.setattr(
+        tq, "select_transformer_quant_scheme", lambda target, mode, family = None: TQ_FP8
+    )
     seen: dict = {}
 
     def _mk(scheme, fast_accum = None):
@@ -458,13 +460,17 @@ def test_quantize_transformer_applies_and_marks(monkeypatch):
 
 
 def test_quantize_transformer_none_when_unsupported(monkeypatch):
-    monkeypatch.setattr(tq, "select_transformer_quant_scheme", lambda target, mode: None)
+    monkeypatch.setattr(
+        tq, "select_transformer_quant_scheme", lambda target, mode, family = None: None
+    )
     pipe = types.SimpleNamespace(transformer = types.SimpleNamespace())
     assert quantize_transformer(pipe, _target(), mode = "auto") is None
 
 
 def test_quantize_transformer_tolerates_failure(monkeypatch):
-    monkeypatch.setattr(tq, "select_transformer_quant_scheme", lambda target, mode: TQ_INT8)
+    monkeypatch.setattr(
+        tq, "select_transformer_quant_scheme", lambda target, mode, family = None: TQ_INT8
+    )
     monkeypatch.setattr(tq, "_make_quant_config", lambda scheme: "cfg")
     tqz = types.ModuleType("torchao.quantization")
 
@@ -480,3 +486,60 @@ def test_quantize_transformer_tolerates_failure(monkeypatch):
     pipe = types.SimpleNamespace(transformer = types.SimpleNamespace())
     # A quantise failure returns None (caller falls back to GGUF), never raises.
     assert quantize_transformer(pipe, _target(), mode = "int8") is None
+
+
+# ── family scheme deny (measured model-level breakage) ────────────────────────
+
+
+def test_family_deny_auto_skips_fp8_for_qwen(monkeypatch):
+    # B200 with every scheme available: auto must NOT pick fp8 / nvfp4 / mxfp8 for the
+    # Qwen DiT (per-row fp8 renders black frames on it; see _FAMILY_SCHEME_DENY) and
+    # falls through the ladder to int8, which measures excellent on Qwen.
+    _stub_torch(monkeypatch, cc = (10, 0))
+    _allow(monkeypatch, {TQ_FP8, TQ_NVFP4, TQ_MXFP8, TQ_INT8})
+    assert select_transformer_quant_scheme(_target(), "auto", family = "qwen-image") == TQ_INT8
+    assert select_transformer_quant_scheme(_target(), "auto", family = "qwen-image-edit") == TQ_INT8
+
+
+def test_family_deny_refuses_explicit_fp8_for_qwen(monkeypatch):
+    # An explicit fp8 request on qwen-image returns None (same contract as an
+    # unsupported scheme: the caller builds the GGUF pipeline instead). int8 stays
+    # honored on qwen, and fp8 stays honored on families outside the deny table.
+    _stub_torch(monkeypatch, cc = (10, 0))
+    _allow(monkeypatch, {TQ_FP8, TQ_INT8})
+    assert select_transformer_quant_scheme(_target(), "fp8", family = "qwen-image") is None
+    assert select_transformer_quant_scheme(_target(), "int8", family = "qwen-image") == TQ_INT8
+    assert select_transformer_quant_scheme(_target(), "fp8", family = "z-image") == TQ_FP8
+
+
+def test_family_deny_no_family_keeps_ladder(monkeypatch):
+    # Without a family (or an unknown one) the ladder is unchanged: fp8 first on B200.
+    _stub_torch(monkeypatch, cc = (10, 0))
+    _allow(monkeypatch, {TQ_FP8, TQ_INT8})
+    assert select_transformer_quant_scheme(_target(), "auto") == TQ_FP8
+    assert select_transformer_quant_scheme(_target(), "auto", family = "sdxl") == TQ_FP8
+
+
+def test_quantize_transformer_threads_family(monkeypatch):
+    # quantize_transformer passes the family down to the selector, so a denied
+    # (family, scheme) pair never reaches torchao.
+    _stub_torch(monkeypatch, cc = (10, 0))
+    _allow(monkeypatch, {TQ_FP8, TQ_INT8})
+    pipe = types.SimpleNamespace(transformer = types.SimpleNamespace())
+    called = {}
+    tqz = types.ModuleType("torchao.quantization")
+
+    def _quantize(
+        module,
+        config,
+        filter_fn = None,
+    ):
+        called["scheme"] = True
+
+    tqz.quantize_ = _quantize
+    tqz.Int8DynamicActivationInt8WeightConfig = lambda: "int8-cfg"
+    tqz.Float8DynamicActivationFloat8WeightConfig = lambda **kw: "fp8-cfg"
+    tqz.PerRow = lambda: "per-row"
+    monkeypatch.setitem(sys.modules, "torchao.quantization", tqz)
+    assert quantize_transformer(pipe, _target(), mode = "fp8", family = "qwen-image") is None
+    assert called == {}