Hugging Face's logo

Spaces:
YashKotak
/
devatlas-ai-gpu
Running on Zero

App Files Community
devatlas-ai-gpu / scripts /fetch_i2v_model.py
YashKotak's picture
YashKotak
Upload Python project
6ef0a95
Raw History Blame Contribute Delete
35 kB
"""Fetch the LOCAL image-to-video model snapshot used by the AI service.
Run this ONCE on a machine that has the disk and bandwidth:
python scripts/fetch_i2v_model.py # download / repair
python scripts/fetch_i2v_model.py --verify # check only, download nothing
python scripts/fetch_i2v_model.py --force # re-fetch anything suspect
It downloads `Lightricks/LTX-Video` (an open-weights, image-to-video diffusion
model) into `LOCAL_DIFFUSION_MODEL_DIR` (default `models/ltx-video`). After
this, video generation with `VIDEO_PROVIDER=local_diffusion` runs completely
offline: the service loads the snapshot with `local_files_only=True` and never
touches the network.
This is the ONLY place in the project that downloads model weights. There is
no third-party AI API anywhere in the generation path.
WHY THIS SCRIPT IS NOT A `snapshot_download` CALL
------------------------------------------------
`huggingface_hub.snapshot_download()` fetches each file over a SINGLE HTTP
stream. This host's link delivers ~0.3-1.9 MB/s per stream and drops
connections mid-transfer (observed: `IncompleteRead(7285660 bytes read, 517027676
more expected)`). The T5 text encoder is four ~4.8 GB shards, so a single
stream needs over an hour per shard, and the one that dies at 90% forces the
whole shard to restart. The retry budget then expires and the snapshot stays
permanently short - the exact "folder exists but is incomplete" failure, which
surfaced as:
local I2V model snapshot ... is incomplete; missing:
text_encoder/model-00003-of-00004.safetensors, ...
Measured on this link: 1 stream = 1.9 MB/s, 8 streams = 3.0 MB/s. So the fix is
to keep `snapshot_download` out of the transfer path and drive plain HTTP range
requests ourselves:
* the authoritative file list + byte sizes come from the Hub API, not from
guessing globs, so "complete" is defined by the publisher, not by us;
* each file is fetched as N concurrent byte ranges into a `.part` file, with
a `.parts.json` ledger recording which ranges already landed, so an
interrupted run resumes at byte granularity instead of restarting;
* every file is validated against the published size AND the published
sha256 (the blob ETag) before it is moved into place;
* a file that is already present and correct is never touched.
Size: roughly 27 GB on disk (fp32 weights). Peak RAM during inference is much
lower than that because the pipeline loads in bfloat16 and decodes the VAE in
chunks, but the host still needs ~14 GB free RAM. See README for guidance on
smaller local models.
"""
from __future__ import annotations
import hashlib
import json
import os
import shutil
import subprocess
import sys
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
# LTX-Video ships its weights in fp32; the pipeline casts to bfloat16/fp16 on
# load, so these are the files the snapshot must contain.
#
# ALLOW_PATTERNS is the DOWNLOAD filter.
#
# The sharded `*.safetensors.index.json` manifests must be listed EXPLICITLY.
# The obvious-looking glob `text_encoder/model-*.safetensors` does not match
# them: the manifest is `model.safetensors.index.json` - a `.` (not a `-`)
# after "model", and a `.json` (not `.safetensors`) suffix. Omitting the
# manifests downloads every multi-GB shard and still leaves the snapshot
# UNLOADABLE, because transformers cannot map weights onto shard files without
# the manifest. That is exactly the "folder exists but is incomplete" failure.
ALLOW_PATTERNS = [
"model_index.json",
"scheduler/scheduler_config.json",
"tokenizer/*",
"text_encoder/config.json",
"text_encoder/model*.safetensors",
"text_encoder/model.safetensors.index.json",
"transformer/config.json",
"transformer/diffusion_pytorch_model*.safetensors",
"transformer/diffusion_pytorch_model.safetensors.index.json",
"vae/config.json",
"vae/diffusion_pytorch_model.safetensors",
]
# The repo also publishes ~150 GB of single-file checkpoints at its root
# (ltxv-13b-*, ltx-video-2b-*, spatial/temporal upscalers) plus demo media.
# None of them are part of the diffusers pipeline. `allow_patterns` already
# excludes them; listing them in `ignore_patterns` as well keeps them excluded
# even if the allow filter above is ever loosened.
IGNORE_PATTERNS = [
"media/*",
"*.md",
"*.txt",
"*.gif",
"*.png",
"*.mp4",
"ltxv-*.safetensors",
"ltx-video-*.safetensors",
"*.pth",
"*.ckpt",
]
# Exact files that must exist AND be non-empty for the pipeline to load.
# (Globs drive the download; verification needs exact names.)
REQUIRED_FILES = [
"model_index.json",
"scheduler/scheduler_config.json",
"tokenizer/tokenizer_config.json",
"tokenizer/spiece.model",
"text_encoder/config.json",
"text_encoder/model.safetensors.index.json",
"transformer/config.json",
"transformer/diffusion_pytorch_model.safetensors.index.json",
"vae/config.json",
"vae/diffusion_pytorch_model.safetensors",
]
# Manifests that enumerate the shards of a sharded checkpoint. Every shard a
# manifest names must be present and non-empty.
SHARD_MANIFESTS = [
"text_encoder/model.safetensors.index.json",
"transformer/diffusion_pytorch_model.safetensors.index.json",
]
# Target size of one range request, and how many run at once.
#
# Chunks are deliberately small: the resume ledger records progress per chunk,
# so a 32 MB chunk means an interrupted run loses at most 32 MB of work instead
# of restarting a multi-GB shard. 8 concurrent chunks measured ~3.0 MB/s on
# this link versus ~1.9 MB/s for a single stream, and each is retried
# independently, so one dropped connection costs a chunk rather than a shard.
RANGE_CHUNK_BYTES = int(os.environ.get("I2V_RANGE_CHUNK_BYTES", str(32 * 1024 * 1024)))
RANGES_IN_FLIGHT = int(os.environ.get("I2V_RANGES_IN_FLIGHT", "8"))
# Read timeout for a single range request. Kept short so a wedged socket is
# retried quickly instead of blocking a worker slot until the OS gives up.
RANGE_TIMEOUT = float(os.environ.get("I2V_RANGE_TIMEOUT", "30"))
# Per-chunk retry budget. A retry re-requests only the bytes still missing, so
# a long flaky transfer still converges.
RANGE_ATTEMPTS = int(os.environ.get("I2V_RANGE_ATTEMPTS", "40"))
# Live progress, in bytes, shared with the stall watchdog.
# The watchdog MUST use this rather than the size of the files on disk: the
# `.part` scratch file is pre-allocated to the full published size with
# `truncate()`, so its on-disk length is constant from the first second and a
# directory-size watchdog would declare a perfectly healthy transfer wedged.
_PROGRESS = [0]
def _wanted(
rel: str,
allow: list | None = None,
ignore: list | None = None,
) -> bool:
"""True when `rel` is part of the diffusers pipeline snapshot.
The allow/ignore lists are parameters (defaulting to the video model's) so
`scripts/fetch_image_model.py` can reuse this exact downloader for a
different repository instead of forking a second, flakier copy.
"""
from fnmatch import fnmatch
ignore = IGNORE_PATTERNS if ignore is None else ignore
allow = ALLOW_PATTERNS if allow is None else allow
if any(fnmatch(rel, pat) for pat in ignore):
return False
return any(fnmatch(rel, pat) for pat in allow)
def fetch_manifest(
repo: str,
allow: list | None = None,
ignore: list | None = None,
) -> dict:
"""Ask the Hub for the authoritative file list, sizes, sha256s and commit sha.
Returns `{"repo", "revision", "files": {rel: size}, "oids": {rel: sha256}}`.
The revision is returned because the whole snapshot must come from ONE
commit: mixing files from two revisions of a diffusers repo produces a
folder that passes every structural check and still fails to load.
The sha256 of every LFS blob is captured HERE, at the pinned revision, and
carried in the manifest as the single source of truth for integrity. It
must not be re-read from the CDN's `ETag` at download time: the redirect
that serves the bytes lands on an edge node whose `ETag` is NOT the
content hash, so a correct download is rejected and thrown away. That bug
discarded ~2.6 GB of perfectly good weights and made this snapshot look
permanently unfetchable.
"""
from huggingface_hub import HfApi
info = HfApi().model_info(repo, files_metadata=True)
files = {}
oids = {}
for sibling in info.siblings or []:
rel = sibling.rfilename
if not _wanted(rel, allow, ignore):
continue
size = getattr(sibling, "size", None)
if size is None:
raise RuntimeError(f"Hub did not report a size for {repo}:{rel}")
files[rel] = int(size)
# `lfs.oid` (or the sibling's own sha for a non-LFS blob) is the
# content-addressed sha256 the Hub published for this exact revision.
blob_sha = getattr(sibling, "lfs", None)
oid = getattr(blob_sha, "sha256", None) if blob_sha is not None else None
if oid is None:
oid = getattr(sibling, "blob_id", None)
if isinstance(oid, str) and len(oid) == 64:
oids[rel] = oid
if not files:
raise RuntimeError(
f"{repo} returned no files matching the snapshot allow-list"
)
return {"repo": repo, "revision": info.sha, "files": files, "oids": oids}
def resolve_target() -> Path:
"""Absolute snapshot directory, from settings, resolved against this file."""
from app.config.settings import get_settings
settings = get_settings()
target = Path(settings.local_diffusion_model_dir)
if not target.is_absolute():
target = Path.cwd() / target
return target
def resolve_repo() -> str:
from app.config.settings import get_settings
return get_settings().local_diffusion_model_id or "Lightricks/LTX-Video"
def verify_snapshot(
root: Path,
manifest: dict | None = None,
required: list | None = None,
shard_manifests: list | None = None,
) -> list:
"""Return every problem that would make this snapshot unloadable.
A completed transfer is NOT proof of a usable snapshot, and neither is a
full-looking directory: a truncated shard is still a file, and
`from_pretrained` then fails deep inside safetensors with a confusing
error. This is the honest check that runs before the script claims success.
Checks, in increasing cost:
1. the small config/manifest files are present and non-empty,
2. every shard a manifest promises is present and non-empty,
3. every file the Hub published is present, and when `manifest` is
supplied its on-disk size matches the published size exactly.
Level 3 is what catches a half-written shard that level 2 waves through.
"""
problems: list = []
required = REQUIRED_FILES if required is None else required
shard_manifests = SHARD_MANIFESTS if shard_manifests is None else shard_manifests
for rel in required:
path = root / rel
if not path.is_file():
problems.append(f"missing: {rel}")
elif path.stat().st_size == 0:
problems.append(f"empty (0 bytes): {rel}")
for manifest_rel in shard_manifests:
manifest_path = root / manifest_rel
if not manifest_path.is_file():
continue # already reported by the REQUIRED_FILES pass
try:
manifest_doc = json.loads(manifest_path.read_text(encoding="utf-8"))
except (OSError, ValueError) as exc:
problems.append(f"unreadable manifest: {manifest_rel} ({exc})")
continue
shards = sorted(set((manifest_doc.get("weight_map") or {}).values()))
if not shards:
problems.append(f"manifest lists no shards: {manifest_rel}")
folder = manifest_rel.split("/")[0]
for shard in shards:
shard_path = manifest_path.parent / shard
if not shard_path.is_file():
problems.append(f"missing shard: {folder}/{shard}")
elif shard_path.stat().st_size == 0:
problems.append(f"empty shard (0 bytes): {folder}/{shard}")
# Authoritative size check against the Hub's published file list.
for rel, expected in sorted((manifest or {}).get("files", {}).items()):
path = root / rel
if not path.is_file():
problems.append(f"missing: {rel}")
continue
actual = path.stat().st_size
if actual != expected:
problems.append(
f"wrong size: {rel} ({actual} bytes on disk, "
f"{expected} expected)"
)
return problems
# ---------------------------------------------------------------------------
# Resumable parallel-range transfer
# ---------------------------------------------------------------------------
#
# State on disk, per file being fetched:
#
# <name>.part pre-allocated to the published size; ranges are written
# into it at their true offsets
# <name>.parts.json {"repo", "revision", "size", "done": [[start, end], ..]}
# the ranges that already landed, merged
#
# A range is only recorded as done after its bytes are written, so the ledger
# can under-report progress after a crash but never over-report it: re-fetching
# a range is harmless, skipping one that was never written would not be.
def _file_url(repo: str, revision: str, rel: str) -> str:
from urllib.parse import quote
return f"https://huggingface.co/{repo}/resolve/{revision}/{quote(rel)}"
def _merge_ranges(ranges: list, size: int) -> list:
"""Merge, sort and clamp `[start, end)` ranges."""
out: list[list[int]] = []
for start, end in sorted((int(a), int(b)) for a, b in ranges if b > a):
start = max(0, min(start, size))
end = max(0, min(end, size))
if end <= start:
continue
if out and start <= out[-1][1]:
out[-1][1] = max(out[-1][1], end)
else:
out.append([start, end])
return out
def _gaps(done: list, size: int) -> list:
"""Byte ranges of `size` not covered by `done`."""
gaps: list[list[int]] = []
cursor = 0
for start, end in done:
if start > cursor:
gaps.append([cursor, start])
cursor = max(cursor, end)
if cursor < size:
gaps.append([cursor, size])
return gaps
def _split(gap: list, parts: int) -> list:
"""Split one gap into at most `parts` equal ranges (used by tests)."""
start, end = gap
span = end - start
if parts <= 1 or span <= 1:
return [[start, end]]
parts = min(parts, span)
step = span // parts
out: list[list[int]] = []
cursor = start
for _ in range(parts):
stop = min(cursor + step, end)
out.append([cursor, stop])
cursor = stop
if cursor < end:
out.append([cursor, end])
return out
def _covered(done: list) -> int:
return sum(int(b) - int(a) for a, b in done)
def _load_ledger(path: Path, repo: str, revision: str, size: int) -> list:
"""Ranges already fetched for this file, or [] to start over.
A ledger written for a different repo, revision or expected size describes
a different file, so it is discarded rather than trusted.
"""
try:
state = json.loads(path.read_text(encoding="utf-8"))
except (OSError, ValueError):
return []
if (
state.get("repo") != repo
or state.get("revision") != revision
or int(state.get("size", -1)) != size
):
return []
try:
return _merge_ranges([tuple(r) for r in state.get("done", [])], size)
except (TypeError, ValueError):
return []
def _save_ledger(path: Path, repo: str, revision: str, size: int, done: list) -> None:
"""Write the ledger atomically so a crash cannot leave it corrupt."""
payload = {
"repo": repo,
"revision": revision,
"size": size,
"done": [[int(a), int(b)] for a, b in done],
}
tmp = path.with_name(path.name + ".tmp")
tmp.write_text(json.dumps(payload), encoding="utf-8")
os.replace(tmp, path)
def _expected_sha256(session, url: str) -> str | None:
"""DEPRECATED fallback only - do not use this to verify a download.
The Hub's CDN redirect does not end on a node whose `ETag` is the content
hash: for these blobs the header carries an unrelated digest, so a byte-
perfect download is reported as corrupt and DISCARDED. `fetch_manifest`
captures the authoritative `lfs.oid` at the pinned revision instead, and
that is what the integrity gate compares against.
Kept only so the video script's own callers keep importing successfully.
"""
try:
resp = session.head(url, allow_redirects=True, timeout=RANGE_TIMEOUT)
etag = (resp.headers.get("ETag") or "").strip().strip('"')
resp.close()
except Exception:
return None
if len(etag) == 64 and all(c in "0123456789abcdef" for c in etag):
return etag
return None
def _sha256_file(path: Path, cancel: threading.Event) -> str | None:
"""sha256 of a file, or None if the run was cancelled mid-hash."""
digest = hashlib.sha256()
with path.open("rb") as handle:
while True:
if cancel.is_set():
return None
block = handle.read(8 * 1024 * 1024)
if not block:
break
digest.update(block)
return digest.hexdigest()
def _fetch_range(session, url, handle, start, end, total, cancel) -> int:
"""Stream `[start, end)` into `handle` at its true offset.
Returns the number of bytes written. A short read is NOT an error: the
socket simply died mid-stream, and the caller resumes this chunk from
where it stopped instead of restarting the whole multi-GB shard.
"""
handle.seek(start)
written = 0
resp = session.get(
url,
headers={"Range": f"bytes={start}-{end - 1}"},
stream=True,
timeout=(RANGE_TIMEOUT, RANGE_TIMEOUT),
)
try:
if resp.status_code == 200 and (start != 0 or end != total):
# Server ignored the Range header and streamed the whole blob.
raise RuntimeError("server ignored the Range header")
if resp.status_code not in (200, 206):
raise RuntimeError(f"HTTP {resp.status_code}")
for block in resp.iter_content(1024 * 1024):
if cancel.is_set():
raise RuntimeError("cancelled")
if not block:
continue
handle.write(block)
written += len(block)
_PROGRESS[0] += len(block)
finally:
resp.close()
return written
def _chunk_plan(done: list, size: int) -> list:
"""Missing bytes as fixed-size chunks, in order.
Fixed-size chunks (rather than "split the file into N ranges") are what make
the ledger meaningful: progress is recorded per chunk, so an interrupted
run resumes from the last completed chunk boundary instead of redoing the
whole file.
"""
plan: list[list[int]] = []
for gap_start, gap_end in _gaps(done, size):
cursor = gap_start
while cursor < gap_end:
stop = min(cursor + RANGE_CHUNK_BYTES, gap_end)
plan.append([cursor, stop])
cursor = stop
return plan
def _download_file(repo, revision, rel, dest: Path, size: int, force=False, expect_sha256: str | None = None) -> str:
"""Fetch one file, resuming whatever a previous run left behind.
Returns "cached", "downloaded", "repaired" or "cancelled". Raises
RuntimeError when the finished bytes do not match the published
size/sha256, so a corrupt file is never moved into place and later loaded
as if it were real.
`expect_sha256` is the Hub's content-addressed hash for this exact
revision (from the manifest), NOT the CDN `ETag`: the two disagree on
these blobs, and trusting the `ETag` rejects - and deletes - transfers
that are in fact byte-perfect.
"""
import requests
if dest.is_file() and dest.stat().st_size == size and not force:
return "cached"
part = dest.with_name(dest.name + ".part")
ledger = dest.with_name(dest.name + ".parts.json")
url = _file_url(repo, revision, rel)
dest.parent.mkdir(parents=True, exist_ok=True)
# A final file of the WRONG size is a truncated transfer, not part of the
# snapshot. Move it aside rather than deleting it, so a failed repair never
# leaves the user with less than they started with.
had_dest = dest.is_file()
if had_dest and dest.stat().st_size != size:
dest.replace(dest.with_name(dest.name + ".bad"))
done = _load_ledger(ledger, repo, revision, size)
if not part.exists() or part.stat().st_size != size:
# A ledger without a correctly sized scratch file cannot be trusted.
done = []
with part.open("wb") as handle:
handle.truncate(size)
print(
f" {rel}: {size / 1e9:.2f} GB "
f"({_covered(done) / 1e9:.2f} GB already fetched)",
flush=True,
)
cancel = threading.Event()
lock = threading.Lock()
before = _PROGRESS[0]
session = requests.Session()
session.headers["User-Agent"] = "devatlas-fetch-i2v/1.0"
plan = _chunk_plan(done, size)
if plan:
def _one(chunk):
start, end = chunk
for attempt in range(1, RANGE_ATTEMPTS + 1):
if cancel.is_set():
return
try:
with part.open("r+b") as handle:
wrote = _fetch_range(
session, url, handle, start, end, size, cancel
)
if wrote == end - start:
with lock:
done[:] = _merge_ranges(done + [[start, end]], size)
_save_ledger(ledger, repo, revision, size, list(done))
return
start += wrote # resume this chunk, never restart the file
except Exception as exc:
if attempt >= RANGE_ATTEMPTS:
cancel.set()
print(
f" chunk {start}-{end} of {rel} failed after "
f"{RANGE_ATTEMPTS} attempts: {exc}",
flush=True,
)
return
time.sleep(min(2 ** (attempt // 4), 15))
# A bounded pool keeps every socket busy without spawning one thread
# per 32 MB chunk (a 4.8 GB shard would otherwise need ~150 threads).
with ThreadPoolExecutor(max_workers=RANGES_IN_FLIGHT) as pool:
for fut in [pool.submit(_one, c) for c in plan]:
fut.result()
gained = _PROGRESS[0] - before
if gained:
print(f" +{gained / 1e9:.2f} GB this run", flush=True)
if cancel.is_set() or _covered(done) != size:
return "cancelled"
# Integrity gate. A correct size alone cannot catch a corrupted file, so
# verify against the sha256 the Hub published for this exact revision.
#
# The expected value comes from the manifest (`fetch_manifest`), NOT from a
# fresh `ETag` lookup: the CDN edge node that actually serves these bytes
# reports an unrelated digest in `ETag`, so the old check rejected and
# DELETED every correct download. That is why this snapshot could never be
# completed, which in turn is why image generation silently fell back to
# re-showing the companion's avatar instead of generating a new picture.
want = expect_sha256
if want:
got = _sha256_file(part, cancel)
if got is None:
return "cancelled"
if got != want:
# Discard the bad scratch file so the next run starts this file
# clean rather than trusting the corrupt ranges forever.
for victim in (part, ledger):
try:
victim.unlink()
except OSError:
pass
raise RuntimeError(
f"{rel}: sha256 mismatch (expected {want}, got {got}); "
"corrupt data discarded - re-run to retry this file"
)
os.replace(part, dest)
try:
ledger.unlink()
except OSError:
pass
return "repaired" if had_dest else "downloaded"
def _clean_partial_downloads(root: Path, keep_parts: bool = False) -> int:
"""Delete partial download leftovers, returning bytes reclaimed.
Older runs used `huggingface_hub`, which leaves partial chunks behind as
`.<repo filename>.<chunk index>` (e.g.
`.p_text_encoder_model-00002-of-00004.safetensors.0`) and resumed transfers
as `<filename>.incomplete`. Those are not part of the snapshot and can be
tens of GB, so a folder that "looks complete" on disk can still be missing
real shards.
`keep_parts=True` preserves this script's OWN `<name>.part` /
`<name>.parts.json` resume state, which is the whole point of the range
ledger: deleting it would throw away gigabytes of already-fetched bytes and
make a long flaky transfer impossible to finish.
"""
reclaimed = 0
for path in sorted(root.rglob("*")):
if not path.is_file():
continue
name = path.name
is_chunk = name.startswith(".") and ".safetensors." in name
is_incomplete = name.endswith(".incomplete")
is_xet_chunk = name.endswith(".incomplete") and "=" in name
if keep_parts and (name.endswith(".part") or name.endswith(".parts.json")):
continue
if not (is_chunk or is_incomplete or is_xet_chunk):
continue
try:
reclaimed += path.stat().st_size
path.unlink()
except OSError:
continue
return reclaimed
def _free_bytes(path: Path) -> int:
"""Free space on the volume holding `path`, 0 when it cannot be read."""
try:
usage = shutil.disk_usage(str(path))
except OSError:
return 0
return int(usage.free)
def _check_disk_space(target: Path, manifest: dict) -> list:
"""Report files that will not fit, before starting a multi-GB transfer.
A download that dies at 99% because the volume filled up is exactly the
kind of failure that leaves a permanently incomplete snapshot, so this is
checked up front rather than discovered at the last shard.
"""
needed = 0
for rel, size in manifest["files"].items():
dest = target / rel
part = dest.with_name(dest.name + ".part")
for candidate in (dest, part):
if candidate.is_file():
needed -= candidate.stat().st_size
needed += size
free = _free_bytes(target)
if needed > free:
return [
f"not enough free disk space: need {needed / 1e9:.1f} GB, "
f"have {free / 1e9:.1f} GB on the volume holding {target}"
]
return []
def download_snapshot(manifest: dict, target: Path, force=False) -> list:
"""Fetch every file in `manifest` into `target`. Returns error strings.
Files already present at the published size are skipped, so this repairs
an incomplete snapshot without re-downloading the ~20 GB that is already
correct. Larger files go first: they dominate the wall clock, so starting
them early means a cancelled run still banks the most progress.
"""
repo, revision = manifest["repo"], manifest["revision"]
files = manifest["files"]
# Authoritative per-file sha256 captured by `fetch_manifest` at this exact
# revision. Passed to every download so the integrity gate never has to
# guess a hash from a CDN response header.
oids = manifest.get("oids") or {}
errors: list = []
order = sorted(files.items(), key=lambda kv: kv[1], reverse=True)
for index, (rel, size) in enumerate(order, start=1):
dest = target / rel
label = rel if len(rel) < 58 else "..." + rel[-55:]
print(f"[{index}/{len(order)}] {label}", flush=True)
try:
status = _download_file(
repo, revision, rel, dest, size,
force=force, expect_sha256=oids.get(rel),
)
except Exception as exc:
errors.append(f"{rel}: {exc}")
print(f" FAILED: {exc}", flush=True)
continue
if status == "cancelled":
errors.append(f"{rel}: transfer did not complete")
elif status != "cached":
print(f" {status}", flush=True)
return errors
def _worker(repo: str, target: str, stall_seconds: int) -> int:
"""Run one download attempt in a CHILD process, with a stall watchdog.
Invoked as `python scripts/fetch_i2v_model.py --_worker <repo> <target>
<stall-seconds>`. Each attempt runs out-of-process so the parent can kill a
wedged transfer outright: a dead socket raises nothing, it just stops making
progress while holding its connection open, so the only reliable detector is
watching the target directory stop growing from a separate THREAD.
"""
root = Path(target)
root.mkdir(parents=True, exist_ok=True)
stall = {"last": time.time(), "bytes": -1}
def watch() -> None:
# Progress is measured from the transfer's own byte counter, NOT from
# the size of the files in `root`. The `.part` scratch file is
# pre-allocated to the full published size, so its length is constant
# from the first second - a directory-size watchdog would kill a
# perfectly healthy download and restart it forever.
while True:
time.sleep(10)
done = _PROGRESS[0]
if done != stall["bytes"]:
stall["bytes"] = done
stall["last"] = time.time()
print(f" {done / 1e9:.2f} GB transferred", flush=True)
elif time.time() - stall["last"] > stall_seconds:
print(
f" no transfer progress for {stall_seconds}s - wedged",
flush=True,
)
os._exit(3)
threading.Thread(target=watch, daemon=True).start()
try:
manifest = fetch_manifest(repo)
except Exception as exc:
print(f" could not read the Hub manifest: {exc}", flush=True)
return 2
# A previous huggingface_hub run may have left unusable partials. This
# script's own `.part` / `.parts.json` resume state is preserved.
reclaimed = _clean_partial_downloads(root, keep_parts=True)
if reclaimed:
print(f" reclaimed {reclaimed / 1e9:.2f} GB of stale partials", flush=True)
errors = download_snapshot(manifest, root)
if errors:
print(f" {len(errors)} file(s) did not complete", flush=True)
return 2
return 0
def main() -> int:
if len(sys.argv) > 1 and sys.argv[1] == "--_worker":
# The child indexes four values out of argv, so a truncated invocation
# must fail with a readable message here. Bare indexing raises
# IndexError, and the parent - which can only read a child exit code -
# counts that as "attempt failed" and burns every attempt without a
# single byte being downloaded.
if len(sys.argv) != 5:
print(
"usage: fetch_i2v_model.py --_worker <repo> <target> <stall-seconds>",
file=sys.stderr,
)
return 64
return _worker(sys.argv[2], sys.argv[3], int(sys.argv[4]))
verify_only = "--verify" in sys.argv
force = "--force" in sys.argv
repo = resolve_repo()
target = resolve_target()
target.mkdir(parents=True, exist_ok=True)
print(f"Fetching {repo} -> {target}")
try:
manifest = fetch_manifest(repo)
except Exception as exc:
print(f"\nCould not reach the Hugging Face Hub: {exc}")
print("Check the network/proxy, then re-run.")
return 1
print(
f"Revision {manifest['revision']} - {len(manifest['files'])} files, "
f"{sum(manifest['files'].values()) / 1e9:.1f} GB total."
)
print("Re-running is safe: complete files are skipped, partial ones resume.\n")
space_problems = _check_disk_space(target, manifest)
if space_problems:
for problem in space_problems:
print(f" - {problem}")
return 1
if verify_only:
problems = verify_snapshot(target, manifest)
if problems:
print("Snapshot is INCOMPLETE:")
for problem in problems:
print(f" - {problem}")
return 1
print(f"Snapshot is complete and verified at: {target}")
return 0
# Each attempt is a child process so a wedged transfer can be killed
# outright. Retried because a ~27 GB transfer on a slow link is the single
# most likely thing here to drop; every attempt resumes from the byte ranges
# already on disk, so no progress is ever lost.
attempts = 6
stall_seconds = 300
script = os.path.abspath(__file__)
failed = False
for attempt in range(1, attempts + 1):
print(f"\nDownload attempt {attempt}/{attempts}...", flush=True)
proc = subprocess.Popen(
[sys.executable, script, "--_worker", repo, str(target),
str(stall_seconds)],
cwd=str(target.parents[1]) if len(target.parents) > 1 else None,
)
code = proc.wait()
if code == 0:
break
if code == 3:
print("\nAttempt: transfer stalled (no disk progress).", flush=True)
else:
print(f"\nAttempt failed (exit {code}).", flush=True)
if attempt == attempts:
print(f"\nDownload failed after {attempts} attempts.")
print("Re-run this script to resume. Nothing already fetched is lost.")
failed = True
break
print("Retrying in 5s...", flush=True)
time.sleep(5)
if failed:
return 1
# A completed transfer is not the same as a loadable snapshot. Verify
# against the published file list and the shard manifests, and fail loudly
# rather than leaving a half-usable snapshot behind.
problems = verify_snapshot(target, manifest)
if problems:
print("\nSnapshot is STILL INCOMPLETE:")
for problem in problems:
print(f" - {problem}")
print("\nRe-run this script to finish the download.")
return 1
print(f"\nDone. Model snapshot is complete and verified at: {target}")
print("Set VIDEO_PROVIDER=local_diffusion in python-ai/.env to use it.")
return 0
if __name__ == "__main__":
raise SystemExit(main())