Refactor build_pipe to take a model key with fp16-variant fallback

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

1 files changed, 16 insertions, 9 deletions

diff --git a/generate.py b/generate.py
index eb20dc1..e0d6355 100644
--- a/generate.py
+++ b/generate.py
@@ -42,7 +42,7 @@ def resolve_model(key):
     return MODELS[key]
 
 
-def build_pipe():
+def build_pipe(model_key):
     if not torch.xpu.is_available():
         # Fail loud and specific. A silent CPU fallback would run but take
         # minutes per image, which is worse than a clear error.
@@ -54,15 +54,22 @@ def build_pipe():
         )
         sys.exit(2)
 
+    repo = resolve_model(model_key)
     vae = AutoencoderKL.from_pretrained(VAE, torch_dtype=torch.float16)
-    pipe = StableDiffusionXLPipeline.from_pretrained(
-        MODEL,
-        vae=vae,
-        torch_dtype=torch.float16,
-        use_safetensors=True,
-        variant="fp16",
-    ).to("xpu")
-    return pipe
+    try:
+        # SDXL base publishes an fp16 variant (smaller download).
+        pipe = StableDiffusionXLPipeline.from_pretrained(
+            repo, vae=vae, torch_dtype=torch.float16,
+            use_safetensors=True, variant="fp16",
+        )
+    except (ValueError, OSError, EnvironmentError):
+        # Many fine-tunes do not publish a separate fp16 variant; load the
+        # default files at fp16 dtype instead.
+        print(f"no fp16 variant for {repo}, loading default files", file=sys.stderr)
+        pipe = StableDiffusionXLPipeline.from_pretrained(
+            repo, vae=vae, torch_dtype=torch.float16, use_safetensors=True,
+        )
+    return pipe.to("xpu")
 
 
 def generate(pipe, prompt, negative, steps, guidance, seed, width, height):