Spaces:

Opera8
/

LTX-2-3

Paused

LTX-2-3 / model_resolver.py

Transfer Bot

Moved to Hugging Face automatically

87d4eda 7 days ago

9.6 kB

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

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