Studio: gate gallery serve/export on ownership; keep image progress active until persisted; reserve diffusion training before the dataset scan

studio/backend/core/inference/diffusion.py

@@ -2534,12 +2534,13 @@ class DiffusionBackend:
                     self._reset_step_cache(state.pipe)
 
                 self._gen = gen
-                try:
-                    # inference_mode is faster than no_grad and numerically identical here.
-                    with torch.inference_mode():
-                        images = pipe(**kwargs).images
-                finally:
-                    self._gen = None
+                # inference_mode is faster than no_grad and numerically identical here.
+                with torch.inference_mode():
+                    images = pipe(**kwargs).images
+                # Keep progress ACTIVE through the post-denoise work below (the compile-cache save can
+                # take a moment) -- don't null _gen here. The route persists the image AFTER this
+                # returns, so a reload's mount probe that read idle now would refresh the gallery
+                # before the result exists. The outer finally clears _gen on every exit (return/raise).
                 # A cancelled denoise returns a partial/garbage image; don't persist it.
                 if cancel.is_set():
                     raise RuntimeError(DIFFUSION_CANCELLED_MSG)
@@ -2569,8 +2570,9 @@ class DiffusionBackend:
                 with self._lock:
                     if self._active_generate_cancel is cancel:
                         self._active_generate_cancel = None
-                    # Drop the published progress state, covering a setup-time error that skips
-                    # the inner finally. Safe under _generate_lock.
+                    # Sole clear of the published progress state, on every exit (return, cancel, or a
+                    # setup/denoise error), so it stays active through post-denoise work but a crashed
+                    # generation never leaves the UI stuck "active". Safe under _generate_lock.
                     self._gen = None
 
     def generate_progress(self) -> dict[str, Any]:

studio/backend/core/inference/image_gallery.py

@@ -136,6 +136,17 @@ def _read_meta(path: Path) -> Optional[dict[str, Any]]:
     return meta
 
 
+def owned_image_path(image_id: str) -> Optional[Path]:
+    """Resolve an id to its PNG only when it is a Studio-owned image (a readable recipe chunk),
+    else None. The serve route uses this instead of image_path() so a guessed stem for a
+    hand-dropped foreign PNG -- which list_images/delete/clear already treat as not ours -- can't
+    be streamed out. Mirrors the delete/clear ownership guard."""
+    path = image_path(image_id)
+    if path is None or _read_meta(path) is None:
+        return None
+    return path
+
+
 def _mtime(path: Path) -> float:
     try:
         return path.stat().st_mtime

studio/backend/core/inference/video_gallery.py

@@ -82,7 +82,9 @@ def transcode(video_id: str, fmt: str) -> Optional[bytes]:
     """Re-encode a stored MP4 for the Download menu: "webm" (VP9) or "gif". Returns the bytes, or
     None when the id doesn't resolve. Raises RuntimeError on missing codec/deps (route 501s). MP4
     downloads stream the original via /file, not here."""
-    path = video_path(video_id)
+    # Ownership-gate like /file: only transcode a Studio-owned clip (readable sidecar), so a
+    # guessed stem for a foreign/orphan MP4 the gallery hides can't be re-encoded out either.
+    path = owned_video_path(video_id)
     if path is None:
         return None
     normalized = fmt.strip().lower()
@@ -206,6 +208,17 @@ def _read_meta(sidecar: Path) -> Optional[dict[str, Any]]:
     return meta
 
 
+def owned_video_path(video_id: str) -> Optional[Path]:
+    """Resolve an id to its MP4 only when it is a Studio-owned clip (a readable sidecar), else
+    None. The serve and export routes use this instead of video_path() so a guessed stem for a
+    hand-dropped/orphan MP4 -- which list_videos/delete/clear already treat as not ours -- can't
+    be streamed or transcoded out. Mirrors the delete/clear ownership guard."""
+    path = video_path(video_id)
+    if path is None or _read_meta(_sidecar_path(video_id)) is None:
+        return None
+    return path
+
+
 def _mtime(path: Path) -> float:
     try:
         return path.stat().st_mtime

studio/backend/routes/inference.py

@@ -14382,6 +14382,13 @@ async def load_diffusion_model(
         raise HTTPException(status_code = 409, detail = str(exc))
 
 
+# Count of finished generations still writing their PNG/gallery records. generate-progress reports
+# active while this is > 0 so a reload's mount probe never reads idle between the denoise finishing
+# (engine drops _gen) and the image reaching the gallery, which would refresh the gallery before the
+# record exists. Mutated only on the event loop (around the persist await), so no lock is needed.
+_diffusion_persist_active = 0
+
+
 @studio_router.post("/images/generate", response_model = DiffusionGenerateResponse)
 async def generate_diffusion_image(
     request: DiffusionGenerateRequest, current_subject: str = Depends(get_current_subject)
@@ -14500,11 +14507,18 @@ async def generate_diffusion_image(
             )
         return records
 
+    # Hold generate-progress "active" across the persist so a concurrent reload's mount probe can't
+    # see idle and refresh the gallery before these records exist. Set synchronously right after the
+    # engine returned (no await between, so no idle gap), cleared in the finally.
+    global _diffusion_persist_active
+    _diffusion_persist_active += 1
     try:
         records = await asyncio.to_thread(_persist)
     except Exception as exc:
         logger.error("diffusion.persist_failed: %s", exc)
         raise HTTPException(status_code = 500, detail = "Failed to save the generated image.")
+    finally:
+        _diffusion_persist_active -= 1
 
     return DiffusionGenerateResponse(images = [GalleryImage(**r) for r in records])
 
@@ -14548,7 +14562,9 @@ async def get_gallery_image_file(
 ):
     from core.inference import image_gallery
 
-    path = await asyncio.to_thread(image_gallery.image_path, image_id)
+    # Ownership-gate the serve like delete/clear: resolve only a Studio-owned PNG (readable recipe),
+    # so a guessed stem for a hand-dropped foreign PNG the listing hides can't be streamed out.
+    path = await asyncio.to_thread(image_gallery.owned_image_path, image_id)
     if path is None:
         raise HTTPException(status_code = 404, detail = "Image not found.")
     data = await asyncio.to_thread(path.read_bytes)
@@ -14623,7 +14639,13 @@ async def diffusion_load_progress(current_subject: str = Depends(get_current_sub
 @studio_router.get("/images/generate-progress", response_model = DiffusionGenerateProgressResponse)
 async def diffusion_generate_progress(current_subject: str = Depends(get_current_subject)):
     from core.inference.diffusion_engine_router import get_active_diffusion_engine
-    return DiffusionGenerateProgressResponse(**get_active_diffusion_engine().generate_progress())
+
+    progress = get_active_diffusion_engine().generate_progress()
+    # A finished generation still persisting its gallery record counts as active, so a reload's
+    # mount probe keeps polling instead of refreshing the gallery before the image lands.
+    if _diffusion_persist_active > 0 and not progress["active"]:
+        progress = {**progress, "active": True}
+    return DiffusionGenerateProgressResponse(**progress)
 
 
 # ──────────────────────────────────────────────────────────────────────────

studio/backend/routes/training.py

@@ -1374,36 +1374,39 @@ async def start_diffusion_training(
         _preflight_gated_base, config.get("base_model", ""), config.get("hf_token")
     )
 
-    # Preflight the dataset too: a missing/empty/uncaptionable data_dir otherwise fails inside
-    # the spawned trainer AFTER the user's model was evicted. Same discovery the trainer runs,
-    # so the two cannot disagree.
     from core.training import diffusion_train_common as _dtc
 
-    try:
-        await asyncio.to_thread(
-            _dtc.discover_image_caption_pairs,
-            config["data_dir"],
-            instance_prompt = config.get("instance_prompt") or None,
-            caption_column = config.get("caption_column") or "text",
-            # Decode-probe every image now (cheap PIL header check) so a corrupt/zero-byte upload
-            # 400s BEFORE _free_gpu_for_diffusion_training() tears down the user's models, rather
-            # than crashing the spawned trainer post-eviction.
-            verify_images = True,
-        )
-    except (FileNotFoundError, ValueError) as e:
-        raise HTTPException(status_code = 400, detail = str(e))
-
     service = get_diffusion_training_service()
-    # Reserve the training slot BEFORE freeing residents: is_active() otherwise flips true only at
-    # service.start(), after the free, so a concurrent /images/load or /video/load would pass its
-    # training guard during the free-then-spawn window and double-allocate VRAM. reserve() is a
-    # compare-and-set: a second overlapping /diffusion/start raises RuntimeError (-> 409) before
-    # freeing anything, so two starts never both tear down residents. unreserve() runs in the
-    # finally ONLY when THIS request reserved, so a rejected second request can't clear the claim.
+    # Reserve the training slot BEFORE the dataset preflight (not just before freeing residents):
+    # is_active() otherwise flips true only at service.start(), so during this scan -- which
+    # decode-probes every image and can take noticeable time on a large folder -- a concurrent
+    # upload/caption/delete would pass _require_diffusion_dataset_mutable() and mutate the dataset
+    # the trainer is about to read (training the wrong data, or a missing file mid-step), and a
+    # concurrent /images/load or /video/load would pass its guard and double-allocate VRAM.
+    # reserve() is a compare-and-set: a second overlapping /diffusion/start raises RuntimeError
+    # (-> 409) before touching anything. unreserve() runs in the finally ONLY when THIS request
+    # reserved, so a rejected second request can't clear the claim, and any preflight failure below
+    # rolls the reservation back.
     reserved = False
     try:
         service.reserve()
         reserved = True
+        # Preflight the dataset: a missing/empty/uncaptionable data_dir otherwise fails inside the
+        # spawned trainer AFTER the user's model was evicted. Same discovery the trainer runs, so
+        # the two cannot disagree.
+        try:
+            await asyncio.to_thread(
+                _dtc.discover_image_caption_pairs,
+                config["data_dir"],
+                instance_prompt = config.get("instance_prompt") or None,
+                caption_column = config.get("caption_column") or "text",
+                # Decode-probe every image now (cheap PIL header check) so a corrupt/zero-byte
+                # upload 400s BEFORE _free_gpu_for_diffusion_training() tears down the user's
+                # models, rather than crashing the spawned trainer post-eviction.
+                verify_images = True,
+            )
+        except (FileNotFoundError, ValueError) as e:
+            raise HTTPException(status_code = 400, detail = str(e))
         # Free resident GPU workloads (export / Images pipeline / chat) before the trainer loads
         # its own pipeline. Offload the blocking teardown (engine unload waits on generation
         # locks; export subprocess join can take seconds) to a worker thread so the event loop

studio/backend/routes/video.py

@@ -272,7 +272,9 @@ async def get_gallery_video_file(
 ):
     from core.inference import video_gallery
 
-    path = await asyncio.to_thread(video_gallery.video_path, video_id)
+    # Ownership-gate the serve like delete/clear: resolve only a Studio-owned MP4 (readable
+    # sidecar), so a guessed stem for a foreign/orphan clip the listing hides can't be streamed out.
+    path = await asyncio.to_thread(video_gallery.owned_video_path, video_id)
     if path is None:
         raise HTTPException(status_code = 404, detail = "Video not found.")
     from fastapi.responses import FileResponse

studio/backend/tests/test_diffusion_backend.py

@@ -541,6 +541,40 @@ def test_generate_progress_cleared_on_setup_error(fake_runtime, tmp_path, monkey
     assert backend.generate_progress()["active"] is False
 
 
+def test_generate_progress_active_through_compile_cache_save(fake_runtime, tmp_path, monkeypatch):
+    # Post-denoise work (the compile-cache save) still runs before the route persists the image, so
+    # progress must stay active through it; clearing early would let a reload's mount probe read idle
+    # and refresh the gallery before the result exists.
+    from core.inference import diffusion as dmod
+
+    (tmp_path / "model.gguf").write_bytes(b"weights")
+    backend = DiffusionBackend()
+    backend.load_pipeline(
+        str(tmp_path),
+        gguf_filename = "model.gguf",
+        base_repo = "base/repo",
+        family_override = "z-image",
+        hf_token = "hf_secret",
+    )
+
+    seen = {}
+
+    def fake_save(ctx, *, logger = None):
+        seen["progress"] = backend.generate_progress()
+        return True
+
+    monkeypatch.setattr(dmod.compile_cache, "register_shape", lambda *a, **k: None)
+    monkeypatch.setattr(dmod.compile_cache, "save", fake_save)
+
+    gen = backend.generate(prompt = "a sloth", steps = 4)
+    assert len(gen["images"]) == 1
+    # Still active while the compile-cache save ran (after the denoise, before generate() returns).
+    assert seen["progress"]["active"] is True
+    assert seen["progress"]["total_steps"] == 4
+    # And cleared once the generation returns.
+    assert backend.generate_progress()["active"] is False
+
+
 def test_dense_speed_auto_defers_compile_to_third_generation(fake_runtime, tmp_path, monkeypatch):
     # Dense models with speed unset stay bit-identical eager for the first two generations; the
     # 3rd engages the `default` profile mid-session (repeated use amortises the one-time compile),

studio/backend/tests/test_diffusion_routes.py

@@ -111,6 +111,10 @@ class _FakeBackend:
             "repo_id": "x/z-image",
         }
 
+    def generate_progress(self):
+        # Idle by default; the persist-window override lives in the route, not here.
+        return {"active": False, "step": 0, "total_steps": 0, "fraction": 0.0, "eta_seconds": None}
+
     def unload(self):
         self.loaded = False
         return _unloaded_status()
@@ -197,6 +201,13 @@ def client(monkeypatch, tmp_path):
         "image_path",
         lambda i: (tmp_path / f"{i}.png") if i in store else None,
     )
+    # The serve route resolves through owned_image_path (ownership-gated); the fake store only ever
+    # holds owned records, so a stem not in it is treated as foreign and refused, like the real guard.
+    monkeypatch.setattr(
+        gallery_module,
+        "owned_image_path",
+        lambda i: (tmp_path / f"{i}.png") if i in store else None,
+    )
     monkeypatch.setattr(gallery_module, "delete", lambda i: store.pop(i, None) is not None)
     monkeypatch.setattr(gallery_module, "clear", _clear)
 
@@ -241,6 +252,49 @@ def test_load_generate_status_unload_roundtrip(client):
     assert client.get("/api/inference/images/status").json()["loaded"] is False
 
 
+def test_gallery_serve_refuses_unowned_id(client):
+    # The serve route resolves through the ownership guard, so a guessed stem for a PNG the gallery
+    # does not own (no record) is a 404, not a stream of foreign bytes.
+    assert client.get("/api/inference/images/gallery/family-photo/file").status_code == 404
+
+
+def test_generate_holds_progress_active_during_persist(client, monkeypatch):
+    # generate-progress must stay active while a finished generation is still writing its gallery
+    # record, so a concurrent reload's mount probe keeps polling instead of refreshing the gallery
+    # before the image lands. Probe the persist counter from inside the save call, and confirm it
+    # is cleared afterwards.
+    import core.inference.image_gallery as gallery_module
+    import routes.inference as inf
+
+    client.post(
+        "/api/inference/images/load",
+        json = {
+            "model_path": "unsloth/Z-Image-Turbo-GGUF",
+            "gguf_filename": "z-image-turbo-Q4_K_S.gguf",
+            "base_repo": "unsloth/Z-Image-base",
+        },
+    )
+
+    # Idle before any generation.
+    assert client.get("/api/inference/images/generate-progress").json()["active"] is False
+
+    seen = {}
+    real_save = gallery_module.save
+
+    def _probe_save(image, meta):
+        seen["during"] = inf._diffusion_persist_active
+        return real_save(image, meta)
+
+    monkeypatch.setattr(gallery_module, "save", _probe_save)
+
+    gen = client.post("/api/inference/images/generate", json = {"prompt": "a sloth", "seed": 7})
+    assert gen.status_code == 200
+    # Active while the record was being persisted, and back to idle once the route returned.
+    assert seen["during"] >= 1
+    assert inf._diffusion_persist_active == 0
+    assert client.get("/api/inference/images/generate-progress").json()["active"] is False
+
+
 def test_load_rejects_untrusted_base_repo(client):
     # A trusted GGUF model_path paired with an untrusted remote base_repo is rejected at the
     # route (validate runs before the GPU handoff), so an authenticated client cannot make the

studio/backend/tests/test_diffusion_training.py

@@ -431,6 +431,49 @@ def test_route_start_reserves_before_freeing_gpu(client, monkeypatch):
     assert "unreserve" in client._fake.calls
 
 
+def test_route_start_reserves_before_scanning_dataset(client, monkeypatch):
+    # The dataset preflight scan decode-probes every image and can take noticeable time on a large
+    # folder. It must run AFTER the slot is reserved (is_active -> true), so a concurrent
+    # upload/caption/delete guarded by _require_diffusion_dataset_mutable() can't mutate the dataset
+    # the trainer is about to read during that window. Assert reserve precedes the scan and the
+    # service is active while the scan is in flight.
+    order: list = []
+
+    def _record_scan(data_dir, **kw):
+        client._fake.calls.append("scan")
+        order.append(f"scan_active={client._fake.is_active()}")
+        return [("img.png", "caption")]
+
+    monkeypatch.setattr(
+        "core.training.diffusion_train_common.discover_image_caption_pairs", _record_scan
+    )
+
+    r = client.post("/api/train/diffusion/start", json = _BODY)
+    assert r.status_code == 200, r.text
+    assert client._fake.calls.index("reserve") < client._fake.calls.index("scan")
+    assert order == ["scan_active=True"]
+    assert "unreserve" in client._fake.calls
+
+
+def test_route_start_unreserves_when_dataset_preflight_fails(client, monkeypatch):
+    # A dataset preflight failure AFTER the reservation must roll it back (unreserve), or a rejected
+    # start would leave training permanently "active" and keep blocking loads and dataset edits.
+    def _bad_scan(data_dir, **kw):
+        raise ValueError("no captioned images found")
+
+    monkeypatch.setattr(
+        "core.training.diffusion_train_common.discover_image_caption_pairs", _bad_scan
+    )
+
+    r = client.post("/api/train/diffusion/start", json = _BODY)
+    assert r.status_code == 400
+    assert "no captioned images" in r.json()["detail"]
+    # Reserved, then rolled back; never started, and no longer active.
+    assert "reserve" in client._fake.calls and "unreserve" in client._fake.calls
+    assert "start" not in client._fake.calls
+    assert client._fake.is_active() is False
+
+
 def test_service_reserve_marks_active_and_rolls_back():
     # The real service: reserve() flips is_active true before any proc exists (so a load guard
     # refuses during the free window), and unreserve() clears it without a live proc, so a failed

studio/backend/tests/test_image_gallery.py

@@ -128,6 +128,22 @@ def test_image_path_rejects_unsafe_ids():
     assert gallery.image_path("missing") is None
 
 
+def test_owned_image_path_serves_only_owned_pngs():
+    # A hand-dropped foreign PNG resolves via image_path (safe stem, on disk) but must NOT be
+    # served: owned_image_path applies the same recipe check as delete/clear, so the serve route
+    # can't stream a file the listing hides.
+    foreign = gallery.gallery_dir() / "family-photo.png"
+    _img().save(foreign, format = "PNG")
+    assert gallery.image_path("family-photo") is not None  # resolvable...
+    assert gallery.owned_image_path("family-photo") is None  # ...but not ours to serve
+
+    ours = gallery.save(_img(), _meta(prompt = "ours"))
+    assert gallery.owned_image_path(ours["id"]) is not None
+    # Unsafe / missing ids resolve to nothing, like image_path.
+    assert gallery.owned_image_path("../../etc/passwd") is None
+    assert gallery.owned_image_path("missing") is None
+
+
 def test_list_skips_foreign_pngs(tmp_path):
     # A PNG without our recipe chunk (user dropped a file) is ignored.
     foreign = gallery.gallery_dir() / "foreign.png"

studio/backend/tests/test_video_gallery.py

@@ -108,6 +108,30 @@ def test_video_path_returns_mp4_for_saved_id():
     assert path is not None and path.name == f"{record['id']}.mp4"
 
 
+def test_owned_video_path_serves_only_owned_clips():
+    # A hand-dropped orphan MP4 resolves via video_path (safe stem, on disk) but must NOT be
+    # served: owned_video_path applies the same sidecar check as delete/clear, so the serve and
+    # export routes can't stream/transcode a clip the listing hides.
+    orphan = gallery.gallery_dir() / "recording.mp4"
+    orphan.write_bytes(_mp4())
+    assert gallery.video_path("recording") is not None  # resolvable...
+    assert gallery.owned_video_path("recording") is None  # ...but not ours to serve
+
+    ours = gallery.save(_mp4(), _meta(prompt = "ours"))
+    assert gallery.owned_video_path(ours["id"]) is not None
+    assert gallery.owned_video_path("../../etc/passwd") is None
+    assert gallery.owned_video_path("missing") is None
+
+
+def test_transcode_refuses_orphan_mp4():
+    # Export starts from the same resolver as /file, so a guessed stem for an orphan MP4 (no
+    # readable sidecar) must not be re-encoded out either.
+    orphan = gallery.gallery_dir() / "recording.mp4"
+    orphan.write_bytes(_real_mp4_bytes())
+    assert gallery.transcode("recording", "gif") is None
+    assert gallery.transcode("recording", "webm") is None
+
+
 def test_delete_removes_both_files():
     record = gallery.save(_mp4(), _meta(prompt = "a"))
     gallery.save(_mp4(), _meta(prompt = "b"))

studio/backend/tests/test_video_routes.py

@@ -578,6 +578,15 @@ def test_file_endpoint_404_for_bad_id(client):
     assert resp.status_code == 404
 
 
+def test_serve_and_export_refuse_orphan_mp4(client, tmp_path):
+    # A hand-dropped orphan MP4 (no readable sidecar) is hidden by the listing; the serve and
+    # export routes resolve through the ownership guard, so a guessed stem can neither stream nor
+    # transcode it out.
+    (tmp_path / "recording.mp4").write_bytes(b"\x00\x00\x00\x18ftypmp42")
+    assert client.get("/api/inference/video/gallery/recording/file").status_code == 404
+    assert client.get("/api/inference/video/gallery/recording/export?format=gif").status_code == 404
+
+
 def test_delete_and_clear(client):
     client.post(
         "/api/inference/video/load",