import io
import modal
import os
import subprocess
import sys
import time
import shutil
from datetime import datetime
from typing import Dict, List, Tuple

# --- Modal Setup ---
app = modal.App("hunyuan-training")
volume = modal.Volume.from_name(
    "hunyuan-volume", create_if_missing=True
)

# Using a Python-based CUDA image as the base
image = (
    modal.Image.from_registry("pytorch/pytorch:2.0.1-cuda11.7-cudnn8-devel")
    .apt_install(
        "git",
        "wget",
        "curl",
        "build-essential"
    )
    .pip_install(
        "tomli",
        "toml",
        "diffusers",
        "transformers",
        "accelerate",
        "safetensors",
        "tensorboard",
        "ftfy",
        "protobuf==3.20.1",
        "deepspeed",
        "huggingface_hub[hf_transfer]",
    )
    .run_commands(
        "wget https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.1.post4/flash_attn-2.7.1.post4+cu11torch2.4cxx11abiFALSE-cp310-cp310-linux_x86_64.whl",
        "pip install flash_attn-2.7.1.post4+cu11torch2.4cxx11abiFALSE-cp310-cp310-linux_x86_64.whl"
    )
    .env({
        "HF_HUB_ENABLE_HF_TRANSFER": "1",
        "CUDA_HOME": "/usr/local/cuda"
    })
)

# --- Constants ---
GIT_REPO = "https://github.com/tdrussell/diffusion-pipe.git"
REMOTE_ROOT = "/root/volume"

# --- Helper Functions ---
@app.function(image=image)
def get_python_version():
    print(f"Python version: {sys.version}")
    print(f"Executable: {sys.executable}")
    return sys.version

@app.function(image=image, volumes={REMOTE_ROOT: volume})
def hf_download_repo(repo_id: str, model_type: str, ignore_patterns: List[str] = None):
    from huggingface_hub import snapshot_download

    snapshot_download(
        repo_id=repo_id,
        local_dir=f"/root/volume/models/{model_type}",
        ignore_patterns=ignore_patterns or [],
    )
    print(f"Repository {repo_id} downloaded to /root/volume/models/{model_type}")

@app.function(image=image, volumes={REMOTE_ROOT: volume})
def hf_download_file(repo_id: str, filename: str, model_type: str):
    from huggingface_hub import hf_hub_download

    hf_hub_download(
        repo_id=repo_id,
        filename=filename,
        local_dir=f"/root/volume/models/{model_type}",
    )
    print(f"File {filename} from {repo_id} downloaded to /root/volume/models/{model_type}")

# --- Modal Functions ---
@app.function(image=image, volumes={REMOTE_ROOT: volume})
def download_models():
    models_to_download_from_hf_repos: List[Tuple[str, str, List[str]]] = [
        ("Kijai/llava-llama-3-8b-text-encoder-tokenizer", "llm"),
        ("openai/clip-vit-large-patch14", "clip", ["flax_model.msgpack", "pytorch_model.bin", "tf_model.h5"]),  # Exclude these files
    ]

    models_to_download_from_hf_files: List[Tuple[str, str, str]] = [
        ("Kijai/HunyuanVideo_comfy", "hunyuan_video_720_cfgdistill_fp8_e4m3fn.safetensors", "hunyuan"),
        ("Kijai/HunyuanVideo_comfy", "hunyuan_video_vae_bf16.safetensors", "hunyuan"),
    ]

    # Download entire HF repos
    print("Downloading entire Hugging Face repositories...")
    for repo in models_to_download_from_hf_repos:
        if len(repo) == 3:
            repo_id, model_type, ignore_patterns = repo
        elif len(repo) == 2:
            repo_id, model_type = repo
            ignore_patterns = []
        else:
            raise ValueError(f"Unexpected tuple length: {len(repo)}")
    
        hf_download_repo.remote(repo_id, model_type, ignore_patterns)

    # Download individual files from Hugging Face
    print("Downloading individual files from Hugging Face...")
    for repo_id, filename, model_type in models_to_download_from_hf_files:
        hf_download_file.remote(repo_id, filename, model_type)

@app.function(
    image=image,
    gpu="L40S",
    volumes={REMOTE_ROOT: volume},
    timeout=86400,
)
def train():
    import os
    import subprocess

    subprocess.run(["nvcc", "--version"], check=True)
    subprocess.run(["nvidia-smi"], check=True)

    # Clone the diffusion-pipe repository
    if not os.path.exists('/root/diffusion-pipe'):
        subprocess.run(['git', 'clone', '--recurse-submodules', GIT_REPO, '/root/diffusion-pipe'], check=True)
    else:
        print("diffusion-pipe repository already cloned.")

    # Change directory to diffusion-pipe
    os.chdir('/root/diffusion-pipe')

    # Install requirements
    subprocess.run(['pip', 'install', '-r', 'requirements.txt'], check=True)

    # Prepare 'config.toml' and 'dataset.toml' with adjusted paths

    # Prepare 'config.toml'
    config_toml_content = """
# Dataset config file.
output_dir = '/root/volume/training_output'
dataset = 'dataset.toml'

# Training settings
epochs = 256
micro_batch_size_per_gpu = 2
pipeline_stages = 1
gradient_accumulation_steps = 2
gradient_clipping = 1.0
warmup_steps = 128

# eval settings
eval_every_n_epochs = 8
eval_before_first_step = true
eval_micro_batch_size_per_gpu = 2
eval_gradient_accumulation_steps = 1

# misc settings
save_every_n_epochs = 5
checkpoint_every_n_minutes = 30
activation_checkpointing = true
partition_method = 'parameters'
save_dtype = 'bfloat16'
caching_batch_size = 1
steps_per_print = 1
video_clip_mode = 'single_middle'

[model]
type = 'hunyuan-video'
transformer_path = '/root/volume/models/hunyuan/hunyuan_video_720_cfgdistill_fp8_e4m3fn.safetensors'
vae_path = '/root/volume/models/hunyuan/hunyuan_video_vae_bf16.safetensors'
llm_path = '/root/volume/models/llm'
clip_path = '/root/volume/models/clip'
dtype = 'bfloat16'
transformer_dtype = 'float8'
timestep_sample_method = 'logit_normal'

[adapter]
type = 'lora'
rank = 32
dtype = 'bfloat16'

[optimizer]
type = 'adamw_optimi'
lr = 4e-4
betas = [0.9, 0.99]
weight_decay = 0.02
eps = 6e-4
"""
    with open('config.toml', 'w') as f:
        f.write(config_toml_content)

    # Prepare 'dataset.toml'
    dataset_toml_content = """
# Resolution settings.
resolutions = [1024]

# Aspect ratio bucketing settings
enable_ar_bucket = true
min_ar = 0.5
max_ar = 2.5
num_ar_buckets = 7

# Frame buckets (1 is for images)
frame_buckets = [1]

[[directory]]
# Path to your dataset
path = '/root/volume/training_data/images'
num_repeats = 4
"""
    with open('dataset.toml', 'w') as f:
        f.write(dataset_toml_content)

    # Run the training command
    os.environ['NCCL_P2P_DISABLE'] = "1"
    os.environ['NCCL_IB_DISABLE'] = "1"
    os.environ['CUDA_VISIBLE_DEVICES'] = "0"

    print("Starting training...")
    subprocess.run(['deepspeed', '--num_gpus=1', 'train.py', '--deepspeed', '--config', 'config.toml'], check=True)
    print("Training completed.")

# --- Entry Points ---
@app.local_entrypoint()
def get_python_version_entrypoint():
    get_python_version.remote()

@app.local_entrypoint()
def download_models_entrypoint():
    # First, download all necessary models
    print("Downloading models...")
    download_models.remote()
    print("Models downloaded.")
    
@app.local_entrypoint()
def train_entrypoint():
    # Start training
    train.remote()

As you can see, only subprocess.run(['deepspeed', '--num_gpus=1', 'train.py', '--deepspeed', '--config', 'config.toml'], check= True) in the train() does need to be run on the GPU, and the others are just a waste of GPU running time. Fix it please.