fix bug

app.py

@@ -15,6 +15,7 @@ from NaRCan_model import Homography, Siren
 from util import get_mgrid, apply_homography, jacobian, VideoFitting, TestVideoFitting
 
 
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 def get_example():
     case = [
@@ -116,9 +117,9 @@ def NaRCan_make_video(edit_canonical, pth_path, frames_path):
     # load NaRCan model
     checkpoint_g_old = torch.load(os.path.join(pth_path, "homography_g.pth"))
     checkpoint_g = torch.load(os.path.join(pth_path, "mlp_g.pth"))
-    g_old = Homography(hidden_features=256, hidden_layers=2).cuda()
+    g_old = Homography(hidden_features=256, hidden_layers=2).to(device)
     g = Siren(in_features=3, out_features=2, hidden_features=256,
-              hidden_layers=5, outermost_linear=True).cuda()
+              hidden_layers=5, outermost_linear=True).to(device)
     
     g_old.load_state_dict(checkpoint_g_old)
     g.load_state_dict(checkpoint_g)
@@ -135,7 +136,7 @@ def NaRCan_make_video(edit_canonical, pth_path, frames_path):
     videoloader = DataLoader(v, batch_size=1, pin_memory=True, num_workers=0)
 
     model_input, ground_truth = next(iter(videoloader))
-    model_input, ground_truth = model_input[0].cuda(), ground_truth[0].cuda()
+    model_input, ground_truth = model_input[0].to(device), ground_truth[0].to(device)
 
     myoutput = None
     data_len = len(os.listdir(frames_path))
@@ -156,7 +157,7 @@ def NaRCan_make_video(edit_canonical, pth_path, frames_path):
             # use canonical to reconstruct
             w, h = v.W, v.H
             canonical_img = np.array(edit_canonical.convert('RGB'))
-            canonical_img = torch.from_numpy(canonical_img).float().cuda()
+            canonical_img = torch.from_numpy(canonical_img).float().to(device)
             h_c, w_c = canonical_img.shape[:2]
             grid_new = xy_.clone()
             grid_new[..., 1] = xy_[..., 0] / 1.5
@@ -204,7 +205,7 @@ def edit_with_pnp(input_video, prompt, num_steps, guidance_scale, seed, n_prompt
         pipe = StableDiffusionControlNetPipeline.from_pretrained(
             "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16
         )
-        pipe.to("cuda")
+        pipe.to(device)
         # lineart
         processor = LineartDetector.from_pretrained("lllyasviel/Annotators")
         processor_partial = partial(processor, coarse=False)
@@ -231,7 +232,7 @@ def edit_with_pnp(input_video, prompt, num_steps, guidance_scale, seed, n_prompt
         pipe = StableDiffusionControlNetPipeline.from_pretrained(
             "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16
         )
-        pipe.to("cuda")
+        pipe.to(device)
         # canny
         canonical_image = cv2.imread(image_path)
         canonical_image = cv2.cvtColor(canonical_image, cv2.COLOR_BGR2RGB)