enable zerogpu

app.py

@@ -312,6 +312,7 @@ def get_ref_anno(ref):
                 point_labels=input_label,
                 multimask_output=False,
             )
+            print("finished SAM")
             hand_mask = masks[0]
             masked_img = img * hand_mask[..., None] + 255 * (1 - hand_mask[..., None])
             ref_pose = visualize_hand(keypts, masked_img)
@@ -323,51 +324,48 @@ def get_ref_anno(ref):
 
     @spaces.GPU(duration=120)
     def make_ref_cond(
-        img,
-        keypts,
-        hand_mask,
-        device=device,
-        target_size=(256, 256),
-        latent_size=(32, 32),
+        image
     ):
-        image_transform = Compose(
-            [
-                ToTensor(),
-                Resize(target_size),
-                Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
-            ]
-        )
-        image = image_transform(img).to(device)
-        kpts_valid = check_keypoints_validity(keypts, target_size)
-        heatmaps = torch.tensor(
-            keypoint_heatmap(
-                scale_keypoint(keypts, target_size, latent_size), latent_size, var=1.0
-            )
-            * kpts_valid[:, None, None],
-            dtype=torch.float,
-            device=device,
-        )[None, ...]
-        mask = torch.tensor(
-            cv2.resize(
-                hand_mask.astype(int),
-                dsize=latent_size,
-                interpolation=cv2.INTER_NEAREST,
-            ),
-            dtype=torch.float,
-            device=device,
-        ).unsqueeze(0)[None, ...]
+        print("ready to run autoencoder")
         latent = opts.latent_scaling_factor * autoencoder.encode(image[None, ...]).sample()
-        return image[None, ...], heatmaps, mask, latent
-
-    image, heatmaps, mask, latent = make_ref_cond(
-        img,
-        keypts,
-        hand_mask,
+        return image[None, ...], latent
+
+    image_transform = Compose(
+        [
+            ToTensor(),
+            Resize(opts.image_size),
+            Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
+        ]
+    )
+    image = image_transform(img).to(device)
+    kpts_valid = check_keypoints_validity(keypts, opts.image_size)
+    heatmaps = torch.tensor(
+        keypoint_heatmap(
+            scale_keypoint(keypts, opts.image_size, opts.latent_size), opts.latent_size, var=1.0
+        )
+        * kpts_valid[:, None, None],
+        dtype=torch.float,
         device=device,
-        target_size=opts.image_size,
-        latent_size=opts.latent_size,
+    )[None, ...]
+    mask = torch.tensor(
+        cv2.resize(
+            hand_mask.astype(int),
+            dsize=opts.latent_size,
+            interpolation=cv2.INTER_NEAREST,
+        ),
+        dtype=torch.float,
+        device=device,
+    ).unsqueeze(0)[None, ...]
+    image, latent = make_ref_cond(
+        image,
+        # keypts,
+        # hand_mask,
+        # device=device,
+        # target_size=opts.image_size,
+        # latent_size=opts.latent_size,
     )
-    
+    print("finished autoencoder")
+
     if not REF_POSE_MASK:
         heatmaps = torch.zeros_like(heatmaps)
         mask = torch.zeros_like(mask)