enable zerogpu

app.py

@@ -256,60 +256,15 @@ hands = mp_hands.Hands(
     min_detection_confidence=0.1,
 )
 
-# @spaces.GPU(duration=60)
-# def make_ref_cond(
-#     image
-# ):
-#     print("ready to run autoencoder")
-#     # print(f"image.device: {image.device}, type(image): {type(image)}")
-#     os.environ["CUDA_VISIBLE_DEVICES"] = "0"
-#     torch.cuda.set_device(0)
-#     image = image.to("cuda")
-#     latent = opts.latent_scaling_factor * autoencoder.encode(image[None, ...]).sample()
-#     return image[None, ...], latent
-
-# @spaces.GPU(duration=60)
-# def make_ref_cond(image):
-#     def initialize_and_process(image):
-#         os.environ["CUDA_VISIBLE_DEVICES"] = "0"
-#         torch.cuda.set_device(0)
-#         print("Initializing autoencoder in worker process")
-#         image = image.to("cuda")
-#         latent = opts.latent_scaling_factor * autoencoder.encode(image[None, ...]).sample()
-#         return image[None, ...], latent
-
-#     from multiprocessing import Process, Queue
-#     queue = Queue()
-
-#     def worker(image, queue):
-#         result = initialize_and_process(image)
-#         queue.put(result)
-
-#     process = Process(target=worker, args=(image, queue))
-#     process.start()
-#     process.join()
-
-#     return queue.get()
-
-@spaces.GPU(duration=60)
-def make_ref_cond(image):
-    result = []
-
-    def initialize_and_process(image):
-        os.environ["CUDA_VISIBLE_DEVICES"] = "0"
-        torch.cuda.set_device(0)
-
-        print("Initializing autoencoder in worker thread")
-        image = image.to("cuda")
-        latent = opts.latent_scaling_factor * autoencoder.encode(image[None, ...]).sample()
-        result.append((image[None, ...], latent))
-
-    from threading import Thread
-    thread = Thread(target=initialize_and_process, args=(image,))
-    thread.start()
-    thread.join()
-
-    return result[0]
+def make_ref_cond(
+    image
+):
+    print("ready to run autoencoder")
+    # print(f"image.device: {image.device}, type(image): {type(image)}")
+    image = image.to("cuda")
+    latent = opts.latent_scaling_factor * autoencoder.encode(image[None, ...]).sample()
+    return image[None, ...], latent
+
 
 def get_ref_anno(ref):
     print("inside get_ref_anno")

vqvae.py

@@ -49,6 +49,7 @@ class Autoencoder(nn.Module):
         # embedding space
         self.post_quant_conv = nn.Conv2d(emb_channels, z_channels, 1)
 
+    @spaces.GPU
     def encode(self, img: torch.Tensor) -> "GaussianDistribution":
         """
         ### Encode images to latent representation