update

app.py

@@ -22,6 +22,7 @@ import cv2
 import subprocess
 from PIL import Image
 from datetime import datetime
+import spaces
 
 ACTION_KEYS = [
     "inventory",
@@ -141,7 +142,7 @@ SUNFLOWERS_IMAGE = "assets/sunflower_plains.png"
 SUNFLOWERS_RAIN_IMAGE = "assets/rain_sunflower_plains.png"
 
 DEFAULT_IMAGE = ICE_PLAINS_IMAGE
-device = "cuda:0"
+device = torch.device('cuda')
 
 def save_video(frames, path="output.mp4", fps=10):
     h, w, _ = frames[0].shape
@@ -171,11 +172,14 @@ def run(cfg: DictConfig):
     
     memory_frames.append(load_image_as_tensor(DEFAULT_IMAGE))
 
-    _ = algo.interactive(memory_frames[0],
-                                actions[0],
-                                poses[0],
-                                memory_curr_frame,
-                                device="cuda:0")
+    @spaces.GPU()
+    def run_interactive(first_frame, action, first_pose, curr_frame, device):
+        new_frame = algo.interactive(first_frame,
+                                        action,
+                                        first_pose, 
+                                        curr_frame, 
+                                        device=device)
+        return new_frame
 
     def set_denoising_steps(denoising_steps, sampling_timesteps_state):
         algo.sampling_timesteps = denoising_steps
@@ -190,11 +194,11 @@ def run(cfg: DictConfig):
         global memory_curr_frame
         for i in range(len(actions)):
             memory_curr_frame += 1
-            new_frame = algo.interactive(memory_frames[0],
+            new_frame = run_interactive(memory_frames[0],
                                           actions[i],
                                           None, 
                                           memory_curr_frame, 
-                                          device="cuda:0")
+                                          device=device)
 
             memory_frames.append(new_frame)
 
@@ -222,11 +226,11 @@ def run(cfg: DictConfig):
         memory_curr_frame = 0
         input_history = ""
 
-        _ = algo.interactive(memory_frames[0],
+        _ = run_interactive(memory_frames[0],
                                     actions[0],
                                     poses[0],
                                     memory_curr_frame,
-                                    device="cuda:0")
+                                    device=device)
         return input_history, DEFAULT_IMAGE
 
     def on_image_click(SELECTED_IMAGE):
@@ -235,6 +239,12 @@ def run(cfg: DictConfig):
         reset()
         return SELECTED_IMAGE
 
+    _ = run_interactive(memory_frames[0],
+                                actions[0],
+                                poses[0],
+                                memory_curr_frame,
+                                device)
+
     css = """
     h1 {
         text-align: center;