diff options
Diffstat (limited to 'generate.py')
| -rw-r--r-- | generate.py | 25 |
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): |