diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index d162559fc9..e71d7aa256 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/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). diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index f6a9ea2766..567cc39c26 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/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"}, diff --git a/studio/backend/tests/test_diffusion_train_perf.py b/studio/backend/tests/test_diffusion_train_perf.py index d430601abd..1477e24874 100644 --- a/studio/backend/tests/test_diffusion_train_perf.py +++ b/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)