diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000000000000000000000000000000000000..ba9dcd1b37de90d75ab76879c9d847487be7f8a5 --- /dev/null +++ b/.gitattributes @@ -0,0 +1,46 @@ +*.7z filter=lfs diff=lfs merge=lfs -text +*.arrow filter=lfs diff=lfs merge=lfs -text +*.bin filter=lfs diff=lfs merge=lfs -text +*.bz2 filter=lfs diff=lfs merge=lfs -text +*.ckpt filter=lfs diff=lfs merge=lfs -text +*.ftz filter=lfs diff=lfs merge=lfs -text +*.gz filter=lfs diff=lfs merge=lfs -text +*.h5 filter=lfs diff=lfs merge=lfs -text +*.joblib filter=lfs diff=lfs merge=lfs -text +*.lfs.* filter=lfs diff=lfs merge=lfs -text +*.mlmodel filter=lfs diff=lfs merge=lfs -text +*.model filter=lfs diff=lfs merge=lfs -text +*.msgpack filter=lfs diff=lfs merge=lfs -text +*.npy filter=lfs diff=lfs merge=lfs -text +*.npz filter=lfs diff=lfs merge=lfs -text +*.onnx filter=lfs diff=lfs merge=lfs -text +*.ot filter=lfs diff=lfs merge=lfs -text +*.parquet filter=lfs diff=lfs merge=lfs -text +*.pb filter=lfs diff=lfs merge=lfs -text +*.pickle filter=lfs diff=lfs merge=lfs -text +*.pkl filter=lfs diff=lfs merge=lfs -text +*.pt filter=lfs diff=lfs merge=lfs -text +*.pth filter=lfs diff=lfs merge=lfs -text +*.rar filter=lfs diff=lfs merge=lfs -text +*.safetensors filter=lfs diff=lfs merge=lfs -text +saved_model/**/* filter=lfs diff=lfs merge=lfs -text +*.tar.* filter=lfs diff=lfs merge=lfs -text +*.tar filter=lfs diff=lfs merge=lfs -text +*.tflite filter=lfs diff=lfs merge=lfs -text +*.tgz filter=lfs diff=lfs merge=lfs -text +*.wasm filter=lfs diff=lfs merge=lfs -text +*.xz filter=lfs diff=lfs merge=lfs -text +*.zip filter=lfs diff=lfs merge=lfs -text +*.zst filter=lfs diff=lfs merge=lfs -text +*tfevents* filter=lfs diff=lfs merge=lfs -text +examples/gr1_stir_pot/episode/video_ego_view.mp4 filter=lfs diff=lfs merge=lfs -text +examples/gr1_stir_pot/input.mp4 filter=lfs diff=lfs merge=lfs -text +examples/openarm_box_convey/episode/mask_left.mp4 filter=lfs diff=lfs merge=lfs -text +examples/openarm_box_convey/episode/object_mask_left.mp4 filter=lfs diff=lfs merge=lfs -text +examples/openarm_box_convey/episode/video_left.mp4 filter=lfs diff=lfs merge=lfs -text +examples/openarm_box_convey/input.mp4 filter=lfs diff=lfs merge=lfs -text +examples/openarm_box_convey/reference.png filter=lfs diff=lfs merge=lfs -text +examples/rby1_bottle_pour/episode/mask_left.mp4 filter=lfs diff=lfs merge=lfs -text +examples/rby1_bottle_pour/episode/video_left.mp4 filter=lfs diff=lfs merge=lfs -text +examples/rby1_bottle_pour/input.mp4 filter=lfs diff=lfs merge=lfs -text +examples/rby1_bottle_pour/reference.png filter=lfs diff=lfs merge=lfs -text diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..90fc132a6e0ee8685ba2a04e052f7f6a700c9482 --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +__pycache__/ +prep_assets.py +assets/ diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..cbee9d418324999774d84b724dd2fef36a19c6da --- /dev/null +++ b/README.md @@ -0,0 +1,16 @@ +--- +title: RoboTryOn +emoji: 🤖 +colorFrom: gray +colorTo: indigo +sdk: gradio +sdk_version: 5.50.0 +python_version: "3.10" +app_file: app.py +startup_duration_timeout: 1h +pinned: false +short_description: Human videos to dexterous robot videos via virtual try-on +--- + +Demo of [RoboTryOn](https://github.com/ryan-seungyong-lee/RoboTryOn) ([paper](https://arxiv.org/abs/2611.05163)). +`robotryon/`, `scripts/infer_video.py` and `configs/` are copies of the code repository. diff --git a/app.py b/app.py new file mode 100644 index 0000000000000000000000000000000000000000..bf88fc019e99ffd94e3312fb0d853603f8606cc3 --- /dev/null +++ b/app.py @@ -0,0 +1,135 @@ +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() diff --git a/configs/image/infer.yaml b/configs/image/infer.yaml new file mode 100644 index 0000000000000000000000000000000000000000..aa884f0d8bea88e000dd28c1389056e92dd4f948 --- /dev/null +++ b/configs/image/infer.yaml @@ -0,0 +1,32 @@ +# Autoregressive keyframe translation (paper Eq. 3.2). See scripts/infer_keyframes.py. + +model: + base: black-forest-labs/FLUX.2-klein-base-9B +checkpoint: null # trained checkpoint dir (FFT: has transformer/; LoRA: has pytorch_lora_weights.safetensors) + +episodes: [] # preprocessed human episode directories (with keyframes.json) +reference: {} # {view: robot-only reference image}; a single path is used for every view +output_dir: runs/keyframes + +views: auto # auto (joint stereo when the episode has two views) | stereo | mono +ref_skeleton: false # false: reuse the previous keyframe's latent as reference tokens + # true: decode it, paint the current skeleton, re-encode +steps: 30 +guidance_scale: 3.5 # true CFG against the empty prompt +seed: 2024 # keyframe t uses seed + t + +image: + long_edge: 1024 + aspect: [16, 9] + crop_anchor: bottom + +masks: # inference dilation = (dilate_min + dilate_max) // 2 + bonus = 58 px + dilate_min: 3 + dilate_max: 71 + bonus: 21 + obj_erode: 9 # px, erode the object mask before subtracting it from the hole + +caption: # per-keyframe hand-state prompt (robotryon.captions) + cuts: # closure (rad) / dorsumness (deg) edges -> open|half-closed|closed fist, palm|edge-on|back + closure: [0.95, 1.15] + facing: [49.0, 121.0] diff --git a/configs/video/infer.yaml b/configs/video/infer.yaml new file mode 100644 index 0000000000000000000000000000000000000000..732b603b0e50ef57f7e0c9b2a2883fa8c503cce5 --- /dev/null +++ b/configs/video/infer.yaml @@ -0,0 +1,35 @@ +# Video synthesis from translated robot keyframes and the human skeleton flow (paper App. B.4). +model: + base: Lightricks/LTX-2.5-Diffusers + transformer_subfolder: transformer_full + lora: outputs/video/stage2/checkpoint-001000 # directory with pytorch_lora_weights.safetensors + attention_backend: native + +input: + episode: null # human episode dir: meta.json, keypoints_.npz, keyframes.json + keyframes_dir: null # image-model output for this episode: keyframe__.png + keyframes.json + views: null # default: all views in meta.json + stereo: true # generate two views jointly in one sequence; false = one view at a time + +output: + dir: outputs/video/infer + +text: + template: "{caption}" + embodiment: null # fills {embodiment} in the template (target robot) + +fps: 20 +resolution: [1024, 576] # output, per view +skeleton_size: [512, 288] # as in training: reference downscale 1 at 512x288, 2 at 1024x576 +mask_human: true # true: skeleton flow on black | false: skeleton drawn on the unmasked human video +frames_range: [49, 145] # clip length limits (8k+1) +seed: 0 # clip i uses seed + i + +two_stage: + mode: same_res # same_res: 8 steps + 3 refine steps at 1024x576 + # upsample: 8 steps at 512x288 -> latent upsampler x2 -> 3 steps at 1024x576 + +skeleton: + draw: + thickness: 2 + radius: 3 diff --git a/core.py b/core.py new file mode 100644 index 0000000000000000000000000000000000000000..d0b4ab9f7868ef383639b2537bb733b66b38ae42 --- /dev/null +++ b/core.py @@ -0,0 +1,194 @@ +"""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 cv2 +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], view: str | None = None) -> 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, view or 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 match_brightness(image: np.ndarray, human: np.ndarray, mask: np.ndarray, + margin: int = 15) -> np.ndarray: + """Match the keyframe's lightness (mean and spread of L in Lab) to the human frame. + + Only the background outside the human-arm mask (grown by ``margin`` px) is compared, so the + robot itself does not drive the statistics; the same correction is applied to every pixel. + """ + if human.shape[:2] != image.shape[:2]: + human = cv2.resize(human, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_AREA) + if mask.shape[:2] != image.shape[:2]: + mask = cv2.resize(mask, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST) + k = np.ones((2 * margin + 1, 2 * margin + 1), np.uint8) + bg = cv2.dilate((mask > 0).astype(np.uint8), k) == 0 + if bg.sum() < 1000: + return image + lab = cv2.cvtColor(image, cv2.COLOR_RGB2LAB).astype(np.float32) + ref = cv2.cvtColor(human, cv2.COLOR_RGB2LAB).astype(np.float32) + src_l, ref_l = lab[..., 0][bg], ref[..., 0][bg] + gain = ref_l.std() / max(src_l.std(), 1e-3) + lab[..., 0] = (lab[..., 0] - src_l.mean()) * gain + ref_l.mean() + return cv2.cvtColor(np.clip(lab, 0, 255).astype(np.uint8), cv2.COLOR_LAB2RGB) + + +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"] + if self.clips[clip]["source"] != "self-recorded": + # downloaded human videos (VITRA, EgoDex): the generated keyframes drift in overall + # brightness, so match each one to its human frame before stage 2 reads it + for v in ep.views: + human = human_frames(ep, cfg, frames, v) + for t, h in zip(frames, human): + f = kf_dir / f"keyframe_{v}_{t:06d}.png" + mask = np.asarray(Image.open(kf_dir / f"mask_{v}_{t:06d}.png").convert("L")) + img = np.asarray(Image.open(f).convert("RGB")) + Image.fromarray(match_brightness(img, h, mask)).save(f) + 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") diff --git a/packages.txt b/packages.txt new file mode 100644 index 0000000000000000000000000000000000000000..20645e641240cb419f5fc66c14c1447e91daf669 --- /dev/null +++ b/packages.txt @@ -0,0 +1 @@ +ffmpeg diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..e9308bc26c6b0ed8327a9058cb68f41021871be6 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,14 @@ +torch==2.8.0 +torchvision==0.23.0 +diffusers @ git+https://github.com/huggingface/diffusers@90b4e34e79a86ec5e7f2437634fe95ecd2108796 +transformers==5.15.0 +peft==0.20.0 +accelerate==1.14.0 +safetensors +huggingface_hub +numpy<2 +scipy +opencv-python-headless +pillow +pyyaml +tqdm diff --git a/robotryon/__init__.py b/robotryon/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/robotryon/captions.py b/robotryon/captions.py new file mode 100644 index 0000000000000000000000000000000000000000..91fabde041cb1f2026000862091d6266e62ed363 --- /dev/null +++ b/robotryon/captions.py @@ -0,0 +1,146 @@ +"""Text prompts: per-hand state words and the embodiment phrase. + +Image model. A keyframe's prompt states two facts per hand that a 2D skeleton conveys weakly at +the sizes hands occupy in the frame -- how closed the hand is and which side faces the camera:: + + Match the hand pose to the drawn skeleton. + Left hand: open, palm toward camera. Right hand: closed fist, back of hand toward camera. + +``closure`` is the MCP + PIP flexion averaged over the four fingers (radians); +``dorsumness`` is 180 minus the angle between the knuckle-plane normal and the ray from the hand +to the camera (degrees, high = back of the hand toward the lens). Both are read from the +3D hand in the camera frame -- the reconstructed human hand, or the robot hand from forward +kinematics, with the same formula -- and cut into words (``HandStateCuts``). A hand that is not +tracked is ``not visible``. + +Each episode stores both values per view and frame in ``hand_state_.npz`` (``closure``, +``dorsumness``, ``(T, 2)`` float32, NaN where a hand is not tracked), written by the preprocessing +for human clips (reconstructed hands) and robot episodes (forward kinematics), so training and +inference read the same quantity. + +Video model. The prompt is the task caption (``text.template``). +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np + +INSTRUCTION = "Match the hand pose to the drawn skeleton." +CLOSURE = ("open", "half-closed", "closed fist") +FACING = ("palm toward camera", "edge-on", "back of hand toward camera") +NOT_VISIBLE = "not visible" + +# MediaPipe / MANO 21-joint layout: (MCP, PIP, DIP, tip) per finger. +_CHAINS = {"index": (5, 6, 7, 8), "middle": (9, 10, 11, 12), "ring": (13, 14, 15, 16), + "little": (17, 18, 19, 20)} + + +def _angle(a, b, c) -> float: + """Flexion at b: 0 when a-b-c is straight.""" + u, v = a - b, c - b + u = u / max(np.linalg.norm(u), 1e-9) + v = v / max(np.linalg.norm(v), 1e-9) + return float(np.pi - np.arccos(np.clip(np.dot(u, v), -1.0, 1.0))) + + +def hand_fields(j21: np.ndarray, hand: str) -> tuple[float, float]: + """``(21, 3)`` joints of one hand in the camera frame (camera at the origin, looking down +z) + -> ``(closure, dorsumness)``; NaN when the hand is not finite.""" + j = np.asarray(j21, np.float64) + if not np.isfinite(j).all(): + return float("nan"), float("nan") + w = j[0] + closure = np.mean([_angle(w, j[a], j[b]) + _angle(j[a], j[b], j[c]) + for a, b, c, _ in _CHAINS.values()]) + n = np.cross(j[5] - w, j[17] - w) + ray = -j.mean(0) # hand -> camera + nn, nr = np.linalg.norm(n), np.linalg.norm(ray) + if nn < 1e-9 or nr < 1e-9: + return float(closure), float("nan") + d = 180.0 - float(np.degrees(np.arccos(np.clip(np.dot(n, ray) / (nn * nr), -1.0, 1.0)))) + # the normal leaves through the palm on the right hand and the back on the left + return float(closure), (d if hand == "left" else 180.0 - d) + + +@dataclass +class HandStateCuts: + """Edges that turn ``closure`` / ``dorsumness`` into words (lower, upper).""" + closure: tuple[float, float] = (0.95, 1.15) + facing: tuple[float, float] = (49.0, 121.0) + + @classmethod + def from_config(cls, cfg) -> "HandStateCuts": + cfg = cfg or {} + return cls(tuple(float(x) for x in cfg.get("closure", cls.closure)), + tuple(float(x) for x in cfg.get("facing", cls.facing))) + + def words(self, closure: float, dorsumness: float) -> str: + if not (np.isfinite(closure) and np.isfinite(dorsumness)): + return NOT_VISIBLE + ci = int(closure >= self.closure[0]) + int(closure >= self.closure[1]) + fi = int(dorsumness >= self.facing[0]) + int(dorsumness >= self.facing[1]) + return f"{CLOSURE[ci]}, {FACING[fi]}" + + +def hand_state_prompt(left: str, right: str) -> str: + """The image prompt from the two hands' words.""" + return f"{INSTRUCTION} Left hand: {left}. Right hand: {right}." + + +def all_prompts() -> list[str]: + """Every prompt training can produce: each pair of hand descriptions, the hand-dropout prompt + (the instruction alone) and the empty (unconditional) prompt.""" + cells = [f"{c}, {f}" for c in CLOSURE for f in FACING] + [NOT_VISIBLE] + return [hand_state_prompt(a, b) for a in cells for b in cells] + [INSTRUCTION, ""] + + +def sample_prompt(words: tuple[str, str], rng: np.random.Generator, p_uncond: float, + p_drop: float) -> str: + """Training prompt: empty with ``p_uncond``, hand words dropped with ``p_drop``.""" + u = float(rng.random()) + if u < p_uncond: + return "" + if u < p_uncond + p_drop: + return INSTRUCTION + return hand_state_prompt(*words) + + +def hand_state_arrays(joints3d: np.ndarray, valid: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """``(T, 2, 21|17, 3)`` camera-frame joints and ``(T, 2)`` validity -> ``closure`` and + ``dorsumness`` ``(T, 2)`` float32 (NaN where invalid).""" + from .embodiment.retarget import to21 + j = to21(np.asarray(joints3d, np.float64)) + T = j.shape[0] + clo = np.full((T, 2), np.nan, np.float32) + dor = np.full((T, 2), np.nan, np.float32) + for t in range(T): + for h, name in enumerate(("left", "right")): + if valid[t, h]: + clo[t, h], dor[t, h] = hand_fields(j[t, h], name) + return clo, dor + + +def save_hand_state(path, closure: np.ndarray, dorsumness: np.ndarray) -> None: + np.savez_compressed(path, closure=closure.astype(np.float32), dorsumness=dorsumness.astype(np.float32)) + + +def load_hand_state(path) -> tuple[np.ndarray, np.ndarray]: + z = np.load(path) + return z["closure"], z["dorsumness"] + + +def words_at(states: list[tuple[np.ndarray, np.ndarray]], t: int, + cuts: HandStateCuts) -> tuple[str, str]: + """Words of both hands at frame ``t`` from one or more views' hand states (averaged).""" + out = [] + for h in range(2): + vals = [(c[t, h], d[t, h]) for c, d in states + if t < len(c) and np.isfinite(c[t, h]) and np.isfinite(d[t, h])] + if not vals: + out.append(NOT_VISIBLE) + continue + cl, dn = np.mean(np.asarray(vals, np.float64), axis=0) + out.append(cuts.words(float(cl), float(dn))) + return out[0], out[1] diff --git a/robotryon/common/__init__.py b/robotryon/common/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/robotryon/common/config.py b/robotryon/common/config.py new file mode 100644 index 0000000000000000000000000000000000000000..66e496663a555ab5537120bdee335372c98ca66a --- /dev/null +++ b/robotryon/common/config.py @@ -0,0 +1,115 @@ +"""YAML configs with dotted command-line overrides. + +Every script takes ``--config path.yaml`` followed by any number of ``key.sub=value`` +overrides, e.g.:: + + python scripts/train_image.py --config configs/image/klein9b_lora.yaml \ + train.max_steps=20 data.root=/path/to/data + +A config may name a parent with ``_base_: other.yaml`` (relative to itself); the child is merged +on top of it. Values are parsed with YAML rules, so ``1e-5``, ``true``, ``[1, 2]`` and ``null`` +work as expected on the command line. +""" + +from __future__ import annotations + +import argparse +import copy +from pathlib import Path +from typing import Any, Sequence + +import yaml + + +class Config(dict): + """A dict whose keys are also attributes (recursively).""" + + def __getattr__(self, key: str) -> Any: + try: + return self[key] + except KeyError as e: + raise AttributeError(key) from e + + def __setattr__(self, key: str, value: Any) -> None: + self[key] = value + + def get_path(self, dotted: str, default: Any = None) -> Any: + node: Any = self + for k in dotted.split("."): + if not isinstance(node, dict) or k not in node: + return default + node = node[k] + return node + + def to_dict(self) -> dict: + return _plain(self) + + +def _wrap(x: Any) -> Any: + if isinstance(x, dict): + return Config({k: _wrap(v) for k, v in x.items()}) + if isinstance(x, list): + return [_wrap(v) for v in x] + return x + + +def _plain(x: Any) -> Any: + if isinstance(x, dict): + return {k: _plain(v) for k, v in x.items()} + if isinstance(x, list): + return [_plain(v) for v in x] + return x + + +def merge(base: dict, over: dict) -> dict: + out = copy.deepcopy(base) + for k, v in over.items(): + if isinstance(v, dict) and isinstance(out.get(k), dict): + out[k] = merge(out[k], v) + else: + out[k] = copy.deepcopy(v) + return out + + +def _set_dotted(d: dict, dotted: str, value: Any) -> None: + keys = dotted.split(".") + for k in keys[:-1]: + d = d.setdefault(k, {}) + d[keys[-1]] = value + + +def _parse_value(raw: str) -> Any: + v = yaml.safe_load(raw) + if isinstance(v, str): # YAML 1.1 reads "1e-5" as a string; accept it as a float + try: + return float(v) + except ValueError: + pass + return v + + +def load_config(path: str | Path, overrides: Sequence[str] = ()) -> Config: + path = Path(path) + cfg = yaml.safe_load(path.read_text()) or {} + base = cfg.pop("_base_", None) + if base is not None: + cfg = merge(load_config(path.parent / base).to_dict(), cfg) + for item in overrides: + if "=" not in item: + raise ValueError(f"override must be key=value, got {item!r}") + key, raw = item.split("=", 1) + _set_dotted(cfg, key, _parse_value(raw)) + return _wrap(cfg) + + +def parse_args(description: str = "", argv: Sequence[str] | None = None) -> Config: + """``--config file.yaml [key=value ...]`` -> config.""" + p = argparse.ArgumentParser(description=description) + p.add_argument("--config", required=True, help="YAML config file") + p.add_argument("overrides", nargs="*", help="dotted overrides, e.g. train.lr=1e-5") + a = p.parse_args(argv) + return load_config(a.config, a.overrides) + + +def save_config(cfg: dict, path: str | Path) -> None: + Path(path).write_text(yaml.safe_dump(_plain(cfg), sort_keys=False)) diff --git a/robotryon/embodiment/__init__.py b/robotryon/embodiment/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..594adb3ad1fb2168879c08bb89ed0326ccd72de4 --- /dev/null +++ b/robotryon/embodiment/__init__.py @@ -0,0 +1,25 @@ +"""Robot embodiments: kinematics, cameras, silhouettes and hand retargeting from a YAML spec. + +``load_embodiment("my_robot.yaml")`` returns a :class:`MujocoEmbodiment`; see +``configs/embodiments/gr1.yaml`` for every field. +""" + +from .base import OFFSCREEN_TOLERANCE, Embodiment +from .camera import CameraModel +from .spec import EmbodimentSpec, load_spec, parse_spec + + +def load_embodiment(path_or_name): + """Embodiment from a YAML path, or from ``configs/embodiments/.yaml`` (relative to + the working directory) when given a bare name.""" + from pathlib import Path + + from .mujoco_embodiment import MujocoEmbodiment # mujoco is an optional dependency + p = Path(path_or_name) + if not p.exists() and p.suffix == "": + p = Path("configs/embodiments") / f"{path_or_name}.yaml" + return MujocoEmbodiment(load_spec(p)) + + +__all__ = ["Embodiment", "CameraModel", "EmbodimentSpec", "OFFSCREEN_TOLERANCE", + "load_embodiment", "load_spec", "parse_spec"] diff --git a/robotryon/embodiment/base.py b/robotryon/embodiment/base.py new file mode 100644 index 0000000000000000000000000000000000000000..662920c8886a1311b27d16ca88712cb3fd0e5ab6 --- /dev/null +++ b/robotryon/embodiment/base.py @@ -0,0 +1,144 @@ +"""The interface every robot is used through. + +The generative models never see joint angles, only images, a mask of the arm and hand, and a +17-keypoint hand skeleton. What a robot must provide is therefore small: + +* ``fk(q)``: wrist pose and the 17 hand keypoints of each hand in the robot base frame; +* ``camera_pose(q, view)`` and ``project(points, view)``: the head camera(s), to draw the + skeleton of a robot frame and to place a human hand recorded from the same viewpoint; +* ``silhouette(q, view)``: the rendered robot, which together with SAM3 gives the robot mask; +* ``retarget_hand(human_kp3d)``: the robot hand posed like a reconstructed human hand. + +Arrays over hands always use the order ``(left, right)``; a hand the robot does not have is NaN. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + +import numpy as np + +from robotryon.skeleton.topology import NUM_HANDS, NUM_KEYPOINTS + +from .spec import HANDS, EmbodimentSpec + +# A keypoint off screen is still drawn (OpenCV clips the bone) unless it lies further than +# this many frame sizes outside the image, where the projection is degenerate. +OFFSCREEN_TOLERANCE = 3.0 + + +class Embodiment(ABC): + spec: EmbodimentSpec + + @property + def name(self) -> str: + return self.spec.name + + @property + def views(self) -> list[str]: + return self.spec.views + + @property + def joint_names(self) -> list[str | None]: + """Layout of the recorded joint vector ``q`` (``None`` = slot not in the model).""" + return self.spec.joints + + # --- kinematics ---------------------------------------------------------------------- + @abstractmethod + def fk(self, q) -> dict[str, dict[str, np.ndarray]]: + """Forward kinematics in the base frame, metres. + + Args: + q: joint vector in ``joint_names`` order, or a ``{joint: value}`` dict. + Returns: + ``{hand: {"wrist_pos": (3,), "wrist_rot": (3, 3), "keypoints": (17, 3)}}`` for each + hand the robot has (``"left"``/``"right"``); keypoint 0 is the wrist and + ``wrist_rot`` columns are the wrist frame's axes in the base frame. + """ + + @property + def arm_joint_slices(self) -> dict[str, slice]: + """Slice of the joint/action vector holding each arm's joints (YAML ``arm_joints``).""" + return {h: slice(*r) for h, r in self.spec.arm_joints.items()} + + def hand_keypoints(self, q) -> np.ndarray: + """``(2, 17, 3)`` keypoints of ``fk(q)`` in (left, right) order, NaN for a missing hand.""" + out, f = empty_hands(), self.fk(q) + for hi, h in enumerate(HANDS): + if h in f: + out[hi] = f[h]["keypoints"] + return out + + @abstractmethod + def camera_pose(self, q, view: str) -> np.ndarray: + """4x4 camera-to-base pose of ``view`` (optical frame) at joints ``q``. + + ``q=None`` uses the default pose (useful for fixed or stereo-relative quantities). + """ + + @abstractmethod + def silhouette(self, q, view: str, image_size=None) -> np.ndarray: + """Rendered robot arm+hand mask, bool ``(H, W)``.""" + + @abstractmethod + def retarget_hand(self, human_kp3d: np.ndarray, valid=None, prev: dict | None = None): + """Robot hand posed like a human hand (vector retargeting, see ``retarget``). + + Args: + human_kp3d: ``(2, 21|17, 3)`` human joints (MediaPipe order) in any frame. + valid: ``(2,)`` bool, hands to retarget (default: finite ones). + prev: the joints returned for the previous frame of the clip; they warm-start the + optimization and its smoothing term (None for the first frame). + Returns: + ``(2, 17, 3)`` robot keypoints in the same frame (NaN for skipped hands) and a + ``{joint: value}`` dict of the retargeted robot hand joints. + """ + + # --- cameras ------------------------------------------------------------------------- + def camera(self, view: str): + return self.spec.cameras[view].intrinsics + + def stereo_extrinsic(self, view_a: str, view_b: str) -> np.ndarray: + """4x4 transform from ``view_b``'s camera frame to ``view_a``'s.""" + Ta, Tb = self.camera_pose(None, view_a), self.camera_pose(None, view_b) + return np.linalg.inv(Ta) @ Tb + + def project(self, points, view: str, q=None, frame: str = "base", image_size=None): + """``(..., 3)`` points -> ``(..., 2)`` pixels and ``(...)`` depth-ok mask. + + ``frame="base"`` needs ``q`` (the camera may move with the head); + ``frame="camera"`` takes points already in ``view``'s optical frame. + """ + pts = np.asarray(points, np.float64) + if frame == "base": + T = np.linalg.inv(self.camera_pose(q, view)) + pts = pts @ T[:3, :3].T + T[:3, 3] + elif frame != "camera": + raise ValueError(f"frame must be 'base' or 'camera', got {frame!r}") + return self.camera(view).project(pts, image_size) + + def keypoints_2d(self, points, view: str, image_size, q=None, frame: str = "base"): + """Project ``(2, 17, 3)`` keypoints into the skeleton format. + + Returns ``kp`` float32 ``(2, 17, 2)`` (0 where invisible) and ``vis`` bool ``(2, 17)``: + in front of the camera and within ``OFFSCREEN_TOLERANCE`` frames of the image. + """ + px, ok = self.project(points, view, q=q, frame=frame, image_size=image_size) + W, H = image_size + m = OFFSCREEN_TOLERANCE + u, v = px[..., 0], px[..., 1] + with np.errstate(invalid="ignore"): + vis = ok & (u > -m * W) & (u < (1 + m) * W) & (v > -m * H) & (v < (1 + m) * H) + kp = np.where(vis[..., None], px, 0.0).astype(np.float32) + return kp, vis + + def keypoints_in_view(self, q, view: str, image_size): + """Projected skeleton ``(kp, vis)`` of the robot at joints ``q`` (see keypoints_2d).""" + return self.keypoints_2d(self.hand_keypoints(q), view, image_size, q=q) + + +def empty_hands(d: int = 3) -> np.ndarray: + return np.full((NUM_HANDS, NUM_KEYPOINTS, d), np.nan) + + +__all__ = ["Embodiment", "HANDS", "OFFSCREEN_TOLERANCE", "empty_hands"] diff --git a/robotryon/embodiment/camera.py b/robotryon/embodiment/camera.py new file mode 100644 index 0000000000000000000000000000000000000000..aff1696086d39738818ec657190ad222cc58a37e --- /dev/null +++ b/robotryon/embodiment/camera.py @@ -0,0 +1,103 @@ +"""Camera intrinsics: projection of camera-frame points to pixels and back to rays. + +A camera is calibrated at one resolution (``width`` x ``height``); every method takes the +image size it should work at and rescales the intrinsics, so the same calibration serves a +video stored at any resolution with the same aspect. Points are in the OpenCV optical frame +(x right, y down, z forward). Two lens models are supported: + +* ``pinhole``: OpenCV plumb-bob distortion ``[k1, k2, p1, p2, k3]`` (empty = rectified). +* ``fisheye``: OpenCV equidistant model ``[k1, k2, k3, k4]`` (zeros = ideal equidistant). +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +import cv2 +import numpy as np + + +@dataclass +class CameraModel: + width: int + height: int + fx: float + fy: float + cx: float + cy: float + model: str = "pinhole" + dist: tuple = () + # Additive pixel offset applied after projection (a fitted principal-point residual), + # given at the calibration resolution. + pp_offset: tuple = (0.0, 0.0) + # Points closer than this to the lens are inside the robot's own head and are not drawn. + min_depth: float = 0.15 + extra: dict = field(default_factory=dict) + + def __post_init__(self): + if self.model not in ("pinhole", "fisheye"): + raise ValueError(f"unknown camera model {self.model!r}") + + @classmethod + def from_dict(cls, d: dict) -> "CameraModel": + known = {"width", "height", "fx", "fy", "cx", "cy", "model", "dist", "pp_offset", + "min_depth"} + kw = {k: d[k] for k in known if k in d} + kw["dist"] = tuple(float(x) for x in d.get("dist", ()) or ()) + kw["pp_offset"] = tuple(float(x) for x in d.get("pp_offset", (0.0, 0.0))) + return cls(**kw, extra={k: v for k, v in d.items() if k not in known}) + + # --- intrinsics at a given resolution ------------------------------------------------- + def _scale(self, image_size): + if image_size is None: + return 1.0, 1.0 + return image_size[0] / self.width, image_size[1] / self.height + + def K(self, image_size=None) -> np.ndarray: + sx, sy = self._scale(image_size) + return np.array([[self.fx * sx, 0.0, self.cx * sx], + [0.0, self.fy * sy, self.cy * sy], + [0.0, 0.0, 1.0]]) + + def D(self) -> np.ndarray: + n = 4 if self.model == "fisheye" else 5 + d = np.zeros(n) + d[:len(self.dist)] = self.dist + return d + + def offset(self, image_size=None) -> np.ndarray: + sx, sy = self._scale(image_size) + return np.array([self.pp_offset[0] * sx, self.pp_offset[1] * sy]) + + # --- projection ------------------------------------------------------------------------ + def project(self, pts_cam: np.ndarray, image_size=None): + """``(..., 3)`` camera-frame points -> ``(..., 2)`` pixels and ``(...)`` depth-ok mask. + + Points behind ``min_depth`` get NaN pixels. + """ + pts = np.asarray(pts_cam, np.float64) + shape = pts.shape[:-1] + flat = pts.reshape(-1, 3) + ok = np.isfinite(flat).all(1) & (flat[:, 2] > self.min_depth) + px = np.full((flat.shape[0], 2), np.nan) + if ok.any(): + obj = flat[ok].reshape(-1, 1, 3) + K, D = self.K(image_size), self.D() + if self.model == "fisheye": + uv, _ = cv2.fisheye.projectPoints(obj, np.zeros(3), np.zeros(3), K, D) + else: + uv, _ = cv2.projectPoints(obj, np.zeros(3), np.zeros(3), K, D) + px[ok] = uv.reshape(-1, 2) + self.offset(image_size) + return px.reshape(shape + (2,)), ok.reshape(shape) + + def normalized(self, pixels: np.ndarray, image_size=None) -> np.ndarray: + """``(..., 2)`` pixels -> ``(..., 2)`` undistorted normalized coordinates (x/z, y/z).""" + uv = np.asarray(pixels, np.float64) + shape = uv.shape[:-1] + uv = (uv.reshape(-1, 2) - self.offset(image_size)).reshape(-1, 1, 2) + K, D = self.K(image_size), self.D() + if self.model == "fisheye": + n = cv2.fisheye.undistortPoints(uv, K, D) + else: + n = cv2.undistortPoints(uv, K, D) + return n.reshape(shape + (2,)) diff --git a/robotryon/embodiment/mujoco_embodiment.py b/robotryon/embodiment/mujoco_embodiment.py new file mode 100644 index 0000000000000000000000000000000000000000..f667362e242809c34d51d53d85406f3a69be3363 --- /dev/null +++ b/robotryon/embodiment/mujoco_embodiment.py @@ -0,0 +1,320 @@ +"""MuJoCo implementation of :class:`Embodiment`, generic over the YAML description. + +The robot model (MJCF, or URDF which MuJoCo also parses) is supplied by the user; nothing +robot-specific is coded here. Forward kinematics is ``mj_kinematics`` on a pose built from +the recorded joint vector: unlisted joints stay at ``qpos0``, coupled passive joints are +filled from their sources and ``fixed_joints`` are applied last. A head camera mounted on a +model frame therefore follows the recorded head joints. + +Pin ``mujoco==3.11.0``: with 3.12 site-based camera poses silently collapse (every point +lands at camera depth ~0), so a different version raises unless +``ROBOTRYON_ALLOW_ANY_MUJOCO=1``. + +Rendering (``silhouette``) needs an OpenGL backend; set ``MUJOCO_GL=egl`` (GPU) or +``MUJOCO_GL=osmesa`` (CPU) before the first import of mujoco. +""" + +from __future__ import annotations + +import os + +import cv2 +import numpy as np + +from . import retarget as RT +from .base import Embodiment, empty_hands +from .spec import HANDS, EmbodimentSpec, load_spec + +MUJOCO_PIN = "3.11.0" +_SIL_HIDDEN_GROUP = 5 + + +def _import_mujoco(): + import mujoco + if mujoco.__version__ != MUJOCO_PIN and os.environ.get("ROBOTRYON_ALLOW_ANY_MUJOCO") != "1": + raise RuntimeError( + f"mujoco {mujoco.__version__} != {MUJOCO_PIN}: newer versions break site-based " + f"camera poses without an error. Install mujoco=={MUJOCO_PIN} or set " + f"ROBOTRYON_ALLOW_ANY_MUJOCO=1.") + return mujoco + + +class MujocoEmbodiment(Embodiment): + def __init__(self, spec: EmbodimentSpec | str): + self.spec = load_spec(spec) if isinstance(spec, (str, os.PathLike)) else spec + mj = self._mj = _import_mujoco() + if not os.path.exists(self.spec.model_path): + raise FileNotFoundError(f"robot model not found: {self.spec.model_path}") + self.model = mj.MjModel.from_xml_path(self.spec.model_path) + self.data = mj.MjData(self.model) + + names = set(n for n in self.spec.joints if n) | set(self.spec.fixed_joints) + names |= {c[0] for c in self.spec.coupling} | {c[2] for c in self.spec.coupling} + names |= {j for js in self.spec.retarget.joints.values() for j in js} + self._adr = {n: self._joint_adr(n) for n in names} + self._rt_bounds = {h: self._joint_limits(js) for h, js in self.spec.retarget.joints.items()} + self._rt_palm: dict[str, float] = {} + self._kp = {h: [(*self._frame(n), off) for n, off in + zip(self.spec.keypoints[h], self.spec.keypoint_offsets[h])] + for h in self.spec.hands} + # the wrist pose: the explicit wrist_frame, else keypoint 0 (with its offset) + self._wrist = {h: (*self._frame(self.spec.wrist_frame[h]), + self.spec.keypoint_offsets[h][0] + if self.spec.wrist_frame[h] == self.spec.keypoints[h][0] else None) + for h in self.spec.hands} + self._mount = {v: (self._frame(c.mount.frame) if c.mount.frame else None) + for v, c in self.spec.cameras.items()} + self._renderers: dict = {} + + # --- name resolution ----------------------------------------------------------------- + def _joint_adr(self, name: str) -> int: + jid = self._mj.mj_name2id(self.model, self._mj.mjtObj.mjOBJ_JOINT, name) + if jid < 0: + raise ValueError(f"joint {name!r} not in {self.spec.model_path}") + return int(self.model.jnt_qposadr[jid]) + + def _joint_limits(self, names) -> tuple[np.ndarray, np.ndarray]: + """``(lo, hi)`` of each joint: the model's limits ([-pi, pi] when unlimited), narrowed + by ``retarget.limits``.""" + lo, hi = [], [] + for n in names: + jid = self._mj.mj_name2id(self.model, self._mj.mjtObj.mjOBJ_JOINT, n) + r = self.model.jnt_range[jid] + ok = bool(self.model.jnt_limited[jid]) and r[1] > r[0] + a, b = (float(r[0]), float(r[1])) if ok else (-np.pi, np.pi) + if n in self.spec.retarget.limits: + ua, ub = self.spec.retarget.limits[n] + a, b = max(a, ua), min(b, ub) + if b <= a: + raise ValueError(f"retarget.limits for {n} [{ua}, {ub}] leaves no range " + f"inside the model limits") + lo.append(a) + hi.append(b) + return np.asarray(lo), np.asarray(hi) + + def _frame(self, name: str) -> tuple[str, int]: + """``site:x`` / ``body:x`` / ``x`` (site first, then body) -> (kind, id).""" + kind, _, bare = name.rpartition(":") + T = self._mj.mjtObj + for k in ([kind] if kind else ["site", "body"]): + oid = self._mj.mj_name2id(self.model, T.mjOBJ_SITE if k == "site" else T.mjOBJ_BODY, + bare) + if oid >= 0: + return k, oid + raise ValueError(f"frame {name!r} not in {self.spec.model_path}") + + def _pos(self, f): + p = (self.data.site_xpos if f[0] == "site" else self.data.xpos)[f[1]].copy() + if len(f) > 2 and f[2] is not None: # point fixed in the frame + p = p + self._rot(f) @ f[2] + return p + + def _rot(self, f): + m = self.data.site_xmat if f[0] == "site" else self.data.xmat + return m[f[1]].reshape(3, 3).copy() + + # --- kinematics ---------------------------------------------------------------------- + def joint_dict(self, q) -> dict[str, float]: + """Joint vector (``spec.joints`` order) or dict -> complete ``{joint: value}``.""" + if q is None: + d = {} + elif isinstance(q, dict): + d = {k: float(v) for k, v in q.items()} + else: + q = np.asarray(q, np.float64).reshape(-1) + if len(q) != len(self.spec.joints): + raise ValueError(f"expected {len(self.spec.joints)} joints, got {len(q)}") + d = {n: float(v) * self.spec.joint_scale.get(n, 1.0) for n, v in zip(self.spec.joints, q) if n} + for joint, ratio, src in self.spec.coupling: + if src in d: + d[joint] = ratio * d[src] + d.update(self.spec.fixed_joints) + return d + + def set_state(self, q) -> None: + self.data.qpos[:] = self.model.qpos0 + for n, v in self.joint_dict(q).items(): + adr = self._adr.get(n) + if adr is None: + adr = self._adr[n] = self._joint_adr(n) + self.data.qpos[adr] = v + self._mj.mj_kinematics(self.model, self.data) + + def _hand_points(self) -> np.ndarray: + """Keypoints of the current state, base frame, with the per-finger length scale. + + The scale pulls each finger's distal points toward its knuckle + (``p <- p_knuckle + s * (p - p_knuckle)``), a fitted correction for model fingers that + render longer than the real hand. + """ + out = empty_hands() + for hi, h in enumerate(HANDS): + if h not in self._kp: + continue + p = np.stack([self._pos(f) for f in self._kp[h]]) + for finger, (k, distal) in RT.CHAINS17.items(): + s = self.spec.finger_scale.get(finger) + if s is not None: + p[list(distal)] = p[k] + s * (p[list(distal)] - p[k]) + out[hi] = p + return out + + def fk(self, q) -> dict[str, dict[str, np.ndarray]]: + self.set_state(q) + kp = self._hand_points() + return {h: {"wrist_pos": self._pos(self._wrist[h]), "wrist_rot": self._rot(self._wrist[h]), + "keypoints": kp[HANDS.index(h)]} + for h in HANDS if h in self._wrist} + + def hand_keypoints(self, q) -> np.ndarray: + self.set_state(q) + return self._hand_points() + + def _camera_pose_now(self, view: str) -> np.ndarray: + m = self.spec.cameras[view].mount + if m.extrinsic is not None: + return m.extrinsic.copy() + f = self._mount[view] + R_frame = self._rot(f) + T = np.eye(4) + T[:3, :3] = R_frame @ m.rotation.T + T[:3, 3] = self._pos(f) + R_frame @ m.translation + return T + + def camera_pose(self, q, view: str) -> np.ndarray: + self.set_state(q) + return self._camera_pose_now(view) + + def keypoints_in_view(self, q, view: str, image_size): + """One FK pass -> projected skeleton ``kp (2,17,2)``, ``vis (2,17)`` of a robot frame.""" + self.set_state(q) + T = np.linalg.inv(self._camera_pose_now(view)) + pts = self._hand_points() @ T[:3, :3].T + T[:3, 3] + return self.keypoints_2d(pts, view, image_size, frame="camera") + + # --- rendering ----------------------------------------------------------------------- + def _renderer(self, view: str, image_size, scale: float): + key = (view, tuple(image_size), scale) + if key in self._renderers: + return self._renderers[key] + mj, model = self._mj, self.model + inc, exc = tuple(self.spec.silhouette_bodies), tuple(self.spec.silhouette_exclude) + for g in range(model.ngeom): + body = mj.mj_id2name(model, mj.mjtObj.mjOBJ_BODY, int(model.geom_bodyid[g])) or "" + show = body.startswith(inc) and not (exc and body.startswith(exc)) + model.geom_group[g] = 0 if show else _SIL_HIDDEN_GROUP + # Render a symmetric wide pinhole, then remap each output pixel through the real + # camera's (distorted, off-centre) ray so the silhouette matches the projection. + W, H = image_size + w, h = max(8, round(W * scale)), max(8, round(H * scale)) + cam = self.camera(view) + u, v = np.meshgrid((np.arange(w) + 0.5) * W / w - 0.5, (np.arange(h) + 0.5) * H / h - 0.5) + n = cam.normalized(np.stack([u, v], -1), image_size) + f = cam.K(image_size)[1, 1] * h / H + ext = np.abs(n.reshape(-1, 2)).max(0) * 1.03 + Wp, Hp = int(2 * np.ceil(ext[0] * f)) + 2, int(2 * np.ceil(ext[1] * f)) + 2 + mx = (n[..., 0] * f + Wp / 2.0 - 0.5).astype(np.float32) + my = (n[..., 1] * f + Hp / 2.0 - 0.5).astype(np.float32) + model.vis.global_.offwidth = max(int(model.vis.global_.offwidth), Wp) + model.vis.global_.offheight = max(int(model.vis.global_.offheight), Hp) + model.vis.map.znear = 0.001 + model.vis.quality.offsamples = 0 # MSAA would blend segmentation ids + r = mj.Renderer(model, height=Hp, width=Wp) + r.enable_segmentation_rendering() + opt = mj.MjvOption() + opt.geomgroup[:] = 0 + opt.geomgroup[0] = 1 + opt.sitegroup[:] = 0 + mcam = mj.MjvCamera() + mj.mjv_defaultCamera(mcam) + self._renderers[key] = (r, opt, mcam, mx, my, f, Hp) + return self._renderers[key] + + def silhouette(self, q, view: str, image_size=None, scale: float = 0.5) -> np.ndarray: + cam = self.camera(view) + image_size = tuple(image_size or (cam.width, cam.height)) + r, opt, mcam, mx, my, f, Hp = self._renderer(view, image_size, scale) + self.set_state(q) + T = self._camera_pose_now(view) + r.update_scene(self.data, camera=mcam, scene_option=opt) + near, far = 0.02, 6.0 + half_h = near * (Hp / 2.0) / f + for gc in r.scene.camera: + gc.pos[:] = T[:3, 3] + gc.forward[:] = T[:3, 2] + gc.up[:] = -T[:3, 1] + gc.frustum_center = 0.0 + gc.frustum_bottom, gc.frustum_top = -half_h, half_h + gc.frustum_near, gc.frustum_far = near, far + gc.orthographic = 0 + seg = r.render()[:, :, 0] + sil = cv2.remap((seg >= 0).astype(np.uint8), mx, my, cv2.INTER_NEAREST, + borderMode=cv2.BORDER_CONSTANT, borderValue=0) + return cv2.resize(sil, image_size, interpolation=cv2.INTER_NEAREST) > 0 + + # --- retargeting --------------------------------------------------------------------- + def _hand_vectors(self, hand: str, x: np.ndarray): + """Retargeting vectors (wrist frame) of one robot hand at hand-joint values ``x``.""" + self.set_state(dict(zip(self.spec.retarget.joints[hand], x))) + return RT.local_vectors(self._hand_points()[HANDS.index(hand)]) + + def _palm(self, hand: str) -> float: + if hand not in self._rt_palm: + lo, hi = self._rt_bounds[hand] + self.set_state(dict(zip(self.spec.retarget.joints[hand], np.clip(0.0, lo, hi)))) + self._rt_palm[hand] = RT.palm_size(self._hand_points()[HANDS.index(hand)]) + return self._rt_palm[hand] + + def _solve_hand(self, hand: str, h17: np.ndarray, x_prev=None): + """AnyTeleop vector retargeting of one hand: hand joints (model order) or None.""" + from scipy.optimize import least_squares + + rt = self.spec.retarget + target = RT.local_vectors(h17) + if target is None: + return None + scale = rt.scaling if rt.scaling is not None else self._palm(hand) / max(RT.palm_size(h17), 1e-6) + target = target * scale + lo, hi = self._rt_bounds[hand] + x0 = np.clip(np.zeros(len(lo)) if x_prev is None else np.asarray(x_prev, np.float64), lo, hi) + x0 = np.clip(x0, lo + 1e-9, hi - 1e-9) + w = np.sqrt(rt.smoothing) + + def residual(x): + v = self._hand_vectors(hand, x) + r = np.zeros(target.size) if v is None else (v - target).ravel() + return r if x_prev is None else np.concatenate([r, w * (x - x0)]) + + sol = least_squares(residual, x0, bounds=(lo, hi), loss="huber", f_scale=rt.huber_delta, + diff_step=1e-4, max_nfev=200) + return sol.x + + def retarget_hand(self, human_kp3d, valid=None, prev=None): + j21 = RT.to21(human_kp3d) + if valid is None: + valid = np.isfinite(j21).all(axis=(-1, -2)) + rt = self.spec.retarget + q = {} + for hi, h in enumerate(HANDS): + joints = rt.joints.get(h) + if not valid[hi] or not joints: + continue + x_prev = None + if prev and all(n in prev for n in joints): + x_prev = [prev[n] for n in joints] + x = self._solve_hand(h, RT.human17(j21[hi]), x_prev) + if x is not None: + q.update({n: float(v) for n, v in zip(joints, x)}) + out = empty_hands() + if not q: + return out, q + robot = self.hand_keypoints(q) + for hi, h in enumerate(HANDS): + if valid[hi] and all(n in q for n in rt.joints.get(h) or [None]): + placed = RT.place_on_human(robot[hi], j21[hi], rt.bone_match) + if placed is not None: + out[hi] = placed + return out, q + + +__all__ = ["MujocoEmbodiment", "MUJOCO_PIN"] diff --git a/robotryon/embodiment/retarget.py b/robotryon/embodiment/retarget.py new file mode 100644 index 0000000000000000000000000000000000000000..3d629bad335e5691c33f3a17f4f8ca53106dda79 --- /dev/null +++ b/robotryon/embodiment/retarget.py @@ -0,0 +1,122 @@ +"""Retarget a reconstructed human hand onto a robot hand (AnyTeleop vector retargeting). + +The skeleton that conditions the generative models must look like the target robot's hand, +so a human hand is not drawn directly. Instead (Appendix B.2.2): + +1. Both hands are expressed in a wrist frame (origin at the wrist, x toward the knuckle + centroid, z along the knuckle-plane normal) built identically from the human joints and + from the robot's 17 keypoints. +2. Following the vector retargeting of AnyTeleop (Qin et al., 2023), the robot hand joints are + optimized so that the robot's wrist-to-fingertip and wrist-to-middle-joint vectors match the + human's, scaled by the ratio of the palm sizes, within the joint limits of the robot model, + with a small penalty on the change from the previous frame (``VECTORS17``, + ``MujocoEmbodiment.retarget_hand``). +3. Robot forward kinematics at those joints gives the robot hand's 17 keypoints. Bone + matching stretches each robot finger to the human's knuckle distance and finger length + while keeping the robot's directions (the robot's own curl), so the fingertips reach what + the human touches. +4. The result is placed at the human wrist frame, in the human's camera frame. + +The frame helpers and the vector set live here; the optimization needs FK and is done by the +embodiment. +""" + +from __future__ import annotations + +import numpy as np + +from robotryon.skeleton.topology import MEDIAPIPE21_TO_17 + +# MediaPipe / MANO 21-joint layout: wrist, then four joints per finger, proximal to distal. +CHAINS21 = {"thumb": (1, 2, 3, 4), "index": (5, 6, 7, 8), "middle": (9, 10, 11, 12), + "ring": (13, 14, 15, 16), "little": (17, 18, 19, 20)} +# 17-keypoint layout (robotryon.skeleton.topology): (knuckle, distal points) per finger. +CHAINS17 = {"thumb": (1, (2, 3, 4)), "index": (5, (6, 7)), "middle": (8, (9, 10)), + "ring": (11, (12, 13)), "little": (14, (15, 16))} +FINGERS4 = ("index", "middle", "ring", "little") + +# Robot wrist frame from the 17-keypoint layout: wrist, index knuckle, little knuckle. +WRIST17, INDEX17, LITTLE17 = 0, 5, 14 + +# Retargeting vectors as (origin, target) pairs of the 17-keypoint layout: wrist to every +# fingertip and wrist to the joint before it (thumb distal, index/middle/ring/little middle). +TIPS17 = tuple(distal[-1] for _, distal in CHAINS17.values()) +MIDS17 = (3, 6, 9, 12, 15) +VECTORS17 = tuple((WRIST17, t) for t in TIPS17) + tuple((WRIST17, m) for m in MIDS17) + + +def to21(kp: np.ndarray) -> np.ndarray: + """``(..., 17|21, 3)`` -> 21 joints. Missing DIP joints are put midway between PIP and + tip, which leaves every finger length and direction used below unchanged.""" + kp = np.asarray(kp, np.float64) + if kp.shape[-2] == 21: + return kp + out = np.zeros(kp.shape[:-2] + (21, 3)) + out[..., MEDIAPIPE21_TO_17, :] = kp + for pip, dip, tip in ((6, 7, 8), (10, 11, 12), (14, 15, 16), (18, 19, 20)): + out[..., dip, :] = 0.5 * (out[..., pip, :] + out[..., tip, :]) + return out + + +def wrist_frame(w, imcp, lmcp): + """Right-handed frame (columns = axes) at the wrist, or None when degenerate.""" + x = (imcp + lmcp) * 0.5 - w + n = np.cross(imcp - w, lmcp - w) + nx, nn = np.linalg.norm(x), np.linalg.norm(n) + if nx < 1e-9 or nn < 1e-9: + return None + x, n = x / nx, n / nn + return np.stack([x, np.cross(n, x), n], axis=1) + + +def palm_size(p17: np.ndarray) -> float: + """Mean wrist-to-knuckle distance of the four fingers (17-keypoint layout).""" + return float(np.mean([np.linalg.norm(p17[CHAINS17[f][0]] - p17[WRIST17]) for f in FINGERS4])) + + +def local_vectors(p17: np.ndarray): + """``(17, 3)`` keypoints -> ``(len(VECTORS17), 3)`` vectors in the hand's wrist frame, + or None when the frame is degenerate.""" + R = wrist_frame(p17[WRIST17], p17[INDEX17], p17[LITTLE17]) + if R is None: + return None + o, t = zip(*VECTORS17) + return (p17[list(t)] - p17[list(o)]) @ R + + +def human17(j21: np.ndarray) -> np.ndarray: + """``(21, 3)`` human joints -> the 17-keypoint layout.""" + return np.asarray(j21, np.float64)[list(MEDIAPIPE21_TO_17)] + + +def bone_match(local17: np.ndarray, j21: np.ndarray) -> np.ndarray: + """Stretch robot finger segments (wrist-frame 17 keypoints) to the human's lengths. + + Each knuckle is moved along its own direction to the human wrist-to-knuckle distance, then + the rest of the finger is scaled about the knuckle so knuckle-to-tip matches the human's. + Directions stay the robot's. + """ + out = local17.copy() + for f, (k, distal) in CHAINS17.items(): + hm, ht = CHAINS21[f][0], CHAINS21[f][-1] + r_palm, h_palm = np.linalg.norm(local17[k]), np.linalg.norm(j21[hm] - j21[0]) + if r_palm > 1e-6 and h_palm > 1e-6: + out[k] = local17[k] * (h_palm / r_palm) + r_fin = np.linalg.norm(local17[distal[-1]] - local17[k]) + h_fin = np.linalg.norm(j21[ht] - j21[hm]) + s = h_fin / r_fin if (r_fin > 1e-6 and h_fin > 1e-6) else 1.0 + for i in distal: + out[i] = out[k] + s * (local17[i] - local17[k]) + return out + + +def place_on_human(robot17: np.ndarray, j21: np.ndarray, match_bones: bool = True): + """Move robot keypoints (any frame) onto the human wrist frame; None if degenerate.""" + Rr = wrist_frame(robot17[WRIST17], robot17[INDEX17], robot17[LITTLE17]) + Rh = wrist_frame(j21[0], j21[5], j21[17]) + if Rr is None or Rh is None: + return None + local = (robot17 - robot17[WRIST17]) @ Rr + if match_bones: + local = bone_match(local, j21) + return j21[0] + local @ Rh.T diff --git a/robotryon/embodiment/spec.py b/robotryon/embodiment/spec.py new file mode 100644 index 0000000000000000000000000000000000000000..93c12dd366426636f607e5e094e46de655b6bac3 --- /dev/null +++ b/robotryon/embodiment/spec.py @@ -0,0 +1,201 @@ +"""Embodiment description parsed from YAML (see ``configs/embodiments/gr1.yaml``). + +Everything that differs between robots lives here as data: the kinematic model file, the +layout of the recorded joint vector, passive joints coupled to actuated ones, the link that +stands for each of the 17 hand keypoints, the cameras, the links that make up the robot's +silhouette, and which hand joints the human-hand retargeting moves. Parsing is pure Python +so configs can be validated without a physics engine. + +Strings may contain ``{s}``, which is replaced by each hand's side token (``hands.sides``), so +a mirror-symmetric robot is described once for both hands. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import numpy as np +import yaml + +from robotryon.skeleton.topology import NUM_KEYPOINTS + +from .camera import CameraModel + +HANDS = ("left", "right") + + +@dataclass +class CameraMount: + """Where a camera sits: on a model frame that moves with the robot, or fixed in the base. + + ``frame`` is ``site:``/``body:`` (or a bare name, resolved site first); + ``rotation`` maps the frame's axes to the optical frame (``p_opt = R @ p_frame``); + ``translation`` offsets the lens origin in the mount frame (e.g. a stereo baseline). + ``extrinsic`` instead gives a fixed 4x4 camera-to-base pose. + """ + frame: str | None = None + rotation: np.ndarray = field(default_factory=lambda: np.eye(3)) + translation: np.ndarray = field(default_factory=lambda: np.zeros(3)) + extrinsic: np.ndarray | None = None + + +@dataclass +class CameraSpec: + intrinsics: CameraModel + mount: CameraMount + video_key: str | None = None + + +@dataclass +class RetargetSpec: + """Vector retargeting of a human hand onto the robot hand (``robotryon.embodiment.retarget``). + + ``joints``: per hand, the actuated hand joints the optimization moves; coupled joints follow + their source through ``coupling`` and the limits come from the robot model. + ``scaling``: factor on the human vectors, or None for the ratio of the palm sizes. + ``smoothing``: weight of the squared change from the previous frame's joints. + ``huber_delta``: Huber threshold (m) of the vector residuals. + ``limits``: ``{joint: (lo, hi)}`` narrowing the model's joint limits for the retargeting, e.g. + to keep finger joints out of hyperextension the model allows but the real hand never uses. + """ + joints: dict[str, list[str]] = field(default_factory=dict) + limits: dict[str, tuple[float, float]] = field(default_factory=dict) + scaling: float | None = None + smoothing: float = 4e-3 + huber_delta: float = 0.02 + bone_match: bool = True + + +@dataclass +class EmbodimentSpec: + name: str + model_path: str + joints: list[str | None] + joint_scale: dict[str, float] # joint -> model value / recorded value + arm_joints: dict[str, tuple[int, int]] # hand -> [start, stop) in the joint vector + coupling: list[tuple[str, float, str]] # (joint, ratio, source joint) + fixed_joints: dict[str, float] + keypoints: dict[str, list[str]] # hand -> 17 frame names + keypoint_offsets: dict[str, list] # hand -> 17 offsets (m, in the frame) or None + wrist_frame: dict[str, str] # hand -> frame giving the wrist pose + finger_scale: dict[str, float] + cameras: dict[str, CameraSpec] + silhouette_bodies: list[str] + silhouette_exclude: list[str] + retarget: RetargetSpec = field(default_factory=RetargetSpec) + raw: dict = field(default_factory=dict) + + @property + def hands(self) -> list[str]: + return [h for h in HANDS if h in self.keypoints] + + @property + def views(self) -> list[str]: + return list(self.cameras) + + +def _fmt(x: Any, s: str) -> Any: + return x.replace("{s}", s) if isinstance(x, str) else x + + +def _per_hand(value, hand: str): + """``value`` or ``value[hand]`` when it is given per hand.""" + if isinstance(value, dict) and hand in value: + return value[hand] + return value + + +def parse_spec(cfg: dict, base_dir: str | Path | None = None) -> EmbodimentSpec: + hands_cfg = cfg.get("hands", {}) + sides = hands_cfg.get("sides", {"left": "left", "right": "right"}) + + keypoints, keypoint_offsets, wrist_frame = {}, {}, {} + for hand, tok in sides.items(): + if hand not in HANDS: + raise ValueError(f"hands.sides keys must be left/right, got {hand!r}") + entries = (hands_cfg.get(hand) or {}).get("keypoints") or hands_cfg.get("keypoints") + if entries is None: + raise ValueError(f"no keypoints given for the {hand} hand") + if len(entries) != NUM_KEYPOINTS: + raise ValueError(f"{hand} hand needs {NUM_KEYPOINTS} keypoints, got {len(entries)}") + # an entry is a frame name, or {frame: name, offset: [x, y, z]} for a point fixed in it + names = [_fmt(e["frame"] if isinstance(e, dict) else e, tok) for e in entries] + keypoint_offsets[hand] = [np.asarray(_per_hand(e["offset"], hand), np.float64).reshape(3) + if isinstance(e, dict) and e.get("offset") is not None else None + for e in entries] + keypoints[hand] = names + wf = (hands_cfg.get(hand) or {}).get("wrist_frame") or hands_cfg.get("wrist_frame") + wrist_frame[hand] = _fmt(wf, tok) if wf else names[0] + + coupling = [] + for c in cfg.get("coupling", []) or []: + templated = "{s}" in c["joint"] or "{s}" in c["source"] + for tok in (sides.values() if templated else [""]): + coupling.append((_fmt(c["joint"], tok), float(c["ratio"]), _fmt(c["source"], tok))) + + cameras = {} + for view, c in (cfg.get("cameras") or {}).items(): + m = c.get("mount", {}) or {} + ext = c.get("extrinsic") + mount = CameraMount( + frame=m.get("frame"), + rotation=np.asarray(m.get("rotation", np.eye(3)), np.float64).reshape(3, 3), + translation=np.asarray(m.get("translation", [0, 0, 0]), np.float64).reshape(3), + extrinsic=None if ext is None else np.asarray(ext, np.float64).reshape(4, 4)) + if mount.frame is None and mount.extrinsic is None: + raise ValueError(f"camera {view!r} needs mount.frame or extrinsic") + cameras[view] = CameraSpec(CameraModel.from_dict(c["intrinsics"]), mount, + c.get("video_key")) + + rt = cfg.get("retarget", {}) or {} + if "drivers" in rt: + raise ValueError("retarget.drivers (percentile retargeting) was replaced by vector " + "retargeting: list the hand joints under retarget.joints instead") + rt_joints = {hand: [_fmt(j, tok) for j in (_per_hand(rt.get("joints") or [], hand) or [])] + for hand, tok in sides.items()} + limits = {} + for j, r in (rt.get("limits") or {}).items(): + for hand, tok in sides.items(): + rr = _per_hand(r, hand) + limits[_fmt(j, tok)] = (float(rr[0]), float(rr[1])) + scaling = rt.get("scaling", "auto") + retarget = RetargetSpec( + joints=rt_joints, + limits=limits, + scaling=None if scaling in (None, "auto") else float(scaling), + smoothing=float(rt.get("smoothing", 4e-3)), + huber_delta=float(rt.get("huber_delta", 0.02)), + bone_match=bool(rt.get("bone_match", True))) + + model_path = str(cfg["model_path"]) + if base_dir is not None and not Path(model_path).is_absolute(): + model_path = str(Path(base_dir) / model_path) + sil = cfg.get("silhouette", {}) or {} + toks = list(sides.values()) + expand = lambda xs: sorted({_fmt(x, t) for x in xs or [] for t in toks}) # noqa: E731 + return EmbodimentSpec( + name=str(cfg.get("name", "robot")), + model_path=model_path, + joints=[j["name"] if isinstance(j, dict) else j for j in cfg.get("joints", [])], + joint_scale={j["name"]: float(j["scale"]) for j in cfg.get("joints", []) + if isinstance(j, dict) and "scale" in j}, + arm_joints={h: (int(r[0]), int(r[1])) for h, r in (cfg.get("arm_joints") or {}).items()}, + coupling=coupling, + fixed_joints={k: float(v) for k, v in (cfg.get("fixed_joints") or {}).items()}, + keypoints=keypoints, + keypoint_offsets=keypoint_offsets, + wrist_frame=wrist_frame, + finger_scale={k: float(v) for k, v in (hands_cfg.get("finger_scale") or {}).items()}, + cameras=cameras, + silhouette_bodies=expand(sil.get("bodies")), + silhouette_exclude=expand(sil.get("exclude")), + retarget=retarget, + raw=cfg, + ) + + +def load_spec(path: str | Path) -> EmbodimentSpec: + path = Path(path) + return parse_spec(yaml.safe_load(path.read_text()), base_dir=path.parent) diff --git a/robotryon/idm/__init__.py b/robotryon/idm/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1b17c0cfa1628958840bbfa9128a581cb3b2155a --- /dev/null +++ b/robotryon/idm/__init__.py @@ -0,0 +1,5 @@ +"""Inverse dynamics model: training data, model, action labelling and wrist correction. + +Torch-dependent modules (``model``, ``data``, ``labeling``) are imported explicitly by the +scripts; ``wrist_correction`` only needs numpy/scipy. +""" diff --git a/robotryon/idm/data.py b/robotryon/idm/data.py new file mode 100644 index 0000000000000000000000000000000000000000..8e81065a46a0fcf0498ce8fd5fab12c3eda97f66 --- /dev/null +++ b/robotryon/idm/data.py @@ -0,0 +1,319 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# Adapted from GR00T-Dreams (Apache-2.0): gr00t/data/dataset.py (LeRobot v2.1 reader), +# gr00t/data/transform/{video,state_action,mirror}.py, gr00t/model/transforms_idm.py and the +# multi-embodiment mixture. Rewritten as one compact pipeline; the eval-mode image path and +# the q99 normalisation are numerically identical to the original. +"""IDM training data: LeRobot v2.1 datasets of robot videos with joint actions. + +Each embodiment is declared in the config (no robot list in code):: + + name: openarm_rh56f1 # also the key of its statistics in the checkpoint metadata + embodiment_id: 29 # row of the per-embodiment action encoder/decoder + root: /path/to/lerobot_dataset + views: [camera_ego_left, camera_ego_right] # LeRobot video keys; one key = mono + view_ids: [0, 1, 2, 3] # view-embedding id of [t0 v0, t0 v1, t1 v0, t1 v1] + action_groups: [neck_joints, left_arm_joints, ...] # modality.json groups, raw order + hand_groups: [left_hand_joints, right_hand_joints] # smoothed with the short SG window + norm: q99 # or q99_const (near-constant joints -> 0) + video_size: [512, 288] # dataset frame size; labelled videos are resized to it + mirror: {prob, view_swap, swap, sign} # optional left/right reflection augmentation + +A sample is frames (t, t+K) of every view and the chunk a_{t:t+K}, normalised to [-1, 1] with +the dataset's q01/q99 and zero-padded to the model's action_dim. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import torch +from torch.utils.data import Dataset + +IMAGE_SIZE = 224 +_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32) +_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32) +CONST_EPS = 1e-3 + + +# ----------------------------------------------------------------------------- normalisation +def normalize(x: torch.Tensor, q01: torch.Tensor, q99: torch.Tensor, mode: str) -> torch.Tensor: + q01, q99 = q01.to(x.dtype), q99.to(x.dtype) + if mode == "q99": # constant dims (q01 == q99) pass through + live = q01 != q99 + out = x.clone() + elif mode == "q99_const": # near-constant dims are mapped to 0 + live = (q99 - q01) > CONST_EPS + out = torch.zeros_like(x) + else: + raise ValueError(f"unknown norm mode {mode!r}") + out[..., live] = 2 * (x[..., live] - q01[live]) / (q99[live] - q01[live]) - 1 + return out.clamp(-1, 1) + + +def unnormalize(x: torch.Tensor, q01: torch.Tensor, q99: torch.Tensor, mode: str) -> torch.Tensor: + q01, q99 = q01.to(x.dtype), q99.to(x.dtype) + out = (x + 1) / 2 * (q99 - q01) + q01 + if mode == "q99_const": + out = torch.where((q99 - q01) > CONST_EPS, out, ((q01 + q99) / 2).expand_as(x)) + return out + + +def group_stats(stats: dict, groups: list[str]) -> tuple[torch.Tensor, torch.Tensor]: + """Concatenate the per-group q01/q99 (metadata["statistics"]["action"]) in group order.""" + q01 = np.concatenate([stats[g]["q01"] for g in groups]) + q99 = np.concatenate([stats[g]["q99"] for g in groups]) + return torch.tensor(q01.tolist()), torch.tensor(q99.tolist()) + + +# ----------------------------------------------------------------------------- images +def preprocess_frames(frames: np.ndarray, train: bool = False, fov_range=None, + color_jitter=None) -> np.ndarray: + """(N, H, W, 3) uint8 -> (N, 3, 224, 224) float32, ImageNet-normalised. + + All N frames (every view and both timesteps of one sample) share one random draw, so the + stereo geometry and the temporal pair stay consistent. Train-time augmentation: field of + view (random zoom-in crop or zoom-out black padding at a random offset) and colour jitter. + """ + import torchvision.transforms.v2 as T + import torchvision.transforms.v2.functional as TF + + x = torch.from_numpy(np.ascontiguousarray(frames)).to(torch.float32).div(255.0).permute(0, 3, 1, 2) + if train and fov_range is not None: + x = _random_fov(x, *fov_range) + x = TF.resize(x, [IMAGE_SIZE, IMAGE_SIZE], interpolation=T.InterpolationMode.BILINEAR, + antialias=True) + if train and color_jitter is not None: + x = T.ColorJitter(*color_jitter)(x) + u8 = (x.permute(0, 2, 3, 1) * 255).to(torch.uint8).numpy() + # Channels-last memory (a transposed view, no copy): cuDNN picks its convolution kernel by + # layout, and this is the layout the IDM was trained and evaluated with. + return ((u8.astype(np.float32) / 255.0 - _MEAN) / _STD).transpose(0, 3, 1, 2) + + +def _random_fov(x: torch.Tensor, lo: float, hi: float) -> torch.Tensor: + import torchvision.transforms.v2.functional as TF + + H, W = x.shape[-2:] + s = float(torch.empty(1).uniform_(lo, hi)) + if s <= 1.0: + h, w = max(1, round(H * s)), max(1, round(W * s)) + top, left = int(torch.randint(0, H - h + 1, (1,))), int(torch.randint(0, W - w + 1, (1,))) + return x[..., top:top + h, left:left + w] + h, w = round(H * s), round(W * s) + top, left = int(torch.randint(0, h - H + 1, (1,))), int(torch.randint(0, w - W + 1, (1,))) + return TF.pad(x, [left, top, w - W - left, h - H - top], fill=0.0) + + +def model_inputs(images: np.ndarray, spec: dict, horizon: int, action_dim: int) -> dict: + """The per-sample fields every batch needs besides the action target.""" + n_native = native_dim(spec) + mask = np.zeros((horizon, action_dim), dtype=bool) + mask[:, :n_native] = True + view_ids = list(spec["view_ids"]) + assert len(view_ids) == len(images), f"view_ids {view_ids} for {len(images)} images" + return dict(images=images, view_ids=np.asarray(view_ids, np.int64), n_images=len(images), + embodiment_id=int(spec["embodiment_id"]), actions_mask=mask) + + +def collate(samples: list[dict]) -> dict: + out = {} + for k in samples[0]: + vals = [s[k] for s in samples] + cat = k in ("images", "view_ids") + out[k] = torch.from_numpy(np.concatenate(vals) if cat else np.stack(vals)) \ + if isinstance(vals[0], np.ndarray) else torch.tensor(vals) + return out + + +# ----------------------------------------------------------------------------- LeRobot reader +class LeRobotMeta: + """meta/{info,modality,episodes,stats} of a LeRobot v2.1 dataset, reduced to what the IDM + needs: per-group column slices of the action, video paths and episode lengths.""" + + def __init__(self, root: str | Path): + self.root = Path(root) + m = self.root / "meta" + self.info = json.loads((m / "info.json").read_text()) + self.modality = json.loads((m / "modality.json").read_text()) + self.episodes = [json.loads(l) for l in (m / "episodes.jsonl").read_text().splitlines() if l.strip()] + self.stats = json.loads((m / "stats.json").read_text()) + + def action_slice(self, group: str) -> tuple[str, int, int]: + g = self.modality["action"][group] + return g.get("original_key") or "action", g["start"], g["end"] + + def action_statistics(self, groups: list[str]) -> dict: + out = {} + for g in groups: + col, a, b = self.action_slice(g) + out[g] = {k: list(np.asarray(v)[a:b]) for k, v in self.stats[col].items()} + return out + + def parquet(self, ep: int) -> Path: + return self.root / self.info["data_path"].format( + episode_chunk=ep // self.info["chunks_size"], episode_index=ep) + + def video(self, ep: int, view: str) -> Path: + key = self.modality["video"][view].get("original_key") or view + return self.root / self.info["video_path"].format( + episode_chunk=ep // self.info["chunks_size"], episode_index=ep, video_key=key) + + +def native_dim(spec: dict) -> int: + return int(spec["action_dim"]) + + +def frames_at(path: Path, timestamps: np.ndarray) -> np.ndarray: + """Frames nearest to the given timestamps (seconds), (N, H, W, 3) uint8 RGB.""" + import decord + + vr = decord.VideoReader(str(path)) + ts = vr.get_frame_timestamp(range(len(vr)))[:, :1] + return vr.get_batch(np.abs(ts - timestamps).argmin(axis=0)).asnumpy() + + +class Mirror: + """Left/right reflection: flip the frames, exchange the two eyes, swap left/right joint + groups (with a per-joint sign) and negate the joints listed in ``sign``. Declared per + embodiment because the sign map depends on the kinematics.""" + + def __init__(self, cfg: dict, groups: list[str], views: list[str]): + self.prob = float(cfg.get("prob", 0.0)) + self.view_perm = list(range(len(views))) + for a, b in cfg.get("view_swap", []): + ia, ib = views.index(a), views.index(b) + self.view_perm[ia], self.view_perm[ib] = ib, ia + self.swaps = [(groups.index(a), groups.index(b), s) for a, b, s in cfg.get("swap", [])] + self.signs = {groups.index(g): np.asarray(s) for g, s in (cfg.get("sign") or {}).items()} + + def __call__(self, frames: list[np.ndarray], parts: list[np.ndarray]): + if self.prob <= 0 or float(torch.rand(())) >= self.prob: + return frames, parts + frames = [np.ascontiguousarray(frames[i][..., ::-1, :]) for i in self.view_perm] + parts = list(parts) + for ia, ib, s in self.swaps: + s = 1.0 if s is None else np.asarray(s) + parts[ia], parts[ib] = parts[ib] * s, parts[ia] * s + for i, s in self.signs.items(): + parts[i] = parts[i] * s + return frames, parts + + +class LeRobotIDMDataset(Dataset): + """Every frame t of every episode is a sample: frames (t, t+K) of all views (clamped to the + episode), and actions a_t .. a_{t+H-1} (the last action repeats past the end).""" + + def __init__(self, spec: dict, model_cfg: dict, data_cfg: dict, train: bool = True): + self.spec, self.train = spec, train + self.meta = LeRobotMeta(spec["root"]) + self.groups = list(spec["action_groups"]) + self.H, self.A = int(model_cfg["action_horizon"]), int(model_cfg["action_dim"]) + self.K = int(data_cfg["frame_gap"]) + eps = self.meta.episodes[: data_cfg.get("max_episodes") or None] + self.steps = [(e["episode_index"], t, e["length"]) for e in eps for t in range(e["length"])] + self.statistics = self.meta.action_statistics(self.groups) + self.q01, self.q99 = group_stats(self.statistics, self.groups) + widths = [self.meta.action_slice(g)[2] - self.meta.action_slice(g)[1] for g in self.groups] + assert sum(widths) == native_dim(spec) <= self.A, (spec["name"], widths, self.A) + self.mirror = Mirror(spec.get("mirror") or {}, self.groups, list(spec["views"])) + self.fov = spec.get("fov_range", data_cfg.get("fov_range")) + self.jitter = data_cfg.get("color_jitter") + self._cache = (None, None) + + def __len__(self): + return len(self.steps) + + def metadata(self) -> dict: + return {"statistics": {"action": self.statistics}, "spec": dict(self.spec)} + + def _episode(self, ep: int): + if self._cache[0] != ep: + import pandas as pd + self._cache = (ep, pd.read_parquet(self.meta.parquet(ep))) + return self._cache[1] + + def __getitem__(self, i: int) -> dict: + ep, t, L = self.steps[i] + df = self._episode(ep) + obs = np.clip(np.array([t, t + self.K]), 0, L - 1) + ts = df["timestamp"].to_numpy()[obs] + frames = [frames_at(self.meta.video(ep, v), ts) for v in self.spec["views"]] + rows = np.minimum(np.arange(t, t + self.H), L - 1) + parts = [] + for g in self.groups: + col, a, b = self.meta.action_slice(g) + parts.append(np.stack(df[col].to_numpy()[rows])[:, a:b].astype(np.float64)) + if self.train: + frames, parts = self.mirror(frames, parts) + # (v, t) -> images ordered [t0 v0, t0 v1, t1 v0, t1 v1] + stacked = np.stack(frames, 1).reshape(-1, *frames[0].shape[1:]) + images = preprocess_frames(stacked, self.train, self.fov, self.jitter) + act = normalize(torch.from_numpy(np.concatenate(parts, 1)), self.q01, self.q99, + self.spec.get("norm", "q99")).numpy() + sample = model_inputs(images, self.spec, self.H, self.A) + sample["actions"] = np.pad(act, ((0, 0), (0, self.A - act.shape[1]))).astype(np.float32) + return sample + + +class MixtureDataset(Dataset): + """Weighted mixture of embodiments. Each epoch, embodiment k fills round(w_k * N) of the + N = sum(len) slots with a fresh random permutation of its samples (repeated if up-weighted).""" + + def __init__(self, datasets: list[LeRobotIDMDataset], weights: list[float], seed: int = 42): + self.datasets = datasets + w = np.asarray(weights, np.float64) + self.weights = w / w.sum() + self.seed, self.N = seed, sum(len(d) for d in datasets) + self.epoch, self.table = None, None + self.set_epoch(0) + + def set_epoch(self, epoch: int) -> None: + if epoch == self.epoch: + return + self.epoch = epoch + rng = np.random.default_rng(self.seed + epoch) + raw = self.weights * self.N + slots = np.floor(raw).astype(int) + for i in np.argsort(-(raw - slots))[: self.N - slots.sum()]: + slots[i] += 1 + rows = [] + for k, (ds, n_slots) in enumerate(zip(self.datasets, slots)): + reps, rem = divmod(int(n_slots), len(ds)) + local = np.concatenate([rng.permutation(len(ds)) for _ in range(reps)] + + [rng.permutation(len(ds))[:rem]]) + rows.append(np.stack([np.full(n_slots, k), local], 1)) + self.table = np.concatenate(rows).astype(np.int64) + + def __len__(self): + return self.N + + def __getitem__(self, i): + k, j = self.table[i] + return self.datasets[k][j] + + +def embodiment_specs(data_cfg: dict) -> dict[str, dict]: + """Embodiments declared in the config, keyed by name, with ``data.roots`` overrides applied.""" + roots = data_cfg.get("roots") or {} + specs = {} + for e in data_cfg["embodiments"]: + spec = dict(e) + spec["root"] = roots.get(spec["name"], spec.get("root")) + specs[spec["name"]] = spec + return specs diff --git a/robotryon/idm/dit.py b/robotryon/idm/dit.py new file mode 100644 index 0000000000000000000000000000000000000000..25341fd3a9e02e4d5de7b3edab7692032bb931b8 --- /dev/null +++ b/robotryon/idm/dit.py @@ -0,0 +1,143 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# Adapted from GR00T-Dreams (Apache-2.0), gr00t/model/action_head/cross_attention_dit.py. +# Changes: plain nn.Module instead of diffusers ModelMixin/ConfigMixin, unused options removed. +# Parameter names are unchanged so GR00T-Dreams IDM checkpoints load as-is. +"""Diffusion transformer (DiT) of the flow-matching action head, and the self-attention +transformer that mixes the visual tokens before the DiT cross-attends to them.""" + +from __future__ import annotations + +from typing import Optional + +import torch +import torch.nn.functional as F +from diffusers.models.attention import Attention, FeedForward +from diffusers.models.embeddings import TimestepEmbedding, Timesteps +from torch import nn + + +class TimestepEncoder(nn.Module): + def __init__(self, embedding_dim: int): + super().__init__() + self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=1) + self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) + + def forward(self, timesteps): + dtype = next(self.parameters()).dtype + return self.timestep_embedder(self.time_proj(timesteps).to(dtype)) + + +class AdaLayerNorm(nn.Module): + def __init__(self, embedding_dim: int, norm_eps: float = 1e-5): + super().__init__() + self.silu = nn.SiLU() + self.linear = nn.Linear(embedding_dim, embedding_dim * 2) + self.norm = nn.LayerNorm(embedding_dim, norm_eps, False) + + def forward(self, x, temb): + scale, shift = self.linear(self.silu(temb)).chunk(2, dim=1) + return self.norm(x) * (1 + scale[:, None]) + shift[:, None] + + +class BasicTransformerBlock(nn.Module): + """Attention (self, or cross when encoder states are given) + feed-forward, pre-norm.""" + + def __init__(self, dim: int, num_attention_heads: int, attention_head_dim: int, + dropout: float = 0.0, activation_fn: str = "geglu", attention_bias: bool = False, + upcast_attention: bool = False, norm_elementwise_affine: bool = True, + norm_type: str = "layer_norm", norm_eps: float = 1e-5, final_dropout: bool = False): + super().__init__() + self.norm_type = norm_type + if norm_type == "ada_norm": + self.norm1 = AdaLayerNorm(dim) + else: + self.norm1 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) + self.attn1 = Attention(query_dim=dim, heads=num_attention_heads, dim_head=attention_head_dim, + dropout=dropout, bias=attention_bias, cross_attention_dim=None, + upcast_attention=upcast_attention, out_bias=True) + self.norm3 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine) + self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn, + final_dropout=final_dropout, inner_dim=None, bias=True) + self.final_dropout = nn.Dropout(dropout) if final_dropout else None + + def forward(self, hidden_states, encoder_hidden_states=None, temb=None): + if self.norm_type == "ada_norm": + h = self.norm1(hidden_states, temb) + else: + h = self.norm1(hidden_states) + attn = self.attn1(h, encoder_hidden_states=encoder_hidden_states, attention_mask=None) + if self.final_dropout: + attn = self.final_dropout(attn) + hidden_states = attn + hidden_states + return self.ff(self.norm3(hidden_states)) + hidden_states + + +class DiT(nn.Module): + """Action-token transformer conditioned on the flow time (AdaLN) and on the visual tokens. + With ``interleave_self_attention`` the odd blocks are self-attention only.""" + + def __init__(self, num_attention_heads: int = 16, attention_head_dim: int = 64, + output_dim: int = 1024, num_layers: int = 8, dropout: float = 0.2, + attention_bias: bool = True, activation_fn: str = "gelu-approximate", + upcast_attention: bool = False, norm_type: str = "ada_norm", + norm_elementwise_affine: bool = False, norm_eps: float = 1e-5, + final_dropout: bool = True, interleave_self_attention: bool = True): + super().__init__() + self.inner_dim = num_attention_heads * attention_head_dim + self.interleave_self_attention = interleave_self_attention + self.timestep_encoder = TimestepEncoder(self.inner_dim) + self.transformer_blocks = nn.ModuleList([ + BasicTransformerBlock(self.inner_dim, num_attention_heads, attention_head_dim, + dropout=dropout, activation_fn=activation_fn, + attention_bias=attention_bias, upcast_attention=upcast_attention, + norm_type=norm_type, norm_elementwise_affine=norm_elementwise_affine, + norm_eps=norm_eps, final_dropout=final_dropout) + for _ in range(num_layers)]) + self.norm_out = nn.LayerNorm(self.inner_dim, elementwise_affine=False, eps=1e-6) + self.proj_out_1 = nn.Linear(self.inner_dim, 2 * self.inner_dim) + self.proj_out_2 = nn.Linear(self.inner_dim, output_dim) + + def forward(self, hidden_states, encoder_hidden_states, timestep): + temb = self.timestep_encoder(timestep) + hidden_states = hidden_states.contiguous() + encoder_hidden_states = encoder_hidden_states.contiguous() + for idx, block in enumerate(self.transformer_blocks): + cross = not (idx % 2 == 1 and self.interleave_self_attention) + hidden_states = block(hidden_states, encoder_hidden_states if cross else None, temb) + shift, scale = self.proj_out_1(F.silu(temb)).chunk(2, dim=1) + hidden_states = self.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None] + return self.proj_out_2(hidden_states) + + +class SelfAttentionTransformer(nn.Module): + def __init__(self, num_attention_heads: int = 16, attention_head_dim: int = 64, + num_layers: int = 4, dropout: float = 0.2, attention_bias: bool = True, + activation_fn: str = "gelu-approximate", upcast_attention: bool = False, + final_dropout: bool = True): + super().__init__() + inner = num_attention_heads * attention_head_dim + self.transformer_blocks = nn.ModuleList([ + BasicTransformerBlock(inner, num_attention_heads, attention_head_dim, dropout=dropout, + activation_fn=activation_fn, attention_bias=attention_bias, + upcast_attention=upcast_attention, final_dropout=final_dropout) + for _ in range(num_layers)]) + + def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None): + hidden_states = hidden_states.contiguous() + for block in self.transformer_blocks: + hidden_states = block(hidden_states) + return hidden_states diff --git a/robotryon/idm/labeling.py b/robotryon/idm/labeling.py new file mode 100644 index 0000000000000000000000000000000000000000..e48213dd6bde1585f88121f9e4c7baf1840cce97 --- /dev/null +++ b/robotryon/idm/labeling.py @@ -0,0 +1,99 @@ +"""Action labels for a generated robot video (paper Sec. 3.4, App. B.2.6). + +1. Sliding window, stride 1: at every frame t the IDM reads frames (t, t+K) (clamped to the + clip) and predicts the chunk a_t .. a_{t+H-1}. +2. Chunk averaging: frame u appears in the H chunks that start at u-H+1 .. u; its label is the + mean of those predictions (fewer at the clip ends). +3. Savitzky-Golay smoothing (polynomial order 2) per joint group: a long window on the arm-like + groups, a short one on the hand groups so that grasp onsets stay sharp. +4. Optional wrist correction (``wrist_correction.py``). +5. Proprioceptive state from the actions: q_t = a_{t-d} (the robot reaches a commanded pose + d frames later); the first d states repeat a_0. +""" + +from __future__ import annotations + +import numpy as np +import torch +from scipy.signal import savgol_filter + +from .data import collate, model_inputs, preprocess_frames, unnormalize + + +class ChunkAccumulator: + """Average every chunk prediction that covers a frame.""" + + def __init__(self, length: int, action_dim: int, horizon: int): + self.L, self.H = int(length), int(horizon) + self._sum = np.zeros((self.L, action_dim), dtype=np.float64) + self._cnt = np.zeros((self.L, 1), dtype=np.float64) + + def add(self, start: int, chunk: np.ndarray) -> None: + n = min(self.H, len(chunk), self.L - start) + if n > 0: + self._sum[start:start + n] += chunk[:n] + self._cnt[start:start + n] += 1.0 + + def result(self) -> np.ndarray: + if (self._cnt == 0).any(): + raise RuntimeError("some frames are covered by no chunk") + return (self._sum / self._cnt).astype(np.float32) + + +def smooth_actions(actions: np.ndarray, groups: dict[str, slice], hand_groups, window: int, + hand_window: int, order: int = 2) -> np.ndarray: + """SG filter per group along time; hand groups use ``hand_window``. A window longer than the + clip (or not above the order) leaves the group unchanged.""" + out = np.array(actions, dtype=np.float32, copy=True) + for name, sl in groups.items(): + w = int(hand_window if name in hand_groups else window) | 1 + if order < w <= len(out): + out[:, sl] = savgol_filter(out[:, sl], w, order, axis=0, mode="nearest") + return out + + +def state_from_action(actions: np.ndarray, delay: int) -> np.ndarray: + """q_t = a_{t-delay}; q_t = a_0 for t < delay.""" + a = np.asarray(actions) + d = min(int(delay), len(a)) + return np.concatenate([np.repeat(a[:1], d, axis=0), a[: len(a) - d]], axis=0) + + +def group_slices(spec: dict, widths: dict[str, int]) -> dict[str, slice]: + out, s = {}, 0 + for g in spec["action_groups"]: + out[g] = slice(s, s + widths[g]) + s += widths[g] + return out + + +@torch.no_grad() +def predict_actions(model, views: list[np.ndarray], spec: dict, stats: dict, frame_gap: int, + batch_size: int, generator: torch.Generator | None = None) -> np.ndarray: + """Chunk-averaged raw actions (T, native_dim) for one clip. + + ``views``: one (T, H, W, 3) uint8 RGB array per view in ``spec["views"]`` order, already at + the embodiment's ``video_size``. ``stats``: per-group action statistics of this embodiment. + """ + H, A = model.cfg["action_horizon"], model.cfg["action_dim"] + groups = list(spec["action_groups"]) + q01 = torch.tensor(np.concatenate([stats[g]["q01"] for g in groups]).tolist()) + q99 = torch.tensor(np.concatenate([stats[g]["q99"] for g in groups]).tolist()) + n_native = len(q01) + T = min(len(v) for v in views) + device = next(model.parameters()).device + acc = ChunkAccumulator(T, n_native, H) + + def sample(t): + obs = np.clip(np.array([t, t + frame_gap]), 0, T - 1) + frames = np.stack([v[obs] for v in views], 1).reshape(-1, *views[0].shape[1:]) + return model_inputs(preprocess_frames(frames), spec, H, A) + + for s in range(0, T, batch_size): + batch = collate([sample(t) for t in range(s, min(s + batch_size, T))]) + batch = {k: v.to(device) for k, v in batch.items()} + pred = model.get_action(batch, generator=generator).cpu()[..., :n_native] + raw = unnormalize(pred, q01, q99, spec.get("norm", "q99")).numpy() + for b, chunk in enumerate(raw): + acc.add(s + b, chunk) + return acc.result() diff --git a/robotryon/idm/model.py b/robotryon/idm/model.py new file mode 100644 index 0000000000000000000000000000000000000000..6ef78155ccaf75544c00308aea73d0f5ccfd77f9 --- /dev/null +++ b/robotryon/idm/model.py @@ -0,0 +1,327 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# Adapted from GR00T-Dreams (Apache-2.0): gr00t/model/idm.py, +# gr00t/model/action_head/flow_matching_action_head_idm.py, multimodal_projector.py. +# Changes: Depth Anything V2 vision tower instead of SigLIP, no hydra, plain safetensors I/O, +# image tokens placed from a per-sample image count instead of token-id arrays. +"""Inverse dynamics model h_psi (paper Sec. 3.4, App. B.2.6). + +Given two frames o_t and o_{t+K} (K = 20, 1 s at 20 fps; one or two camera views each), the IDM +predicts the action chunk a_{t:t+K} by flow matching: + + images --DINOv2 trunk of Depth Anything V2-L--> 16x16 patch tokens (CLS dropped) + --linear + 2x(2x2 pool MLP)--> 16 tokens per image, + view embedding + tokens --self-attention transformer--> visual context (left zero-padded to 112 tokens) + noisy chunk a^tau --per-embodiment action encoder--> DiT (AdaLN on tau, cross-attention to + the context) --per-embodiment decoder--> velocity v; loss ||v - (a - eps)||^2 + +One model is shared by all embodiments. Actions are zero-padded to ``action_dim`` (the widest +robot) and the padded dimensions are masked out of the input and the loss; the action encoder +and decoder have separate weights per embodiment id. At inference the chunk is integrated +from Gaussian noise with ``num_inference_timesteps`` Euler steps. +""" + +from __future__ import annotations + +import json +import re +from pathlib import Path + +import torch +import torch.nn.functional as F +from torch import nn +from torch.distributions import Beta + +from .dit import DiT, SelfAttentionTransformer + +DEFAULT_MODEL_CFG = dict( + vision_tower="depth-anything/Depth-Anything-V2-Large-hf", + vision_hidden_size=1024, hidden_size=1024, + action_dim=66, action_horizon=20, + max_num_embodiments=32, max_num_views=6, max_seq_len=1024, + max_vl_tokens=112, tokens_per_image=16, + frozen_vision_layers=[11], + num_inference_timesteps=16, num_timestep_buckets=1000, + noise_beta_alpha=1.5, noise_beta_beta=1.0, noise_s=0.999, + dit=dict(num_layers=8, num_attention_heads=16, attention_head_dim=64, dropout=0.2, + output_dim=1024, interleave_self_attention=True), + vl_self_attention=dict(num_layers=4, num_attention_heads=16, attention_head_dim=64, + dropout=0.2), +) + + +# ----------------------------------------------------------------------------- vision tower +class DepthAnythingTower(nn.Module): + """The DINOv2 trunk of Depth Anything V2 (the DPT neck and depth head are discarded). + Input: ImageNet-normalised 224x224 images. Output: (N, 256, D) post-layernorm patch tokens + of the last layer, CLS token removed.""" + + def __init__(self, name: str, pretrained: bool = True): + super().__init__() + from transformers import DepthAnythingConfig, DepthAnythingForDepthEstimation + + if pretrained: + model = DepthAnythingForDepthEstimation.from_pretrained(name) + else: + model = DepthAnythingForDepthEstimation._from_config(DepthAnythingConfig.from_pretrained(name)) + self.backbone = model.backbone + # Only read with bool_masked_pos, which the backbone never passes: no gradient (DDP). + self.backbone.embeddings.mask_token.requires_grad_(False) + + def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: + return self.backbone(pixel_values).feature_maps[-1][:, 1:] + + +class _DownSample(nn.Module): + """Concatenate 2x2 neighbouring tokens of a square grid: (N, h*w, C) -> (N, h*w/4, 4C).""" + + def forward(self, x): + n, L, c = x.shape + h = w = int(L ** 0.5) + x = x.reshape(n, h, w, c) + if w % 2 or h % 2: + x = F.pad(x, (0, 0, 0, h % 2, 0, w % 2)) + n, w, h, c = x.shape + x = x.view(n, w, h // 2, c * 2).permute(0, 2, 1, 3).contiguous() + return x.view(n, h // 2, w // 2, c * 4).reshape(n, -1, c * 4) + + +class MultimodalProjector(nn.Module): + """``mlp_doubledownsample``: 256 patch tokens -> 64 -> 16 tokens.""" + + def __init__(self, mm_hidden: int, hidden: int): + super().__init__() + self.layers = nn.Sequential( + _DownSample(), nn.LayerNorm(mm_hidden * 4), nn.Linear(mm_hidden * 4, mm_hidden), + nn.GELU(), nn.Linear(mm_hidden, mm_hidden), + _DownSample(), nn.LayerNorm(mm_hidden * 4), nn.Linear(mm_hidden * 4, hidden), + nn.GELU(), nn.Linear(hidden, hidden)) + + def forward(self, x): + return self.layers(x) + + +# ----------------------------------------------------------------------------- action head +class CategorySpecificLinear(nn.Module): + def __init__(self, num_categories, input_dim, hidden_dim): + super().__init__() + self.W = nn.Parameter(0.02 * torch.randn(num_categories, input_dim, hidden_dim)) + self.b = nn.Parameter(torch.zeros(num_categories, hidden_dim)) + + def forward(self, x, cat_ids): + return torch.bmm(x, self.W[cat_ids]) + self.b[cat_ids].unsqueeze(1) + + +class CategorySpecificMLP(nn.Module): + def __init__(self, num_categories, input_dim, hidden_dim, output_dim): + super().__init__() + self.layer1 = CategorySpecificLinear(num_categories, input_dim, hidden_dim) + self.layer2 = CategorySpecificLinear(num_categories, hidden_dim, output_dim) + + def forward(self, x, cat_ids): + return self.layer2(F.relu(self.layer1(x, cat_ids)), cat_ids) + + +def sinusoidal(timesteps: torch.Tensor, dim: int) -> torch.Tensor: + """(B, T) -> (B, T, dim) sin/cos encoding.""" + timesteps = timesteps.float() + half = dim // 2 + exponent = -torch.arange(half, dtype=torch.float, device=timesteps.device) * ( + torch.log(torch.tensor(10000.0)) / half) + freqs = timesteps.unsqueeze(-1) * exponent.exp() + return torch.cat([torch.sin(freqs), torch.cos(freqs)], dim=-1) + + +class MultiEmbodimentActionEncoder(nn.Module): + def __init__(self, action_dim, hidden_size, num_embodiments): + super().__init__() + self.hidden_size = hidden_size + self.W1 = CategorySpecificLinear(num_embodiments, action_dim, hidden_size) + self.W2 = CategorySpecificLinear(num_embodiments, 2 * hidden_size, hidden_size) + self.W3 = CategorySpecificLinear(num_embodiments, hidden_size, hidden_size) + + def forward(self, actions, timesteps, cat_ids): + B, T, _ = actions.shape + a_emb = self.W1(actions, cat_ids) + tau = sinusoidal(timesteps.unsqueeze(1).expand(-1, T), self.hidden_size).to(a_emb.dtype) + x = self.W2(torch.cat([a_emb, tau], dim=-1), cat_ids) + return self.W3(x * torch.sigmoid(x), cat_ids) + + +class IDM(nn.Module): + """Flow-matching inverse dynamics model. See the module docstring. + + Batch (see ``data.collate``): ``images`` (N, 3, 224, 224) for all samples concatenated in + order [t0 view0, t0 view1, t1 view0, ...]; ``view_ids`` (N,); ``n_images`` (B,); + ``embodiment_id`` (B,); ``actions_mask`` (B, H, A) bool, True on the robot's native + dimensions; ``actions`` (B, H, A) normalised targets (training only). + """ + + def __init__(self, cfg: dict | None = None, pretrained_vision: bool = True): + super().__init__() + c = {**DEFAULT_MODEL_CFG, **(cfg or {})} + self.cfg = c + D = c["hidden_size"] + self.vision = DepthAnythingTower(c["vision_tower"], pretrained=pretrained_vision) + self.vision_projector = nn.Linear(c["vision_hidden_size"], D) + self.mm_projector = MultimodalProjector(c["vision_hidden_size"], D) + self.view_embedding = nn.Embedding(c["max_num_views"], D) + self.position_embedding = nn.Embedding(c["max_seq_len"], D) + nn.init.normal_(self.view_embedding.weight, std=0.02) + nn.init.normal_(self.position_embedding.weight, std=0.02) + self.vl_self_attention = SelfAttentionTransformer(**c["vl_self_attention"]) + self.dit = DiT(**c["dit"]) + self.action_encoder = MultiEmbodimentActionEncoder(c["action_dim"], D, c["max_num_embodiments"]) + self.action_decoder = CategorySpecificMLP(c["max_num_embodiments"], D, D, c["action_dim"]) + self.beta_dist = Beta(c["noise_beta_alpha"], c["noise_beta_beta"]) + # One ViT layer is kept frozen, as in the trained checkpoint. + for i in c["frozen_vision_layers"]: + self.vision.backbone.encoder.layer[i].requires_grad_(False) + + # -- encoding ------------------------------------------------------------------------- + def encode(self, batch) -> torch.Tensor: + """Visual context (B, max_vl_tokens, D): each sample's image tokens right-aligned, + zeros in front (the zeros are attended to, exactly as in training).""" + x = self.mm_projector(self.vision_projector(self.vision(batch["images"]))) + x = x + self.view_embedding(batch["view_ids"]).unsqueeze(1) + n_tok = batch["n_images"] * self.cfg["tokens_per_image"] + L = self.cfg["max_vl_tokens"] + assert int(n_tok.max()) <= L, f"{int(n_tok.max())} image tokens > max_vl_tokens={L}" + pos = torch.arange(L, device=x.device) + mask = (pos[None] >= (L - n_tok)[:, None]).unsqueeze(-1).expand(-1, -1, x.shape[-1]) + vl = torch.zeros(len(n_tok), L, x.shape[-1], dtype=x.dtype, device=x.device) + return self.vl_self_attention(vl.masked_scatter(mask, x)) + + def _velocity(self, context, actions, t_disc, t_dit, emb_id): + sa = self.action_encoder(actions, t_disc, emb_id) + T = sa.shape[1] + sa = sa + self.position_embedding(torch.arange(T, device=sa.device)).unsqueeze(0) + out = self.dit(sa, context, t_dit) + return self.action_decoder(out, emb_id)[:, -T:] + + # -- training ------------------------------------------------------------------------- + def forward(self, batch) -> torch.Tensor: + """Flow-matching loss over the native (unpadded) action dimensions.""" + context = self.encode(batch) + actions, mask = batch["actions"], batch["actions_mask"].to(batch["actions"].dtype) + noise = torch.randn_like(actions) + s = self.cfg["noise_s"] + t = (s - self.beta_dist.sample([actions.shape[0]]).to(actions.device, actions.dtype)) / s + t = t[:, None, None] + noisy = ((1 - t) * noise + t * actions) * mask + t_disc = (t[:, 0, 0] * self.cfg["num_timestep_buckets"]).long() + pred = self._velocity(context, noisy, t_disc, t_disc, batch["embodiment_id"]) + loss = F.mse_loss(pred, actions - noise, reduction="none") * mask + return loss.sum() / mask.sum() + + # -- inference ------------------------------------------------------------------------ + @torch.inference_mode() + def get_action(self, batch, generator: torch.Generator | None = None) -> torch.Tensor: + """Integrate the flow from noise; returns normalised, padded actions (B, H, A).""" + emb_id = batch["embodiment_id"] + images = batch["images"] + B, H, A = len(emb_id), self.cfg["action_horizon"], self.cfg["action_dim"] + x = torch.randn((B, H, A), dtype=images.dtype, device=images.device, generator=generator) + mask = batch["actions_mask"].to(x.dtype) + x = x * mask + context = self.encode(batch) + n = self.cfg["num_inference_timesteps"] + for i in range(n): + t = int(i / float(n) * self.cfg["num_timestep_buckets"]) + v = self._velocity(context, x, (torch.ones(B) * t).to(x.device), + torch.tensor([t], device=x.device, dtype=torch.long), emb_id) + x = (x + (1.0 / n) * v) * mask + return x + + # -- I/O ------------------------------------------------------------------------------ + def save_pretrained(self, out_dir: str | Path, metadata: dict | None = None) -> None: + from safetensors.torch import save_file + + out = Path(out_dir) + out.mkdir(parents=True, exist_ok=True) + sd = {k: v.detach().contiguous().cpu() for k, v in self.state_dict().items()} + save_file(sd, str(out / "model.safetensors"), metadata={"format": "pt"}) + (out / "idm_config.json").write_text(json.dumps(self.cfg, indent=1)) + if metadata is not None: + (out / "metadata.json").write_text(json.dumps(metadata, indent=1)) + + @classmethod + def from_pretrained(cls, ckpt_dir: str | Path, device="cpu") -> "IDM": + """Load a checkpoint written by ``save_pretrained`` or by the GR00T-Dreams IDM trainer.""" + from safetensors.torch import load_file + + d = Path(ckpt_dir) + if (d / "idm_config.json").exists(): + cfg = json.loads((d / "idm_config.json").read_text()) + else: + cfg = _legacy_config(json.loads((d / "config.json").read_text())) + model = cls(cfg, pretrained_vision=False) + sd = load_file(str(d / "model.safetensors")) + if any(k.startswith("action_head.") for k in sd): + sd = _legacy_state_dict(sd) + model.load_state_dict(sd, strict=True) + return model.to(device).eval() + + +def load_metadata(ckpt_dir: str | Path) -> dict: + """Per-embodiment normalisation statistics saved with the checkpoint, keyed by name.""" + d = Path(ckpt_dir) + for p in (d / "metadata.json", d / "experiment_cfg" / "metadata.json"): + if p.exists(): + return json.loads(p.read_text()) + raise FileNotFoundError(f"no metadata.json in {d}") + + +def _legacy_config(c: dict) -> dict: + h = c["action_head_cfg"]["config"] + keep = ("num_layers", "num_attention_heads", "attention_head_dim", "dropout") + return dict( + vision_tower=h["siglip_model_cfg"]["pretrained_model_name_or_path"], + vision_hidden_size=h["siglip_hidden_size"], hidden_size=h["hidden_size"], + action_dim=h["action_dim"], action_horizon=h["action_horizon"], + max_num_embodiments=h.get("max_num_embodiments", 32), max_num_views=h["max_num_views"], + max_seq_len=h.get("max_seq_len", 1024), + num_inference_timesteps=h["num_inference_timesteps"], + num_timestep_buckets=h["num_timestep_buckets"], + noise_beta_alpha=h["noise_beta_alpha"], noise_beta_beta=h["noise_beta_beta"], + noise_s=h["noise_s"], + dit={**{k: h["diffusion_model_cfg"][k] for k in keep}, + "output_dim": h["diffusion_model_cfg"]["output_dim"], + "interleave_self_attention": h["diffusion_model_cfg"]["interleave_self_attention"]}, + vl_self_attention={k: h["vl_self_attention_cfg"][k] for k in keep}, + ) + + +_LEGACY_RENAMES = ( + (r"^action_head\.siglip_model\.vision_model\.backbone\.", "vision.backbone."), + (r"^action_head\.model\.", "dit."), + (r"^action_head\.vl_self_attention_model\.", "vl_self_attention."), + (r"^action_head\.", ""), +) + + +def _legacy_state_dict(sd: dict) -> dict: + out = {} + for k, v in sd.items(): + if k.endswith(("siglip_model.logit_scale", "siglip_model.logit_bias")): + continue + for pat, rep in _LEGACY_RENAMES: + k2 = re.sub(pat, rep, k) + if k2 != k: + k = k2 + break + out[k] = v + return out diff --git a/robotryon/idm/wrist_correction.py b/robotryon/idm/wrist_correction.py new file mode 100644 index 0000000000000000000000000000000000000000..73d9f4f61411bba5773d832f082445a356655070 --- /dev/null +++ b/robotryon/idm/wrist_correction.py @@ -0,0 +1,229 @@ +"""Wrist correction of IDM action labels (paper Sec. 3.4, App. B.2.6). + +The IDM reads joint angles from pixels and places the wrist less accurately than the 3D hand +reconstruction of the human source video. For each hand we therefore move the wrist, and only +the wrist, toward the reconstruction: + +1. The reconstructed human wrist (camera frame) is mapped to the robot base frame with the + head-camera extrinsics T_base_cam. +2. The robot hand implied by the IDM joints (FK) is scaled by ``hand_scale`` + about its wrist; the target wrist is the one at which the scaled hand's end-effector keypoint + coincides with the reconstructed wrist. +3. Targets are kept only on high-confidence frames. The offset d = target - IDM wrist is + cleaned: excursion rejection (hysteresis around a running median of accepted offsets), + linear fill of short gaps, zero-phase 2nd-order Butterworth low-pass per usable segment. +4. Damped least squares IK on the arm joints of that side: wrist position to the IDM wrist + d, + wrist orientation held at the FK orientation of the IDM arm joints (weighted rotation + residual), damping toward the IDM arm joints, numerical Jacobian of FK, warm-started from + the previous frame of the segment. +5. The solution is blended into the IDM joints with weights ramping 0 -> 1 over ``ramp`` frames + at segment edges (not at the clip ends), then the arm tracks get SG smoothing. Hand joints + and all other joints are left untouched. + +Kinematics interface required (minimal protocol; ``robotryon.embodiment.Embodiment`` is +expected to satisfy it):: + + class Kinematics(Protocol): + arm_joint_slices: dict[str, slice] # {"left": slice, "right": slice} into the action + def fk(self, q: np.ndarray) -> dict[str, dict[str, np.ndarray]]: + # {"left": {"wrist_pos": (3,), "wrist_rot": (3, 3), "keypoints": (17, 3)}, "right": ...} + # robot base frame, metres; keypoints use robotryon.skeleton order (0 = wrist) + +Hands are indexed 0 = left, 1 = right, as in ``robotryon.skeleton``. +""" + +from __future__ import annotations + +from typing import Callable + +import numpy as np +from scipy.signal import butter, filtfilt, savgol_filter + +SIDES = ("left", "right") + +DEFAULTS = dict( + hand_scale=1.0, ee_keypoint=0, min_confidence=0.5, + reject_enter_mm=150.0, reject_exit_mm=80.0, reject_hard_mm=500.0, reject_ref_n=15, + gap_fill=6, cutoff_hz=2.0, fps=20.0, ramp=8, smooth_window=9, smooth_order=2, + ik_iters=30, w_rot=2000.0, damping=0.5, step_cap=0.08, fd_eps=1e-3, +) + + +# ----------------------------------------------------------------------------- track cleaning +def segments(mask: np.ndarray) -> list[tuple[int, int]]: + """Half-open [a, b) runs of True.""" + out, a = [], None + for i, m in enumerate(mask): + if m and a is None: + a = i + if not m and a is not None: + out.append((a, i)) + a = None + if a is not None: + out.append((a, len(mask))) + return out + + +def reject_excursions(d, ok, enter, exit_, hard, ref_n=15): + """Refuse frames whose offset leaves the level it has been holding (running median of the + last ``ref_n`` accepted offsets); a refusal lasts until the offset is back within ``exit_``. + ``hard`` caps |d| before a level is established.""" + good, buf, ref = ok.copy(), [], None + for t in range(len(d)): + if not ok[t]: + continue + if np.linalg.norm(d[t]) > hard: + good[t] = False + continue + if ref is None: + buf.append(d[t]) + if len(buf) >= 5: + ref = np.median(np.array(buf[-ref_n:]), axis=0) + continue + dev = float(np.linalg.norm(d[t] - ref)) + if dev > enter or (dev > exit_ and not good[max(t - 1, 0)]): + good[t] = False + else: + buf.append(d[t]) + ref = np.median(np.array(buf[-ref_n:]), axis=0) + return good + + +def fill_short_gaps(track, ok, max_gap): + """Linearly interpolate refused runs of <= max_gap frames between usable frames (in place); + returns the new usable mask.""" + ok = ok.copy() + for a, b in segments(~ok): + if a == 0 or b == len(ok) or b - a > max_gap: + continue + for k in range(track.shape[1]): + track[a:b, k] = np.interp(np.arange(a, b), [a - 1, b], [track[a - 1, k], track[b, k]]) + ok[a:b] = True + return ok + + +def lowpass(track, ok, cutoff_hz, fps): + """Zero-phase Butterworth per usable segment (segments too short for filtfilt are kept).""" + b, a = butter(2, cutoff_hz / (fps / 2.0), "low") + out = track.copy() + for lo, hi in segments(ok): + if hi - lo > 3 * max(len(a), len(b)) + 2: + out[lo:hi] = np.stack([filtfilt(b, a, track[lo:hi, c]) for c in range(track.shape[1])], 1) + return out + + +def ramp_weights(ok, ramp): + """1 inside usable segments, ramping linearly at segment edges that are not clip ends.""" + T, w = len(ok), np.zeros(len(ok)) + for a, b in segments(ok): + for t in range(a, b): + up = 1.0 if a == 0 else (t - a + 1) / ramp + dn = 1.0 if b == T else (b - t) / ramp + w[t] = min(1.0, up, dn) + return w + + +# ----------------------------------------------------------------------------- IK +def rotvec(R_cur, R_ref): + """Axis-angle of R_cur R_ref^T.""" + from scipy.spatial.transform import Rotation + + return Rotation.from_matrix(R_cur @ R_ref.T).as_rotvec() + + +def ik_wrist(fk: Callable, q, side, sl, target_mm, R_ref, q_ref, iters=30, w_rot=2000.0, + damping=0.5, step_cap=0.08, fd_eps=1e-3): + """Damped least squares on the joints ``sl``: wrist position (mm) to ``target_mm``, + orientation pinned to ``R_ref`` (``w_rot`` mm per rad), damping toward ``q_ref``. + Returns (q, position residual in mm before the last step).""" + s = np.array(q, dtype=np.float64) + n = sl.stop - sl.start + + def res(x): + h = fk(x)[side] + return np.concatenate([np.asarray(h["wrist_pos"]) * 1000.0 - target_mm, + w_rot * rotvec(np.asarray(h["wrist_rot"]), R_ref)]) + + lam, last = damping * 1e4, np.inf + for _ in range(iters): + r = res(s) + if not np.all(np.isfinite(r)): + return np.array(q, dtype=np.float64), np.inf + last = float(np.linalg.norm(r[:3])) + J = np.zeros((6, n)) + for j in range(n): + s2 = s.copy() + s2[sl.start + j] += fd_eps + J[:, j] = (res(s2) - r) / fd_eps + if not np.all(np.isfinite(J)): + return np.array(q, dtype=np.float64), np.inf + dq = np.linalg.solve(J.T @ J + (1.0 + lam) * np.eye(n), -J.T @ r - lam * (s[sl] - q_ref[sl])) + norm = np.linalg.norm(dq) + if norm > step_cap: + dq *= step_cap / norm + s[sl] += dq + if norm < 1e-4: + break + return s, last + + +# ----------------------------------------------------------------------------- main entry +def camera_to_base(points_cam: np.ndarray, T_base_cam: np.ndarray) -> np.ndarray: + """(T, 2, 3) camera-frame points -> base frame; T_base_cam is (4, 4) or (T, 4, 4).""" + Tm = np.broadcast_to(np.asarray(T_base_cam, np.float64), (len(points_cam), 4, 4)) + return np.einsum("tij,thj->thi", Tm[:, :3, :3], points_cam) + Tm[:, None, :3, 3] + + +def correct_wrists(actions: np.ndarray, kin, wrist_base: np.ndarray, confidence: np.ndarray, + cfg: dict | None = None) -> tuple[np.ndarray, dict]: + """Correct the arm joints of ``actions`` (T, D) so the wrists follow ``wrist_base``. + + wrist_base: (T, 2, 3) reconstructed human wrists in the robot base frame (metres, NaN when + missing); confidence: (T, 2). Returns the corrected actions and per-side statistics. + """ + c = {**DEFAULTS, **(cfg or {})} + act = np.asarray(actions, np.float64) + T = len(act) + out = act.copy() + fk_idm = [kin.fk(act[t]) for t in range(T)] + info = {} + for hi, side in enumerate(SIDES): + sl = kin.arm_joint_slices[side] + sl = sl if isinstance(sl, slice) else slice(*sl) + p_idm = np.stack([np.asarray(f[side]["wrist_pos"]) for f in fk_idm]) * 1000.0 + R_idm = [np.asarray(f[side]["wrist_rot"]) for f in fk_idm] + # end effector of the scaled IDM hand, relative to its wrist + ee = np.stack([np.asarray(f[side]["keypoints"])[c["ee_keypoint"]] for f in fk_idm]) * 1000.0 + target = wrist_base[:, hi] * 1000.0 - c["hand_scale"] * (ee - p_idm) + d = target - p_idm + ok = (confidence[:, hi] >= c["min_confidence"]) & np.isfinite(d).all(1) + n_conf = int(ok.sum()) + d = np.where(ok[:, None], d, 0.0) + ok = reject_excursions(d, ok, c["reject_enter_mm"], c["reject_exit_mm"], + c["reject_hard_mm"], c["reject_ref_n"]) + n_rej = n_conf - int(ok.sum()) + ok = fill_short_gaps(d, ok, c["gap_fill"]) + d = lowpass(d, ok, c["cutoff_hz"], c["fps"]) + q_ik, prev, resid = act[:, sl].copy(), None, [] + for t in range(T): + if not ok[t]: + prev = None + continue + cur = act[t].copy() + if prev is not None: + cur[sl] = prev + new, r = ik_wrist(kin.fk, cur, side, sl, p_idm[t] + d[t], R_idm[t], act[t], + c["ik_iters"], c["w_rot"], c["damping"], c["step_cap"], c["fd_eps"]) + if not np.isfinite(r): + ok[t], prev = False, None + continue + q_ik[t] = prev = new[sl] + resid.append(r) + w = ramp_weights(ok, c["ramp"]) + blended = act[:, sl] + w[:, None] * (q_ik - act[:, sl]) + if T > c["smooth_window"]: + blended = savgol_filter(blended, c["smooth_window"], c["smooth_order"], axis=0, mode="nearest") + out[:, sl] = blended + info[side] = dict(confident_frames=n_conf, rejected=n_rej, corrected_frames=int((w > 0).sum()), + ik_residual_mm_median=float(np.median(resid)) if resid else None) + return out.astype(np.asarray(actions).dtype), info diff --git a/robotryon/image/__init__.py b/robotryon/image/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4c095dd2dba9335f113ee782b76a6f68315af6b9 --- /dev/null +++ b/robotryon/image/__init__.py @@ -0,0 +1,14 @@ +"""Keyframe translation f_theta (paper Sec. 3.2): FLUX.2 [klein] 9B Base as a reference-guided +inpainting model. + +Modules: + fill_channels widened patch embedding, channel-concatenated [noisy | conditioned frame | mask] + stereo_ids RoPE ids with the camera view on the 4th axis + episodes data layout access and model-resolution geometry + dataset training samples (reference window, robot-only exemplar, masks, skeleton) + text prompts, caption dropout and the text-embedding cache + training flow-matching loss + model FFT / LoRA build, save and load + sampler true-CFG Euler sampler returning latent tokens + ar_keyframes autoregressive translation with latent reuse +""" diff --git a/robotryon/image/ar_keyframes.py b/robotryon/image/ar_keyframes.py new file mode 100644 index 0000000000000000000000000000000000000000..cfc66f23d94044bfe2572d14a8e8045462fd7a79 --- /dev/null +++ b/robotryon/image/ar_keyframes.py @@ -0,0 +1,161 @@ +"""Autoregressive keyframe translation (paper Eq. 3.2, App. B.2.3). + +Human keyframes are translated in temporal order. Keyframe t_j is conditioned on + +* the conditioned human frame Phi(o^h; M^h, S^h): the human arm mask (convex hull per arm, + fixed dilation, eroded object removed) blacked out and the retargeted robot-hand skeleton drawn + on top, plus the mask itself as an input channel; +* the reference: for the first keyframe, the encoded robot-only reference image supplied by the + user (it selects the embodiment); for every later keyframe, the DENOISED LATENT TOKENS of the + previous keyframe, reused directly as reference tokens (no decode / re-encode). + +``ref_skeleton: true`` replaces latent reuse by a pixel-space reference: the previous keyframe is +decoded, the current keyframe's skeleton is painted on it, and the result is re-encoded (the first +keyframe's reference image is painted the same way). + +Stereo episodes are translated jointly (both views in one token sequence) by default; ``views: +mono`` translates each view as its own chain. +""" +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +import torch +from PIL import Image + +from ..masks import MaskAug, inference_mask +from ..skeleton import condition_frame, draw_skeleton, paint +from .episodes import Episode, FrameReader, load_view_frame +from .fill_channels import pack +from .sampler import LatentCodec, ViewCondition, denoise +from ..captions import HandStateCuts, hand_state_arrays, hand_state_prompt, load_hand_state, words_at + + +@dataclass +class TranslatedKeyframe: + image: np.ndarray # uint8 (H, W, 3) robot keyframe + cond: np.ndarray # uint8 (H, W, 3) conditioned human frame + mask: np.ndarray # uint8 (H, W) {0, 1} inference mask + tokens: torch.Tensor # (1, h*w, 128) denoised latent tokens (next keyframe's reference) + + +class KeyframeTranslator: + def __init__(self, pipe, cfg, device="cuda"): + self.pipe, self.cfg, self.device = pipe, cfg, device + self.codec = LatentCodec(pipe.vae, device) + self.reader = FrameReader() + self.dilate = MaskAug.from_config(cfg.masks).infer_dilate() + self.ref_skeleton = bool(cfg.get("ref_skeleton", False)) + self._prompts: dict[str, tuple] = {} + + def prompt_fn(self, ep: Episode, episode_dir: Path): + """``prompt(t)``: the instruction and each hand's closure / facing words at keyframe t + (``robotryon.captions``), read from the episode's ``hand_state_.npz`` (or computed + from ``hands.npz`` when absent).""" + cuts = HandStateCuts.from_config(self.cfg.caption.get("cuts")) + states = [] + for i, v in enumerate(ep.views): + f = episode_dir / f"hand_state_{v}.npz" + if f.exists(): + states.append(load_hand_state(f)) + if not states: + z = np.load(episode_dir / "hands.npz") + states = [hand_state_arrays(z["joints3d"][i], z["valid"][i].astype(bool)) + for i in range(z["joints3d"].shape[0])] + return lambda t: hand_state_prompt(*words_at(states, int(t), cuts)) + + def encode_prompt(self, text: str): + if text not in self._prompts: + self._prompts[text] = self.pipe.encode_prompt(prompt=text, device=self.device) + return self._prompts[text] + + def condition(self, ep: Episode, view: str, t: int): + """(frame, conditioned frame, mask, skeleton image) of one human keyframe at model resolution.""" + c = self.cfg + f = load_view_frame(self.reader, ep, view, t, int(c.image.long_edge), + tuple(c.image.aspect) if c.image.get("aspect") else None, + c.image.get("crop_anchor", "bottom")) + H, W = f.image.shape[:2] + mask = inference_mask(f.hulls, W, H, self.dilate, f.object_mask, int(c.masks.obj_erode)) + skel = draw_skeleton(f.kp, f.vis, W, H) + cond = condition_frame(f.image, mask, skel) + return f.image, cond, mask, skel + + def reference(self, view: str, prev: dict, references: dict, skel: np.ndarray) -> torch.Tensor: + """Reference tokens of one view for the current keyframe.""" + if view in prev and not self.ref_skeleton: + return prev[view].tokens # latent reuse + if view in prev: + img = prev[view].image + else: + H, W = skel.shape[:2] + img = np.asarray(Image.fromarray(references[view]).resize((W, H), Image.LANCZOS)) + return pack(self.codec.encode(paint(img, skel) if self.ref_skeleton else img)) + + @torch.no_grad() + def translate_chain(self, views: list[str], keyframes: list[int], references: dict[str, np.ndarray], + prompt, conditions) -> dict: + """One chain over ``keyframes``; ``views`` are denoised jointly. + + ``prompt`` is a string or ``prompt(t)`` giving each keyframe its own prompt. + ``conditions(view, t)`` returns ``(frame, conditioned frame, mask, skeleton image)``. + Returns ``{view: {t: TranslatedKeyframe}}``. + """ + c = self.cfg + prompt_at = prompt if callable(prompt) else (lambda t: prompt) + neg, neg_ids = self.encode_prompt("") if float(c.guidance_scale) > 1.0 else (None, None) + prev: dict[str, TranslatedKeyframe] = {} + out: dict[str, dict] = {v: {} for v in views} + for t in keyframes: + conds, frames = [], [] + for v in views: + _, cond, mask, skel = conditions(v, t) + mask_t = torch.from_numpy(mask.astype(np.float32))[None, None].to(self.device) + conds.append(ViewCondition(self.codec.encode(cond), mask_t, + self.reference(v, prev, references, skel))) + frames.append((cond, mask)) + pos, pos_ids = self.encode_prompt(prompt_at(t)) + gen = torch.Generator(device=self.device).manual_seed(int(c.seed) + int(t)) + tokens = denoise(self.pipe.transformer, self.pipe.scheduler, conds, pos, pos_ids, neg, neg_ids, + steps=int(c.steps), guidance_scale=float(c.guidance_scale), generator=gen) + for v, tok, (cond, mask) in zip(views, tokens, frames): + H, W = mask.shape + img = self.codec.decode_tokens(tok, H // 16, W // 16) + out[v][t] = prev[v] = TranslatedKeyframe(img, cond, mask, tok) + return out + + def translate_episode(self, episode_dir: str | Path, references: dict[str, np.ndarray], + out_dir: str | Path) -> Path: + """Translate every keyframe of a preprocessed human episode and write the results.""" + c = self.cfg + ep = Episode.load(episode_dir) + kf = json.loads((Path(episode_dir) / "keyframes.json").read_text())["indices"] + keyframes = sorted({int(i) for i in kf} | {0}) + stereo = len(ep.views) >= 2 and c.views in ("auto", "stereo") + if c.views == "stereo" and not stereo: + raise ValueError(f"{episode_dir}: views=stereo needs two camera views, found {ep.views}") + chains = [ep.views[:2]] if stereo else [[v] for v in ep.views] + missing = [v for v in ep.views if v not in references] + if missing: + raise ValueError(f"no reference image for view(s) {missing}") + prompt_at = self.prompt_fn(ep, Path(episode_dir)) + prompts = {int(t): prompt_at(t) for t in keyframes} + out = Path(out_dir) / ep.path.name + out.mkdir(parents=True, exist_ok=True) + for views in chains: + res = self.translate_chain(views, keyframes, references, prompt_at, + lambda v, t: self.condition(ep, v, t)) + for v, per_t in res.items(): + for t, k in per_t.items(): + Image.fromarray(k.image).save(out / f"keyframe_{v}_{t:06d}.png") + Image.fromarray(k.cond).save(out / f"cond_{v}_{t:06d}.png") + Image.fromarray(k.mask * 255).save(out / f"mask_{v}_{t:06d}.png") + (out / "keyframes.json").write_text(json.dumps({ + "indices": keyframes, "views": ep.views, "joint_stereo": stereo, "prompts": prompts, + "ref_skeleton": self.ref_skeleton, "steps": int(c.steps), + "guidance_scale": float(c.guidance_scale), + "seed": int(c.seed)}, indent=1)) + return out diff --git a/robotryon/image/dataset.py b/robotryon/image/dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..b37e7359ce092895dfc8f2aa5c9487638ac3dc95 --- /dev/null +++ b/robotryon/image/dataset.py @@ -0,0 +1,227 @@ +"""Training samples for the keyframe-translation model (paper Sec. 3.2, App. B.2.3). + +A sample is a robot frame o_t of one episode, seen by one camera (mono) or both (stereo): + +* target o_t +* cond Phi(o_t; M_t, S_t): the randomised shape-agnostic mask blacked out, the augmented + hand skeleton (rendered on the fly from keypoints) drawn on top +* mask M_t +* reference o_t' with t' drawn uniformly from [t - 5 s, t - 2 s] of the same episode (frame 0 if + the target is too early). With probability ``p_robot_only`` the reference is instead a + "robot-only" exemplar: a frame of ANOTHER episode of the same embodiment with + everything but the robot (its tight mask, eroded) filled with white. This is + the reference the first keyframe gets at inference. +* prompt the hand-state prompt of the target frame (each hand's closure / facing words from the + episode's ``hand_state_.npz``; see ``robotryon.captions``), as a cached + embedding. + +With probability ``ref_skeleton_p`` the target's (augmented) skeleton is also painted on the +reference, sequential or robot-only alike; otherwise the reference is left clean. + +The two views of a stereo sample share every random draw (mask dilations, skeleton corruption, +reference frame, field-of-view crop): they are two views of one instant. +""" +from __future__ import annotations + +import numpy as np +import torch +from PIL import Image + +from ..masks import MaskAug, fov_augment, robot_only +from ..skeleton import SkeletonAugment, condition_frame, draw_skeleton, paint +from .episodes import Episode, FrameReader, crop_resize, load_view_frame, output_geometry, scan_root +from ..captions import HandStateCuts, all_prompts, load_hand_state, sample_prompt, words_at +from .text import TextEmbedCache + + +def _to_tensor(img: Image.Image) -> torch.Tensor: + a = np.asarray(img, dtype=np.float32) + if a.ndim == 2: + return torch.from_numpy(a / 255.0)[None] + return torch.from_numpy(a / 127.5 - 1.0).permute(2, 0, 1).contiguous() + + +class KeyframeDataset(torch.utils.data.Dataset): + """Indexed by ``(episode_index, frame, view_indices, seed)`` tuples from :class:`FillBatchSampler`.""" + + def __init__(self, cfg, text_cache: TextEmbedCache | None): + d = cfg.data + roots = d.roots if isinstance(d.roots, (list, tuple)) else [d.roots] + self.root_weights = [] + self.episodes: list[Episode] = [] + for i, r in enumerate(roots): + path, w = (r.path, float(r.get("weight", 1.0))) if isinstance(r, dict) else (r, 1.0) + eps = scan_root(path, i) + if not eps: + raise FileNotFoundError(f"no episodes (*/meta.json) under {path}") + self.episodes += eps + self.root_weights.append(w) + self.long_edge = int(cfg.image.long_edge) + self.aspect = tuple(cfg.image.aspect) if cfg.image.get("aspect") else None + self.anchor = cfg.image.get("crop_anchor", "bottom") + self.ref_window = tuple(d.ref_window_sec) + self.p_robot_only = float(d.p_robot_only) + self.ref_skeleton_p = float(d.ref_skeleton_p) + self.robot_only_erode = int(d.robot_only.erode) + self.robot_only_bg = int(d.robot_only.background) + self.robot_only_min_area = float(d.robot_only.min_area) + self.bonus_tags = set(cfg.masks.get("bonus_embodiments", []) or []) + self.mask_aug = MaskAug.from_config(cfg.masks) + self.fov = cfg.masks.fov + self.skel_aug = SkeletonAugment.from_config(cfg.skeleton.augment) + self.thickness = tuple(cfg.skeleton.thickness) + self.cap = cfg.caption + self.text_cache = text_cache + self.reader: FrameReader | None = None + self.by_embodiment: dict[str, list[int]] = {} + for i, ep in enumerate(self.episodes): + self.by_embodiment.setdefault(ep.embodiment, []).append(i) + + def __len__(self) -> int: + return len(self.episodes) + + def hand_states(self, ep: Episode, views: list[str]) -> list: + return [load_hand_state(ep.path / f"hand_state_{v}.npz") for v in views + if (ep.path / f"hand_state_{v}.npz").exists()] + + def all_prompts(self) -> set[str]: + return set(all_prompts()) + + # -- reference ------------------------------------------------------------------------------ + def reference_index(self, ep: Episode, t: int, rng: np.random.Generator) -> int: + lo, hi = self.ref_window + newest = t - int(round(lo * ep.fps)) + oldest = max(0, t - int(round(hi * ep.fps))) + if newest < oldest: + return 0 + return int(rng.integers(oldest, newest + 1)) + + def draw_exemplar(self, ep_idx: int, views: list[str], rng) -> tuple[int, int] | None: + """(episode, frame) of another episode of the same embodiment with a visible robot.""" + ep = self.episodes[ep_idx] + pool = [i for i in self.by_embodiment.get(ep.embodiment, []) + if i != ep_idx and all(v in self.episodes[i].views for v in views)] + if not pool: + return None + ex = int(pool[int(rng.integers(0, len(pool)))]) + ex_ep = self.episodes[ex] + for _ in range(8): + t = int(rng.integers(0, ex_ep.num_frames)) + m = self.reader.binary(ex_ep.mask(views[0]), t) + if m.mean() >= self.robot_only_min_area: + return ex, t + return None + + def robot_only(self, ex: int, t: int, view: str, size) -> np.ndarray: + ep = self.episodes[ex] + rgb = self.reader.rgb(ep.video(view), t) + src = (rgb.shape[1], rgb.shape[0]) + box, _ = output_geometry(src[0], src[1], self.long_edge, self.aspect, self.anchor) + img = crop_resize(rgb, box, size, src) + m = crop_resize(self.reader.binary(ep.mask(view), t).astype(np.uint8) * 255, box, size, src) + return robot_only(img, m > 127, self.robot_only_erode, self.robot_only_bg) + + # -- one view ------------------------------------------------------------------------------- + def view_streams(self, ep: Episode, view: str, t: int, ref_t: int, exemplar, ref_skeleton: bool, + rng) -> dict: + f = load_view_frame(self.reader, ep, view, t, self.long_edge, self.aspect, self.anchor) + H, W = f.image.shape[:2] + mask, _ = self.mask_aug.train_mask(f.hulls, W, H, rng, bonus=ep.embodiment in self.bonus_tags, + object_mask=f.object_mask) + kp, vis = self.skel_aug(f.kp, f.vis, rng) + th = int(rng.integers(self.thickness[0], self.thickness[1] + 1)) + skel = draw_skeleton(kp, vis, W, H, thickness=th, radius=max(1, int(round(th * 1.5)))) + cond = condition_frame(f.image, mask, skel) + if exemplar is not None: + ref = self.robot_only(exemplar[0], exemplar[1], view, (W, H)) + else: + rgb = self.reader.rgb(ep.video(view), ref_t) + src = (rgb.shape[1], rgb.shape[0]) + box, _ = output_geometry(src[0], src[1], self.long_edge, self.aspect, self.anchor) + ref = crop_resize(rgb, box, (W, H), src) + if ref_skeleton: + ref = paint(ref, skel) + return {"target": f.image, "cond": cond, "mask": mask * 255, "reference": ref, + "has_hull": bool(f.hulls)} + + def __getitem__(self, item) -> dict: + ep_idx, t, view_idx, seed = item + if self.reader is None: + self.reader = FrameReader() + ep = self.episodes[ep_idx] + views = [ep.views[i] for i in view_idx] + rng = np.random.default_rng([int(seed), 0]) + exemplar = None + if float(rng.random()) < self.p_robot_only: + exemplar = self.draw_exemplar(ep_idx, views, rng) + ref_skeleton = float(rng.random()) < self.ref_skeleton_p + for attempt in range(8): # prefer frames where the arm is visible + ref_t = self.reference_index(ep, t, rng) + out = [self.view_streams(ep, v, t, ref_t, exemplar, ref_skeleton, + np.random.default_rng([int(seed), 1, attempt])) for v in views] + if out[0]["has_hull"]: + break + t = int(rng.integers(0, ep.num_frames)) + keys = ("target", "cond", "mask", "reference") + imgs = [Image.fromarray(np.asarray(o[k], np.uint8)) for o in out for k in keys] + imgs, _ = fov_augment(imgs, rng, float(self.fov.crop_min), float(self.fov.scale_max), + float(self.fov.get("flip_prob", 0.0))) + n = len(keys) + sample = {k: torch.stack([_to_tensor(imgs[v * n + j]) for v in range(len(views))]) + for j, k in enumerate(keys)} + words = words_at(self.hand_states(ep, views), t, HandStateCuts.from_config(self.cap.get("cuts"))) + prompt = sample_prompt(words, rng, float(self.cap.p_uncond), float(self.cap.p_drop_hands)) + sample["prompt"] = prompt + if self.text_cache is not None: + sample["prompt_embeds"] = self.text_cache.load(prompt) + sample["robot_only"] = exemplar is not None + sample["ref_skeleton"] = ref_skeleton + return sample + + +def collate(batch: list[dict]) -> dict: + out = {k: torch.stack([b[k] for b in batch]) for k in ("target", "cond", "mask", "reference")} + if "prompt_embeds" in batch[0]: + out["prompt_embeds"] = torch.stack([b["prompt_embeds"] for b in batch]) + out["prompt"] = [b["prompt"] for b in batch] + out["robot_only"] = [b["robot_only"] for b in batch] + return out + + +class FillBatchSampler(torch.utils.data.Sampler): + """Endless stream of homogeneous batches for one rank. + + A batch is all-stereo or all-mono (the token sequence length depends on it): stereo with + probability ``stereo_p`` (paper: stereo : mono = 3 : 1) when stereo episodes exist. Items pick + a data root by weight, an episode uniformly within it, a frame uniformly, and for mono a view + uniformly. Each rank draws its own stream from ``(seed, rank)``. + """ + + def __init__(self, dataset: KeyframeDataset, batch_size: int, stereo_p: float, seed: int, + rank: int = 0): + self.ds, self.bs, self.stereo_p = dataset, int(batch_size), float(stereo_p) + self.rng = np.random.default_rng([int(seed), int(rank)]) + eps = dataset.episodes + n_roots = len(dataset.root_weights) + self.pools = {s: [[i for i, e in enumerate(eps) if e.root == r and (len(e.views) >= 2 or not s)] + for r in range(n_roots)] for s in (True, False)} + self.has_stereo = any(self.pools[True]) + + def _root_probs(self, stereo: bool) -> np.ndarray: + w = np.array([wt if self.pools[stereo][r] else 0.0 + for r, wt in enumerate(self.ds.root_weights)], np.float64) + return w / w.sum() + + def __iter__(self): + while True: + stereo = self.has_stereo and float(self.rng.random()) < self.stereo_p + p = self._root_probs(stereo) + batch = [] + for r in self.rng.choice(len(p), size=self.bs, p=p): + pool = self.pools[stereo][int(r)] + e = pool[int(self.rng.integers(0, len(pool)))] + ep = self.ds.episodes[e] + t = int(self.rng.integers(0, ep.num_frames)) + views = (0, 1) if stereo else (int(self.rng.integers(0, len(ep.views))),) + batch.append((e, t, views, int(self.rng.integers(0, 2 ** 31)))) + yield batch diff --git a/robotryon/image/episodes.py b/robotryon/image/episodes.py new file mode 100644 index 0000000000000000000000000000000000000000..67387db57bfa2097845699f5d0ebd2959f85a9e0 --- /dev/null +++ b/robotryon/image/episodes.py @@ -0,0 +1,173 @@ +"""Episode directories in the shared RoboTryOn data layout, and random frame access. + +:: + + // + meta.json {"fps", "views", "num_frames", "caption", "embodiment"} + video_.mp4 RGB + keypoints_.npz kp (T, 2, 17, 2) px at video resolution, vis (T, 2, 17) + mask_.mp4 tight arm + hand mask (gray, > 127 = masked) + object_mask_.mp4 optional manipulated-object mask +""" +from __future__ import annotations + +import json +from collections import OrderedDict +from dataclasses import dataclass, field +from pathlib import Path + +import cv2 +import numpy as np +from PIL import Image + + +@dataclass +class Episode: + path: Path + views: list[str] + num_frames: int + fps: float + caption: str = "" + embodiment: str = "" + root: int = 0 + extra: dict = field(default_factory=dict) + + @classmethod + def load(cls, path: str | Path, root: int = 0) -> "Episode": + path = Path(path) + m = json.loads((path / "meta.json").read_text()) + return cls(path=path, views=list(m["views"]), num_frames=int(m["num_frames"]), + fps=float(m["fps"]), caption=str(m.get("caption", "")), + embodiment=str(m.get("embodiment", "")), root=root, extra=m) + + def video(self, view: str) -> Path: + return self.path / f"video_{view}.mp4" + + def mask(self, view: str) -> Path: + return self.path / f"mask_{view}.mp4" + + def object_mask(self, view: str) -> Path | None: + p = self.path / f"object_mask_{view}.mp4" + return p if p.exists() else None + + def keypoints(self, view: str) -> Path: + return self.path / f"keypoints_{view}.npz" + + +def scan_root(root: str | Path, root_index: int = 0) -> list[Episode]: + return [Episode.load(p.parent, root_index) for p in sorted(Path(root).glob("*/meta.json"))] + + +class FrameReader: + """Small LRU of open ``cv2.VideoCapture`` objects and keypoint files (one per worker).""" + + def __init__(self, max_open: int = 16): + self._caps: OrderedDict[str, cv2.VideoCapture] = OrderedDict() + self._kps: OrderedDict[str, tuple[np.ndarray, np.ndarray]] = OrderedDict() + self.max_open = max_open + + def _cap(self, path: Path) -> cv2.VideoCapture: + key = str(path) + if key in self._caps: + self._caps.move_to_end(key) + return self._caps[key] + cap = cv2.VideoCapture(key) + if not cap.isOpened(): + raise FileNotFoundError(f"cannot open video {key}") + self._caps[key] = cap + while len(self._caps) > self.max_open: + self._caps.popitem(last=False)[1].release() + return cap + + def rgb(self, path: Path, t: int) -> np.ndarray: + cap = self._cap(path) + cap.set(cv2.CAP_PROP_POS_FRAMES, int(t)) + ok, bgr = cap.read() + if not ok: + raise RuntimeError(f"cannot decode frame {t} of {path}") + return cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) + + def binary(self, path: Path, t: int) -> np.ndarray: + """Gray mask frame -> bool ``(H, W)`` (> 127).""" + return self.rgb(path, t).max(axis=-1) > 127 + + def keypoints(self, path: Path) -> tuple[np.ndarray, np.ndarray]: + key = str(path) + if key not in self._kps: + z = np.load(key) + self._kps[key] = (z["kp"].astype(np.float32), z["vis"].astype(bool)) + while len(self._kps) > 4 * self.max_open: + self._kps.popitem(last=False) + return self._kps[key] + + +def output_geometry(src_w: int, src_h: int, long_edge: int, aspect: tuple[int, int] | None, + anchor: str = "bottom") -> tuple[tuple[int, int, int, int], tuple[int, int]]: + """Crop box ``(x0, y0, w, h)`` in source pixels and output size ``(W, H)``. + + A source with a different aspect ratio than ``aspect`` is cropped to it first ("bottom" keeps + the bottom edge, where the arms enter; "center" crops symmetrically). The long edge is then + resized to ``long_edge`` and both sides snapped to multiples of 16 (VAE /8 and 2x2 packing). + """ + x0, y0, w, h = 0, 0, src_w, src_h + if aspect is not None and abs(src_w / src_h - aspect[0] / aspect[1]) > 1e-3: + if src_w / src_h < aspect[0] / aspect[1]: # too tall: crop height + h = int(round(src_w * aspect[1] / aspect[0])) + y0 = src_h - h if anchor == "bottom" else (src_h - h) // 2 + else: # too wide: crop width, centred + w = int(round(src_h * aspect[0] / aspect[1])) + x0 = (src_w - w) // 2 + s = long_edge / max(w, h) + snap = lambda v: max(16, int(round(v / 16) * 16)) # noqa: E731 + return (x0, y0, w, h), (snap(w * s), snap(h * s)) + + +def crop_resize(img: np.ndarray, box, size, src_size, nearest: bool = False) -> np.ndarray: + """Crop ``box`` (source-frame pixels) out of ``img`` and resize to ``size``. + + ``img`` may have a different resolution than the source frame ``src_size`` (e.g. a mask video + stored smaller); the box is rescaled accordingly. RGB uses PIL Lanczos (as the original + pipeline did), masks use nearest neighbour. + """ + fx, fy = img.shape[1] / src_size[0], img.shape[0] / src_size[1] + x0, y0 = int(round(box[0] * fx)), int(round(box[1] * fy)) + x1, y1 = int(round((box[0] + box[2]) * fx)), int(round((box[1] + box[3]) * fy)) + sub = img[y0:y1, x0:x1] + if nearest: + return cv2.resize(sub.astype(np.uint8), size, interpolation=cv2.INTER_NEAREST) + pil = Image.fromarray(sub) + return np.asarray(pil if pil.size == tuple(size) else pil.resize(size, Image.LANCZOS)) + + +@dataclass +class ViewFrame: + """One camera frame mapped to the model resolution. + + image: uint8 ``(H, W, 3)``; hulls: per-arm convex hulls in output pixels; kp/vis: hand + keypoints in output pixels; object_mask: uint8 ``(H, W)`` or None. + """ + image: np.ndarray + hulls: list + kp: np.ndarray + vis: np.ndarray + object_mask: np.ndarray | None + + +def load_view_frame(reader: FrameReader, ep: Episode, view: str, t: int, long_edge: int, + aspect=None, anchor: str = "bottom", with_object: bool = True) -> ViewFrame: + from ..masks import arm_hulls, scale_hulls + rgb = reader.rgb(ep.video(view), t) + src = (rgb.shape[1], rgb.shape[0]) + box, size = output_geometry(src[0], src[1], long_edge, aspect, anchor) + sx, sy = size[0] / box[2], size[1] / box[3] + image = crop_resize(rgb, box, size, src) + m = reader.binary(ep.mask(view), t) + mx, my = src[0] / m.shape[1], src[1] / m.shape[0] # mask px -> source px + hulls = scale_hulls(arm_hulls(m), mx * sx, my * sy, (box[0] / mx, box[1] / my)) + kp, vis = reader.keypoints(ep.keypoints(view)) + kp_t = (kp[t] - np.array(box[:2], np.float32)) * np.array([sx, sy], np.float32) + om = None + if with_object and ep.object_mask(view) is not None: + om = crop_resize(reader.binary(ep.object_mask(view), t).astype(np.uint8), box, size, src, + nearest=True) + return ViewFrame(image, hulls, kp_t, vis[t].copy(), om) diff --git a/robotryon/image/fill_channels.py b/robotryon/image/fill_channels.py new file mode 100644 index 0000000000000000000000000000000000000000..b8689574ee6bf465b2bf5a1c0d54e062b51163fa --- /dev/null +++ b/robotryon/image/fill_channels.py @@ -0,0 +1,92 @@ +"""Channel-concatenated inpainting conditioning for the FLUX.2 [klein] transformer. + +The conditioned frame and the mask are pixel-aligned with the target, so they ride the CHANNEL +axis of the target tokens (FLUX-Fill style):: + + target token = [ noisy latent (128) | conditioned-frame latent (128) | folded mask (256) ] + +The patch embedding ``x_embedder`` is ``Linear(128, inner_dim)`` in the base model and is applied +once to the whole token sequence, so it is widened to 512 input features. The new columns are +zero-initialised: before training the widened model reproduces the base model exactly, and a +conditioning stream it cannot use degrades to "ignored" rather than to noise. The widened +embedding is always trained in full (FFT and LoRA). + +Reference tokens are spatially unaligned with the target and ride the SEQUENCE axis instead; they +carry only the 128 real latent channels and are zero-padded to 512, which contributes exactly +what the base model would because the new columns start at zero. +""" +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.nn.functional as F + +BASE_IN = 128 # 32 VAE latent channels x 2x2 patchify +MASK_IN = 256 # 8x8 pixel fold -> 64 channels at H/8, then 2x2 patchify +FILL_IN = BASE_IN + BASE_IN + MASK_IN # 512 + + +def widen_x_embedder(transformer, new_in: int = FILL_IN) -> nn.Linear: + """Grow ``transformer.x_embedder`` in place (old columns kept, new ones zero). Idempotent. + + Also records the new width in the model config (``in_channels=512, out_channels=128``) so a + full checkpoint written with ``save_pretrained`` reloads with ``from_pretrained`` directly. + """ + old: nn.Linear = transformer.x_embedder + if old.in_features == new_in: + return old + if old.in_features != BASE_IN: + raise ValueError(f"x_embedder has {old.in_features} inputs, expected {BASE_IN} or {new_in}") + new = nn.Linear(new_in, old.out_features, bias=old.bias is not None, + device=old.weight.device, dtype=old.weight.dtype) + with torch.no_grad(): + new.weight.zero_() + new.weight[:, :BASE_IN].copy_(old.weight) + if old.bias is not None: + new.bias.copy_(old.bias) + transformer.x_embedder = new + out_ch = transformer.config.out_channels or transformer.config.in_channels + transformer.register_to_config(in_channels=new_in, out_channels=out_ch) + return new + + +def fold_mask_to_latent(mask: torch.Tensor, vae_scale: int = 8) -> torch.Tensor: + """``(B, 1, H, W)`` mask in [0, 1] -> ``(B, 256, H/16, W/16)``. + + The mask is not interpolated: each 8x8 pixel block is folded into 64 channels (as in FLUX + Fill) so sub-latent edge precision survives, then the same 2x2 patchify as the latents. + """ + b, c, h, w = mask.shape + if c != 1: + raise ValueError(f"expected a single-channel mask, got {c}") + if h % (vae_scale * 2) or w % (vae_scale * 2): + raise ValueError(f"mask {h}x{w} must be a multiple of {vae_scale * 2}") + lh, lw = h // vae_scale, w // vae_scale + m = mask.view(b, lh, vae_scale, lw, vae_scale) + m = m.permute(0, 2, 4, 1, 3).reshape(b, vae_scale * vae_scale, lh, lw) + m = m.view(b, vae_scale * vae_scale, lh // 2, 2, lw // 2, 2) + return m.permute(0, 1, 3, 5, 2, 4).reshape(b, vae_scale * vae_scale * 4, lh // 2, lw // 2) + + +def pad_condition_tokens(packed: torch.Tensor, new_in: int = FILL_IN) -> torch.Tensor: + """Zero-pad ``(B, N, 128)`` reference tokens to the widened embedding width.""" + if packed.shape[-1] == new_in: + return packed + return F.pad(packed, (0, new_in - packed.shape[-1])) + + +def pack(latents: torch.Tensor) -> torch.Tensor: + """``(B, C, h, w)`` -> ``(B, h*w, C)`` tokens (same as ``Flux2KleinPipeline._pack_latents``).""" + b, c, h, w = latents.shape + return latents.reshape(b, c, h * w).permute(0, 2, 1) + + +def unpack(tokens: torch.Tensor, h: int, w: int) -> torch.Tensor: + """Inverse of :func:`pack`.""" + b, n, c = tokens.shape + return tokens.permute(0, 2, 1).reshape(b, c, h, w) + + +def fill_tokens(noisy: torch.Tensor, cond: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + """Target tokens ``[noisy | cond | mask]`` from latents ``(B,128,h,w)`` and pixel mask ``(B,1,H,W)``.""" + return torch.cat([pack(noisy), pack(cond), pack(fold_mask_to_latent(mask.to(cond.dtype)))], dim=-1) diff --git a/robotryon/image/model.py b/robotryon/image/model.py new file mode 100644 index 0000000000000000000000000000000000000000..6563a3bbb2fe4069c029d4d19a3e77713b979166 --- /dev/null +++ b/robotryon/image/model.py @@ -0,0 +1,138 @@ +"""Build, save and load the keyframe-translation model f_theta (FLUX.2 [klein] 9B Base). + +Two training modes share everything but the trainable set: + +* ``fft``: the whole transformer is trained (FSDP, paper Table 4). Checkpoint: + ``/transformer/`` written by ``save_pretrained`` (sharded safetensors, config records the + widened ``in_channels=512``). +* ``lora``: LoRA adapters on the attention / fused QKV-MLP projections of every block plus the + widened ``x_embedder`` trained in full. Checkpoint: ``/pytorch_lora_weights.safetensors`` + (diffusers LoRA format) and ``/x_embedder.safetensors``. + +``load_pipeline`` recognises either checkpoint type. +""" +from __future__ import annotations + +from pathlib import Path + +import torch +from safetensors.torch import load_file, save_file + +from .fill_channels import FILL_IN, widen_x_embedder + +LORA_FILE = "pytorch_lora_weights.safetensors" +X_EMBEDDER_FILE = "x_embedder.safetensors" +NUM_SINGLE_BLOCKS = 24 + + +def lora_targets() -> list[str]: + """Attention projections of the double-stream blocks, the fused QKV+MLP input projection of + the single-stream blocks, and the single-stream attention output projections.""" + return (["to_k", "to_q", "to_v", "to_out.0", "to_qkv_mlp_proj"] + + [f"single_transformer_blocks.{i}.attn.to_out" for i in range(NUM_SINGLE_BLOCKS)]) + + +def load_transformer(base: str, dtype: torch.dtype): + from diffusers import Flux2Transformer2DModel + return Flux2Transformer2DModel.from_pretrained(base, subfolder="transformer", torch_dtype=dtype) + + +def build_trainable(base: str, mode: str, lora_cfg=None): + """Transformer ready for training (x_embedder widened, trainable set marked). + + fft: loaded in fp32 (FSDP keeps fp32 master shards and computes in bf16), all parameters + trainable. lora: base weights frozen in bf16, adapters and x_embedder trainable in fp32. + """ + if mode == "fft": + tr = load_transformer(base, torch.float32) + widen_x_embedder(tr, FILL_IN) + tr.requires_grad_(True) + elif mode == "lora": + from peft import LoraConfig + tr = load_transformer(base, torch.bfloat16) + tr.requires_grad_(False) + tr.add_adapter(LoraConfig(r=int(lora_cfg.rank), lora_alpha=int(lora_cfg.alpha), + lora_dropout=float(lora_cfg.get("dropout", 0.0)), + init_lora_weights="gaussian", target_modules=lora_targets())) + widen_x_embedder(tr, FILL_IN).requires_grad_(True) + for p in tr.parameters(): + if p.requires_grad: + p.data = p.data.float() + else: + raise ValueError(f"train.mode must be 'fft' or 'lora', got {mode!r}") + return tr + + +def save_lora(transformer, out_dir: str | Path) -> None: + """LoRA tensors in diffusers format plus the full widened x_embedder.""" + from diffusers import Flux2KleinPipeline + from peft.utils import get_peft_model_state_dict + out_dir = Path(out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + lora = get_peft_model_state_dict(transformer) + lora = {k: v.detach().to(torch.bfloat16).cpu().contiguous() for k, v in lora.items()} + Flux2KleinPipeline.save_lora_weights(out_dir, transformer_lora_layers=lora) + save_file({"x_embedder.weight": transformer.x_embedder.weight.detach().float().cpu().contiguous()}, + str(out_dir / X_EMBEDDER_FILE)) + + +def save_full(transformer, out_dir: str | Path, state_dict: dict | None = None) -> None: + """Full transformer as a diffusers model directory ``/transformer`` (bf16 shards). + + ``state_dict`` is the gathered full state dict when training under FSDP (only the main + process calls this); the unwrapped module supplies the config. + """ + import json + from huggingface_hub import split_torch_state_dict_into_shards + d = Path(out_dir) / "transformer" + d.mkdir(parents=True, exist_ok=True) + sd = state_dict if state_dict is not None else transformer.state_dict() + sd = {k: v.detach().to(torch.bfloat16).cpu().contiguous() for k, v in sd.items()} + split = split_torch_state_dict_into_shards( + sd, filename_pattern="diffusion_pytorch_model{suffix}.safetensors", max_shard_size="5GB") + for fname, keys in split.filename_to_tensors.items(): + save_file({k: sd[k] for k in keys}, str(d / fname), metadata={"format": "pt"}) + if split.is_sharded: + index = {"metadata": split.metadata, "weight_map": split.tensor_to_filename} + (d / "diffusion_pytorch_model.safetensors.index.json").write_text(json.dumps(index, indent=2)) + transformer.save_config(d) + + +def checkpoint_kind(ckpt: str | Path) -> str: + ckpt = Path(ckpt) + if (ckpt / LORA_FILE).exists(): + return "lora" + if (ckpt / "transformer" / "config.json").exists(): + return "fft" + raise FileNotFoundError(f"{ckpt} holds neither {LORA_FILE} nor transformer/config.json") + + +def load_pipeline(base: str, ckpt: str | Path, device: str = "cuda", text_encoder: bool = True): + """``Flux2KleinPipeline`` in bf16 with the trained keyframe-translation weights. + + ``text_encoder=False`` skips the text encoder (prompts encoded beforehand). + + LoRA: widen ``x_embedder`` BEFORE loading the adapter (it was fitted to the 512-wide + embedding), copy the trained embedding, then load the adapter unfused. + """ + from diffusers import Flux2KleinPipeline + ckpt = Path(ckpt) + kind = checkpoint_kind(ckpt) + kw = {} if text_encoder else {"text_encoder": None, "tokenizer": None} + if kind == "fft": + from diffusers import Flux2Transformer2DModel + tr = Flux2Transformer2DModel.from_pretrained(ckpt / "transformer", torch_dtype=torch.bfloat16) + pipe = Flux2KleinPipeline.from_pretrained(base, transformer=tr, torch_dtype=torch.bfloat16, **kw) + pipe.to(device) + else: + pipe = Flux2KleinPipeline.from_pretrained(base, torch_dtype=torch.bfloat16, **kw).to(device) + widen_x_embedder(pipe.transformer, FILL_IN) + w = load_file(str(ckpt / X_EMBEDDER_FILE))["x_embedder.weight"] + with torch.no_grad(): + pipe.transformer.x_embedder.weight.copy_(w.to(pipe.transformer.x_embedder.weight.dtype)) + pipe.load_lora_weights(str(ckpt), weight_name=LORA_FILE) + if pipe.transformer.x_embedder.in_features != FILL_IN: + raise ValueError(f"{ckpt}: x_embedder has {pipe.transformer.x_embedder.in_features} inputs") + pipe.transformer.requires_grad_(False) + return pipe + diff --git a/robotryon/image/sampler.py b/robotryon/image/sampler.py new file mode 100644 index 0000000000000000000000000000000000000000..082bc74de7fb2fba96cdfb612c69602ff5186135 --- /dev/null +++ b/robotryon/image/sampler.py @@ -0,0 +1,121 @@ +"""Inpainting sampler for keyframe translation (true classifier-free guidance, Euler flow matching). + +``Flux2KleinPipeline.__call__`` cannot express channel-concatenated conditioning, so the denoising +loop is written out here. The token layout must match training exactly (see ``fill_channels`` and +``stereo_ids``):: + + [ target_v (T=0, L=v): noisy | conditioned frame | mask ] for every view v + [ reference_v (T=10, L=v): 128 latent channels, zero-padded ] for every view v + +The schedule is the pipeline's: ``sigmas = linspace(1, 1/steps, steps)`` shifted by the empirical +``mu`` of the total target token count (both views for stereo). FLUX.2 [klein] Base has no guidance +embedding, so guidance is true CFG against the empty prompt: +``v = v_uncond + s * (v_cond - v_uncond)``. + +The sampler returns the final latent TOKENS. In autoregressive inference these tokens are fed +directly as the next keyframe's reference tokens, without decoding and re-encoding (App. B.2.3). +""" +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np +import torch + +from .fill_channels import BASE_IN, FILL_IN, fold_mask_to_latent, pack, pad_condition_tokens, unpack +from .stereo_ids import fill_ids + + +class LatentCodec: + """Pixels <-> normalised, 2x2-patchified FLUX.2 VAE latents ``(1, 128, H/16, W/16)``.""" + + def __init__(self, vae, device): + self.vae, self.device = vae, device + self.mean = vae.bn.running_mean.view(1, -1, 1, 1).to(device) + self.std = torch.sqrt(vae.bn.running_var.view(1, -1, 1, 1) + vae.config.batch_norm_eps).to(device) + + @staticmethod + def to_tensor(img: np.ndarray) -> torch.Tensor: + """uint8 ``(H, W, 3)`` -> float ``(1, 3, H, W)`` in [-1, 1].""" + return torch.from_numpy(np.asarray(img).astype(np.float32) / 127.5 - 1).permute(2, 0, 1)[None] + + @staticmethod + def patchify(lat: torch.Tensor) -> torch.Tensor: + b, c, h, w = lat.shape + lat = lat.view(b, c, h // 2, 2, w // 2, 2).permute(0, 1, 3, 5, 2, 4) + return lat.reshape(b, c * 4, h // 2, w // 2) + + @staticmethod + def unpatchify(lat: torch.Tensor) -> torch.Tensor: + b, c, h, w = lat.shape + lat = lat.reshape(b, c // 4, 2, 2, h, w).permute(0, 1, 4, 2, 5, 3) + return lat.reshape(b, c // 4, h * 2, w * 2) + + def encode_tensor(self, px: torch.Tensor) -> torch.Tensor: + lat = self.vae.encode(px.to(self.device, self.vae.dtype)).latent_dist.mode() + return (self.patchify(lat) - self.mean) / self.std + + def encode(self, img: np.ndarray) -> torch.Tensor: + return self.encode_tensor(self.to_tensor(img)) + + def decode_tokens(self, tokens: torch.Tensor, h: int, w: int) -> np.ndarray: + """``(1, h*w, 128)`` tokens -> uint8 ``(16h, 16w, 3)`` image.""" + lat = unpack(tokens, h, w) * self.std + self.mean + img = self.vae.decode(self.unpatchify(lat).to(self.vae.dtype), return_dict=False)[0] + return ((img[0].float().permute(1, 2, 0).clamp(-1, 1).cpu().numpy() + 1) * 127.5).astype(np.uint8) + + +@dataclass +class ViewCondition: + """Conditioning of one camera view. + + cond_latent: ``(1, 128, h, w)`` latent of the conditioned frame Phi(o; M, S). + mask: ``(1, 1, 16h, 16w)`` pixel mask in [0, 1]. + reference_tokens: ``(1, h*w, 128)`` reference tokens (an encoded image, or the previous + keyframe's denoised tokens). + """ + cond_latent: torch.Tensor + mask: torch.Tensor + reference_tokens: torch.Tensor + + +@torch.no_grad() +def denoise(transformer, scheduler, views: list[ViewCondition], prompt_embeds, text_ids, + neg_embeds=None, neg_text_ids=None, steps: int = 30, guidance_scale: float = 3.5, + generator: torch.Generator | None = None) -> list[torch.Tensor]: + """Jointly denoise one or more views. Returns one ``(1, h*w, 128)`` token tensor per view.""" + from diffusers.pipelines.flux2.pipeline_flux2_klein import compute_empirical_mu + device, dtype = prompt_embeds.device, transformer.dtype + n = len(views) + _, _, h, w = views[0].cond_latent.shape + S = h * w + cond = [torch.cat([pack(v.cond_latent.float()), + pack(fold_mask_to_latent(v.mask.to(device, torch.float32)))], dim=-1).to(dtype) + for v in views] + refs = [pad_condition_tokens(v.reference_tokens.to(dtype), FILL_IN) for v in views] + ids = fill_ids(h, w, n).to(device) + lats = [torch.randn((1, S, BASE_IN), generator=generator, device=device, dtype=dtype) + for _ in range(n)] + + use_cfg = guidance_scale is not None and guidance_scale > 1.0 and neg_embeds is not None + sigmas = np.linspace(1.0, 1 / steps, steps) + scheduler.set_timesteps(sigmas=sigmas, device=device, mu=compute_empirical_mu(S * n, steps)) + scheduler.set_begin_index(0) + for t in scheduler.timesteps: + hidden = torch.cat([torch.cat([lats[i], cond[i]], dim=-1) for i in range(n)] + refs, dim=1) + tt = t.expand(1).to(dtype) / 1000 + pred = transformer(hidden_states=hidden, timestep=tt, guidance=None, + encoder_hidden_states=prompt_embeds, txt_ids=text_ids, img_ids=ids, + return_dict=False)[0] + if use_cfg: + neg = transformer(hidden_states=hidden, timestep=tt, guidance=None, + encoder_hidden_states=neg_embeds, txt_ids=neg_text_ids, img_ids=ids, + return_dict=False)[0] + pred = neg + guidance_scale * (pred - neg) + # One Euler step per view against the same t: rewind the scheduler's step counter. + idx = scheduler.step_index + for i in range(n): + scheduler._step_index = idx + lats[i] = scheduler.step(pred[:, i * S:(i + 1) * S], t, lats[i], + return_dict=False)[0].to(dtype) + return lats diff --git a/robotryon/image/stereo_ids.py b/robotryon/image/stereo_ids.py new file mode 100644 index 0000000000000000000000000000000000000000..f8ef2269d8bb27986cb661ad6d5434e796dcdce9 --- /dev/null +++ b/robotryon/image/stereo_ids.py @@ -0,0 +1,34 @@ +"""RoPE position ids for one or two camera views in one token sequence. + +FLUX.2 gives every image token a 4D position ``(T, H, W, L)``. The pipeline sets ``L = 0`` for all +image tokens (only text tokens use it), so for images ``L`` is a free axis; stereo generation +uses it as the camera-view index. ``T`` keeps its learned meaning: targets at ``T = 0`` and the +reference at ``T = 10``. The token order is:: + + [target_v0, target_v1, ..., reference_v0, reference_v1, ...] + +With one view this is byte-identical to the stock ``_prepare_latent_ids`` + ``_prepare_image_ids``. +""" +from __future__ import annotations + +import torch + +TARGET_T = 0 +REFERENCE_T = 10 + + +def view_image_ids(sizes: list[tuple[int, int]], view_ids: list[int], + t_coords: list[int]) -> torch.Tensor: + """``sizes[i] = (h, w)`` token grid of image i -> ``(1, sum h*w, 4)`` int64 position ids.""" + if not (len(sizes) == len(view_ids) == len(t_coords)): + raise ValueError("one size, view id and t coordinate per image") + out = [torch.cartesian_prod(torch.tensor([int(t)]), torch.arange(h), torch.arange(w), + torch.tensor([int(v)])) + for (h, w), v, t in zip(sizes, view_ids, t_coords)] + return torch.cat(out, dim=0).unsqueeze(0) + + +def fill_ids(h: int, w: int, n_views: int) -> torch.Tensor: + """Ids for ``n_views`` targets followed by ``n_views`` references, all on an ``h x w`` grid.""" + return view_image_ids([(h, w)] * (2 * n_views), list(range(n_views)) * 2, + [TARGET_T] * n_views + [REFERENCE_T] * n_views) diff --git a/robotryon/image/text.py b/robotryon/image/text.py new file mode 100644 index 0000000000000000000000000000000000000000..d2ff7e9a6fe220fdca4ffb7ad30d414ad0c6897a --- /dev/null +++ b/robotryon/image/text.py @@ -0,0 +1,56 @@ +"""The on-disk text-embedding cache of the image model's prompts (``robotryon.captions``). + +Qwen3 prompt embeddings are ``(512, 12288)`` per prompt, too large to keep for every prompt in +memory, so they are computed once at trainer start-up and stored as one safetensors file per +prompt, named by a hash of the text. +""" +from __future__ import annotations + +import hashlib +from pathlib import Path + +import torch +from safetensors.torch import load_file, save_file + +MAX_SEQUENCE_LENGTH = 512 + + +class TextEmbedCache: + """One ``.safetensors`` file per prompt holding ``prompt_embeds`` ``(L, D)``.""" + + def __init__(self, root: str | Path): + self.root = Path(root) + + def path(self, prompt: str) -> Path: + return self.root / f"{hashlib.sha1(prompt.encode('utf-8')).hexdigest()}.safetensors" + + def missing(self, prompts) -> list[str]: + return sorted({p for p in prompts if not self.path(p).exists()}) + + def load(self, prompt: str) -> torch.Tensor: + p = self.path(prompt) + if not p.exists(): + raise FileNotFoundError(f"no cached embedding for prompt {prompt!r}") + return load_file(str(p))["prompt_embeds"] + + @torch.no_grad() + def fill(self, text_pipe, prompts, device, batch_size: int = 16) -> int: + """Encode and store the given prompts with a ``Flux2KleinPipeline`` holding the text encoder.""" + self.root.mkdir(parents=True, exist_ok=True) + prompts = list(prompts) + for i in range(0, len(prompts), batch_size): + chunk = prompts[i:i + batch_size] + emb, _ = text_pipe.encode_prompt(prompt=chunk, device=device, + max_sequence_length=MAX_SEQUENCE_LENGTH) + for p, e in zip(chunk, emb): + tmp = self.path(p).with_suffix(".tmp") + save_file({"prompt_embeds": e.to(torch.bfloat16).contiguous().cpu()}, str(tmp)) + tmp.replace(self.path(p)) + return len(prompts) + + +def text_ids_for(prompt_embeds: torch.Tensor) -> torch.Tensor: + """``(B, L, D)`` -> ``(B, L, 4)`` text position ids (same as ``Flux2KleinPipeline._prepare_text_ids``).""" + b, n, _ = prompt_embeds.shape + ids = torch.cartesian_prod(torch.arange(1), torch.arange(1), torch.arange(1), torch.arange(n)) + return ids.unsqueeze(0).expand(b, -1, -1).to(prompt_embeds.device) diff --git a/robotryon/image/training.py b/robotryon/image/training.py new file mode 100644 index 0000000000000000000000000000000000000000..6f1b39d135803e2b6e3070e9244edfd570d40a80 --- /dev/null +++ b/robotryon/image/training.py @@ -0,0 +1,48 @@ +"""Flow-matching objective of the keyframe-translation model (paper Eq. 3.1). + +With ``z = E_img(o_t)`` and ``z_tau = (1 - tau) z + tau eps``, the transformer predicts +``eps - z`` from the target tokens ``[z_tau | E_img(Phi(o_t; M_t, S_t)) | fold(M_t)]`` and the +reference tokens ``E_img(o_t')``. ``tau`` is drawn uniformly from the base model's training +sigma grid, one value per sample shared by its views. For stereo samples the loss is the mean +over the two views. +""" +from __future__ import annotations + +import torch + +from .fill_channels import FILL_IN, fill_tokens, pack, pad_condition_tokens, unpack +from .sampler import LatentCodec +from .stereo_ids import fill_ids +from .text import text_ids_for + + +def flow_matching_loss(transformer, codec: LatentCodec, sched_sigmas: torch.Tensor, + sched_timesteps: torch.Tensor, batch: dict, device) -> torch.Tensor: + B, V = batch["target"].shape[:2] + with torch.no_grad(): + enc = [{k: codec.encode_tensor(batch[k][:, v].to(device)) + for k in ("target", "cond", "reference")} for v in range(V)] + z0 = enc[0]["target"] + _, _, h, w = z0.shape + n = sched_timesteps.numel() + idx = (torch.rand(B) * n).long().clamp(max=n - 1) + sigma = sched_sigmas[idx].to(device=device, dtype=z0.dtype).view(B, 1, 1, 1) + timestep = sched_timesteps[idx].to(device) + + targets, tokens = [], [] + for v, e in enumerate(enc): + noise = torch.randn_like(e["target"]) + noisy = (1.0 - sigma) * e["target"] + sigma * noise + tokens.append(fill_tokens(noisy, e["cond"], batch["mask"][:, v].to(device))) + targets.append(noise - e["target"]) + tokens += [pad_condition_tokens(pack(e["reference"]), FILL_IN) for e in enc] + hidden = torch.cat(tokens, dim=1).to(torch.bfloat16) + ids = fill_ids(h, w, V).to(device).expand(B, -1, -1) + prompt_embeds = batch["prompt_embeds"].to(device) + pred = transformer(hidden_states=hidden, timestep=timestep / 1000, guidance=None, + encoder_hidden_states=prompt_embeds, txt_ids=text_ids_for(prompt_embeds), + img_ids=ids, return_dict=False)[0] + S = h * w + losses = [((unpack(pred[:, v * S:(v + 1) * S], h, w).float() - targets[v].float()) ** 2) + .reshape(B, -1).mean(1).mean() for v in range(V)] + return sum(losses) / V diff --git a/robotryon/masks/__init__.py b/robotryon/masks/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9c215dc64adcae929b52c2c5527be5b3eb922b56 --- /dev/null +++ b/robotryon/masks/__init__.py @@ -0,0 +1,14 @@ +"""Shape-agnostic arm/hand masks for keyframe translation (paper Sec. 3.2, App. B.2.3). + +The mask must hide the silhouette of whatever embodiment it covers, so it is never the raw +segmentation: every arm component becomes its convex hull, which is then dilated. During training +the dilation is random per arm, the manipulated object is cut out of the hole with some +probability, and the whole sample is jointly rescaled (field-of-view crop / pad). At inference the +dilation is fixed and the eroded object mask is always subtracted. +""" +from .hull import arm_hulls, fill_hulls, scale_hulls +from .augment import MaskAug, fov_augment, inference_mask +from .reference import robot_only + +__all__ = ["arm_hulls", "fill_hulls", "scale_hulls", "MaskAug", "fov_augment", "inference_mask", + "robot_only"] diff --git a/robotryon/masks/augment.py b/robotryon/masks/augment.py new file mode 100644 index 0000000000000000000000000000000000000000..90985ebb70aff6386859e53fe1ddd9120ab7d575 --- /dev/null +++ b/robotryon/masks/augment.py @@ -0,0 +1,131 @@ +"""Train-time mask randomisation, the fixed inference mask, and the joint field-of-view augmentation. + +Training (``MaskAug.train_mask``): + * one random dilation per arm hull, so two arms in one frame already carry different + hole-to-arm ratios and the model cannot read the hand size off the hole; + * episodes flagged as "bonus" (thin human arms) get a raised dilation floor, because a hull + hugging a thin arm leaves no room for a thicker robot arm; + * with probability ``p_object_cut``, where an object mask exists, the eroded object is cut out + of the hole (a concave hole), which is what inference always does. + +Inference (``inference_mask``): every hull dilated by a fixed kernel (the middle of the training +range plus the bonus), then the eroded object mask is subtracted so the object is preserved. +""" +from __future__ import annotations + +from dataclasses import dataclass + +import cv2 +import numpy as np +from PIL import Image + +from .hull import fill_hulls + + +def _ellipse(k: int) -> np.ndarray: + k = int(k) | 1 + return cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (k, k)) + + +@dataclass +class MaskAug: + dilate_min: int = 3 + dilate_max: int = 71 + bonus: int = 21 # added to the dilation floor for "bonus" (human) episodes + p_object_cut: float = 0.5 + obj_erode_min: int = 3 + obj_erode_max: int = 21 + min_keep_frac: float = 0.4 # reject a cut that leaves less than this fraction of the hole + + @classmethod + def from_config(cls, cfg) -> "MaskAug": + cfg = cfg or {} + return cls(**{k: cfg[k] for k in cls.__dataclass_fields__ if k in cfg}) + + def sample_dilate(self, rng: np.random.Generator, bonus: bool = False) -> int: + lo = self.dilate_min + (self.bonus if bonus else 0) + return int(rng.integers(lo, max(lo + 1, self.dilate_max) + 1)) | 1 + + def sample_obj_erode(self, rng: np.random.Generator) -> int: + return int(rng.integers(self.obj_erode_min, self.obj_erode_max + 1)) | 1 + + def infer_dilate(self) -> int: + """Fixed inference dilation: middle of the training range plus the bonus.""" + return (self.dilate_min + self.dilate_max) // 2 + self.bonus + + def train_mask(self, polys: list[np.ndarray], width: int, height: int, + rng: np.random.Generator, bonus: bool = False, + object_mask: np.ndarray | None = None) -> tuple[np.ndarray, int]: + """Randomised training hole. Returns ``(mask uint8 {0,1}, object erode kernel or 0)``.""" + mask = np.zeros((height, width), np.uint8) + for one in fill_hulls(polys, width, height): + mask |= cv2.dilate(one, _ellipse(self.sample_dilate(rng, bonus))) + if object_mask is None: + return mask, 0 + if float(rng.random()) >= self.p_object_cut: + return mask, 0 + om = (np.asarray(object_mask) > 0).astype(np.uint8) + if om.sum() == 0: + return mask, 0 + e = self.sample_obj_erode(rng) + om = cv2.erode(om, _ellipse(e)) + cut = (mask.astype(bool) & ~om.astype(bool)).astype(np.uint8) + hull_px = int(mask.sum()) + if cut.sum() > 0 and (hull_px == 0 or cut.sum() / hull_px > self.min_keep_frac): + return cut, e + return mask, 0 + + +def inference_mask(polys: list[np.ndarray], width: int, height: int, dilate: int, + object_mask: np.ndarray | None = None, obj_erode: int = 9) -> np.ndarray: + """Fixed-dilation hull mask minus the eroded object -> ``(H, W)`` uint8 {0, 1}.""" + mask = np.zeros((height, width), np.uint8) + for one in fill_hulls(polys, width, height): + mask |= one + if mask.any(): + mask = cv2.dilate(mask, _ellipse(dilate)) + if object_mask is not None: + om = (np.asarray(object_mask) > 0).astype(np.uint8) + if obj_erode > 0: + om = cv2.erode(om, _ellipse(obj_erode)) + mask = (mask.astype(bool) & ~om.astype(bool)).astype(np.uint8) + return mask + + +def fov_augment(imgs: list[Image.Image], rng: np.random.Generator, crop_min: float, + scale_max: float = 1.0, flip_prob: float = 0.0) -> tuple[list[Image.Image], float]: + """One shared field-of-view draw for every stream of a sample (all views included). + + ``scale ~ U[crop_min, max(1, scale_max)]``. Below 1 a ``scale``-sized window is cropped and + resized back; above 1 the image is pasted at a random offset onto a black canvas ``scale`` + times larger and resized back, so the scene shrinks and black borders appear (the mask is + black there too: "not a hole"). The model thus learns that the camera's field of view is not + fixed. + """ + if not imgs: + return imgs, 1.0 + w, h = imgs[0].size + if any(im.size != (w, h) for im in imgs): + raise ValueError("fov_augment needs equally sized images") + scale = float(rng.uniform(crop_min, max(1.0, float(scale_max)))) + if scale > 1.0 + 1e-6: + cw, ch = int(round(w * scale)), int(round(h * scale)) + left = int(rng.integers(0, cw - w + 1)) + top = int(rng.integers(0, ch - h + 1)) + flip = float(rng.random()) < flip_prob + out = [] + for im in imgs: + canvas = Image.new(im.mode, (cw, ch), 0) + canvas.paste(im, (left, top)) + im2 = canvas.resize((w, h), Image.LANCZOS) + out.append(im2.transpose(Image.FLIP_LEFT_RIGHT) if flip else im2) + return out, scale + cw, ch = max(8, int(round(w * scale))), max(8, int(round(h * scale))) + left = int(rng.integers(0, w - cw + 1)) if w > cw else 0 + top = int(rng.integers(0, h - ch + 1)) if h > ch else 0 + flip = float(rng.random()) < flip_prob + out = [] + for im in imgs: + im2 = im.crop((left, top, left + cw, top + ch)).resize((w, h), Image.LANCZOS) + out.append(im2.transpose(Image.FLIP_LEFT_RIGHT) if flip else im2) + return out, scale diff --git a/robotryon/masks/hull.py b/robotryon/masks/hull.py new file mode 100644 index 0000000000000000000000000000000000000000..d485f44d18c6e53b48a5dad5e666e66ab22f50fb --- /dev/null +++ b/robotryon/masks/hull.py @@ -0,0 +1,53 @@ +"""Convex hull per arm. + +A segmentation silhouette gives the embodiment away (a robot arm has cables and a different +outline than a human arm). Each connected component of the mask is therefore replaced by its +convex hull. Small gaps are closed before the component split, so a limb and a thin cable that +pinches to a one-pixel thread stay one component instead of producing a separate cable-shaped +polygon. +""" +from __future__ import annotations + +import cv2 +import numpy as np + +PRE_DILATE = 9 # closing kernel before the component split (px at mask resolution) +MIN_COMPONENT_PX = 256 # smaller components are segmentation noise, not an arm + + +def arm_hulls(mask: np.ndarray, pre_dilate: int = PRE_DILATE, + min_component_px: int = MIN_COMPONENT_PX) -> list[np.ndarray]: + """Binary mask ``(H, W)`` -> list of ``(K, 2)`` float32 hull polygons in mask pixels.""" + m = (np.asarray(mask) > 0).astype(np.uint8) + if m.sum() == 0: + return [] + if pre_dilate > 0: + k = pre_dilate | 1 + m = cv2.dilate(m, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (k, k))) + n, lab = cv2.connectedComponents(m) + polys = [] + for i in range(1, n): + ys, xs = np.where(lab == i) + if ys.size < min_component_px: + continue + pts = np.stack([xs, ys], axis=1).astype(np.int32) + polys.append(cv2.convexHull(pts).reshape(-1, 2).astype(np.float32)) + return polys + + +def scale_hulls(polys: list[np.ndarray], sx: float, sy: float, + offset: tuple[float, float] = (0.0, 0.0)) -> list[np.ndarray]: + """Map polygons from mask pixels to output pixels: ``(p - offset) * (sx, sy)``.""" + off = np.asarray(offset, np.float32) + s = np.array([sx, sy], np.float32) + return [(p - off) * s for p in polys] + + +def fill_hulls(polys: list[np.ndarray], width: int, height: int) -> list[np.ndarray]: + """Rasterise each polygon separately -> list of ``(H, W)`` uint8 {0, 1} masks.""" + out = [] + for p in polys: + one = np.zeros((height, width), np.uint8) + cv2.fillConvexPoly(one, p.astype(np.int32), 1) + out.append(one) + return out diff --git a/robotryon/masks/reference.py b/robotryon/masks/reference.py new file mode 100644 index 0000000000000000000000000000000000000000..db5b2d9e0b9db2b8757e0073887670005a959f50 --- /dev/null +++ b/robotryon/masks/reference.py @@ -0,0 +1,29 @@ +"""Robot-only reference image: the robot's pixels on a plain white background. + +The first keyframe of a demonstration has no translated predecessor, so its reference is a real +frame of the target robot with everything but the robot removed; this selects the embodiment. +Training uses the same construction for its robot-only exemplars. +""" + +from __future__ import annotations + +import cv2 +import numpy as np + +WHITE = 255 + + +def robot_only(image: np.ndarray, mask: np.ndarray, erode: int = 7, + background: int = WHITE) -> np.ndarray: + """Keep the robot pixels of ``image`` (H, W, 3) where ``mask`` (H, W) is set. + + The mask is eroded by ``erode`` px so no background leaks in at the silhouette; everything + else is filled with ``background``. + """ + m = np.asarray(mask).astype(bool) + if erode > 0: + k = int(erode) | 1 + m = cv2.erode(m.astype(np.uint8), cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (k, k))) > 0 + out = np.full_like(image, background) + out[m] = image[m] + return out diff --git a/robotryon/skeleton/__init__.py b/robotryon/skeleton/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9b61ea5aa1117142a95d3ef539217ffab5b00c74 --- /dev/null +++ b/robotryon/skeleton/__init__.py @@ -0,0 +1,8 @@ +from .augment import SkeletonAugment +from .render import condition_frame, draw_skeleton, draw_skeleton_video, paint +from .topology import EDGES, KEYPOINT_NAMES, MEDIAPIPE21_TO_17, NUM_HANDS, NUM_KEYPOINTS + +__all__ = [ + "SkeletonAugment", "condition_frame", "draw_skeleton", "draw_skeleton_video", "paint", + "EDGES", "KEYPOINT_NAMES", "MEDIAPIPE21_TO_17", "NUM_HANDS", "NUM_KEYPOINTS", +] diff --git a/robotryon/skeleton/augment.py b/robotryon/skeleton/augment.py new file mode 100644 index 0000000000000000000000000000000000000000..84279fc5e898cb66ecc8823f22e561522d6852bd --- /dev/null +++ b/robotryon/skeleton/augment.py @@ -0,0 +1,75 @@ +"""Training-time skeleton augmentation. + +Human keypoints at inference come from a reconstructed and retargeted hand, which is imprecise. +To make the generative models treat the skeleton as a coarse guide rather than an exact pose, +robot keypoints are corrupted during training: + +* each keypoint is displaced by up to half the length of the finger segment it belongs to, +* the whole skeleton is dropped with probability ``p_drop_all``, +* each keypoint is dropped independently with probability ``p_drop_keypoint``. + +For video, the displacement is sampled once per clip and per keypoint and then held fixed in +time, so the skeleton flow stays temporally coherent; keypoint dropout is also per clip. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np + +from .topology import segment_length_per_keypoint + + +@dataclass +class SkeletonAugment: + max_shift_frac: float = 0.5 + p_drop_all: float = 0.2 + p_drop_keypoint: float = 0.2 + enabled: bool = True + + @classmethod + def from_config(cls, cfg) -> "SkeletonAugment": + if cfg is None: + return cls(enabled=False) + return cls(**{k: cfg[k] for k in ("max_shift_frac", "p_drop_all", "p_drop_keypoint", + "enabled") if k in cfg}) + + def __call__(self, kp: np.ndarray, vis: np.ndarray, rng: np.random.Generator, + temporal: bool | None = None) -> tuple[np.ndarray, np.ndarray]: + """Corrupt keypoints. + + Args: + kp: ``(H, 17, 2)`` for one frame or ``(T, H, 17, 2)`` for a clip, in pixels. + vis: matching visibility, ``(H, 17)`` or ``(T, H, 17)``. + rng: random generator. + temporal: treat the first axis as time. Inferred from ``kp.ndim == 4`` if omitted. + + Returns: + ``(kp, vis)``: new arrays, the inputs are not modified. + """ + kp = np.array(kp, dtype=np.float32, copy=True) + vis = np.array(vis, dtype=bool, copy=True) + if not self.enabled: + return kp, vis + if temporal is None: + temporal = kp.ndim == 4 + shape = kp.shape[1:-1] if temporal else kp.shape[:-1] # (H, 17) + + if rng.random() < self.p_drop_all: + vis[...] = False + return kp, vis + + # Segment length from the clip median, so a single bad frame cannot blow up the shift. + seg = segment_length_per_keypoint(kp, vis) + if temporal: + seg = np.median(np.where(vis, seg, np.nan), axis=0) + seg = np.nan_to_num(seg, nan=0.0) + radius = self.max_shift_frac * seg * np.sqrt(rng.random(shape)) # uniform in a disc + angle = rng.uniform(0.0, 2 * np.pi, shape) + shift = np.stack([radius * np.cos(angle), radius * np.sin(angle)], -1).astype(np.float32) + kp = kp + (shift[None] if temporal else shift) + + keep = rng.random(shape) >= self.p_drop_keypoint + vis &= keep[None] if temporal else keep + return kp, vis diff --git a/robotryon/skeleton/render.py b/robotryon/skeleton/render.py new file mode 100644 index 0000000000000000000000000000000000000000..2e56c576c50adfa242bd0b6a47911c0ebdf7571f --- /dev/null +++ b/robotryon/skeleton/render.py @@ -0,0 +1,75 @@ +"""Draw hand skeletons as 2D images. + +The same renderer is used for every embodiment and for humans, so the skeleton carries pose +only and never reveals what kind of hand produced it. +""" + +from __future__ import annotations + +import cv2 +import numpy as np + +from .topology import COLORS, EDGES, GROUP_OF + + +def draw_skeleton(kp: np.ndarray, vis: np.ndarray, width: int, height: int, + thickness: int = 2, radius: int = 3, supersample: int = 2) -> np.ndarray: + """Render one frame of hand skeletons on a black canvas. + + Args: + kp: ``(H, 17, 2)`` pixel coordinates. + vis: ``(H, 17)`` visibility. Invisible or non-finite keypoints are skipped together with + their edges. + width, height: canvas size in pixels. + + Returns: + ``(height, width, 3)`` uint8 RGB image. + """ + s = max(1, int(supersample)) + canvas = np.zeros((height * s, width * s, 3), np.uint8) + # Same stroke convention as the original renderer: sizes are given on the supersampled + # canvas at half scale, so thickness=2, radius=3 is the production look. + lw, dr = max(1, thickness * s // 2), max(1, radius * s // 2) + kp = np.asarray(kp, np.float64) + vis = np.asarray(vis, bool) & np.isfinite(kp).all(-1) + for h in range(kp.shape[0]): + pts = np.round(kp[h] * s).astype(np.int64) if vis[h].any() else None + if pts is None: + continue + for a, b, g in EDGES: + if vis[h, a] and vis[h, b]: + cv2.line(canvas, tuple(pts[a]), tuple(pts[b]), COLORS[g], lw, cv2.LINE_AA) + for j in range(kp.shape[1]): + if vis[h, j]: + cv2.circle(canvas, tuple(pts[j]), dr, COLORS[GROUP_OF[j]], -1, cv2.LINE_AA) + if s > 1: + canvas = cv2.resize(canvas, (width, height), interpolation=cv2.INTER_AREA) + return canvas + + +def draw_skeleton_video(kp: np.ndarray, vis: np.ndarray, width: int, height: int, + **kw) -> np.ndarray: + """``(T, H, 17, 2)`` keypoints -> ``(T, height, width, 3)`` uint8 skeleton video.""" + return np.stack([draw_skeleton(kp[t], vis[t], width, height, **kw) + for t in range(kp.shape[0])]) + + +def paint(frame: np.ndarray, skeleton: np.ndarray) -> np.ndarray: + """Overlay a rendered skeleton on a frame; skeleton pixels win wherever drawn.""" + out = frame.copy() + lit = skeleton.max(axis=-1) > 0 + out[lit] = skeleton[lit] + return out + + +def condition_frame(frame: np.ndarray, mask: np.ndarray | None, + skeleton: np.ndarray) -> np.ndarray: + """Phi(o; M, S) = ((1 - M) * o) (+) S: black out the mask region, then draw the skeleton. + + The skeleton is drawn even slightly outside the mask so fingertips are never clipped. + ``mask=None`` leaves the frame unmasked (mask-free inference). + """ + out = frame.copy() + if mask is not None: + out[np.asarray(mask).astype(bool)] = 0 + return paint(out, skeleton) diff --git a/robotryon/skeleton/topology.py b/robotryon/skeleton/topology.py new file mode 100644 index 0000000000000000000000000000000000000000..12b56886bf0486c1a23dfbea05288d9549386d70 --- /dev/null +++ b/robotryon/skeleton/topology.py @@ -0,0 +1,90 @@ +"""Hand keypoint topology shared by every embodiment. + +Each hand has 17 keypoints in MediaPipe order with the four distal (DIP) joints removed: + + 0 wrist + 1, 2, 3, 4 thumb: base (CMC), MCP, IP, tip + 5, 6, 7 index: knuckle (MCP), middle (PIP), tip + 8, 9, 10 middle + 11, 12, 13 ring + 14, 15, 16 little + +Robot hands map their links onto the same 17 points, so hands with different degrees of +freedom share one skeleton. Arrays of keypoints use shape ``(..., 2, 17, D)``: hand 0 is the +left hand, hand 1 is the right hand. +""" + +from __future__ import annotations + +import numpy as np + +NUM_KEYPOINTS = 17 +NUM_HANDS = 2 + +KEYPOINT_NAMES = ( + "wrist", + "thumb_cmc", "thumb_mcp", "thumb_ip", "thumb_tip", + "index_mcp", "index_pip", "index_tip", + "middle_mcp", "middle_pip", "middle_tip", + "ring_mcp", "ring_pip", "ring_tip", + "little_mcp", "little_pip", "little_tip", +) + +# MediaPipe-21 index of each of the 17 keypoints (the DIP joints 7/11/15/19 are dropped). +MEDIAPIPE21_TO_17 = (0, 1, 2, 3, 4, 5, 6, 8, 9, 10, 12, 13, 14, 16, 17, 18, 20) + +GROUP_OF = ( + "palm", + "thumb", "thumb", "thumb", "thumb", + "index", "index", "index", + "middle", "middle", "middle", + "ring", "ring", "ring", + "little", "little", "little", +) + +# (parent, child, group). The palm connects the wrist to the thumb base and the little knuckle +# and links the knuckles to one another; each finger is a chain from its knuckle to its tip. +EDGES: tuple[tuple[int, int, str], ...] = ( + (0, 1, "palm"), (0, 5, "palm"), (5, 8, "palm"), (8, 11, "palm"), (11, 14, "palm"), + (0, 14, "palm"), + (1, 2, "thumb"), (2, 3, "thumb"), (3, 4, "thumb"), + (5, 6, "index"), (6, 7, "index"), + (8, 9, "middle"), (9, 10, "middle"), + (11, 12, "ring"), (12, 13, "ring"), + (14, 15, "little"), (15, 16, "little"), +) + +COLORS = { + "palm": (255, 170, 60), "thumb": (255, 232, 70), "index": (110, 235, 120), + "middle": (95, 205, 255), "ring": (185, 150, 255), "little": (255, 125, 205), +} + +# Finger segments a keypoint belongs to, used to scale the training-time perturbation. +# A keypoint is displaced by at most half the length of the shortest finger segment it touches. +_FINGER_EDGES = tuple((a, b) for a, b, g in EDGES if g != "palm") + + +def segment_length_per_keypoint(kp: np.ndarray, vis: np.ndarray | None = None) -> np.ndarray: + """Length of the finger segment each keypoint belongs to. + + Args: + kp: ``(..., 17, D)`` keypoints of one or more hands. + vis: optional ``(..., 17)`` visibility; invisible endpoints do not contribute. + + Returns: + ``(..., 17)`` lengths. The wrist uses its distance to the nearest knuckle. A keypoint + without any valid segment gets 0, so it is never displaced. + """ + kp = np.asarray(kp, dtype=np.float64) + if vis is None: + vis = np.isfinite(kp).all(-1) + out = np.full(kp.shape[:-1], np.inf) + edges = _FINGER_EDGES + ((0, 5), (0, 1)) + for a, b in edges: + ok = vis[..., a] & vis[..., b] + d = np.linalg.norm(np.nan_to_num(kp[..., a, :] - kp[..., b, :]), axis=-1) + d = np.where(ok, d, np.inf) + out[..., a] = np.minimum(out[..., a], d) + out[..., b] = np.minimum(out[..., b], d) + out[~np.isfinite(out)] = 0.0 + return out diff --git a/robotryon/video/__init__.py b/robotryon/video/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3e20e055fde7765849fb88fa699ca534e5edc101 --- /dev/null +++ b/robotryon/video/__init__.py @@ -0,0 +1,15 @@ +"""Video model g_phi: an in-context LoRA on LTX-2.5 that synthesizes the robot video between +translated keyframes under the guidance of the hand-skeleton flow (paper Sec. 3.3). + +Modules: + conditioning training sequence layout, keyframe sampling, flow-matching loss + dataset random training clips with on-the-fly skeleton rendering and augmentation + rope4d fourth RoPE axis distinguishing the two stereo views + pipeline diffusers pipeline subclass and the two-stage generation of one clip + chunks splitting long videos at keyframes and concatenating the clips + text caption embeddings (cached once; the text encoder is never co-resident) + training gradient checkpointing and LoRA weight I/O + io mp4 / png helpers + +Heavy dependencies (torch, diffusers) are imported by the submodules that need them. +""" diff --git a/robotryon/video/chunks.py b/robotryon/video/chunks.py new file mode 100644 index 0000000000000000000000000000000000000000..3bea856679a19aa56caf7edadf143b097c995177 --- /dev/null +++ b/robotryon/video/chunks.py @@ -0,0 +1,164 @@ +"""Splitting a long demonstration into clips the video model can generate at once. + +The video VAE takes clips of ``8k + 1`` frames, and the model is trained on clips of at most +145 frames (about 7 s at 20 fps). A longer demonstration is therefore cut at its keyframes: +each clip runs from one keyframe to a later one, holds every keyframe in between, and the next +clip starts at the keyframe the previous one ended on. Adjacent clips thus **share their +boundary keyframe**, so the arm pose cannot jump at the seam, and the concatenation drops the +duplicated boundary frame of every clip after the first. + +A clip is ``8k + 1`` frames long, which usually reaches past its last keyframe; those trailing +frames are anchored to nothing and are discarded at concatenation. Two corner cases follow the +original generation code: + +* two consecutive keyframes further apart than the cap: the clip is stretched past the cap + rather than leaving a hole in the video (logged as off-distribution); +* the end of the video leaves less than the minimum clip length: a short tail clip of at least + 9 frames is generated instead of dropping the end of the demonstration, pulling the last + keyframe onto the clip's final frame if the ``8k + 1`` constraint requires it. +""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass + +import numpy as np + +VAE_TEMPORAL = 8 + + +def allowed_frame_counts(lo: int, hi: int) -> list[int]: + """Clip lengths in ``[lo, hi]`` the VAE accepts (``F = 8k + 1``).""" + return [f for f in range(lo, hi + 1) if (f - 1) % VAE_TEMPORAL == 0] + + +def next_8k1(n: int) -> int: + """Smallest ``8k + 1`` that is >= n.""" + return n if (n - 1) % VAE_TEMPORAL == 0 else n + (VAE_TEMPORAL - (n - 1) % VAE_TEMPORAL) + + +def prev_8k1(n: int) -> int: + """Largest ``8k + 1`` that is <= n (0 if none).""" + return n - (n - 1) % VAE_TEMPORAL if n >= 1 else 0 + + +def latent_index(pixel: int, num_frames: int) -> int: + """Latent frame holding a pixel frame: latent 0 is pixel 0, latent j>0 is [8j-7, 8j].""" + lat_n = (num_frames - 1) // VAE_TEMPORAL + 1 + return 0 if pixel == 0 else min((pixel + 7) // VAE_TEMPORAL, lat_n - 1) + + +@dataclass +class Chunk: + """One generated clip: ``frames`` frames starting at ``start`` of the episode. + + ``kf_slots`` index into the episode's keyframe list, ``kf_local`` are the same keyframes + as frame indices inside the clip (first is always 0). + """ + start: int + frames: int + kf_slots: list[int] + kf_local: list[int] + + def to_dict(self) -> dict: + return asdict(self) + + +def _tail_chunk(kf: list[int], i: int, start: int, avail: int, log) -> Chunk | None: + """A short last clip covering what is left after keyframe ``i``.""" + F = prev_8k1(avail) + if F < 9: + return None + ids = [t for t in range(i, len(kf)) if kf[t] <= start + F - 1] or [i] + if ids[-1] + 1 < len(kf): + ids.append(ids[-1] + 1) # include the last keyframe, pulled to F-1 + local = [min(kf[t] - start, F - 1) for t in ids] + if len(set(local)) < 2: + return None + log(f" tail clip: frames {start}..{start + F - 1} ({F}f, below the training minimum) " + f"keyframes {ids} @ {local}") + return Chunk(start, F, ids, local) + + +def plan_chunks(kf: list[int], total: int, frames_range=(49, 145), log=print) -> list[Chunk]: + """Cut a ``total``-frame video with keyframes ``kf`` (sorted frame indices) into clips. + + Each clip starts at a keyframe and ends at a later one, holding as many keyframes as fit + in ``frames_range[1]`` frames; its length is the smallest allowed ``8k + 1`` that covers + them (at least ``frames_range[0]``). + """ + choices = allowed_frame_counts(*frames_range) + lo, hi = choices[0], choices[-1] + out: list[Chunk] = [] + i = 0 + while i < len(kf) - 1: + j = i + 1 + while j + 1 < len(kf) and (kf[j + 1] - kf[i] + 1) <= hi: + j += 1 + span = kf[j] - kf[i] + 1 + start = kf[i] + if span > hi: + # Two consecutive keyframes are further apart than the cap: stretch the clip. + F = next_8k1(span) + if start + F > total: + F = prev_8k1(total - start) + if F < lo: + tc = _tail_chunk(kf, i, start, total - start, log) + if tc is not None: + out.append(tc) + break + log(f" keyframes {i}->{j} are {span} frames apart, above the {hi}-frame cap: " + f"generating {F} frames") + out.append(Chunk(start, F, [i, j], [0, min(kf[j] - start, F - 1)])) + i = j + continue + avail = total - start + F = next((f for f in choices if f >= max(span, lo) and f <= avail), None) + if F is None: + # The video ends before a long-enough clip: keep the start (frame 0 of a clip is + # the substituted keyframe) and pull the last keyframe onto the clip's last frame. + F = max((f for f in choices if f <= avail), default=None) + if F is None or F < lo: + tc = _tail_chunk(kf, i, start, avail, log) + if tc is not None: + out.append(tc) + break + log(f" last clip: {span} -> {F} frames, last keyframe pulled to {F - 1}") + local = [min(k - start, F - 1) for k in kf[i:j + 1]] + out.append(Chunk(start, F, list(range(i, j + 1)), local)) + i = j + return out + + +def drop_colliding_keyframes(kf_local: list[int], num_frames: int) -> list[int]: + """Positions (into ``kf_local``) to keep so no two keyframes share a latent frame. + + Short clips can put two keyframes into the same latent slab (17 frames, keyframes at + [0, 11, 16] -> latents [0, 2, 2]), which the pipeline cannot represent. The first and last + keyframes are kept; colliding interior ones are dropped. + """ + if len(kf_local) <= 2: + return list(range(len(kf_local))) + last = latent_index(kf_local[-1], num_frames) + keep, seen = [], set() + for n, loc in enumerate(kf_local): + li = latent_index(loc, num_frames) + if n < len(kf_local) - 1 and (li in seen or li == last): + continue + seen.add(li) + keep.append(n) + return keep + + +def concat_chunks(chunks: list[Chunk], clips: list[np.ndarray]) -> np.ndarray: + """Join generated clips ``(F, H, W, 3)`` into one video. + + Clip 0 contributes frames ``0..kf_local[-1]``; every later clip starts one frame in, since + its frame 0 is the boundary keyframe the previous clip already ended on. Frames after a + clip's last keyframe are dropped. + """ + parts = [] + for n, (ch, clip) in enumerate(zip(chunks, clips)): + lo = ch.kf_local[0] + (1 if n else 0) + parts.append(clip[lo:ch.kf_local[-1] + 1]) + return np.concatenate(parts, axis=0) diff --git a/robotryon/video/conditioning.py b/robotryon/video/conditioning.py new file mode 100644 index 0000000000000000000000000000000000000000..1d9eaa87f2682536288fe433a2f08649c0cc2f98 --- /dev/null +++ b/robotryon/video/conditioning.py @@ -0,0 +1,150 @@ +"""Assembling one training step of the in-context video model g_phi (paper Eq. 3.3). + +A training sequence holds three groups of tokens, in the order of diffusers' +``LTX2InContextPipeline``:: + + [ target tokens | keyframe tokens | skeleton-flow reference tokens ] + +* **target**: the clip ``w = E_vid(o_{1:T})``, noised by flow matching to ``w_tau``. The model + predicts the velocity ``eps - w`` on these tokens. +* **keyframes** ``o_T``: each keyframe is one pixel frame, VAE-encoded on its own (exactly what + inference can build from a PNG), clean (timestep 0), and placed at its true time by RoPE + coordinates ``[f, f+1) / fps``. The keyframe at frame 0 instead *replaces* latent frame 0 of + the target, which encodes exactly that one frame; it is excluded from the loss. +* **reference**: the rendered skeleton flow ``S_{1:T}``, VAE-encoded like the clip, clean, at + ``1 / ref_downscale`` resolution with its spatial coordinates scaled back onto the target + grid. + +Stereo clips hold both views in every group, each view encoded separately and laid end to end +inside the group; a fourth RoPE axis (rope4d.py) tells the views apart. Mono clips use plain +three-axis coordinates. The transformer is bidirectional, so a keyframe appended at the end +of the sequence conditions every frame. +""" + +from __future__ import annotations + +import torch + +AUDIO_TOKENS = 1 # LTX-2 is audio-video; training feeds a one-token silent audio stream + + +def sample_keyframe_indices(num_frames: int, count_weights: dict[int, float], + include_last_p: float, generator: torch.Generator | None = None + ) -> list[int]: + """Keyframes of one training clip (paper App. B.2.4): 1 to 4, the first frame always. + + The count is drawn from ``count_weights``; the remaining keyframes sit at random frames. + With probability ``include_last_p`` one of them is the clip's last frame, which is where + inference always places the shared boundary keyframe of a chunk. + """ + counts = sorted(count_weights) + probs = torch.tensor([float(count_weights[c]) for c in counts]) + n = min(counts[int(torch.multinomial(probs, 1, generator=generator))], num_frames) + idx = {0} + if n >= 2 and torch.rand((), generator=generator) < include_last_p: + idx.add(num_frames - 1) + pool = torch.randperm(num_frames - 1, generator=generator) + 1 # frames 1..F-1 + for f in pool.tolist(): + if len(idx) >= n: + break + idx.add(f) + return sorted(idx) + + +def _per_view(coords: torch.Tensor, num_views: int) -> list[torch.Tensor]: + if num_views == 1: + return [coords] # mono: no view axis, pretrained RoPE untouched + from .rope4d import add_view_coord + return [add_view_coord(coords, v) for v in range(num_views)] + + +def build_video_coords(rope, num_latent_frames: int, height: int, width: int, + ref_downscale: int, device, fps: float, keyframe_indices: list[int], + num_views: int = 1) -> torch.Tensor: + """RoPE coordinates for ``[target | keyframes | reference]``, matching ``build_step``. + + ``height``/``width`` are one view's latent grid. Keyframe tokens are one pixel frame wide + at ``[f, f+1)`` (already divided by fps, like ``prepare_video_coords``). + """ + tgt = rope.prepare_video_coords(1, num_latent_frames, height, width, device=device, fps=fps) + parts = _per_view(tgt, num_views) + if keyframe_indices: + one = rope.prepare_video_coords(1, 1, height, width, device=device, fps=fps) + for f in keyframe_indices: + c = one.clone() + c[:, 0, :, 0] = float(f) / fps + c[:, 0, :, 1] = float(f + 1) / fps + parts += _per_view(c, num_views) + if ref_downscale == 1: + ref = tgt.clone() + else: + ref = rope.prepare_video_coords(1, num_latent_frames, height // ref_downscale, + width // ref_downscale, device=device, fps=fps) + ref[:, 1:3] *= ref_downscale + parts += _per_view(ref, num_views) + return torch.cat(parts, dim=2) + + +def build_step(target: torch.Tensor, reference: torch.Tensor, + keyframes: dict[int, torch.Tensor], sigma: torch.Tensor, + tokens_per_frame: int, num_views: int = 1): + """Noise the target and lay out one sequence. + + Args: + target: ``(1, V * seq, C)`` packed clean target tokens, view-major. + reference: ``(1, V * ref_seq, C)`` packed skeleton-flow tokens. + keyframes: ``{pixel frame: (1, V * tokens_per_frame, C)}``, one-frame encodes. + sigma: ``(1,)`` flow-matching noise level in [0, 1]. + tokens_per_frame: spatial tokens of one view's latent frame. + + Returns: + ``(hidden, timestep, velocity, pinned, appended)``: the input tokens, per-token + timesteps (``sigma * 1000`` on noised tokens, 0 on clean ones), the velocity target + for the first ``V * seq`` outputs, a mask of target tokens replaced by the frame-0 + keyframe, and the appended keyframe indices in sequence order. + """ + b, seq, _ = target.shape + per_view = seq // num_views + noise = torch.randn_like(target) + s = sigma.view(-1, 1, 1).to(target.dtype) + noisy = (1.0 - s) * target + s * noise + pinned = torch.zeros(b, seq, dtype=torch.bool, device=target.device) + if 0 in keyframes: # frame 0 replaces latent frame 0 of each view + noisy = noisy.clone() + for v in range(num_views): + dst, src = v * per_view, v * tokens_per_frame + noisy[:, dst:dst + tokens_per_frame] = keyframes[0][:, src:src + tokens_per_frame] + pinned[:, dst:dst + tokens_per_frame] = True + appended = sorted(f for f in keyframes if f != 0) + hidden = torch.cat([noisy] + [keyframes[f] for f in appended] + [reference], dim=1) + ts = (sigma.view(-1, 1) * 1000.0).to(target.dtype).expand(b, hidden.shape[1]).clone() + ts[:, :seq] = torch.where(pinned, torch.zeros_like(ts[:, :seq]), ts[:, :seq]) + ts[:, seq:] = 0.0 + return hidden, ts, noise - target, pinned, appended + + +def flow_matching_loss(pred: torch.Tensor, velocity: torch.Tensor, + pinned: torch.Tensor) -> torch.Tensor: + """MSE between predicted and true velocity over the noised target tokens.""" + keep = (~pinned).unsqueeze(-1) + diff = (pred.float() - velocity.float()) ** 2 * keep + return diff.sum() / keep.expand_as(diff).sum().clamp(min=1) + + +def shifted_logit_normal_sigma(seq_len: int, device=None, std: float = 1.0, eps: float = 1e-3, + uniform_prob: float = 0.1) -> torch.Tensor: + """Noise level for one sample, as LTX-2's trainer draws it. + + Adapted from LTX-2 ``ltx_trainer.timestep_samplers.ShiftedLogitNormalTimestepSampler`` + (Apache-2.0): a logit-normal whose mean shifts with the token count, stretched between its + 0.5th and 99.9th percentiles, mixed with 10% uniform samples. + """ + m = (2.05 - 0.95) / (4096 - 1024) + mu = m * seq_len + (0.95 - m * 1024) + x = torch.sigmoid(torch.randn((1,), device=device) * std + mu) + p999 = torch.sigmoid(torch.tensor(mu + 3.0902 * std, device=device)) + p005 = torch.sigmoid(torch.tensor(mu - 2.5758 * std, device=device)) + z = (x - p005) / (p999 - p005) + z = torch.clamp(torch.where(z >= eps, z, 2 * eps - z), 0, 1) + uniform = (1 - eps) * torch.rand((1,), device=device) + eps + return torch.where(torch.rand((1,), device=device) > uniform_prob, z, uniform) diff --git a/robotryon/video/dataset.py b/robotryon/video/dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..d3d4eec0f9ae91ab5434fba89d4fa66ff990edee --- /dev/null +++ b/robotryon/video/dataset.py @@ -0,0 +1,167 @@ +"""Training clips for the video model: robot video, its skeleton flow, and its caption. + +Episodes follow the shared format (``//meta.json``, ``video_.mp4``, +``keypoints_.npz``). One item is a random window of one episode: + +* length ``F`` drawn uniformly from the ``8k + 1`` values in ``frames_range`` that fit, + start drawn uniformly; +* stereo episodes give both views (``V = 2``), except with probability ``mono_p`` where one + random view is used (``V = 1``), so one model learns stereo and mono generation; +* one crop rectangle (bottom-anchored, scale in ``[min_scale, 1]``) and one horizontal flip + for every view and for the skeleton. A mirrored stereo pair also swaps its views, since the + reflected right image is what a left camera would see; +* the skeleton flow ``S_{1:T}`` is rendered on the fly from the keypoints, after the paper's + augmentation (``robotryon.skeleton.SkeletonAugment``: shift every keypoint by up to half its + finger segment, drop the whole skeleton with p=0.2, each keypoint with p=0.2), held fixed + over the clip. The views of a stereo clip are augmented jointly (the hands of both views + share one drop-all decision). + +The skeleton is drawn directly at ``skeleton_size`` (the reference resolution), the same +resolution inference renders at. +""" + +from __future__ import annotations + +import json +import random +from dataclasses import dataclass +from pathlib import Path + +import cv2 +import numpy as np +import torch +from torch.utils.data import Dataset + +from ..skeleton import SkeletonAugment, draw_skeleton_video +from .chunks import allowed_frame_counts +from .io import read_frames + + +@dataclass +class Episode: + path: Path + views: tuple[str, ...] + num_frames: int + caption: str + root_index: int + + +def caption_text(template: str, meta: dict, embodiment: str | None = None) -> str: + """The text prompt: ``template`` filled with the episode's caption (and embodiment).""" + return template.format(caption=meta.get("caption", ""), + embodiment=embodiment or meta.get("embodiment", "")) + + +def load_episodes(roots: list, min_frames: int, fps: float, + caption_template: str = "{caption}") -> tuple[list[Episode], list[float]]: + """All usable episodes under ``roots`` and their sampling weights. + + ``roots`` items are paths or ``{path, weight}``. A root's weight is split over its + episodes in proportion to their length. + """ + episodes, weights = [], [] + for ri, r in enumerate(roots): + path, w = (r, 1.0) if isinstance(r, str) else (r["path"], float(r.get("weight", 1.0))) + found = [] + for m in sorted(Path(path).glob("*/meta.json")): + meta = json.loads(m.read_text()) + if abs(float(meta["fps"]) - fps) > 1e-6: + raise ValueError(f"{m}: fps {meta['fps']} != {fps}; resample the episode first") + if int(meta["num_frames"]) < min_frames: + continue + found.append(Episode(m.parent, tuple(meta["views"]), int(meta["num_frames"]), + caption_text(caption_template, meta), ri)) + if not found: + raise FileNotFoundError(f"no episode with >= {min_frames} frames under {path}") + total = sum(e.num_frames for e in found) + episodes += found + weights += [w * e.num_frames / total for e in found] + return episodes, weights + + +def crop_rect(W: int, H: int, aspect: float, scale: float, rng: random.Random): + """A crop of the given aspect, bottom-anchored (arms enter from the bottom edge).""" + bw, bh = (int(round(H * aspect)), H) if W / H > aspect else (W, int(round(W / aspect))) + w = max(8, int(round(bw * scale))) + h = min(bh, max(8, int(round(w / aspect)))) + return rng.randint(0, max(0, W - w)), H - h, w, h + + +def crop_frames(frames: np.ndarray, rect, flip: bool, size) -> np.ndarray: + x, y, w, h = rect + out = np.empty((len(frames), size[1], size[0], 3), np.uint8) + for i, f in enumerate(frames): + c = f[y:y + h, x:x + w] + out[i] = cv2.resize(c[:, ::-1] if flip else c, size, interpolation=cv2.INTER_AREA) + return out + + +def crop_keypoints(kp: np.ndarray, rect, flip: bool, size) -> np.ndarray: + """Map pixel keypoints through the same crop, flip and resize as the frames.""" + x, y, w, h = rect + out = np.array(kp, np.float32, copy=True) + out[..., 0] = (out[..., 0] - x + 0.5) * size[0] / w - 0.5 + out[..., 1] = (out[..., 1] - y + 0.5) * size[1] / h - 0.5 + if flip: + out[..., 0] = size[0] - 1 - out[..., 0] + return out + + +def to_tensor(a: np.ndarray) -> torch.Tensor: + """``(V, F, H, W, 3)`` uint8 RGB -> ``(V, 3, F, H, W)`` float in [-1, 1].""" + return torch.from_numpy(np.ascontiguousarray(a)).permute(0, 4, 1, 2, 3).float() / 127.5 - 1 + + +class VideoClipDataset(Dataset): + """Random training windows. Items are unbatched dicts (clip lengths differ).""" + + def __init__(self, episodes: list[Episode], out_size, skeleton_size, frames_range, + min_scale: float, flip_p: float, mono_p: float, augment: SkeletonAugment, + draw_kw: dict | None = None, fps: float = 20.0, seed: int = 0): + self.episodes = episodes + self.out_size, self.skeleton_size = tuple(out_size), tuple(skeleton_size) + self.choices = allowed_frame_counts(*frames_range) + self.min_scale, self.flip_p, self.mono_p = min_scale, flip_p, mono_p + self.augment, self.draw_kw = augment, dict(draw_kw or {}) + self.fps, self.seed, self.epoch = fps, seed, 0 + + def set_epoch(self, epoch: int) -> None: + self.epoch = epoch + + def __len__(self) -> int: + return len(self.episodes) + + def __getitem__(self, idx: int) -> dict: + ep = self.episodes[idx] + rng = random.Random(hash((self.seed, self.epoch, idx)) & 0xFFFFFFFF) + nrng = np.random.default_rng(rng.getrandbits(64)) + F = rng.choice([f for f in self.choices if f <= ep.num_frames]) + start = rng.randint(0, ep.num_frames - F) + views = list(ep.views) + if len(views) > 1 and rng.random() < self.mono_p: + views = [rng.choice(views)] + + frames = [read_frames(ep.path / f"video_{v}.mp4", start, F) for v in views] + kps, vis = [], [] + for v in views: + z = np.load(ep.path / f"keypoints_{v}.npz") + kps.append(z["kp"][start:start + F]) + vis.append(z["vis"][start:start + F].astype(bool)) + + H, W = frames[0].shape[1:3] + rect = crop_rect(W, H, self.out_size[0] / self.out_size[1], + rng.uniform(self.min_scale, 1.0), rng) + flip = rng.random() < self.flip_p + video = np.stack([crop_frames(f, rect, flip, self.out_size) for f in frames]) + kp = np.concatenate([crop_keypoints(k, rect, flip, self.skeleton_size) for k in kps], 1) + kp, vs = self.augment(kp, np.concatenate(vis, 1), nrng, temporal=True) + sw, sh = self.skeleton_size + skel = np.stack([draw_skeleton_video(kp[:, 2 * i:2 * i + 2], vs[:, 2 * i:2 * i + 2], + sw, sh, **self.draw_kw) + for i in range(len(views))]) + if flip and len(views) == 2: # a mirrored stereo pair swaps cameras + video, skel = video[::-1], skel[::-1] + return {"video": to_tensor(video), "condition": to_tensor(skel), + "caption": ep.caption, "num_frames": F, "fps": self.fps, + "meta": {"episode": ep.path.name, "start": start, "views": len(views), + "flip": flip}} diff --git a/robotryon/video/io.py b/robotryon/video/io.py new file mode 100644 index 0000000000000000000000000000000000000000..45ab2b18c35150ca04555fc0a0585c7315722af8 --- /dev/null +++ b/robotryon/video/io.py @@ -0,0 +1,52 @@ +"""Minimal mp4 / png I/O for the video model (RGB uint8 arrays, OpenCV underneath).""" + +from __future__ import annotations + +from pathlib import Path + +import cv2 +import numpy as np + + +def read_frames(path: Path, start: int, count: int, size=None) -> np.ndarray: + """Frames ``[start, start + count)`` as ``(count, H, W, 3)`` RGB, optionally resized + (INTER_AREA) to ``size = (W, H)``.""" + cap = cv2.VideoCapture(str(path)) + cap.set(cv2.CAP_PROP_POS_FRAMES, start) + frames = [] + for _ in range(count): + ok, f = cap.read() + if not ok: + break + if size is not None: + f = cv2.resize(f, tuple(size), interpolation=cv2.INTER_AREA) + frames.append(f[..., ::-1]) + cap.release() + if len(frames) < count: + raise RuntimeError(f"{path}: wanted {count} frames from {start}, got {len(frames)}") + return np.ascontiguousarray(np.stack(frames)) + + +def num_frames(path: Path) -> int: + cap = cv2.VideoCapture(str(path)) + n = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + cap.release() + return n + + +def read_image(path: Path) -> np.ndarray: + """PNG -> ``(H, W, 3)`` RGB uint8 at its own resolution.""" + img = cv2.imread(str(path)) + if img is None: + raise FileNotFoundError(path) + return np.ascontiguousarray(img[..., ::-1]) + + +def write_video(frames: np.ndarray, path: Path, fps: float) -> None: + """``(F, H, W, 3)`` RGB uint8 -> mp4 (mp4v).""" + path.parent.mkdir(parents=True, exist_ok=True) + h, w = frames.shape[1:3] + wr = cv2.VideoWriter(str(path), cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h)) + for f in frames: + wr.write(np.ascontiguousarray(f[..., ::-1])) + wr.release() diff --git a/robotryon/video/pipeline.py b/robotryon/video/pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..b673a44d78c28dc830f3340c2cd3fddf58e017bc --- /dev/null +++ b/robotryon/video/pipeline.py @@ -0,0 +1,271 @@ +"""Video synthesis with the in-context LoRA (paper Eq. 3.4 and App. B.4). + +``StereoIcLoraPipeline`` extends diffusers' ``LTX2InContextPipeline`` with two things: + +* **stereo**: ``conditions`` and ``reference_conditions`` are per-view lists. Each view is + prepared by the parent and the token groups are interleaved by role, + ``[base(L) base(R) | kf(L) kf(R) | ref(L) ref(R)]``, so the reference tokens stay at the end + of the sequence where the parent's attention mask expects them. A fourth RoPE axis + (rope4d.py) tells the views apart; mono uses the pipeline's own coordinates. +* **exact keyframe times**: diffusers places a keyframe at a latent index, which snaps it to + pixel frame ``8k - 7``. Training places keyframes at their exact pixel frame ``[f, f+1)``, + so the appended keyframe coordinates are rewritten to the requested pixel frames. The + keyframe at frame 0 replaces latent frame 0 and needs no rewrite. + +``generate`` runs the two-stage recipe on one clip. Both stages use the full LTX-2.5 +transformer with the IC-LoRA and Lightricks' distilled LoRA, the published distilled sigma +schedules (8 steps, then 3 refinement steps from sigma 0.909) and neutral guidance. Both stages +are conditioned on the same keyframes and skeleton flow. ``two_stage.mode``: + +* ``same_res``: stage 1 already at the output resolution; stage 2 refines at that resolution. +* ``upsample``: stage 1 at half the output resolution, the stage-1 latents are upsampled x2 by + the LTX-2.5 spatial latent upsampler, and stage 2 refines at the output resolution. + +The skeleton flow is always given at one ``skeleton_size`` (512x288); each stage encodes it +with ``reference_downscale = stage width / skeleton width`` (1 at 512x288, 2 at 1024x576), +exactly as the two training stages did. +""" + +from __future__ import annotations + +import glob +import os + +import numpy as np +import torch +from diffusers.pipelines.ltx2 import LTX2InContextPipeline + +from .chunks import latent_index +from .rope4d import add_view_coord, install_view_axis_rope, set_view_coord + +MAX_SEQ = 256 # text tokens; training used the same +DISTILLED_LORA = "ltx-2.5-22b-distilled-lora-450-bf16.safetensors" + + +def snapshot_dir(repo_or_path: str) -> str: + """A local directory for a Hugging Face repo id (newest cached snapshot) or a path.""" + if os.path.isdir(repo_or_path): + return repo_or_path + from huggingface_hub import snapshot_download + from huggingface_hub.constants import HF_HUB_CACHE + slug = "models--" + repo_or_path.replace("/", "--") + hits = sorted(glob.glob(os.path.join(HF_HUB_CACHE, slug, "snapshots", "*"))) + return hits[-1] if hits else snapshot_download(repo_or_path) + + +class StereoIcLoraPipeline(LTX2InContextPipeline): + """``LTX2InContextPipeline`` with a view axis and exact-frame keyframes.""" + + num_views = 1 + exact_kf_px: list[int] | None = None # pixel frames of the appended keyframes + + def _unpack_latents(self, latents, num_frames, height, width, patch_size=1, + patch_size_t=1): + """Unpack the views' base tokens into a batch ``(V, C, F, H, W)`` for decoding.""" + per_view = num_frames * height * width + views = latents.shape[1] // per_view + if views < 2: + return super()._unpack_latents(latents, num_frames, height, width, patch_size, + patch_size_t) + return torch.cat([super(StereoIcLoraPipeline, self)._unpack_latents( + latents[:, v * per_view:(v + 1) * per_view], num_frames, height, width, + patch_size, patch_size_t) for v in range(views)], dim=0) + + def _retime_keyframes(self, kf_coords, n_kf, kw): + """Move the appended keyframe blocks to their exact pixel frames ``[f, f+1) / fps``.""" + if not self.exact_kf_px or not n_kf: + return kf_coords + fps = float(kw["frame_rate"]) + tpf = (int(kw["height"]) // 32) * (int(kw["width"]) // 32) + if n_kf != len(self.exact_kf_px) * tpf: + raise ValueError(f"{n_kf} appended keyframe tokens do not match " + f"{len(self.exact_kf_px)} keyframes of {tpf} tokens") + out = kf_coords.clone() + for j, f in enumerate(self.exact_kf_px): + sl = slice(j * tpf, (j + 1) * tpf) + out[:, 0, sl, 0] = float(f) / fps + out[:, 0, sl, 1] = float(f + 1) / fps + return out + + def prepare_latents(self, conditions=None, reference_conditions=None, **kw): + views = self.num_views + conditions = conditions or [None] * views + reference_conditions = reference_conditions or [None] * views + lat_in = kw.pop("latents", None) # stage 2: stage-1 (upsampled) latents per view + + def view_latents(v): + if lat_in is None: + return None + if lat_in.dim() == 5 and lat_in.shape[0] == views: + return lat_in[v:v + 1] + raise ValueError(f"latents {tuple(lat_in.shape)} do not split into {views} views") + + self._single_view_coords = True # reference coords from the parent stay 3-axis + try: + per = [super(StereoIcLoraPipeline, self).prepare_latents( + conditions=conditions[v], reference_conditions=reference_conditions[v], + latents=view_latents(v), **kw) for v in range(views)] + finally: + self._single_view_coords = False + + groups = {k: [[] for _ in range(3)] for k in ("lat", "mask", "clean")} + kf_co, ref_co, ref_masks, ref_counts = [], [], [], [] + for v, (lat, mask, clean, acoords, nref, rmask, rcounts) in enumerate(per): + n_app = 0 if acoords is None else acoords.shape[2] + n_kf, n_base = n_app - nref, lat.shape[1] - n_app + bounds = (slice(0, n_base), slice(n_base, n_base + n_kf), + slice(n_base + n_kf, lat.shape[1])) + for name, src in (("lat", lat), ("mask", mask), ("clean", clean)): + for g, sl in enumerate(bounds): + groups[name][g].append(src[:, sl]) + if acoords is not None: + a4 = set_view_coord(acoords, v) if views > 1 else acoords + kf_co.append(self._retime_keyframes(a4[:, :, :n_kf], n_kf, kw)) + ref_co.append(a4[:, :, n_kf:]) + if rmask is not None: + ref_masks.append(rmask) + ref_counts.extend(rcounts or []) + + def flat(g): + return torch.cat(g[0] + g[1] + g[2], dim=1) + coords = torch.cat(kf_co + ref_co, dim=2) if (kf_co or ref_co) else None + return (flat(groups["lat"]), flat(groups["mask"]), flat(groups["clean"]), coords, + sum(x.shape[1] for x in groups["lat"][2]), + torch.cat(ref_masks, dim=1) if ref_masks else None, ref_counts) + + +def _install_stereo(pipe: StereoIcLoraPipeline) -> None: + """View-axis RoPE, and a base grid that covers every view (view axis stamped per view).""" + install_view_axis_rope(pipe.transformer) + pipe._single_view_coords = False + rope = pipe.transformer.rope + orig = rope.prepare_video_coords + + def widened(*args, **kwargs): + plain = orig(*args, **kwargs) + if pipe._single_view_coords: + return plain + return torch.cat([add_view_coord(plain, v) for v in range(pipe.num_views)], dim=2) + + rope.prepare_video_coords = widened + + +def build_pipeline(lora: str, *, base: str = "Lightricks/LTX-2.5-Diffusers", + transformer_subfolder: str = "transformer_full", views: int = 1, + attention_backend: str | None = "native", upsample: bool = False, + device: str = "cuda") -> StereoIcLoraPipeline: + """Load LTX-2.5 + IC-LoRA + distilled LoRA, without the text encoder. + + Args: + lora: checkpoint directory holding ``pytorch_lora_weights.safetensors`` (or the file). + views: views generated jointly in one sequence (2 = stereo, 1 = mono). + upsample: also load the spatial latent upsampler (``two_stage.mode: upsample``). + """ + from diffusers import LTX2VideoTransformer3DModel + + root = snapshot_dir(base) + pipe = StereoIcLoraPipeline.from_pretrained( + root, transformer=LTX2VideoTransformer3DModel.from_pretrained( + root, subfolder=transformer_subfolder, dtype=torch.bfloat16), + text_encoder=None, tokenizer=None, dtype=torch.bfloat16) + pipe.num_views = int(views) + if views > 1: + _install_stereo(pipe) + lora_dir, name = ((os.path.dirname(lora), os.path.basename(lora)) if os.path.isfile(lora) + else (lora, "pytorch_lora_weights.safetensors")) + pipe.load_lora_weights(lora_dir, weight_name=name, adapter_name="iclora") + if not any("lora" in k for k, _ in pipe.transformer.named_parameters()): + raise RuntimeError(f"no LoRA weights attached from {lora}") + pipe.load_lora_weights(root, weight_name=DISTILLED_LORA, adapter_name="distilled") + pipe.set_adapters(["iclora", "distilled"], [1.0, 1.0]) + if attention_backend: + pipe.transformer.set_attention_backend(attention_backend) + pipe.vae.enable_tiling() + pipe.vae.enable_slicing() + pipe.to(device) + pipe.upsampler = None + if upsample: + from diffusers.pipelines.ltx2 import LTX2LatentUpsamplePipeline + from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel + model = LTX2LatentUpsamplerModel.from_pretrained( + root, subfolder="latent_upsampler", dtype=torch.bfloat16).to(device) + pipe.upsampler = LTX2LatentUpsamplePipeline(vae=pipe.vae, latent_upsampler=model) + return pipe + + +def _u8(img: torch.Tensor) -> np.ndarray: + """``(3, H, W)`` in [-1, 1] -> ``(1, H, W, 3)`` uint8 (keyframe convention).""" + return ((img.clamp(-1, 1) + 1) * 127.5).byte().permute(1, 2, 0)[None].cpu().numpy() + + +def _f01(video: torch.Tensor) -> np.ndarray: + """``(3, F, H, W)`` in [-1, 1] -> ``(F, H, W, 3)`` float in [0, 1] (reference convention).""" + return ((video.clamp(-1, 1) + 1) / 2).permute(1, 2, 3, 0).float().cpu().numpy() + + +def _to_uint8(out: torch.Tensor, views: int) -> torch.Tensor: + """Pipeline ``pt`` output in [0, 1], ``(V, F, 3, H, W)`` -> ``(V, F, H, W, 3)`` uint8.""" + v = out if out.shape[0] == views else out[0] + frames = [x if x.shape[0] == 3 else x.permute(1, 0, 2, 3) for x in v] # -> (3, F, H, W) + return torch.stack([(x.float().clamp(0, 1) * 255).byte().permute(1, 2, 3, 0).cpu() + for x in frames]) + + +@torch.no_grad() +def generate(pipe: StereoIcLoraPipeline, *, skeleton: torch.Tensor, + keyframes: list[tuple[int, torch.Tensor]], text: tuple, num_frames: int, + size: tuple[int, int], fps: float, seed: int = 0, + mode: str = "same_res") -> torch.Tensor: + """Synthesize one clip. + + Args: + skeleton: ``(V, 3, F, h, w)`` skeleton flow in [-1, 1] at the reference resolution. + keyframes: ``[(pixel frame, (V, 3, H, W) in [-1, 1])]`` sorted, first at frame 0. Any + resolution: the pipeline resizes them for each stage. + text: ``(prompt_embeds, mask, negative_embeds, negative_mask)`` (see text.py). + size: output ``(width, height)`` of one view. + mode: ``same_res`` or ``upsample``. + + Returns: + ``(V, F, height, width, 3)`` uint8 RGB. + """ + from diffusers.pipelines.ltx2 import LTX2ReferenceCondition, LTX2VideoCondition + from diffusers.pipelines.ltx2.utils import (DISTILLED_SIGMA_VALUES, + STAGE_2_DISTILLED_SIGMA_VALUES) + if mode not in ("same_res", "upsample"): + raise ValueError(f"two_stage.mode must be same_res or upsample, got {mode!r}") + views = int(skeleton.shape[0]) + if views != pipe.num_views: + raise ValueError(f"{views} views given to a {pipe.num_views}-view pipeline") + px = [p for p, _ in keyframes] + lat = [latent_index(p, num_frames) for p in px] + if px != sorted(set(px)) or len(set(lat)) != len(lat): + raise ValueError(f"keyframes must be sorted and in distinct latent frames: {px}") + pipe.exact_kf_px = [p for p, li in zip(px, lat) if li != 0] + + W, H = size + skel_w = skeleton.shape[-1] + W1, H1 = (W // 2, H // 2) if mode == "upsample" else (W, H) + pe, pm, npe, npm = (x.cuda() for x in text) + kw = dict(prompt_embeds=pe, prompt_attention_mask=pm, negative_prompt_embeds=npe, + negative_prompt_attention_mask=npm, + conditions=[[LTX2VideoCondition(frames=_u8(img[v]), index=li, crf=0) + for li, (_, img) in zip(lat, keyframes)] for v in range(views)], + reference_conditions=[[LTX2ReferenceCondition(frames=_f01(skeleton[v]))] + for v in range(views)], + reference_downscale_factor=W1 // skel_w, max_sequence_length=MAX_SEQ, + height=H1, width=W1, num_frames=num_frames, frame_rate=fps, + generator=torch.Generator("cuda").manual_seed(seed), return_dict=False, + audio_latents=torch.zeros(1, pipe.audio_latent_channels, 1, + pipe.audio_mel_bins // pipe.audio_vae_mel_compression_ratio, + device="cuda", dtype=torch.float32), + guidance_scale=1.0, stg_scale=0.0, modality_scale=1.0, audio_guidance_scale=1.0, + audio_stg_scale=0.0, audio_modality_scale=1.0) + + lat1 = pipe(sigmas=list(DISTILLED_SIGMA_VALUES), output_type="latent", **kw)[0] + if mode == "upsample": + lat1 = pipe.upsampler(latents=lat1, output_type="latent", return_dict=False)[0] + s2 = list(STAGE_2_DISTILLED_SIGMA_VALUES) + kw.update(height=H, width=W, reference_downscale_factor=W // skel_w) + out = pipe(sigmas=s2, latents=lat1, noise_scale=s2[0], output_type="pt", **kw)[0] + return _to_uint8(out, views) diff --git a/robotryon/video/rope4d.py b/robotryon/video/rope4d.py new file mode 100644 index 0000000000000000000000000000000000000000..fb2dc479d8f6245fa398b6cfa75dba2398d6e3aa --- /dev/null +++ b/robotryon/video/rope4d.py @@ -0,0 +1,129 @@ +"""A fourth RoPE axis for the camera view. + +Both stereo views are generated in one sequence: each view is encoded by the VAE on its own +and the token runs are concatenated. To let attention relate the same (t, h, w) position in the +two cameras, every token also carries a view coordinate. + +Passing four axes to the stock LTX-2 rotary embedding would re-cut its channel-to-frequency +map (``dim // (2 * num_axes)`` frequencies per axis) and perturb every pretrained head at once. +Instead the t/h/w frequencies stay exactly where they were, and the view rotation is written +into the last ``VIEW_CHANNELS_PER_HEAD`` channels of every head's slice (64 of 2048 channels), +using the lowest frequency of the table. With two views the coordinates sit a quarter turn +apart. Three-axis coordinates fall through to the original implementation unchanged, so mono +clips see the pretrained embedding exactly. +""" + +from __future__ import annotations + +import types + +import torch + +VIEW_CHANNELS_PER_HEAD = 2 +BASE_VIEWS = 2 + + +def view_channel_indices(dim: int, num_heads: int, per_head: int, device=None) -> torch.Tensor: + """Channels taken by the view axis: the last ``per_head`` of each head's slice, head-major.""" + head_half = (dim // 2) // num_heads + if per_head >= head_half: + raise ValueError(f"per_head={per_head} would take a head's whole {head_half} channels") + return torch.tensor([h * head_half + head_half - per_head + c + for h in range(num_heads) for c in range(per_head)], + dtype=torch.long, device=device) + + +def _forward_with_view_axis(self, coords: torch.Tensor, device=None): + """Drop-in for ``LTX2AudioVideoRotaryPosEmbed.forward`` accepting a 4th (view) axis.""" + device = device or coords.device + if coords.shape[1] != 4: + return self._forward_3axis(coords, device) + if self.rope_type != "split": + raise NotImplementedError(f"view axis needs rope_type='split', got {self.rope_type!r}") + if coords.ndim == 4: # [start, end) -> midpoints + start, end = coords.chunk(2, dim=-1) + coords = ((start + end) / 2.0).squeeze(-1) + spatial, view = coords[:, :3], coords[:, 3] + + # t/h/w exactly as the original: three axes, dim // 6 frequencies each + max_positions = (self.base_num_frames, self.base_height, self.base_width) + grid = torch.stack([spatial[:, i] / max_positions[i] for i in range(3)], dim=-1).to(device) + freqs_dtype = torch.float64 if self.double_precision else torch.float32 + pow_indices = torch.pow(self.theta, torch.linspace(0.0, 1.0, steps=self.dim // 6, + dtype=freqs_dtype, device=device)) + freqs_base = (pow_indices * torch.pi / 2.0).to(dtype=torch.float32) + angles = ((grid.unsqueeze(-1) * 2 - 1) * freqs_base).transpose(-1, -2).flatten(2) + + # pad the angle (cos 1 / sin 0 in the original) to dim // 2 + pad = self.dim // 2 - angles.shape[-1] + if pad < 0: + raise RuntimeError(f"angle width {angles.shape[-1]} exceeds dim//2") + if pad > 0: + angles = torch.cat([torch.zeros_like(angles[:, :, :pad]), angles], dim=-1) + + # the view rotation in the tail channels of every head + idx = view_channel_indices(self.dim, self.num_attention_heads, VIEW_CHANNELS_PER_HEAD, + device=device) + view_angle = ((view.to(device).unsqueeze(-1) / BASE_VIEWS * 2 - 1) + * freqs_base[:VIEW_CHANNELS_PER_HEAD]) + angles = angles.clone() + angles[:, :, idx] = view_angle.repeat(1, 1, self.num_attention_heads).to(angles.dtype) + + b, n = angles.shape[:2] + cos = angles.cos().reshape(b, n, self.num_attention_heads, -1).swapaxes(1, 2) + sin = angles.sin().reshape(b, n, self.num_attention_heads, -1).swapaxes(1, 2) + return cos, sin + + +def install_view_axis_rope(transformer_or_rope) -> None: + """Patch the video rope (only) to accept view coordinates. Idempotent.""" + rope = getattr(transformer_or_rope, "rope", transformer_or_rope) + if getattr(rope, "_view_axis_installed", False): + return + rope._forward_3axis = types.MethodType(type(rope).forward, rope) + rope.forward = types.MethodType(_forward_with_view_axis, rope) + rope._view_axis_installed = True + + +def add_view_coord(coords: torch.Tensor, view: int) -> torch.Tensor: + """``(B, 3, N, 2)`` -> ``(B, 4, N, 2)`` with the view axis constant at ``[view, view+1)``.""" + b, _, n, _ = coords.shape + v = torch.empty(b, 1, n, 2, dtype=coords.dtype, device=coords.device) + v[..., 0] = float(view) + v[..., 1] = float(view + 1) + return torch.cat([coords, v], dim=1) + + +def set_view_coord(coords: torch.Tensor, view: int) -> torch.Tensor: + """Set the view axis, appending it if the coordinates have only three axes.""" + if coords.shape[1] != 4: + return add_view_coord(coords, view) + out = coords.clone() + out[:, 3, :, 0] = float(view) + out[:, 3, :, 1] = float(view + 1) + return out + + +@torch.no_grad() +def check_bands_preserved(rope, device=None, atol: float = 1e-3) -> dict: + """Verify the patched rope changes only the view channels (run before training). + + Channels count as changed when they differ by more than ``atol``: the 3-axis and 4-axis + paths are computed separately and can differ by float rounding on some CPUs.""" + device = device or torch.device("cpu") + base = rope.prepare_video_coords(1, 4, 9, 16, device=device, fps=20.0) + cos3, sin3 = rope._forward_3axis(base, device) + cos4, sin4 = rope(add_view_coord(base, 0)) + idx = view_channel_indices(rope.dim, rope.num_attention_heads, VIEW_CHANNELS_PER_HEAD, + device=cos3.device) + keep = torch.ones(rope.dim // 2, dtype=torch.bool, device=cos3.device) + keep[idx] = False + keep = keep.view(rope.num_attention_heads, -1)[None, :, None, :] + d_cos, d_sin = (cos4 - cos3).abs(), (sin4 - sin3).abs() + changed = ((d_cos > atol) | (d_sin > atol)).any(dim=2).squeeze(0) + return {"n_changed": int(changed.sum()), + "n_expected": rope.num_attention_heads * VIEW_CHANNELS_PER_HEAD, + "heads_touched": int(changed.any(dim=-1).sum()), + "num_heads": rope.num_attention_heads, + "max_abs_diff_off_view": float(torch.max((d_cos * keep).max(), + (d_sin * keep).max()))} diff --git a/robotryon/video/text.py b/robotryon/video/text.py new file mode 100644 index 0000000000000000000000000000000000000000..f653242a732448e12aa50256523cb01fc949799a --- /dev/null +++ b/robotryon/video/text.py @@ -0,0 +1,69 @@ +"""Caption embeddings for the video model, computed once so training and inference never hold +the 23 GB Gemma text encoder next to the transformer. + +Two forms, because the two code paths enter the model at different points: + +* **training** feeds the transformer directly and needs the text *connector outputs* + ``{caption: {"video", "audio", "mask"}}``; +* **inference** calls the diffusers pipeline, which runs the connectors itself, and needs the + *text-encoder outputs* ``{"max_seq", "embeds": {caption: (embeds, mask, neg, neg_mask)}}`` + with the pipeline's default negative prompt. + +Both are encoded with ``max_sequence_length = 256``, the value training used. +""" + +from __future__ import annotations + +import torch + +from .pipeline import MAX_SEQ, StereoIcLoraPipeline, snapshot_dir + + +def _text_pipeline(base: str, device: str): + """The pipeline with only tokenizer and text encoder loaded.""" + pipe = StereoIcLoraPipeline.from_pretrained( + snapshot_dir(base), transformer=None, vae=None, connectors=None, audio_vae=None, + vocoder=None, duration_head=None, dtype=torch.bfloat16) + pipe.text_encoder.to(device) + return pipe + + +@torch.no_grad() +def encode_for_inference(captions: list[str], base: str = "Lightricks/LTX-2.5-Diffusers", + device: str = "cuda") -> dict: + """``{caption: (embeds, mask, negative embeds, negative mask)}``, exactly as the + pipeline's ``__call__`` would compute them.""" + from diffusers.pipelines.ltx2.utils import DEFAULT_NEGATIVE_PROMPT + pipe = _text_pipeline(base, device) + out = {} + for c in dict.fromkeys(captions): + pe, pm, npe, npm = pipe.encode_prompt( + prompt=c, negative_prompt=DEFAULT_NEGATIVE_PROMPT, do_classifier_free_guidance=True, + num_videos_per_prompt=1, max_sequence_length=MAX_SEQ, device=device) + out[c] = tuple(x.cpu() for x in (pe, pm, npe, npm)) + del pipe + torch.cuda.empty_cache() + return out + + +@torch.no_grad() +def encode_for_training(captions: list[str], base: str = "Lightricks/LTX-2.5-Diffusers", + device: str = "cuda") -> dict: + """Connector outputs for the training loop.""" + from diffusers.pipelines.ltx2 import LTX2TextConnectors + pipe = _text_pipeline(base, device) + conn = LTX2TextConnectors.from_pretrained(snapshot_dir(base), subfolder="connectors", + torch_dtype=torch.bfloat16).to(device).eval() + side = getattr(pipe.tokenizer, "padding_side", "left") + out = {} + for c in dict.fromkeys(captions): + pe, pm, _, _ = pipe.encode_prompt(prompt=c, negative_prompt=None, + do_classifier_free_guidance=False, + num_videos_per_prompt=1, max_sequence_length=MAX_SEQ, + device=device) + v, a, m = conn(pe, pm, padding_side=side) + out[c] = {"video": v[0].cpu(), "audio": a[0].cpu(), "mask": m[0].cpu()} + del pipe, conn + torch.cuda.empty_cache() + return out + diff --git a/robotryon/video/training.py b/robotryon/video/training.py new file mode 100644 index 0000000000000000000000000000000000000000..75a25bd0455f05368bcf931e18e7b6ae5ecc0b77 --- /dev/null +++ b/robotryon/video/training.py @@ -0,0 +1,80 @@ +"""Training plumbing for the video LoRA: gradient checkpointing under DDP and LoRA weight I/O. + +diffusers' own (non-reentrant) checkpointing does not reproduce the LTX-2 block on recompute +under DDP, so each transformer block is wrapped in a reentrant checkpoint here. Reentrant +checkpointing needs the block inputs to require grad, hence the hook on ``proj_in``, and pairs +with DDP ``static_graph``. Wrapping inserts ``.block.`` into parameter names; saved LoRA +weights strip it so they load into a plain transformer. +""" + +from __future__ import annotations + +import re + +import torch +from torch import nn + +LORA_TARGETS = ["attn1.to_k", "attn1.to_q", "attn1.to_v", "attn1.to_out.0", + "attn2.to_k", "attn2.to_q", "attn2.to_v", "attn2.to_out.0", + "ff.net.0.proj", "ff.net.2"] # video modules only (official IC-LoRA) + + +class CheckpointedBlock(nn.Module): + def __init__(self, block: nn.Module): + super().__init__() + self.block = block + + def forward(self, *args, **kwargs): + if not self.training or not torch.is_grad_enabled(): + return self.block(*args, **kwargs) + keys = list(kwargs) + + def run(*pos): + return self.block(*pos[:len(args)], **dict(zip(keys, pos[len(args):]))) + return torch.utils.checkpoint.checkpoint(run, *args, *[kwargs[k] for k in keys], + use_reentrant=True) + + +def enable_checkpointing(transformer) -> None: + blocks = transformer.transformer_blocks + for i, b in enumerate(blocks): + if not isinstance(b, CheckpointedBlock): + blocks[i] = CheckpointedBlock(b) + + def hook(_m, _inp, out): + if isinstance(out, torch.Tensor): + out.requires_grad_(True) + transformer.proj_in.register_forward_hook(hook) + transformer.audio_proj_in.register_forward_hook(hook) + + +def lora_state_dict(model) -> dict: + """PEFT LoRA weights with the checkpoint wrapper's ``.block.`` removed.""" + from peft.utils import get_peft_model_state_dict + return {re.sub(r"(transformer_blocks\.\d+)\.block\.", r"\1.", k): v + for k, v in get_peft_model_state_dict(model).items()} + + +def set_lora_state(model, want: dict) -> None: + """Inverse of ``lora_state_dict``: load LoRA weights, re-inserting ``.block.`` if wrapped.""" + from peft import set_peft_model_state_dict + cur = lora_state_dict(model) + if set(want) != set(cur): + raise ValueError(f"LoRA keys differ: {len(set(want) - set(cur))} only in file, " + f"{len(set(cur) - set(want))} only in model (rank or targets changed?)") + wrapped = any(".block." in n for n, _ in model.named_parameters()) + set_peft_model_state_dict(model, {re.sub(r"(transformer_blocks\.\d+)\.", r"\1.block.", k) + if wrapped else k: v for k, v in want.items()}) + + +def load_lora_into(model, path) -> int: + """Initialise the model's LoRA from a saved ``pytorch_lora_weights.safetensors`` (stage 2 + starts from stage 1). Verifies by value that every tensor was loaded.""" + from safetensors.torch import load_file + want = {k.removeprefix("transformer."): v for k, v in load_file(str(path)).items()} + set_lora_state(model, want) + after = lora_state_dict(model) + bad = [k for k in want if not torch.equal(after[k].float().cpu(), want[k].float().cpu())] + if bad: + raise RuntimeError(f"{len(bad)} LoRA tensors did not load, e.g. {bad[0]}") + return len(want) diff --git a/scripts/infer_video.py b/scripts/infer_video.py new file mode 100644 index 0000000000000000000000000000000000000000..c9d159fb9752663d63799d810c2dbb0c1e80990e --- /dev/null +++ b/scripts/infer_video.py @@ -0,0 +1,174 @@ +"""Synthesize the robot video of one human episode (paper Sec. 3.3, Eq. 3.4, App. B.2.4 / B.4). + + python scripts/infer_video.py --config configs/video/infer.yaml \ + input.episode=path/to/human/ep0 input.keyframes_dir=path/to/keyframes/ep0 \ + output.dir=path/to/video/ep0 two_stage.mode=same_res + +Inputs: + /meta.json fps (must equal ``fps``), views, num_frames, caption + /keypoints_.npz retargeted hand keypoints at video resolution + /video_.mp4 the human video (its frame size; its pixels too with + ``mask_human: false``) + /keyframes.json {"indices": [...], ...} written by scripts/infer_keyframes.py + (falls back to /keyframes.json) + /keyframe__.png translated robot keyframe per index + +The skeleton flow S^h_{1:T} is rendered from the keypoints at ``skeleton_size`` on black +(``mask_human: true``, paper default). With ``mask_human: false`` it is drawn on the unmasked human +video instead; the keyframes always come from the masked image condition. The episode is +cut at its keyframes into clips of at most ``frames_range[1]`` frames that share their boundary +keyframe (``robotryon.video.chunks``); each clip is generated in two stages from its keyframes +and skeleton flow, with both stereo views in one sequence by default. + +Outputs in ``output.dir``: + video_.mp4 the robot video, ``resolution`` per view, ``fps`` frames per second, + frames keyframes[0] .. keyframes[-1] of the episode + chunks.json the clip plan (start, length, keyframes of every clip) + meta.json episode metadata in the shared format (for the IDM) +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import torch + +from robotryon.common.config import parse_args, save_config +from robotryon.skeleton import draw_skeleton_video, paint +from robotryon.video.chunks import concat_chunks, drop_colliding_keyframes, plan_chunks +from robotryon.video.dataset import caption_text +from robotryon.video.io import read_image, write_video +from robotryon.video.pipeline import build_pipeline, generate +from robotryon.video.text import encode_for_inference + + +def to_unit(a: np.ndarray) -> torch.Tensor: + """uint8 ``(..., H, W, 3)`` -> float in [-1, 1] with channels moved before the frame axes.""" + t = torch.from_numpy(np.ascontiguousarray(a)).float() / 127.5 - 1 + return t.permute(2, 0, 1) if a.ndim == 3 else t.permute(3, 0, 1, 2) + + +def skeleton_flow(ep: Path, view: str, total: int, size, draw_kw: dict) -> np.ndarray: + """``(T, h, w, 3)`` uint8 skeleton flow of one view at ``size``.""" + import cv2 + cap = cv2.VideoCapture(str(ep / f"video_{view}.mp4")) + W, H = cap.get(cv2.CAP_PROP_FRAME_WIDTH), cap.get(cv2.CAP_PROP_FRAME_HEIGHT) + cap.release() + z = np.load(ep / f"keypoints_{view}.npz") + kp = z["kp"][:total].astype(np.float32).copy() + kp[..., 0] = (kp[..., 0] + 0.5) * size[0] / W - 0.5 + kp[..., 1] = (kp[..., 1] + 0.5) * size[1] / H - 0.5 + return draw_skeleton_video(kp, z["vis"][:total].astype(bool), size[0], size[1], **draw_kw) + + +def human_frames(ep: Path, view: str, total: int, size) -> np.ndarray: + """``(T, h, w, 3)`` uint8 RGB frames of the human video of one view, resized to ``size``.""" + import cv2 + cap = cv2.VideoCapture(str(ep / f"video_{view}.mp4")) + frames = [] + while len(frames) < total: + ok, f = cap.read() + if not ok: + break + frames.append(cv2.cvtColor(cv2.resize(f, tuple(size), interpolation=cv2.INTER_AREA), + cv2.COLOR_BGR2RGB)) + cap.release() + if len(frames) < total: + raise SystemExit(f"{ep / f'video_{view}.mp4'}: {len(frames)} frames, expected {total}") + return np.stack(frames) + + +def reference_flow(ep: Path, view: str, total: int, cfg) -> np.ndarray: + """Reference video of one view: the skeleton flow, on black or on the human video.""" + skel = skeleton_flow(ep, view, total, cfg.skeleton_size, cfg.skeleton.draw) + if cfg.get("mask_human", True): + return skel + human = human_frames(ep, view, total, cfg.skeleton_size) + return np.stack([paint(h, s) for h, s in zip(human, skel)]) + + +def plan(cfg): + """``(episode dir, views, view groups, frame count, keyframes, clips)`` of one episode.""" + ep = Path(cfg.input.episode) + meta = json.loads((ep / "meta.json").read_text()) + if abs(float(meta["fps"]) - cfg.fps) > 1e-6: + raise SystemExit(f"episode is {meta['fps']} fps, the model generates {cfg.fps}: " + f"resample the episode first") + views = list(cfg.input.views or meta["views"]) + groups = [views] if cfg.input.stereo and len(views) == 2 else [[v] for v in views] + total = int(meta["num_frames"]) + kf_json = Path(cfg.input.keyframes_dir) / "keyframes.json" + kf_json = kf_json if kf_json.exists() else ep / "keyframes.json" + kf = sorted(set(int(i) for i in json.loads(kf_json.read_text())["indices"])) + chunks = plan_chunks(kf, total, cfg.frames_range) + if not chunks: + raise SystemExit(f"no clip can be planned from keyframes {kf} in {total} frames") + return ep, meta, views, groups, total, kf, chunks + + +def episode_caption(cfg) -> str: + meta = json.loads((Path(cfg.input.episode) / "meta.json").read_text()) + return caption_text(cfg.text.template, meta, cfg.text.embodiment) + + +def load_pipeline(cfg, views: int, device: str = "cuda"): + return build_pipeline(cfg.model.lora, base=cfg.model.base, + transformer_subfolder=cfg.model.transformer_subfolder, + views=views, attention_backend=cfg.model.attention_backend, + upsample=cfg.two_stage.mode == "upsample", device=device) + + +def synthesize(cfg, pipe, text) -> Path: + """Generate and write the robot video of ``cfg.input.episode`` with a built pipeline and the + encoded caption ``text`` (``robotryon.video.text.encode_for_inference``).""" + ep, meta, views, groups, total, kf, chunks = plan(cfg) + kf_dir, out = Path(cfg.input.keyframes_dir), Path(cfg.output.dir) + out.mkdir(parents=True, exist_ok=True) + print(f"{ep.name}: {total} frames, keyframes {kf}, {len(chunks)} clips, views {groups}, " + f"mode {cfg.two_stage.mode}, mask_human {cfg.get('mask_human', True)}") + skel = {v: reference_flow(ep, v, total, cfg) for v in views} + + clips: dict[str, list[np.ndarray]] = {v: [] for v in views} + for ci, ch in enumerate(chunks): + keep = drop_colliding_keyframes(ch.kf_local, ch.frames) + print(f" clip {ci}: frames {ch.start}..{ch.start + ch.frames - 1} ({ch.frames}f), " + f"keyframes at {[ch.kf_local[k] for k in keep]}", flush=True) + for group in groups: + sk = torch.stack([to_unit(skel[v][ch.start:ch.start + ch.frames]) for v in group]) + kfs = [(ch.kf_local[k], torch.stack([ + to_unit(read_image(kf_dir / f"keyframe_{v}_{kf[ch.kf_slots[k]]:06d}.png")) + for v in group])) for k in keep] + video = generate(pipe, skeleton=sk, keyframes=kfs, text=text, + num_frames=ch.frames, size=tuple(cfg.resolution), fps=cfg.fps, + seed=cfg.seed + ci, mode=cfg.two_stage.mode) + for vi, v in enumerate(group): + clips[v].append(video[vi].numpy()) + + for v in views: + full = concat_chunks(chunks, clips[v]) + write_video(full, out / f"video_{v}.mp4", cfg.fps) + (out / "chunks.json").write_text(json.dumps( + {"keyframes": kf, "chunks": [c.to_dict() for c in chunks]}, indent=1)) + (out / "meta.json").write_text(json.dumps( + {"fps": cfg.fps, "views": views, "num_frames": int(full.shape[0]), + "caption": meta.get("caption", ""), "embodiment": cfg.text.embodiment or "robot", + "source_episode": ep.name, "first_frame": kf[0]}, indent=1)) + save_config(cfg, out / "infer_config.yaml") + print(f"wrote {out}: {len(views)} x {full.shape[0]} frames at " + f"{cfg.resolution[0]}x{cfg.resolution[1]}, {cfg.fps} fps") + return out + + +def main() -> int: + cfg = parse_args(__doc__) + groups = plan(cfg)[3] + caption = episode_caption(cfg) + text = encode_for_inference([caption], cfg.model.base)[caption] + synthesize(cfg, load_pipeline(cfg, len(groups[0])), text) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())