From be0bd0006473f4ecd346e9e0f4588d51b980ef35 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 13 Jul 2026 12:28:52 +0000 Subject: [PATCH] Studio: gate local-pipeline image tagging on a real family; validate video sidecars before delete/clear --- .../backend/core/inference/video_gallery.py | 27 ++++++++++++++++++- studio/backend/routes/models.py | 19 ++++++++++++- .../backend/tests/test_local_model_format.py | 23 +++++++++++++--- studio/backend/tests/test_video_gallery.py | 26 ++++++++++++++++++ 4 files changed, 89 insertions(+), 6 deletions(-) diff --git a/studio/backend/core/inference/video_gallery.py b/studio/backend/core/inference/video_gallery.py index 2116c1ca32..292d0a72b8 100644 --- a/studio/backend/core/inference/video_gallery.py +++ b/studio/backend/core/inference/video_gallery.py @@ -172,6 +172,26 @@ def _sidecar_path(video_id: str) -> Path: return gallery_dir() / f"{video_id}.json" +# The sidecar keys a genuine Studio record always carries (save() always writes them). delete() and +# clear() treat a pair as owned only when its sidecar has all of these, so a hand-dropped MP4 with a +# parseable-but-empty ("{}") or partial JSON sidecar -- which list_videos already hides via the +# GalleryVideo schema filter -- is neither counted as ours nor destroyed. Mirrors +# image_gallery._REQUIRED_META: a key-presence check (the route still owns full schema/value-type +# validation), aligned with GalleryVideo's required stored fields. +_REQUIRED_META = ( + "prompt", + "width", + "height", + "num_frames", + "fps", + "duration_s", + "steps", + "guidance", + "seed", + "created_at", +) + + def _read_meta(sidecar: Path) -> Optional[dict[str, Any]]: try: raw = sidecar.read_text(encoding = "utf-8") @@ -181,7 +201,12 @@ def _read_meta(sidecar: Path) -> Optional[dict[str, Any]]: meta = json.loads(raw) except (ValueError, TypeError): return None - return meta if isinstance(meta, dict) else None + # A parseable dict is not enough to claim ownership: a foreign sidecar (e.g. "{}") or one from a + # different schema lacks these keys. Require them so delete()/clear() never destroy a clip the + # gallery never surfaced (mirrors image_gallery._read_meta). + if not isinstance(meta, dict) or any(k not in meta for k in _REQUIRED_META): + return None + return meta def _mtime(path: Path) -> float: diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 82b018d741..d8226473ce 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -3332,7 +3332,24 @@ def _local_model_task(model: "LocalModelInfo") -> Optional[str]: return _VIDEO_GEN_TASK except Exception: pass - return "text-to-image" + # The Images load path resolves the family via detect_family_for_pick and REJECTS a pick + # whose id / name / checkpoint filename carries no supported image-family token + # (diffusion.py validate_load_request), 400ing AFTER it has already evicted the GPU owner. + # A bare model_index.json directory alone (a generically named on-device pipeline) is not + # enough. Tag text-to-image only when that same family detection succeeds, so the picker + # never advertises a local pipeline the load will always reject. Detection uses the same + # _local_family_needles the video branch does (leaf name / id / sole-file, not the raw + # path), so a family token in a parent directory can't spuriously tag it. + try: + from core.inference.diffusion_families import detect_family + for needle in _local_family_needles(model): + if detect_family(needle) is not None: + return "text-to-image" + return None + except Exception: + # Detection unavailable (import/exec error): fall back to the prior permissive tag + # rather than hiding a possibly-loadable pipeline. + return "text-to-image" return None diff --git a/studio/backend/tests/test_local_model_format.py b/studio/backend/tests/test_local_model_format.py index dcc948ab30..75b8f20161 100644 --- a/studio/backend/tests/test_local_model_format.py +++ b/studio/backend/tests/test_local_model_format.py @@ -175,13 +175,28 @@ def _local( ) -def test_local_task_tags_diffusers_pipeline_dir(tmp_path): - # A local diffusers pipeline (top-level model_index.json) is an image model even - # though its model_format is not "gguf": tag it so the Images picker keeps it. +def test_local_task_tags_family_named_pipeline_dir(tmp_path): + # A local diffusers pipeline (top-level model_index.json) whose id resolves to a supported + # image family loads fine, so tag it so the Images picker keeps it. + d = tmp_path / "flux-pipeline" + _touch(d / "model_index.json") + _touch(d / "unet" / "diffusion_pytorch_model.safetensors") + assert ( + models_route._local_model_task(_local(d, model_id = "black-forest-labs/FLUX.1-dev")) + == "text-to-image" + ) + + +def test_local_task_none_for_familyless_pipeline_dir(tmp_path): + # A generically named on-device pipeline (top-level model_index.json, no family token in its + # id / name / filename) is UNLOADABLE: the Images load path resolves no family via + # detect_family_for_pick and 400s after evicting the GPU owner. It must stay untagged so the + # picker never advertises a row that always fails; model_index.json alone is not enough. d = tmp_path / "my-local-pipeline" _touch(d / "model_index.json") _touch(d / "unet" / "diffusion_pytorch_model.safetensors") - assert models_route._local_model_task(_local(d)) == "text-to-image" + assert models_route._local_is_diffusers(_local(d)) is True + assert models_route._local_model_task(_local(d)) is None def test_local_task_tags_diffusers_by_family_id(tmp_path): diff --git a/studio/backend/tests/test_video_gallery.py b/studio/backend/tests/test_video_gallery.py index 1b140d00de..e6a3d6b361 100644 --- a/studio/backend/tests/test_video_gallery.py +++ b/studio/backend/tests/test_video_gallery.py @@ -34,6 +34,8 @@ def _meta(**over): "width": 1024, "height": 576, "num_frames": 49, + "fps": 24, + "duration_s": 2.0, "steps": 30, "guidance": 6.0, "seed": 7, @@ -218,6 +220,30 @@ def test_list_skips_corrupt_sidecar(): assert [r["prompt"] for r in listed] == ["ours"] +def test_clear_preserves_mp4_with_present_but_invalid_sidecar(): + # A hand-dropped MP4 whose sidecar PARSES as JSON but lacks the required recipe keys (e.g. "{}") + # is hidden by list_videos (fails the GalleryVideo schema filter), so clear must not destroy it + # while removing the owned pair. Regression for the sidecar-validation gap. + directory = gallery.gallery_dir() + (directory / "foreign.mp4").write_bytes(_mp4()) + (directory / "foreign.json").write_text("{}", encoding = "utf-8") + gallery.save(_mp4(), _meta(prompt = "ours")) + assert gallery.clear() == 1 + assert (directory / "foreign.mp4").exists() + + +def test_delete_refuses_mp4_with_present_but_invalid_sidecar(): + # A per-id delete of an MP4 whose sidecar parses but is missing required keys must refuse it: + # the gallery never surfaced it, so a guessed id must not destroy it. + directory = gallery.gallery_dir() + (directory / "foreign.mp4").write_bytes(_mp4()) + (directory / "foreign.json").write_text( + json.dumps({"prompt": "x"}), encoding = "utf-8" + ) # partial sidecar (no width/seed/...) + assert gallery.delete("foreign") is False + assert (directory / "foreign.mp4").exists() + + def test_valid_callback_paginates_over_accepted_records(): # ``valid`` must filter before pagination, so offset/limit/has_more count over accepted records; # else a leading bad record returns a short page with more remaining and stalls scroll.