| """The two demo stages, run with the RoboTryOn code as released (``robotryon/``, |
| ``scripts/infer_video.py``, ``configs/``). |
| |
| The human clips are preprocessed in advance for every embodiment (hand reconstruction, |
| retargeting to the robot hand, arm masks, keyframe selection) and their prompts are encoded in |
| advance, so neither the embodiment models nor the text encoders are needed here. |
| """ |
| from __future__ import annotations |
|
|
| import importlib.util |
| import json |
| import os |
| import subprocess |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| from huggingface_hub import snapshot_download |
| from PIL import Image, ImageDraw, ImageFont |
|
|
| from robotryon.common.config import load_config |
| from robotryon.image.ar_keyframes import KeyframeTranslator |
| from robotryon.image.episodes import Episode, FrameReader, load_view_frame |
| from robotryon.image.model import load_pipeline as load_image_pipeline |
|
|
| HERE = Path(__file__).resolve().parent |
| ASSETS = HERE / "assets" |
| _spec = importlib.util.spec_from_file_location("infer_video", HERE / "scripts" / "infer_video.py") |
| iv = importlib.util.module_from_spec(_spec) |
| _spec.loader.exec_module(iv) |
|
|
| CKPT_REPO = "RyanL22/RoboTryOn" |
| ASSETS_REPO = "RyanL22/RoboTryOn-demo-assets" |
| FLUX_REPO = "black-forest-labs/FLUX.2-klein-base-9B" |
| LTX_REPO = "Lightricks/LTX-2.5-Diffusers" |
| |
| FLUX_REVISION = "32773329fbe7e81a90ef971740e8ba4b0364ecf3" |
| LTX_REVISION = "a6de4b5354f078db24d9cf4778c14846788aea3d" |
| FLUX_FILES = ["model_index.json", "scheduler/*", "transformer/*", "vae/*"] |
| LTX_FILES = ["model_index.json", "scheduler/*", "transformer_full/*", "vae/*", "audio_vae/*", |
| "vocoder/*", "connectors/*", "diffusion_decoder/*", "duration_head/*", |
| "ltx-2.5-22b-distilled-lora-450-bf16.safetensors"] |
| IMAGE_CONFIG = HERE / "configs" / "image" / "infer.yaml" |
| VIDEO_CONFIG = HERE / "configs" / "video" / "infer.yaml" |
|
|
|
|
| def load_manifest() -> dict: |
| if not (ASSETS / "manifest.json").exists(): |
| snapshot_download(ASSETS_REPO, repo_type="dataset", local_dir=ASSETS) |
| return json.loads((ASSETS / "manifest.json").read_text()) |
|
|
|
|
| def _dir(env: str, repo: str, files: list[str] | None, revision: str | None = None) -> str: |
| return os.environ.get(env) or snapshot_download(repo, allow_patterns=files, revision=revision) |
|
|
|
|
| def human_frames(ep: Episode, cfg, frames: list[int]) -> list[np.ndarray]: |
| """The human frames at ``frames``, cropped and resized as the image model sees them.""" |
| c, reader = cfg.image, FrameReader() |
| return [load_view_frame(reader, ep, ep.views[0], t, int(c.long_edge), |
| tuple(c.aspect) if c.get("aspect") else None, |
| c.get("crop_anchor", "bottom")).image for t in frames] |
|
|
|
|
| def keyframe_grid(human: list[np.ndarray], robot: list[np.ndarray], frames: list[int], |
| cols: int = 4, width: int = 384) -> Image.Image: |
| """Human keyframes above their translations, ``cols`` keyframes per row pair.""" |
| h = round(width * human[0].shape[0] / human[0].shape[1]) |
| rows = (len(frames) + cols - 1) // cols |
| n = min(cols, len(frames)) |
| pad = 4 |
| grid = Image.new("RGB", (n * width + (n + 1) * pad, rows * (2 * h + pad) + (rows + 1) * pad), |
| (255, 255, 255)) |
| draw = ImageDraw.Draw(grid) |
| font = ImageFont.load_default(size=18) |
| for i, (a, b, t) in enumerate(zip(human, robot, frames)): |
| x = pad + (i % cols) * (width + pad) |
| y = pad + (i // cols) * (2 * h + 2 * pad) |
| grid.paste(Image.fromarray(a).resize((width, h), Image.LANCZOS), (x, y)) |
| grid.paste(Image.fromarray(b).resize((width, h), Image.LANCZOS), (x, y + h + pad)) |
| draw.text((x + 6, y + 4), f"frame {t}", fill=(255, 255, 255), font=font, |
| stroke_width=2, stroke_fill=(0, 0, 0)) |
| return grid |
|
|
|
|
| def h264(src: Path, dst: Path) -> Path: |
| subprocess.run(["ffmpeg", "-y", "-loglevel", "error", "-i", str(src), "-c:v", "libx264", |
| "-crf", "18", "-pix_fmt", "yuv420p", str(dst)], check=True) |
| return dst |
|
|
|
|
| class Demo: |
| """Both models, kept on the CPU between calls and moved to the GPU for one stage.""" |
|
|
| def __init__(self, load_models: bool = True): |
| self.manifest = load_manifest() |
| self.clips = {c["key"]: c for c in self.manifest["clips"]} |
| self.embodiments = {e["key"]: e for e in self.manifest["embodiments"]} |
| self.episodes = self.manifest["episodes"] |
| if not load_models: |
| return |
| self.prompts = torch.load(ASSETS / "prompts.pt") |
| self.negative = torch.load(ASSETS / "negative_prompt.pt") |
| self.ckpt = Path(_dir("ROBOTRYON_CKPT", CKPT_REPO, ["image/*", "video/*"])) |
| self.flux = _dir("FLUX_DIR", FLUX_REPO, FLUX_FILES, FLUX_REVISION) |
| self.ltx = _dir("LTX_DIR", LTX_REPO, LTX_FILES, LTX_REVISION) |
| self.image = load_image_pipeline(self.flux, self.ckpt / "image", device="cpu", |
| text_encoder=False) |
| self.video = iv.load_pipeline(self.video_config([]), views=1, device="cpu") |
|
|
| def episode(self, clip: str, emb: str) -> dict: |
| return self.episodes[f"{emb}/{clip}"] |
|
|
| def image_config(self, overrides: list[str]): |
| return load_config(IMAGE_CONFIG, [f"model.base={self.flux}", |
| f"checkpoint={self.ckpt / 'image'}", *overrides]) |
|
|
| def video_config(self, overrides: list[str]): |
| return load_config(VIDEO_CONFIG, [f"model.base={self.ltx}", |
| f"model.lora={self.ckpt / 'video'}", *overrides]) |
|
|
| def translate(self, clip: str, emb: str, out_root: str | Path) -> tuple[Path, Image.Image]: |
| """Stage 1: keyframe translation. Returns the keyframe directory and the grid.""" |
| e = self.episode(clip, emb) |
| cfg = self.image_config([f"output_dir={out_root}", *e["image_overrides"]]) |
| ep_dir = HERE / e["dir"] |
| ep = Episode.load(ep_dir) |
| ref = np.asarray(Image.open(HERE / self.embodiments[emb]["reference"]).convert("RGB")) |
| self.image.to("cuda") |
| try: |
| tr = KeyframeTranslator(self.image, cfg) |
| tr._prompts.update({k: tuple(x.to("cuda") for x in self.prompts[k]) |
| for k in e["prompts"]}) |
| kf_dir = tr.translate_episode(ep_dir, {v: ref for v in ep.views}, out_root) |
| finally: |
| self.image.to("cpu") |
| torch.cuda.empty_cache() |
| frames = json.loads((kf_dir / "keyframes.json").read_text())["indices"] |
| view = ep.views[0] |
| robot = [np.asarray(Image.open(kf_dir / f"keyframe_{view}_{t:06d}.png").convert("RGB")) |
| for t in frames] |
| grid = keyframe_grid(human_frames(ep, cfg, frames), robot, frames) |
| grid.save(kf_dir / "grid.jpg", quality=92) |
| return kf_dir, grid |
|
|
| def synthesize(self, clip: str, emb: str, kf_dir: str | Path, out_dir: str | Path) -> Path: |
| """Stage 2: video synthesis from the stage-1 keyframes. Returns an h264 mp4.""" |
| e = self.episode(clip, emb) |
| cfg = self.video_config([f"input.episode={HERE / e['dir']}", |
| f"input.keyframes_dir={kf_dir}", f"output.dir={out_dir}"]) |
| pe, pm = torch.load(HERE / self.clips[clip]["caption_embeds"]) |
| text = (pe, pm, *self.negative) |
| self.video.to("cuda") |
| try: |
| out = iv.synthesize(cfg, self.video, text) |
| finally: |
| self.video.to("cpu") |
| torch.cuda.empty_cache() |
| view = json.loads((out / "meta.json").read_text())["views"][0] |
| return h264(out / f"video_{view}.mp4", out / "robot.mp4") |
|
|
