update

app.py

@@ -7,6 +7,13 @@ import math
 import spaces 
 import torch 
 
+if torch.backends.mps.is_available():
+    DEVICE = "mps"
+elif torch.cuda.is_available():
+    DEVICE = "cuda"
+else:
+    DEVICE = "cpu"
+
 edit_file = hf_hub_download(repo_id="stabilityai/cosxl", filename="cosxl_edit.safetensors")
 normal_file = hf_hub_download(repo_id="stabilityai/cosxl", filename="cosxl.safetensors")
 
@@ -30,11 +37,11 @@ pipe_edit = CosStableDiffusionXLInstructPix2PixPipeline.from_single_file(
     edit_file, num_in_channels=8
 )
 pipe_edit.scheduler = EDMEulerScheduler(sigma_min=0.002, sigma_max=120.0, sigma_data=1.0, prediction_type="v_prediction")
-pipe_edit.to("cuda")
+pipe_edit.to(DEVICE)
 
 pipe_normal = StableDiffusionXLPipeline.from_single_file(normal_file, torch_dtype=torch.float16)
 pipe_normal.scheduler = EDMEulerScheduler(sigma_min=0.002, sigma_max=120.0, sigma_data=1.0, prediction_type="v_prediction")
-pipe_normal.to("cuda")
+pipe_normal.to(DEVICE)
 
 @spaces.GPU
 def run_normal(prompt, negative_prompt="", guidance_scale=7, progress=gr.Progress(track_tqdm=True)):