scripts/download_sd_model.py

"""
Download Stable Diffusion v1.5 to ./weights/sd-v1-5 for offline use.
Run once (with internet). App auto-uses this folder if present.

  python -m scripts.download_sd_model

Set HF_TOKEN=your_token if behind firewall. Can also use "Download model" in the app sidebar.
"""

import os
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
MODEL_ID = "runwayml/stable-diffusion-v1-5"
DEFAULT_LOCAL_DIR = ROOT / "weights" / "sd-v1-5"


def download_sd_model(local_dir: str | Path | None = None, token: str | None = None) -> str:
    """Download runwayml/stable-diffusion-v1-5 to local_dir. Returns path on success, raises on failure."""
    from huggingface_hub import snapshot_download
    out_dir = Path(local_dir or DEFAULT_LOCAL_DIR)
    out_dir.mkdir(parents=True, exist_ok=True)
    tok = token or os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
    snapshot_download(
        repo_id=MODEL_ID,
        local_dir=str(out_dir),
        token=tok,
    )
    return str(out_dir.resolve())


def main() -> None:
    out_dir = os.environ.get("SD_MODEL_PATH", str(DEFAULT_LOCAL_DIR))
    try:
        path = download_sd_model(local_dir=out_dir)
        print(f"Done. App will use: {path}")
        print("Run: streamlit run app.py")
    except Exception as e:
        print(f"Download failed: {e}", file=sys.stderr)
        print("Set HF_TOKEN=your_token if behind firewall (huggingface.co/settings/tokens)", file=sys.stderr)
        raise SystemExit(1) from e


if __name__ == "__main__":
    main()