init

app.py

@@ -43,6 +43,8 @@ from src.diffusion.flow_matching.scheduling import LinearScheduler
 from PIL import Image
 import gradio as gr
 import tempfile
+import spaces
+
 from huggingface_hub import snapshot_download
 
 
@@ -65,9 +67,9 @@ def load_model(weight_dict, denoiser):
 
 class Pipeline:
     def __init__(self, vae, denoiser, conditioner, resolution):
-        self.vae = vae.cuda()
-        self.denoiser = denoiser.cuda()
-        self.conditioner = conditioner.cuda()
+        self.vae = vae
+        self.denoiser = denoiser
+        self.conditioner = conditioner
         self.conditioner.compile()
         self.resolution = resolution
         self.tmp_dir = tempfile.TemporaryDirectory(prefix="traj_gifs_")
@@ -76,6 +78,7 @@ class Pipeline:
     def __del__(self):
         self.tmp_dir.cleanup()
 
+    @spaces.GPU
     @torch.no_grad()
     @torch.autocast(device_type="cuda", dtype=torch.bfloat16)
     def __call__(self, y, num_images, seed, image_height, image_width, num_steps, guidance, timeshift, order):
@@ -93,8 +96,7 @@ class Pipeline:
         self.denoiser.decoder_patch_scaling_h = image_height / 512
         self.denoiser.decoder_patch_scaling_w = image_width / 512
         xT = torch.randn((num_images, 3, image_height, image_width), device="cpu", dtype=torch.float32,
-                         generator=generator)
-        xT = xT.to("cuda")
+                         generator=generator).cuda()
         with torch.no_grad():
             condition, uncondition = conditioner([y,]*num_images)