adjusted for zeroGPU

demo.py

@@ -15,6 +15,7 @@ from PIL import Image, ImageOps
 from safetensors.torch import load_file
 from torch.nn import functional as F
 from torchdiffeq import odeint_adjoint as odeint
+import spaces
 
 from echoflow.common import instantiate_class_from_config, unscale_latents
 from echoflow.common.models import (
@@ -238,7 +239,7 @@ def preprocess_mask(mask):
     # Convert back to numpy array
     return np.array(mask_pil)
 
-
+@spaces.GPU
 def generate_latent_image(mask, class_selection, sampling_steps=50):
     """Generate a latent image based on mask, class selection, and sampling steps"""
 
@@ -300,7 +301,7 @@ def generate_latent_image(mask, class_selection, sampling_steps=50):
 
     return latent_image  # B x C x H x W
 
-
+@spaces.GPU
 def decode_images(latents, vae):
     """Decode latent representations to pixel space using a VAE.
 
@@ -400,7 +401,7 @@ def check_privacy(latent_image_numpy, class_selection):
             f"✅ **Success:** Generated image passed privacy check (corr = {corr:.4f} / tau = {tau:.4f})",
         )
 
-
+@spaces.GPU
 def generate_animation(
     latent_image, ejection_fraction, sampling_steps=50, cfg_scale=1.0
 ):

requirements.txt

@@ -11,4 +11,5 @@ torch==2.2.2
 torchdiffeq==0.2.4
 xformers==0.0.25.post1
 timm==0.9.16
-accelerate==0.34.2
+accelerate==0.34.2
+spaces==0.33.1