fdanyone/assets.py

"""Locate the model files used by 4DAnyone.

Every published file is anchored by one immutable Hugging Face revision and
downloaded on demand. ``fdanyone.download`` fetches missing files;
the resolvers here only locate them.
"""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING

from fdanyone.errors import AssetError

if TYPE_CHECKING:
    from fdanyone.config import ModeSettings

HF_REPO_ID = "AntResearch/4DAnyone"
HF_REVISION = "7850985888b56aabf09e69480b73248f1a76bcbe"

BIREFNET_REPO_ID = "ZhengPeng7/BiRefNet"
BIREFNET_REVISION = "e2bf8e4460fc8fa32bba5ea4d94b3233d367b0e4"
BIREFNET_DIR = "birefnet"
BIREFNET_FILES = (
    "BiRefNet_config.py",
    "birefnet.py",
    "config.json",
    "model.safetensors",
)

CHECKPOINT = "4danyone/model.safetensors"
MHR70_REGRESSOR = "4danyone/smplx_to_goliath70.pt"
WAN_VAE = "4danyone/Wan2.2_VAE.pth"
TEXT_ENCODER = "4danyone/models_t5_umt5-xxl-enc-bf16.pth"
TOKENIZER_DIR = "4danyone/umt5-xxl"
TOKENIZER_FILES = tuple(
    f"{TOKENIZER_DIR}/{name}"
    for name in ("special_tokens_map.json", "spiece.model", "tokenizer.json", "tokenizer_config.json")
)

GVHMR_CHECKPOINT = "gvhmr/gvhmr_siga24_release.ckpt"
HMR2_CHECKPOINT = "gvhmr/epoch=10-step=25000.ckpt"
VITPOSE_CHECKPOINT = "gvhmr/vitpose-h-multi-coco.pth"
YOLO_CHECKPOINT = "gvhmr/yolov8x.pt"
PERCEPTUAL_VGG19 = "perceptual/imagenet-vgg-verydeep-19-conv.safetensors"
TAEW2_2 = "fps-assets/taehv/taew2_2.pth"
TAEW2_2_URL = (
    "https://raw.githubusercontent.com/madebyollin/taehv/"
    "e743234f3217ab3d1570f65642ab06596d1bd7c5/taew2_2.pth"
)
TAEW2_2_SHA256 = "d053e216ca50e2bb837bbcd79b85f0366bea00e5938025572382a773b74c559a"
TURBO_LORA = (
    "fps-assets/turbo-lora/LoRAs/Wan22-Turbo/"
    "Wan22_TI2V_5B_Turbo_lora_rank_64_fp16.safetensors"
)
TURBO_LORA_REPO_ID = "Kijai/WanVideo_comfy"
TURBO_LORA_REVISION = "86c2b0442e01eeee630b48fd7efc0cd37af03252"
TURBO_LORA_REPO_FILE = (
    "LoRAs/Wan22-Turbo/Wan22_TI2V_5B_Turbo_lora_rank_64_fp16.safetensors"
)
TURBO_LORA_SHA256 = "0ace5244e3d1256f884662c261b017249796cf5b95f05d5ed93cc02a478967b8"

SMPLX_MODEL = "body_models/smplx/SMPLX_NEUTRAL.npz"

# Space patch (sync_vendor.sh): the exported prompt embedding replaces the
# UMT5-XXL encoder and its tokenizer, and reconstruction never runs here.
MODEL_FILES = (
    CHECKPOINT,
    MHR70_REGRESSOR,
    WAN_VAE,
    GVHMR_CHECKPOINT,
    HMR2_CHECKPOINT,
    VITPOSE_CHECKPOINT,
    YOLO_CHECKPOINT,
)

EXAMPLE_FILES = (
    "data/source/pexels/10331522-uhd_2160_4096_25fps.mp4",
    "data/source/pexels/2785536-uhd_2160_3840_25fps.mp4",
    "data/source/pexels/5435720-uhd_2160_4096_25fps.mp4",
    "data/source/pexels/5885633-hd_1080_1920_25fps.mp4",
    "data/source/pexels/5999210-uhd_2160_4096_25fps.mp4",
    "data/source/pexels/6980035-uhd_2160_4096_30fps.mp4",
    "data/source/pexels/7080903-hd_1080_1920_30fps.mp4",
    "data/source/pexels/7480858-uhd_2160_3840_25fps.mp4",
)

# Upstream GVHMR resolves its model files relative to its own checkout, so the
# install commands link each downloaded file to the location GVHMR expects.
GVHMR_LINKS = (
    (GVHMR_CHECKPOINT, "inputs/checkpoints/gvhmr/gvhmr_siga24_release.ckpt"),
    (HMR2_CHECKPOINT, "inputs/checkpoints/hmr2/epoch=10-step=25000.ckpt"),
    (VITPOSE_CHECKPOINT, "inputs/checkpoints/vitpose/vitpose-h-multi-coco.pth"),
    (YOLO_CHECKPOINT, "inputs/checkpoints/yolo/yolov8x.pt"),
    (SMPLX_MODEL, "inputs/checkpoints/body_models/smplx/SMPLX_NEUTRAL.npz"),
)


@dataclass(frozen=True, slots=True)
class BaseAssets:
    vae: Path
    """Wan VAE checkpoint."""
    text_encoder: Path | None
    """UMT5 text-encoder checkpoint; ``None`` when a prompt embedding replaces it."""
    tokenizer: Path | None
    """UMT5 tokenizer directory; ``None`` when a prompt embedding replaces it."""
    tiny_decoder: Path | None
    """Pinned TAEW2.2 checkpoint when turbo mode needs it."""
    turbo_lora: Path | None
    """Pinned Wan2.2 Turbo-LoRA when turbo mode needs it."""


def _require_file(path: Path, label: str, command: str) -> Path:
    resolved = path.expanduser().resolve()
    if not resolved.is_file():
        raise AssetError(f"{label} does not exist: {resolved}. Run `python {command}` to install it.")
    return resolved


def resolve_checkpoint(path: str | Path | None = None, model_dir: str | Path = "models") -> Path:
    if path is not None:
        resolved = Path(path).expanduser().resolve()
        if not resolved.is_file():
            raise AssetError(f"Checkpoint override does not exist: {resolved}")
        return resolved
    return _require_file(Path(model_dir) / CHECKPOINT, "Checkpoint", "scripts/download_model.py")


def resolve_regressor(path: str | Path | None = None, model_dir: str | Path = "models") -> Path:
    if path is not None:
        resolved = Path(path).expanduser().resolve()
        if not resolved.is_file():
            raise AssetError(f"MHR70 regressor override does not exist: {resolved}")
        return resolved
    return _require_file(Path(model_dir) / MHR70_REGRESSOR, "MHR70 regressor", "scripts/download_model.py")


def resolve_foreground_model(model_dir: str | Path = "models") -> Path:
    root = Path(model_dir).expanduser() / BIREFNET_DIR
    for relative in BIREFNET_FILES:
        _require_file(root / relative, "BiRefNet file", "scripts/download_model.py")
    return root.resolve()


def resolve_perceptual_vgg19(model_dir: str | Path = "models") -> Path:
    """Resolve the converted VGG-19 weights used by perceptual reconstruction."""

    return _require_file(
        Path(model_dir) / PERCEPTUAL_VGG19,
        "Perceptual VGG-19 weights",
        "scripts/download_model.py",
    )


def resolve_base_assets(
    model_dir: str | Path,
    settings: ModeSettings,
    *,
    have_prompt_embedding: bool = False,
) -> BaseAssets:
    """Resolve the local VAE, T5 encoder, and tokenizer.

    The encoder and its tokenizer serve only the single fixed prompt, so a
    caller that already holds the exported embedding needs neither and must
    not be asked for the 11 GB checkpoint.
    """

    root = Path(model_dir).expanduser()
    text_encoder: Path | None = None
    tokenizer: Path | None = None
    if not have_prompt_embedding:
        for relative in TOKENIZER_FILES:
            _require_file(root / relative, "Tokenizer file", "scripts/download_model.py")
        text_encoder = _require_file(root / TEXT_ENCODER, "Text encoder", "scripts/download_model.py")
        tokenizer = (root / TOKENIZER_DIR).expanduser().resolve()
    return BaseAssets(
        vae=_require_file(root / WAN_VAE, "VAE", "scripts/download_model.py"),
        text_encoder=text_encoder,
        tokenizer=tokenizer,
        tiny_decoder=(
            _require_file(root / TAEW2_2, "TAEW2.2 checkpoint", "scripts/download_model.py")
            if settings.tiny_decoders
            else None
        ),
        turbo_lora=(
            _require_file(root / TURBO_LORA, "Turbo-LoRA", "scripts/download_model.py")
            if settings.turbo_lora
            else None
        ),
    )