scripts/infer_bimanual_singleview_single_segment.py

"""
Single-segment full-length autoregressive rollout for bimanual singleview DreamDojo.

For each episode:
  - Start from the first GT frame.
  - Rollout autoregressively (chunk_size actions per step) for the ENTIRE episode
    length — no GT reset between segments. Pure autoregressive.
  - GT video and actions come from the dataset pipeline (properly normalized/resized).

Output: full_gt.mp4, full_pred.mp4, full_merged.mp4, metrics.json.
"""

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_aloha")
    parser.add_argument("--dataset-path", type=str, required=True,
                        help="Comma-separated list of task directories")
    parser.add_argument("--save-dir", type=str, required=True)
    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=10)
    parser.add_argument("--output-dir", type=str, default=None)
    parser.add_argument("--guidance", type=float, default=0)
    parser.add_argument("--save-video-only", action="store_true", default=True,
                        help="Only save the generated prediction video (no GT/merged/metrics). Default on.")
    parser.add_argument("--save-full", dest="save_video_only", action="store_false",
                        help="Also save full_gt.mp4, full_merged.mp4 and metrics.json.")

    # 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)

    # 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)

    # 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)

    return parser.parse_args()


def build_cache_config_from_args(args):
    """Build cache config from CLI args (mirrors infer_humanoid_singleview_full_episode.py)."""
    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, "fastercache_enabled", False):
        return FasterCacheConfig(
            start_step=args.fastercache_start_step,
            model_interval=args.fastercache_model_interval,
            block_interval=args.fastercache_block_interval,
        )
    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,
        )
    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=13,
        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

    paths = [p.strip() for p in args.dataset_path.split(",") if p.strip()]
    valid_paths = [p for p in paths if list(Path(p).glob("data/*/*.parquet"))]

    dataset = MultiVideoActionDataset(
        num_frames=13,
        dataset_path=valid_paths,
        data_split="full",
        single_base_index=False,
        restrict_len=None,
        deterministic_uniform_sampling=False,
    )
    # ds_idx -> task name (dir basename), aligned with MultiVideoActionDataset order.
    task_names = [Path(p).name for p in valid_paths]
    return dataset, task_names


def get_episode_plan(dataset, chunk_size, task_names=None):
    """
    For each episode, collect data_ids stepping through by chunk_size.
    Each data_id provides chunk_size normalized actions via the dataset pipeline.
    """
    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]
            ])
            task_name = task_names[ds_idx] if task_names else None
            plan.append({
                "ds_idx": ds_idx,
                "task_name": task_name,
                "traj_id": int(traj_id),
                "traj_length": traj_length,
                "timestep_interval": int(timestep_interval),
                "segment_data_ids": segment_indices,
            })

    return plan


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, task_names = build_dataset(args)

    print("Planning episodes...")
    plan = get_episode_plan(dataset, args.chunk_size, task_names)
    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"]
        task_name = ep_info.get("task_name")
        # Match DreamGen singleview naming: {task}__episode_{id:06d} (unique across tasks).
        if task_name:
            ep_dir_name = f"{task_name}__episode_{traj_id:06d}"
        else:
            ep_dir_name = f"episode_{traj_id:06d}"
        ep_save_dir = save_root / ep_dir_name

        if (ep_save_dir / "full_pred.mp4").exists():
            print(f"[{ep_idx}] {ep_dir_name} already exists, skipping.")
            continue

        num_chunks = len(ep_info["segment_data_ids"])
        print(f"[{ep_idx}] {ep_dir_name} traj_id={traj_id}, traj_length={ep_info['traj_length']}, "
              f"chunks={num_chunks}")

        if num_chunks == 0:
            print("  No chunks, skipping.")
            continue

        ep_save_dir.mkdir(parents=True, exist_ok=True)

        # Get first frame from first segment
        first_sample = dataset[ep_info["segment_data_ids"][0]]
        img_array = first_sample["video"].transpose(0, 1)[:1]  # (1, C, H, W)

        gt_frames = []
        chunk_videos = []
        first_round = True

        for chunk_idx, data_id in enumerate(ep_info["segment_data_ids"]):
            sample = dataset[data_id]

            video_tensor = sample["video"]
            gt_video_chunk = video_tensor.permute(1, 2, 3, 0).numpy()
            gt_frames.append(gt_video_chunk)

            actions = sample["action"][:args.chunk_size]
            if isinstance(actions, torch.Tensor):
                actions = actions.numpy()

            if actions.shape[0] != args.chunk_size:
                print(f"    chunk {chunk_idx}: only {actions.shape[0]} actions (need {args.chunk_size}), stopping.")
                break

            lam_video = sample.get("lam_video", None)
            current_lam_video = None
            if lam_video is not None and len(lam_video) >= args.chunk_size * 2:
                current_lam_video = lam_video[:args.chunk_size * 2]

            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.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).float()
                if isinstance(actions, np.ndarray)
                else actions,
                guidance=args.guidance,
                num_video_frames=num_video_frames,
                num_latent_conditional_frames=1,
                resolution="480,640",
                seed=chunk_idx,
                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()
            )

            # Pure autoregressive: use last predicted frame as next input
            img_array = video_clamped[-1]
            chunk_videos.append(video_clamped)

            print(f"    chunk {chunk_idx+1}/{num_chunks} done")

        if not chunk_videos:
            continue

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

        if args.save_video_only:
            # Trim prediction to the real GT episode length in the DOWNSAMPLED frame
            # space (dreamdojo samples every `timestep_interval` raw frames), so the
            # generated clip covers ~the full episode like DreamGen does.
            interval = max(1, int(ep_info.get("timestep_interval", 1)))
            expected_frames = -(-int(ep_info["traj_length"]) // interval)  # ceil div
            gt_total = min(len(concat_pred), expected_frames)
            concat_pred = concat_pred[:gt_total]
            mediapy.write_video(str(ep_save_dir / "full_pred.mp4"), concat_pred, fps=args.save_fps)
            print(f"  -> saved {len(concat_pred)} frames (video only)")
            continue

        gt_list = [gt_frames[0]] + [
            gt_frames[i][:args.chunk_size] for i in range(1, len(gt_frames))
        ]
        concat_gt = np.concatenate(gt_list, axis=0)

        min_len = min(len(concat_pred), len(concat_gt))
        concat_pred = concat_pred[:min_len]
        concat_gt = concat_gt[:min_len]

        mediapy.write_video(str(ep_save_dir / "full_pred.mp4"), concat_pred, 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, concat_pred], axis=2)
        mediapy.write_video(str(ep_save_dir / "full_merged.mp4"), concat_merged, fps=args.save_fps)

        x_batch = torch.clamp(torch.from_numpy(concat_pred.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)
        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()

        with open(ep_save_dir / "metrics.json", "w") as f:
            json.dump({
                "psnr": psnr_val, "ssim": ssim_val, "lpips": lpips_val,
                "num_chunks": len(chunk_videos),
                "total_frames_pred": len(concat_pred),
                "total_frames_gt": ep_info["traj_length"],
                "trajectory_id": traj_id,
                "task_name": task_name,
                "mode": "single_segment_full_rollout",
            }, f, indent=2)

        all_psnr.append(psnr_val)
        all_ssim.append(ssim_val)
        all_lpips.append(lpips_val)
        print(f"  -> {len(concat_pred)} 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": "single_segment_full_rollout",
        }
        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()