import os
import subprocess

HF_TOKEN = os.getenv("HF_TOKEN")
HEADERS = f'--header="Authorization: Bearer {HF_TOKEN}"'
BASE_URLS = {
    "text_encoder": "https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/text_encoder/",
    "clip": "https://huggingface.co/openai/clip-vit-large-patch14/resolve/main/",
    "stable_diffusion": "https://huggingface.co/stabilityai/stable-diffusion-3.5-large/resolve/main/"
}

FILES = {
    "text_encoder": [
        "model.fp16.safetensors",
        "model.safetensors"
    ],
    "clip": [
        "model.safetensors"
    ],
    "stable_diffusion": [
        "sd3.5_large.safetensors"
    ]
}

DOWNLOAD_PATHS = {
    "text_encoder": "models/text_encoder",
    "clip": "models/clip",
    "stable_diffusion": "models/Stable-diffusion"
}

def download_models():
    for category, files in FILES.items():
        path = os.path.join("/stable-diffusion-webui-forge", DOWNLOAD_PATHS[category])
        os.makedirs(path, exist_ok=True)
        for file in files:
            url = BASE_URLS[category] + file
            command = f'wget {HEADERS} -O {os.path.join(path, file)} {url}'
            print(f"Downloading {file} to {path}")
            subprocess.run(command, shell=True, check=True)

if __name__ == "__main__":
    print("Starting model download...")
    download_models()
    print("Download completed. Launching WebUI...")
    subprocess.run(["python3", "launch.py", "--share", "--disable-torch-cuda"])