| """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 |
| cond: np.ndarray |
| mask: np.ndarray |
| tokens: torch.Tensor |
|
|
|
|
| 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_<view>.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 |
| 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 |
|
|
