update

app.py

@@ -17,6 +17,14 @@ from src.pix2pix_turbo import Pix2Pix_Turbo
 
 model = Pix2Pix_Turbo("sketch_to_image_stochastic")
 
+if torch.backends.mps.is_available():
+  device = "mps"
+  #torch_dtype = torch.float32
+elif torch.cuda.is_available():
+  device = "cuda"
+else:
+  device = "cpu"
+
 style_list = [
     {
         "name": "No Style",
@@ -86,7 +94,8 @@ def run(image, prompt, prompt_template, style_name, seed, val_r):
     image_pil = TF.to_pil_image(image_t.to(torch.float32))
     print(f"r_val={val_r}, seed={seed}")
     with torch.no_grad():
-        c_t = image_t.unsqueeze(0).cuda().float()
+        #c_t = image_t.unsqueeze(0).cuda().float()
+        c_t = image_t.unsqueeze(0).to(device).float()
         torch.manual_seed(seed)
         B,C,H,W = c_t.shape
         noise = torch.randn((1,4,H//8, W//8), device=c_t.device)