Close the load-versus-training-start race, and two picker fixes

- The image and video load guards read is_active() and only then selected an
  engine, acquired the arbiter and registered the load. A /train/diffusion/start
  reserving inside that window freed residents the load had not registered yet,
  so the trainer came up beside a brand-new pipeline. The service already had
  exactly the right pattern for this in dataset_mutation, so gpu_load_admission
  mirrors it: reserve() refuses while an admission is open, an admission refuses
  once a start is reserved, both decided under the one lock. The span is only the
  registration, since begin_load returns as soon as the load is registered and
  _free_gpu_for_diffusion_training preempts an in-flight load from that point.
  Chat is deliberately not covered: its load spans an eviction plus a multi-minute
  GGUF load, and it admits models that fit beside training by design, which is a
  different contract from the diffusion pipeline's all-or-nothing one.

- Hugging Face gives the LTX-2 family the image-to-video pipeline_tag (both
  Lightricks/LTX-2 and unsloth/LTX-2.3-GGUF report it), so a text-to-video-only
  filter dropped the flagship audio family out of Video Hub search while the rest
  of the app routed it to Video.

- Task-scoped quant fit sized picks against the LARGEST visible device while
  resolve_diffusion_device_target returns a bare "cuda" and torch places on the
  current one. On a heterogeneous host that recommended a checkpoint sized for the
  bigger card and then loaded it onto the smaller one. Fit now uses the device the
  load actually lands on; identical on a homogeneous host.

studio/backend/core/training/diffusion_training_service.py

@@ -228,6 +228,10 @@ class DiffusionTrainingService:
         # A start refuses while any is open and a mutation refuses once a start is reserved, both
         # decided under _lock, so neither can slip through the other's check-then-act window.
         self._dataset_mutations = 0
+        # GPU load admissions in flight (an image/video/chat load between its training guard and
+        # the moment it registers with the arbiter). Same two-sided rule as the dataset mutations:
+        # a start refuses while one is open, and an admission refuses once a start is reserved.
+        self._gpu_admissions = 0
         self._proc: Any = None
         self._stop_queue: Any = None
         self._pump: Optional[threading.Thread] = None
@@ -262,6 +266,15 @@ class DiffusionTrainingService:
                     "The training images are being changed right now. Wait for that to finish, "
                     "then start the run."
                 )
+            if self._gpu_admissions:
+                # A load already passed its training guard and is about to take the GPU. Reserving
+                # now would free residents it has not registered yet, so the trainer and a
+                # brand-new pipeline would allocate together. Refusing is safe: the admission is
+                # held only across the load's registration, not the load itself.
+                raise RuntimeError(
+                    "A model is being loaded onto the GPU right now. Wait for that to finish, "
+                    "then start the run."
+                )
             self._reserved = True
 
     def unreserve(self) -> None:
@@ -295,6 +308,34 @@ class DiffusionTrainingService:
             with self._lock:
                 self._dataset_mutations = max(0, self._dataset_mutations - 1)
 
+    @contextlib.contextmanager
+    def gpu_load_admission(self):
+        """Hold the GPU-admission interlock across a load's guard -> arbiter -> registration.
+
+        The load guards read ``is_active()`` and only THEN acquire the arbiter and register the
+        load, so a start reserving inside that gap freed residents the load had not registered
+        yet and the trainer came up beside a brand-new pipeline. Registering the admission under
+        the same lock ``reserve()`` uses closes it from both sides, exactly like
+        ``dataset_mutation``: this raises once a start is reserved or running, and ``reserve()``
+        raises while an admission is open, so neither waits on the other.
+
+        The span is deliberately short. ``begin_load`` returns as soon as the load is registered
+        (the download and build run on a daemon thread), and from that point
+        ``_free_gpu_for_diffusion_training`` preempts the in-flight load, so holding this for the
+        whole load would block starts for minutes to no purpose."""
+        with self._lock:
+            if self._reserved or (self._proc is not None and self._proc.is_alive()):
+                raise TrainingActiveError(
+                    "Diffusion training is running, so the GPU is in use. Stop the run before "
+                    "loading a model."
+                )
+            self._gpu_admissions += 1
+        try:
+            yield
+        finally:
+            with self._lock:
+                self._gpu_admissions = max(0, self._gpu_admissions - 1)
+
     def start(self, config: dict) -> str:
         """Validate ``config``, spawn the trainer, and start pumping its events.
 

studio/backend/routes/inference.py

@@ -23,7 +23,7 @@ from loggers import get_logger
 import asyncio
 import threading
 import weakref
-from contextlib import ExitStack
+from contextlib import ExitStack, contextmanager
 
 
 import re as _re
@@ -16000,6 +16000,30 @@ def _diffusion_training_active() -> bool:
         return False
 
 
+@contextmanager
+def _diffusion_training_admission():
+    """Hold the diffusion trainer's GPU-admission interlock for this load's registration.
+
+    The guards below only cover the instant they run. A load then selects its engine, acquires
+    the arbiter and registers with the backend, and a ``/train/diffusion/start`` reserving inside
+    that window frees residents this load has not registered yet, so the trainer comes up beside
+    a brand-new pipeline. Registering the admission under the same lock ``reserve()`` takes makes
+    the two mutually exclusive: this raises (409) once a start is reserved, and a start raises
+    while an admission is open.
+
+    Fails open on an import error, like the guards it complements. Covers the DIFFUSION trainer
+    only; the LLM trainer admits loads that fit beside it, which is a different contract."""
+    try:
+        from core.training.diffusion_training_service import get_diffusion_training_service
+
+        service = get_diffusion_training_service()
+    except Exception:  # noqa: BLE001 -- unknowable state never blocks a load
+        yield
+        return
+    with service.gpu_load_admission():
+        yield
+
+
 def _guard_diffusion_load_against_training() -> None:
     """Refuse loading an image model while a training run is active. Unlike chat,
     a diffusion pipeline's VRAM can't be cheaply estimated before the load, so the
@@ -16178,8 +16202,14 @@ async def load_diffusion_model(
         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 is marked in-flight,
-            # finds nothing to cancel, and both loaders allocate VRAM at once.
-            status_dict = await asyncio.to_thread(acquire_for, DIFFUSION, _begin_load)
+            # finds nothing to cancel, and both loaders allocate VRAM at once. The training admission wraps
+            # the same span for the OTHER competitor: a diffusion-training start reserving here would free
+            # residents this load has not registered yet (see _diffusion_training_admission).
+            def _acquire_and_begin():
+                with _diffusion_training_admission():
+                    return acquire_for(DIFFUSION, _begin_load)
+
+            status_dict = await asyncio.to_thread(_acquire_and_begin)
         else:
             # A CPU-only native load never touches the GPU, so it neither acquires nor is tracked by the
             # arbiter. But switching here FROM a previous diffusers/GPU load leaves DIFFUSION still marked as

studio/backend/routes/video.py

@@ -183,8 +183,16 @@ async def load_video_model(
             # Register the in-flight load UNDER the arbiter lock (not after acquire_for returns): otherwise a
             # competing Images/chat acquire in that gap evicts VIDEO before the load is marked in-flight,
             # finds nothing to cancel, and both loaders allocate VRAM at once. Mirrors the images/load
-            # handoff.
-            status_dict = await asyncio.to_thread(acquire_for, VIDEO, _begin_load)
+            # handoff. The training admission wraps the same span for the OTHER competitor: a
+            # diffusion-training start reserving here would free residents this load has not
+            # registered yet (see _diffusion_training_admission).
+            from routes.inference import _diffusion_training_admission
+
+            def _acquire_and_begin():
+                with _diffusion_training_admission():
+                    return acquire_for(VIDEO, _begin_load)
+
+            status_dict = await asyncio.to_thread(_acquire_and_begin)
         else:
             await asyncio.to_thread(release, VIDEO)
             status_dict = await asyncio.to_thread(_begin_load)

studio/backend/tests/test_diffusion_training.py

@@ -1877,3 +1877,37 @@ def test_diffusion_seed_is_bounded_to_torch_range():
     # The extremes torch does accept stay valid.
     for good in (2**64 - 1, -(2**63)):
         assert DiffusionTrainingStartRequest(**request, seed = good).seed == good
+
+
+def test_gpu_load_admission_and_reserve_exclude_each_other():
+    # The load guards read is_active() and only THEN acquire the arbiter and register the load, so a
+    # start reserving inside that window freed residents the load had not registered yet and the
+    # trainer came up beside a brand-new pipeline. The admission closes it from both sides, exactly
+    # like the dataset interlock.
+    from core.training.diffusion_training_service import TrainingActiveError
+
+    svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target)
+
+    # A start cannot reserve while a load is registering.
+    with svc.gpu_load_admission():
+        with pytest.raises(RuntimeError, match = "loaded onto the GPU"):
+            svc.reserve()
+    # ...and the admission is released afterwards, so the start goes through.
+    svc.reserve()
+    try:
+        # A load cannot register while a start is reserved, and it says why.
+        with pytest.raises(TrainingActiveError, match = "Diffusion training is running"):
+            with svc.gpu_load_admission():
+                pass
+    finally:
+        svc.unreserve()
+
+    # Nested/concurrent admissions are counted, not boolean: the first exit must not open the door
+    # while a second load is still registering.
+    with svc.gpu_load_admission():
+        with svc.gpu_load_admission():
+            pass
+        with pytest.raises(RuntimeError, match = "loaded onto the GPU"):
+            svc.reserve()
+    svc.reserve()
+    svc.unreserve()

studio/frontend/src/features/model-picker/components/model-selector/pickers.tsx

@@ -1208,7 +1208,12 @@ export const IMAGE_GEN_TASKS = [
 // Video-generation pipeline tasks: handled by the Video page, never loadable as
 // chat models. The backend reports "text-to-video" for video-diffusion GGUFs. The
 // Video page reuses this as its picker's `task` filter, so it lives here.
-export const VIDEO_GEN_TASKS = ["text-to-video"] as const;
+// image-to-video is here because that is the pipeline_tag Hugging Face gives the LTX-2 family
+// (both Lightricks/LTX-2 and unsloth/LTX-2.3-GGUF report it, alongside text-to-video in their
+// tag list), so a text-to-video-only filter dropped the flagship audio family out of Video Hub
+// search while the rest of the app routed it to Video. Locally cached video GGUFs are unaffected
+// either way: the backend tags those text-to-video itself.
+export const VIDEO_GEN_TASKS = ["text-to-video", "image-to-video"] as const;
 
 // Diffusion GGUF archs the Images backend can't assemble yet (SD/SDXL/PixArt/Wan/...). The
 // backend tags them with this task so the chat picker hides them (they die with "unknown model
@@ -2135,11 +2140,14 @@ export function HubModelPicker({
       ),
     [cachedModels, downloadedSort, loadTimes, task, catalog],
   );
-  // Task-scoped loads put the whole pipeline on one device, so quant fit must use the
-  // largest device, not the multi-GPU sum. Chat keeps the sum (llama.cpp splits layers).
+  // Task-scoped loads put the whole pipeline on ONE device, so quant fit must use a single
+  // device, not the multi-GPU sum. Specifically the device the load will land on (the lowest
+  // visible ordinal), not the largest one: on a heterogeneous host sizing against the bigger
+  // card recommends a checkpoint that then OOMs the smaller card it actually loads onto.
+  // Chat keeps the sum (llama.cpp splits layers).
   const expanderGpuGb = gpu.available
     ? task
-      ? gpu.maxDeviceMemoryGb
+      ? gpu.loadDeviceMemoryGb || gpu.maxDeviceMemoryGb
       : gpu.memoryTotalGb
     : undefined;
 

studio/frontend/src/hooks/use-gpu-info.ts

@@ -10,8 +10,14 @@ export interface GpuInfo {
   name: string;
   memoryTotalGb: number;
   /** Largest single device's VRAM. Image/video loads live on ONE device (no
-   * tensor split), so their fit math must use this, not the multi-GPU sum. */
+   * tensor split), so their fit math must use a single device, not the multi-GPU sum. */
   maxDeviceMemoryGb: number;
+  /** VRAM of the device an image/video load will actually land on: the lowest visible
+   * ordinal, since resolve_diffusion_device_target() returns a bare "cuda" and torch places
+   * on the current device. On a heterogeneous host this is NOT maxDeviceMemoryGb, and sizing
+   * a pick against the larger card would recommend a checkpoint that then OOMs the smaller
+   * one it loads onto. */
+  loadDeviceMemoryGb: number;
   cpuCore: number;
   cpuThread: number;
   systemRamAvailableGb: number;
@@ -36,6 +42,7 @@ const DEFAULT_GPU: GpuInfo = {
   name: "Unknown",
   memoryTotalGb: 0,
   maxDeviceMemoryGb: 0,
+  loadDeviceMemoryGb: 0,
   cpuCore: 0,
   cpuThread: 0,
   systemRamAvailableGb: 0,
@@ -83,6 +90,14 @@ function toGpuInfo(data: SystemInfoResponse | null): GpuInfo {
     name: devices[0]?.name ?? "Unknown",
     memoryTotalGb: devices.reduce((sum, d) => sum + (d.memory_total_gb ?? 0), 0),
     maxDeviceMemoryGb: devices.reduce((max, d) => Math.max(max, d.memory_total_gb ?? 0), 0),
+    // Lowest visible ordinal = torch's current device = where the pipeline lands. The list is
+    // already filtered by any CUDA_VISIBLE_DEVICES mask, so the minimum index is right whether
+    // the reported indices are physical or relative.
+    loadDeviceMemoryGb:
+      devices.reduce(
+        (pick, d) => ((d.index ?? 0) < (pick.index ?? 0) ? d : pick),
+        devices[0],
+      )?.memory_total_gb ?? 0,
   };
 }