File size: 5,703 Bytes
1454d47 | 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 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 | 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"] # A100: ~34 s per keyframe
def stage2_duration(clip_label, emb_name, *_):
return 40 + 60 * DEMO.episode(*_sel(clip_label, emb_name))["num_clips"] # A100: ~60 s per clip
@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") # LTX-2.5 needs more than a half H200
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()
|