scripts/fetch_i2v_model.py

"""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())