import os
import subprocess
from huggingface_hub import hf_hub_download

model_dir_stablediffusion = "models/Stable-diffusion/"
os.makedirs(model_dir_stablediffusion, exist_ok=True)

model_dir_controlnet = "models/ControlNet/"
os.makedirs(model_dir_controlnet, exist_ok=True)

model_dir_ESRGAN = "models/ESRGAN"
os.makedirs(model_dir_ESRGAN, exist_ok=True)

model_dir_VAEapprox = "models/VAE-approx"
os.makedirs(model_dir_VAEapprox, exist_ok=True)

model_dir_lora = "models/Lora"
os.makedirs(model_dir_lora, exist_ok=True)

model_dir_karlo = "models/karlo"
os.makedirs(model_dir_karlo, exist_ok=True)

model_path = hf_hub_download(repo_id="MattiaSavoia/ImageEnhancerModels", filename="juggernaut_reborn.safetensors", local_dir=model_dir_stablediffusion)
model_path = hf_hub_download(repo_id="MattiaSavoia/ImageEnhancerModels", filename="control_v11f1e_sd15_tile.pth", local_dir=model_dir_controlnet)
model_path = hf_hub_download(repo_id="MattiaSavoia/ImageEnhancerModels", filename="4x-UltraSharp.pth", local_dir=model_dir_ESRGAN)
model_path = hf_hub_download(repo_id="MattiaSavoia/ImageEnhancerModels", filename="model.pt", local_dir=model_dir_VAEapprox)
model_path = hf_hub_download(repo_id="MattiaSavoia/ImageEnhancerModels", filename="more_details.safetensors", local_dir=model_dir_lora)
model_path = hf_hub_download(repo_id="MattiaSavoia/ImageEnhancerModels", filename="SDXLrender_v2.0.safetensors", local_dir=model_dir_lora)
model_path = hf_hub_download(repo_id="MattiaSavoia/ImageEnhancerModels", filename="ViT-L-14_stats.th", local_dir=model_dir_karlo)

# Clona le estensioni se non esistono già
extensions_dir = "extensions"
os.makedirs(extensions_dir, exist_ok=True)

# Clona multidiffusion-upscaler
multidiffusion_dir = os.path.join(extensions_dir, "multidiffusion-upscaler-for-automatic1111")
if not os.path.exists(multidiffusion_dir):
    subprocess.run(["git", "clone", "https://github.com/pkuliyi2015/multidiffusion-upscaler-for-automatic1111.git", multidiffusion_dir])

# Clona controlnet
controlnet_dir = os.path.join(extensions_dir, "sd-webui-controlnet")
if not os.path.exists(controlnet_dir):
    subprocess.run(["git", "clone", "https://github.com/Mikubill/sd-webui-controlnet.git", controlnet_dir])

# Hugging Face Spaces usa la variabile d'ambiente PORT
port = os.getenv("PORT", "7860")

# Aggiungi il parametro --skip-torch-cuda-test per forzare l'uso della CPU
subprocess.run(["python3", "launch.py", "--skip-torch-cuda-test", "--port", port, "--api"])