"""Minimal Any4D + Wan 2.1 Fun 1.3B InP inference. Reuses the repo's diffusion sampling logic (Any4DPipeline.generate) — no custom sampler. Builds Any4DWanPipeline directly (skipping Any4DModel) so it can run without instantiating training-only components. Flow: 1. Build Any4DWanPipeline (loads DiT weights from checkpoint). 2. Load N-view videos + trajectory. 3. VAE-encode the first-frame-only pixel video for each view (rgb_ref latent). 4. Populate a4d_latent with masks (first-frame-only conditioning). 5. Call pipe.generate() — uses WanFlowMatchScheduler (shift=5, s-space). 6. VAE-decode the predicted v{N}[:, 0:16] rgb slots and save as mp4. `--video_paths` (one mp4 per view) and `--traj_npy` ((num_frames, 2) ego trajectory in metres, lateral+forward) are both required. Use `custom/eval/dump_val_sample_wan.py` to extract them from a validation sample. Text conditioning is fixed at the precomputed UMT5-XXL embedding in `custom/wan/default_text_wan.pkl` (matching training, which sets `text_encoder_path=''`). There is no `--prompt` flag because no runtime text encoder is wired in — a free-form prompt string would not affect the output. Usage (inside cosmos-predict2-local docker): python -m custom.eval.infer_wan \\ --ckpt_path /path/to/iter_NNNNN.pt \\ --vae_path /path/to/Wan2.1_VAE.pth \\ --out_dir /path/to/outputs \\ --video_paths /path/CAMERA_01.mp4 /path/CAMERA_05.mp4 ... \\ --traj_npy /path/to/traj.npy """ import argparse import copy import os import pickle import sys import numpy as np import torch import torch.nn as nn import torch.nn.functional as F ANY4D_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) if ANY4D_ROOT not in sys.path: sys.path.insert(0, ANY4D_ROOT) os.chdir(ANY4D_ROOT) def _init_dist_single_gpu(): if torch.distributed.is_initialized(): return import random torch.distributed.init_process_group( backend='gloo', init_method=f'tcp://127.0.0.1:{random.randint(29000, 30000)}', world_size=1, rank=0, ) # ------------------------------- build pipeline --------------------------------- def _build_wan_pipe_cfg(vae_path: str, chunk_duration: int, state_t: int): """Return a fresh copy of ANY4D_WAN_PIPELINE_1_3B aligned with the experiment's pixel-frame chunking and latent-frame count (default ANY4D_WAN_PIPELINE_1_3B is 41 pixel -> 11 latent; 6-view driving runs at 21 -> 6).""" from cosmos_predict2.configs.base.defaults.model import ANY4D_WAN_PIPELINE_1_3B pipe_cfg = copy.deepcopy(ANY4D_WAN_PIPELINE_1_3B) pipe_cfg.tokenizer.vae_pth = vae_path pipe_cfg.tokenizer.chunk_duration = chunk_duration pipe_cfg.state_t = state_t if pipe_cfg.extra_nets is None: pipe_cfg.extra_nets = nn.ModuleDict() return pipe_cfg def _build_any4d_config(num_views: int, num_frames: int): from custom.any4d.a4d_config import Any4DConfig from custom.any4d import a4d_surgery video_entries = [dict( rgb0=(0, 16, 'load', 'load'), rgb0_input_mask=(16, 17, 'load', None), rgb0_output_mask=(17, 18, 'load', None), rgb0_ref=(18, 34, 'load', None), )] for v in range(1, num_views): prev = 0 if v == 1 else 1 # all later views copy from rgb1 (matches 6-cam experiment) video_entries.append({ f'rgb{v}': (0, 16, f'copy:rgb{prev}/load', f'copy:rgb{prev}/load'), f'rgb{v}_input_mask': (16, 17, f'copy:rgb{prev}_input_mask/load', None), f'rgb{v}_output_mask': (17, 18, f'copy:rgb{prev}_output_mask/load', None), f'rgb{v}_ref': (18, 34, f'copy:rgb{prev}_ref/load', None), }) cfg = Any4DConfig( vidar_active=True, any4d_active=True, use_wan_backbone=True, num_views=num_views, video_entries=video_entries, lowdim_adaln_entries=dict( traj=(0, num_frames, 0, 2, 'zero/load'), traj_input_mask=(0, num_frames, 2, 3, 'zero/load'), ), video_concat_mode='view', video_proj_mode='per_view', view_timestep_mode='per_view', block_gate_fix=True, harmonize_streams=True, harmonize_frames=True, load_modals=['rgb', 'traj'], ) return a4d_surgery.validate_config(cfg) def build_pipe(ckpt_path: str, vae_path: str, num_views: int, num_frames: int): _init_dist_single_gpu() from custom.wan.a4d_pipe_wan import build_wan_pipeline state_t = (num_frames - 1) // 4 + 1 pipe_cfg = _build_wan_pipe_cfg(vae_path, chunk_duration=num_frames, state_t=state_t) cfg = _build_any4d_config(num_views=num_views, num_frames=num_frames) print(f'[build] loading DiT weights from {ckpt_path} ...') pipe = build_wan_pipeline( pipe_cfg, dit_path=ckpt_path, text_encoder_path='', vae_path=vae_path, any4d_config=cfg, ) pipe.any4d_config = cfg return pipe, cfg # ------------------------------- data loading --------------------------------- def _load_input_pixel_video(path, target_h, target_w, num_frames, num_input_frames, device, dtype): """(3, Tp, H, W) in [0, 1], with the first num_input_frames real, the rest zeros.""" assert 1 <= num_input_frames <= num_frames import imageio.v3 as iio frames = [] for i in range(num_input_frames): frame_np = iio.imread(path, index=i) t = torch.from_numpy(np.asarray(frame_np)).permute(2, 0, 1).float() / 255.0 frames.append(t) real = torch.stack(frames, dim=1) # (3, num_input_frames, H0, W0) real = F.interpolate(real.unsqueeze(0), size=(num_input_frames, target_h, target_w), mode='trilinear', align_corners=False).squeeze(0) zeros = torch.zeros((3, num_frames - num_input_frames, target_h, target_w)) return torch.cat([real, zeros], dim=1).to(device=device, dtype=dtype) def _vae_encode_wan_ref(tokenizer, pixel_video_01, num_input_frames): """Build Wan Fun InP's ref latent. pixel_video_01 shape: (B, 3, Tp, H, W) in [0, 1], with the first num_input_frames being the real conditioning frames and the rest zero. Straight `* 2 - 1` would map pad frames to BLACK (-1) in Wan VAE's [-1, 1] input space — mismatched with Wan's training distribution where pad frames are 0 (GRAY) in [-1, 1]. Convert only the conditioning frames to [-1, 1] and keep the rest at 0. """ ref_m11 = torch.zeros_like(pixel_video_01) ref_m11[:, :, :num_input_frames] = pixel_video_01[:, :, :num_input_frames] * 2.0 - 1.0 return tokenizer.encode(ref_m11) def _load_default_text(device, dtype): """Load the precomputed UMT5-XXL embeddings (positive + negative).""" path = os.path.join(ANY4D_ROOT, 'custom/wan/default_text_wan.pkl') assert os.path.isfile(path), ( f'Missing {path} — run:\n' f' python -m custom.wan.build_default_text_wan \\\n' f' --umt5_tokenizer_path /path/to/Wan2.1-Fun-1.3B-InP/google/umt5-xxl \\\n' f' --umt5_weights_path /path/to/models_t5_umt5-xxl-enc-bf16.pth') with open(path, 'rb') as f: d = pickle.load(f) out = { 't5_text_embeddings': d['t5_text_embeddings'].to(device=device, dtype=dtype), 't5_text_mask': d['t5_text_mask'].to(device=device, dtype=dtype), } if 't5_text_embeddings_neg' not in d: raise RuntimeError( 'default_text_wan.pkl is missing the negative embedding; rebuild it with the ' 'updated custom/wan/build_default_text_wan.py so CFG uses the UMT5("") tensor ' 'rather than a raw zero vector.') out['neg_t5_text_embeddings'] = d['t5_text_embeddings_neg'].to(device=device, dtype=dtype) out['neg_t5_text_mask'] = d['t5_text_mask_neg'].to(device=device, dtype=dtype) return out def build_data_batch(pipe, video_paths, num_views, num_frames, num_input_frames, target_h, target_w, fps, traj): """Construct a data_batch compatible with Any4DPipeline.generate(). num_input_frames: number of leading pixel frames to use as conditioning. The remaining pixel frames are predicted. Must satisfy `(num_input_frames - 1) % 4 == 0` so the conditioning region aligns with Wan VAE's 4x temporal compression (no partially-conditioned latent frames). Text conditioning uses the precomputed UMT5 embedding from `custom/wan/default_text_wan.pkl` (same path training takes when `text_encoder_path=''`). There is no runtime text encoder, so a free-form prompt string has no effect on generation. """ device, dtype = torch.device('cuda'), torch.bfloat16 Tp = num_frames Tl = (Tp - 1) // 4 + 1 Hl, Wl = target_h // 8, target_w // 8 assert (num_input_frames - 1) % 4 == 0, \ f'num_input_frames must be 1 mod 4 to align with VAE 4x temporal compression; got {num_input_frames}' cond_lat = (num_input_frames - 1) // 4 + 1 # latent frames covered entirely by conditioning assert len(video_paths) == num_views, \ f'video_paths must have exactly num_views ({num_views}) entries, got {len(video_paths)}' pix = [_load_input_pixel_video(vp, target_h, target_w, Tp, num_input_frames, device, dtype) for vp in video_paths] pix = [p.unsqueeze(0) for p in pix] tok = pipe.tokenizer ref_latents = [_vae_encode_wan_ref(tok, p, num_input_frames) for p in pix] rgb_latents = ref_latents a4d_latent = {} for v, (rgb_z, ref_z) in enumerate(zip(rgb_latents, ref_latents)): a4d_latent[f'rgb{v}'] = rgb_z a4d_latent[f'rgb{v}_ref'] = ref_z imask = torch.zeros((1, 1, Tl, Hl, Wl), device=device, dtype=dtype) imask[:, :, :cond_lat] = 1.0 a4d_latent[f'rgb{v}_input_mask'] = imask a4d_latent[f'rgb{v}_output_mask'] = 1.0 - imask a4d_latent[f'rgb{v}_supervise_mask'] = 1.0 - imask a4d_latent[f'rgb{v}_ref_input_mask'] = torch.ones_like(imask) a4d_latent[f'rgb{v}_ref_output_mask'] = torch.zeros_like(imask) a4d_latent[f'rgb{v}_ref_supervise_mask'] = torch.zeros_like(imask) a4d_latent['traj'] = traj.to(device=device, dtype=dtype) a4d_latent['traj_input_mask'] = torch.ones((1, Tp, 1), device=device, dtype=dtype) a4d_latent['traj_output_mask'] = torch.zeros((1, Tp, 1), device=device, dtype=dtype) a4d_latent['traj_supervise_mask'] = torch.zeros((1, Tp, 1), device=device, dtype=dtype) a4d_latent['video_dims'] = [(Tl, Hl, Wl)] * num_views data_batch = { 'fps': torch.tensor([fps], dtype=dtype, device=device), 'padding_mask': torch.zeros((1, 1, target_h, target_w), dtype=dtype, device=device), 'prompt': [''], # placeholder; actual text conditioning comes from the pkl below 'a4d_raw': {}, 'a4d_latent': a4d_latent, 'anydata': {}, } data_batch.update(_load_default_text(device, dtype)) return data_batch # ------------------------------- sampling + decode --------------------------------- def run_inference(pipe, data_batch, num_steps, guidance, seed): with torch.no_grad(): samples = pipe.generate( data_batch, seed=seed, num_sampling_steps=num_steps, guidance=guidance, is_negative_prompt=True, phase='val', cond_aug_sigma=0.0, ) return samples def decode_and_save(pipe, samples, out_dir, num_views, fps=10): import imageio os.makedirs(out_dir, exist_ok=True) y0 = samples['y0_pred_streams'] for v in range(num_views): v_key = f'v{v}' if v_key not in y0: continue lat_full = y0[v_key] lat_rgb = lat_full[:, 0:16].to(dtype=torch.bfloat16) with torch.no_grad(): raw = pipe.tokenizer.decode(lat_rgb) raw = raw.clamp(-1, 1).squeeze(0) raw = (raw + 1.0) / 2.0 pix = (raw * 255.0).to(torch.uint8).permute(1, 2, 3, 0).cpu().numpy() out_path = os.path.join(out_dir, f'pred_v{v}.mp4') imageio.mimwrite(out_path, pix, fps=fps, quality=7) print(f'[save] {out_path} shape={pix.shape} mean={pix.mean():.1f} std={pix.std():.1f}') # -------------------------------- entry point --------------------------------- def main(): p = argparse.ArgumentParser() p.add_argument('--ckpt_path', type=str, required=True, help='Path to a Wan-format DiT checkpoint (.pt or .safetensors)') p.add_argument('--vae_path', type=str, required=True, help='Path to Wan2.1_VAE.pth (or a dir containing it)') p.add_argument('--out_dir', type=str, required=True) p.add_argument('--video_paths', nargs='+', required=True, help='Per-view input mp4 paths (one per view, len == --num_views). ' 'Only the first frame is used for conditioning.') p.add_argument('--traj_npy', type=str, required=True, help='Path to a (num_frames, 2) .npy of (lateral, forward) metres ' 'relative to the first ego pose.') p.add_argument('--num_views', type=int, default=6) p.add_argument('--num_frames', type=int, default=21, help='Total pixel frames per view; latent frames = (num_frames - 1) / 4 + 1.') p.add_argument('--num_input_frames', type=int, default=9, help='Number of leading pixel frames used as conditioning (the rest are ' 'predicted). Must be 1 mod 4 (1, 5, 9, 13, ...) so the conditioning ' 'region maps cleanly to whole latent frames.') p.add_argument('--target_h', type=int, default=320) p.add_argument('--target_w', type=int, default=512) p.add_argument('--fps', type=float, default=10.0) p.add_argument('--num_steps', type=int, default=35) p.add_argument('--guidance', type=float, default=5.0) p.add_argument('--seed', type=int, default=0) # Note: no --prompt. Text conditioning is the precomputed default_text_wan.pkl # embedding (same as training); a runtime prompt string would have no effect # because no text encoder is wired up. args = p.parse_args() device, dtype = torch.device('cuda'), torch.bfloat16 assert len(args.video_paths) == args.num_views, \ f'--video_paths must have exactly --num_views ({args.num_views}) entries, got {len(args.video_paths)}' for vp in args.video_paths: assert os.path.isfile(vp), f'Missing per-view video: {vp}' video_paths = list(args.video_paths) traj_np = np.load(args.traj_npy) assert traj_np.shape == (args.num_frames, 2), \ f'--traj_npy must be shape ({args.num_frames}, 2), got {traj_np.shape}' traj = torch.from_numpy(traj_np).unsqueeze(0).to(device=device, dtype=dtype) pipe, cfg = build_pipe(args.ckpt_path, args.vae_path, num_views=args.num_views, num_frames=args.num_frames) print('[data] building data_batch ...') data_batch = build_data_batch( pipe, video_paths, num_views=args.num_views, num_frames=args.num_frames, num_input_frames=args.num_input_frames, target_h=args.target_h, target_w=args.target_w, fps=args.fps, traj=traj, ) print(f'[infer] pipe.generate(num_steps={args.num_steps}, guidance={args.guidance}, seed={args.seed}) ...') samples = run_inference(pipe, data_batch, args.num_steps, args.guidance, args.seed) print('[decode] VAE decode + save mp4 ...') decode_and_save(pipe, samples, args.out_dir, num_views=args.num_views, fps=int(args.fps)) print('[done]') if __name__ == '__main__': main()