Studio: close arbiter load-registration race and surface native progress + local pipeline folders

Publish native sd.cpp generate progress (_gen) before LoRA resolution so a reload probe reads active during setup, matching the diffusers path.

Register the diffusion/video GPU load under the arbiter lock (acquire_for now takes a register callback) so a competing acquire cannot evict an owner before its load is marked in-flight and let two loaders allocate VRAM at once.

Admit local diffusers pipeline folders (root model_index.json, weights in component subdirs) in the local model scan so they reach task tagging and the On Device picker.

studio/backend/core/inference/gpu_arbiter.py

@@ -12,7 +12,7 @@ under the lock, so a transfer is atomic vs other acquires.
 from __future__ import annotations
 
 import threading
-from typing import Optional
+from typing import Any, Callable, Optional
 
 from loggers import get_logger
 
@@ -64,8 +64,17 @@ def _evict_video() -> None:
 _EVICTORS = {CHAT: _evict_chat, DIFFUSION: _evict_diffusion, VIDEO: _evict_video}
 
 
-def acquire_for(owner: str) -> None:
-    """Make ``owner`` the sole GPU owner, evicting the other if it holds it."""
+def acquire_for(owner: str, register: Optional[Callable[[], Any]] = None) -> Any:
+    """Make ``owner`` the sole GPU owner, evicting the other if it holds it.
+
+    ``register``, if given, runs under the arbiter lock right after ownership transfers,
+    and its return value is returned. Registering the in-flight load HERE -- not after
+    ``acquire_for`` returns -- closes the window where a competing acquire could evict this
+    owner before its load is marked in-flight: eviction would then find nothing to cancel
+    and both loaders would allocate VRAM at once. ``register`` must be quick (it holds the
+    lock) and must not re-enter the arbiter. If it raises, ownership stays with ``owner`` --
+    matching the pre-register behaviour where a failed load left the handoff in place.
+    """
     global _owner
     if owner not in _EVICTORS:
         raise ValueError(f"unknown GPU owner: {owner!r}")
@@ -74,6 +83,7 @@ def acquire_for(owner: str) -> None:
             logger.info("gpu_arbiter: evicting %s for %s", _owner, owner)
             _EVICTORS[_owner]()
         _owner = owner
+        return register() if register is not None else None
 
 
 def release(owner: str) -> None:

studio/backend/core/inference/sd_cpp_backend.py

@@ -764,6 +764,12 @@ class SdCppDiffusionBackend:
                     self._state = None
                     raise RuntimeError(DIFFUSION_NOT_LOADED_MSG)
                 self._active_generate_cancel = cancel
+                # Publish an active (step 0) state now, before the slow pre-generate setup
+                # (LoRA listing/download), so a reload's progress probe doesn't read idle
+                # while this generation already holds _generate_lock and let a second generate
+                # queue behind it. The parsed sd-cli progress lines advance this step count.
+                # Mirrors DiffusionBackend.generate, which publishes _gen before its setup.
+                self._gen = _SdGen(total_steps = int(steps))
             try:
                 if seed is None:
                     seed = int.from_bytes(os.urandom(6), "big") & ((1 << 53) - 1)
@@ -789,7 +795,6 @@ class SdCppDiffusionBackend:
                     lora_resolved = diffusion_lora.resolve_specs(
                         active_loras, hf_token = state.hf_token, cancel_event = cancel
                     )
-                self._gen = _SdGen(total_steps = int(steps))
                 if state.mode == "server" and state.server is not None:
                     images, seeds = self._generate_server(
                         state,

studio/backend/routes/inference.py

@@ -14338,10 +14338,36 @@ async def load_diffusion_model(
         # engine name -- else we'd evict a resident chat model for a load that can't use the GPU.
         device = await asyncio.to_thread(lambda: resolve_diffusion_device_target().device)
         needs_gpu = device != "cpu"
+
+        def _begin_load():
+            # Kicks the (slow) load onto a background thread and returns at once (the client
+            # polls images/load-progress); begin_load itself validates network-free.
+            return engine.begin_load(
+                request.model_path,
+                gguf_filename = request.gguf_filename,
+                base_repo = request.base_repo,
+                family_override = request.family_override,
+                hf_token = request.hf_token,
+                cpu_offload = request.cpu_offload,
+                memory_mode = request.memory_mode,
+                speed_mode = request.speed_mode,
+                text_encoder_quant = request.text_encoder_quant,
+                transformer_quant = request.transformer_quant,
+                transformer_quant_fast_accum = request.transformer_quant_fast_accum,
+                transformer_prequant_path = request.transformer_prequant_path,
+                attention_backend = request.attention_backend,
+                transformer_cache = request.transformer_cache,
+                transformer_cache_threshold = request.transformer_cache_threshold,
+                model_kind = kind,
+            )
+
         if needs_gpu:
-            # Then kick the (slow) load onto a background thread and return at once --
-            # the client polls images/load-progress.
-            await asyncio.to_thread(acquire_for, DIFFUSION)
+            # Register the in-flight load UNDER the arbiter lock (not after acquire_for
+            # returns): a competing Video/chat acquire in that gap would otherwise evict
+            # DIFFUSION before begin_load marks a load in-flight, so eviction finds nothing
+            # to cancel and both loaders allocate VRAM at once. begin_load returns at once,
+            # so the lock is held only briefly.
+            status_dict = await asyncio.to_thread(acquire_for, DIFFUSION, _begin_load)
         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
@@ -14349,25 +14375,7 @@ async def load_diffusion_model(
             # "evict" this CPU model for no reason. Release that stale ownership -- release()
             # is owner-guarded, so it's a no-op when diffusion never owned the GPU.
             await asyncio.to_thread(release, DIFFUSION)
-        status_dict = await asyncio.to_thread(
-            engine.begin_load,
-            request.model_path,
-            gguf_filename = request.gguf_filename,
-            base_repo = request.base_repo,
-            family_override = request.family_override,
-            hf_token = request.hf_token,
-            cpu_offload = request.cpu_offload,
-            memory_mode = request.memory_mode,
-            speed_mode = request.speed_mode,
-            text_encoder_quant = request.text_encoder_quant,
-            transformer_quant = request.transformer_quant,
-            transformer_quant_fast_accum = request.transformer_quant_fast_accum,
-            transformer_prequant_path = request.transformer_prequant_path,
-            attention_backend = request.attention_backend,
-            transformer_cache = request.transformer_cache,
-            transformer_cache_threshold = request.transformer_cache_threshold,
-            model_kind = kind,
-        )
+            status_dict = await asyncio.to_thread(_begin_load)
         return DiffusionStatusResponse(**annotate_status(status_dict))
     except (ValueError, FileNotFoundError) as exc:
         raise HTTPException(status_code = 400, detail = redact_native_paths(str(exc)))

studio/backend/routes/models.py

@@ -338,7 +338,15 @@ def _scan_models_dir(models_dir: Path, *, limit: int | None = None) -> List[Loca
             has_config = (child / "config.json").exists() or (
                 child / "adapter_config.json"
             ).exists()
-            has_model_files = has_gguf or has_non_gguf_weights or has_config
+            # A standard diffusers PIPELINE folder keeps its weights/configs in component
+            # subdirs (transformer/, vae/, ...) and carries only model_index.json at the
+            # root, so the checks above miss it. The Images/Video load path accepts such a
+            # local pipeline dir, so admit it here too (task tagging then classifies it via
+            # _local_is_diffusers); otherwise it is hidden from the On Device picker.
+            has_pipeline_index = (child / "model_index.json").is_file()
+            has_model_files = (
+                has_gguf or has_non_gguf_weights or has_config or has_pipeline_index
+            )
         except OSError:
             # Skip unreadable children rather than failing the scan.
             continue

studio/backend/routes/video.py

@@ -112,26 +112,36 @@ async def load_video_model(
         # Take the GPU from chat only for a non-CPU load; a CPU load never touches GPU memory,
         # so key off the device. Release stale VIDEO ownership on a CPU load (owner-guarded no-op).
         device = await asyncio.to_thread(lambda: resolve_diffusion_device_target().device)
+
+        def _begin_load():
+            # Kicks the (slow) load onto a background thread and returns at once;
+            # begin_load itself validates network-free.
+            return backend.begin_load(
+                request.model_path,
+                gguf_filename = request.gguf_filename,
+                base_repo = request.base_repo,
+                family_override = request.family_override,
+                hf_token = request.hf_token,
+                memory_mode = request.memory_mode,
+                speed_mode = request.speed_mode,
+                attention_backend = request.attention_backend,
+                transformer_cache = request.transformer_cache,
+                transformer_cache_threshold = request.transformer_cache_threshold,
+                transformer_quant = request.transformer_quant,
+                text_encoder_quant = request.text_encoder_quant,
+                model_kind = kind,
+            )
+
         if device != "cpu":
-            await asyncio.to_thread(acquire_for, VIDEO)
+            # Register the in-flight load UNDER the arbiter lock (not after acquire_for
+            # returns): a competing Images/chat acquire in that gap would otherwise evict
+            # VIDEO before begin_load marks a load in-flight, so eviction finds nothing to
+            # cancel and both loaders allocate VRAM at once. begin_load returns at once, so
+            # the lock is held only briefly. Mirrors the images/load handoff.
+            status_dict = await asyncio.to_thread(acquire_for, VIDEO, _begin_load)
         else:
             await asyncio.to_thread(release, VIDEO)
-        status_dict = await asyncio.to_thread(
-            backend.begin_load,
-            request.model_path,
-            gguf_filename = request.gguf_filename,
-            base_repo = request.base_repo,
-            family_override = request.family_override,
-            hf_token = request.hf_token,
-            memory_mode = request.memory_mode,
-            speed_mode = request.speed_mode,
-            attention_backend = request.attention_backend,
-            transformer_cache = request.transformer_cache,
-            transformer_cache_threshold = request.transformer_cache_threshold,
-            transformer_quant = request.transformer_quant,
-            text_encoder_quant = request.text_encoder_quant,
-            model_kind = kind,
-        )
+            status_dict = await asyncio.to_thread(_begin_load)
         return VideoStatusResponse(**status_dict)
     except (ValueError, FileNotFoundError) as exc:
         raise HTTPException(status_code = 400, detail = redact_native_paths(str(exc)))

studio/backend/tests/test_diffusion_routes.py

@@ -789,7 +789,13 @@ def _force_engine(monkeypatch, backend, *, engine_name, device):
         devmod, "resolve_diffusion_device_target", lambda: _types.SimpleNamespace(device = device)
     )
     acquired: list = []
-    monkeypatch.setattr(gpu_arbiter, "acquire_for", lambda role: acquired.append(role))
+
+    def _fake_acquire(role, register = None):
+        # Mirror the real arbiter: record the handoff and run the (registered) load under it.
+        acquired.append(role)
+        return register() if register is not None else None
+
+    monkeypatch.setattr(gpu_arbiter, "acquire_for", _fake_acquire)
     return acquired
 
 

studio/backend/tests/test_gpu_arbiter.py

@@ -104,3 +104,75 @@ def test_evict_chat_unloads_a_still_loading_chat_backend(monkeypatch):
     arb._evict_chat()
 
     assert unloaded == [True]  # still-loading chat backend was unloaded, not skipped
+
+
+def test_register_runs_under_ownership_and_returns_result(calls):
+    # A register callback runs after ownership transfers (owner already set) and its
+    # return value is forwarded -- the route uses this to register the in-flight load.
+    seen_owner: list = []
+
+    def register():
+        seen_owner.append(arb.current_owner())
+        return "status-dict"
+
+    result = arb.acquire_for(arb.DIFFUSION, register)
+    assert result == "status-dict"
+    assert seen_owner == [arb.DIFFUSION]
+    assert arb.current_owner() == arb.DIFFUSION
+
+
+def test_register_failure_leaves_ownership_in_place(calls):
+    # A failing register (e.g. begin_load reporting a load already in progress) propagates
+    # but must not drop ownership -- the prior handoff (chat already evicted) stands.
+    arb.acquire_for(arb.CHAT)
+
+    def register():
+        raise RuntimeError("A diffusion load is already in progress.")
+
+    with pytest.raises(RuntimeError):
+        arb.acquire_for(arb.DIFFUSION, register)
+    assert calls == ["evict-chat"]
+    assert arb.current_owner() == arb.DIFFUSION
+
+
+def test_competing_acquire_blocks_until_register_completes(monkeypatch):
+    # The window this closes: while DIFFUSION registers its load, a competing VIDEO acquire
+    # must not evict DIFFUSION until the load is marked in-flight. Holding the lock across
+    # register makes the competitor wait, so eviction never races an unregistered load.
+    import threading
+    import time
+
+    monkeypatch.setattr(arb, "_owner", None)
+    evicted: list = []
+    monkeypatch.setitem(arb._EVICTORS, arb.DIFFUSION, lambda: evicted.append("evict-diffusion"))
+    monkeypatch.setitem(arb._EVICTORS, arb.VIDEO, lambda: evicted.append("evict-video"))
+
+    in_register = threading.Event()
+    release_register = threading.Event()
+
+    def register():
+        in_register.set()
+        # Hold the arbiter lock here; a competing acquire_for(VIDEO) must block until we return.
+        assert release_register.wait(2.0)
+        return "loading"
+
+    loader = threading.Thread(target = lambda: arb.acquire_for(arb.DIFFUSION, register))
+    loader.start()
+    assert in_register.wait(2.0)
+
+    competitor_done = threading.Event()
+    threading.Thread(
+        target = lambda: (arb.acquire_for(arb.VIDEO), competitor_done.set()),
+    ).start()
+
+    # The competitor cannot evict DIFFUSION while register still holds the lock.
+    time.sleep(0.1)
+    assert evicted == []
+    assert not competitor_done.is_set()
+
+    # Let register finish; ownership is now safely registered, so the competitor proceeds.
+    release_register.set()
+    loader.join(2.0)
+    assert competitor_done.wait(2.0)
+    assert evicted == ["evict-diffusion"]
+    assert arb.current_owner() == arb.VIDEO

studio/backend/tests/test_local_model_format.py

@@ -117,6 +117,25 @@ def test_scan_models_dir_classifies_root_gguf_with_config(tmp_path):
     assert row.model_format == "gguf"
 
 
+def test_scan_models_dir_surfaces_diffusers_pipeline_folder(tmp_path):
+    # A standard diffusers PIPELINE folder keeps its weights/configs in component subdirs
+    # (transformer/, vae/, ...) and carries only model_index.json at the root. The Images/Video
+    # load path accepts such a local pipeline dir, so the scan must surface it -- otherwise the
+    # weights-in-subdirs layout is missed and it never reaches task tagging / the On Device
+    # picker. It is not a GGUF, so model_format stays None (task tagging classifies it later).
+    root = tmp_path / "models"
+    pipe = root / "my-pipeline"
+    _touch(pipe / "model_index.json")
+    _touch(pipe / "transformer" / "config.json")
+    _touch(pipe / "transformer" / "diffusion_pytorch_model.safetensors")
+    _touch(pipe / "vae" / "diffusion_pytorch_model.safetensors")
+
+    rows = {Path(m.path).name: m for m in models_route._scan_models_dir(root)}
+
+    assert "my-pipeline" in rows
+    assert rows["my-pipeline"].model_format is None
+
+
 # ── Images picker task tag for local (non-GGUF) diffusers models ──────────────
 from models.models import LocalModelInfo  # noqa: E402
 

studio/backend/tests/test_sd_cpp_backend.py

@@ -279,6 +279,32 @@ def test_generate_progress_tracks_parsed_steps():
     assert b.generate_progress()["step"] == 4
 
 
+def test_generate_publishes_progress_before_lora_resolution(monkeypatch):
+    # Native LoRA resolution (listing/downloading a not-yet-cached adapter) happens during the
+    # pre-generate setup while _generate_lock is already held. A reload/progress probe in that
+    # window must read ACTIVE, not idle, or the UI queues a second generate behind the first.
+    # So _gen is published before LoRA resolution, mirroring the diffusers path.
+    from core.inference import diffusion_lora
+
+    eng = _FakeEngine()
+    b = _loaded_backend(engine = eng)
+    monkeypatch.setattr(diffusion_lora, "supports_lora", lambda **_k: True)
+
+    seen: dict = {}
+
+    def _resolve(active, *, hf_token = None, cancel_event = None):
+        # Mid-setup: the in-flight generation must already be reported as active.
+        seen["progress"] = b.generate_progress()
+        return []
+
+    monkeypatch.setattr(diffusion_lora, "resolve_specs", _resolve)
+
+    out = b.generate(prompt = "a fox", width = 64, height = 64, steps = 8, loras = [("some/lora", 1.0)])
+    assert out["images"]
+    assert seen["progress"]["active"] is True
+    assert seen["progress"]["total_steps"] == 8
+
+
 # ── load validation + binary install ──────────────────────────────────────────
 
 

studio/backend/tests/test_video_routes.py

@@ -263,7 +263,13 @@ def test_load_happy_path_and_arbiter_acquired(client, monkeypatch):
         devmod, "resolve_diffusion_device_target", lambda: types.SimpleNamespace(device = "cuda")
     )
     acquired: list = []
-    monkeypatch.setattr(gpu_arbiter, "acquire_for", lambda role: acquired.append(role))
+
+    def _fake_acquire(role, register = None):
+        # Mirror the real arbiter: record the handoff and run the (registered) load under it.
+        acquired.append(role)
+        return register() if register is not None else None
+
+    monkeypatch.setattr(gpu_arbiter, "acquire_for", _fake_acquire)
 
     resp = client.post(
         "/api/inference/video/load",