release

demo_img.py

@@ -28,17 +28,18 @@ def img2vid(model_type, img0, img1, frame_ratio, iters):
     img1_t = img2tensor(img1).to(device)
     inputs = [img0_t, img1_t]
     
-    if device == 'cpu':
+    
+    if device == 'cuda':
+        anchor_resolution = 1024 * 512
+        anchor_memory = 1500 * 1024**2
+        anchor_memory_bias = 2500 * 1024**2
+        vram_avail = torch.cuda.get_device_properties(device).total_memory
+    else':
         # Do not resize in cpu mode
         anchor_resolution = 8192*8192
         anchor_memory = 1
         anchor_memory_bias = 0
         vram_avail = 1
-    elif device == 'cuda':
-        anchor_resolution = 1024 * 512
-        anchor_memory = 1500 * 1024**2
-        anchor_memory_bias = 2500 * 1024**2
-        vram_avail = torch.cuda.get_device_properties(device).total_memory
     embt = torch.tensor(1/2).float().view(1, 1, 1, 1).to(device)
 
     inputs = check_dim_and_resize(inputs)

demo_vid.py

@@ -27,17 +27,17 @@ def vid2vid(model_type, video, iters):
     inputs = []
     h = int(vcap.get(cv2.CAP_PROP_FRAME_WIDTH))
     w = int(vcap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-    if device == 'cpu':
+    if device == 'cuda':
+        anchor_resolution = 1024 * 512
+        anchor_memory = 1500 * 1024**2
+        anchor_memory_bias = 2500 * 1024**2
+        vram_avail = torch.cuda.get_device_properties(device).total_memory
+    else':
         # Do not resize in cpu mode
         anchor_resolution = 8192*8192
         anchor_memory = 1
         anchor_memory_bias = 0
         vram_avail = 1
-    elif device == 'cuda':
-        anchor_resolution = 1024 * 512
-        anchor_memory = 1500 * 1024**2
-        anchor_memory_bias = 2500 * 1024**2
-        vram_avail = torch.cuda.get_device_properties(device).total_memory
     
     scale = anchor_resolution / (h * w) * np.sqrt((vram_avail - anchor_memory_bias) / anchor_memory)
     scale = 1 if scale > 1 else scale