diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 073feb6305..5e33d78984 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/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]: diff --git a/studio/backend/core/inference/image_gallery.py b/studio/backend/core/inference/image_gallery.py index 77c4b88e58..131c1eafb1 100644 --- a/studio/backend/core/inference/image_gallery.py +++ b/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 diff --git a/studio/backend/core/inference/video_gallery.py b/studio/backend/core/inference/video_gallery.py index 1cddb94b7b..860912fb3d 100644 --- a/studio/backend/core/inference/video_gallery.py +++ b/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 diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 9807f2af60..7a01cb3c76 100644 --- a/studio/backend/routes/inference.py +++ b/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) # ────────────────────────────────────────────────────────────────────────── diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 1173fb21b0..ccc242bf9a 100644 --- a/studio/backend/routes/training.py +++ b/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 diff --git a/studio/backend/routes/video.py b/studio/backend/routes/video.py index b1ab370a45..1f1fb65e97 100644 --- a/studio/backend/routes/video.py +++ b/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 diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index a3c7fb4bb1..0fa63d3f76 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/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), diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index 322c188534..0b225f0376 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/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 diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index 22003af1ee..e015a2d661 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/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 diff --git a/studio/backend/tests/test_image_gallery.py b/studio/backend/tests/test_image_gallery.py index a87bd6324a..466398cf13 100644 --- a/studio/backend/tests/test_image_gallery.py +++ b/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" diff --git a/studio/backend/tests/test_video_gallery.py b/studio/backend/tests/test_video_gallery.py index fdc6f93f02..d072dd44da 100644 --- a/studio/backend/tests/test_video_gallery.py +++ b/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")) diff --git a/studio/backend/tests/test_video_routes.py b/studio/backend/tests/test_video_routes.py index f7bc076225..c127bc0cbe 100644 --- a/studio/backend/tests/test_video_routes.py +++ 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",