add mps to diffusers

dottxt-ai · Jun 7, 2023 · 28da4c5 · 28da4c5
1 parent acda7f2
commit 28da4c5
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/outlines/models/hf_diffusers.py b/outlines/models/hf_diffusers.py
@@ -46,6 +46,8 @@ def call_stable_diffusion_pipeline(
     pipe = StableDiffusionPipeline.from_pretrained(model_name)
     if torch.cuda.is_available():
         pipe = pipe.to("cuda")
+    elif torch.backends.mps.is_available():
+        pipe = pipe.to("mps")
 
     images = pipe(prompt, num_images_per_prompt=samples).images
     if not isinstance(images, list):