Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,465 Bytes
246c106 bce7a55 246c106 dded309 246c106 dded309 246c106 9912547 246c106 13dce27 246c106 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
import gradio as gr
import spaces
import numpy as np
from PIL import Image
import cv2
from sim.simulator import GenieSimulator
RES = 512
image = Image.open("sim/assets/langtable_prompt/frame_06.png")
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=3,
action_stride=1,
domain='language_table',
)
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))
image = genie.reset()
image = cv2.resize(image, (RES, RES))
image = Image.fromarray(image)
# Example model: takes a direction and returns a random image
def model(direction: str, genie=genie):
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 = 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):
print(f"User clicked: {direction}")
new_image = model(direction) # Get a new image from the model
return new_image
if __name__ == '__main__':
with gr.Blocks() as demo:
with gr.Row():
image_display = gr.Image(value=image, type="pil", label="Generated Image")
with gr.Row():
up = gr.Button("↑ Up")
with gr.Row():
left = gr.Button("← Left")
down = gr.Button("↓ Down")
right = gr.Button("→ Right")
# Define button interactions
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')
demo.launch()
|