MPS support

app.py

@@ -17,18 +17,27 @@ from share_btn import community_icon_html, loading_icon_html, share_js
 # load pipelines
 # sd_model_id = "runwayml/stable-diffusion-v1-5"
 sd_model_id = "stabilityai/stable-diffusion-2-1-base"
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", torch_dtype=torch.float16)
-pipe = SemanticStableDiffusionImg2ImgPipeline_DPMSolver.from_pretrained(sd_model_id,vae=vae,torch_dtype=torch.float16, safety_checker=None, requires_safety_checker=False).to(device)
+#device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+torch_dtype = torch.float16
+if torch.cuda.is_available():
+  device = "cuda"
+elif torch.backends.mps.is_available():
+  device = "mps"
+  torch_dtype = torch.float32
+else:
+  device = "cpu"
+vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", torch_dtype=torch_dtype)
+pipe = SemanticStableDiffusionImg2ImgPipeline_DPMSolver.from_pretrained(sd_model_id,vae=vae,torch_dtype=torch_dtype, safety_checker=None, requires_safety_checker=False).to(device)
 pipe.scheduler = DPMSolverMultistepSchedulerInject.from_pretrained(sd_model_id, subfolder="scheduler"
                                                              , algorithm_type="sde-dpmsolver++", solver_order=2)
 
 blip_processor = AutoProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
-blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base",torch_dtype=torch.float16).to(device)
+blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base",torch_dtype=torch_dtype).to(device)
 
 ## IMAGE CPATIONING ##
 def caption_image(input_image):
-    inputs = blip_processor(images=input_image, return_tensors="pt").to(device, torch.float16)
+    inputs = blip_processor(images=input_image, return_tensors="pt").to(device, torch_dtype)
     pixel_values = inputs.pixel_values
 
     generated_ids = blip_model.generate(pixel_values=pixel_values, max_length=50)
@@ -228,7 +237,7 @@ def randomize_seed_fn(seed, is_random):
 
 def seed_everything(seed):
     torch.manual_seed(seed)
-    torch.cuda.manual_seed(seed)
+#    torch.cuda.manual_seed(seed)
     random.seed(seed)
     np.random.seed(seed)
 
@@ -902,4 +911,4 @@ with gr.Blocks(css="style.css") as demo:
     )
 
 demo.queue(default_concurrency_limit=1)
-demo.launch()
+demo.launch()