scripts/infer_single_arm_singleview_autoregressive.py

"""
Pure autoregressive rollout for single-arm singleview DreamDojo (DROID dataset).

For each episode:
  1. Take ONLY the first GT frame as conditioning.
  2. Use the full action sequence for the episode.
  3. Generate all frames autoregressively in chunks — each chunk uses the LAST predicted
     frame from the previous chunk (no GT re-anchoring).
  4. Output: one video per episode matching the GT episode length.
"""

import argparse
import json
import sys
from pathlib import Path

import mediapy
import numpy as np
import piq
import torch
import torchvision

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT / "models" / "DreamDojo"))


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoints-dir", type=str, required=True)
    parser.add_argument("--experiment", type=str, default="dreamdojo_2b_480_640_single_arm_sv")
    parser.add_argument("--dataset-path", type=str, required=True)
    parser.add_argument("--save-dir", type=str, required=True)
    parser.add_argument("--num-frames", type=int, default=97,
                        help="Frames per segment (model generates num_frames-1 new frames + 1 conditioning)")
    parser.add_argument("--chunk-size", type=int, default=12)
    parser.add_argument("--num-episodes", type=int, default=None)
    parser.add_argument("--save-fps", type=int, default=15)
    parser.add_argument("--output-dir", type=str, default=None)
    parser.add_argument("--guidance", type=float, default=0.0)

    # WorldCache
    parser.add_argument("--worldcache-enabled", action="store_true")
    parser.add_argument("--worldcache-num-steps", type=int, default=35)
    parser.add_argument("--worldcache-rel-l1-thresh", type=float, default=0.03)
    parser.add_argument("--worldcache-ret-ratio", type=float, default=0.4)
    parser.add_argument("--worldcache-probe-depth", type=int, default=4)
    parser.add_argument("--worldcache-motion-sensitivity", type=float, default=5.0)

    # DiCache
    parser.add_argument("--dicache-enabled", action="store_true")
    parser.add_argument("--dicache-num-steps", type=int, default=35)
    parser.add_argument("--dicache-rel-l1-thresh", type=float, default=0.08)
    parser.add_argument("--dicache-ret-ratio", type=float, default=0.2)
    parser.add_argument("--dicache-probe-depth", type=int, default=2)

    # FasterCache
    parser.add_argument("--fastercache-enabled", action="store_true")
    parser.add_argument("--fastercache-start-step", type=int, default=0)
    parser.add_argument("--fastercache-model-interval", type=int, default=5)
    parser.add_argument("--fastercache-block-interval", type=int, default=3)
    return parser.parse_args()


def build_cache_config_from_args(args):
    """Build a cache strategy config from CLI flags (WorldCache / DiCache / FasterCache)."""
    from methods.cache_strategy.common import WorldCacheConfig, DiCacheConfig, FasterCacheConfig

    if getattr(args, "worldcache_enabled", False):
        return WorldCacheConfig(
            num_steps=args.worldcache_num_steps,
            rel_l1_thresh=args.worldcache_rel_l1_thresh,
            ret_ratio=args.worldcache_ret_ratio,
            probe_depth=args.worldcache_probe_depth,
            motion_sensitivity=args.worldcache_motion_sensitivity,
        )
    if getattr(args, "dicache_enabled", False):
        return DiCacheConfig(
            num_steps=args.dicache_num_steps,
            rel_l1_thresh=args.dicache_rel_l1_thresh,
            ret_ratio=args.dicache_ret_ratio,
            probe_depth=args.dicache_probe_depth,
        )
    if getattr(args, "fastercache_enabled", False):
        return FasterCacheConfig(
            start_step=args.fastercache_start_step,
            model_interval=args.fastercache_model_interval,
            block_interval=args.fastercache_block_interval,
        )
    return None


def build_model(args):
    from cosmos_predict2.action_conditioned_config import ActionConditionedSetupArguments
    from cosmos_predict2.config import MODEL_CHECKPOINTS
    from cosmos_predict2._src.predict2.inference.video2world import Video2WorldInference

    setup_args = ActionConditionedSetupArguments(
        model="2B/robot/action-cond",
        config_file="cosmos_predict2/_src/predict2/action/configs/action_conditioned/config.py",
        checkpoints_dir=args.checkpoints_dir,
        experiment=args.experiment,
        num_frames=args.num_frames,
        dataset_path=args.dataset_path,
        save_dir=args.save_dir,
        output_dir=args.output_dir or args.save_dir,
        num_samples=1,
        data_split="full",
        single_base_index=False,
    )

    checkpoints_dir = Path(args.checkpoints_dir)
    last_checkpoint_file = checkpoints_dir / "latest_checkpoint.txt"
    if not last_checkpoint_file.exists():
        parent_file = checkpoints_dir.parent / "latest_checkpoint.txt"
        if parent_file.exists():
            checkpoints_dir = checkpoints_dir.parent
            last_checkpoint_file = parent_file

    if not last_checkpoint_file.exists():
        raise FileNotFoundError(f"Could not find latest_checkpoint.txt in {args.checkpoints_dir} or its parent.")

    with open(last_checkpoint_file) as f:
        last_checkpoint = f.read().strip()
    checkpoint_iter_dir = checkpoints_dir / last_checkpoint

    from examples.action_conditioned import resolve_checkpoint_path
    checkpoint_path = resolve_checkpoint_path(checkpoint_iter_dir)

    checkpoint = MODEL_CHECKPOINTS[setup_args.model_key]
    experiment = setup_args.experiment or checkpoint.experiment
    cache_config = build_cache_config_from_args(args)

    video2world_cli = Video2WorldInference(
        experiment_name=experiment,
        ckpt_path=checkpoint_path,
        s3_credential_path="",
        context_parallel_size=setup_args.context_parallel_size,
        config_file=setup_args.config_file,
        experiment_opts=[],
        cache_config=cache_config,
    )

    return video2world_cli, checkpoint_iter_dir.name


def build_dataset(args):
    from groot_dreams.dataloader import MultiVideoActionDataset

    # Fixed 13-frame windows (matching the model's state_t / chunk pipeline),
    # so all_steps provides dense base indices up to traj_length-13. The
    # chunk_size-strided plan then walks the full episode. (args.num_frames is
    # kept for CLI compatibility but does not bound the rollout length.)
    dataset = MultiVideoActionDataset(
        num_frames=13,
        dataset_path=args.dataset_path,
        data_split="full",
        single_base_index=False,
        restrict_len=None,
        deterministic_uniform_sampling=False,
    )
    return dataset


def get_episode_plan(dataset, chunk_size):
    """Group dataset indices by episode. Return plan with chunk-strided segment indices.

    Steps through each episode in increments of chunk_size timesteps so the
    autoregressive rollout covers the ENTIRE episode (matching the input video
    length), rather than stopping after a single num_frames window.
    """
    episodes = {}

    global_offset = 0
    for ds_idx, ds in enumerate(dataset.datasets):
        lerobot_ds = ds.lerobot_dataset
        for local_idx, (traj_id, base_index) in enumerate(lerobot_ds.all_steps):
            key = (ds_idx, int(traj_id))
            if key not in episodes:
                episodes[key] = []
            episodes[key].append((global_offset + local_idx, int(base_index)))
        global_offset += len(ds)

    delta_indices = dataset.datasets[0].lerobot_dataset.modality_configs["video"].delta_indices
    timestep_interval = delta_indices[1] - delta_indices[0]
    stride_raw = chunk_size * timestep_interval

    plan = []
    for key, steps in episodes.items():
        ds_idx, traj_id = key
        steps_sorted = sorted(steps, key=lambda x: x[1])
        if not steps_sorted:
            continue

        segment_indices = []
        next_base = 0
        for global_id, base_idx in steps_sorted:
            if base_idx >= next_base:
                segment_indices.append(global_id)
                next_base = base_idx + stride_raw

        if segment_indices:
            traj_length = int(dataset.datasets[ds_idx].lerobot_dataset.trajectory_lengths[
                np.where(dataset.datasets[ds_idx].lerobot_dataset.trajectory_ids == traj_id)[0][0]
            ])
            plan.append({
                "ds_idx": ds_idx,
                "traj_id": int(traj_id),
                "traj_length": traj_length,
                "segment_data_ids": segment_indices,
            })

    return plan


def generate_autoregressive(video2world_cli, first_frame, all_actions, chunk_size, all_lam_video=None, guidance=0):
    """
    Pure autoregressive rollout from a single GT frame.
    Each chunk uses the last predicted frame as conditioning (no GT).
    """
    img_array = first_frame
    chunk_video = []
    first_round = True

    for i in range(0, len(all_actions), chunk_size):
        actions_chunk = all_actions[i: i + chunk_size]
        if actions_chunk.shape[0] != chunk_size:
            break

        current_lam_video = None
        if all_lam_video is not None:
            current_lam_video = all_lam_video[i * 2: (i + chunk_size) * 2]
            if current_lam_video is not None and len(current_lam_video) < chunk_size * 2:
                current_lam_video = None

        if not first_round:
            img_tensor = torchvision.transforms.functional.to_tensor(img_array).unsqueeze(0) * 255.0
        else:
            img_tensor = img_array
        first_round = False

        num_video_frames = actions_chunk.shape[0] + 1
        vid_input = torch.cat(
            [img_tensor, torch.zeros_like(img_tensor).repeat(num_video_frames - 1, 1, 1, 1)], dim=0
        )
        vid_input = vid_input.to(torch.uint8)
        vid_input = vid_input.unsqueeze(0).permute(0, 2, 1, 3, 4)

        video = video2world_cli.generate_vid2world(
            prompt="",
            input_path=vid_input,
            action=torch.from_numpy(actions_chunk).float()
            if isinstance(actions_chunk, np.ndarray)
            else actions_chunk,
            guidance=guidance,
            num_video_frames=num_video_frames,
            num_latent_conditional_frames=1,
            resolution="480,640",
            seed=i,
            negative_prompt="The video captures a scene with low visual quality, blurring, jittering, or distortion.",
            lam_video=current_lam_video,
        )

        video_normalized = (video - (-1)) / (1 - (-1))
        video_clamped = (
            (torch.clamp(video_normalized[0], 0, 1) * 255).to(torch.uint8).permute(1, 2, 3, 0).cpu().numpy()
        )
        img_array = video_clamped[-1]
        chunk_video.append(video_clamped)

    if not chunk_video:
        return None

    chunk_list = [chunk_video[0]] + [
        chunk_video[i][:chunk_size] for i in range(1, len(chunk_video))
    ]
    return np.concatenate(chunk_list, axis=0)


def main():
    args = parse_args()

    from cosmos_oss.init import init_environment, cleanup_environment
    init_environment()
    torch.enable_grad(False)

    print("Building model...")
    video2world_cli, iter_name = build_model(args)

    print("Building dataset...")
    dataset = build_dataset(args)

    print("Planning episodes...")
    plan = get_episode_plan(dataset, args.chunk_size)
    total_episodes = len(plan)
    num_episodes = min(args.num_episodes or total_episodes, total_episodes)
    print(f"Total episodes: {total_episodes}, processing: {num_episodes}")

    save_root = Path(args.save_dir) / iter_name
    save_root.mkdir(parents=True, exist_ok=True)

    all_psnr, all_ssim, all_lpips = [], [], []

    for ep_idx in range(num_episodes):
        ep_info = plan[ep_idx]
        traj_id = ep_info["traj_id"]
        traj_length = ep_info["traj_length"]
        ep_save_dir = save_root / f"episode_{traj_id:06d}"

        if (ep_save_dir / "full_pred.mp4").exists() and (ep_save_dir / "metrics.json").exists():
            print(f"[{ep_idx}/{num_episodes}] episode_{traj_id:06d} already exists, loading metrics.")
            try:
                with open(ep_save_dir / "metrics.json") as f:
                    m = json.load(f)
                if m.get("psnr") is not None:
                    all_psnr.append(m["psnr"])
                    all_ssim.append(m["ssim"])
                    all_lpips.append(m["lpips"])
            except (json.JSONDecodeError, KeyError):
                pass
            continue

        num_segments = len(ep_info["segment_data_ids"])
        print(f"[{ep_idx}/{num_episodes}] traj_id={traj_id}, "
              f"length={traj_length}, segments={num_segments}")

        # Load all segments for this episode
        all_actions_parts = []
        all_lam_parts = []
        gt_segments = []

        for seg_idx, data_id in enumerate(ep_info["segment_data_ids"]):
            sample = dataset[data_id]
            actions = sample["action"][:args.chunk_size]
            if isinstance(actions, torch.Tensor):
                actions = actions.numpy()
            all_actions_parts.append(actions)

            lam = sample.get("lam_video", None)
            if lam is not None:
                all_lam_parts.append(lam[:args.chunk_size * 2])

            gt_seg = sample["video"].permute(1, 2, 3, 0).numpy()
            gt_segments.append(gt_seg)

        # First frame from the first segment (the only GT we use)
        first_sample = dataset[ep_info["segment_data_ids"][0]]
        first_frame = first_sample["video"].transpose(0, 1)[:1]  # (1, C, H, W)

        # Concatenate all actions into one continuous stream
        full_actions = np.concatenate(all_actions_parts, axis=0)

        # Concatenate lam_video if available
        full_lam = None
        if all_lam_parts:
            full_lam = torch.cat(all_lam_parts, dim=0) if isinstance(all_lam_parts[0], torch.Tensor) else None

        # Pure autoregressive rollout from first GT frame only
        pred_video = generate_autoregressive(
            video2world_cli, first_frame, full_actions, args.chunk_size, full_lam,
            guidance=args.guidance,
        )

        if pred_video is None:
            print(f"  Skipping episode {traj_id}: could not generate any frames.")
            continue

        # Build GT by concatenating segments: first segment full, subsequent
        # segments contribute only their first chunk_size frames (segments are
        # chunk_size-strided so the tail would otherwise overlap).
        gt_list = [gt_segments[0]] + [
            gt_segments[i][:args.chunk_size] for i in range(1, len(gt_segments))
        ]
        concat_gt = np.concatenate(gt_list, axis=0)

        # Trim to same length
        min_len = min(len(pred_video), len(concat_gt))
        pred_video = pred_video[:min_len]
        concat_gt = concat_gt[:min_len]

        ep_save_dir.mkdir(parents=True, exist_ok=True)
        mediapy.write_video(str(ep_save_dir / "full_pred.mp4"), pred_video, fps=args.save_fps)
        mediapy.write_video(str(ep_save_dir / "full_gt.mp4"), concat_gt, fps=args.save_fps)
        concat_merged = np.concatenate([concat_gt, pred_video], axis=2)
        mediapy.write_video(str(ep_save_dir / "full_merged.mp4"), concat_merged, fps=args.save_fps)

        # Compute metrics
        x_batch = torch.clamp(torch.from_numpy(pred_video.copy()) / 255.0, 0, 1).permute(0, 3, 1, 2)
        y_batch = torch.clamp(torch.from_numpy(concat_gt.copy()) / 255.0, 0, 1).permute(0, 3, 1, 2)
        try:
            psnr_val = piq.psnr(x_batch, y_batch).mean().item()
            ssim_val = piq.ssim(x_batch, y_batch).mean().item()
            lpips_val = piq.LPIPS()(x_batch, y_batch).mean().item()
        except (RuntimeError, ValueError) as e:
            print(f"  Metrics failed: {e}")
            psnr_val = ssim_val = lpips_val = None

        with open(ep_save_dir / "metrics.json", "w") as f:
            json.dump({
                "psnr": psnr_val, "ssim": ssim_val, "lpips": lpips_val,
                "total_frames_pred": len(pred_video),
                "total_frames_gt": traj_length,
                "trajectory_id": traj_id,
                "mode": "autoregressive",
            }, f, indent=2)

        if psnr_val is not None:
            all_psnr.append(psnr_val)
            all_ssim.append(ssim_val)
            all_lpips.append(lpips_val)
        print(f"  -> {len(pred_video)} frames, PSNR={psnr_val:.2f}, SSIM={ssim_val:.4f}, LPIPS={lpips_val:.4f}")

    if all_psnr:
        summary = {
            "mean_psnr": sum(all_psnr) / len(all_psnr),
            "mean_ssim": sum(all_ssim) / len(all_ssim),
            "mean_lpips": sum(all_lpips) / len(all_lpips),
            "num_episodes_processed": len(all_psnr),
            "mode": "autoregressive",
        }
        with open(save_root / "all_summary.json", "w") as f:
            json.dump(summary, f, indent=2)
        print(f"\n=== Summary ({len(all_psnr)} episodes) ===")
        print(f"PSNR: {summary['mean_psnr']:.3f}")
        print(f"SSIM: {summary['mean_ssim']:.4f}")
        print(f"LPIPS: {summary['mean_lpips']:.4f}")

    cleanup_environment()


if __name__ == "__main__":
    main()