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:
parent
ef6f18c1ba
commit
a274308330
2 changed files with 50 additions and 1 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 == []
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue