fix

app.py

@@ -65,31 +65,35 @@ def handle_image_selection(image_name, state):
     print(f"User selected image: {image_name}")
     return initialize_simulator(image_name, state)
 
-if __name__ == '__main__':
-    genie = GenieSimulator(
-            image_encoder_type='temporalvae',
-            image_encoder_ckpt='stabilityai/stable-video-diffusion-img2vid',
-            quantize=False,
-            backbone_type='stmar',
-            backbone_ckpt='data/mar_ckpt/langtable',
-            prompt_horizon=PROMPT_HORIZON,
-            action_stride=1,
-            domain='language_table',
-            device="cpu"
-        )
+@spaces.GPU
+def init_model():
+    genie =  GenieSimulator(
+                image_encoder_type='temporalvae',
+                image_encoder_ckpt='stabilityai/stable-video-diffusion-img2vid',
+                quantize=False,
+                backbone_type='stmar',
+                backbone_ckpt='data/mar_ckpt/langtable',
+                prompt_horizon=PROMPT_HORIZON,
+                action_stride=1,
+                domain='language_table',
+            )
+
+    image = Image.open("sim/assets/langtable_prompt/frame_06.png")
+    prompt_image = np.tile(
+        np.array(image), (genie.prompt_horizon, 1, 1, 1)
+    ).astype(np.uint8)
+    prompt_action = np.zeros(
+        (genie.prompt_horizon, genie.action_stride, 2)
+    ).astype(np.float32)
+    genie.set_initial_state((prompt_image, prompt_action))
+    genie.device = "cuda"
+    image = genie.reset()
+    return genie
 
+if __name__ == '__main__':
+    genie = init_model()
 
     with gr.Blocks() as demo:
-        image = Image.open("sim/assets/langtable_prompt/frame_06.png")
-        prompt_image = np.tile(
-            np.array(image), (genie.prompt_horizon, 1, 1, 1)
-        ).astype(np.uint8)
-        prompt_action = np.zeros(
-            (genie.prompt_horizon, genie.action_stride, 2)
-        ).astype(np.float32)
-        genie.set_initial_state((prompt_image, prompt_action))
-        genie.device = "cuda"
-        image = genie.reset()
         genie_instance = gr.State({'genie': genie})
 
         with gr.Row():

sim/simulator.py

@@ -245,6 +245,7 @@ class GenieSimulator(LearnedSimulator):
         if self.gauss_act_perturb_scale is not None:
             action = np.random.normal(action, self.gauss_act_perturb_scale)
 
+        self.backbone = backbone.to(device=self.device).eval()
         # encoding
         input_latent_states = torch.cat([
             self.cached_latent_frames,