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 _raise_if_403(err: Exception) -> None:
    """Re-raise with a clear message if the error is a 403 from the Hub."""
    if getattr(err, "response", None) is not None:
        status = getattr(err.response, "status_code", None)
        if status == 403:
            raise RuntimeError(
                "403 Forbidden from Hugging Face Hub. "
                "Set HF_TOKEN (Settings → Variables and secrets in this Space, or env var locally). "
                "Get a token at huggingface.co/settings/tokens (read access)."
            ) from err
    if "403" in str(err).lower() or "forbidden" in str(err).lower():
        raise RuntimeError(
            "403 Forbidden from Hugging Face Hub. "
            "Set HF_TOKEN (Settings → Variables and secrets in this Space, or env var locally). "
            "Get a token at huggingface.co/settings/tokens (read access)."
        ) from err


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")
    try:
        snapshot_download(
            repo_id=MODEL_ID,
            local_dir=str(out_dir),
            token=tok,
        )
    except Exception as e:
        _raise_if_403(e)
        raise
    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()