From 4d137f35e724e98128916709c9fcf154190c744c Mon Sep 17 00:00:00 2001 From: "Danilo M." Date: Wed, 8 Jul 2026 15:44:59 +0200 Subject: Refactor build_pipe to take a model key with fp16-variant fallback Co-Authored-By: Claude Opus 4.8 --- generate.py | 25 ++++++++++++++++--------- 1 file changed, 16 insertions(+), 9 deletions(-) (limited to 'generate.py') 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): -- cgit v1.2.3