updatE

app.py

@@ -65,34 +65,31 @@ def handle_image_selection(image_name, state):
     print(f"User selected image: {image_name}")
     return initialize_simulator(image_name, state)
 
-@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',
-                device="cpu"
-            )
+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"
+        )
+
+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()
 
-    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:
         genie_instance = gr.State({'genie': genie})
         genie.device = "cuda"