fix

.gradio/certificate.pem

@@ -0,0 +1,31 @@
+-----BEGIN CERTIFICATE-----
+MIIFazCCA1OgAwIBAgIRAIIQz7DSQONZRGPgu2OCiwAwDQYJKoZIhvcNAQELBQAw
+TzELMAkGA1UEBhMCVVMxKTAnBgNVBAoTIEludGVybmV0IFNlY3VyaXR5IFJlc2Vh
+cmNoIEdyb3VwMRUwEwYDVQQDEwxJU1JHIFJvb3QgWDEwHhcNMTUwNjA0MTEwNDM4
+WhcNMzUwNjA0MTEwNDM4WjBPMQswCQYDVQQGEwJVUzEpMCcGA1UEChMgSW50ZXJu
+ZXQgU2VjdXJpdHkgUmVzZWFyY2ggR3JvdXAxFTATBgNVBAMTDElTUkcgUm9vdCBY
+MTCCAiIwDQYJKoZIhvcNAQEBBQADggIPADCCAgoCggIBAK3oJHP0FDfzm54rVygc
+h77ct984kIxuPOZXoHj3dcKi/vVqbvYATyjb3miGbESTtrFj/RQSa78f0uoxmyF+
+0TM8ukj13Xnfs7j/EvEhmkvBioZxaUpmZmyPfjxwv60pIgbz5MDmgK7iS4+3mX6U
+A5/TR5d8mUgjU+g4rk8Kb4Mu0UlXjIB0ttov0DiNewNwIRt18jA8+o+u3dpjq+sW
+T8KOEUt+zwvo/7V3LvSye0rgTBIlDHCNAymg4VMk7BPZ7hm/ELNKjD+Jo2FR3qyH
+B5T0Y3HsLuJvW5iB4YlcNHlsdu87kGJ55tukmi8mxdAQ4Q7e2RCOFvu396j3x+UC
+B5iPNgiV5+I3lg02dZ77DnKxHZu8A/lJBdiB3QW0KtZB6awBdpUKD9jf1b0SHzUv
+KBds0pjBqAlkd25HN7rOrFleaJ1/ctaJxQZBKT5ZPt0m9STJEadao0xAH0ahmbWn
+OlFuhjuefXKnEgV4We0+UXgVCwOPjdAvBbI+e0ocS3MFEvzG6uBQE3xDk3SzynTn
+jh8BCNAw1FtxNrQHusEwMFxIt4I7mKZ9YIqioymCzLq9gwQbooMDQaHWBfEbwrbw
+qHyGO0aoSCqI3Haadr8faqU9GY/rOPNk3sgrDQoo//fb4hVC1CLQJ13hef4Y53CI
+rU7m2Ys6xt0nUW7/vGT1M0NPAgMBAAGjQjBAMA4GA1UdDwEB/wQEAwIBBjAPBgNV
+HRMBAf8EBTADAQH/MB0GA1UdDgQWBBR5tFnme7bl5AFzgAiIyBpY9umbbjANBgkq
+hkiG9w0BAQsFAAOCAgEAVR9YqbyyqFDQDLHYGmkgJykIrGF1XIpu+ILlaS/V9lZL
+ubhzEFnTIZd+50xx+7LSYK05qAvqFyFWhfFQDlnrzuBZ6brJFe+GnY+EgPbk6ZGQ
+3BebYhtF8GaV0nxvwuo77x/Py9auJ/GpsMiu/X1+mvoiBOv/2X/qkSsisRcOj/KK
+NFtY2PwByVS5uCbMiogziUwthDyC3+6WVwW6LLv3xLfHTjuCvjHIInNzktHCgKQ5
+ORAzI4JMPJ+GslWYHb4phowim57iaztXOoJwTdwJx4nLCgdNbOhdjsnvzqvHu7Ur
+TkXWStAmzOVyyghqpZXjFaH3pO3JLF+l+/+sKAIuvtd7u+Nxe5AW0wdeRlN8NwdC
+jNPElpzVmbUq4JUagEiuTDkHzsxHpFKVK7q4+63SM1N95R1NbdWhscdCb+ZAJzVc
+oyi3B43njTOQ5yOf+1CceWxG1bQVs5ZufpsMljq4Ui0/1lvh+wjChP4kqKOJ2qxq
+4RgqsahDYVvTH9w7jXbyLeiNdd8XM2w9U/t7y0Ff/9yi0GE44Za4rF2LN9d11TPA
+mRGunUHBcnWEvgJBQl9nJEiU0Zsnvgc/ubhPgXRR4Xq37Z0j4r7g1SgEEzwxA57d
+emyPxgcYxn/eR44/KJ4EBs+lVDR3veyJm+kXQ99b21/+jh5Xos1AnX5iItreGCc=
+-----END CERTIFICATE-----

app copy.py

@@ -0,0 +1,117 @@
+import gradio as gr
+import spaces
+
+
+import gradio as gr
+import numpy as np
+from PIL import Image
+import cv2
+from sim.simulator import GenieSimulator
+import os
+
+if not os.path.exists("data/mar_ckpt/langtable"):
+    # download from google drive
+    import gdown
+    gdown.download_folder("https://drive.google.com/drive/u/2/folders/1XU87cRqV-IMZA6RLiabIR_uZngynvUFN")
+    os.system("mkdir -p data/mar_ckpt/; mv langtable data/mar_ckpt/")
+
+RES = 512
+PROMPT_HORIZON = 3
+IMAGE_DIR = "sim/assets/langtable_prompt/"
+
+# Load available images
+available_images = sorted([img for img in os.listdir(IMAGE_DIR) if img.endswith(".png")])
+
+
+@spaces.GPU
+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), (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)
+
+@spaces.GPU
+def model(direction, state):
+    if direction == 'right':
+        action = np.array([0, 0.05])
+    elif direction == 'left':
+        action = np.array([0, -0.05])
+    elif direction == 'down':
+        action = np.array([0.05, 0])
+    elif direction == 'up':
+        action = np.array([-0.05, 0])
+    else:
+        raise ValueError(f"Invalid direction: {direction}")
+    next_image = state['genie'].step(action)['pred_next_frame']
+    next_image = cv2.resize(next_image, (RES, RES))
+    return Image.fromarray(next_image)
+
+@spaces.GPU
+def handle_input(direction, state):
+    print(f"User clicked: {direction}")
+    new_image = model(direction, state)
+    return new_image
+
+@spaces.GPU
+def handle_image_selection(image_name, state):
+    print(f"User selected image: {image_name}")
+    return initialize_simulator(image_name, state)
+
+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)
+    return genie
+
+if __name__ == '__main__':
+
+    with gr.Blocks() as demo:
+        genie   = init_model()
+        genie_instance = gr.State({
+                'genie': genie})
+        with gr.Row():
+            image_selector = gr.Dropdown(
+                choices=available_images, value=available_images[0], label="Select an Image"
+            )
+            select_button = gr.Button("Load Image")
+
+        with gr.Row():
+            image_display = gr.Image(type="pil", label="Generated Image")
+
+        select_button.click(
+            fn=handle_image_selection, inputs=[image_selector, genie_instance], outputs=image_display, show_progress='hidden'
+        )
+
+        with gr.Row():
+            up = gr.Button("↑ Up")
+        with gr.Row():
+            left = gr.Button("← Left")
+            down = gr.Button("↓ Down")
+            right = gr.Button("→ Right")
+
+
+        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()

app.py

@@ -23,19 +23,30 @@ IMAGE_DIR = "sim/assets/langtable_prompt/"
 available_images = sorted([img for img in os.listdir(IMAGE_DIR) if img.endswith(".png")])
 
 
-@spaces.GPU
-def initialize_simulator(image_name, state):
+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',
+)
+
+# Helper function to reset GenieSimulator with the selected image
+def initialize_simulator(image_name):
     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):
+# Example model: takes a direction and returns a random image
+def model(direction: str):
     if direction == 'right':
         action = np.array([0, 0.05])
     elif direction == 'left':
@@ -46,48 +57,24 @@ 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)
 
+# Gradio function to handle user input
 @spaces.GPU
-def handle_input(direction, state):
+def handle_input(direction):
     print(f"User clicked: {direction}")
-    new_image = model(direction, state)
+    new_image = model(direction)  # Get a new image from the model
     return new_image
 
-@spaces.GPU
-def handle_image_selection(image_name, state):
+# Gradio function to handle image selection
+def handle_image_selection(image_name):
     print(f"User selected image: {image_name}")
-    return initialize_simulator(image_name, state)
-
-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)
-    return genie
+    return initialize_simulator(image_name)
 
 if __name__ == '__main__':
-
     with gr.Blocks() as demo:
-        genie   = init_model()
-        genie_instance = gr.State({
-                'genie': genie})
         with gr.Row():
             image_selector = gr.Dropdown(
                 choices=available_images, value=available_images[0], label="Select an Image"
@@ -97,10 +84,6 @@ if __name__ == '__main__':
         with gr.Row():
             image_display = gr.Image(type="pil", label="Generated Image")
 
-        select_button.click(
-            fn=handle_image_selection, inputs=[image_selector, genie_instance], outputs=image_display, show_progress='hidden'
-        )
-
         with gr.Row():
             up = gr.Button("↑ Up")
         with gr.Row():
@@ -108,10 +91,13 @@ if __name__ == '__main__':
             down = gr.Button("↓ Down")
             right = gr.Button("→ Right")
 
+        # Define interactions
+        select_button.click(
+            fn=handle_image_selection, inputs=image_selector, outputs=image_display
+        )
+        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()
+    demo.launch(share=True)