Serialize the GPU handoffs, gate DiT training on a GPU, and keep 3.9 installable

Six review findings, three of them evict-then-fail orderings:

- The chat load reclaimed the GPU without telling the arbiter it existed. A
  chat load holds no llama-server process until its GGUF has downloaded,
  which is minutes, so a competing Images/Video acquire in that window
  found nothing to cancel, took the GPU, and the chat load then spawned
  onto the same device. It now registers an in-flight marker through
  acquire_for's register hook (under the arbiter lock, as the image and
  video loads do), the evictor cancels a marked load, and the route undoes
  itself if ownership moved while it loaded.
- The Hub-download conflict check ran after that handoff, so a GGUF the
  download manager already owns destroyed the resident Images/Video
  pipeline and then 409'd, having loaded nothing. It moves above the
  handoff, together with the marker it handshakes with.
- The image load released the engine router's transition lock before
  registering the load, so a second load choosing the other engine could
  unload the still-idle engine this one captured; the load then landed on a
  deactivated engine, where generate, status, unload and the arbiter's
  evictor can no longer reach it. Registration now happens under that lock
  and refuses if the engine changed.
- Training a DiT family on a host with no GPU was accepted: nf4 is not a
  CPU fallback, its 4-bit load goes through bitsandbytes, which requires
  CUDA, XPU or MPS. The start unloaded the working Images pipeline, pulled
  the text encoders, and only then died in the child. Rejected before the
  teardown now, and /info stops advertising a precision that always 400s.
  SDXL keeps its documented fp32-on-CPU path.
- Both diffusion pages kept the routed-pick marker forever, so re-picking
  the same checkpoint (after chat evicted it) neither loaded nor cleared
  the query string. The marker is released once the query is gone. The
  Images key also carried a stray NUL byte, which made the file read as
  binary to grep and other tooling.
- diffusers dropped Python 3.9 in 0.38, so the unconditional >=0.39.0 pin
  left pip no candidate at all on 3.9 and made every install that composes
  the huggingface extras unresolvable there. The floor is conditional now.

Also fixes tests that were already red on the branch: two hand-built
request fakes had gone stale against fields this branch added, and the
handoff-ordering test only failed on a host with fewer than two GPUs.

pyproject.toml

@@ -92,7 +92,11 @@ huggingfacenotorch = [
     # workflow raises without it), the cache_context child-registry behavior, and the
     # Flux2 / Z-Image pipelines. An unversioned entry let an upgrade keep an older
     # diffusers, so advertised models failed to load until the user upgraded by hand.
-    "diffusers>=0.39.0",
+    # diffusers dropped 3.9 in 0.38, and we still support it (requires-python >= 3.9),
+    # so the floor is conditional: pinning it outright leaves pip no candidate at all
+    # on 3.9 and the whole extra becomes unresolvable.
+    "diffusers>=0.39.0 ; python_version >= '3.10'",
+    "diffusers ; python_version < '3.10'",
     "transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.5.0",
     "trl>=0.18.2,!=0.19.0,<=0.24.0",
     "sentence-transformers",

studio/backend/core/inference/diffusion_engine_router.py

@@ -22,7 +22,7 @@ from __future__ import annotations
 
 import os
 import threading
-from typing import Any, Optional
+from typing import Any, Callable, Optional
 
 from core.inference.diffusion_device import resolve_diffusion_device_target
 from core.inference.diffusion_families import DiffusionFamily, family_sd_cpp_supported
@@ -124,6 +124,24 @@ def _activate(name: str, reason: Optional[str]) -> Any:
         return get_active_diffusion_engine()
 
 
+def begin_load_on(expected_engine: Any, start: Callable[[], Any]) -> Any:
+    """Run ``start`` under the transition lock, refusing if the engine changed since selection.
+
+    A load route selects its engine, then yields (device probe, arbiter acquire) before it
+    registers the load. A second /images/load picking the OTHER engine can transition in that
+    gap and unload the still-idle engine this request captured, which would then load a model
+    nothing can reach: generate / status / unload and the arbiter's evictor all resolve through
+    get_active_diffusion_engine(). Re-checking under the same lock the switch takes makes
+    selection and registration one operation.
+    """
+    with _transition_lock:
+        if expected_engine is not get_active_diffusion_engine():
+            raise RuntimeError(
+                "The diffusion engine changed while this load was starting. Retry the load."
+            )
+        return start()
+
+
 def select_and_activate_engine(
     fam: DiffusionFamily,
     *,

studio/backend/core/inference/gpu_arbiter.py

@@ -32,10 +32,15 @@ def _evict_chat() -> None:
     from core.inference import get_inference_backend
     from routes.inference import get_llama_cpp_backend
 
+    from core.inference.llama_cpp import chat_load_active
+
     llama = get_llama_cpp_backend()
     # is_active (process exists), not is_loaded (exists AND healthy): a chat model still starting
     # up holds VRAM but isn't healthy, so is_loaded would skip it and let the load race diffusion.
-    if llama.is_active:
+    # chat_load_active too: an HF load has no process until its GGUF downloaded, so is_active
+    # alone found nothing to cancel and let that load spawn onto the GPU we just granted away.
+    # unload_model sets the cancel event the download loop polls, so the pending load aborts.
+    if llama.is_active or chat_load_active():
         llama.unload_model()
     orchestrator = get_inference_backend()
     if orchestrator.active_model_name:

studio/backend/core/inference/llama_cpp.py

@@ -1382,6 +1382,36 @@ def gguf_load_in_flight(hf_repo: Optional[str]):
                 _LOADS_IN_FLIGHT[key] = remaining
 
 
+# Chat loads in flight, repo-agnostic (local paths and safetensors included). The repo-keyed
+# counter above answers the download manager; this one answers the GPU arbiter, which has to know
+# a chat load exists before llama-server is spawned.
+_CHAT_LOADS_IN_FLIGHT = 0
+
+
+@contextlib.contextmanager
+def chat_load_in_flight():
+    """Track a chat load from before its download until the load call returns."""
+    global _CHAT_LOADS_IN_FLIGHT
+    with _LOADS_IN_FLIGHT_LOCK:
+        _CHAT_LOADS_IN_FLIGHT += 1
+    try:
+        yield
+    finally:
+        with _LOADS_IN_FLIGHT_LOCK:
+            _CHAT_LOADS_IN_FLIGHT = max(0, _CHAT_LOADS_IN_FLIGHT - 1)
+
+
+def chat_load_active() -> bool:
+    """Whether a chat load is in flight, spawned or not.
+
+    ``is_active`` only covers a live llama-server process, which an HF load does not have
+    until its GGUF finished downloading -- minutes during which the arbiter's evictor would
+    find nothing to cancel and grant the GPU to a second workload.
+    """
+    with _LOADS_IN_FLIGHT_LOCK:
+        return _CHAT_LOADS_IN_FLIGHT > 0
+
+
 def hf_gguf_load_in_flight(hf_repo: str) -> bool:
     """Return whether a GGUF load is active for hf_repo."""
     key = (hf_repo or "").strip().lower()

studio/backend/core/training/diffusion_train_common.py

@@ -356,15 +356,49 @@ def bf16_unsupported_reason(resolved_family: str) -> Optional[str]:
     return None
 
 
+def dit_accelerator_missing_reason(resolved_family: str) -> Optional[str]:
+    """Reason a DiT family cannot train on this host at all, else None. Never raises.
+
+    nf4 is not a CPU fallback: the 4-bit base load goes through diffusers' bitsandbytes
+    quantizer, whose validate_environment raises "No GPU found. A GPU is needed for
+    quantization." unless CUDA, XPU or MPS is present. Without this gate a GPU-less host
+    accepts the default nf4 start, evicts the resident Images pipeline, downloads the text
+    encoders, and only then dies in the child. SDXL keeps its own fp32-on-CPU path.
+    """
+    if (resolved_family or "").strip().lower() not in _DIT_TRAIN_FAMILIES:
+        return None
+    try:
+        import torch
+        xpu = getattr(torch, "xpu", None)
+        mps = getattr(torch, "mps", None)
+        if (
+            torch.cuda.is_available()
+            or bool(xpu is not None and xpu.is_available())
+            or bool(mps is not None and mps.is_available())
+        ):
+            return None
+    except Exception:  # noqa: BLE001 -- probe failure must not block a start
+        return None
+    return (
+        "Training the DiT families needs a GPU: even the 4-bit (nf4) base load requires "
+        "CUDA, XPU or MPS, and this host has none. Train SDXL here, or use a GPU machine."
+    )
+
+
 def training_precision_preflight_error(resolved_family: str, base_precision: str) -> Optional[str]:
     """Reason the requested DiT precision cannot run on this host, else None -- checked by the
     start route BEFORE evicting resident GPU workloads (the trainer's own checks fire only in the
     child, after eviction). Four gates, all mirroring _resolve_base_precision so a doomed run is
-    rejected before teardown: the bf16-GPU requirement (bf16_unsupported_reason); the dense
+    rejected before teardown: the bf16-GPU requirement (bf16_unsupported_reason); a host with no
+    accelerator at all (dit_accelerator_missing_reason, which covers nf4 too); the dense
     precisions (bf16/int8/fp8/mxfp8) requiring a CUDA GPU; an explicit int8 needing a FUNCTIONAL
     torchao (its _int8_quantize_base has no fallback); and an explicit mxfp8 needing a Blackwell
     (sm100+) GPU (its MX GEMM has no kernel below sm100). Never raises."""
     reason = bf16_unsupported_reason(resolved_family)
+    if reason:
+        return reason
+    # No accelerator at all: every DiT precision is out, nf4 included.
+    reason = dit_accelerator_missing_reason(resolved_family)
     if reason:
         return reason
     fam = (resolved_family or "").strip().lower()
@@ -430,7 +464,11 @@ def family_train_infos() -> list[dict[str, Any]]:
         # DiT option that always 400s. Otherwise drop any scheme this family's DiT corrupts (fp8 on
         # Qwen-Image: activation outliers exceed fp8's range; the inference path denies the same
         # set), so the UI never offers a mode normalized() would reject.
-        dit_block = bf16_unsupported_reason(name) if is_dit else None
+        dit_block = (
+            bf16_unsupported_reason(name) or dit_accelerator_missing_reason(name)
+            if is_dit
+            else None
+        )
         if not is_dit or dit_block:
             fam_modes: list[str] = []
         else:

studio/backend/routes/inference.py

@@ -4395,7 +4395,7 @@ async def _load_model_impl(
         # validation so a doomed load (bad id, unsupported gpu_ids on GGUF, training 409) can't
         # evict a working image/video model and then error. Mirrors the image/video loaders,
         # which validate before acquire_for.
-        from core.inference.gpu_arbiter import acquire_for, CHAT
+        from core.inference.gpu_arbiter import acquire_for, current_owner, CHAT
 
         # ── Already-loaded check: skip reload if the exact model is active ──
         backend = get_inference_backend()
@@ -4606,44 +4606,56 @@ async def _load_model_impl(
             gpu_memory_mode = request.gpu_memory_mode,
         )
 
+        # Mark the load and refuse one the download manager already owns, BEFORE the eviction
+        # below: this 409 leaves nothing loaded, so checking it afterwards destroyed a working
+        # Images/Video pipeline for a load that could never start. The marker/check order is the
+        # handshake with the download manager, so both move together. It also has to run after
+        # pass-through argument inheritance, since a carried --no-mmproj changes the companion
+        # requirement exactly as it does for the load.
+        if config.is_gguf and config.gguf_hf_repo:
+            from core.inference.llama_cpp import gguf_load_in_flight
+            gguf_load_stack.enter_context(gguf_load_in_flight(config.gguf_hf_repo))
+
+            from core.inference.llama_cpp import _hub_download_blocks_gguf_load
+            if await asyncio.to_thread(
+                _hub_download_blocks_gguf_load,
+                config.gguf_hf_repo,
+                config.gguf_variant,
+                require_mmproj = bool(
+                    config.is_vision and not extra_args_disable_mmproj(extra_llama_args)
+                ),
+                hf_token = request.hf_token,
+            ):
+                raise HTTPException(
+                    status_code = 409,
+                    detail = (
+                        f"'{model_log_label}' is currently being downloaded "
+                        "by the download manager. Wait for the download to "
+                        "finish (or cancel it), then load the model."
+                    ),
+                )
+
         # Load now known viable (valid identifier, gpu_ids ok, fits alongside any active
         # training): reclaim the GPU for chat, evicting a resident Images/Video pipeline. Doing
         # this only here -- not before the validation above -- keeps a doomed load from evicting
         # a working image/video model and then erroring. No-op when chat already owns the GPU.
-        await asyncio.to_thread(acquire_for, CHAT)
+        # The in-flight marker is entered UNDER the arbiter lock (the `register` hook), as the
+        # image and video loads do: a chat load holds no llama-server process until its GGUF has
+        # downloaded, so a competing Images/Video acquire in that window found nothing to evict
+        # and both then allocated VRAM at once. With the marker, that evictor cancels this load.
+        from core.inference.llama_cpp import chat_load_in_flight
+
+        await asyncio.to_thread(
+            acquire_for,
+            CHAT,
+            lambda: gguf_load_stack.enter_context(chat_load_in_flight()),
+        )
 
         # ── GGUF path: load via llama-server ──────────────────────
         if config.is_gguf:
             llama_backend = get_llama_cpp_backend()
             unsloth_backend = get_inference_backend()
 
-            if config.gguf_hf_repo:
-                from core.inference.llama_cpp import gguf_load_in_flight
-                gguf_load_stack.enter_context(gguf_load_in_flight(config.gguf_hf_repo))
-
-            # Block cache writes that would race the download manager. This runs
-            # after pass-through argument inheritance so a carried --no-mmproj
-            # changes the companion requirement exactly as it does for the load.
-            if config.gguf_hf_repo:
-                from core.inference.llama_cpp import _hub_download_blocks_gguf_load
-                if await asyncio.to_thread(
-                    _hub_download_blocks_gguf_load,
-                    config.gguf_hf_repo,
-                    config.gguf_variant,
-                    require_mmproj = bool(
-                        config.is_vision and not extra_args_disable_mmproj(extra_llama_args)
-                    ),
-                    hf_token = request.hf_token,
-                ):
-                    raise HTTPException(
-                        status_code = 409,
-                        detail = (
-                            f"'{model_log_label}' is currently being downloaded "
-                            "by the download manager. Wait for the download to "
-                            "finish (or cancel it), then load the model."
-                        ),
-                    )
-
             # Keep the resident model alive until every active generation finishes;
             # the caller's lifecycle gate blocks new starts.
             await _wait_for_model_switch_idle(current_request_counted = current_request_counted)
@@ -4799,6 +4811,20 @@ async def _load_model_impl(
                     detail = f"Failed to load GGUF model: {model_log_label if native_grant_backed else config.display_name}",
                 )
 
+            # An Images/Video acquire can land in the gap between the acquire above and
+            # load_model clearing the cancel event, so its cancellation is lost and this load
+            # spawns anyway. Ownership survives that gap: whoever took the GPU keeps it, and this
+            # load undoes itself rather than leaving two models resident on one device.
+            if current_owner() != CHAT:
+                await asyncio.to_thread(llama_backend.unload_model)
+                raise HTTPException(
+                    status_code = 409,
+                    detail = (
+                        "An image or video model took the GPU while this model was loading, "
+                        "so the load was cancelled. Unload that model, then try again."
+                    ),
+                )
+
             logger.info(
                 f"Loaded GGUF model via llama-server: {model_log_label if native_grant_backed else config.identifier}"
             )
@@ -16028,6 +16054,7 @@ async def load_diffusion_model(
     from core.inference.diffusion_device import resolve_diffusion_device_target
     from core.inference.diffusion_engine_router import (
         annotate_status,
+        begin_load_on,
         select_and_activate_engine,
     )
     from core.inference.gpu_arbiter import acquire_for, release, DIFFUSION
@@ -16077,7 +16104,7 @@ async def load_diffusion_model(
         device = await asyncio.to_thread(lambda: resolve_diffusion_device_target().device)
         needs_gpu = device != "cpu"
 
-        def _begin_load():
+        def _start_engine_load():
             # Kicks the (slow) load onto a background thread and returns at once (client polls
             # images/load-progress); begin_load itself validates network-free.
             return engine.begin_load(
@@ -16100,6 +16127,12 @@ async def load_diffusion_model(
                 loras = [(s.id, s.weight) for s in request.loras] if request.loras else None,
             )
 
+        def _begin_load():
+            # Under the router's transition lock, refusing if a competing load switched engines
+            # since select_and_activate_engine above: begin_load on a deactivated engine leaves a
+            # resident model that generate / status / unload and the evictor can no longer reach.
+            return begin_load_on(engine, _start_engine_load)
+
         if needs_gpu:
             # Register the in-flight load UNDER the arbiter lock (not after acquire_for returns):
             # otherwise a competing Video/chat acquire in that gap evicts DIFFUSION before the load

studio/backend/tests/test_diffusion_base_precision.py

@@ -219,17 +219,28 @@ def test_training_precision_preflight_error(monkeypatch):
     assert training_precision_preflight_error("sdxl", "int8") is None
     assert training_precision_preflight_error("", "int8") is None
 
-    # On a CUDA-ABSENT host, bf16_unsupported_reason exempts CPU-only, but the DiT trainer's dense
-    # precisions still require CUDA (mirroring _resolve_base_precision), so bf16/int8/fp8/mxfp8 for
-    # a DiT family are rejected UP FRONT rather than after eviction. nf4/auto (and SDXL) still pass.
+    # On a host with NO accelerator every DiT precision is rejected up front rather than after
+    # eviction -- nf4 included, since its 4-bit load goes through bitsandbytes, which needs a GPU.
+    # SDXL (its own fp32-on-CPU path) still passes.
     monkeypatch.setattr(common, "has_functional_torchao", lambda: True)
     monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
+    monkeypatch.setattr(torch.xpu, "is_available", lambda: False)
+    monkeypatch.setattr(torch.mps, "is_available", lambda: False)
     for dense in ("bf16", "int8", "fp8", "mxfp8"):
         reason = training_precision_preflight_error("flux.1", dense)
         assert reason is not None and "CUDA" in reason
+    for mode in ("nf4", "auto"):
+        reason = training_precision_preflight_error("flux.1", mode)
+        assert reason is not None and "GPU" in reason
+    assert training_precision_preflight_error("sdxl", "bf16") is None
+
+    # An accelerator that is not CUDA (XPU here) satisfies the 4-bit load, so nf4/auto pass again
+    # while the dense CUDA-only precisions stay rejected.
+    monkeypatch.setattr(torch.xpu, "is_available", lambda: True)
     assert training_precision_preflight_error("flux.1", "nf4") is None
     assert training_precision_preflight_error("flux.1", "auto") is None
-    assert training_precision_preflight_error("sdxl", "bf16") is None
+    assert training_precision_preflight_error("flux.1", "bf16") is not None
+    monkeypatch.setattr(torch.xpu, "is_available", lambda: False)
 
     # mxfp8 needs a Blackwell (sm100+) GPU: its MX GEMM has no kernel below sm100 and would raise at
     # the first training step, AFTER a full dense-transformer load. On a CUDA GPU that is older than
@@ -609,3 +620,36 @@ def test_assert_trusted_base_model_rejects_local_non_pipeline(tmp_path):
     # An untrusted remote base is still rejected by the trust gate.
     with pytest.raises(ValueError, match = "untrusted"):
         common._assert_trusted_base_model("evil/base")
+
+def test_dit_accelerator_missing_reason_and_info_hide_train_without_a_gpu(monkeypatch):
+    # Clicking Start on a GPU-less host evicted the resident Images pipeline, downloaded the text
+    # encoders, and only then died in the child: diffusers' bitsandbytes quantizer refuses 4-bit
+    # without CUDA/XPU/MPS. Reject it up front, and stop /info advertising a mode that always 400s.
+    import torch
+
+    from core.training.diffusion_train_common import (
+        _DIT_TRAIN_FAMILIES,
+        dit_accelerator_missing_reason,
+        family_train_infos,
+    )
+
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
+    monkeypatch.setattr(torch.xpu, "is_available", lambda: False)
+    monkeypatch.setattr(torch.mps, "is_available", lambda: False)
+    assert "GPU" in (dit_accelerator_missing_reason("flux.1") or "")
+    # SDXL and unknown families keep their own paths.
+    assert dit_accelerator_missing_reason("sdxl") is None
+    assert dit_accelerator_missing_reason("") is None
+
+    infos = {info["name"]: info for info in family_train_infos()}
+    for name, info in infos.items():
+        if name in _DIT_TRAIN_FAMILIES:
+            assert info["precision_modes"] == []
+            assert "GPU" in info["vram_note"]
+            assert info["supports_compile"] is False
+        else:
+            assert info["precision_modes"] != [] or name not in _DIT_TRAIN_FAMILIES
+
+    # Any accelerator clears it (MPS here, which bitsandbytes accepts).
+    monkeypatch.setattr(torch.mps, "is_available", lambda: True)
+    assert dit_accelerator_missing_reason("flux.1") is None

studio/backend/tests/test_diffusion_engine_router.py

@@ -283,3 +283,48 @@ def test_activate_serializes_switch_and_concurrent_query(monkeypatch):
     t.join(2.0)
     q.join(2.0)
     assert switch_done.is_set() and query_done.is_set()
+
+
+def test_begin_load_on_refuses_an_engine_that_was_switched_away(monkeypatch):
+    # A load selects its engine, then yields (device probe, arbiter acquire) before registering.
+    # A second load choosing the OTHER engine transitions in that gap and unloads the still-idle
+    # engine this one captured, so registering there would leave a model that generate / status /
+    # unload and the arbiter's evictor can no longer reach (they all resolve the ACTIVE engine).
+    diffusers = SimpleNamespace(name = "diffusers")
+    sd_cpp = SimpleNamespace(name = "sd_cpp")
+    active = {"engine": diffusers}
+    monkeypatch.setattr(r, "get_active_diffusion_engine", lambda: active["engine"])
+
+    started: list[str] = []
+    assert r.begin_load_on(diffusers, lambda: started.append("ok") or "status") == "status"
+    assert started == ["ok"]
+
+    # A competing request switched the active engine after this one captured `diffusers`.
+    active["engine"] = sd_cpp
+    with pytest.raises(RuntimeError, match = "engine changed"):
+        r.begin_load_on(diffusers, lambda: started.append("leaked"))
+    assert started == ["ok"]
+
+
+def test_begin_load_on_holds_the_transition_lock_while_registering(monkeypatch):
+    # The check and the registration must be one operation: taken under the same lock a switch
+    # takes, so no _activate can slip between them.
+    import threading
+
+    engine = SimpleNamespace(name = "diffusers")
+    monkeypatch.setattr(r, "get_active_diffusion_engine", lambda: engine)
+
+    inside = threading.Event()
+    release = threading.Event()
+
+    def _slow_start():
+        inside.set()
+        release.wait(2.0)
+        return "status"
+
+    t = threading.Thread(target = lambda: r.begin_load_on(engine, _slow_start))
+    t.start()
+    assert inside.wait(2.0)
+    assert r._transition_lock.locked()
+    release.set()
+    t.join(2.0)

studio/backend/tests/test_diffusion_training.py

@@ -1583,6 +1583,34 @@ def test_start_gated_base_without_access_is_400_and_keeps_gpu(client, monkeypatc
     assert client._fake.started_with is None
 
 
+def test_route_start_refuses_a_dit_family_without_a_gpu_before_freeing(client, monkeypatch):
+    # nf4 is not a CPU fallback: the 4-bit base load needs CUDA/XPU/MPS. Without a pre-teardown
+    # gate the default Train pick on a GPU-less host unloaded the working Images pipeline, pulled
+    # the text encoders, and only then failed in the child.
+    import torch
+
+    import routes.training as tr
+
+    freed = []
+    monkeypatch.setattr(tr, "_free_gpu_for_diffusion_training", lambda: freed.append(1))
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
+    monkeypatch.setattr(torch.xpu, "is_available", lambda: False)
+    monkeypatch.setattr(torch.mps, "is_available", lambda: False)
+
+    r = client.post(
+        "/api/train/diffusion/start",
+        json = {**_BODY, "base_model": "black-forest-labs/FLUX.1-dev"},
+    )
+    assert r.status_code == 400
+    assert "GPU" in r.json()["detail"]
+    assert freed == []
+    assert client._fake.started_with is None
+
+    # SDXL trains fp32 on CPU (that is its documented fallback), so it is not gated.
+    r2 = client.post("/api/train/diffusion/start", json = _BODY)
+    assert r2.status_code == 200, r2.text
+
+
 def test_start_ungated_base_preflight_is_noop(client, monkeypatch):
     # A reachable base (HEAD 200) proceeds to start normally.
     import urllib.request
@@ -1590,6 +1618,12 @@ def test_start_ungated_base_preflight_is_noop(client, monkeypatch):
     import routes.training as tr
 
     monkeypatch.setattr(tr, "_free_gpu_for_diffusion_training", lambda: None)
+    # This asserts the gated-repo probe is a no-op, not device support: pin the precision
+    # preflight so a GPU-less test host does not 400 for an unrelated reason.
+    monkeypatch.setattr(
+        "core.training.diffusion_train_common.training_precision_preflight_error",
+        lambda fam, prec: None,
+    )
     monkeypatch.setattr(urllib.request, "urlopen", lambda req, timeout = None: object())
     r = client.post(
         "/api/train/diffusion/start",

studio/backend/tests/test_gguf_load_cache_reuse.py

@@ -628,6 +628,24 @@ class TestLoadHubDownloadExclusion:
         with gguf_load_in_flight(None):
             assert not hf_gguf_load_in_flight("")
 
+    def test_chat_load_marker_is_repo_agnostic_and_nests(self):
+        # The GPU arbiter needs to know a chat load exists before llama-server is spawned, for
+        # local paths and safetensors too, so this marker carries no repo key.
+        from core.inference.llama_cpp import chat_load_active, chat_load_in_flight
+
+        assert not chat_load_active()
+        with chat_load_in_flight():
+            assert chat_load_active()
+            with chat_load_in_flight():
+                assert chat_load_active()
+            assert chat_load_active()
+        assert not chat_load_active()
+
+        with pytest.raises(RuntimeError):
+            with chat_load_in_flight():
+                raise RuntimeError("boom")
+        assert not chat_load_active()
+
     def test_marker_cleared_on_exception(self):
         with pytest.raises(RuntimeError):
             with gguf_load_in_flight(REPO):
@@ -808,27 +826,26 @@ class TestLoadHubDownloadExclusion:
 
         asyncio.run(scenario())
 
-    def test_load_marker_precedes_hub_guard_and_unload(self):
+    def test_load_marker_precedes_hub_guard_which_precedes_the_gpu_handoff(self):
         source = (Path(__file__).resolve().parent.parent / "routes" / "inference.py").read_text()
-        # _load_model_impl has more than one `if config.is_gguf:`, so anchor on
-        # the branch that actually owns the load marker rather than the first
-        # one in the file, which belongs to an earlier check.
-        marker = source.index("enter_context(gguf_load_in_flight")
-        gguf_branch_start = source.rindex("if config.is_gguf:", 0, marker)
-        gguf_branch = source[gguf_branch_start:]
+        impl = source[source.index("async def _load_model_impl") :]
 
-        # The gguf_load_in_flight marker must be entered before the hub-download
-        # guard and the unload so a concurrent load can't race the download
-        # manager. The llama_extra_args inheritance moved out of the branch into
-        # _resolve_inherited_extra_args, which must still run BEFORE it: the
-        # inherited value (e.g. a carried --no-mmproj) shapes the guard's
-        # require_mmproj. Anchor on the call form so the assertion pins the
-        # endpoint's call site, not the function definition.
-        assert source.index("= _resolve_inherited_extra_args(") < gguf_branch_start
+        # One chain, in this order:
+        # - _resolve_inherited_extra_args first: the inherited value (e.g. a carried
+        #   --no-mmproj) shapes the guard's require_mmproj.
+        # - the gguf_load_in_flight marker before the hub-download guard: that pair is the
+        #   handshake that keeps a load and the download manager off the same files.
+        # - both before the CHAT handoff: the guard's 409 loads nothing, so checking it after
+        #   the handoff destroyed a resident Images/Video pipeline for a load that could never
+        #   start. The handoff registers its own marker under the arbiter lock.
+        # - the resident unload last.
+        # Anchored on call forms so each assertion pins a call site, not a definition.
         assert (
-            gguf_branch.index("enter_context(gguf_load_in_flight")
-            < gguf_branch.index("_hub_download_blocks_gguf_load")
-            < gguf_branch.index("unsloth_backend.unload_model")
+            impl.index("= _resolve_inherited_extra_args(")
+            < impl.index("enter_context(gguf_load_in_flight")
+            < impl.index("_hub_download_blocks_gguf_load")
+            < impl.index("enter_context(chat_load_in_flight")
+            < impl.index("unsloth_backend.unload_model")
         )
         llama_source = (
             Path(__file__).resolve().parent.parent / "core" / "inference" / "llama_cpp.py"

studio/backend/tests/test_gpu_arbiter.py

@@ -203,3 +203,49 @@ def test_competing_acquire_blocks_until_register_completes(monkeypatch):
     assert competitor_done.wait(2.0)
     assert evicted == ["evict-diffusion"]
     assert arb.current_owner() == arb.VIDEO
+
+
+def test_evict_chat_cancels_a_chat_load_that_has_not_spawned_yet(monkeypatch):
+    # An HF chat load has no llama-server process until its GGUF finished downloading, which is
+    # minutes. Gating only on is_active let the evictor find nothing to cancel, grant the GPU to
+    # the image/video load, and the chat load then spawned onto the same device: two big models
+    # allocating at once. The in-flight marker is what makes that load cancellable.
+    import core.inference as core_inference
+    import routes.inference as routes_inference
+    from core.inference.llama_cpp import chat_load_in_flight
+
+    unloaded: list[bool] = []
+
+    class _FakeLlama:
+        is_active = False  # nothing spawned yet: the download is still running
+        is_loaded = False
+
+        def unload_model(self):
+            unloaded.append(True)
+
+        def _wait_for_vram_settle(self, *, since_kill):
+            pass
+
+    class _FakeOrchestrator:
+        active_model_name = None
+
+        def unload_model(self, name):
+            pass
+
+        def _shutdown_subprocess(self, timeout = 5.0):
+            pass
+
+    monkeypatch.setattr(routes_inference, "get_llama_cpp_backend", lambda: _FakeLlama())
+    monkeypatch.setattr(core_inference, "get_inference_backend", lambda: _FakeOrchestrator())
+
+    # No load in flight: nothing to cancel.
+    arb._evict_chat()
+    assert unloaded == []
+
+    with chat_load_in_flight():
+        arb._evict_chat()
+    assert unloaded == [True]
+
+    # The marker is released with the load, so a later eviction is a no-op again.
+    arb._evict_chat()
+    assert unloaded == [True]

studio/backend/tests/test_gpu_selection.py

@@ -1127,6 +1127,7 @@ class TestRouteErrors(unittest.TestCase):
             gguf_hf_repo = None,
             gguf_file = "/tmp/test.gguf",
             gguf_mmproj_file = None,
+            gguf_mtp_file = None,
             gguf_variant = None,
             identifier = "unsloth/test.gguf",
             display_name = "unsloth/test.gguf",
@@ -1136,6 +1137,10 @@ class TestRouteErrors(unittest.TestCase):
             has_audio_input = False,
         )
         acquired = []
+        # Make [0, 1] invalid on any host (a duplicate id is rejected everywhere, GPUs or not):
+        # the point of the test is the ORDER -- validation before the handoff -- not this
+        # machine's device count.
+        request.gpu_ids = [0, 0]
         with (
             patch.object(
                 inference_route,
@@ -1145,7 +1150,10 @@ class TestRouteErrors(unittest.TestCase):
             patch.object(inference_route, "_guard_chat_load_against_training", return_value = None),
             patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread),
             patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext),
-            patch.object(arb, "acquire_for", lambda owner: acquired.append(owner)),
+            # The chat handoff passes a `register` hook (the in-flight marker), so accept it.
+            patch.object(
+                arb, "acquire_for", lambda owner, register = None: acquired.append(owner)
+            ),
         ):
             with self.assertRaises(HTTPException) as exc_info:
                 asyncio.run(
@@ -1162,6 +1170,126 @@ class TestRouteErrors(unittest.TestCase):
         self.assertEqual(exc_info.exception.status_code, 400)
         self.assertEqual(acquired, [])  # no CHAT handoff before the doomed load errored
 
+    def test_inference_route_checks_hub_download_conflict_before_the_handoff(self):
+        # A GGUF the download manager is fetching 409s and loads nothing, so that check has to
+        # run BEFORE the CHAT handoff: afterwards, it destroyed the resident Images/Video
+        # pipeline for a load that could never start.
+        import core.inference.gpu_arbiter as arb
+        import core.inference.llama_cpp as llama_cpp
+
+        inference_route = _load_route_module(
+            "inference_route_module_for_hub_conflict_test",
+            "routes/inference.py",
+        )
+        request = LoadRequest(model_path = "unsloth/Qwen3-4B-GGUF", gguf_variant = "Q4_K_M")
+        model_config = SimpleNamespace(
+            is_gguf = True,
+            is_lora = False,
+            gguf_hf_repo = "unsloth/Qwen3-4B-GGUF",
+            gguf_file = None,
+            gguf_mmproj_file = None,
+            gguf_variant = "Q4_K_M",
+            identifier = "unsloth/Qwen3-4B-GGUF",
+            display_name = "Qwen3-4B",
+            is_vision = False,
+            is_audio = False,
+            audio_type = None,
+            has_audio_input = False,
+        )
+        acquired = []
+        with (
+            patch.object(
+                inference_route,
+                "ModelConfig",
+                SimpleNamespace(from_identifier = lambda **_kwargs: model_config),
+            ),
+            patch.object(inference_route, "_guard_chat_load_against_training", return_value = None),
+            patch.object(inference_route, "_resolve_inherited_extra_args", lambda *a, **k: None),
+            patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread),
+            patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext),
+            patch.object(llama_cpp, "_hub_download_blocks_gguf_load", lambda *a, **k: True),
+            patch.object(arb, "acquire_for", lambda *a, **k: acquired.append(a[0])),
+        ):
+            with self.assertRaises(HTTPException) as exc_info:
+                asyncio.run(
+                    inference_route._load_model_impl(
+                        request,
+                        SimpleNamespace(
+                            app = SimpleNamespace(
+                                state = SimpleNamespace(llama_parallel_slots = 1),
+                            ),
+                        ),
+                        current_subject = "test-user",
+                    )
+                )
+        self.assertEqual(exc_info.exception.status_code, 409)
+        self.assertIn("download", exc_info.exception.detail.lower())
+        self.assertEqual(acquired, [])  # nothing evicted for a load that cannot start
+
+    def test_inference_route_marks_the_chat_load_under_the_arbiter_lock(self):
+        # A chat load holds no llama-server process until its GGUF downloaded, so the arbiter is
+        # told about it through acquire_for's `register` hook (which runs under the arbiter lock).
+        # Passing no register left a competing Images/Video acquire with nothing to cancel.
+        import core.inference.gpu_arbiter as arb
+        import core.inference.llama_cpp as llama_cpp
+
+        inference_route = _load_route_module(
+            "inference_route_module_for_chat_marker_test",
+            "routes/inference.py",
+        )
+        request = LoadRequest(model_path = "unsloth/Qwen3-4B-GGUF", gguf_variant = "Q4_K_M")
+        model_config = SimpleNamespace(
+            is_gguf = True,
+            is_lora = False,
+            gguf_hf_repo = "unsloth/Qwen3-4B-GGUF",
+            gguf_file = None,
+            gguf_mmproj_file = None,
+            gguf_variant = "Q4_K_M",
+            identifier = "unsloth/Qwen3-4B-GGUF",
+            display_name = "Qwen3-4B",
+            is_vision = False,
+            is_audio = False,
+            audio_type = None,
+            has_audio_input = False,
+        )
+        marked = []
+
+        def _acquire(owner, register = None):
+            # Under the arbiter lock the evictor must already be able to see this load.
+            if register is not None:
+                register()
+                marked.append(llama_cpp.chat_load_active())
+            raise RuntimeError("stop the load here")
+
+        with (
+            patch.object(
+                inference_route,
+                "ModelConfig",
+                SimpleNamespace(from_identifier = lambda **_kwargs: model_config),
+            ),
+            patch.object(inference_route, "_guard_chat_load_against_training", return_value = None),
+            patch.object(inference_route, "_resolve_inherited_extra_args", lambda *a, **k: None),
+            patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread),
+            patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext),
+            patch.object(llama_cpp, "_hub_download_blocks_gguf_load", lambda *a, **k: False),
+            patch.object(arb, "acquire_for", _acquire),
+        ):
+            with self.assertRaises(HTTPException):
+                asyncio.run(
+                    inference_route._load_model_impl(
+                        request,
+                        SimpleNamespace(
+                            app = SimpleNamespace(
+                                state = SimpleNamespace(llama_parallel_slots = 1),
+                            ),
+                        ),
+                        current_subject = "test-user",
+                    )
+                )
+        self.assertEqual(marked, [True])
+        # The marker is scoped to the request: it must not outlive the failed load.
+        self.assertFalse(llama_cpp.chat_load_active())
+
     def test_training_route_returns_400_for_invalid_gpu_ids(self):
         training_route = _load_route_module(
             "training_route_module_for_test",

studio/frontend/src/features/images/images-page.tsx

@@ -1838,12 +1838,16 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
     // to load the other page's checkpoint as its own kind of model.
     if (!active) return;
     const wanted = routeSearch.model;
-    // Key on the model AND the quant: this page stays mounted across navigation, so a
-    // model-only marker made every later pick of the same repo a no-op -- including a
-    // different quant, or the same one after chat evicted it -- without even clearing
-    // the query string, so the click did nothing at all.
-    const key = wanted ? `${wanted}${routeSearch.quant ?? ""}` : null;
-    if (!wanted || !key || handledRouteModel.current === key) return;
+    // Key on the model AND the quant, and release the marker once the query is gone. This
+    // page stays mounted, so a marker that outlived the query made re-picking the same
+    // checkpoint (after chat evicted it) a click that neither loaded nor cleared the URL.
+    // It only has to survive the re-renders between consuming a pick and the clear landing.
+    if (!wanted) {
+      handledRouteModel.current = null;
+      return;
+    }
+    const key = `${wanted}|${routeSearch.quant ?? ""}`;
+    if (handledRouteModel.current === key) return;
     handledRouteModel.current = key;
     void navigateSelf({ to: "/images", search: {}, replace: true });
     void loadOrStage(

studio/frontend/src/features/video/video-page.tsx

@@ -1170,10 +1170,15 @@ export function VideoPage({ active = true }: { active?: boolean }) {
     // load an image checkpoint as a video model.
     if (!active) return;
     const wanted = routeSearch.model;
-    // Model AND quant, as on the Images page: this page stays mounted, so a model-only
-    // marker turned every later pick of the same repo into a silent no-op.
-    const key = wanted ? `${wanted} ${routeSearch.quant ?? ""}` : null;
-    if (!wanted || !key || handledRouteModel.current === key) return;
+    // Model AND quant, released once the query is gone, as on the Images page: this page
+    // stays mounted, so a marker that outlived the query turned re-picking the same
+    // checkpoint into a click that neither loaded nor cleared the URL.
+    if (!wanted) {
+      handledRouteModel.current = null;
+      return;
+    }
+    const key = `${wanted}|${routeSearch.quant ?? ""}`;
+    if (handledRouteModel.current === key) return;
     handledRouteModel.current = key;
     void navigateSelf({ to: "/video", search: {}, replace: true });
     void loadOrStage(