From d979763b7234b298b2e7ba64cb69199794306eb9 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Sat, 4 Jul 2026 01:05:35 +0000
Subject: [PATCH 1/4] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 studio/backend/core/inference/sd_cpp_backend.py |  5 ++++-
 studio/backend/tests/test_sd_cpp_backend.py     | 14 +++++++++++---
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py
index 7632551661..d8cb21d409 100644
--- a/studio/backend/core/inference/sd_cpp_backend.py
+++ b/studio/backend/core/inference/sd_cpp_backend.py
@@ -893,7 +893,10 @@ class SdCppDiffusionBackend:
                 lora_stage = Path(server_lora_dir) / f"gen_{os.urandom(6).hex()}"
                 materialized = diffusion_lora.materialize_native_dir(lora_resolved, lora_stage)
                 lora_payload = [
-                    {"path": f"{lora_stage.name}/{Path(m.path).name}", "multiplier": float(m.weight)}
+                    {
+                        "path": f"{lora_stage.name}/{Path(m.path).name}",
+                        "multiplier": float(m.weight),
+                    }
                     for m in materialized
                 ]
         try:
diff --git a/studio/backend/tests/test_sd_cpp_backend.py b/studio/backend/tests/test_sd_cpp_backend.py
index 0b52ea6e14..be2fc22a4b 100644
--- a/studio/backend/tests/test_sd_cpp_backend.py
+++ b/studio/backend/tests/test_sd_cpp_backend.py
@@ -691,7 +691,11 @@ def _fake_materialize(resolved, dest):
     return out
 
 
-def _patch_lora(monkeypatch, resolved, supported = True):
+def _patch_lora(
+    monkeypatch,
+    resolved,
+    supported = True,
+):
     from core.inference import diffusion_lora as dl
 
     monkeypatch.setattr(dl, "supports_lora", lambda **k: supported)
@@ -706,7 +710,9 @@ def test_generate_oneshot_applies_loras_via_prompt_tags(monkeypatch):
 
     eng = _FakeEngine()
     b = _loaded_backend(engine = eng)  # mode = "oneshot"
-    _patch_lora(monkeypatch, [dl.ResolvedLora("id1", "myalias", "/x/a.safetensors", "safetensors", 0.8)])
+    _patch_lora(
+        monkeypatch, [dl.ResolvedLora("id1", "myalias", "/x/a.safetensors", "safetensors", 0.8)]
+    )
     b.generate(prompt = "a fox", steps = 4, seed = 1, loras = [("id1", 0.8)])
     _, params, _, _ = eng.calls[0]
     assert params.lora_dir is not None and params.lora_apply_mode == "auto"
@@ -725,7 +731,9 @@ def test_generate_server_stages_loras_and_sends_structured_field(monkeypatch, tm
     servers: list = []
     _run_server_load(monkeypatch, b, servers)
     servers[0].lora_dir = str(tmp_path)
-    _patch_lora(monkeypatch, [dl.ResolvedLora("id1", "myalias", "/x/a.safetensors", "safetensors", 0.7)])
+    _patch_lora(
+        monkeypatch, [dl.ResolvedLora("id1", "myalias", "/x/a.safetensors", "safetensors", 0.7)]
+    )
     b.generate(prompt = "x", steps = 4, seed = 1, batch_size = 1, loras = [("id1", 0.7)])
     payload = servers[0].payloads[0]
     assert "lora" in payload and len(payload["lora"]) == 1

From 3bde714cd944ce027a244d2dbc6dcb9e6ebce473 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Sat, 4 Jul 2026 01:06:09 +0000
Subject: [PATCH 2/4] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 studio/backend/core/inference/sd_cpp_backend.py |  5 ++++-
 studio/backend/tests/test_sd_cpp_backend.py     | 14 +++++++++++---
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py
index 7632551661..d8cb21d409 100644
--- a/studio/backend/core/inference/sd_cpp_backend.py
+++ b/studio/backend/core/inference/sd_cpp_backend.py
@@ -893,7 +893,10 @@ class SdCppDiffusionBackend:
                 lora_stage = Path(server_lora_dir) / f"gen_{os.urandom(6).hex()}"
                 materialized = diffusion_lora.materialize_native_dir(lora_resolved, lora_stage)
                 lora_payload = [
-                    {"path": f"{lora_stage.name}/{Path(m.path).name}", "multiplier": float(m.weight)}
+                    {
+                        "path": f"{lora_stage.name}/{Path(m.path).name}",
+                        "multiplier": float(m.weight),
+                    }
                     for m in materialized
                 ]
         try:
diff --git a/studio/backend/tests/test_sd_cpp_backend.py b/studio/backend/tests/test_sd_cpp_backend.py
index 0b52ea6e14..be2fc22a4b 100644
--- a/studio/backend/tests/test_sd_cpp_backend.py
+++ b/studio/backend/tests/test_sd_cpp_backend.py
@@ -691,7 +691,11 @@ def _fake_materialize(resolved, dest):
     return out
 
 
-def _patch_lora(monkeypatch, resolved, supported = True):
+def _patch_lora(
+    monkeypatch,
+    resolved,
+    supported = True,
+):
     from core.inference import diffusion_lora as dl
 
     monkeypatch.setattr(dl, "supports_lora", lambda **k: supported)
@@ -706,7 +710,9 @@ def test_generate_oneshot_applies_loras_via_prompt_tags(monkeypatch):
 
     eng = _FakeEngine()
     b = _loaded_backend(engine = eng)  # mode = "oneshot"
-    _patch_lora(monkeypatch, [dl.ResolvedLora("id1", "myalias", "/x/a.safetensors", "safetensors", 0.8)])
+    _patch_lora(
+        monkeypatch, [dl.ResolvedLora("id1", "myalias", "/x/a.safetensors", "safetensors", 0.8)]
+    )
     b.generate(prompt = "a fox", steps = 4, seed = 1, loras = [("id1", 0.8)])
     _, params, _, _ = eng.calls[0]
     assert params.lora_dir is not None and params.lora_apply_mode == "auto"
@@ -725,7 +731,9 @@ def test_generate_server_stages_loras_and_sends_structured_field(monkeypatch, tm
     servers: list = []
     _run_server_load(monkeypatch, b, servers)
     servers[0].lora_dir = str(tmp_path)
-    _patch_lora(monkeypatch, [dl.ResolvedLora("id1", "myalias", "/x/a.safetensors", "safetensors", 0.7)])
+    _patch_lora(
+        monkeypatch, [dl.ResolvedLora("id1", "myalias", "/x/a.safetensors", "safetensors", 0.7)]
+    )
     b.generate(prompt = "x", steps = 4, seed = 1, batch_size = 1, loras = [("id1", 0.7)])
     payload = servers[0].payloads[0]
     assert "lora" in payload and len(payload["lora"]) == 1

From 39bc37712b97d291a2068d77e07b818549994670 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sat, 4 Jul 2026 01:27:32 +0000
Subject: [PATCH 3/4] Address review: fp32 latent cache stats +
 strict-JSON-safe progress floats

- Latent caches (DiT + SDXL) now hold the posterior mean/std in fp32 and draw the
  per-step sample in fp32, casting only the result to the training dtype. This
  matches the in-loop path (encode fp32 -> sample fp32 -> cast) exactly instead of
  sampling in bf16; the cache is tiny so the doubled RAM is negligible.
- The training service nulls non-finite floats (NaN/Inf loss, avg_loss,
  learning_rate) at its single ingestion point so status snapshots and persisted
  run records stay strict-JSON serializable; the metric history skips non-finite
  loss points. Test covers NaN/Inf progress followed by a finite point.
---
 .../core/training/diffusion_dit_trainer.py    | 25 ++++++----
 .../core/training/diffusion_lora_trainer.py   | 26 ++++++----
 .../training/diffusion_training_service.py    | 47 ++++++++++++++-----
 .../backend/tests/test_diffusion_training.py  | 39 +++++++++++++++
 4 files changed, 106 insertions(+), 31 deletions(-)

diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py
index 047583f881..0d7f244364 100644
--- a/studio/backend/core/training/diffusion_dit_trainer.py
+++ b/studio/backend/core/training/diffusion_dit_trainer.py
@@ -640,8 +640,11 @@ def _load_pixel_tensor_planned(path, resolution, center_crop, u_left, u_top, fli
 
 def _build_latent_cache(spec, vae, image_paths, cfg, device, weight_dtype, on_event, check_stop):
     """Precompute the per-image latent posterior cache: for each planned crop/flip variant,
-    encode once and store the affine (A, B) pair on CPU (pinned when possible) in the
-    training dtype. Returns None if the build was interrupted by a stop request."""
+    encode once and store the affine (A, B) pair on CPU (pinned when possible) in fp32. The
+    stats stay fp32 so the per-step sample happens in fp32 and only the RESULT is cast to
+    weight_dtype, matching the in-loop path (encode fp32 -> sample/normalise fp32 ->
+    .to(weight_dtype)); fp32 doubles the cache RAM over bf16 but the cache is tiny (a handful
+    of latents per image). Returns None if the build was interrupted by a stop request."""
 
     plan = _plan_cache_variants(
         len(image_paths), cfg.cache_variants, cfg.center_crop, cfg.random_flip, cfg.seed
@@ -650,7 +653,9 @@ def _build_latent_cache(spec, vae, image_paths, cfg, device, weight_dtype, on_ev
     def _hold(t):
         if t is None:
             return None
-        t = t.to(weight_dtype).cpu()
+        import torch
+
+        t = t.to(torch.float32).cpu()
         if device == "cuda":
             try:
                 t = t.pin_memory()
@@ -680,10 +685,12 @@ def _build_latent_cache(spec, vae, image_paths, cfg, device, weight_dtype, on_ev
     return cache
 
 
-def _sample_cached_latents(cache, idxs, variant_rng, device):
+def _sample_cached_latents(cache, idxs, variant_rng, device, weight_dtype):
     """Draw one latent per index from the cache: pick a variant, then sample the posterior
     (A + B * randn) when the family is stochastic. Fresh noise per step, exactly like an
-    in-loop ``latent_dist.sample()``."""
+    in-loop ``latent_dist.sample()``. The cached stats are fp32, so the sample is drawn in
+    fp32 and only the RESULT is cast to weight_dtype (matching the in-loop path's
+    ``encode_latents(...).to(weight_dtype)``)."""
     import torch
 
     parts_a, parts_b = [], []
@@ -694,9 +701,9 @@ def _sample_cached_latents(cache, idxs, variant_rng, device):
         parts_b.append(b)
     lat_a = torch.cat(parts_a).to(device, non_blocking = True)
     if parts_b[0] is None:
-        return lat_a
+        return lat_a.to(dtype = weight_dtype)
     lat_b = torch.cat(parts_b).to(device, non_blocking = True)
-    return lat_a + lat_b * torch.randn_like(lat_a)
+    return (lat_a + lat_b * torch.randn_like(lat_a)).to(dtype = weight_dtype)
 
 
 def _should_compile(cfg, base_is_bnb, device) -> bool:
@@ -962,7 +969,9 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto
         for _ in range(cfg.gradient_accumulation_steps):
             idxs = [rng.randrange(n_images) for _ in range(batch_size)]
             if latent_cache is not None:
-                latents = _sample_cached_latents(latent_cache, idxs, variant_rng, device)
+                latents = _sample_cached_latents(
+                    latent_cache, idxs, variant_rng, device, weight_dtype
+                )
             else:
                 px = torch.stack(
                     [
diff --git a/studio/backend/core/training/diffusion_lora_trainer.py b/studio/backend/core/training/diffusion_lora_trainer.py
index 7c7ca79d51..d8480d5ab1 100644
--- a/studio/backend/core/training/diffusion_lora_trainer.py
+++ b/studio/backend/core/training/diffusion_lora_trainer.py
@@ -177,11 +177,14 @@ def _build_sdxl_latent_cache(
     vae, vae_scale, image_paths, cfg, device, weight_dtype, on_event, check_stop
 ):
     """Precompute the per-image latent posterior cache: for each planned crop/flip variant,
-    encode once and store ``(A, B, time_ids)`` on CPU in the training dtype. ``A`` and ``B``
-    are the affine posterior parameters (mean/std with the VAE scale folded in) so a per-step
-    sample is ``A + B * randn`` -- distribution-identical to an in-loop ``latent_dist.sample()``
-    -- and ``time_ids`` is the SDXL micro-conditioning for the crop. Returns None if the build
-    was interrupted by a stop request. ``vae_scale`` is read before the VAE is freed."""
+    encode once and store ``(A, B, time_ids)`` on CPU in fp32. ``A`` and ``B`` are the affine
+    posterior parameters (mean/std with the VAE scale folded in) so a per-step sample is
+    ``A + B * randn`` -- distribution-identical to an in-loop ``latent_dist.sample()`` -- and
+    ``time_ids`` is the SDXL micro-conditioning for the crop. The stats stay fp32 so the
+    per-step sample happens in fp32 and only the RESULT is cast to weight_dtype, matching the
+    in-loop path (encode fp32 -> sample fp32 -> scale -> .to(weight_dtype)); fp32 doubles the
+    cache RAM over bf16 but the cache is tiny (a handful of latents per image). Returns None if
+    the build was interrupted by a stop request. ``vae_scale`` is read before the VAE is freed."""
     import torch
 
     plan = _plan_cache_variants(
@@ -189,7 +192,7 @@ def _build_sdxl_latent_cache(
     )
 
     def _hold(t):
-        t = t.to(weight_dtype).cpu()
+        t = t.to(torch.float32).cpu()
         if device == "cuda":
             try:
                 t = t.pin_memory()
@@ -222,8 +225,10 @@ def _build_sdxl_latent_cache(
 def _sample_sdxl_cached_latents(cache, idxs, variant_rng, device, weight_dtype):
     """Draw one latent + its time_ids per index from the cache: pick a variant, then sample
     the posterior (A + B * randn) with fresh noise per step, exactly like an in-loop
-    ``latent_dist.sample() * vae_scale``. Returns ``(latents, batch_time_ids)`` already on
-    ``device`` in the training dtype (scale + dtype are folded into the cache)."""
+    ``latent_dist.sample() * vae_scale``. The cached stats are fp32, so the sample is drawn in
+    fp32 and only the RESULT is cast to weight_dtype (matching the in-loop path). Returns
+    ``(latents, batch_time_ids)`` already on ``device`` in the training dtype (scale is folded
+    into the cache)."""
     import torch
 
     parts_a, parts_b, tid_rows = [], [], []
@@ -237,7 +242,7 @@ def _sample_sdxl_cached_latents(cache, idxs, variant_rng, device, weight_dtype):
         tid_rows.append(time_ids)
     lat_a = torch.cat(parts_a).to(device, non_blocking = True)
     lat_b = torch.cat(parts_b).to(device, non_blocking = True)
-    latents = lat_a + lat_b * torch.randn_like(lat_a)
+    latents = (lat_a + lat_b * torch.randn_like(lat_a)).to(dtype = weight_dtype)
     batch_time_ids = torch.tensor(tid_rows, device = device, dtype = weight_dtype)
     return latents, batch_time_ids
 
@@ -443,7 +448,8 @@ def run_diffusion_lora_training(
             for _ in range(cfg.gradient_accumulation_steps):
                 idx, img_paths, captions = _next_batch()
                 if latent_cache is not None:
-                    # Scale + dtype are already folded into the cache, so do not re-apply.
+                    # Scale is folded into the cache; the sampler draws in fp32 and casts the
+                    # result to weight_dtype (matching the in-loop path below).
                     latents, batch_time_ids = _sample_sdxl_cached_latents(
                         latent_cache, idx, variant_rng, device, weight_dtype
                     )
diff --git a/studio/backend/core/training/diffusion_training_service.py b/studio/backend/core/training/diffusion_training_service.py
index fa44cceabe..4fe072105d 100644
--- a/studio/backend/core/training/diffusion_training_service.py
+++ b/studio/backend/core/training/diffusion_training_service.py
@@ -17,6 +17,7 @@ runs a scripted target on a thread.
 
 from __future__ import annotations
 
+import math
 import multiprocessing as mp
 import threading
 import time
@@ -31,6 +32,21 @@ _CTX = mp.get_context("spawn")
 _TERMINAL = ("complete", "error")
 
 
+def _finite_or_none(value: Any) -> Optional[float]:
+    """Coerce a numeric progress field to a finite float, or None. A divergent run (or a
+    grad clip that returns inf) can push loss / grad_norm to NaN or +/-Infinity, and those
+    are invalid in strict JSON -- FastAPI's encoder would emit the JS-only NaN/Infinity
+    tokens that break a strict client parse. Nulling them here (the single service ingestion
+    point both trainers feed) keeps every status snapshot and persisted record JSON-safe."""
+    if value is None:
+        return None
+    try:
+        f = float(value)
+    except (TypeError, ValueError):
+        return None
+    return f if math.isfinite(f) else None
+
+
 def _run_diffusion_child(*, event_queue: Any, stop_queue: Any, config: dict) -> None:
     # Imported lazily so this module (and the route layer) stays torch-free at import.
     from .diffusion_lora_trainer import run_diffusion_training_process
@@ -98,17 +114,12 @@ def _append_metric(state: dict[str, Any], step: Any, loss: Any, lr: Any) -> None
         return
     if istep <= 0 or loss is None:
         return
-    try:
-        floss = float(loss)
-    except (TypeError, ValueError):
+    floss = _finite_or_none(loss)
+    if floss is None:  # non-numeric or non-finite (NaN/Inf): skip, keep the curve JSON-safe
         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
+    # lr may be None (sparse LR series) or non-finite; a non-finite lr is nulled, not dropped,
+    # so a bad lr point never taints the (loss-driven) history.
+    flr = _finite_or_none(lr)
     steps = state["metric_steps"]
     losses = state["metric_loss"]
     lrs = state["metric_lr"]
@@ -308,13 +319,23 @@ class DiffusionTrainingService:
                 # training state, surface the text.
                 s["message"] = str(ev.get("message", "warning"))
             elif etype == "progress":
+                # Null any non-finite float (NaN/Inf from a divergent step or an inf grad
+                # norm) so the JSON status stays strict-parseable; a missing key keeps the
+                # last value, a present-but-non-finite one becomes None.
+                loss = _finite_or_none(ev["loss"]) if "loss" in ev else s["loss"]
+                avg_loss = _finite_or_none(ev["avg_loss"]) if "avg_loss" in ev else s["avg_loss"]
+                learning_rate = (
+                    _finite_or_none(ev["learning_rate"])
+                    if "learning_rate" in ev
+                    else s["learning_rate"]
+                )
                 s.update(
                     status = "running",
                     step = ev.get("step", s["step"]),
                     total_steps = ev.get("total_steps", s["total_steps"]),
-                    loss = ev.get("loss", s["loss"]),
-                    avg_loss = ev.get("avg_loss", s["avg_loss"]),
-                    learning_rate = ev.get("learning_rate", s["learning_rate"]),
+                    loss = loss,
+                    avg_loss = avg_loss,
+                    learning_rate = learning_rate,
                     message = "Training...",
                 )
                 # Fold optional perf fields (emitted by the trainers) so the UI can show
diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py
index f3b7d2b58a..6e2f9dea32 100644
--- a/studio/backend/tests/test_diffusion_training.py
+++ b/studio/backend/tests/test_diffusion_training.py
@@ -189,6 +189,45 @@ def test_apply_event_transitions():
     assert svc.status()["status"] == "error" and svc.status()["message"] == "boom"
 
 
+def test_progress_nulls_non_finite_floats_for_strict_json():
+    # A divergent step (or an inf grad norm) can push loss / avg_loss / learning_rate to
+    # NaN or Infinity, which strict JSON forbids. The service must null those so the status
+    # snapshot and the metric history stay strict-JSON serializable.
+    import json
+    import math
+
+    svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target)
+    svc._apply_event(
+        {
+            "type": "progress",
+            "step": 3,
+            "total_steps": 10,
+            "loss": float("nan"),
+            "avg_loss": float("inf"),
+            "learning_rate": float("-inf"),
+        }
+    )
+    snap = svc.status()
+    assert snap["loss"] is None
+    assert snap["avg_loss"] is None
+    assert snap["learning_rate"] is None
+    # The non-finite point is skipped in the history, so the loss series stays clean.
+    assert snap["metric_loss"] == []
+    assert snap["metric_steps"] == []
+    # strict JSON (allow_nan=False) round-trips without a ValueError from NaN/Infinity.
+    json.dumps(snap, allow_nan = False)
+
+    # A finite point after the bad one is recorded and preserved verbatim.
+    svc._apply_event(
+        {"type": "progress", "step": 4, "total_steps": 10, "loss": 0.5, "learning_rate": 1e-4}
+    )
+    snap2 = svc.status()
+    assert snap2["loss"] == 0.5
+    assert snap2["metric_loss"] == [0.5] and snap2["metric_steps"] == [4]
+    assert math.isfinite(snap2["learning_rate"])
+    json.dumps(snap2, allow_nan = False)
+
+
 def test_terminal_events_clear_model_load_flag():
     # A stop or error during model load emits complete/error WITHOUT a preceding
     # model_load_completed, so the terminal update must reset in_model_load or the

From 8190f649c1b3d915740ceb6f0a4f8c682f66ce50 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sat, 4 Jul 2026 01:27:36 +0000
Subject: [PATCH 4/4] Address review: gate auto int8 on torchao, scope dense
 validation to DiT

- base_precision="auto" only picks int8 when torchao is importable (the int8
  quantize has no runtime fallback, unlike fp8); otherwise the middle band falls
  back to nf4. Threaded as a parameter so the policy stays pure.
- The dense-mode validation (prequant base / bf16 compute) now applies only to
  DiT families: sdxl ignores base_precision entirely, so a leftover value can no
  longer fail an SDXL run. The mode-name validity check still runs everywhere.
---
 .../core/training/diffusion_dit_trainer.py    | 23 ++++--
 .../core/training/diffusion_train_common.py   |  6 +-
 .../tests/test_diffusion_base_precision.py    | 74 +++++++++++++++++++
 3 files changed, 96 insertions(+), 7 deletions(-)

diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py
index 5221425e4b..ce815f8268 100644
--- a/studio/backend/core/training/diffusion_dit_trainer.py
+++ b/studio/backend/core/training/diffusion_dit_trainer.py
@@ -278,7 +278,9 @@ def _apply_fp8_training(transformer, on_event) -> bool:
         return False
 
 
-def _pick_auto_precision(prequant, device, free_gb, dense_gb, capability, has_fp8) -> str:
+def _pick_auto_precision(
+    prequant, device, free_gb, dense_gb, capability, has_fp8, has_torchao = True
+) -> str:
     """Pure policy for base_precision="auto": nf4 for a prequant base or no CUDA; else the
     fastest dense mode whose weights + headroom (activations, optimizer, cache) fit the
     free VRAM at decision time. bf16 + regional compile is the measured speed winner
@@ -287,16 +289,18 @@ def _pick_auto_precision(prequant, device, free_gb, dense_gb, capability, has_fp
     the same hardware. int8 must still materialise the full bf16 transformer before
     ``quantize_`` shrinks it module-by-module, so its band requires the dense-load
     transient (1.15x dense) to fit -- what int8 buys in that band is steady-state
-    headroom for activations and the latent cache, not load-time memory.
-    ``capability``/``has_fp8`` remain parameters so the policy can be revisited per GPU
-    generation without changing callers."""
+    headroom for activations and the latent cache, not load-time memory. int8 also needs
+    torchao at runtime (``_int8_quantize_base`` has no fallback, unlike fp8), so auto only
+    picks it when torchao is importable and drops to nf4 otherwise. ``capability``/``has_fp8``
+    remain parameters so the policy can be revisited per GPU generation without changing
+    callers."""
     _ = capability, has_fp8
     if prequant or device != "cuda" or not free_gb or not dense_gb:
         return "nf4"
     if free_gb > dense_gb * 1.5:
         return "bf16"
     if free_gb > dense_gb * 1.15:
-        return "int8"
+        return "int8" if has_torchao else "nf4"
     return "nf4"
 
 
@@ -323,6 +327,11 @@ def _resolve_base_precision(cfg, spec, device) -> str:
     free_gb = None
     capability = None
     has_fp8 = False
+    # int8 quantization has no runtime fallback, so gate the auto pick on torchao being
+    # importable (find_spec avoids the cost/side-effects of an actual import).
+    import importlib.util
+
+    has_torchao = importlib.util.find_spec("torchao") is not None
     if device == "cuda":
         try:
             import torch
@@ -332,7 +341,9 @@ def _resolve_base_precision(cfg, spec, device) -> str:
             has_fp8 = hasattr(torch, "float8_e4m3fn")
         except Exception:  # noqa: BLE001 -- probe failure -> the safe mode
             pass
-    return _pick_auto_precision(prequant, device, free_gb, spec.dense_bf16_gb, capability, has_fp8)
+    return _pick_auto_precision(
+        prequant, device, free_gb, spec.dense_bf16_gb, capability, has_fp8, has_torchao
+    )
 
 
 # ── FLUX.1-dev ────────────────────────────────────────────────────────────────
diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py
index 56b1020cd5..8d168203c4 100644
--- a/studio/backend/core/training/diffusion_train_common.py
+++ b/studio/backend/core/training/diffusion_train_common.py
@@ -320,7 +320,11 @@ class DiffusionLoraConfig:
         base_precision = str(self.base_precision or "nf4").strip().lower()
         if base_precision not in ("nf4", "bf16", "int8", "fp8", "auto"):
             raise ValueError("base_precision must be one of nf4 / bf16 / int8 / fp8 / auto")
-        if base_precision in ("bf16", "int8", "fp8"):
+        # base_precision is a DiT-only lever (nf4/bf16/int8/fp8/auto for the transformer
+        # load); SDXL uses its own mixed_precision path and ignores base_precision entirely,
+        # so the dense-mode gates (prequant base / non-bf16 compute) apply only to the DiT
+        # families. The mode-name validity check above still runs for every family.
+        if resolved_family != "sdxl" and base_precision in ("bf16", "int8", "fp8"):
             if repo_is_prequantized(self.base_model):
                 raise ValueError(
                     f"base_precision={base_precision!r} needs a dense base repo, but "
diff --git a/studio/backend/tests/test_diffusion_base_precision.py b/studio/backend/tests/test_diffusion_base_precision.py
index af6789b199..1f8ca9318f 100644
--- a/studio/backend/tests/test_diffusion_base_precision.py
+++ b/studio/backend/tests/test_diffusion_base_precision.py
@@ -32,6 +32,9 @@ from models.training import DiffusionTrainingStartRequest
 # family from their names alone, so normalized() runs without a network call.
 _FLUX_DENSE = "black-forest-labs/FLUX.1-dev"
 _Z_PREQUANT = "unsloth/Z-Image-Turbo-unsloth-bnb-4bit"
+# An SDXL base whose name LOOKS prequant (bnb-4bit): SDXL ignores base_precision, so the
+# dense-mode gates must not fire for it even with a dense mode + fp16 compute.
+_SDXL_PREQUANT_NAME = "some/sdxl-model-bnb-4bit"
 
 
 def _cfg(base_model = _FLUX_DENSE, **kw) -> DiffusionLoraConfig:
@@ -66,6 +69,32 @@ def test_base_precision_validation():
     assert _cfg(base_model = _Z_PREQUANT, base_precision = "auto").normalized().base_precision == "auto"
 
 
+def test_base_precision_gates_skip_sdxl():
+    # SDXL ignores base_precision, so the dense-mode gates (prequant base / non-bf16 compute)
+    # must not fire for it: a prequant-looking SDXL name with base_precision="bf16" does not
+    # raise, and the mode is still stored lowered.
+    norm = _cfg(base_model = _SDXL_PREQUANT_NAME, base_precision = "bf16").normalized()
+    assert norm.resolved_family == "sdxl"
+    assert norm.base_precision == "bf16"
+
+    # The non-bf16-compute gate is also skipped for SDXL (fp16 is a valid SDXL mixed
+    # precision), even with a dense base_precision requested.
+    norm2 = _cfg(
+        base_model = "stabilityai/stable-diffusion-xl-base-1.0",
+        base_precision = "int8",
+        mixed_precision = "fp16",
+    ).normalized()
+    assert norm2.resolved_family == "sdxl"
+
+    # The mode-name validity check still runs for SDXL: an unknown mode is rejected.
+    with pytest.raises(ValueError, match = "base_precision"):
+        _cfg(base_model = _SDXL_PREQUANT_NAME, base_precision = "banana").normalized()
+
+    # The gates STILL fire for a DiT family: a prequant DiT base with a dense mode raises.
+    with pytest.raises(ValueError, match = "dense base repo"):
+        _cfg(base_model = _Z_PREQUANT, base_precision = "bf16").normalized()
+
+
 # ── repo_is_prequantized heuristic + trainer alias ────────────────────────────
 @pytest.mark.parametrize(
     "repo, expected",
@@ -107,6 +136,10 @@ def test_pick_auto_precision_policy_table():
 
     # Middle band (30 > 23.8 * 1.15 = 27.4, but not > 23.8 * 1.5 = 35.7) -> int8.
     assert p(False, "cuda", 30, 23.8, (10, 0), True) == "int8"
+    # int8 needs torchao at runtime (no fallback), so the int8 band drops to nf4 when
+    # torchao is not importable while the bf16 band is unaffected.
+    assert p(False, "cuda", 30, 23.8, (10, 0), True, False) == "nf4"
+    assert p(False, "cuda", 140, 23.8, (10, 0), True, False) == "bf16"
     # int8 still materialises the full bf16 transformer before quantize_ shrinks it, so
     # free VRAM below the dense-load transient (25 < 27.4) must fall back to nf4 even
     # though the QUANTIZED weights would have fit.
@@ -140,6 +173,47 @@ def test_resolve_auto_requires_bf16_compute():
     assert dit._resolve_base_precision(cfg, spec, "cuda") == "nf4"
 
 
+def test_resolve_auto_int8_band_gates_on_torchao(monkeypatch):
+    # The int8 auto band needs torchao at runtime; when torchao is not importable
+    # _resolve_base_precision must fall to nf4 instead of picking an int8 that would crash
+    # in _int8_quantize_base. Drive the probe into the int8 band and toggle torchao.
+    import importlib.util as _ilu
+    import torch
+
+    spec = dit._SPECS["flux.1"]  # dense_bf16_gb = 23.8
+    cfg = _cfg(base_precision = "auto", mixed_precision = "bf16")
+    real_find_spec = _ilu.find_spec
+
+    def _no_torchao(name, *args, **kwargs):
+        if name == "torchao":
+            return None  # simulate torchao not installed
+        return real_find_spec(name, *args, **kwargs)
+
+    def _has_torchao(name, *args, **kwargs):
+        if name == "torchao":
+            return object()  # simulate torchao installed
+        return real_find_spec(name, *args, **kwargs)
+
+    class _FakeCuda:
+        # Free VRAM in the int8 band (30 > 23.8 * 1.15) but below the bf16 band.
+        @staticmethod
+        def mem_get_info():
+            return (int(30 * 1e9), int(80 * 1e9))
+
+        @staticmethod
+        def get_device_capability():
+            return (10, 0)
+
+    monkeypatch.setattr(torch, "cuda", _FakeCuda)
+
+    monkeypatch.setattr(_ilu, "find_spec", _no_torchao)
+    assert dit._resolve_base_precision(cfg, spec, "cuda") == "nf4"
+
+    # With torchao importable the same band picks int8.
+    monkeypatch.setattr(_ilu, "find_spec", _has_torchao)
+    assert dit._resolve_base_precision(cfg, spec, "cuda") == "int8"
+
+
 # ── _fp8_module_filter ────────────────────────────────────────────────────────
 def test_fp8_module_filter():
     lin = nn.Linear(64, 64)