#!/usr/bin/env python3
# FILE: install_deps.py
# Description: Installs all dependencies and downloads models for GhostPack F1 Pro on Hugging Face Spaces with H200 GPU
# Version: 1.0.2
# Timestamp: 2025-07-02 04:50 CDT
# Author: Grok 3, built by xAI
# NOTE: Installs PyTorch 2.4.1 with CUDA 12.1, dependencies, and downloads models to /data/models
#       Requires HF_TOKEN set as environment variable or secret
#       Run this before app.py
#       Logs to /data/install_deps.log

import os
import sys
import subprocess
import logging

# Set up logging
logging.basicConfig(
    filename="/data/install_deps.log",
    level=logging.DEBUG,
    format="%(asctime)s %(levelname)s:%(message)s",
)
logger = logging.getLogger(__name__)

# Create log directory
os.makedirs("/data", exist_ok=True)
os.chmod("/data", 0o775)
logger.info("Starting dependency installation and model download")

# Function to run shell commands
def run_command(command, error_message):
    try:
        result = subprocess.run(command, check=True, capture_output=True, text=True)
        logger.info(f"Command succeeded: {' '.join(command)}\n{result.stdout}")
        print(f"Command succeeded: {' '.join(command)}")
        return True
    except subprocess.CalledProcessError as e:
        logger.error(f"{error_message}: {e}\n{e.stderr}")
        print(f"{error_message}: {e}")
        sys.exit(1)

# Update pip
logger.info("Upgrading pip...")
run_command(
    [sys.executable, "-m", "pip", "install", "--upgrade", "pip"],
    "ERROR: Failed to upgrade pip"
)

# Install PyTorch with CUDA 12.1
logger.info("Installing PyTorch 2.4.1 with CUDA 12.1...")
run_command(
    [sys.executable, "-m", "pip", "install", "torch==2.4.1+cu121", "--extra-index-url", "https://download.pytorch.org/whl/cu121"],
    "ERROR: Failed to install PyTorch"
)

# Install dependencies
logger.info("Installing required dependencies...")
with open("requirements.txt", "w") as f:
    f.write("""transformers==4.44.2
fastapi==0.115.0
uvicorn==0.30.6
gradio==4.44.0
python-multipart==0.0.9
diffusers==0.30.3
pydantic==2.9.2
einops==0.8.0
numpy==1.26.4
pillow==10.4.0
requests==2.32.3
colorama==0.4.6
# Note: diffusers_helper must be included as a folder in the Space's root directory (/diffusers_helper)
""")
run_command(
    [sys.executable, "-m", "pip", "install", "-r", "requirements.txt"],
    "ERROR: Failed to install dependencies"
)

# Install optional dependencies
logger.info("Installing optional dependencies (xformers, sage-attn, flash-attn)...")
for pkg, warn in [
    ("xformers", "WARNING: Failed to install xformers"),
    ("sage-attn", "WARNING: Failed to install sage-attn"),
    ("flash-attn", "WARNING: Failed to install flash-attn")
]:
    try:
        run_command([sys.executable, "-m", "pip", "install", pkg], warn)
    except SystemExit:
        logger.warning(warn)
        print(warn)

# Download models
logger.info("Downloading models to /data/models...")
MODEL_DIR = "/data/models"
HF_TOKEN = os.getenv('HF_TOKEN', 'your-hf-token')  # Set in Spaces secrets

os.makedirs(MODEL_DIR, exist_ok=True)
os.chmod(MODEL_DIR, 0o775)
logger.info(f"Created model directory: {MODEL_DIR}")

from transformers import LlamaModel, CLIPTextModel, LlamaTokenizerFast, CLIPTokenizer, SiglipImageProcessor, SiglipVisionModel
from diffusers import AutoencoderKLHunyuanVideo
from diffusers_helper.models.hunyuan_video_packed import HunyuanVideoTransformer3DModelPacked

def download_model(model_class, model_name, subfolder=None, **kwargs):
    try:
        logger.info(f"Downloading {model_name} (subfolder: {subfolder}) to {MODEL_DIR}")
        model = model_class.from_pretrained(
            model_name,
            subfolder=subfolder,
            token=HF_TOKEN,
            cache_dir=MODEL_DIR,
            local_files_only=False,
            **kwargs
        )
        logger.info(f"Successfully downloaded {model_name} (subfolder: {subfolder})")
    except Exception as e:
        logger.error(f"Failed to download {model_name} (subfolder: {subfolder}): {e}")
        print(f"Error: Failed to download {model_name} (subfolder: {subfolder}): {e}")
        sys.exit(1)

# Download HunyuanVideo components
try:
    download_model(LlamaModel, "hunyuanvideo-community/HunyuanVideo", subfolder="text_encoder", torch_dtype=torch.float16)
    download_model(CLIPTextModel, "hunyuanvideo-community/HunyuanVideo", subfolder="text_encoder_2", torch_dtype=torch.float16)
    download_model(LlamaTokenizerFast, "hunyuanvideo-community/HunyuanVideo", subfolder="tokenizer")
    download_model(CLIPTokenizer, "hunyuanvideo-community/HunyuanVideo", subfolder="tokenizer_2")
    download_model(AutoencoderKLHunyuanVideo, "hunyuanvideo-community/HunyuanVideo", subfolder="vae", torch_dtype=torch.float16)
except Exception as e:
    logger.error(f"Failed to download HunyuanVideo components: {e}")
    print(f"Error: Failed to download HunyuanVideo components: {e}")
    sys.exit(1)

# Download FramePack components
try:
    download_model(SiglipImageProcessor, "lllyasviel/flux_redux_bfl", subfolder="feature_extractor")
    download_model(SiglipVisionModel, "lllyasviel/flux_redux_bfl", subfolder="image_encoder", torch_dtype=torch.float16)
    download_model(HunyuanVideoTransformer3DModelPacked, "lllyasviel/FramePack_F1_I2V_HY_20250503", torch_dtype=torch.bfloat16)
except Exception as e:
    logger.error(f"Failed to download FramePack components: {e}")
    print(f"Error: Failed to download FramePack components: {e}")
    sys.exit(1)

logger.info("Dependency installation and model download completed successfully")
print("Dependency installation and model download completed successfully")