scripts/inference_trex.py

#!/usr/bin/env python3
"""Standalone T-Rex inference (no websocket server).

Loads your LoRA checkpoint, runs causal video+action diffusion on dataset frames,
and saves the *predicted* future video as MP4.

Usage:
    CUDA_VISIBLE_DEVICES=0 python scripts/inference_trex.py \
        --model_path /scratch1/home/zhicao/dreamzero/checkpoints/dreamzero_trex_wan22_lora/checkpoint-5000 \
        --dataset_path /scratch1/home/zhicao/dreamzero/data/trex_small \
        --episode 0 \
        --num_chunks 8 \
        --output_dir results_trex_infer

Why not server/client?
  The official server (socket_test_optimized_AR.py) + test_client_AR.py split
  model hosting from the robot/eval loop for multi-GPU distributed inference.
  For offline debugging you only need GrootSimPolicy + lazy_joint_forward_causal.
"""

from __future__ import annotations

import argparse
import glob
import os
import time

import cv2
import imageio
import numpy as np
import pyarrow.parquet as pq
import torch
import torch._dynamo
import torch.distributed as dist
from einops import rearrange
from tianshou.data import Batch

torch._dynamo.config.disable = True

from groot.vla.data.schema import EmbodimentTag
from groot.vla.data.transform import ComposedModalityTransform
from groot.vla.model.n1_5.sim_policy import GrootSimPolicy

# Modality keys (must match transform_trex / modality_config_trex)
VIDEO_KEYS = [
    "video.head_left",
    "video.left_wrist",
    "video.right_wrist",
]
STATE_KEYS = {
    "state.left_arm": (0, 7),
    "state.left_hand": (7, 29),
    "state.right_arm": (29, 36),
    "state.right_hand": (36, 58),
}
VIDEO_FOLDERS = {
    "video.head_left": "observation.images.head_left",
    "video.left_wrist": "observation.images.left_wrist",
    "video.right_wrist": "observation.images.right_wrist",
}

# Causal chunk schedule (same idea as test_client_AR.py / DROID server)
RELATIVE_OFFSETS = [-23, -16, -8, 0]
ACTION_HORIZON = 24


def get_expected_video_resolution(policy: GrootSimPolicy) -> tuple[int, int]:
    """Return (height, width) that eval_transform VideoToTensor expects."""
    cfg = policy.trained_model.action_head.config
    target_h = getattr(cfg, "target_video_height", None)
    target_w = getattr(cfg, "target_video_width", None)
    if target_h is not None and target_w is not None:
        return int(target_h), int(target_w)

    eval_transform = getattr(policy, "eval_transform", None)
    if isinstance(eval_transform, ComposedModalityTransform):
        for t in eval_transform.transforms:
            res = getattr(t, "original_resolutions", None)
            if res:
                w, h = next(iter(res.values()))
                return int(h), int(w)
    return 160, 320


def resize_frames(frames: np.ndarray, target_h: int, target_w: int) -> np.ndarray:
    """Resize (H,W,C) or (T,H,W,C) uint8 frames to (target_h, target_w)."""
    if frames.ndim == 3:
        if (frames.shape[0], frames.shape[1]) == (target_h, target_w):
            return frames
        return cv2.resize(frames, (target_w, target_h), interpolation=cv2.INTER_LINEAR)
    return np.stack(
        [cv2.resize(f, (target_w, target_h), interpolation=cv2.INTER_LINEAR) for f in frames],
        axis=0,
    )


class TrexEpisode:
    """One T-Rex episode from LeRobot v2 layout."""

    def __init__(self, dataset_root: str, episode_index: int):
        pq_path = os.path.join(
            dataset_root,
            "data",
            f"chunk-{episode_index // 1000:03d}",
            f"episode_{episode_index:06d}.parquet",
        )
        if not os.path.isfile(pq_path):
            raise FileNotFoundError(pq_path)
        self.table = pq.read_table(pq_path)
        self.length = self.table.num_rows
        self.episode_index = episode_index
        self.root = dataset_root

        self.video_dirs = {}
        for key, folder in VIDEO_FOLDERS.items():
            pattern = os.path.join(
                dataset_root,
                "videos",
                "**",
                folder,
                f"episode_{episode_index:06d}.mp4",
            )
            hits = sorted(glob.glob(pattern, recursive=True))
            if not hits:
                raise FileNotFoundError(f"No video for {key}: {pattern}")
            self.video_dirs[key] = hits[0]

        print(
            f"TrexEpisode {episode_index}: {self.length} steps, "
            f"{len(self.video_dirs)} cameras"
        )

    def get_task(self, row: int) -> str:
        try:
            return str(self.table.column("annotation.task")[row].as_py())
        except Exception:
            return ""

    def get_state(self, row: int) -> np.ndarray:
        return np.array(self.table.column("observation.state")[row].as_py(), dtype=np.float64)

    def get_frame(self, row: int, video_key: str) -> np.ndarray:
        cap = cv2.VideoCapture(self.video_dirs[video_key])
        cap.set(cv2.CAP_PROP_POS_FRAMES, row)
        ok, frame = cap.read()
        cap.release()
        if not ok:
            raise RuntimeError(f"Failed frame {row} from {self.video_dirs[video_key]}")
        return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)


def build_chunk_schedule(total_frames: int, num_chunks: int) -> list[list[int]]:
    """First chunk: 1 frame; later chunks: 4 frames ending at anchor."""
    chunks: list[list[int]] = []
    anchor = 0
    for i in range(num_chunks):
        if i == 0:
            indices = [0]
        else:
            indices = [max(anchor + off, 0) for off in RELATIVE_OFFSETS]
        if indices[-1] >= total_frames:
            break
        chunks.append(indices)
        anchor += ACTION_HORIZON
    return chunks


def build_obs(
    episode: TrexEpisode,
    frame_indices: list[int],
    prompt: str,
    video_height: int,
    video_width: int,
) -> dict:
    obs: dict = {}
    anchor = frame_indices[-1]
    state = episode.get_state(anchor)

    for key in VIDEO_KEYS:
        frames = np.stack([episode.get_frame(i, key) for i in frame_indices], axis=0)
        frames = resize_frames(frames.astype(np.uint8), video_height, video_width)
        if len(frame_indices) == 1:
            obs[key] = frames[0]  # (H, W, 3)
        else:
            obs[key] = frames  # (T, H, W, 3)

    for key, (s, e) in STATE_KEYS.items():
        obs[key] = state[s:e].reshape(1, -1).astype(np.float64)

    obs["annotation.task"] = prompt
    return obs


def decode_video_latents(policy: GrootSimPolicy, video_chunks: list[torch.Tensor]) -> np.ndarray:
    """Concat latent chunks along time, VAE decode -> (T, H, W, 3) uint8."""
    if not video_chunks:
        raise ValueError("No video chunks to decode")
    cat = torch.cat(video_chunks, dim=2)
    ah = policy.trained_model.action_head
    frames = ah.vae.decode(
        cat,
        tiled=ah.tiled,
        tile_size=(ah.tile_size_height, ah.tile_size_width),
        tile_stride=(ah.tile_stride_height, ah.tile_stride_width),
    )
    frames = rearrange(frames, "B C T H W -> B T H W C")[0]
    frames = ((frames.float() + 1) * 127.5).clip(0, 255).cpu().numpy().astype(np.uint8)
    return frames


def save_mp4(path: str, frames: np.ndarray, fps: int = 5) -> None:
    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
    imageio.mimsave(path, list(frames), fps=fps, codec="libx264")
    print(f"Saved {len(frames)} frames -> {path}")


def run(args: argparse.Namespace) -> None:
    if not dist.is_initialized():
        os.environ.setdefault("MASTER_ADDR", "localhost")
        os.environ.setdefault("MASTER_PORT", "29500")
        dist.init_process_group(backend="gloo", world_size=1, rank=0)

    print(f"Loading checkpoint: {args.model_path}")
    policy = GrootSimPolicy(
        embodiment_tag=EmbodimentTag.TREX,
        model_path=args.model_path,
        device=args.device,
    )
    print("Model loaded.")
    video_height, video_width = get_expected_video_resolution(policy)
    print(f"Resizing input video to {video_height}x{video_width} (HxW) for eval_transform")

    episode = TrexEpisode(args.dataset_path, args.episode)
    prompt = episode.get_task(0) if args.use_dataset_prompt else args.prompt
    schedule = build_chunk_schedule(episode.length, args.num_chunks)
    print(f"Prompt: {prompt!r}")
    print(f"Running {len(schedule)} causal chunks: {schedule}")

    os.makedirs(args.output_dir, exist_ok=True)
    video_chunks: list[torch.Tensor] = []
    times = []

    for ci, frame_indices in enumerate(schedule):
        obs = build_obs(episode, frame_indices, prompt, video_height, video_width)
        t0 = time.perf_counter()
        with torch.inference_mode():
            result, video_pred = policy.lazy_joint_forward_causal(Batch(obs=obs))
        elapsed = time.perf_counter() - t0
        times.append(elapsed)
        video_chunks.append(video_pred)

        act = result.act
        print(
            f"  chunk {ci:02d} frames={frame_indices} "
            f"infer={elapsed:.2f}s start_frame={policy.trained_model.action_head.current_start_frame}"
        )
        if ci == 0:
            print(f"    action keys: {[k for k in dir(act) if k.startswith('action.')]}")

    pred_frames = decode_video_latents(policy, video_chunks)
    out_pred = os.path.join(
        args.output_dir,
        f"ep{args.episode:06d}_pred.mp4",
    )
    save_mp4(out_pred, pred_frames, fps=args.fps)

    if args.save_input_clip:
        # Save the conditioning frames (head camera) for reference
        input_frames = []
        max_idx = min(schedule[-1][-1] + 1, episode.length)
        for i in range(max_idx):
            input_frames.append(episode.get_frame(i, "video.head_left"))
        save_mp4(
            os.path.join(args.output_dir, f"ep{args.episode:06d}_input_head_left.mp4"),
            np.stack(input_frames, axis=0),
            fps=args.fps,
        )

    print(f"Avg inference time per chunk: {np.mean(times):.2f}s")
    print(f"Done. Output dir: {os.path.abspath(args.output_dir)}")


def main() -> None:
    p = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    p.add_argument("--model_path", required=True, help="checkpoint-5000 directory")
    p.add_argument(
        "--dataset_path",
        default="/scratch1/home/zhicao/dreamzero/data/trex_small",
    )
    p.add_argument("--episode", type=int, default=0)
    p.add_argument("--num_chunks", type=int, default=8,
                   help="Number of causal chunks (more -> longer predicted video)")
    p.add_argument("--prompt", default="perform the task")
    p.add_argument("--use_dataset_prompt", action="store_true")
    p.add_argument("--device", default="cuda:0")
    p.add_argument("--output_dir", default="results_trex_infer")
    p.add_argument("--fps", type=int, default=5)
    p.add_argument("--save_input_clip", action="store_true",
                   help="Also save input head_left frames for comparison")
    run(p.parse_args())


if __name__ == "__main__":
    main()