import os
import subprocess
import requests
from huggingface_hub import hf_hub_download

# Create directories
os.makedirs("/app/ComfyUI/models/checkpoints", exist_ok=True)
os.makedirs("/app/ComfyUI/models/controlnet", exist_ok=True)
os.makedirs("/app/ComfyUI/models/ipadapter", exist_ok=True)
os.makedirs("/app/ComfyUI/models/pulid", exist_ok=True)
os.makedirs("/app/ComfyUI/models/clip_vision", exist_ok=True)

# List of models to download
models = [
    # SDXL Checkpoint
    {"repo_id": "Lykon/DreamShaper-XL-Turbo", "filename": "dreamshaperXL_turboDpmppSDEKarras.safetensors", "local_path": "/app/ComfyUI/models/checkpoints/sdxl/dreamshaperXL_turboDpmppSDEKarras.safetensors"},
    
    # ControlNet model
    {"repo_id": "thibaud/controlnet-openpose-sdxl-1.0", "filename": "diffusion_pytorch_model.safetensors", "local_path": "/app/ComfyUI/models/controlnet/sdxl/thibaud_xl_openpose.safetensors"},
    
    # IP-Adapter
    {"repo_id": "h94/IP-Adapter", "filename": "ip-adapter_sdxl.safetensors", "local_path": "/app/ComfyUI/models/ipadapter/sdxl/ip-adapter_sdxl.safetensors"},
    
    # PuLID model
    {"repo_id": "l7dev/ip-adapter-pulid-sdxl", "filename": "ip-adapter_pulid_sdxl_fp16.safetensors", "local_path": "/app/ComfyUI/models/pulid/ip-adapter_pulid_sdxl_fp16.safetensors"},
    
    # CLIP Vision
    {"repo_id": "h94/IP-Adapter", "filename": "CLIP-ViT-bigG-14-laion2B-39B-b160k.safetensors", "local_path": "/app/ComfyUI/models/clip_vision/CLIP-ViT-bigG-14-laion2B-39B-b160k.safetensors"}
]

# Download each model
for model in models:
    try:
        os.makedirs(os.path.dirname(model["local_path"]), exist_ok=True)
        
        print(f"Downloading {model['filename']} from {model['repo_id']}...")
        
        # Get token from environment variable if available
        token = os.environ.get("HF_TOKEN", None)
        
        hf_hub_download(
            repo_id=model["repo_id"],
            filename=model["filename"],
            local_dir=os.path.dirname(model["local_path"]),
            local_dir_use_symlinks=False,
            token=token
        )
        
        # Rename if necessary
        downloaded_path = os.path.join(os.path.dirname(model["local_path"]), model["filename"])
        if downloaded_path != model["local_path"]:
            os.rename(downloaded_path, model["local_path"])
            
        print(f"Successfully downloaded {model['filename']}")
    except Exception as e:
        print(f"Error downloading {model['filename']}: {str(e)}")