Spaces:

Opera8
/

LTX-2-3

Paused

File size: 9,597 Bytes

87d4eda

"""Locate (or fetch) the model files this Space needs.

Resolution order for every model:

1. A mounted Hub volume (``/models/ltx`` and ``/models/gemma``).  This is the
   recommended setup: nothing touches the Space's small ephemeral disk.
2. The local Hugging Face cache (a previous download).
3. A fresh download from the Hub (needs ``HF_TOKEN`` for the gated Gemma repo
   and a *lot* of free disk space).

If none of these work a ``ModelSetupError`` is raised with a message that says
exactly what to fix, instead of the cryptic
``No files matching pattern 'tokenizer.model'`` from deep inside ltx-core.

This module deliberately imports neither torch nor ltx so it can be unit-tested
anywhere.
"""

from __future__ import annotations

import fnmatch
import os
import shutil
from dataclasses import dataclass
from pathlib import Path

# --------------------------------------------------------------------------- #
# Configuration (all overridable through environment variables / Space vars)
# --------------------------------------------------------------------------- #
LTX_REPO_ID = os.environ.get("LTX_REPO_ID", "Lightricks/LTX-2.3")
GEMMA_REPO_ID = os.environ.get("GEMMA_REPO_ID", "google/gemma-3-12b-it-qat-q4_0-unquantized")

LTX_MOUNT = os.environ.get("LTX_MODEL_DIR", "/models/ltx")
GEMMA_MOUNT = os.environ.get("GEMMA_MODEL_DIR", "/models/gemma")

DISTILLED_FILENAME = os.environ.get("LTX_DISTILLED_FILENAME", "ltx-2.3-22b-distilled-1.1.safetensors")
UPSCALER_FILENAME = os.environ.get("LTX_UPSCALER_FILENAME", "ltx-2.3-spatial-upscaler-x2-1.1.safetensors")

# Only what the text encoder needs (weights + tokenizer + configs).
GEMMA_ALLOW_PATTERNS = ["*.json", "*.model", "*.safetensors", "*.txt", "*.jinja"]

# On a Space, SPACE_ID is set by Hugging Face (e.g. "Opera8/LTX-2-3").
_SPACE = os.environ.get("SPACE_ID", "<your-username>/<your-space>")

SETUP_HINT = (
    "\n\nHOW TO FIX (pick one):\n"
    "  A) Recommended - mount the models as read-only volumes (no disk used):\n"
    f"       python setup_volumes.py {_SPACE}\n"
    "     or:\n"
    f"       hf spaces volumes set {_SPACE} \\\n"
    f"           -v hf://models/{LTX_REPO_ID}:{LTX_MOUNT}:ro \\\n"
    f"           -v hf://models/{GEMMA_REPO_ID}:{GEMMA_MOUNT}:ro\n"
    "  B) Fallback - download at startup: add a Space *secret* named HF_TOKEN\n"
    "     (a token of an account that accepted the Gemma license) and make sure\n"
    "     the Space has enough free disk space.\n"
    f"  Gemma license: https://huggingface.co/{GEMMA_REPO_ID}\n"
)


class ModelSetupError(RuntimeError):
    """Raised when the model files cannot be found or downloaded."""


@dataclass(frozen=True)
class ModelPaths:
    distilled_checkpoint: str
    spatial_upsampler: str
    gemma_root: str


# --------------------------------------------------------------------------- #
# Small helpers
# --------------------------------------------------------------------------- #
def hf_token() -> str | None:
    return os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") or None


def describe_dir(path: str | os.PathLike, limit: int = 20) -> str:
    """Human readable one-liner about a directory, used for diagnostics."""
    p = Path(path)
    if not p.exists():
        return f"{p}: DOES NOT EXIST (volume not mounted?)"
    if not p.is_dir():
        return f"{p}: exists but is not a directory"
    try:
        names = sorted(x.name for x in p.iterdir())
    except OSError as e:  # pragma: no cover - defensive
        return f"{p}: cannot list ({e})"
    if not names:
        return f"{p}: EMPTY"
    shown = ", ".join(names[:limit]) + (" ..." if len(names) > limit else "")
    return f"{p}: {len(names)} entries -> {shown}"


def find_file(root: str | os.PathLike, name: str) -> Path | None:
    """Return the first file called ``name`` anywhere under ``root`` (or None)."""
    root = Path(root)
    if not root.is_dir():
        return None
    for candidate in root.rglob(name):
        if candidate.is_file():
            return candidate
    return None


def _gemma_is_complete(root: str | os.PathLike) -> bool:
    root = Path(root)
    has_tokenizer = find_file(root, "tokenizer.model") is not None
    has_weights = any(root.rglob("*.safetensors")) if root.is_dir() else False
    return has_tokenizer and has_weights


def _remote_size_bytes(repo_id: str, patterns: list[str], token: str | None) -> int | None:
    """Total size of the files in ``repo_id`` matching ``patterns`` (None if unknown)."""
    try:
        from huggingface_hub import HfApi

        info = HfApi(token=token).model_info(repo_id, files_metadata=True)
        total = 0
        for sibling in info.siblings or []:
            if any(fnmatch.fnmatch(sibling.rfilename, pat) for pat in patterns):
                total += sibling.size or 0
        return total or None
    except Exception:
        return None


def _hub_cache_dir() -> str:
    try:
        from huggingface_hub.constants import HF_HUB_CACHE

        return str(HF_HUB_CACHE)
    except Exception:  # pragma: no cover - very old huggingface_hub
        return os.path.expanduser("~/.cache/huggingface/hub")


def _check_disk_space(needed: int | None, what: str) -> None:
    if not needed:
        return
    target = _hub_cache_dir()
    os.makedirs(target, exist_ok=True)
    free = shutil.disk_usage(target).free
    if free < needed * 1.02:
        raise ModelSetupError(
            f"Not enough disk space to download {what}: need ~{needed / 1024**3:.1f} GB, "
            f"only {free / 1024**3:.1f} GB free in {target}. "
            "A Space's ephemeral disk is too small for these models - mount them as "
            "volumes instead." + SETUP_HINT
        )


def _explain_download_error(exc: Exception, repo_id: str) -> str:
    name = type(exc).__name__
    text = str(exc)
    gated = name in {"GatedRepoError", "RepositoryNotFoundError"} or any(
        code in text for code in ("401", "403", "gated", "restricted", "Access to model")
    )
    if gated and not hf_token():
        return (
            f"Cannot download '{repo_id}': it is gated and no HF_TOKEN is set. "
            f"Accept the license at https://huggingface.co/{repo_id} and add a Space secret "
            "named HF_TOKEN." + SETUP_HINT
        )
    if gated:
        return (
            f"Access to '{repo_id}' was denied. Make sure the account that owns HF_TOKEN has "
            f"accepted the license at https://huggingface.co/{repo_id}." + SETUP_HINT
        )
    return f"Downloading '{repo_id}' failed: {name}: {text}" + SETUP_HINT


# --------------------------------------------------------------------------- #
# Resolvers
# --------------------------------------------------------------------------- #
def _resolve_ltx_file(filename: str) -> str:
    # 1) mounted volume
    direct = Path(LTX_MOUNT) / filename
    if direct.is_file():
        return str(direct)
    nested = find_file(LTX_MOUNT, filename)
    if nested is not None:
        return str(nested)

    print(f"[MODELS] '{filename}' not found under {LTX_MOUNT} -> trying the Hub cache / download.")
    print(f"[MODELS] {describe_dir(LTX_MOUNT)}")

    from huggingface_hub import hf_hub_download

    token = hf_token()
    # 2) local cache
    try:
        return hf_hub_download(LTX_REPO_ID, filename, token=token, local_files_only=True)
    except Exception:
        pass

    # 3) real download
    _check_disk_space(_remote_size_bytes(LTX_REPO_ID, [filename], token), filename)
    try:
        return hf_hub_download(LTX_REPO_ID, filename, token=token)
    except Exception as exc:
        raise ModelSetupError(_explain_download_error(exc, LTX_REPO_ID)) from exc


def _resolve_gemma_root() -> str:
    # 1) mounted volume
    if _gemma_is_complete(GEMMA_MOUNT):
        return GEMMA_MOUNT

    print(f"[MODELS] Gemma (tokenizer.model + weights) not found under {GEMMA_MOUNT} "
          "-> trying the Hub cache / download.")
    print(f"[MODELS] {describe_dir(GEMMA_MOUNT)}")

    from huggingface_hub import snapshot_download

    token = hf_token()
    # 2) local cache
    try:
        root = snapshot_download(
            GEMMA_REPO_ID, allow_patterns=GEMMA_ALLOW_PATTERNS, token=token, local_files_only=True
        )
        if _gemma_is_complete(root):
            return root
    except Exception:
        pass

    # 3) real download
    _check_disk_space(_remote_size_bytes(GEMMA_REPO_ID, GEMMA_ALLOW_PATTERNS, token), "Gemma")
    try:
        root = snapshot_download(GEMMA_REPO_ID, allow_patterns=GEMMA_ALLOW_PATTERNS, token=token)
    except Exception as exc:
        raise ModelSetupError(_explain_download_error(exc, GEMMA_REPO_ID)) from exc

    if not _gemma_is_complete(root):
        raise ModelSetupError(
            f"Gemma download finished but {root} has no tokenizer.model / *.safetensors. "
            f"Check that {GEMMA_REPO_ID} is the right repo." + SETUP_HINT
        )
    return root


def resolve_model_paths() -> ModelPaths:
    """Return usable paths for every model file, or raise ``ModelSetupError``."""
    print("[MODELS] " + describe_dir(LTX_MOUNT))
    print("[MODELS] " + describe_dir(GEMMA_MOUNT))

    paths = ModelPaths(
        distilled_checkpoint=_resolve_ltx_file(DISTILLED_FILENAME),
        spatial_upsampler=_resolve_ltx_file(UPSCALER_FILENAME),
        gemma_root=_resolve_gemma_root(),
    )
    print(f"[MODELS] distilled checkpoint : {paths.distilled_checkpoint}")
    print(f"[MODELS] spatial upsampler    : {paths.spatial_upsampler}")
    print(f"[MODELS] gemma root           : {paths.gemma_root}")
    return paths