''' BVH, Apr 2026. Any4D inference script for anydata datasets -- barebones single-dataset variant. Supports anydata (Webbed) datasets, visualizations, uncertainty heatmaps. Optionally overrides a single camera's RGB frames in-place from a video file (gradio-style), without introducing a second dataloader. For full donor-camera + donor-RGB splicing via a second dataset, use infer_anydata_splice.py. For legacy vidar datasets, use infer_vidar.py. ''' import os import sys sys.path.insert(0, os.getcwd()) # Library imports import argparse import copy import json import numpy as np import time import torch import traceback import warnings from rich import print from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TimeRemainingColumn from torch.utils.data import DataLoader # Optional imports try: import lpips LPIPS_AVAILABLE = True except ImportError: print('[yellow]Warning: lpips not available, metrics calculation will be disabled') LPIPS_AVAILABLE = False try: import skimage.metrics SKIMAGE_AVAILABLE = True except ImportError: print('[yellow]Warning: skimage not available, metrics calculation will be disabled') SKIMAGE_AVAILABLE = False # Internal imports from custom.dataloader.collate import any4d_collate from custom.dataset.anydata_dataset import AnyDataset from custom.eval import infer_utils # Shared utilities (also used by gradio_app.py) np.set_printoptions(precision=3, suppress=True) warnings.filterwarnings('ignore', category=FutureWarning) def add_base_args(parser): '''Register CLI flags shared by infer_anydata.py and infer_anydata_splice.py.''' # Resource options parser.add_argument('--device', type=str, default='cuda', help='Device to use (cuda or cpu)') parser.add_argument('--gpu_id', type=int, default=0, help='GPU device ID to use') # Debug options parser.add_argument('--debug', action='store_true', help='Enable debug mode (single-threaded, no distributed init)') # Experiment & model options parser.add_argument('--exp_cfg', type=str, required=True, help='Path to experiment config file (e.g., custom/experiment/examples_anydata/rdvs4_robot.py)') parser.add_argument('--ckpt_path', type=str, default=None, help='Path to model weights (overrides experiment config if provided)') # Dataset options parser.add_argument('--dset_cfg', type=str, required=True, help='Path to anydata dataset yaml config file') parser.add_argument('--batch_size', type=int, default=1, help='Batch size for dataloader') parser.add_argument('--stop_after', type=int, default=-1, help='If > 0, only evaluate this many scenes before finishing. Default -1 means process all.') parser.add_argument('--shard', type=str, default=None, help='Process only a modulo subset of scenes for multi-GPU parallelism, ' 'format "i/n" (e.g. "0/2"); scenes with (index mod n == i) are run. ' 'Default None = all. Each shard writes its own _aggregated_metrics.json; ' 'merge them by sample_id afterwards.') parser.add_argument('--data_overrides', type=str, default=None, help='JSON dict of dataset overrides (e.g. \'{"length": 33, "resize": [-16, 480]}\')') parser.add_argument('--data_defaults', type=str, default=None, help='JSON dict of dataset defaults (e.g. \'{"frame_rate": 10.0}\')') parser.add_argument('--exp_overrides', type=str, default=None, help='JSON dict of any4d_config overrides applied before build_model ' '(e.g. \'{"block_gate_fix": false}\')') parser.add_argument('--metric_resolution', type=str, default=None, help='If set, resize pred AND gt to a common resolution before PSNR/SSIM/LPIPS, ' 'using the anydata augmentation.resize spec (e.g. \'[-16, 512]\' = longest ' 'side 512, multiple of 16). Default None = score at native resolution.') parser.add_argument('--infer_metrics', type=int, default=0, help='Compute the separate in-code metrics path (calculate_metrics_per_scene + ' '_aggregated_metrics.json + per-scene metrics_*.json). Default 0 = off; ' 'info.json already carries the per-view 11-view metrics regardless. ' 'metric_resolution applies ONLY to this path.') # Sampling options parser.add_argument('--num_samples', type=int, default=2, help='Number of samples per input for uncertainty estimation') parser.add_argument('--num_steps', type=int, default=35, help='Number of diffusion sampling steps (overrides config val_num_steps)') parser.add_argument('--seed', type=int, default=None, help='Fix generation + perturbation seeds for reproducible / paired runs. ' 'Per-scene seed = (seed + crc32(sample_id)) so it is stable across runs, ' 'shard layouts, and checkpoints (scene order does not matter). ' 'Default None = random per run.') parser.add_argument('--guidance', type=float, default=None, help='Classifier-free guidance scale (overrides config val_cfg_scale)') parser.add_argument('--cond_aug_sigma', type=float, default=None, help='Conditioning augmentation sigma (overrides config val_cond_aug_sigma)') # Output / visualization options parser.add_argument('--output_dir', type=str, required=True, help='Directory to save results') parser.add_argument('--save_uncertainty', type=int, default=1, help='Save uncertainty heatmaps') parser.add_argument('--visual_detail', type=int, default=2, help='Level of visual detail (1/2/3)') parser.add_argument('--quality_pix', type=int, default=7, help='Quality for pixel space videos (0-10)') parser.add_argument('--quality_lat', type=int, default=8, help='Quality for latent space videos (0-10)') parser.add_argument('--viz_extra_modes', type=str, default=None, help='Comma-separated list of extra viz modes (e.g., anyact1,anydrive1)') # World model options parser.add_argument('--perturb_action', type=str, default=None, help='Perturb robot actions (e.g., basile1, basile2)') parser.add_argument('--perturb_traj', type=str, default=None, help='Perturb driving trajectory (e.g., basile1, basile2, basile3)') parser.add_argument('--perturb_per_sample', type=int, default=1, help='Generate different perturbation per sample (1=enabled, useful for WM eval)') parser.add_argument('--run_autoregressive', type=int, default=0, help='Run autoregressive inference (1=enabled), disabled by default') parser.add_argument('--cond_frames_raw', type=int, default=None, help='Number of conditioning RGB frames. Sets the conditioning length for the ' 'standard single-forward forecast too (not autoregressive-only); pinning it ' 'bypasses the per-scene random draw so paired/seeded evals stay controlled.') parser.add_argument('--num_segments', type=int, default=2, help='Number of segments for autoregressive inference') parser.add_argument('--extrapolation_strategy', type=str, default='backtrack', choices=('backtrack', 'extrapolate'), help='Action extrapolation strategy for autoregressive inference (backtrack or extrapolate)') def parse_args(): parser = argparse.ArgumentParser( description='Run Any4D diffusion inference on a single anydata dataset') add_base_args(parser) # Lightweight optional donor RGB: overrides one camera's frames from a video file. # No second dataloader; cameras/extrinsics/intrinsics are left untouched. parser.add_argument('--donor_rgb_path', type=str, default=None, help='Path to an RGB video whose frames replace one camera\'s RGB in each batch') parser.add_argument('--donor_rgb_cam', type=int, default=0, help='Camera index in the batch to overwrite with --donor_rgb_path frames (default 0)') args = parser.parse_args() return args def load_dataset(args, exp_cfg=None): ''' Load anydata dataset from yaml config. :param exp_cfg: Experiment config dict (from load_experiment_config) to read data_val_overrides. :return (DataLoader): DataLoader for the dataset ''' print(f'[cyan]Loading anydata dataset from {args.dset_cfg}...') # Start with data_val_overrides from experiment config (if available) data_overrides = {} if exp_cfg is not None: a4d_cfg = exp_cfg.get('any4d_config', {}) val_ov = a4d_cfg.get('data_val_overrides', {}) if val_ov: data_overrides.update(val_ov) print(f'[cyan] from exp config data_val_overrides: {val_ov}') # CLI overrides take precedence if args.data_overrides: cli_ov = json.loads(args.data_overrides) data_overrides.update(cli_ov) print(f'[cyan] from CLI data_overrides: {cli_ov}') if not data_overrides: data_overrides = None else: print(f'[cyan] merged data_overrides: {data_overrides}') data_defaults = None if exp_cfg is not None: a4d_cfg = exp_cfg.get('any4d_config', {}) data_defaults = a4d_cfg.get('data_defaults', None) if args.data_defaults: cli_defaults = json.loads(args.data_defaults) data_defaults = {**(data_defaults or {}), **cli_defaults} dataset = AnyDataset( dataset_dir=args.dset_cfg, phase='test', data_overrides=data_overrides, data_defaults=data_defaults, single_gpu=True, shuffle=False, ) dataloader = DataLoader( dataset, batch_size=args.batch_size, shuffle=False, num_workers=0, pin_memory=True, collate_fn=any4d_collate, ) print(f'[green]Dataset loaded: {len(dataset)} samples') return dataloader def maybe_override_rgb_from_path(data_batch, rgb_path, cam_idx): ''' Lightweight donor RGB: replace one camera's RGB frames in the batch with frames loaded from a video file. Leaves cameras/extrinsics/intrinsics untouched. Modeled on gradio_app.override_batch_with_video, but parameterized on cam_idx and sized from the batch's existing (t, cam_idx) rgb entries (no T/HW inference from elsewhere). ''' import torch.nn.functional as F anydata = data_batch['anydata'] ts_keys = list(anydata.get('rgb', {}).keys()) times = sorted(set(k[0] for k in ts_keys if k[1] == cam_idx)) if not times: cams = sorted(set(k[1] for k in ts_keys)) raise ValueError(f'--donor_rgb_cam={cam_idx} not present in batch rgb cams {cams}') T = len(times) first_key = (times[0], cam_idx) tgt = anydata['rgb'][first_key] target_H, target_W = tgt.shape[-2:] B = tgt.shape[0] print(f'[cyan]Overriding RGB for cam {cam_idx} (T={T}, HxW={target_H}x{target_W}) from {rgb_path}') frames = infer_utils.create_input_from_video( rgb_path, num_frames=T, frame_width=int(target_W), frame_height=int(target_H), center_crop=True, frame_stride=1, frame_offset=0, ) if frames is None: raise RuntimeError(f'Could not load donor RGB from {rgb_path}') # create_input_from_video already resizes, but guard in case of int rounding mismatch. if tuple(frames.shape[-2:]) != (int(target_H), int(target_W)): frames = F.interpolate( frames, size=(int(target_H), int(target_W)), mode='bilinear', align_corners=False) for i, t in enumerate(times): key = (t, cam_idx) src = anydata['rgb'][key] anydata['rgb'][key] = frames[i:i + 1].repeat(B, 1, 1, 1).to(src.device, dtype=src.dtype) # Downstream model code expects timestep entries as tensors (same as gradio does). if 'timestep' in anydata: for k in list(anydata['timestep'].keys()): if isinstance(anydata['timestep'][k], int): anydata['timestep'][k] = torch.tensor([anydata['timestep'][k]]) return data_batch def run_inference_multiple_samples(model, config, data_batch, num_samples, output_dir, perturb_action=None, perturb_traj=None, perturb_per_sample=False, run_autoregressive=False, cond_frames_raw=None, num_segments=4, extrapolation_strategy='backtrack', extra_directives=None, base_seed=None, visual_detail=2): ''' Run inference multiple times to generate samples for uncertainty estimation. :param model: The Any4D model :param config: Model configuration :param data_batch: Input data batch :param num_samples: Number of samples to generate :param output_dir: Output directory for visualizations :param perturb_action: Optional action perturbation method (e.g., 'basile1') :param perturb_traj: Optional trajectory perturbation method (e.g., 'basile3') :param perturb_per_sample: If True, use different perturbation seed per sample (for WM eval) :param run_autoregressive: If True, run autoregressive inference :param cond_frames_raw: Number of raw frames to condition on (for autoregressive inference) :param num_segments: Number of segments for autoregressive inference :param extrapolation_strategy: Strategy for action extrapolation ('backtrack' or 'extrapolate') :param extra_directives: Additional directives to merge (e.g., used_cams, num_pred_views) :param base_seed: Optional fixed base seed (per-scene); None = random per call :param visual_detail: val_visuals_detail level used while saving (1/2/3; CLI --visual_detail) :return (list): List of output dictionaries ''' all_samples = [] # Get directives from config base_directives = {} if hasattr(config, 'val_directives') and isinstance(config.val_directives, dict): base_directives = config.val_directives.copy() if extra_directives: base_directives.update(extra_directives) # Add perturb_action directive if specified if perturb_action is not None: base_directives['perturb_action'] = perturb_action # Add perturb_traj directive if specified if perturb_traj is not None: base_directives['perturb_traj'] = perturb_traj if run_autoregressive: base_directives['run_autoregressive'] = True base_directives['num_segments'] = num_segments base_directives['extrapolation_strategy'] = extrapolation_strategy if cond_frames_raw is not None: base_directives['cond_frames_raw'] = cond_frames_raw if base_seed is None: base_seed = np.random.randint(1000000, 9999999) has_perturbation = perturb_action is not None or perturb_traj is not None for sample_idx in range(num_samples): cur_seed = base_seed + sample_idx print(f'[cyan]Generating sample {sample_idx + 1}/{num_samples} (seed={cur_seed})...') # Copy directives for this sample directives = base_directives.copy() # If perturb_per_sample is enabled, pass a different seed for perturbation # This allows generating different perturbations while keeping same data if perturb_per_sample and has_perturbation: directives['perturb_seed'] = cur_seed print(f'[yellow] Using perturb_seed={cur_seed} for perturbation') # Run model validation step with direct output to final directory with torch.no_grad(): # Use desired visualization detail level (CLI --visual_detail; was hardcoded to 2) original_detail = config.val_visuals_detail config.val_visuals_detail = visual_detail val_dict, loss = model.validation_step( data_batch=copy.deepcopy(data_batch), iteration=-1, dataloader_key='infer', local_path=output_dir, # Save directly to output directory directives=directives, val_iter=sample_idx, seed=cur_seed, ) # Restore original setting config.val_visuals_detail = original_detail all_samples.append(val_dict) return all_samples def setup_runtime(args): ''' Setup device + distributed, load experiment config + model + VAE + text encoder, apply CLI overrides, initialize LPIPS. Returns (model, config, lpips_loss, device, exp_cfg). ''' device = args.device if 'cuda' in device: device = f'cuda:{args.gpu_id}' torch.cuda.set_device(args.gpu_id) # Initialize torch.distributed for single GPU (required by some model code) if not args.debug and not torch.distributed.is_initialized(): import random # Fixed port via env avoids EADDRINUSE when many single-GPU infer procs launch at once. port = int(os.environ.get('INFER_MASTER_PORT', random.randint(29000, 30000))) torch.distributed.init_process_group( backend='gloo', # Use gloo for CPU/single GPU init_method=f'tcp://127.0.0.1:{port}', world_size=1, rank=0 ) elif args.debug: print('[yellow]Debug mode: Skipping distributed initialization') print(f'[bold cyan]=======================================') print(f'[bold cyan] Any4D Inference Script (AnyData)') print(f'[bold cyan]=======================================') print(f'[cyan]Device: {device}') print(f'[cyan]Output directory: {args.output_dir}') # Load experiment configuration exp_cfg = infer_utils.load_experiment_config(args.exp_cfg) # Apply CLI overrides to any4d_config before model build (e.g. ablations like block_gate_fix=False). if args.exp_overrides: cli_exp_ov = json.loads(args.exp_overrides) any4d_cfg = exp_cfg.setdefault('any4d_config', {}) for k, v in cli_exp_ov.items(): old = any4d_cfg.get(k, '') any4d_cfg[k] = v print(f'[yellow]Overriding any4d_config.{k}: {old} -> {v}') # Build model model, config = infer_utils.build_model( exp_cfg, ckpt_path=args.ckpt_path, num_steps=args.num_steps) model = model.to(device) model.device = device # Set path_local for print_run_info (normally set by training job setup) if not hasattr(config.job, 'path_local') or config.job.get('path_local', None) is None: config.job.path_local = os.path.relpath(args.output_dir) # Append test tag to train tag to affect exported file names and visualizations test_tag = os.path.basename(args.output_dir).split('_')[0] print(f'[cyan]Detected test tag: {test_tag}') config.job.undated_name = config.job.undated_name + '_' + test_tag print(f'[cyan]Updated overall experiment tag to: {config.job.undated_name}') # Apply sampling parameter overrides if args.guidance is not None: config.val_cfg_scale = args.guidance print(f'[cyan]Overriding val_cfg_scale: {args.guidance}') if args.cond_aug_sigma is not None: config.val_cond_aug_sigma = args.cond_aug_sigma print(f'[cyan]Overriding val_cond_aug_sigma: {args.cond_aug_sigma}') # Apply viz_extra_modes: CLI override if given, else inherit the exp config's modes (e.g. rwm11_gaia # sets ['anydrive1', 'gaia1']). Don't clobber the config when unset; nogt is opt-in (add explicitly). if args.viz_extra_modes is not None: config.viz_extra_modes = [m.strip() for m in args.viz_extra_modes.split(',')] else: config.viz_extra_modes = list(config.viz_extra_modes or []) print(f'[cyan]viz_extra_modes: {config.viz_extra_modes}') # Set up all model modules (loads DiT weights, VAE, text encoder, transforms) model = infer_utils.setup_model(model, config, ckpt_path=args.ckpt_path) model.eval() # Initialize validation metrics list (normally done in on_validation_start) model.all_val_metrics = [] # Initialize LPIPS loss for metrics calculation if available lpips_loss = None if LPIPS_AVAILABLE: print('[cyan]Initializing LPIPS loss...') lpips_loss = lpips.LPIPS(net='vgg').to(device) lpips_loss.eval() else: print('[yellow]LPIPS not available, skipping metrics calculation') return model, config, lpips_loss, device, exp_cfg def process_one_sample(args, config, model, data_batch, lpips_loss, device, iteration, all_scene_metrics, extra_directives=None): ''' Forward pass + per-scene metrics + uncertainty heatmaps + save for a single data_batch. Appends per-scene metrics to all_scene_metrics in place. Returns scene_metrics. ''' base_seed = infer_utils.derive_scene_seed(args.seed, data_batch, iteration) all_samples = run_inference_multiple_samples( model, config, data_batch, args.num_samples, args.output_dir, perturb_action=args.perturb_action, perturb_traj=args.perturb_traj, perturb_per_sample=args.perturb_per_sample, run_autoregressive=(args.run_autoregressive == 1), cond_frames_raw=args.cond_frames_raw, num_segments=args.num_segments, extrapolation_strategy=args.extrapolation_strategy, extra_directives=extra_directives, base_seed=base_seed, visual_detail=args.visual_detail, ) metric_resolution = json.loads(args.metric_resolution) if args.metric_resolution else None # Optional in-code metrics path (--infer_metrics). Default OFF: info.json already carries the # native-res 11-view block, so this rgb-multi-view recompute is redundant unless explicitly asked. scene_metrics = {} if args.infer_metrics and LPIPS_AVAILABLE and SKIMAGE_AVAILABLE and lpips_loss is not None: scene_metrics = infer_utils.calculate_metrics_per_scene( all_samples, data_batch, config, lpips_loss, model, device, metric_resolution=metric_resolution) all_scene_metrics.append(scene_metrics) avg_psnr = scene_metrics.get('psnr', {}).get('avg') if scene_metrics else None if avg_psnr is not None: print(f'[green]Scene {iteration} (avg over {len(scene_metrics["psnr"]) - 1} views): ' f'PSNR={scene_metrics["psnr"]["avg"]:.2f}, SSIM={scene_metrics["ssim"]["avg"]:.4f}, ' f'LPIPS={scene_metrics["lpips"]["avg"]:.4f}') uncertainty_videos = {} diversity_metrics = {} if args.num_samples > 1: entry_names = ['rgb0'] # Only calculate diversity for rgb0 uncertainty_videos, diversity_metrics = infer_utils.calculate_uncertainty_heatmaps( all_samples, entry_names, a4d_vae=model.vae) infer_utils.save_inference_results( args, config, model, iteration, data_batch, all_samples, uncertainty_videos, diversity_metrics, scene_metrics) return scene_metrics def finalize_run(args, all_scene_metrics, num_to_process, start_time): '''Aggregate + save metrics across scenes, print summary.''' if len(all_scene_metrics) > 0: mr = json.loads(args.metric_resolution) if args.metric_resolution else None infer_utils.aggregate_and_save_metrics(all_scene_metrics, args.output_dir, metric_resolution=mr, seed=args.seed) elapsed_time = time.time() - start_time print(f'[bold green]=======================================') print(f'[bold green] Inference Complete!') print(f'[bold green]=======================================') print(f'[green]Total time: {elapsed_time:.2f}s') if num_to_process > 0: print(f'[green]Average time per sample: {elapsed_time / num_to_process:.2f}s') def main(args): start_time = time.time() model, config, lpips_loss, device, exp_cfg = setup_runtime(args) dataloader = load_dataset(args, exp_cfg=exp_cfg) num_total = len(dataloader) num_to_process = num_total if args.stop_after < 0 else min(args.stop_after, num_total) shard_i, shard_n = 0, 1 if args.shard: shard_i, shard_n = (int(x) for x in args.shard.split('/')) assert 0 <= shard_i < shard_n, f'invalid --shard {args.shard}' num_in_shard = len(range(shard_i, num_to_process, shard_n)) print(f'[cyan]Processing {num_in_shard} scenes (shard {shard_i}/{shard_n} of {num_to_process} total)...') all_scene_metrics = [] with Progress( SpinnerColumn(), TextColumn('[progress.description]{task.description}'), BarColumn(), TextColumn('[progress.percentage]{task.percentage:>3.0f}%'), TimeRemainingColumn(), ) as progress: task = progress.add_task('[cyan]Processing samples...', total=num_in_shard) for iteration, data_batch in enumerate(dataloader): if iteration >= num_to_process: break if shard_n > 1 and (iteration % shard_n) != shard_i: continue try: if args.donor_rgb_path: data_batch = maybe_override_rgb_from_path( data_batch, args.donor_rgb_path, args.donor_rgb_cam) process_one_sample( args, config, model, data_batch, lpips_loss, device, iteration, all_scene_metrics, extra_directives=None) progress.update(task, advance=1) except Exception as e: print(f'[red]Error processing sample {iteration}: {e}') print(f'[red]Traceback: {traceback.format_exc()}') print(f'[yellow]Skipping...') progress.update(task, advance=1) continue finalize_run(args, all_scene_metrics, num_in_shard, start_time) if __name__ == '__main__': args = parse_args() main(args)