enable zerogpu

app.py

@@ -256,17 +256,40 @@ 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
-):
-    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
+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()
 
 def get_ref_anno(ref):
     print("inside get_ref_anno")