From 27987c8bdc0325f663afd81813f6a9348ab47621 Mon Sep 17 00:00:00 2001
From: "Danilo M." <danix@danix.xyz>
Date: Wed, 8 Jul 2026 15:40:25 +0200
Subject: Add model registry and resolve_model with valid-key errors

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 generate.py      | 14 +++++++++++++-
 test_generate.py | 29 +++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+), 1 deletion(-)
 create mode 100644 test_generate.py

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
diff --git a/test_generate.py b/test_generate.py
new file mode 100644
index 0000000..c2425d5
--- /dev/null
+++ b/test_generate.py
@@ -0,0 +1,29 @@
+import generate
+
+
+def test_registry_integrity():
+    assert generate.DEFAULT_MODEL in generate.MODELS
+    for key, repo in generate.MODELS.items():
+        assert isinstance(key, str) and key
+        assert isinstance(repo, str) and "/" in repo
+
+
+def test_resolve_known_key():
+    assert generate.resolve_model("sdxl") == generate.MODELS["sdxl"]
+    assert generate.resolve_model("realvis") == generate.MODELS["realvis"]
+
+
+def test_resolve_unknown_key_raises_listing_valid():
+    try:
+        generate.resolve_model("nope")
+    except SystemExit as e:
+        assert "sdxl" in str(e) and "realvis" in str(e) and "pony" in str(e)
+    else:
+        assert False, "expected SystemExit for unknown key"
+
+
+if __name__ == "__main__":
+    test_registry_integrity()
+    test_resolve_known_key()
+    test_resolve_unknown_key_raises_listing_valid()
+    print("ok")
-- 
cgit v1.2.3