''' BVH, Nov 2025. Any4D / AnyView / DVS (dynamic view synthesis) inference script. Supports vidar datasets, MP4 videos, visualizations, and uncertainty heatmaps. ''' 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.dataset.vidar_dataset import VidarDataset 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 parse_args(): parser = argparse.ArgumentParser( description='Run Any4D diffusion inference on vidar dataset') # 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/basile/rdvs2_mix_b16.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 vidar 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.') # 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('--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)') parser.add_argument('--exp_overrides', type=str, default=None, help='JSON dict of any4d_config overrides applied before build_model ' '(e.g. \'{"traj_multiplier": 2.0}\')') # Video override options parser.add_argument('--override_tarfile', type=str, default=None, help='Path to tarfile to override input video and/or camera') parser.add_argument('--override_video', type=int, default=0) parser.add_argument('--override_camera', type=int, default=0) # 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=1, help='Level of visual detail (1-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)') parser.add_argument('--seed', type=int, default=None, help='Fix generation + perturbation seeds for reproducible / paired runs. ' 'Per-scene seed = (seed + crc32(sample_id)), stable across runs and ' 'checkpoints. Default None = random per run.') # 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 raw frames to condition on (for autoregressive inference)') 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)') args = parser.parse_args() return args def load_dataset(args): ''' Load vidar dataset from yaml config. :return (DataLoader): DataLoader for the dataset ''' print(f'[cyan]Loading dataset from {args.dset_cfg}...') dataset = VidarDataset( dataset_dir=args.dset_cfg, phase='test', ) dataloader = DataLoader( dataset, batch_size=args.batch_size, shuffle=False, num_workers=0, pin_memory=True, ) print(f'[green]Dataset loaded: {len(dataset)} samples') return dataloader 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', base_seed=None): ''' 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') :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() # 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 original_detail = config.val_visuals_detail config.val_visuals_detail = 2 # Good detail level for inference 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 main(args): start_time = time.time() # Setup device 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 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') 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 traj_multiplier=2.0). 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 / _s3_path (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 override if specified if args.viz_extra_modes: config.viz_extra_modes = [m.strip() for m in args.viz_extra_modes.split(',')] print(f'[cyan]Overriding 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') # Load dataset dataloader = load_dataset(args) # Determine how many samples to process num_total = len(dataloader) num_to_process = num_total if args.stop_after < 0 else min(args.stop_after, num_total) print(f'[cyan]Processing {num_to_process} samples...') # Parse video override options # TODO: Implement tarfile support for video/camera overrides if args.override_tarfile: print(f'[yellow]Note: tarfile override not yet fully implemented: {args.override_tarfile}') # Iterate over dataset with progress bar 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_to_process) skips = 0 for iteration, data_batch in enumerate(dataloader): if iteration < skips: print(f'[yellow]Skipping sample {iteration} (skips={skips})') continue if iteration >= num_to_process + skips: break try: base_seed = infer_utils.derive_scene_seed(args.seed, data_batch, iteration) # Run inference multiple times 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, base_seed=base_seed, ) # Calculate metrics per scene if LPIPS is available (per-entry + avg schema) scene_metrics = {} if 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) all_scene_metrics.append(scene_metrics) # Print scene metrics (headline = avg over scored views) avg_psnr = scene_metrics.get('psnr', {}).get('avg') if scene_metrics else None if avg_psnr is not None: print(f'[green]Scene {iteration}: PSNR={scene_metrics["psnr"]["avg"]:.2f}, ' f'SSIM={scene_metrics["ssim"]["avg"]:.4f}, LPIPS={scene_metrics["lpips"]["avg"]:.4f}') # Calculate uncertainty if multiple samples (only for rgb0) 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) # Save additional results (metrics, uncertainty videos) infer_utils.save_inference_results( args, config, model, iteration, data_batch, all_samples, uncertainty_videos, diversity_metrics, scene_metrics) 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 # Aggregate and save metrics across all scenes if len(all_scene_metrics) > 0: infer_utils.aggregate_and_save_metrics(all_scene_metrics, args.output_dir) 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') print(f'[green]Average time per sample: {elapsed_time / num_to_process:.2f}s') if __name__ == '__main__': args = parse_args() main(args)