Check the cancel event between predownload files and probe local dirs for LTX-2.3

A warm-cache predownload sweep never consults the cancel event (each cached
file returns instantly), so an unload during it was ignored until a cold
file hit the network. The 2.3 detection also only probed bare-file local
repos; resolve directory repos through the same child resolver the loader
uses so their base pull is scoped too.
This commit is contained in:
Daniel Han 2026-07-05 05:25:03 +00:00
commit a274308330
2 changed files with 50 additions and 1 deletions

View file

@ -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

View file

@ -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 == []