enable zerogpu

app.py

@@ -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 (
@@ -359,7 +359,6 @@ def get_ref_anno(ref):
         target_size=opts.image_size,
         latent_size=opts.latent_size,
     )
-    print("ready to go to autoencoder")
     latent = opts.latent_scaling_factor * autoencoder.encode(image).sample()
     if not REF_POSE_MASK:
         heatmaps = torch.zeros_like(heatmaps)

vqvae.py

@@ -20,6 +20,7 @@ from typing import List
 import torch
 import torch.nn.functional as F
 from torch import nn
+import spaces
 
 
 class Autoencoder(nn.Module):
@@ -72,6 +73,7 @@ class Autoencoder(nn.Module):
         # Decode the image of shape `[batch_size, channels, height, width]`
         return self.decoder(z)
 
+    @spaces.GPU(duration=120)
     def forward(self, x):
         posterior = self.encode(x)
         z = posterior.sample()