| import os |
| import tempfile |
|
|
| import gradio as gr |
| import spaces |
|
|
| from core import HERE, Demo |
|
|
| DEMO = Demo(load_models=os.environ.get("ROBOTRYON_UI_ONLY") != "1") |
| CLIPS = DEMO.manifest["clips"] |
| EMB_KEY = {e["name"]: e["key"] for e in DEMO.manifest["embodiments"]} |
| EMB_NAME = {v: k for k, v in EMB_KEY.items()} |
| CLIP_LABEL = {c["key"]: f'{c["name"]} ({c["source"]})' for c in CLIPS} |
| CLIP_KEY = {v: k for k, v in CLIP_LABEL.items()} |
| CODE = "https://github.com/ryan-seungyong-lee/RoboTryOn" |
|
|
|
|
| def _p(rel): |
| return str(HERE / rel) |
|
|
|
|
| def _sel(clip_label, emb_name): |
| return CLIP_KEY[clip_label], EMB_KEY[emb_name] |
|
|
|
|
| def stage1_duration(clip_label, emb_name, *_): |
| return 20 + 25 * DEMO.episode(*_sel(clip_label, emb_name))["num_keyframes"] |
|
|
|
|
| def stage2_duration(clip_label, emb_name, *_): |
| return 40 + 60 * DEMO.episode(*_sel(clip_label, emb_name))["num_clips"] |
|
|
|
|
| @spaces.GPU(duration=stage1_duration) |
| def _stage1(clip_label, emb_name, progress=gr.Progress(track_tqdm=True)): |
| clip, emb = _sel(clip_label, emb_name) |
| kf_dir, _ = DEMO.translate(clip, emb, tempfile.mkdtemp(prefix="keyframes_")) |
| return str(kf_dir / "grid.jpg"), (clip, emb, str(kf_dir)) |
|
|
|
|
| @spaces.GPU(duration=stage2_duration, size="xlarge") |
| def _stage2(clip_label, emb_name, kf_dir, progress=gr.Progress(track_tqdm=True)): |
| clip, emb = _sel(clip_label, emb_name) |
| return str(DEMO.synthesize(clip, emb, kf_dir, tempfile.mkdtemp(prefix="video_"))) |
|
|
|
|
| def stage1(clip_label, emb_name): |
| return _stage1(clip_label, emb_name) |
|
|
|
|
| def stage2(clip_label, emb_name, state): |
| if not state or tuple(state[:2]) != _sel(clip_label, emb_name): |
| raise gr.Error("Run stage 1 (keyframe translation) for this video and embodiment first.") |
| return _stage2(clip_label, emb_name, state[2]) |
|
|
|
|
| def pick_clip(evt: gr.SelectData): |
| return CLIP_LABEL[CLIPS[evt.index]["key"]] |
|
|
|
|
| def show_clip(clip_label, emb_name, state): |
| clip = DEMO.clips[CLIP_KEY[clip_label]] |
| keep = state and tuple(state[:2]) == _sel(clip_label, emb_name) |
| return (_p(clip["video"]), clip["caption"], |
| gr.update(selected_index=[c["key"] for c in CLIPS].index(clip["key"])), |
| *((gr.update(), gr.update()) if keep else (None, None))) |
|
|
|
|
| def show_embodiment(clip_label, emb_name, state): |
| keep = state and tuple(state[:2]) == _sel(clip_label, emb_name) |
| return (_p(DEMO.embodiments[EMB_KEY[emb_name]]["reference"]), |
| *((gr.update(), gr.update()) if keep else (None, None))) |
|
|
|
|
| def load_example(clip_label, emb_name): |
| clip, emb = _sel(clip_label, emb_name) |
| ex = next(e for e in DEMO.manifest["examples"] if e["clip"] == clip and e["embodiment"] == emb) |
| return _p(ex["grid"]), _p(ex["video"]), (clip, emb, _p(ex["keyframes"])) |
|
|
|
|
| CSS = """ |
| #clips .thumbnail-item { aspect-ratio: 16 / 9 !important; } |
| #clips .thumbnail-item img { object-fit: contain !important; } |
| #clips .caption-label { font-size: 12px; } |
| """ |
|
|
| first = CLIPS[0] |
| with gr.Blocks(title="RoboTryOn", css=CSS) as demo: |
| gr.Markdown( |
| "# RoboTryOn\n" |
| "**Synthesizing Dexterous Robot Demonstrations from Human Videos via Virtual Try-On** 路 " |
| f"[Paper](https://arxiv.org/abs/2611.05163) 路 " |
| f"[Project page](https://ryan-seungyong-lee.github.io/RoboTryOn) 路 [Code]({CODE})\n\n" |
| "The input videos are preprocessed in advance (hand reconstruction, retargeting to each " |
| "robot hand, arm masks) due to the time limits of this Space. " |
| f"To run RoboTryOn on your own videos, see the [code]({CODE}).") |
| state = gr.State(None) |
|
|
| gr.Markdown("### 1. Human video") |
| gallery = gr.Gallery([(_p(c["thumb"]), c["name"]) for c in CLIPS], columns=9, |
| rows=1, height=120, object_fit="contain", allow_preview=False, |
| selected_index=0, show_label=False, elem_id="clips") |
| with gr.Row(): |
| human = gr.Video(_p(first["video"]), label="Human video", interactive=False, height=220, |
| scale=1) |
| with gr.Column(scale=2): |
| clip_dd = gr.Dropdown(list(CLIP_LABEL.values()), value=CLIP_LABEL[first["key"]], |
| label="Selected video") |
| caption = gr.Textbox(first["caption"], label="Task", interactive=False, lines=3) |
|
|
| gr.Markdown("### 2. Embodiment") |
| with gr.Row(): |
| emb_radio = gr.Radio(list(EMB_KEY), value=next(iter(EMB_KEY)), label="Target embodiment") |
| reference = gr.Image(_p(DEMO.manifest["embodiments"][0]["reference"]), |
| label="Reference", interactive=False, height=180) |
|
|
| gr.Markdown("### 3. Generate") |
| with gr.Row(): |
| b1 = gr.Button("Stage 1: Keyframe translation", variant="primary") |
| b2 = gr.Button("Stage 2: Video synthesis (uses the stage-1 keyframes)", variant="primary") |
| grid = gr.Image(label="Keyframes: human (top) and robot (bottom)", type="filepath") |
| video = gr.Video(label="Robot video") |
|
|
| gr.Markdown("### Examples (precomputed)") |
| gr.Examples([[CLIP_LABEL[e["clip"]], EMB_NAME[e["embodiment"]]] for e in DEMO.manifest["examples"]], |
| inputs=[clip_dd, emb_radio], outputs=[grid, video, state], fn=load_example, |
| run_on_click=True, cache_examples=False) |
|
|
| gallery.select(pick_clip, None, clip_dd) |
| clip_dd.change(show_clip, [clip_dd, emb_radio, state], [human, caption, gallery, grid, video]) |
| emb_radio.change(show_embodiment, [clip_dd, emb_radio, state], [reference, grid, video]) |
| b1.click(stage1, [clip_dd, emb_radio], [grid, state]) |
| b2.click(stage2, [clip_dd, emb_radio, state], video) |
|
|
| demo.queue().launch() |
|
|
