# setup.py
"""
This script runs automatically when the Space is built.
It pulls the entire SDXL‑i2i repo into the cache so the runtime never
has to hit the Hub again.
"""

import os
import torch
from diffusers import StableDiffusionXLImg2ImgPipeline

MODEL_ID = "stabilityai/stable-diffusion-xl-1.0"

def download():
    print(f"[setup] Downloading {MODEL_ID}…")
    # CPU‑friendly download; the full checkpoint (~10 GB) will be cached
    StableDiffusionXLImg2ImgPipeline.from_pretrained(
        MODEL_ID,
        torch_dtype=torch.float32,      # no GPU needed for the download
        use_safetensors=True,
        variant=None,                  # 0.x repo, no special variant
    )
    print("[setup] ✅ Download finished")

if __name__ == "__main__":
    # Ensure the cache lives in the default location
    # (you can change HF_HOME or HUGGINGFACE_HUB_CACHE if you want)
    download()