Add model registry and resolve_model with valid-key errors

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

1 files changed, 13 insertions, 1 deletions

diff --git a/generate.py b/generate.py
index 5eed627..eb20dc1 100644
--- a/generate.py
+++ b/generate.py
@@ -24,12 +24,24 @@ from pathlib import Path
 import torch
 from diffusers import StableDiffusionXLPipeline, AutoencoderKL
 
-MODEL = "stabilityai/stable-diffusion-xl-base-1.0"
+MODELS = {
+    "sdxl":    "stabilityai/stable-diffusion-xl-base-1.0",  # default, general
+    "realvis": "SG161222/RealVisXL_V5.0",                   # photoreal, permissive
+    "pony":    "AstraliteHeart/pony-diffusion-v6-xl",       # unrestricted, versatile
+}
+DEFAULT_MODEL = "sdxl"
 # fp16-fix VAE avoids the fp32 upcast that OOMs a 12GB card at decode time.
 VAE = "madebyollin/sdxl-vae-fp16-fix"
 OUT_DIR = Path("/out")
 
 
+def resolve_model(key):
+    if key not in MODELS:
+        valid = ", ".join(sorted(MODELS))
+        raise SystemExit(f"unknown model '{key}'. valid: {valid}")
+    return MODELS[key]
+
+
 def build_pipe():
     if not torch.xpu.is_available():
         # Fail loud and specific. A silent CPU fallback would run but take