styletransfer-ComfyUI-TRELLIS2_Motion/utils/download_utils.py

"""Download utilities for TRELLIS2 model with progress tracking."""

import logging
import os
from pathlib import Path
from typing import Optional

logger = logging.getLogger(__name__)

# TRELLIS2 model configuration
TRELLIS_HF_REPO = "microsoft/TRELLIS-2-4B"
TRELLIS_CACHE_DIR = Path.home() / ".cache" / "huggingface" / "hub"

# Progress bar fallback for when ComfyUI is not available
try:
    import comfy.utils
    COMFY_AVAILABLE = True
except ImportError:
    COMFY_AVAILABLE = False


class ProgressBar:
    """Fallback progress bar when ComfyUI is not available."""

    def __init__(self, total: int):
        self.total = total
        self.current = 0

    def update_absolute(self, current: int, total: Optional[int] = None, desc: str = ""):
        self.current = current
        total = total or self.total
        pct = (current / total * 100) if total > 0 else 0
        logger.info(f"{desc}: {pct:.1f}% ({current}/{total})")

    def update(self, value: int):
        self.current += value
        pct = (self.current / self.total * 100) if self.total > 0 else 0
        logger.info(f"Progress: {pct:.1f}%")


def get_progress_bar(total: int):
    """Get a progress bar instance (ComfyUI or fallback)."""
    if COMFY_AVAILABLE:
        return comfy.utils.ProgressBar(total)
    return ProgressBar(total)


def get_trellis_cache_dir() -> Path:
    """Get the TRELLIS model cache directory."""
    cache_dir = TRELLIS_CACHE_DIR
    cache_dir.mkdir(parents=True, exist_ok=True)
    return cache_dir


def is_trellis_model_downloaded() -> bool:
    """Check if the TRELLIS model is already downloaded."""
    try:
        from huggingface_hub import snapshot_download, HfFileSystem

        fs = HfFileSystem()
        # Check if model files exist in cache
        cache_dir = get_trellis_cache_dir()

        # Look for model files
        model_dir = cache_dir / f"models--{TRELLIS_HF_REPO.replace('/', '--')}"
        if model_dir.exists():
            # Check for key files
            snapshots = model_dir / "snapshots"
            if snapshots.exists():
                # Check any snapshot directory
                for snapshot in snapshots.iterdir():
                    if (snapshot / "config.json").exists():
                        return True
        return False
    except Exception as e:
        logger.warning(f"Could not check TRELLIS model status: {e}")
        return False


def ensure_trellis_model(pbar: Optional[ProgressBar] = None) -> str:
    """
    Ensure the TRELLIS model is available, downloading if necessary.

    Args:
        pbar: Optional progress bar for download progress

    Returns:
        Path to the TRELLIS model directory
    """
    try:
        from huggingface_hub import snapshot_download
    except ImportError:
        raise ImportError(
            "huggingface_hub is required for downloading TRELLIS2 model. "
            "Install with: pip install huggingface_hub"
        )

    if is_trellis_model_downloaded():
        logger.info("TRELLIS2 model found in cache")
        cache_dir = get_trellis_cache_dir()
        model_dir = cache_dir / f"models--{TRELLIS_HF_REPO.replace('/', '--')}"
        # Return the actual model path
        snapshots = model_dir / "snapshots"
        for snapshot in snapshots.iterdir():
            if (snapshot / "config.json").exists():
                return str(snapshot)
        return str(model_dir)

    logger.info(f"Downloading TRELLIS2 model from {TRELLIS_HF_REPO}...")
    logger.info("This may take a while (~8GB download)")

    if pbar:
        pbar.update_absolute(0, 100, "Downloading TRELLIS2 model")

    try:
        model_path = snapshot_download(
            TRELLIS_HF_REPO,
            local_dir_use_symlinks=False,
        )

        if pbar:
            pbar.update_absolute(100, 100, "TRELLIS2 download complete")

        logger.info(f"TRELLIS2 model downloaded to {model_path}")
        return model_path

    except Exception as e:
        logger.error(f"Failed to download TRELLIS2 model: {e}")
        raise RuntimeError(
            f"Failed to download TRELLIS2 model from {TRELLIS_HF_REPO}. "
            f"Error: {e}\n"
            "Try manually downloading with:\n"
            f"  huggingface-cli download {TRELLIS_HF_REPO}"
        )


def check_trellis_dependencies() -> dict:
    """
    Check if all TRELLIS2 dependencies are available.

    Returns:
        Dictionary with status of each dependency
    """
    status = {
        "trellis": False,
        "model_downloaded": False,
        "nvdiffrast": False,
        "cuda": False,
        "rembg": False,
    }

    # Check TRELLIS package
    try:
        import trellis
        status["trellis"] = True
    except ImportError:
        pass

    # Check model
    status["model_downloaded"] = is_trellis_model_downloaded()

    # Check nvdiffrast
    try:
        import nvdiffrast
        status["nvdiffrast"] = True
    except ImportError:
        pass

    # Check CUDA
    try:
        import torch
        status["cuda"] = torch.cuda.is_available()
    except ImportError:
        pass

    # Check rembg
    try:
        import rembg
        status["rembg"] = True
    except ImportError:
        pass

    return status


def print_dependency_status():
    """Print the status of all dependencies."""
    status = check_trellis_dependencies()

    print("\nTRELLIS2 Dependency Status:")
    print("-" * 40)
    for name, available in status.items():
        icon = "OK" if available else "MISSING"
        print(f"  [{icon}] {name}")

    if not all(status.values()):
        print("\nTo install missing dependencies:")
        if not status["trellis"]:
            print("  pip install git+https://github.com/microsoft/TRELLIS.git")
        if not status["nvdiffrast"]:
            print("  pip install git+https://github.com/NVlabs/nvdiffrast.git")
        if not status["rembg"]:
            print("  pip install rembg")
        if not status["model_downloaded"]:
            print(f"  huggingface-cli download {TRELLIS_HF_REPO}")