enable zerogpu

app.py

@@ -210,13 +210,13 @@ if NEW_MODEL:
         learn_sigma=True,
     ).to(device)
     # ckpt_state_dict = torch.load(model_path)['model_state_dict']
-    ckpt_state_dict = torch.load(model_path, map_location=torch.device(device))['ema_state_dict']
+    ckpt_state_dict = torch.load(model_path, map_location='cpu')['ema_state_dict']
     missing_keys, extra_keys = model.load_state_dict(ckpt_state_dict, strict=False)
     model.eval()
     print(missing_keys, extra_keys)
     assert len(missing_keys) == 0
     vae_state_dict = torch.load(vae_path)['state_dict']
-    autoencoder = vqvae.create_model(3, 3, opts.latent_dim).eval().requires_grad_(False).to(device)
+    autoencoder = vqvae.create_model(3, 3, opts.latent_dim).eval().requires_grad_(False) # .to(device)
     missing_keys, extra_keys = autoencoder.load_state_dict(vae_state_dict, strict=False)
     autoencoder.eval()
     assert len(missing_keys) == 0
@@ -243,7 +243,7 @@ else:
     autoencoder.eval()
     assert len(missing_keys) == 0 and len(extra_keys) == 0
 sam_path = hf_hub_download(repo_id="Chaerin5/FoundHand-weights", filename="sam_vit_h_4b8939.pth", token=token)
-sam_predictor = init_sam(ckpt_path=sam_path, device=device)
+sam_predictor = init_sam(ckpt_path=sam_path, device='cpu')
 
 
 print("Mediapipe hand detector and SAM ready...")
@@ -254,7 +254,7 @@ hands = mp_hands.Hands(
     min_detection_confidence=0.1,
 )
 
-# @spaces.GPU(duration=120)
+@spaces.GPU(duration=120)
 def get_ref_anno(ref):
     if ref is None:
         return (
@@ -301,6 +301,7 @@ def get_ref_anno(ref):
             elif keypts[21].sum() != 0:
                 input_point = np.array(keypts[21:22])
                 input_label = np.array([1])
+            print("ready to run SAM")
             masks, _, _ = sam_predictor.predict(
                 point_coords=input_point,
                 point_labels=input_label,