# vid2vid FORK of custom/dataloader/unified_anydrive.py. # Kept as a separate module so the shared unified_anydrive.py stays owned/untouched (avoids a # perpetual merge conflict). Select it per-experiment with dataloader='unified_anydrive_vid2vid'. # Adds, on top of the forecast / world-model base: # - structured-noise phase-source slot fill (rgb{v}_edge_control) via control_maps.populate_edge_control; # - anchor-conditioned view-synthesis task (task_probs.view_synth) + anchor_view plumbing. # It forked from unified_anydrive at this branch's point, so it does NOT track main's later # dataloader refactors (e.g. dyn_view_synth / DVS); use plain 'unified_anydrive' for those. # # Original header: # Dataloader for DRIVING datasets in the UNIFIED (anydata) pipeline (Created by BVH, Apr 2026). # Based on unified_flex (trajectory) and unified_anyact (simplified task logic). # Handles different driving datasets with different # cameras via dataset-specific cam mapping. import numpy as np import torch from rich import print from custom.dataloader.utils import prepare_dense, prepare_traj_anydata, random_choice, random_value, reanchor_cameras_to_view from custom.dataloader.wm_utils import compute_trajectory_from_extrinsics, perturb_traj_basile1, perturb_traj_basile2, perturb_traj_basile3 from custom.utils.constants import VAE_RATIO_THW from custom.wan.control_maps import populate_edge_control # structured-noise phase-slot builders from imaginaire.utils import log # Each driving dataset maps to 4 directional groups: front (=ego), left, right, back. # Each group is a priority-ordered list of camera names. # View assignment: rgb0 = front[0] (ego), # then probably we will do left[0], right[0], back[0], left[1] or right[1], etc. # depending on the number of viewpoints we actually want to use in this batch. # TODO(bvh): inspect datasets to verify cameras, these contain some guesses for now DRIVING_CAM_LAYOUT = { # DDAD: 6 cameras: CAMERA_01=front, 05=front_left, 06=front_right, # 07=rear_left, 08=rear_right, 09=rear. 'ddad': { 'front': ['CAMERA_01'], 'left': ['CAMERA_05', 'CAMERA_07'], 'right': ['CAMERA_06', 'CAMERA_08'], 'back': ['CAMERA_09'], }, # PDv2: 6 cameras: same as DDAD but lowercase 'pdv2': { 'front': ['camera_01'], 'left': ['camera_05', 'camera_07'], 'right': ['camera_06', 'camera_08'], 'back': ['camera_09'], }, # ArgoVerse2sync: 7 cameras # TODO verify names 'argoverse2sync': { 'front': ['ring_front_center'], 'left': ['ring_front_left', 'ring_side_left', 'ring_rear_left'], 'right': ['ring_front_right', 'ring_side_right', 'ring_rear_right'], 'back': [], }, # LyftL5: 7 cameras: 0=front, 5=front_left, 1=front_right # TODO check others 'lyftl5': { 'front': ['0'], 'left': ['5', '3'], 'right': ['1', '2'], 'back': ['4', '6'], }, # Waymo: 5 cameras 'waymo': { 'front': ['0'], 'left': ['1', '3'], # FRONT_LEFT, SIDE_LEFT 'right': ['2', '4'], # FRONT_RIGHT, SIDE_RIGHT 'back': [], }, # PD4D: 19 cameras total, but only 3 egocentric cameras are useful/realistic, # the others are exo views (inward facing from the air) 'pd4d': { 'front': ['yaw-0'], 'left': ['yaw-neg-60'], 'right': ['yaw-60'], 'back': [], }, # GAIA: 11 cameras total (parking scenarios, fisheye). # Both standard and PVM configs use this single layout; _order_cameras # matches whichever cameras are present in the batch. # Front priority: front_wide (standard ego), front_pvm (PVM ego), front_tele (aux). 'gaia': { 'front': ['front_wide', 'front_pvm', 'front_tele'], 'left': ['left_side_forward', 'left_side_pvm', 'left_side_rearward'], 'right': ['right_side_forward', 'right_side_pvm', 'right_side_rearward'], 'back': ['rear_tele', 'rear_pvm'], }, } def _get_ego_camera(dset_name, cam_names): """Return integer index of the ego (front) camera for this dataset. :param dset_name: str, dataset name (e.g. 'DDAD', 'LyftL5'). :param cam_names: list of str, camera names in index order, from metadata['cameras'][b], e.g. ['yaw-0', 'yaw-60', 'yaw-neg-60']. :return: int, camera index of the ego camera. """ # TODO(bvh): debug/verify this more rigorously dset_key = dset_name.lower().split('_')[0] layout = DRIVING_CAM_LAYOUT.get(dset_key, None) assert layout is not None, f'No camera layout for dataset {dset_name!r} (key={dset_key!r})' name_to_idx = {name: idx for idx, name in enumerate(cam_names)} for ego_name in layout['front']: if ego_name in name_to_idx: return name_to_idx[ego_name] # Single-camera samples: when the dataset randomly loads ONE camera per __getitem__ # (e.g. single_sample='random' for random-camera training), there is no front camera # to locate and no choice to make -- the lone camera IS the view. Use it silently; # the front-priority layout simply does not apply, and logging here would fire on # essentially every sample on every rank (~94k ERROR lines/run observed). if len(cam_names) <= 1: return 0 # NOTE(bvh): fallback for genuine multi-camera samples that are missing the front # camera (e.g. IKE_24 DDAD has CAMERA_07/08/09 only). Use first available. Debug-level # + rank0-only so it stays diagnosable without flooding logs across ranks/samples. log.debug( f'Ego camera not found: layout front={layout["front"]}, ' f'available names={cam_names} (dataset={dset_name!r}). ' f'Falling back to first available: {cam_names[0]!r}.', rank0_only=True) return 0 def _order_cameras(dset_name, cam_names, V_max): """Order camera indices: front (ego) first, then interleaved left/right/back. Clamp to V_max. Returns a list of integer camera indices (matching (t, c) dict keys). :param dset_name: str, dataset name. :param cam_names: list of str, camera names in index order. :param V_max: int, max number of viewpoints to use. :return: list of int, camera indices in priority order. """ # TODO(bvh): debug/verify this more rigorously dset_key = dset_name.lower().split('_')[0] layout = DRIVING_CAM_LAYOUT.get(dset_key, None) all_indices = list(range(len(cam_names))) if layout is None: return all_indices[:V_max] name_to_idx = {name: idx for idx, name in enumerate(cam_names)} ordered = [] # Front (ego) first for cam_name in layout['front']: if cam_name in name_to_idx: idx = name_to_idx[cam_name] if idx not in ordered: ordered.append(idx) if not ordered: ordered.append(0) # fallback: index 0 # Interleave left, right, back by priority groups = [layout['left'], layout['right'], layout['back']] max_depth = max((len(g) for g in groups), default=0) for depth in range(max_depth): for group in groups: if depth < len(group): cam_name = group[depth] if cam_name in name_to_idx: idx = name_to_idx[cam_name] if idx not in ordered: ordered.append(idx) # Append any remaining indices not covered by layout for idx in all_indices: if idx not in ordered: ordered.append(idx) return ordered[:V_max] def anydata_to_any4d(data_dict, phase=None, config=None, directives=None, generator=None, train_step=-1, val_step=-1, **leftover): # TODO(bvh): use train_step and/or val_step to orchestrate # synchronized video length variations assert phase in ['train', 'val', 'test'] if directives is None: directives = dict() # V_max from entries (not config.num_views, since concatenation might be involved). # Skip suffixed slots: _mask, _ref (Wan InP), _control (Wan Fun-Control). Only base # rgb{N}/cams{N} entries contribute to the V_max count. _entries_no_mask = [x for x in config.all_highdim_entries if '_mask' not in x and '_ref' not in x and '_control' not in x] V_max = 1 + max([int(x[3:]) for x in _entries_no_mask if x.startswith('rgb')]) has_rgb_entries = any(x.startswith('rgb') for x in _entries_no_mask) has_cams_entries = any(x.startswith('cams') for x in _entries_no_mask) a4d_raw = dict() a4d_latent = dict() anydata_dict = data_dict['anydata'] scale_factors = data_dict['scale_factor'] # (B) tensor of float # NOTE(bvh): dset_name is set by anydata_dataset.py from metadata, available at batch level dset_names = data_dict['dset_name'] # list (B) of str after collation # Extract modalities anydata_rgb = anydata_dict.get('rgb', None) anydata_cams = anydata_dict.get('cams', None) # Fill the structured-noise phase-source slot (rgb{v}_edge_control) from depth / colorized # semantic / Canny-from-RGB (or a per-step source mix) when the experiment declares it but no # precomputed control modality is present. All the logic lives in custom/wan/control_maps.py. anydata_rgb_edge_control = populate_edge_control(anydata_dict, config, phase, train_step, val_step) latent_rgb = anydata_dict.get('lat_rgb', None) latent_cams = anydata_dict.get('lat_cams', None) first_valid = anydata_rgb or latent_rgb assert first_valid is not None, 'Need at least rgb or lat_rgb' keys = list(first_valid.keys()) B = first_valid[keys[0]].shape[0] device = first_valid[keys[0]].device # Available timesteps & camera indices. # NOTE(bvh): (t, c) dict keys use integer indices, NOT camera name strings. # metadata['cameras'][b] maps index -> name for batch element b. avail_times = sorted(set(key[0] for key in anydata_dict['timestep'].keys())) avail_cam_indices = sorted(set(key[1] for key in anydata_dict['timestep'].keys())) T_data = len(avail_times) V_data = len(avail_cam_indices) # TODO(bvh): using first batch element for now, but ordering may differ across batch elements. cam_names = anydata_dict['metadata']['cameras'][0] # Task probabilities (driving: only forecast and world_model) p_fc = config.task_probs.get('forecast', -0.1) p_wm = config.task_probs.get('world_model', -0.1) # View-synthesis (anchor-conditioned) task probability. Guarded: defaults to 0.0 # so non-view-synth runs are byte-for-byte unchanged. When active (Mode B), V_used # is forced to 2: BOTH views are fully GENERATED + supervised (rgb noisy, output_mask # all 1). View 0 is additionally designated the ANCHOR: its dedicated ref slot # ([18:34]) is filled by a4d_vae_wan with view 0's OWN full clean VAE latent # (reference-image broadcast), and that ref latent reaches the Wan Fun-Control DiT's # y/image-conditioning channel [32:48]. View 1's ref is zero. Cross-view consistency # is learned via the 2-view joint attention (view 1 attends to the anchor's clean ref # tokens), exactly like I2V frames attending to the reference frame. The anchor index # is plumbed to a4d_vae_wan via data_dict['meta']['anchor_view'] (see below). p_vs = config.task_probs.get('view_synth', 0.0) bernoulli_fn = lambda p: random_value(device, generator) < p # NOTE(bvh): We always use available cameras up to V_max, in a fixed order per dataset. # Ego forward/front first, then populate with others. # This replaces the shuffle_cams2views / use_views logic from robotics dataloaders. # TODO(bvh): debug/verify this # _get_ego_camera and _order_cameras return integer indices (matching dict keys). ego_cam = _get_ego_camera(dset_names[0], cam_names) used_cams = _order_cameras(dset_names[0], cam_names, V_max) V_used = len(used_cams) T_used = T_data # Re-anchor cameras to a reference view if configured. # This makes plucker rays relative to whichever camera was assigned to the reference view. ref_view = getattr(config, 'reference_view', False) if ref_view not in (False, None) and anydata_cams is not None: ref_cam_idx = used_cams[ref_view] per_timestep = getattr(config, 'ref_view_rel_per_timestep', True) anydata_cams = reanchor_cameras_to_view(anydata_cams, ref_cam_idx, per_frame=per_timestep) # ---- Task selection: forecast + world model (+ optional view-synth) ---- if 'tasks' in directives: is_forecast = 'forecast' in directives['tasks'] is_world_model = 'world_model' in directives['tasks'] is_view_synth = 'view_synth' in directives['tasks'] else: # View-synth is decided FIRST and mutually-exclusively vs the single/forecast # path: a sample is either Mode B (anchor view synth) or the existing Mode A / # forecast behavior. Guarded by p_vs (default 0.0) so legacy runs never enter here. is_view_synth = bernoulli_fn(p_vs) is_forecast = bernoulli_fn(p_fc) if not is_forecast: is_forecast = True # Fallback: always forecast # World model (trajectory conditioning) sampled independently is_world_model = bernoulli_fn(p_wm) # The whole "cap views per-sample" mechanism is ONLY engaged for anchor-mixing runs # (those that configure task_probs.view_synth > 0). For every other config (single-view # forecast, 11-view joint forecast, ...) p_vs == 0 -> this block is skipped entirely and # V_used / used_cams keep their legacy values, so behavior is byte-for-byte unchanged. # anchor_view: index of the view whose ref slot carries its own clean latent (the # I2V/interpolation anchor across the VIEW dimension). -1 => no anchor (ref slot zero # everywhere). Read by a4d_vae_wan to populate rgb{v}_ref. Default -1 keeps every # legacy/forecast Control run with y=0 (anchor absent -> refs zero). anchor_view = -1 view_synth_enabled = p_vs > 0.0 if view_synth_enabled: # View-synth requires exactly 2 views (anchor + target). If fewer cameras are # available in this sample, gracefully fall back to the single-generation path. if is_view_synth and V_used < 2: is_view_synth = False # Mode B (view-synth): exactly 2 views, BOTH generated/supervised. View 0 is the # ANCHOR (its ref slot = its own clean latent; view 1's ref = zero). View 1 attends # to the anchor's clean ref tokens in the 2-view joint attention. # Mode A (single): exactly 1 generated view, NO anchor (ref zero) -> attention stays # cheap (the experiment declares 2 video_entries but only view 0 is populated). if is_view_synth: used_cams = used_cams[:2] V_used = 2 anchor_view = 0 else: used_cams = used_cams[:1] V_used = 1 else: is_view_synth = False # Forecast: condition on first N frames, predict the rest if is_forecast: if 'cond_frames_raw' in directives: cond_frames_raw = directives['cond_frames_raw'] else: if 'cond_frames_choices' in directives: cond_frames_choices = directives['cond_frames_choices'] else: if phase == 'train': max_cond_frames = (T_used * 2) // 3 else: max_cond_frames = (T_used + 1) // 3 cond_frames_choices = list(range(1, max_cond_frames, VAE_RATIO_THW[0])) cond_frames_raw = random_choice(cond_frames_choices, device, generator) else: cond_frames_raw = T_used total_frames_raw = T_used cond_frames_latent = (cond_frames_raw - 1) // VAE_RATIO_THW[0] + 1 total_frames_latent = (total_frames_raw - 1) // VAE_RATIO_THW[0] + 1 log.debug(f"cond_frames_raw: {cond_frames_raw}, total_frames_raw: {total_frames_raw}, " f"cond_frames_latent: {cond_frames_latent}, total_frames_latent: {total_frames_latent}") # World model: trajectory is fully conditioned (input), not predicted if is_world_model: cond_frames_traj = T_used total_frames_traj = 0 else: cond_frames_traj = 0 total_frames_traj = 0 assert is_forecast, 'Forecast task must be active for driving' # ---- Populate Any4D entries ---- for v in range(V_used): # RGB. EVERY view (including view-synth Mode B) shares the temporal forecast split # cond=cond_frames_latent / pred=total -> rgb slot noisy + supervised. There is no # "given" view: the anchor is conditioned purely via its ref slot (filled by # a4d_vae_wan from anchor_view), NOT via the rgb slot. This keeps the rgb path # byte-for-byte identical between legacy/forecast and view-synth runs. cond_frames_rgb = cond_frames_latent # RGB if has_rgb_entries and (anydata_rgb is not None or latent_rgb is not None): (a4d_raw, a4d_latent) = prepare_dense( a4d_raw, a4d_latent, f'rgb{v}', anydata_rgb, latent_rgb, used_cams[v], cond_frames_rgb) # Control video (canny edge map for Wan Fun-Control). All frames are # read-through conditioning, so cond_frames_latent == total_frames_latent. if anydata_rgb_edge_control is not None and f'rgb{v}_edge_control' in config.all_highdim_entries: (a4d_raw, a4d_latent) = prepare_dense( a4d_raw, a4d_latent, f'rgb{v}_edge_control', anydata_rgb_edge_control, None, used_cams[v], total_frames_latent) # Cameras (plucker) if has_cams_entries and (anydata_cams is not None or latent_cams is not None): cond_frames_cams = total_frames_latent (a4d_raw, a4d_latent) = prepare_dense( a4d_raw, a4d_latent, f'cams{v}', anydata_cams, latent_cams, used_cams[v], cond_frames_cams, plucker_orig=False, scale_factor=scale_factors, clip_min=-1.0, clip_max=1.0) # Optional: mark the LAST latent frame as input as well (Wan Fun-InP "start+end" # mode). The default prepare_dense recipe is "first N frames as input" (mask = # [1]*N + [0]*(T-N)); setting cond_last_frame=True flips the last frame to also # be conditioned, producing mask = [1]*N + [0]*(T-N-1) + [1]. This shifts what # a4d_vae_wan._build_wan_ref_latent gray-pads (real pixels at first AND last raw # frames; gray in between) and ensures the final-frame loss is excluded. if directives.get('cond_last_frame', False): for v in range(V_used): for mask_key in (f'rgb{v}_input_mask', f'rgb{v}_output_mask'): if mask_key not in a4d_latent: continue if mask_key.endswith('_input_mask'): a4d_latent[mask_key][:, :, -1:] = 1.0 else: a4d_latent[mask_key][:, :, -1:] = 0.0 # ---- Trajectory from ego camera extrinsics ---- orig_traj = None if anydata_cams is not None: # TODO(bvh): debug/verify this # DEBUG: check ego camera position at t=0 to verify reference_camera re-anchoring if 0: ego_pos_t0 = anydata_cams[(0, ego_cam)].get_center() # (B, 3) log.debug(f'ego_cam={ego_cam} ({cam_names[ego_cam]}), pos@t0={ego_pos_t0[0].tolist()}') # DEBUG traj_multiplier = getattr(config, 'traj_multiplier', 1.0) trajectory = compute_trajectory_from_extrinsics( anydata_dict, avail_times, ego_cam, traj_multiplier) # trajectory: (B, T_used, 2) # Get expected token count from adaln config to pad shorter clips adaln_traj = config.lowdim_adaln_entries.get('traj', None) target_tokens = (adaln_traj[1] - adaln_traj[0]) if adaln_traj else None if 'perturb_traj' in directives: perturb_traj = directives['perturb_traj'] perturb_seed = directives.get('perturb_seed', None) orig_traj = trajectory.clone() if perturb_traj == 'basile1': trajectory = perturb_traj_basile1(orig_traj, cond_frames_raw, seed=perturb_seed) elif perturb_traj == 'basile2': trajectory = perturb_traj_basile2(orig_traj, cond_frames_raw, seed=perturb_seed) elif perturb_traj == 'basile3': trajectory = perturb_traj_basile3(orig_traj, cond_frames_raw, seed=perturb_seed) # NOTE(bvh): trajectory is input-only for now (pred_frames=0). (a4d_raw, a4d_latent) = prepare_traj_anydata( a4d_raw, a4d_latent, trajectory, cond_frames_traj, pred_frames=total_frames_traj, target_tokens=target_tokens) if orig_traj is None: orig_traj = a4d_raw['traj'].clone() # Store results data_dict['a4d_raw'] = a4d_raw data_dict['a4d_latent'] = a4d_latent # Metadata meta_tasks = [] meta_tasks += ['forecast'] if is_forecast else [] meta_tasks += ['world_model'] if is_world_model else [] meta_tasks += ['view_synth'] if is_view_synth else [] data_dict['meta'] = { 'used_modals': ','.join( m for m, v in [('rgb', anydata_rgb or latent_rgb), ('cams', anydata_cams or latent_cams), ('traj', orig_traj)] if v is not None), 'used_cams': ','.join(str(cam) for cam in used_cams), # a bit ugly but we should avoid changing standards 'tasks': ','.join(meta_tasks), # 'cond_modals': 'rgb', # TODO update 'cond_views_cams': V_used, 'cond_views_notcams': V_used, 'cond_frames_raw': cond_frames_raw, 'total_frames_raw': total_frames_raw, 'cond_frames_latent': cond_frames_latent, 'total_frames_latent': total_frames_latent, 'cond_frames_traj': cond_frames_traj, 'total_frames_traj': total_frames_traj, # View index whose ref slot carries its own clean latent (view-synth anchor), # or -1 for none. Read by a4d_vae_wan to fill rgb{v}_ref (anchor=clean latent, # others=zeros). Collated to a (B,) tensor; absent/-1 => all refs zero (legacy). 'anchor_view': anchor_view, } if orig_traj is not None: data_dict['meta']['orig_traj'] = orig_traj return data_dict