jonstreeter-ComfyUI-Reference-Based-Video-Colorization/auto_installer.py

"""
Auto-Installer for Modern Components
Automatically installs missing dependencies when needed.
"""

import sys
import subprocess
import importlib
from typing import Optional

# Cache to avoid repeated installation attempts
_INSTALLATION_CACHE = {}


def check_package_installed(package_name: str) -> bool:
    """Check if a Python package is installed."""
    try:
        importlib.import_module(package_name)
        return True
    except ImportError:
        return False


def install_package(package_spec: str, display_name: Optional[str] = None) -> bool:
    """
    Silently install a package via pip.

    Args:
        package_spec: Package specification (e.g., "timm>=0.9.0" or "git+https://...")
        display_name: Human-readable name for logging

    Returns:
        True if installation successful, False otherwise
    """
    if display_name is None:
        display_name = package_spec.split(">=")[0].split("==")[0]

    # Check cache
    if package_spec in _INSTALLATION_CACHE:
        return _INSTALLATION_CACHE[package_spec]

    print(f"[AutoInstall] Installing {display_name}...")
    try:
        # Run pip install
        subprocess.check_call(
            [sys.executable, "-m", "pip", "install", package_spec],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.PIPE
        )
        print(f"[AutoInstall] ✓ {display_name} installed successfully")
        _INSTALLATION_CACHE[package_spec] = True
        return True
    except subprocess.CalledProcessError as e:
        print(f"[AutoInstall] ✗ Failed to install {display_name}: {e}")
        _INSTALLATION_CACHE[package_spec] = False
        return False
    except Exception as e:
        print(f"[AutoInstall] ✗ Unexpected error installing {display_name}: {e}")
        _INSTALLATION_CACHE[package_spec] = False
        return False


def ensure_timm() -> bool:
    """Ensure timm is installed (for DINOv2)."""
    if check_package_installed('timm'):
        return True

    print("[AutoInstall] timm not found (required for DINOv2)")
    return install_package("timm>=0.9.0", "timm (DINOv2 support)")


def ensure_clip() -> bool:
    """Ensure CLIP is installed."""
    if check_package_installed('clip'):
        return True

    print("[AutoInstall] CLIP not found (required for text-guided colorization)")
    return install_package("git+https://github.com/openai/CLIP.git", "CLIP")


def ensure_color_matcher() -> bool:
    """Ensure color-matcher is installed."""
    # Check for the actual package name
    try:
        from color_matcher import ColorMatcher
        return True
    except ImportError:
        pass

    print("[AutoInstall] color-matcher not found (required for color matching post-processing)")
    return install_package("color-matcher>=0.3.0", "color-matcher")


def ensure_opencv_contrib() -> bool:
    """Ensure opencv-contrib-python is installed (for WLS/Guided filters)."""
    try:
        import cv2
        # Check if ximgproc is available
        if hasattr(cv2, 'ximgproc'):
            return True
    except ImportError:
        pass

    print("[AutoInstall] opencv-contrib-python not found (required for advanced filters)")
    # Uninstall opencv-python first to avoid conflicts
    try:
        subprocess.check_call(
            [sys.executable, "-m", "pip", "uninstall", "-y", "opencv-python"],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL
        )
    except:
        pass

    return install_package("opencv-contrib-python>=4.7.0", "opencv-contrib-python")


def ensure_dependencies_for_encoder(encoder_type: str) -> bool:
    """
    Ensure all dependencies for a given encoder are installed.

    Args:
        encoder_type: Encoder name (e.g., 'dinov2_vitb', 'clip_vitb')

    Returns:
        True if all dependencies are available, False otherwise
    """
    if encoder_type == 'vgg19':
        return True  # No extra dependencies

    elif encoder_type.startswith('dinov2'):
        return ensure_timm()

    elif encoder_type.startswith('clip'):
        return ensure_clip()

    else:
        return True  # Unknown encoder, assume no dependencies


def install_all_modern_components() -> dict:
    """
    Install all modern components at once.

    Returns:
        Dictionary with installation status for each component
    """
    print("[AutoInstall] Installing all modern components...")
    print("[AutoInstall] This may take a few minutes on first run...")

    results = {
        'timm': ensure_timm(),
        'clip': ensure_clip(),
    }

    # Summary
    print("\n[AutoInstall] Installation Summary:")
    for component, success in results.items():
        status = "✓" if success else "✗"
        print(f"[AutoInstall]   {status} {component}")

    all_success = all(results.values())
    if all_success:
        print("[AutoInstall] ✓ All components installed successfully!")
    else:
        print("[AutoInstall] ⚠ Some components failed to install")

    return results


if __name__ == "__main__":
    # When run directly, install everything
    install_all_modern_components()