"""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" # dataset: preprocessed clips, encoded prompts, examples FLUX_REPO = "black-forest-labs/FLUX.2-klein-base-9B" LTX_REPO = "Lightricks/LTX-2.5-Diffusers" # base-model revisions the demo was tested with (later LTX revisions add components) 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")