From dabb12e1987984f44b60ae650d76a5edcb7a2e1d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 13 Jul 2026 04:21:36 +0000 Subject: [PATCH] Studio: gate dataset uploads on the symlink check and surface local video single-file checkpoints --- studio/backend/routes/models.py | 14 ++++++++++++++ studio/backend/routes/training.py | 8 ++++++-- .../backend/tests/test_diffusion_dataset_api.py | 16 ++++++++++++++++ studio/backend/tests/test_local_model_format.py | 15 ++++++++++----- 4 files changed, 46 insertions(+), 7 deletions(-) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index cd45378110..71f4c2f41d 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -3320,6 +3320,20 @@ def _local_is_diffusers(model: "LocalModelInfo") -> bool: return True except Exception: pass + # A single-file VIDEO checkpoint (LTX / Wan / Hunyuan .safetensors, no model_index.json) + # ships no pipeline index and resolves to no IMAGE family either, so the checks above miss + # it. The video load route reinterprets a bare single-file local dir as a single_file load + # (routes/video.py), so it IS loadable and must be surfaced; without tagging it here + # _local_model_task returns task=null and the Video On-Device picker hides it. Match the same + # clean id / name needles (not the raw path) so a parent-dir family token can't spuriously + # match; _local_model_task then routes the video family to the text-to-video task. + try: + from core.inference.video_families import detect_video_family + for needle in (model.model_id, model.display_name, Path(model.id).name): + if needle and detect_video_family(needle) is not None: + return True + except Exception: + pass return False diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index cb36e43a45..3ce107a323 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -1613,11 +1613,15 @@ async def upload_diffusion_dataset( import os import tempfile - from utils.paths import datasets_root from utils.upload_limits import get_upload_limit_bytes, get_upload_limit_label cleaned = _clean_diffusion_dataset_name(name) - folder = datasets_root() / cleaned + # Reject a symlinked dataset directory BEFORE any write. A bare mkdir(exist_ok=True) succeeds + # through an existing name -> external-directory symlink, and the staged upload would then + # write/replace files outside the Studio datasets root through that link. The read/caption/ + # delete endpoints already enforce this containment via _resolve_dataset_folder; the upload + # path must run the same symlink + root-containment check first so writes can never escape. + folder = _resolve_dataset_folder(name, must_exist = False) folder.mkdir(parents = True, exist_ok = True) limit_bytes = get_upload_limit_bytes() diff --git a/studio/backend/tests/test_diffusion_dataset_api.py b/studio/backend/tests/test_diffusion_dataset_api.py index e80b29d6e6..aa59adef10 100644 --- a/studio/backend/tests/test_diffusion_dataset_api.py +++ b/studio/backend/tests/test_diffusion_dataset_api.py @@ -563,6 +563,22 @@ def test_resolve_dataset_folder_rejects_symlink(ds_root, tmp_path): assert exc.value.status_code == 400 +def test_upload_through_symlinked_dataset_cannot_escape_root(client, ds_root, tmp_path): + # End to end: an upload targeting a dataset name that already exists as a symlink to an + # external directory must be refused (400) BEFORE any bytes are written, so the upload can + # never create/replace files outside the datasets root through the link. + external = tmp_path / "external" + external.mkdir() + (ds_root / "linked").symlink_to(external, target_is_directory = True) + + r = _upload(client, "linked", [("intruder.png", _png_bytes())]) + assert r.status_code == 400 + assert "symbolic link" in r.json()["detail"] + # Nothing was written through the link into the external directory. + assert not (external / "intruder.png").exists() + assert not any(external.iterdir()) + + def test_delete_through_symlinked_dataset_cannot_escape_root(client, ds_root, tmp_path): # End to end: a DELETE against an image inside a symlinked dataset dir is refused (400) and the # external file it points at is NOT removed. diff --git a/studio/backend/tests/test_local_model_format.py b/studio/backend/tests/test_local_model_format.py index 3da98ceabb..7163db424b 100644 --- a/studio/backend/tests/test_local_model_format.py +++ b/studio/backend/tests/test_local_model_format.py @@ -180,13 +180,18 @@ def test_local_task_tags_video_pipeline_dir(tmp_path): ) -def test_local_task_video_name_without_pipeline_not_surfaced(tmp_path): - # A dir whose name matches a video family but which is NOT a diffusers pipeline (no - # model_index.json) is not a loadable pipeline, so it must stay untagged -- never surfaced - # to the Video picker, so it can never trigger a pipeline load that evicts then fails. +def test_local_task_tags_video_single_file_checkpoint(tmp_path): + # A dir whose name matches a video family holding a bare single-file .safetensors (no + # model_index.json) IS loadable: the video load route reinterprets a sole single-file local + # pick as a single_file load (routes/video.py), validating BEFORE it touches the GPU. So it + # must be tagged text-to-video and surfaced in the Video On-Device picker -- not left task=null + # and hidden, which would make the advertised-and-loadable checkpoint unusable. d = tmp_path / "ltx-loose" _touch(d / "ltx-2.safetensors") # loose weights, no model_index.json - assert models_route._local_model_task(_local(d, model_id = "Lightricks/LTX-2")) is None + assert ( + models_route._local_model_task(_local(d, model_id = "Lightricks/LTX-2")) + == models_route._VIDEO_GEN_TASK + ) def test_local_task_ignores_family_token_in_parent_path(tmp_path):