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

from __future__ import annotations

import argparse
import json
import re
import sys
import zipfile
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Iterable

import imageio.v2 as imageio
import numpy as np
import torch
from decord import VideoReader, cpu
from PIL import Image
from tqdm import tqdm

REPO_ROOT = Path(__file__).resolve().parents[2]
RAE_SRC_ROOT = REPO_ROOT / "RAE" / "src"
RAE_DECODER_CONFIG = REPO_ROOT / "RAE" / "configs" / "decoder" / "ViTXL"
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))
if str(RAE_SRC_ROOT) not in sys.path:
    sys.path.append(str(RAE_SRC_ROOT))

from omegaconf import DictConfig, OmegaConf  # noqa: E402
from stage1 import RAE  # noqa: E402
from utils.model_utils import instantiate_from_config  # noqa: E402
from utils.train_utils import parse_configs  # noqa: E402

DEFAULT_RAE_ROOT = Path("/mnt/posttrain/zhaoshitian/models/RAE-collections")
DEFAULT_FLUX2_AE_PATH = Path("/mnt/posttrain/zhaoshitian/models/FLUX.2-klein-4B/vae/ae_bfl_format.safetensors")
DEFAULT_FPS = 8.0


@dataclass(frozen=True)
class BackboneSpec:
    encoder_cls: str
    encoder_config_path: Path
    encoder_param_key: str
    encoder_input_size: int
    decoder_relpath: tuple[str, ...]
    stats_relpath: tuple[str, ...]
    output_suffix: str
    encoder_extra_params: dict[str, Any] = field(default_factory=dict)
    is_flux2_ae: bool = False


@dataclass(frozen=True)
class ResolvedRAEConfig:
    config: DictConfig
    source: str
    backbone: str
    encoder_reference: str
    encoder_input_size: int
    decoder_path: str | None
    stats_path: str | None
    output_suffix: str


BACKBONE_SPECS: dict[str, BackboneSpec] = {
    "dinov2": BackboneSpec(
        encoder_cls="Dinov2withNorm",
        encoder_config_path=Path("/mnt/posttrain/zhaoshitian/models/dinov2-with-registers-base"),
        encoder_param_key="dinov2_path",
        encoder_input_size=224,
        decoder_relpath=("decoders", "dinov2", "wReg_base", "ViTXL_n08", "model.pt"),
        stats_relpath=("stats", "dinov2", "wReg_base", "imagenet1k", "stat.pt"),
        output_suffix="dinov2_wreg_base",
        encoder_extra_params={"normalize": True},
    ),
    "mae": BackboneSpec(
        encoder_cls="MAEwNorm",
        encoder_config_path=Path("/mnt/posttrain/zhaoshitian/models/vit-mae-base"),
        encoder_param_key="model_name",
        encoder_input_size=256,
        decoder_relpath=("decoders", "mae", "base_p16", "ViTXL_n08", "model.pt"),
        stats_relpath=("stats", "mae", "base_p16", "ImageNet1k", "stat.pt"),
        output_suffix="mae_base_p16",
    ),
    "siglip2": BackboneSpec(
        encoder_cls="SigLIP2wNorm",
        encoder_config_path=Path("/mnt/posttrain/zhaoshitian/models/siglip2-base-patch16-256"),
        encoder_param_key="model_name",
        encoder_input_size=256,
        decoder_relpath=("decoders", "siglip2", "base_p16_i256", "ViTXL_n08", "model.pt"),
        stats_relpath=("stats", "siglip2", "base_p16_i256", "ImageNet1k", "stat.pt"),
        output_suffix="siglip2_base_p16_i256",
    ),
    "flux2_ae": BackboneSpec(
        encoder_cls="Flux2AutoEncoder",
        encoder_config_path=DEFAULT_FLUX2_AE_PATH,
        encoder_param_key="weight_path",
        encoder_input_size=256,
        decoder_relpath=(),
        stats_relpath=(),
        output_suffix="flux2_ae",
        is_flux2_ae=True,
    ),
}

ENCODER_CLS_TO_BACKBONE = {
    spec.encoder_cls: backbone
    for backbone, spec in BACKBONE_SPECS.items()
}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Extract UCF101 video latents with local RAE backbones or FLUX.2 AE and optionally decode reconstructions."
    )
    parser.add_argument(
        "--config",
        type=Path,
        default=None,
        help="Optional stage-1 or stage-2 YAML. If omitted, the script builds a stage-1 RAE config from --rae-backbone.",
    )
    parser.add_argument(
        "--rae-backbone",
        type=str,
        choices=sorted(BACKBONE_SPECS),
        default="dinov2",
        help="Backbone/autoencoder used when --config is not supplied.",
    )
    parser.add_argument(
        "--rae-root",
        type=Path,
        default=DEFAULT_RAE_ROOT,
        help="Root directory containing RAE decoder and latent normalization weights.",
    )
    parser.add_argument(
        "--encoder-path",
        type=Path,
        default=None,
        help="Optional local encoder path override, or FLUX.2 AE checkpoint path for --rae-backbone flux2_ae.",
    )
    parser.add_argument(
        "--video-data-dir-path",
        type=Path,
        help="Root directory containing video data and data.json. Videos are resolved relative to this directory.",
    )
    parser.add_argument(
        "--dataset-name",
        type=str,
        default="ucf101",
        help="Dataset name used in default output directory paths (e.g. 'bridgedatav2').",
    )
    parser.add_argument(
        "--data-split",
        type=str,
        default=None,
        help="If set, only load entries whose 'split' field matches (e.g. 'raw', 'scripted_raw').",
    )
    parser.add_argument(
        "--data-category",
        type=str,
        default=None,
        help="If set, only load entries whose 'category' field matches (e.g. 'toykitchen1').",
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=None,
        help="Directory to store latent feature `.npz` files. Defaults to a backbone-specific UCF101 directory.",
    )
    parser.add_argument(
        "--reconstruction-dir",
        type=Path,
        default=None,
        help="Directory to store decoded reconstruction videos. Defaults to a backbone-specific UCF101 directory.",
    )
    parser.add_argument(
        "--image-size",
        type=int,
        default=None,
        help="Center-crop size before encoding. Defaults to the selected encoder's input size.",
    )
    parser.add_argument("--encode-batch-size", type=int, default=16, help="Batch size used for latent extraction.")
    parser.add_argument("--decode-batch-size", type=int, default=16, help="Batch size used for latent decoding.")
    parser.add_argument("--device", type=str, default=None, help="Torch device. Defaults to cuda when available.")
    parser.add_argument("--max-videos", type=int, default=None, help="Optional cap for debugging a subset of videos.")
    parser.add_argument("--overwrite", action="store_true", help="Recompute outputs even when they already exist.")
    parser.add_argument("--skip-reconstruction", action="store_true", help="Only save latent features, skip mp4 reconstructions.")
    return parser.parse_args()


def get_device(explicit: str | None) -> torch.device:
    if explicit:
        return torch.device(explicit)
    return torch.device("cuda" if torch.cuda.is_available() else "cpu")


def center_crop_arr(pil_image: Image.Image, image_size: int) -> Image.Image:
    while min(*pil_image.size) >= 2 * image_size:
        pil_image = pil_image.resize(tuple(x // 2 for x in pil_image.size), resample=Image.BOX)

    scale = image_size / min(*pil_image.size)
    pil_image = pil_image.resize(tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC)

    arr = np.array(pil_image)
    crop_y = (arr.shape[0] - image_size) // 2
    crop_x = (arr.shape[1] - image_size) // 2
    return Image.fromarray(arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size])

def load_video_data_entries(
    data_dir: Path,
    *,
    split: str | None = None,
    category: str | None = None,
) -> list[str]:
    """Load video entries from a ``data.json`` file in *data_dir*.

    The JSON must be an array of objects with a ``"video_path"`` key::

        [
          {
            "id": 0,
            "video_path": "raw/00_2022-07-09_14-27-32_traj10_images0.mp4",
            "split": "raw",
            "category": "toykitchen2"
          },
          ...
        ]

    *video_path* values are treated as relative to *data_dir*.

    Parameters
    ----------
    data_dir:
        Directory containing ``data.json``.
    split:
        If given, only return entries whose ``"split"`` field matches.
    category:
        If given, only return entries whose ``"category"`` field matches.

    Returns
    -------
    list[str]
        A flat list of ``video_path`` strings.
    """
    data_json = data_dir / "data.json"
    if not data_json.is_file():
        raise FileNotFoundError(f"data.json not found in {data_dir}")

    with open(data_json, encoding="utf-8") as fh:
        all_entries: list[dict[str, Any]] = json.load(fh)

    entries: list[str] = []
    for entry in all_entries:
        video_path = entry.get("video_path")
        if not video_path:
            continue
        if split is not None and entry.get("split") != split:
            continue
        if category is not None and entry.get("category") != category:
            continue
        entries.append(video_path)

    return entries


def load_split_entries(split_zip: Path, split_name: str) -> list[str]:
    if not split_zip.exists():
        raise FileNotFoundError(f"Split zip not found: {split_zip}")

    split_member = f"ucfTrainTestlist/{split_name}"
    with zipfile.ZipFile(split_zip) as zf:
        try:
            lines = zf.read(split_member).decode("utf-8").splitlines()
        except KeyError as exc:
            raise KeyError(f"Split file {split_member} not found in {split_zip}") from exc

    entries: list[str] = []
    for line in lines:
        line = line.strip()
        if not line:
            continue
        entries.append(line.split()[0])
    return entries


def validate_video_root(video_root: Path) -> None:
    if video_root.exists():
        return
    raise FileNotFoundError(
        f"Video root not found: {video_root}. Please extract UCF101.rar into class folders before running preprocessing."
    )


def load_video_frames(video_path: Path, image_size: int) -> tuple[np.ndarray, float]:
    vr = VideoReader(str(video_path), ctx=cpu(0))
    fps = DEFAULT_FPS
    try:
        avg_fps = float(vr.get_avg_fps())
        if avg_fps > 0:
            fps = avg_fps
    except Exception:
        pass

    frames = []
    for idx in range(len(vr)):
        frame = vr[idx].asnumpy()
        frame = np.array(center_crop_arr(Image.fromarray(frame), image_size), copy=False)
        frames.append(frame)

    if not frames:
        raise ValueError(f"No frames decoded from {video_path}")

    return np.stack(frames, axis=0), fps


def frames_to_tensor(frames: np.ndarray) -> torch.Tensor:
    tensor = torch.from_numpy(frames).permute(0, 3, 1, 2).float()
    return tensor / 255.0


def batched_indices(length: int, batch_size: int) -> Iterable[tuple[int, int]]:
    for start in range(0, length, batch_size):
        yield start, min(start + batch_size, length)


@torch.no_grad()
def encode_frames(rae: nn.Module, frames: torch.Tensor, batch_size: int, device: torch.device) -> torch.Tensor:
    chunks = []
    for start, end in batched_indices(frames.shape[0], batch_size):
        chunk = frames[start:end].to(device, non_blocking=True)
        chunks.append(rae.encode(chunk).cpu())
    return torch.cat(chunks, dim=0)


@torch.no_grad()
def decode_latents(rae: nn.Module, latents: torch.Tensor, batch_size: int, device: torch.device) -> torch.Tensor:
    chunks = []
    for start, end in batched_indices(latents.shape[0], batch_size):
        chunk = latents[start:end].to(device, non_blocking=True)
        chunks.append(rae.decode(chunk).cpu().clamp(0.0, 1.0))
    return torch.cat(chunks, dim=0)


def save_video(frames: torch.Tensor, save_path: Path, fps: float) -> None:
    save_path.parent.mkdir(parents=True, exist_ok=True)
    if frames.dim() != 4:
        raise ValueError(f"Expected video tensor with 4 dims, got {tuple(frames.shape)}")
    if frames.shape[1] in (1, 3):
        frames = frames.permute(0, 2, 3, 1)
    frames_uint8 = frames.mul(255.0).round().clamp(0, 255).byte().numpy()
    writer = imageio.get_writer(str(save_path), fps=fps)
    try:
        for frame in frames_uint8:
            writer.append_data(frame)
    finally:
        writer.close()


def build_output_paths(feature_root: Path, recon_root: Path, relative_video_path: Path) -> tuple[Path, Path]:
    class_name = relative_video_path.parent.name
    video_stem = relative_video_path.stem
    feature_path = feature_root / class_name / f"{video_stem}_patch_tokens.npz"
    recon_path = recon_root / class_name / f"{video_stem}.mp4"
    return feature_path, recon_path


def validate_existing_path(path: Path, label: str) -> None:
    if not path.exists():
        raise FileNotFoundError(f"{label} not found: {path}")


def sanitize_name(value: str) -> str:
    cleaned = re.sub(r"[^a-z0-9]+", "_", value.lower()).strip("_")
    return cleaned or "custom"


def split_name_to_tag(split_name: str) -> str:
    stem = Path(split_name).stem.lower()
    match = re.fullmatch(r"(train|test)list0*([0-9]+)", stem)
    if match:
        return f"{match.group(1)}_split{int(match.group(2))}"
    return sanitize_name(stem)


def default_feature_root(dataset_name: str, split_name: str, output_suffix: str) -> Path:
    return REPO_ROOT / "video_features" / f"{dataset_name}_{split_name_to_tag(split_name)}_{output_suffix}"


def default_reconstruction_root(dataset_name: str, split_name: str, output_suffix: str) -> Path:
    return REPO_ROOT / "video_reconstructions" / f"{dataset_name}_{split_name_to_tag(split_name)}_{output_suffix}"


def extract_rae_config(full_cfg: DictConfig, source_name: str) -> DictConfig:
    rae_config, *_ = parse_configs(full_cfg)
    if rae_config is not None:
        return rae_config
    if "target" in full_cfg:
        return full_cfg
    raise ValueError(f"Config {source_name} does not contain a stage_1 section.")


def resolve_config_reference(raw_value: str, config_path: Path | None, rae_root: Path) -> str:
    if raw_value.startswith("models/"):
        return str(rae_root / Path(raw_value).relative_to("models"))
    if raw_value.startswith("configs/"):
        return str((REPO_ROOT / "RAE" / raw_value).resolve())
    if raw_value.startswith("./") or raw_value.startswith("../"):
        if config_path is not None:
            return str((config_path.parent / raw_value).resolve())
    return raw_value


def normalize_rae_config_paths(rae_config: DictConfig, *, config_path: Path | None, rae_root: Path) -> DictConfig:
    normalized = OmegaConf.create(OmegaConf.to_container(rae_config, resolve=True))
    params = normalized.get("params")
    if params is None:
        normalized["params"] = {}
        params = normalized["params"]

    for key in ("encoder_config_path", "decoder_config_path", "pretrained_decoder_path", "normalization_stat_path"):
        raw_value = params.get(key)
        if isinstance(raw_value, str):
            params[key] = resolve_config_reference(raw_value, config_path, rae_root)

    encoder_params = params.get("encoder_params")
    if isinstance(encoder_params, (dict, DictConfig)):
        for key in ("dinov2_path", "model_name"):
            raw_value = encoder_params.get(key)
            if isinstance(raw_value, str):
                encoder_params[key] = resolve_config_reference(raw_value, config_path, rae_root)

    return normalized


def apply_encoder_override(rae_config: DictConfig, encoder_path: Path | None) -> DictConfig:
    if encoder_path is None:
        return rae_config

    selected_encoder_path = encoder_path.expanduser()
    validate_existing_path(selected_encoder_path, "Encoder override")

    params = rae_config.get("params")
    if params is None:
        rae_config["params"] = {}
        params = rae_config["params"]

    params["encoder_config_path"] = str(selected_encoder_path)
    encoder_cls = params.get("encoder_cls", "")
    encoder_params = params.get("encoder_params")
    if encoder_params is None:
        params["encoder_params"] = {}
        encoder_params = params["encoder_params"]

    if encoder_cls == "Dinov2withNorm":
        encoder_params["dinov2_path"] = str(selected_encoder_path)
    else:
        encoder_params["model_name"] = str(selected_encoder_path)
    return rae_config


def is_local_path_reference(raw_value: str) -> bool:
    return raw_value.startswith("/") or raw_value.startswith("./") or raw_value.startswith("../")


def validate_local_paths_in_rae_config(rae_config: DictConfig) -> None:
    params = rae_config.get("params", {})
    for key, label in (
        ("encoder_config_path", "encoder config path"),
        ("decoder_config_path", "decoder config path"),
        ("pretrained_decoder_path", "decoder checkpoint"),
        ("normalization_stat_path", "latent normalization stats"),
    ):
        raw_value = params.get(key)
        if isinstance(raw_value, str) and is_local_path_reference(raw_value):
            validate_existing_path(Path(raw_value).expanduser(), label)

    encoder_params = params.get("encoder_params")
    if isinstance(encoder_params, (dict, DictConfig)):
        for key, label in (("dinov2_path", "encoder path"), ("model_name", "encoder path")):
            raw_value = encoder_params.get(key)
            if isinstance(raw_value, str) and is_local_path_reference(raw_value):
                validate_existing_path(Path(raw_value).expanduser(), label)


def build_auto_rae_config(backbone: str, *, rae_root: Path, encoder_path: Path | None) -> ResolvedRAEConfig:
    spec = BACKBONE_SPECS[backbone]
    selected_encoder_path = (encoder_path or spec.encoder_config_path).expanduser()

    if spec.is_flux2_ae:
        validate_existing_path(selected_encoder_path, "FLUX.2 AE checkpoint")
        config = OmegaConf.create(
            {
                "target": "data_processing.flux2_ae.wrapper.Flux2AutoEncoderWrapper",
                "params": {"weight_path": str(selected_encoder_path)},
            }
        )
        return ResolvedRAEConfig(
            config=config,
            source="auto",
            backbone=backbone,
            encoder_reference=str(selected_encoder_path),
            encoder_input_size=spec.encoder_input_size,
            decoder_path=str(selected_encoder_path),
            stats_path=None,
            output_suffix=spec.output_suffix,
        )

    decoder_path = (rae_root / Path(*spec.decoder_relpath)).expanduser()
    stats_path = (rae_root / Path(*spec.stats_relpath)).expanduser()

    validate_existing_path(rae_root, "RAE root")
    validate_existing_path(selected_encoder_path, f"{backbone} encoder")
    validate_existing_path(RAE_DECODER_CONFIG, "RAE decoder config")
    validate_existing_path(decoder_path, f"{backbone} decoder checkpoint")
    validate_existing_path(stats_path, f"{backbone} latent normalization stats")

    encoder_params = dict(spec.encoder_extra_params)
    encoder_params[spec.encoder_param_key] = str(selected_encoder_path)
    config = OmegaConf.create(
        {
            "target": "stage1.RAE",
            "params": {
                "encoder_cls": spec.encoder_cls,
                "encoder_config_path": str(selected_encoder_path),
                "encoder_input_size": spec.encoder_input_size,
                "encoder_params": encoder_params,
                "decoder_config_path": str(RAE_DECODER_CONFIG),
                "pretrained_decoder_path": str(decoder_path),
                "noise_tau": 0.0,
                "reshape_to_2d": True,
                "normalization_stat_path": str(stats_path),
            },
        }
    )
    return ResolvedRAEConfig(
        config=config,
        source="auto",
        backbone=backbone,
        encoder_reference=str(selected_encoder_path),
        encoder_input_size=spec.encoder_input_size,
        decoder_path=str(decoder_path),
        stats_path=str(stats_path),
        output_suffix=spec.output_suffix,
    )


def inspect_rae_config(rae_config: DictConfig, *, source: str) -> ResolvedRAEConfig:
    params = rae_config.get("params", {})
    encoder_cls = params.get("encoder_cls", "")
    backbone = ENCODER_CLS_TO_BACKBONE.get(encoder_cls, "custom")
    if backbone in BACKBONE_SPECS:
        output_suffix = BACKBONE_SPECS[backbone].output_suffix
    else:
        output_suffix = sanitize_name(Path(source).stem if source != "auto" else encoder_cls or "custom")

    encoder_params = params.get("encoder_params", {})
    encoder_reference = params.get("encoder_config_path")
    if isinstance(encoder_params, (dict, DictConfig)):
        encoder_reference = encoder_params.get("dinov2_path", encoder_params.get("model_name", encoder_reference))

    return ResolvedRAEConfig(
        config=rae_config,
        source=source,
        backbone=backbone,
        encoder_reference=str(encoder_reference) if encoder_reference is not None else "<unknown>",
        encoder_input_size=int(params.get("encoder_input_size", 256)),
        decoder_path=params.get("pretrained_decoder_path"),
        stats_path=params.get("normalization_stat_path"),
        output_suffix=output_suffix,
    )


def resolve_rae_config(args: argparse.Namespace) -> ResolvedRAEConfig:
    rae_root = args.rae_root.expanduser()
    if args.config is None:
        return build_auto_rae_config(
            args.rae_backbone,
            rae_root=rae_root,
            encoder_path=args.encoder_path.expanduser() if args.encoder_path is not None else None,
        )

    config_path = args.config.expanduser()
    validate_existing_path(config_path, "Config")
    full_cfg = OmegaConf.load(config_path)
    rae_config = extract_rae_config(full_cfg, str(config_path))
    normalized_rae_config = normalize_rae_config_paths(rae_config, config_path=config_path, rae_root=rae_root)
    normalized_rae_config = apply_encoder_override(
        normalized_rae_config,
        args.encoder_path.expanduser() if args.encoder_path is not None else None,
    )
    validate_local_paths_in_rae_config(normalized_rae_config)
    return inspect_rae_config(normalized_rae_config, source=str(config_path))


def resolve_runtime_args(args: argparse.Namespace, resolved: ResolvedRAEConfig) -> None:
    args.video_root = args.video_data_dir_path.expanduser()
    args.rae_root = args.rae_root.expanduser()
    if args.image_size is None:
        args.image_size = resolved.encoder_input_size
    if args.output_dir is None:
        args.output_dir = default_feature_root(args.dataset_name, args.data_split or "data", resolved.output_suffix)
    else:
        args.output_dir = args.output_dir.expanduser()
    if args.reconstruction_dir is None:
        args.reconstruction_dir = default_reconstruction_root(args.dataset_name, args.data_split or "data", resolved.output_suffix)
    else:
        args.reconstruction_dir = args.reconstruction_dir.expanduser()


def print_run_configuration(args: argparse.Namespace, resolved: ResolvedRAEConfig, device: torch.device) -> None:
    print(f"RAE source: {resolved.source}")
    print(f"RAE backbone: {resolved.backbone}")
    print(f"Encoder path/id: {resolved.encoder_reference}")
    if resolved.decoder_path is not None:
        print(f"Decoder checkpoint: {resolved.decoder_path}")
    if resolved.stats_path is not None:
        print(f"Normalization stats: {resolved.stats_path}")
    print(f"Encoder input size: {resolved.encoder_input_size}")
    print(f"Frame crop size: {args.image_size}")
    print(f"Output features: {args.output_dir}")
    if not args.skip_reconstruction:
        print(f"Output reconstructions: {args.reconstruction_dir}")
    print(f"Device: {device}")


def load_rae(rae_config: DictConfig, device: torch.device) -> nn.Module:
    rae: nn.Module = instantiate_from_config(rae_config).to(device)
    rae.eval()
    return rae


def process_video(
    rae: nn.Module,
    video_path: Path,
    relative_video_path: Path,
    feature_path: Path,
    recon_path: Path,
    args: argparse.Namespace,
    device: torch.device,
) -> None:
    need_features = args.overwrite or not feature_path.exists()
    need_reconstruction = not args.skip_reconstruction and (args.overwrite or not recon_path.exists())
    if not need_features and not need_reconstruction:
        return

    frames_np, fps = load_video_frames(video_path, args.image_size)
    frame_tensor = frames_to_tensor(frames_np)
    latents = encode_frames(rae, frame_tensor, args.encode_batch_size, device)
    timesteps = np.linspace(0.0, 1.0, latents.shape[0], dtype=np.float32)

    if need_features:
        feature_path.parent.mkdir(parents=True, exist_ok=True)
        np.savez(feature_path, features=latents.numpy().astype(np.float32), timesteps=timesteps)

    if need_reconstruction:
        recon_frames = decode_latents(rae, latents, args.decode_batch_size, device)
        save_video(recon_frames, recon_path, fps=fps)


def main() -> None:
    args = parse_args()
    device = get_device(args.device)
    resolved_rae = resolve_rae_config(args)
    resolve_runtime_args(args, resolved_rae)

    validate_video_root(args.video_root)
    args.output_dir.mkdir(parents=True, exist_ok=True)
    if not args.skip_reconstruction:
        args.reconstruction_dir.mkdir(parents=True, exist_ok=True)

    print_run_configuration(args, resolved_rae, device)

    entries = load_video_data_entries(
        args.video_data_dir_path,
        split=args.data_split,
        category=args.data_category,
    )

    if args.max_videos is not None:
        entries = entries[: args.max_videos]

    missing = [entry for entry in entries if not (args.video_root / entry).exists()]
    if missing:
        preview = ", ".join(missing[:5])
        raise FileNotFoundError(f"Missing {len(missing)} videos under {args.video_root}. First missing entries: {preview}")

    rae = load_rae(resolved_rae.config, device)

    progress = tqdm(entries, desc="Extracting video latents")
    for entry in progress:
        relative_video_path = Path(entry)
        video_path = args.video_root / relative_video_path
        feature_path, recon_path = build_output_paths(args.output_dir, args.reconstruction_dir, relative_video_path)
        process_video(rae, video_path, relative_video_path, feature_path, recon_path, args, device)
        progress.set_postfix_str(relative_video_path.name)

    print(f"Saved latent features to {args.output_dir}")
    if not args.skip_reconstruction:
        print(f"Saved reconstruction videos to {args.reconstruction_dir}")


if __name__ == "__main__":
    main()