app demo

app.py

@@ -33,6 +33,7 @@
 #       step_fn: src.diffusion.stateful_flow_matching.sampling.ode_step_fn
 import os
 import torch
+import spaces
 import argparse
 from omegaconf import OmegaConf
 from src.models.vae import fp2uint8
@@ -69,6 +70,7 @@ class Pipeline:
         self.diffusion_sampler = diffusion_sampler
         self.resolution = resolution
 
+    @spaces.GPU
     @torch.no_grad()
     @torch.autocast(device_type="cuda", dtype=torch.bfloat16)
     def __call__(self, y, num_images, seed, num_steps, guidance, state_refresh_rate, guidance_interval_min, guidance_interval_max, timeshift):