ZeroGPU

demo.py

@@ -7,6 +7,7 @@ import cv2
 import diffusers
 import gradio as gr
 import numpy as np
+import spaces
 import torch
 from einops import rearrange
 from huggingface_hub import hf_hub_download
@@ -15,7 +16,6 @@ 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 (
@@ -239,6 +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"""
@@ -301,6 +302,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.
@@ -401,6 +403,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