Merge diffusion-train-perf: CI test fixes (diffusers import order, arbiter device pin, sigma-gather skip)

studio/backend/core/training/diffusion_dit_trainer.py

@@ -946,13 +946,23 @@ 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."""
-    import torch
-
     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
+
     rng = random.Random(cfg.seed)
     torch.manual_seed(cfg.seed)
     _FLUX_STATIC.clear()
@@ -970,15 +980,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).

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_train_perf.py

@@ -130,6 +130,9 @@ def test_zimage_collate_list():
 
 # ── index-based sigma gather ──────────────────────────────────────────────────
 def test_gather_sigmas_matches_search_based_gather():
+    # CI installs the backend test deps without diffusers; the scheduler math is what we
+    # are checking, so skip rather than fail there.
+    pytest.importorskip("diffusers")
     from diffusers import FlowMatchEulerDiscreteScheduler
 
     torch.manual_seed(0)