updatE

app.py

@@ -28,18 +28,18 @@ 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, state):
+def initialize_simulator(image_name, genie):
     image_path = os.path.join(IMAGE_DIR, image_name)
     image = Image.open(image_path)
-    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()
+    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()
     reset_image = cv2.resize(reset_image, (RES, RES))
     return Image.fromarray(reset_image)
 
 @spaces.GPU
-def model(direction, state):
+def model(direction, genie):
     if direction == 'right':
         action = np.array([0, 0.05])
     elif direction == 'left':
@@ -50,7 +50,7 @@ def model(direction, state):
         action = np.array([-0.05, 0])
     else:
         raise ValueError(f"Invalid direction: {direction}")
-    next_image = state['genie'].step(action)['pred_next_frame']
+    next_image = genie.step(action)['pred_next_frame']
     next_image = cv2.resize(next_image, (RES, RES))
     return Image.fromarray(next_image)
 
@@ -112,12 +112,12 @@ if __name__ == '__main__':
 
         # Define interactions
         select_button.click(
-            fn=handle_image_selection, inputs=[image_selector, genie_instance], outputs=image_display, show_progress='hidden'
+            fn=handle_image_selection, inputs=[image_selector, genie], 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')
+        up.click(fn=lambda state: handle_input("up", state), inputs=[genie], outputs=image_display, show_progress='hidden')
+        down.click(fn=lambda state: handle_input("down", state), inputs=[genie], outputs=image_display, show_progress='hidden')
+        left.click(fn=lambda state: handle_input("left", state), inputs=[genie], outputs=image_display, show_progress='hidden')
+        right.click(fn=lambda state: handle_input("right", state), inputs=[genie], outputs=image_display, show_progress='hidden')
 
     demo.launch(share=True)