Update app.py

app.py

@@ -24,6 +24,7 @@ import numpy as np
 from huggingface_hub import hf_hub_download
 
 def run(image, mode, sample_c=1.3,  num_samples=3, sample_step=100):
+    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
     parser, parser_up = create_argparser()
     
     args = parser.parse_args()
@@ -70,8 +71,8 @@ def run(image, mode, sample_c=1.3,  num_samples=3, sample_step=100):
             model_ckpt2   , strict=True ) 
 
  
-    model.cuda()
-    model_up.cuda()
+    model.to(device)
+    model_up.to(device)
     model.eval()
     model_up.eval()
  
@@ -120,7 +121,7 @@ def run(image, mode, sample_c=1.3,  num_samples=3, sample_step=100):
                 prompt=model_kwargs,
                 batch_size= args.num_samples,
                 guidance_scale=args.sample_c,
-                device=torch.device('cuda'),
+                device=device,
                 prediction_respacing= str(sample_step),
                 upsample_enabled= False,
                 upsample_temp=0.997,
@@ -140,7 +141,7 @@ def run(image, mode, sample_c=1.3,  num_samples=3, sample_step=100):
                 prompt=model_kwargs,
                 batch_size=args.num_samples,
                 guidance_scale=1,
-                device=torch.device('cuda'),
+                device=device,
                 prediction_respacing= "fast27",
                 upsample_enabled=True,
                 upsample_temp=0.997,