Request: split Q2 GGUF into files smaller than 50 GB

#5
by diegoan - opened

Would it be possible to publish DeepSeek-V4.1-Flash-Q2.gguf as multiple smaller binary parts?

Downloading the current 366 GB single file repeatedly fails during Xet reconstruction, while smaller files from the same repository download successfully.

Splitting the Q2 GGUF into parts, preferably smaller than 50 GB, would make downloads easier to resume and more reliable across different environments. The download_model.sh script could then concatenate the parts and verify the final file, similarly to the existing Q4 workflow.

Thank you in advance!

After all those days, and at least 10 retries, I managed to solve it by running this command trying to escape timeout and retry limits:

export HF_HUB_ETAG_TIMEOUT=300
export HF_HUB_DOWNLOAD_TIMEOUT=600
export HF_XET_CLIENT_CONNECT_TIMEOUT=90
export HF_XET_CLIENT_READ_TIMEOUT=600
export HF_XET_CLIENT_RETRY_BASE_DELAY=5
export HF_XET_CLIENT_RETRY_MAX_ATTEMPTS=100000
export HF_XET_CLIENT_RETRY_MAX_DURATION=36000

hf download antirez/deepseek-v4.1-flash-gguf DeepSeek-V4.1-Flash-Q2.gguf

Hope it can help someone, but much more I hope someone in hf could place himself and point to really make the xet work properly.

PS. Salvatore, non so se leggerai mai questo messaggio ma ci spero, per il futuro potresti caricare i modelli splittati ogni circa 50GB?

I had to write a whole script to download this file because it wasn't working via hf.

python -m pip install -U httpx lz4

python ./hf_xet_wget.py \
          antirez/deepseek-v4.1-flash-gguf \
          DeepSeek-V4.1-Flash-Q2.gguf \
          -o /mnt/raid0/models/DeepSeek-V4.1-Flash-Q2.gguf \
          --workers 16 \
          --window-gib 8 \
          --cache-mib 512
hf_xet_wget.py
#!/usr/bin/env python3
"""
hf_xet_wget.py - standalone Hugging Face Xet downloader.

Downloads Xet-backed files directly through the Xet reconstruction protocol,
without using the Hugging Face CAS Bridge / hf download path.

Usage:
  python hf_xet_wget.py USER/REPO filename.gguf -o /mnt/raid0/models/filename.gguf
  python hf_xet_wget.py USER/REPO path/to/file -o /mnt/raid0/models/file --revision main
  python hf_xet_wget.py org/repo file.bin -o ./file --workers 24 --window-gib 8

Dependencies:
  python -m pip install -U httpx lz4

Auth:
  1) HF_TOKEN environment variable, or
  2) token stored by `hf auth login` (HF_TOKEN_PATH / ~/.cache/huggingface/token)

The program uses:
  Hub resolve HEAD (no redirect) -> X-Xet-Hash / Xet refresh route
  -> Xet /v2/reconstructions/<file_id> for bounded file ranges
  -> signed transfer.xethub URLs with HTTP Range requests
  -> Xorb chunk decompression (None/LZ4/BG4+LZ4)
  -> sequential reconstruction into the destination file.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import sys
import threading
import time
from collections import OrderedDict
from concurrent.futures import Future, ThreadPoolExecutor
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterable
from urllib.parse import quote, urljoin

import httpx

try:
    import lz4.block
except ImportError as exc:
    raise SystemExit("Missing dependency: lz4. Install with: python -m pip install -U lz4") from exc

try:
    import lz4.frame as _lz4_frame  # optional; frame-format LZ4 fallback
except ImportError:
    _lz4_frame = None


HUB = "https://huggingface.co"
DEFAULT_WINDOW_GIB = 8
DEFAULT_WORKERS = 16
DEFAULT_CACHE_MIB = 512
DEFAULT_TIMEOUT = 120.0

# Current Xet header names used by huggingface_hub.
XET_HASH = "X-Xet-Hash"
XET_REFRESH = "X-Xet-Refresh-Route"
XET_ENDPOINT = "X-Xet-Cas-Url"
XET_ACCESS = "X-Xet-Access-Token"
XET_EXPIRATION = "X-Xet-Token-Expiration"


@dataclass(frozen=True)
class RangeDesc:
    xorb_hash: str
    chunk_start: int
    chunk_end: int
    byte_start: int
    byte_end: int
    url: str

    @property
    def key(self) -> tuple[str, int, int]:
        return (self.xorb_hash, self.byte_start, self.byte_end)


class ByteLRU:
    """Small thread-safe byte-bounded LRU cache for serialized Xorb ranges."""

    def __init__(self, max_bytes: int) -> None:
        self.max_bytes = max_bytes
        self._data: OrderedDict[tuple[str, int, int], bytes] = OrderedDict()
        self._size = 0
        self._lock = threading.Lock()

    def get(self, key: tuple[str, int, int]) -> bytes | None:
        with self._lock:
            value = self._data.get(key)
            if value is not None:
                self._data.move_to_end(key)
            return value

    def put(self, key: tuple[str, int, int], value: bytes) -> None:
        with self._lock:
            old = self._data.pop(key, None)
            if old is not None:
                self._size -= len(old)
            self._data[key] = value
            self._size += len(value)
            while self._size > self.max_bytes and self._data:
                _, evicted = self._data.popitem(last=False)
                self._size -= len(evicted)

    def clear(self) -> None:
        with self._lock:
            self._data.clear()
            self._size = 0


class DownloadError(RuntimeError):
    pass


def human_bytes(n: float) -> str:
    units = ["B", "KiB", "MiB", "GiB", "TiB"]
    x = float(n)
    for u in units:
        if x < 1024.0 or u == units[-1]:
            return f"{x:.2f} {u}"
        x /= 1024.0
    return f"{n} B"


def read_token() -> str:
    token = os.environ.get("HF_TOKEN", "").strip()
    if token:
        return token

    token_path = os.environ.get("HF_TOKEN_PATH")
    if token_path:
        p = Path(token_path).expanduser()
    else:
        hf_home = Path(os.environ.get("HF_HOME", Path.home() / ".cache" / "huggingface")).expanduser()
        p = Path(hf_home) / "token"

    try:
        token = p.read_text(encoding="utf-8").strip()
    except FileNotFoundError:
        token = ""

    if not token:
        raise DownloadError(
            "No Hugging Face token found. Set HF_TOKEN or run `hf auth login`. "
            "Xet reconstruction requires a read-capable Hub token."
        )
    return token


def header_ci(headers: httpx.Headers, name: str) -> str | None:
    return headers.get(name)


def get_file_metadata(client: httpx.Client, repo_id: str, filename: str, revision: str, token: str,
                      repo_type: str) -> tuple[str, int, str, str | None]:
    prefix = ""
    if repo_type != "model":
        prefix = {"dataset": "datasets/", "space": "spaces/", "kernel": "kernels/"}[repo_type]
    url = f"{HUB}/{prefix}{quote(repo_id, safe='/')}/resolve/{quote(revision, safe='')}/{quote(filename, safe='/')}"
    headers = {
        "Authorization": f"Bearer {token}",
        "Accept-Encoding": "identity",
        "User-Agent": "hf-xet-wget/1.0",
    }
    # Critical: DO NOT follow the resolve redirect. X-Xet-* headers are on this response.
    r = client.head(url, headers=headers, follow_redirects=False)
    if r.status_code not in (200, 201, 202, 204, 301, 302, 303, 307, 308):
        raise DownloadError(f"Hub HEAD failed: HTTP {r.status_code}: {r.text[:500]}")

    file_hash = header_ci(r.headers, XET_HASH)
    if not file_hash:
        raise DownloadError(
            "The resolve HEAD response did not contain X-Xet-Hash. "
            "This file may not be Xet-backed, or the Hub endpoint did not expose Xet metadata."
        )
    if file_hash.startswith("sha256:"):
        file_hash = file_hash.split(":", 1)[1]
    if len(file_hash) != 64 or any(c not in "0123456789abcdefABCDEF" for c in file_hash):
        raise DownloadError(f"Unexpected Xet file ID: {file_hash!r}")

    size_text = header_ci(r.headers, "X-Linked-Size") or header_ci(r.headers, "Content-Length")
    if not size_text:
        raise DownloadError("Hub response did not provide file size (X-Linked-Size/Content-Length).")
    size = int(size_text)

    refresh = header_ci(r.headers, XET_REFRESH)
    if not refresh:
        link = header_ci(r.headers, "Link") or ""
        for item in link.split(","):
            if 'rel="xet-auth"' in item or "rel=xet-auth" in item:
                left = item.find("<")
                right = item.find(">", left + 1)
                if left >= 0 and right > left:
                    refresh = item[left + 1:right]
                    break

    return file_hash, size, url, refresh


def get_refresh_url(repo_id: str, revision: str, repo_type: str, explicit: str | None) -> str:
    if explicit:
        return urljoin(HUB + "/", explicit)
    plural = {
        "model": "models",
        "dataset": "datasets",
        "space": "spaces",
        "kernel": "kernels",
    }[repo_type]
    return f"{HUB}/api/{plural}/{quote(repo_id, safe='/')}/xet-read-token/{quote(revision, safe='')}"


def refresh_xet_connection(client: httpx.Client, refresh_url: str, token: str) -> tuple[str, str]:
    r = client.get(
        refresh_url,
        headers={"Authorization": f"Bearer {token}", "User-Agent": "hf-xet-wget/1.0"},
        follow_redirects=True,
    )
    if r.status_code != 200:
        raise DownloadError(f"Xet token refresh failed: HTTP {r.status_code}: {r.text[:500]}")
    endpoint = header_ci(r.headers, XET_ENDPOINT)
    access = header_ci(r.headers, XET_ACCESS)
    if not endpoint or not access:
        raise DownloadError("Xet token response missing X-Xet-Cas-Url or X-Xet-Access-Token.")
    return endpoint.rstrip("/"), access


def get_reconstruction(client: httpx.Client, cas_endpoint: str, file_id: str, access_token: str,
                       start: int, end: int) -> dict[str, Any]:
    url = f"{cas_endpoint}/v2/reconstructions/{file_id}"
    r = client.get(
        url,
        headers={
            "Authorization": f"Bearer {access_token}",
            "Range": f"bytes={start}-{end}",
            "Accept-Encoding": "gzip",
            "User-Agent": "hf-xet-wget/1.0",
        },
        follow_redirects=True,
    )
    if r.status_code != 200:
        raise DownloadError(f"Xet reconstruction failed: HTTP {r.status_code}: {r.text[:1000]}")
    try:
        obj = r.json()
    except json.JSONDecodeError as exc:
        raise DownloadError("Xet reconstruction response was not valid JSON.") from exc
    if "terms" not in obj or "xorbs" not in obj:
        raise DownloadError("Xet reconstruction response lacks `terms` or `xorbs`.")
    return obj


def int_pair(obj: dict[str, Any], snake: str, camel: str) -> tuple[int, int]:
    value = obj.get(snake, obj.get(camel))
    if not isinstance(value, dict):
        raise DownloadError(f"Missing range object: {snake}")
    a = value.get("start")
    b = value.get("end")
    if not isinstance(a, int) or not isinstance(b, int):
        raise DownloadError(f"Malformed range: {value!r}")
    return a, b


def normalize_descriptors(reconstruction: dict[str, Any]) -> dict[str, list[RangeDesc]]:
    result: dict[str, list[RangeDesc]] = {}
    xorbs = reconstruction["xorbs"]
    if not isinstance(xorbs, dict):
        raise DownloadError("Malformed `xorbs` in reconstruction response.")

    for xorb_hash, entries in xorbs.items():
        if not isinstance(entries, list):
            raise DownloadError(f"Malformed fetch list for xorb {xorb_hash}")
        out: list[RangeDesc] = []
        for entry in entries:
            if not isinstance(entry, dict):
                raise DownloadError(f"Malformed xorb fetch entry for {xorb_hash}")
            url = entry.get("url")
            ranges = entry.get("ranges")
            if not isinstance(url, str) or not isinstance(ranges, list):
                raise DownloadError(f"Malformed xorb fetch entry for {xorb_hash}")
            for desc in ranges:
                if not isinstance(desc, dict):
                    raise DownloadError("Malformed xorb range descriptor")
                chunks = desc.get("chunks")
                bytes_range = desc.get("bytes")
                if not isinstance(chunks, dict) or not isinstance(bytes_range, dict):
                    raise DownloadError("Malformed xorb range descriptor")
                cs, ce = int_pair(desc, "chunks", "chunks")
                bs = bytes_range.get("start")
                be = bytes_range.get("end")
                if not isinstance(bs, int) or not isinstance(be, int):
                    raise DownloadError("Malformed physical xorb byte range")
                # chunks is half-open [cs, ce); bytes is inclusive [bs, be].
                if cs >= ce or bs > be:
                    raise DownloadError("Invalid xorb range bounds")
                out.append(RangeDesc(xorb_hash, cs, ce, bs, be, url))
        out.sort(key=lambda d: (d.chunk_start, d.chunk_end, d.byte_start))
        result[str(xorb_hash)] = out
    return result


def parse_term(term: dict[str, Any]) -> tuple[str, int, int, int]:
    xorb = term.get("hash")
    rng = term.get("range")
    if not isinstance(xorb, str) or not isinstance(rng, dict):
        raise DownloadError(f"Malformed reconstruction term: {term!r}")
    start = rng.get("start")
    end = rng.get("end")
    unpacked = term.get("unpacked_length")
    if not isinstance(start, int) or not isinstance(end, int) or not isinstance(unpacked, int):
        raise DownloadError(f"Malformed reconstruction term fields: {term!r}")
    if start >= end or unpacked < 0:
        raise DownloadError(f"Invalid reconstruction term bounds: {term!r}")
    return xorb, start, end, unpacked


def get_offset_first(obj: dict[str, Any]) -> int:
    value = obj.get("offset_into_first_range", obj.get("offsetIntoFirstRange", 0))
    if not isinstance(value, int) or value < 0:
        raise DownloadError(f"Invalid offset_into_first_range: {value!r}")
    return value


def descriptors_for_term(descs: list[RangeDesc], chunk_start: int, chunk_end: int) -> list[RangeDesc]:
    selected = [d for d in descs if d.chunk_end > chunk_start and d.chunk_start < chunk_end]
    selected.sort(key=lambda d: d.chunk_start)
    if not selected:
        raise DownloadError(f"No Xorb physical ranges cover chunk range [{chunk_start}, {chunk_end})")

    cursor = chunk_start
    for d in selected:
        if d.chunk_start > cursor:
            raise DownloadError(
                f"Gap in Xorb reconstruction: wanted chunk {cursor}, next descriptor starts at {d.chunk_start}"
            )
        cursor = max(cursor, min(chunk_end, d.chunk_end))
        if cursor >= chunk_end:
            break
    if cursor < chunk_end:
        raise DownloadError(f"Xorb descriptors do not cover full chunk range [{chunk_start}, {chunk_end})")
    return selected


def _lz4_decompress_compat(data: bytes, uncompressed_size: int) -> bytes:
    """Decode Xet LZ4 blocks across implementations.

    Xet's storage stack has emitted LZ4 in several shapes:
      * raw LZ4 block (lz4_flex::block::compress_into)
      * 4-byte little-endian u32 size prefix + raw block
        (lz4_flex::block::compress / compress_prepend_size)
      * LZ4 frame format (early builds / tools)
    Additionally, the byte ranges fetched from the xorb URL sometimes carry
    one or more trailing padding bytes inside the last chunk's payload, which
    trips up the strict LZ4 block decoder (error 4). We therefore also try
    trimming 1..16 trailing bytes with each candidate offset.
    """
    errors: list[str] = []

    def _try(label: str, payload: bytes, want: int) -> bytes | None:
        if not payload:
            return None
        try:
            raw = lz4.block.decompress(payload, uncompressed_size=want)
        except Exception as exc:
            errors.append(f"{label}[want={want}]: {exc}")
            return None
        if len(raw) == want:
            return bytes(raw)
        errors.append(f"{label}[want={want}]: decoded {len(raw)} bytes")
        return None

    # (offset, label) candidates for prefix-stripping.
    offsets: list[tuple[int, str]] = [(0, "raw")]
    if len(data) >= 4:
        offsets.append((4, "u32le"))
    if len(data) >= 8:
        offsets.append((8, "u64le"))

    # Pass 1: try each offset with the header-provided uncompressed_size and
    # one "generous" variant in case the header underreports it.
    want_variants = [uncompressed_size]
    if uncompressed_size + 4096 <= (1 << 24) - 1:
        want_variants.append(uncompressed_size + 4096)

    for off, label in offsets:
        if off >= len(data):
            continue
        for want in want_variants:
            r = _try(f"{label}(strip={off})", data[off:], want)
            if r is not None:
                if want != uncompressed_size:
                    r = r[:uncompressed_size]
                return r

    # Pass 2: LZ4 frame format.
    if _lz4_frame is not None:
        try:
            raw = _lz4_frame.decompress(data)
            if len(raw) >= uncompressed_size:
                return bytes(raw[:uncompressed_size])
            errors.append(f"frame: decoded {len(raw)} bytes")
        except Exception as exc:
            errors.append(f"frame: {exc}")

    # Pass 3: trim 1..16 trailing bytes with each offset/want combination.
    max_trim = min(16, max(0, len(data) - 1))
    for extra in range(1, max_trim + 1):
        for off, label in offsets:
            if off + extra >= len(data):
                continue
            payload = data[off:len(data) - extra]
            if not payload:
                continue
            for want in want_variants:
                r = _try(f"{label}(strip={off},trunc={extra})", payload, want)
                if r is not None:
                    if want != uncompressed_size:
                        r = r[:uncompressed_size]
                    return r

    head = data[:16].hex()
    tail = data[-8:].hex() if len(data) >= 8 else data.hex()
    detail = "\n    ".join(errors[-10:])
    raise DownloadError(
        f"LZ4 decompression failed for {len(data)} compressed bytes -> "
        f"{uncompressed_size} bytes (head={head} tail={tail}). Attempts:\n    {detail}"
    )


def decompress_chunk(compression: int, compressed: bytes, uncompressed_size: int) -> bytes:
    if compression == 0:
        raw = compressed
    elif compression == 1:
        raw = _lz4_decompress_compat(compressed, uncompressed_size)
    elif compression == 2:
        grouped = _lz4_decompress_compat(compressed, uncompressed_size)
        raw = bytearray(uncompressed_size)
        q, rem = divmod(uncompressed_size, 4)
        group_lens = [q + (1 if rem >= i + 1 else 0) for i in range(4)]
        group_offsets = [0]
        for glen in group_lens[:-1]:
            group_offsets.append(group_offsets[-1] + glen)
        positions = [0, 0, 0, 0]
        for i in range(uncompressed_size):
            g = i & 3
            raw[i] = grouped[group_offsets[g] + positions[g]]
            positions[g] += 1
        return bytes(raw)
    else:
        raise DownloadError(f"Unsupported Xorb compression type: {compression}")

    if len(raw) != uncompressed_size:
        raise DownloadError(
            f"Xorb chunk size mismatch: decoded {len(raw)} bytes, expected {uncompressed_size}"
        )
    return raw


def iter_xorb_chunks(data: bytes, first_chunk_index: int) -> Iterable[tuple[int, bytes]]:
    pos = 0
    chunk_index = first_chunk_index
    n = len(data)
    while pos < n:
        if n - pos < 8:
            # A few trailing zero bytes are padding introduced by the physical
            # xorb layout; anything else is a genuine truncation.
            trailing = data[pos:]
            if trailing and all(b == 0 for b in trailing):
                break
            raise DownloadError(
                f"Truncated Xorb chunk header ({n - pos} trailing bytes at offset {pos}, "
                f"hex={trailing.hex()})"
            )
        version = data[pos]
        compressed_size = int.from_bytes(data[pos + 1:pos + 4], "little")
        compression = data[pos + 4]
        uncompressed_size = int.from_bytes(data[pos + 5:pos + 8], "little")
        if version != 0:
            raise DownloadError(f"Unsupported Xorb protocol version: {version}")
        pos += 8
        end = pos + compressed_size
        if end > n:
            raise DownloadError(
                f"Truncated Xorb compressed chunk: header wants {compressed_size} bytes "
                f"at offset {pos - 8}, only {n - pos} available"
            )
        compressed = data[pos:end]
        raw = decompress_chunk(compression, compressed, uncompressed_size)
        yield chunk_index, raw
        chunk_index += 1
        pos = end
    if pos != n and not (pos <= n and all(b == 0 for b in data[pos:])):
        raise DownloadError("Internal Xorb parser error")


_thread_local = threading.local()


def worker_client(timeout: float) -> httpx.Client:
    client = getattr(_thread_local, "client", None)
    if client is None:
        client = httpx.Client(
            timeout=httpx.Timeout(timeout, connect=min(timeout, 30.0)),
            follow_redirects=True,
            limits=httpx.Limits(max_keepalive_connections=4, max_connections=4),
            headers={"User-Agent": "hf-xet-wget/1.0"},
        )
        _thread_local.client = client
    return client


def download_range(desc: RangeDesc, retries: int, timeout: float) -> bytes:
    client = worker_client(timeout)
    # bytes.start..bytes.end in the Xet descriptor are inclusive on both ends.
    length = desc.byte_end - desc.byte_start + 1
    if length <= 0:
        raise DownloadError(f"Empty byte range for {desc.key}: {desc.byte_start}..{desc.byte_end}")
    expected = length
    last_error: Exception | None = None

    for attempt in range(1, retries + 1):
        try:
            r = client.get(
                desc.url,
                headers={"Range": f"bytes={desc.byte_start}-{desc.byte_end}"},
            )
            if r.status_code not in (200, 206):
                if r.status_code in (429, 500, 502, 503, 504):
                    raise httpx.HTTPStatusError(
                        f"transient HTTP {r.status_code}", request=r.request, response=r
                    )
                raise DownloadError(f"Xorb download failed: HTTP {r.status_code}: {r.text[:500]}")
            data = r.content
            if len(data) != expected:
                raise DownloadError(
                    f"Xorb range length mismatch for {desc.key}: got {len(data)}, expected {expected}"
                )
            return data
        except (httpx.HTTPError, DownloadError) as exc:
            last_error = exc
            if attempt >= retries:
                break
            time.sleep(min(2 ** (attempt - 1), 10))

    raise DownloadError(f"Failed Xorb range {desc.key} after {retries} attempts: {last_error}")


def plan_term_prefetch(terms: list[dict[str, Any]], descriptors: dict[str, list[RangeDesc]],
                       start_index: int, count: int) -> list[RangeDesc]:
    out: list[RangeDesc] = []
    seen: set[tuple[str, int, int]] = set()
    for term in terms[start_index:start_index + count]:
        xorb, cs, ce, _ = parse_term(term)
        for d in descriptors_for_term(descriptors.get(xorb, []), cs, ce):
            if d.key not in seen:
                out.append(d)
                seen.add(d.key)
    return out


def reconstruct_window(
    *,
    reconstruction: dict[str, Any],
    output: Any,
    window_start: int,
    window_end_exclusive: int,
    workers: int,
    cache: ByteLRU,
    retries: int,
    timeout: float,
    prefetch_terms: int,
) -> int:
    terms_obj = reconstruction["terms"]
    if not isinstance(terms_obj, list):
        raise DownloadError("Malformed reconstruction `terms`")
    terms = [t for t in terms_obj if isinstance(t, dict)]
    descriptors = normalize_descriptors(reconstruction)
    skip_first = get_offset_first(reconstruction)

    written = 0
    decoded_first_seen = False
    completed_term_index = 0

    executor = ThreadPoolExecutor(max_workers=max(1, workers), thread_name_prefix="xet-range")
    futures: dict[tuple[str, int, int], Future[bytes]] = {}

    def ensure_future(desc: RangeDesc) -> Future[bytes]:
        key = desc.key
        if key in futures:
            return futures[key]
        cached = cache.get(key)
        if cached is not None:
            fut: Future[bytes] = Future()
            fut.set_result(cached)
            futures[key] = fut
            return fut
        fut = executor.submit(download_range, desc, retries, timeout)
        futures[key] = fut
        return fut

    try:
        for term_index, term in enumerate(terms):
            xorb, cs, ce, expected_term_bytes = parse_term(term)
            selected = descriptors_for_term(descriptors.get(xorb, []), cs, ce)

            for d in plan_term_prefetch(terms, descriptors, term_index, prefetch_terms):
                ensure_future(d)

            decoded_term = 0
            for d in selected:
                fut = ensure_future(d)
                blob = fut.result()
                cache.put(d.key, blob)

                first_selected_chunk = max(cs, d.chunk_start)
                last_selected_chunk = min(ce, d.chunk_end)
                if first_selected_chunk >= last_selected_chunk:
                    continue

                for chunk_index, raw in iter_xorb_chunks(blob, d.chunk_start):
                    if chunk_index < first_selected_chunk:
                        continue
                    if chunk_index >= last_selected_chunk:
                        break
                    decoded_term += len(raw)

                    data = raw
                    if not decoded_first_seen:
                        decoded_first_seen = True
                        if skip_first:
                            if skip_first >= len(data):
                                skip_first -= len(data)
                                continue
                            data = data[skip_first:]
                            skip_first = 0

                    remaining = (window_end_exclusive - window_start) - written
                    if remaining <= 0:
                        break
                    if len(data) > remaining:
                        data = data[:remaining]
                    output.write(data)
                    written += len(data)
                    if written >= (window_end_exclusive - window_start):
                        break
                if written >= (window_end_exclusive - window_start):
                    break

            if decoded_term != expected_term_bytes:
                raise DownloadError(
                    f"Term {term_index} decoded {decoded_term} bytes, expected {expected_term_bytes}"
                )
            completed_term_index = term_index + 1
            if written >= (window_end_exclusive - window_start):
                break

        expected_window = window_end_exclusive - window_start
        if written != expected_window:
            raise DownloadError(
                f"Window reconstruction ended at {written} bytes, expected {expected_window}. "
                f"Completed terms: {completed_term_index}/{len(terms)}"
            )
        return written
    finally:
        executor.shutdown(wait=True, cancel_futures=False)


def sha256_file(path: Path, chunk_size: int = 64 * 1024 * 1024) -> str:
    h = hashlib.sha256()
    with path.open("rb", buffering=0) as f:
        while True:
            data = f.read(chunk_size)
            if not data:
                break
            h.update(data)
    return h.hexdigest()


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Direct Hugging Face Xet downloader")
    p.add_argument("repo_id", help="Repository, e.g. antirez/deepseek-v4.1-flash-gguf")
    p.add_argument("filename", help="Repository-relative file path")
    p.add_argument("-o", "--output", required=True, help="Destination file")
    p.add_argument("--revision", default="main", help="Branch/tag/commit (default: main)")
    p.add_argument("--repo-type", choices=["model", "dataset", "space", "kernel"], default="model")
    p.add_argument("--endpoint", default=HUB, help="Hugging Face Hub endpoint (default: https://huggingface.co)")
    p.add_argument("--workers", type=int, default=DEFAULT_WORKERS, help=f"Parallel Xorb HTTP requests (default: {DEFAULT_WORKERS})")
    p.add_argument("--window-gib", type=float, default=DEFAULT_WINDOW_GIB, help=f"Reconstruction window in GiB (default: {DEFAULT_WINDOW_GIB})")
    p.add_argument("--cache-mib", type=int, default=DEFAULT_CACHE_MIB, help=f"Serialized Xorb RAM cache in MiB (default: {DEFAULT_CACHE_MIB})")
    p.add_argument("--prefetch-terms", type=int, default=4, help="Number of terms to prefetch (default: 4)")
    p.add_argument("--retries", type=int, default=8, help="Retries per range (default: 8)")
    p.add_argument("--timeout", type=float, default=DEFAULT_TIMEOUT, help="HTTP timeout seconds (default: 120)")
    p.add_argument("--sha256", help="Verify final SHA-256 against this hexadecimal digest")
    p.add_argument("--force", action="store_true", help="Discard existing partial output and start over")
    p.add_argument("--verbose", action="store_true")
    return p.parse_args()


def main() -> int:
    args = parse_args()
    if args.workers < 1 or args.workers > 128:
        raise SystemExit("--workers must be between 1 and 128")
    if args.window_gib <= 0:
        raise SystemExit("--window-gib must be > 0")
    if args.cache_mib < 0:
        raise SystemExit("--cache-mib must be >= 0")
    if args.prefetch_terms < 1:
        raise SystemExit("--prefetch-terms must be >= 1")

    global HUB
    HUB = args.endpoint.rstrip("/")
    token = read_token()
    output_path = Path(args.output).expanduser().resolve()
    output_path.parent.mkdir(parents=True, exist_ok=True)

    if args.force and output_path.exists():
        output_path.unlink()
    if output_path.exists() and not output_path.is_file():
        raise SystemExit(f"Output is not a regular file: {output_path}")

    client = httpx.Client(timeout=httpx.Timeout(args.timeout, connect=min(args.timeout, 30.0)), follow_redirects=False)
    file_id, file_size, _, refresh_hint = get_file_metadata(
        client, args.repo_id, args.filename, args.revision, token, args.repo_type
    )

    print(f"repo       : {args.repo_id}")
    print(f"file       : {args.filename}")
    print(f"revision   : {args.revision}")
    print(f"size       : {human_bytes(file_size)} ({file_size:,} bytes)")
    print(f"xet file-id: {file_id}")
    print(f"output     : {output_path}")
    print(f"workers    : {args.workers}")
    print(f"window     : {args.window_gib:g} GiB")
    print(f"cache      : {args.cache_mib} MiB")

    existing = output_path.stat().st_size if output_path.exists() else 0
    if existing > file_size:
        raise SystemExit(f"Existing output is larger than remote file: {existing} > {file_size}")
    if existing == file_size:
        print("Output already has the expected size; skipping download.")
    else:
        refresh_url = get_refresh_url(args.repo_id, args.revision, args.repo_type, refresh_hint)
        cache = ByteLRU(args.cache_mib * 1024 * 1024)
        window_bytes = max(1, int(args.window_gib * 1024**3))
        total_start = time.monotonic()
        start = existing

        with output_path.open("r+b" if output_path.exists() else "wb", buffering=1024 * 1024) as out:
            out.seek(start)
            while start < file_size:
                end_exclusive = min(file_size, start + window_bytes)
                end_inclusive = end_exclusive - 1

                cas_endpoint, xet_access_token = refresh_xet_connection(client, refresh_url, token)
                t0 = time.monotonic()
                reconstruction = get_reconstruction(
                    client, cas_endpoint, file_id, xet_access_token, start, end_inclusive
                )
                meta_ms = (time.monotonic() - t0) * 1000

                if args.verbose:
                    print(f"\nwindow {human_bytes(start)} .. {human_bytes(end_exclusive)}; reconstruction {meta_ms:.0f} ms")

                t1 = time.monotonic()
                written = reconstruct_window(
                    reconstruction=reconstruction,
                    output=out,
                    window_start=start,
                    window_end_exclusive=end_exclusive,
                    workers=args.workers,
                    cache=cache,
                    retries=args.retries,
                    timeout=args.timeout,
                    prefetch_terms=args.prefetch_terms,
                )
                out.flush()
                os.fsync(out.fileno())
                start += written
                elapsed = time.monotonic() - total_start
                rate = (start - existing) / elapsed if elapsed > 0 else 0
                print(
                    f"\r{human_bytes(start)} / {human_bytes(file_size)}  "
                    f"{rate / (1024**2):.2f} MiB/s  ({100 * start / file_size:.2f}%)",
                    end="",
                    flush=True,
                )
                cache.clear()
            print()

    final_size = output_path.stat().st_size
    if final_size != file_size:
        raise SystemExit(f"Final size mismatch: got {final_size}, expected {file_size}")

    if args.sha256:
        print("Computing SHA-256...")
        actual = sha256_file(output_path)
        expected = args.sha256.lower().strip()
        print(f"sha256     : {actual}")
        if actual != expected:
            raise SystemExit(f"SHA-256 mismatch: expected {expected}, got {actual}")
        print("SHA-256    : OK")

    client.close()
    print(f"Done: {output_path}")
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except KeyboardInterrupt:
        print("\nInterrupted. The partial output is kept; rerun the same command to resume.", file=sys.stderr)
        raise SystemExit(130)
    except (DownloadError, httpx.HTTPError) as exc:
        print(f"ERROR: {exc}", file=sys.stderr)
        raise SystemExit(1)

Sign up or log in to comment