From 27987c8bdc0325f663afd81813f6a9348ab47621 Mon Sep 17 00:00:00 2001 From: "Danilo M." 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 --- generate.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) (limited to '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 -- cgit v1.2.3