HReynaud commited on
Commit
ce2d514
·
1 Parent(s): ca40f6b
Files changed (1) hide show
  1. demo.py +4 -1
demo.py CHANGED
@@ -7,6 +7,7 @@ import cv2
7
  import diffusers
8
  import gradio as gr
9
  import numpy as np
 
10
  import torch
11
  from einops import rearrange
12
  from huggingface_hub import hf_hub_download
@@ -15,7 +16,6 @@ from PIL import Image, ImageOps
15
  from safetensors.torch import load_file
16
  from torch.nn import functional as F
17
  from torchdiffeq import odeint_adjoint as odeint
18
- import spaces
19
 
20
  from echoflow.common import instantiate_class_from_config, unscale_latents
21
  from echoflow.common.models import (
@@ -239,6 +239,7 @@ def preprocess_mask(mask):
239
  # Convert back to numpy array
240
  return np.array(mask_pil)
241
 
 
242
  @spaces.GPU
243
  def generate_latent_image(mask, class_selection, sampling_steps=50):
244
  """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):
301
 
302
  return latent_image # B x C x H x W
303
 
 
304
  @spaces.GPU
305
  def decode_images(latents, vae):
306
  """Decode latent representations to pixel space using a VAE.
@@ -401,6 +403,7 @@ def check_privacy(latent_image_numpy, class_selection):
401
  f"✅ **Success:** Generated image passed privacy check (corr = {corr:.4f} / tau = {tau:.4f})",
402
  )
403
 
 
404
  @spaces.GPU
405
  def generate_animation(
406
  latent_image, ejection_fraction, sampling_steps=50, cfg_scale=1.0
 
7
  import diffusers
8
  import gradio as gr
9
  import numpy as np
10
+ import spaces
11
  import torch
12
  from einops import rearrange
13
  from huggingface_hub import hf_hub_download
 
16
  from safetensors.torch import load_file
17
  from torch.nn import functional as F
18
  from torchdiffeq import odeint_adjoint as odeint
 
19
 
20
  from echoflow.common import instantiate_class_from_config, unscale_latents
21
  from echoflow.common.models import (
 
239
  # Convert back to numpy array
240
  return np.array(mask_pil)
241
 
242
+
243
  @spaces.GPU
244
  def generate_latent_image(mask, class_selection, sampling_steps=50):
245
  """Generate a latent image based on mask, class selection, and sampling steps"""
 
302
 
303
  return latent_image # B x C x H x W
304
 
305
+
306
  @spaces.GPU
307
  def decode_images(latents, vae):
308
  """Decode latent representations to pixel space using a VAE.
 
403
  f"✅ **Success:** Generated image passed privacy check (corr = {corr:.4f} / tau = {tau:.4f})",
404
  )
405
 
406
+
407
  @spaces.GPU
408
  def generate_animation(
409
  latent_image, ejection_fraction, sampling_steps=50, cfg_scale=1.0