tmp/video-velocity-model/data_processing/tools/sample_for_fvd.py

"""DDP-shard sampling: load a NanoGPT video flow ckpt, generate rollouts from
first-frame latents of the test split, and dump mp4s + a manifest for FVD.

Pairs with `data_processing/cal_fvd.py`. See `RAE/scripts/run_fvd_nanogpt_flux2.sh`
for the end-to-end launcher.
"""

from __future__ import annotations

import argparse
import json
import math
import sys
import time
import zipfile
from contextlib import nullcontext
from pathlib import Path
from typing import Iterable, List, NamedTuple, Optional

import imageio.v2 as imageio
import numpy as np
import torch
import torch.distributed as dist
from numpy.lib import format as npy_format
from omegaconf import OmegaConf
from torch.cuda.amp import autocast
from tqdm import tqdm

REPO_ROOT = Path(__file__).resolve().parents[1]
RAE_SRC_ROOT = REPO_ROOT / "RAE" / "src"
if str(RAE_SRC_ROOT) not in sys.path:
    sys.path.append(str(RAE_SRC_ROOT))

from stage2.transport.init_video import Sampler, create_transport  # noqa: E402
from utils.dist_utils import cleanup_distributed, setup_distributed  # noqa: E402
from utils.model_utils import instantiate_from_config  # noqa: E402
from utils.train_utils import parse_configs  # noqa: E402


class VideoEntry(NamedTuple):
    class_idx: int
    class_name: str
    npz_path: str
    avi_path: str
    num_frames: int


def _read_npz_features_first_dim(npz_path: str, key: str = "features") -> int:
    """Lazy: read only the inner .npy header inside a .npz to get the leading dim.

    Avoids decoding the full features array (~15 MB per UCF101 npz). Raises
    if the file is corrupt or the key is missing.
    """
    with zipfile.ZipFile(npz_path) as zf:
        with zf.open(f"{key}.npy") as f:
            version = npy_format.read_magic(f)
            if version == (1, 0):
                shape, _, _ = npy_format.read_array_header_1_0(f)
            elif version == (2, 0):
                shape, _, _ = npy_format.read_array_header_2_0(f)
            else:
                raise ValueError(f"Unsupported .npy version {version} in {npz_path}")
    return int(shape[0])


@torch.no_grad()
def decode_video_latents(rae, latents: torch.Tensor, decode_batch_size: int) -> torch.Tensor:
    """Inlined copy of train_video.decode_video_latents to avoid pulling train_video's heavy transitive imports."""
    if latents.dim() != 5:
        raise ValueError(f"Expected video latents shaped [B, T, C, H, W], got {tuple(latents.shape)}")
    batch, num_frames, channels, height, width = latents.shape
    flat = latents.reshape(batch * num_frames, channels, height, width)
    chunks = []
    for start in range(0, flat.shape[0], decode_batch_size):
        decoded = rae.decode(flat[start : start + decode_batch_size])
        chunks.append(decoded.clamp(0.0, 1.0).cpu())
    decoded = torch.cat(chunks, dim=0)
    return decoded.view(batch, num_frames, *decoded.shape[1:])


def save_video(tensor: torch.Tensor, save_path: Path, fps: int = 8) -> None:
    """[T, 3, H, W] float in [0,1] -> mp4. Mirrors train_video.save_video for value_range='0-1'."""
    save_path = Path(save_path)
    save_path.parent.mkdir(parents=True, exist_ok=True)
    if tensor.is_cuda:
        tensor = tensor.cpu()
    if tensor.dim() != 4:
        raise ValueError(f"Expected a 4D video tensor, got {tuple(tensor.shape)}")
    if tensor.shape[1] in (1, 3):
        tensor = tensor.permute(0, 2, 3, 1)
    tensor = tensor.mul(255).round().clamp(0, 255).byte()
    arr = tensor.numpy().astype(np.uint8)
    writer = imageio.get_writer(str(save_path), fps=fps)
    try:
        for frame in arr:
            writer.append_data(frame)
    finally:
        writer.close()


def select_video_list(
    feature_root: Path,
    raw_video_root: Path,
    num_videos: int,
    seed: int,
    min_frames: int = 16,
) -> List[VideoEntry]:
    """Walk feature_root, filter (missing .avi, T<min_frames), then deterministic seeded subsample.

    Identical across all DDP ranks so each rank can compute the same selection
    and shard by `[rank::world_size]` without coordination.
    """
    feature_root = Path(feature_root)
    raw_video_root = Path(raw_video_root)
    class_dirs = sorted([p for p in feature_root.iterdir() if p.is_dir()])
    class_name_to_idx = {p.name: i for i, p in enumerate(class_dirs)}

    candidates: List[VideoEntry] = []
    for class_dir in class_dirs:
        class_name = class_dir.name
        class_idx = class_name_to_idx[class_name]
        for npz in sorted(class_dir.glob("*.npz")):
            stem = npz.stem
            video_stem = stem[: -len("_patch_tokens")] if stem.endswith("_patch_tokens") else stem
            avi = raw_video_root / class_name / f"{video_stem}.avi"
            if not avi.exists():
                continue
            try:
                t = _read_npz_features_first_dim(str(npz))
            except (KeyError, OSError, ValueError, zipfile.BadZipFile):
                continue
            if t < min_frames:
                continue
            candidates.append(VideoEntry(class_idx, class_name, str(npz), str(avi), t))

    if not candidates:
        raise RuntimeError(f"No eligible videos under {feature_root} with matching .avi in {raw_video_root}")

    if num_videos >= len(candidates):
        return candidates

    rng = np.random.RandomState(seed)
    picks = sorted(rng.choice(len(candidates), size=num_videos, replace=False).tolist())
    return [candidates[i] for i in picks]


def load_init_latents(entries: List[VideoEntry]) -> torch.Tensor:
    """Stack frame-0 latents into [B, C, H, W] float32."""
    arrs = []
    for e in entries:
        with np.load(e.npz_path, allow_pickle=False) as z:
            arrs.append(np.asarray(z["features"][0], dtype=np.float32))
    return torch.from_numpy(np.stack(arrs, axis=0))


def load_state_dict_for_inference(model: torch.nn.Module, ckpt_path: str, prefer_ema: bool) -> dict:
    state = torch.load(ckpt_path, map_location="cpu")
    if prefer_ema and isinstance(state.get("ema"), dict) and state["ema"]:
        sd = state["ema"]
        which = "ema"
    elif "model" in state:
        sd = state["model"]
        which = "model"
    else:
        sd = state
        which = "raw"
    sd = {k[len("module.") :] if k.startswith("module.") else k: v for k, v in sd.items()}
    missing, unexpected = model.load_state_dict(sd, strict=False)
    return {"which": which, "missing": list(missing), "unexpected": list(unexpected)}


def build_inference_sampler(cfg, args, device: torch.device):
    rae_cfg, model_cfg, transport_cfg, sampler_cfg, _guidance_cfg, misc_cfg, training_cfg, _eval_cfg = parse_configs(cfg)
    if rae_cfg is None or model_cfg is None:
        raise ValueError("Config must provide both stage_1 and stage_2 sections.")

    if args.disable_label_condition:
        if model_cfg.get("params") is None:
            model_cfg["params"] = {}
        model_cfg["params"]["disable_label_condition"] = True

    training = OmegaConf.to_container(training_cfg, resolve=True) if training_cfg is not None else {}
    sampler_params = dict(OmegaConf.to_container(sampler_cfg, resolve=True).get("params", {}))
    seq_len = int(training.get("seq_len", sampler_params.get("num_frames", 16)))
    if "num_frames" not in sampler_params:
        sampler_params["num_frames"] = seq_len
    if args.num_frames is not None:
        sampler_params["num_frames"] = int(args.num_frames)
    if args.num_ode_steps is not None:
        sampler_params["num_steps"] = args.num_ode_steps

    if args.pi_train_frames is not None:
        nf = int(sampler_params["num_frames"])
        if nf <= 1:
            raise ValueError("--pi-train-frames requires num_frames > 1.")
        scale = float(args.pi_train_frames) / float(nf)
        if model_cfg.get("params") is None:
            model_cfg["params"] = {}
        model_cfg["params"]["rope_t_scale"] = scale

    rae_decode_cfg = OmegaConf.create(OmegaConf.to_container(rae_cfg, resolve=True))
    rae_decode_cfg.setdefault("params", {})
    if rae_decode_cfg.get("target") == "stage1.RAE":
        rae_decode_cfg["params"]["load_encoder"] = False
    rae = instantiate_from_config(rae_decode_cfg).to(device).eval()

    model = instantiate_from_config(model_cfg).to(device).eval()
    load_info = load_state_dict_for_inference(model, args.ckpt, prefer_ema=not args.no_ema)

    misc = OmegaConf.to_container(misc_cfg, resolve=True) if misc_cfg is not None else {}
    latent_size = tuple(int(x) for x in misc.get("latent_size", (128, 16, 16)))
    shift_dim = misc.get("time_dist_shift_dim", math.prod(latent_size))
    shift_base = misc.get("time_dist_shift_base", 4096)
    time_dist_shift = math.sqrt(shift_dim / shift_base)

    transport_params = dict(OmegaConf.to_container(transport_cfg, resolve=True).get("params", {}))
    use_trajectory = bool(transport_params.pop("use_trajectory", False))
    transport_params.pop("time_dist_shift", None)
    if not use_trajectory:
        raise RuntimeError("FVD sampling requires use_trajectory=True in transport config.")
    transport_params.setdefault("latent_channels", int(latent_size[0]))

    sampler_mode = OmegaConf.to_container(sampler_cfg, resolve=True).get("mode", "ODE-video")
    if sampler_mode != "ODE-video":
        raise RuntimeError(f"Unsupported sampler mode for FVD: {sampler_mode}. Expected ODE-video.")

    transport = create_transport(
        **transport_params,
        time_dist_shift=time_dist_shift,
        use_trajectory=True,
    )
    transport_sampler = Sampler(transport)
    eval_sampler = transport_sampler.sample_ode_video(**sampler_params)

    decode_batch_size = int(training.get("decode_batch_size", 16))
    num_frames = int(sampler_params["num_frames"])
    return {
        "model": model,
        "rae": rae,
        "eval_sampler": eval_sampler,
        "decode_batch_size": decode_batch_size,
        "num_frames": num_frames,
        "latent_size": latent_size,
        "load_info": load_info,
    }


def batched(it: List, batch_size: int) -> Iterable[List]:
    for start in range(0, len(it), batch_size):
        yield it[start : start + batch_size]


def make_autocast_ctx(precision: str):
    if precision == "bf16":
        return autocast(dtype=torch.bfloat16)
    if precision == "fp16":
        return autocast(dtype=torch.float16)
    return nullcontext()


def run_sampling_shard(
    shard: List[VideoEntry],
    bundle: dict,
    args: argparse.Namespace,
    device: torch.device,
    rank: int,
) -> List[dict]:
    if not shard:
        return []

    model = bundle["model"]
    rae = bundle["rae"]
    eval_sampler = bundle["eval_sampler"]
    decode_bs = bundle["decode_batch_size"]
    num_frames = bundle["num_frames"]
    gen_root = Path(args.output_dir) / "generated_videos"
    manifest: List[dict] = []

    pbar = tqdm(
        total=len(shard),
        desc=f"rank{rank} sampling",
        position=rank,
        disable=rank != 0 and not args.show_all_progress,
    )

    for batch in batched(shard, args.batch_size):
        init = load_init_latents(batch).to(device)  # [B, C, H, W]
        labels = torch.zeros(len(batch), dtype=torch.long, device=device)
        sample_kwargs = {"y": labels}

        ctx = make_autocast_ctx(args.precision)
        with torch.no_grad(), ctx:
            sampled = eval_sampler(init, model.forward, **sample_kwargs)  # [B, T, C, H, W]

        if sampled.dim() != 5 or sampled.shape[1] != num_frames:
            raise RuntimeError(
                f"Expected sampled tensor shape [B, {num_frames}, C, H, W], got {tuple(sampled.shape)}"
            )

        decoded = decode_video_latents(rae, sampled.float(), decode_bs)  # [B, T, 3, H, W] in [0,1]

        for i, e in enumerate(batch):
            out_path = gen_root / e.class_name / f"{Path(e.npz_path).stem}.mp4"
            save_video(decoded[i], out_path, fps=args.sample_fps)
            manifest.append(
                {
                    "gen_path": str(out_path),
                    "avi_path": e.avi_path,
                    "npz_path": e.npz_path,
                    "class_name": e.class_name,
                    "class_idx": e.class_idx,
                    "num_frames": int(num_frames),
                }
            )
        pbar.update(len(batch))
    pbar.close()
    return manifest


def gather_manifest(rank: int, world_size: int, local: List[dict]) -> Optional[List[dict]]:
    if world_size <= 1:
        return local
    container: List[Optional[List[dict]]] = [None] * world_size
    dist.gather_object(local, container if rank == 0 else None, dst=0)
    if rank != 0:
        return None
    merged: List[dict] = []
    for part in container:
        if part:
            merged.extend(part)
    return merged


def write_outputs(rank: int, manifest: Optional[List[dict]], args: argparse.Namespace, run_meta: dict) -> None:
    if rank != 0 or manifest is None:
        return
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    (output_dir / "manifest.json").write_text(json.dumps(manifest, indent=2))
    (output_dir / "run_config.json").write_text(json.dumps(run_meta, indent=2))


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Sample videos from a NanoGPT-FLUX2-AE ckpt for FVD eval.")
    p.add_argument("--config", type=str, required=True)
    p.add_argument("--ckpt", type=str, required=True)
    p.add_argument("--feature-root", type=Path, required=True)
    p.add_argument("--raw-video-root", type=Path, required=True)
    p.add_argument("--output-dir", type=Path, required=True)
    p.add_argument("--num-videos", type=int, default=2048)
    p.add_argument("--batch-size", type=int, default=8)
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--precision", type=str, choices=["bf16", "fp16", "fp32"], default="bf16")
    p.add_argument("--num-ode-steps", type=int, default=None)
    p.add_argument(
        "--num-frames",
        type=int,
        default=None,
        help="Override sampler.params.num_frames; also tightens select_video_list min_frames.",
    )
    p.add_argument(
        "--pi-train-frames",
        type=int,
        default=None,
        help=(
            "Enable Position Interpolation: scale frame-axis RoPE by "
            "pi_train_frames/num_frames. When unset, no PI is applied "
            "(rope_t_scale=1.0)."
        ),
    )
    p.add_argument("--sample-fps", type=int, default=8)
    p.add_argument("--disable-label-condition", action="store_true")
    p.add_argument("--no-ema", action="store_true", help="Force loading 'model' state even if ckpt has an 'ema' entry.")
    p.add_argument("--show-all-progress", action="store_true")
    return p.parse_args()


def main() -> None:
    args = parse_args()
    if not torch.cuda.is_available():
        raise RuntimeError("FVD sampling requires CUDA.")
    rank, world_size, device = setup_distributed()

    cfg = OmegaConf.load(args.config)

    bundle = build_inference_sampler(cfg, args, device)

    entries = select_video_list(
        args.feature_root,
        args.raw_video_root,
        num_videos=args.num_videos,
        seed=args.seed,
        min_frames=bundle["num_frames"],
    )
    shard = entries[rank::world_size]

    if rank == 0:
        info = bundle["load_info"]
        print(
            f"[fvd-sample] rank={rank} world={world_size} ckpt={args.ckpt} "
            f"loaded={info['which']} missing={len(info['missing'])} unexpected={len(info['unexpected'])}",
            flush=True,
        )
        if info["missing"]:
            print(f"[fvd-sample] missing keys: {info['missing'][:8]}{' ...' if len(info['missing'])>8 else ''}")
        if info["unexpected"]:
            print(f"[fvd-sample] unexpected keys: {info['unexpected'][:8]}{' ...' if len(info['unexpected'])>8 else ''}")
        print(
            f"[fvd-sample] selected {len(entries)} videos "
            f"(requested {args.num_videos}); shard for rank0 = {len(shard)}",
            flush=True,
        )

    t0 = time.time()
    local_manifest = run_sampling_shard(shard, bundle, args, device, rank)
    if dist.is_initialized():
        dist.barrier()
    manifest = gather_manifest(rank, world_size, local_manifest)

    run_meta = {
        "ckpt": args.ckpt,
        "config": args.config,
        "feature_root": str(args.feature_root),
        "raw_video_root": str(args.raw_video_root),
        "num_videos_requested": args.num_videos,
        "num_videos_eligible": len(entries),
        "seed": args.seed,
        "batch_size": args.batch_size,
        "precision": args.precision,
        "num_frames": bundle["num_frames"],
        "num_ode_steps_override": args.num_ode_steps,
        "num_frames_override": args.num_frames,
        "pi_train_frames": args.pi_train_frames,
        "rope_t_scale": (
            float(args.pi_train_frames) / float(bundle["num_frames"])
            if args.pi_train_frames is not None
            else 1.0
        ),
        "sample_fps": args.sample_fps,
        "disable_label_condition": bool(args.disable_label_condition),
        "world_size": world_size,
        "elapsed_sec": time.time() - t0,
        "ckpt_load_info": bundle["load_info"],
    }
    write_outputs(rank, manifest, args, run_meta)

    cleanup_distributed()


if __name__ == "__main__":
    main()