diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index 4bd7751f7c..3b64e37dca 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -331,8 +331,21 @@ class VideoBackend: probe = checkpoint_local if probe is None: + # Local repos: a bare file, or a directory whose child the same + # resolver load_pipeline uses picks out. Unresolvable here means + # load_pipeline will surface the real error; keep the wide pull. root = Path(kwargs["repo_id"]).expanduser() - probe = root if root.is_file() else None + if root.is_file(): + probe = root + elif root.is_dir(): + try: + probe = self._resolve_checkpoint_path( + kwargs["repo_id"], + kwargs.get("gguf_filename"), + kwargs.get("hf_token"), + ) + except Exception: # noqa: BLE001 -- surfaced by load_pipeline + probe = None ltx23 = probe is not None and is_ltx23_checkpoint(probe) if ltx23: expected = self._estimate_download_bytes( @@ -463,6 +476,11 @@ class VideoBackend: snapshot_root: Optional[Path] = None for name, _ in files: + # Explicit per-file check: a fully-cached file returns without ever + # consulting the event, so a warm-cache sweep would otherwise run to + # completion after an unload already cancelled this load. + if self._cancel_event.is_set(): + raise RuntimeError(VIDEO_CANCELLED_MSG) local = Path( hf_hub_download_with_xet_fallback( base, name, hf_token, cancel_event = self._cancel_event diff --git a/studio/backend/tests/test_video_backend.py b/studio/backend/tests/test_video_backend.py index c34cee5bf9..209983e4f6 100644 --- a/studio/backend/tests/test_video_backend.py +++ b/studio/backend/tests/test_video_backend.py @@ -528,3 +528,34 @@ def test_base_download_files_ltx23_keeps_only_shared_components(): assert not any( n.startswith(("vae/", "connectors/", "latent_upsampler/", "transformer/")) for n in names ) + + +def test_predownload_base_honors_cancel_between_files(monkeypatch): + # A warm-cache sweep returns each file instantly without consulting the event, + # so the loop must check it explicitly or an unload mid-predownload is ignored. + backend = VideoBackend() + backend._cancel_event.set() + calls: list = [] + monkeypatch.setattr( + "utils.hf_xet_fallback.hf_hub_download_with_xet_fallback", + lambda repo, fn, tok, **kw: (calls.append(fn), f"/cache/{fn}")[1], + ) + + class _Api: + def __init__(self, token = None): + pass + + def model_info(self, repo, files_metadata = True): + return types.SimpleNamespace( + siblings = [ + _sibling("model_index.json", 1), + _sibling("vae/diffusion_pytorch_model.safetensors", 2), + ] + ) + + import huggingface_hub + + monkeypatch.setattr(huggingface_hub, "HfApi", _Api) + with pytest.raises(RuntimeError, match = "cancelled"): + backend._predownload_base("base/repo", None, "pipeline") + assert calls == []