ZeroGPU

demo.py

@@ -243,6 +243,7 @@ def preprocess_mask(mask):
 
 
 @spaces.GPU
+@torch.no_grad()
 def generate_latent_image(mask, class_selection, sampling_steps=50):
     """Generate a latent image based on mask, class selection, and sampling steps"""
 
@@ -306,6 +307,7 @@ def generate_latent_image(mask, class_selection, sampling_steps=50):
 
 
 @spaces.GPU
+@torch.no_grad()
 def decode_images(latents):
     """Decode latent representations to pixel space using a VAE.
 
@@ -384,6 +386,7 @@ def decode_latent_to_pixel(latent_image):
 
 
 @spaces.GPU
+@torch.no_grad()
 def check_privacy(latent_image_numpy, class_selection):
     """Check if the latent image is too similar to database images"""
     latent_image = torch.from_numpy(latent_image_numpy).to(device, dtype=dtype)
@@ -410,6 +413,7 @@ def check_privacy(latent_image_numpy, class_selection):
 
 
 @spaces.GPU
+@torch.no_grad()
 def generate_animation(
     latent_image, ejection_fraction, sampling_steps=50, cfg_scale=1.0
 ):
@@ -491,7 +495,7 @@ def generate_animation(
 
     # print("Synthetic video:", synthetic_video.shape)
 
-    return synthetic_video  # B x C x T x H x W
+    return synthetic_video.detach()  # B x C x T x H x W
 
 
 def decode_animation(latent_animation):