File size: 7,788 Bytes
1f5afc9 | 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 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 | """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")
|