update

app.py

@@ -26,20 +26,20 @@ available_images = sorted([img for img in os.listdir(IMAGE_DIR) if img.endswith(
 
 
 # Helper function to reset GenieSimulator with the selected image
+
 @spaces.GPU
-def initialize_simulator(image_name):
+def initialize_simulator(image_name, state):
     image_path = os.path.join(IMAGE_DIR, image_name)
     image = Image.open(image_path)
-    prompt_image = np.tile(np.array(image), (genie.prompt_horizon, 1, 1, 1)).astype(np.uint8)
-    prompt_action = np.zeros((genie.prompt_horizon - 1, genie.action_stride, 2)).astype(np.float32)
-    genie.set_initial_state((prompt_image, prompt_action))
-    reset_image = genie.reset()
+    prompt_image = np.tile(np.array(image), (state['genie'].prompt_horizon, 1, 1, 1)).astype(np.uint8)
+    prompt_action = np.zeros((state['genie'].prompt_horizon - 1, state['genie'].action_stride, 2)).astype(np.float32)
+    state['genie'].set_initial_state((prompt_image, prompt_action))
+    reset_image = state['genie'].reset()
     reset_image = cv2.resize(reset_image, (RES, RES))
     return Image.fromarray(reset_image)
 
-# Example model: takes a direction and returns a random image
 @spaces.GPU
-def model(direction: str):
+def model(direction, state):
     if direction == 'right':
         action = np.array([0, 0.05])
     elif direction == 'left':
@@ -50,22 +50,20 @@ def model(direction: str):
         action = np.array([-0.05, 0])
     else:
         raise ValueError(f"Invalid direction: {direction}")
-    next_image = genie.step(action)['pred_next_frame']
+    next_image = state['genie'].step(action)['pred_next_frame']
     next_image = cv2.resize(next_image, (RES, RES))
     return Image.fromarray(next_image)
 
-# Gradio function to handle user input
 @spaces.GPU
-def handle_input(direction):
+def handle_input(direction, state):
     print(f"User clicked: {direction}")
-    new_image = model(direction)  # Get a new image from the model
+    new_image = model(direction, state)
     return new_image
 
-# Gradio function to handle image selection
 @spaces.GPU
-def handle_image_selection(image_name):
+def handle_image_selection(image_name, state):
     print(f"User selected image: {image_name}")
-    return initialize_simulator(image_name)
+    return initialize_simulator(image_name, state)
 
 if __name__ == '__main__':
     genie = GenieSimulator(
@@ -77,8 +75,10 @@ if __name__ == '__main__':
             prompt_horizon=PROMPT_HORIZON,
             action_stride=1,
             domain='language_table',
+            device="cpu"
         )
 
+
     with gr.Blocks() as demo:
         image = Image.open("sim/assets/langtable_prompt/frame_06.png")
         prompt_image = np.tile(
@@ -88,7 +88,10 @@ if __name__ == '__main__':
             (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():
             image_selector = gr.Dropdown(
                 choices=available_images, value=available_images[0], label="Select an Image"
@@ -107,11 +110,12 @@ if __name__ == '__main__':
 
         # Define interactions
         select_button.click(
-            fn=handle_image_selection, inputs=image_selector, outputs=image_display
+            fn=handle_image_selection, inputs=[image_selector, genie_instance], outputs=image_display, show_progress='hidden'
         )
-        up.click(fn=lambda: handle_input("up"), outputs=image_display, show_progress='hidden')
-        down.click(fn=lambda: handle_input("down"), outputs=image_display, show_progress='hidden')
-        left.click(fn=lambda: handle_input("left"), outputs=image_display, show_progress='hidden')
-        right.click(fn=lambda: handle_input("right"), outputs=image_display, show_progress='hidden')
+
+        up.click(fn=lambda state: handle_input("up", state), inputs=[genie_instance], outputs=image_display, show_progress='hidden')
+        down.click(fn=lambda state: handle_input("down", state), inputs=[genie_instance], outputs=image_display, show_progress='hidden')
+        left.click(fn=lambda state: handle_input("left", state), inputs=[genie_instance], outputs=image_display, show_progress='hidden')
+        right.click(fn=lambda state: handle_input("right", state), inputs=[genie_instance], outputs=image_display, show_progress='hidden')
 
     demo.launch(share=True)