# Created by BVH & VG, Jun - Aug 2025. # Manages the pipeline and high-level training step (loss) and validation step (sampling) flow. import importlib import os from itertools import chain import numpy as np import torch from einops import rearrange from megatron.core import parallel_state from custom.dataloader.transforms import DataTransforms # NOTE: vidar imports (GenericModel, read_config) are deferred to setup_modules() # to avoid requiring vidar when using the anydata pipeline. from custom.eval.metrics import unroll_a4d from custom.utils.io import save_json import wandb from cosmos_predict2.functional.batch_ops import batch_mul from cosmos_predict2.models.video2world_model import Predict2Video2WorldModel, Predict2Video2WorldModelConfig # from custom.any4d.a4d_config import Any4DConfig # avoid due to circular import from custom.any4d.a4d_autoreg import AutoregressiveHandler from custom.any4d.a4d_logistics import ( add_dicts, assemble_clean_predictions, assemble_network_inputs, dict_prune_memory_recursive, dict_to_cuda_recursive, impute_data_batch_masks, multiply_dicts, pack_streams_from_entries, sanitize_nan_inf_streams, unpack_entries_from_streams, validate_masks, verify_entry_shapes, ) from custom.any4d.a4d_network import Any4DDiT from custom.any4d.a4d_metrics import calculate_metrics_a4d from custom.any4d.a4d_visuals import create_save_visuals_a4d from custom.eval.metrics import aggregate_all_metrics, average_metrics_per_dataset from custom.eval.visuals import tiered_detail from custom.any4d.a4d_logging import print_run_info from custom.utils.logging import print_batch_summary from custom.utils.debug import check_nan_dict from custom.wandb import prep_video from imaginaire.utils import log, misc from custom.dataloader.wm_utils import perturb_action_basile2, perturb_action_basile1 def apply_cond_aug(x0_streams, sigma_range): '''Apply conditioning augmentation: add Gaussian noise to conditioning (input) latents. Samples sigma independently per batch element from uniform[lo, hi]. Returns (x0_streams_noised, cond_aug_sigmas) where sigmas is (B,) or None if disabled.''' ca_lo, ca_hi = sigma_range if ca_hi <= 0.0: return (x0_streams, None) B = x0_streams['v0'].shape[0] device = x0_streams['v0'].device cond_aug_sigmas = ca_lo + (ca_hi - ca_lo) * torch.rand(B, device=device) # TODO(bvh): debug/verify cond aug rigorously # TODO(bvh): keep mask channels clean x0_streams_noised = {} for k, v in x0_streams.items(): sigma_shape = [B] + [1] * (v.dim() - 1) x0_streams_noised[k] = v + cond_aug_sigmas.view(sigma_shape) * torch.randn_like(v) return (x0_streams_noised, cond_aug_sigmas) # NOTE(bvh): This is the stepping stone from pure vanilla functionality to full Any4D: # Run the default, untouched Cosmos model, but apply vidar data loading & processing & Any4D visualizations. class VidarModel(Predict2Video2WorldModel): def __init__( self, config ): super().__init__(config) self.config = config assert self.config.vidar_active == True # Set collate verbosity from config from custom.dataloader.collate import set_collate_quiet set_collate_quiet(getattr(config, 'collate_quiet', False)) @property def _s3_path(self): '''Full S3 path for this run, derived from job config. Empty string if S3 sync disabled.''' s3_root = getattr(self.config.job, 's3_root', '') if s3_root: return f'{s3_root}/{self.config.job.path_local}' return '' # New method # used @ train & inference def setup_modules( self, ): ''' Instantiate transforms (vidar GenericModel or anydata DataTransforms) and VAE from self.config. Called once at train & inference startup before any data flows. ''' # Set up data transforms if self.config.data_library == 'vidar': # Legacy path for VidarDataset (deprecated, requires vidar YAML config) if self.config.transforms != '': cfg_path = f'custom/config/vidar/transforms/{self.config.transforms}.yaml' from vidar.arch.models.generic.GenericModel import GenericModel from vidar.utils.config import read_config self.transforms = GenericModel(read_config(cfg_path)) else: self.transforms = None else: # Anydata path: stateless transform engine (defaults/overrides managed by a4d_config). self.transforms = DataTransforms() # Set up VAE. Wan-specific VAE wrappers live under custom/wan/ (filename ending in # `_wan`); the original Cosmos-side wrappers live under custom/vae/. if getattr(self.config, 'data_only_disable_model', False): # Skip the real (heavyweight) VAE; use shape-preserving subsample/tile stub. from custom.vae.fake import VAE as FakeVAE self.vae = FakeVAE('cuda', self.config, self.pipe) elif self.config.vae != '': pkg = 'custom.wan' if self.config.vae.endswith('_wan') else 'custom.vae' module = importlib.import_module(f'{pkg}.{self.config.vae}') vae = getattr(module, 'VAE') self.vae = vae('cuda', self.config, self.pipe) else: self.vae = None # Set up dataloader (= essentially data-to-model interface, formerly known as vidar2a4d) if self.config.dataloader != '': dl_module = importlib.import_module(f"custom.dataloader.{self.config.dataloader}") if self.config.data_library == 'vidar': self.dataloader = dl_module.vidar_to_any4d else: self.dataloader = dl_module.anydata_to_any4d else: self.dataloader = None # used @ train & inference def run_modules( self, data_batch: dict[str, torch.Tensor], phase: str = 'train', train_step: int = -1, val_step: int = -1, directives: dict = dict(), ): ''' Prepares data batch for Any4D (device, transforms, dataloader, masks, VAE encoding, etc.) NOTE(bvh): This is called from training_step() after collate, so the batch dimension is active already, and generally B > 1. ''' data_batch['is_preprocessed'] = data_batch['is_preprocessed'][0].item() data_key = 'anydata' if 'anydata' in data_batch else 'vidar' # Apply transforms to the batch (this may override some parameters such as FPS). if data_key == 'anydata': data_batch = self._apply_transforms_anydata(data_batch, phase) else: data_batch = self._apply_transforms_vidar(data_batch, phase) # Convert data to Any4D model format if method is specified. # This step reads from the data library key and creates/writes # to the 'a4d_raw' AND/OR 'a4d_latent' keys. if self.dataloader is not None: data_batch = self.dataloader(data_batch, phase=phase, config=self.config, directives=directives, train_step=train_step, val_step=val_step) # For VRAM efficiency: data_batch[data_key] = dict_prune_memory_recursive(data_batch[data_key], threshold=1024) if self.config.check_nan: check_nan_dict(data_batch.get('a4d_raw', {}), 'a4d_raw (before VAE)', extra=f'dset={data_batch.get("dset_name", "?")}') # Impute potentially missing input, output, and supervise masks. # This step reads from and writes to the 'a4d_raw' AND 'a4d_latent' keys. if self.config.any4d_active: data_batch = impute_data_batch_masks(data_batch, self.config) # Encode (remaining) pixels / text / etc. into latents if needed. # This step reads from the 'a4d_raw' key and creates/writes # to the 'a4d_latent' key, for missing latent entries only. if self.vae is not None: data_batch = self.vae.encode_a4d_batch(data_batch, phase=phase) if train_step <= 50 and train_step % 13 == 1: print_batch_summary(data_batch, label=f'data_batch (step {train_step}, {phase})') if self.config.check_nan: check_nan_dict(data_batch.get('a4d_latent', {}), 'a4d_latent (after VAE)', extra=f'dset={data_batch.get("dset_name", "?")}') # Verify all encoded entry shapes (consistency with config) if self.config.any4d_active: verify_entry_shapes(data_batch['a4d_latent'], self.config) # In case of Cosmos video2world, this step simply prepares vanilla keys # to maintain backward compatibility. if not(self.config.any4d_active): if 'rgb0' in data_batch['a4d_raw']: # likely not present if loading cached data data_batch['video'] = data_batch['a4d_raw']['rgb0'] data_batch['video_latent'] = data_batch['a4d_latent']['rgb0'] data_batch['num_conditional_frames'] = data_batch['meta']['cond_frames_latent'] # ^ I believe num_conditional_frames is in latent space? (Cosmos code is ambiguous here) # Ensure everything is moved to GPU. data_batch = dict_to_cuda_recursive(data_batch) return data_batch def _apply_transforms_anydata(self, data_batch, phase): ''' Anydata path: per-sample params are already resolved in data_batch (extrinsics already reanchored in __getitem__). Align the temporal augs to plain ints -- frame_stride / frame_trim from sample[0] (uniform within batch), optionally broadcast rank 0 across ranks (sync_temporal_transforms, train only) to avoid forward-pass time-shearing -- then run the transforms (subsample + trim frames, scale fps). clip/scale_depth resolved here. ''' phase_overrides = self.config.data_train_overrides if 'train' in phase else self.config.data_val_overrides defaults = self.config.data_defaults # Temporal augs -> plain ints: sample[0] aligns within batch; broadcast aligns across ranks. frame_stride = int(data_batch['frame_stride'][0]) frame_trim = int(data_batch.get('frame_trim', [0])[0] or 0) if (self.config.sync_temporal_transforms and 'train' in phase and torch.distributed.is_initialized() and torch.distributed.get_world_size() > 1): synced = torch.tensor([frame_stride, frame_trim], device='cuda') torch.distributed.broadcast(synced, src=0) frame_stride, frame_trim = int(synced[0]), int(synced[1]) # frame_offset shifts which window is kept, not its length, so it needs no cross-rank broadcast. frame_offset = float(data_batch.get('frame_offset', [0.0])[0] or 0.0) clip_depth = phase_overrides.get('clip_depth', defaults.get('clip_depth', None)) scale_depth = phase_overrides.get('scale_depth', defaults.get('scale_depth', None)) if self.transforms is not None: data_batch['anydata'], data_batch['fps'] = self.transforms.prepare_batch( data_batch['anydata'], fps=data_batch['fps'], frame_stride=frame_stride, frame_trim=frame_trim, frame_offset=frame_offset, clip_depth=clip_depth, scale_depth=scale_depth, ) return data_batch def _apply_transforms_vidar(self, data_batch, phase): ''' Legacy vidar path: overrides applied here (VidarDataset getters are short-circuited). Kept as-is; new functionality belongs in the anydata path above. ''' phase_overrides = self.config.data_train_overrides if 'train' in phase else self.config.data_val_overrides rc_key = 'zero_origin' B = len(data_batch[rc_key]) rc_value = phase_overrides.get(rc_key, None) if rc_value is None: rc_value = data_batch[rc_key][0] # to align within batch frame_stride = phase_overrides.get('frame_stride', None) if frame_stride is None: frame_stride = data_batch['frame_stride'][0] # to align within batch # Homogenize certain values across the batch (use the first sample). for i in range(B): data_batch[rc_key][i] = rc_value # single str or bool data_batch['frame_stride'][i] = frame_stride # single int # Apply frame rate / prompt / scale_factor overrides if specified. override_fps = phase_overrides.get('frame_rate', None) if override_fps is not None: data_batch['fps'][:] = override_fps # (B) tensor of float override_prompt = phase_overrides.get('prompt', None) if override_prompt is not None: data_batch['prompt'][:] = override_prompt # (B) tensor of str override_sf = phase_overrides.get('scale_factor', None) if override_sf is not None: data_batch['scale_factor'][:] = override_sf # (B) tensor of float # Apply important frame rate correction (per sample separately). data_batch['fps'] = data_batch['fps'] / data_batch['frame_stride'] # ^ (B) tensor of float / single int => (B) tensor of float # Apply data augmentations on GPU if class is specified (legacy GenericModel). if self.transforms is not None: data_batch['vidar'] = self.transforms.prepare_batch( data_batch['vidar'], zero_origin=rc_value, stride=frame_stride, ) return data_batch def is_image_batch( self, data_batch: dict[str, torch.Tensor] ) -> bool: return False def _update_train_stats( self, data_batch: dict[str, torch.Tensor] ) -> None: self.pipe.dit.accum_video_sample_counter += data_batch['a4d_latent']['rgb0'].shape[0] * self.data_parallel_size # NOTE(bvh): ^ This value is around 4575657222473777152 upon starting finetuning, # which seems impossibly large? def on_train_start( self, memory_format: torch.memory_format = torch.preserve_format ) -> None: super().on_train_start(memory_format) # NOTE(bvh): To understand different parallelism groups & ranks correctly. # Simply loop over all methods in parallel_state and print the getters that return a rank or size. my_rank_td = torch.distributed.get_rank() world_size_td = torch.distributed.get_world_size() for method in sorted(dir(parallel_state)): if method.startswith('get_') and (method.endswith('_rank') or method.endswith('_size')) \ and not('_prev_' in method or '_next_' in method): my_value = getattr(parallel_state, method)() if my_value is None: continue # Gather across ranks for more concise printing. my_tensor = torch.tensor([my_value], device='cuda') gathered_tensors = [torch.zeros_like(my_tensor) for _ in range(world_size_td)] torch.distributed.all_gather(gathered_tensors, my_tensor) all_values = [tensor.item() for tensor in gathered_tensors] # if my_rank_td == 0: # this is rank 0 only by default log.debug(f'parallel_state.{method}() = {all_values}') # Overridden by Any4DModel def training_step( self, data_batch: dict[str, torch.Tensor], iteration: int, local_path: str = None, ) -> tuple[dict, torch.Tensor]: ''' TODO expand docstring :param iteration (int): Batch index, follows training step. :param local_path (str): Directory to save results to. ''' train_step = iteration my_rank_td = torch.distributed.get_rank() # Call Predict2Video2WorldModel.training_step() (train_dict, kendall_loss) = super().training_step(data_batch, iteration, local_path) # In the beginning, we want to save visuals twice as frequently. max_detail = self.config.train_visuals_detail interval = self.config.train_visuals_interval want_visuals = (interval >= 1 and max_detail >= 1 and (train_step % interval == 1 or train_step <= 2 or (train_step <= interval * 10 and train_step % (interval // 2) == 1))) if want_visuals: # NOTE(bvh): input_mask and output_mask should already be constructed by anydata_to_a4d, # but supervise_mask needs to be inferred here for visuals: data_batch['a4d_latent']['rgb0_supervise_mask'] = data_batch['a4d_latent']['rgb0_output_mask'] # Replicate other expected Any4D structure for metrics & visuals as well: x0_rgb0 = data_batch['a4d_latent']['rgb0'].detach().clone() y0_pred_rgb0 = train_dict['model_pred'].x0.detach().clone() sigma = train_dict['sigma'].detach().clone() train_dict.update({ 'xt_streams': { 'v0': y0_pred_rgb0, # (B, 16, T, H, W) }, 'xt_entries': { 'rgb0': y0_pred_rgb0, # (B, 16, T, H, W) }, 'y0_pred_streams': { 'v0': y0_pred_rgb0, # (B, 16, T, H, W) }, 'y0_pred_entries': { 'rgb0': y0_pred_rgb0, # (B, 16, T, H, W) }, 'x0_streams': { 'v0': x0_rgb0, # (B, 16, T, H, W) }, 'x0_entries': { 'rgb0': x0_rgb0, # (B, 16, T, H, W) }, 'sigmas': { 'v0': sigma, # (B, 1) }, }) # NOTE(yams_any4d 2026-07-13): skip the video gallery (VAE-decode of predicted RGB + # mp4 assembly/upload) and psnr/ssim metrics when we only care about the head. The # backbone is frozen, so generating video previews every interval is wasted training # time — keep only the cheap 2D V4Head panels. skip_vid = getattr(self.config, 'skip_video_visuals', False) metrics_dict = {} if skip_vid else calculate_metrics_a4d( 'train', train_step, data_batch, train_dict, self.vae, self.config) # NOTE(yams_any4d): V4Head panels (PCA / keypoints / bin grids), v4-trainer style. if getattr(self.config, 'v4head_enabled', False) and 'v4head_out' in train_dict: from custom.any4d.v4_head_viz import make_v4head_visuals _dit = self.pipe.dit _head = getattr(_dit, 'v4head', None) or getattr(getattr(_dit, 'module', _dit), 'v4head', None) if _head is not None: _v4h_panels = make_v4head_visuals( _head, train_dict['v4head_out'], data_batch, f'{local_path}/visuals/train/v4head', train_step) else: _v4h_panels = {} else: _v4h_panels = {} if skip_vid: nice_videos = dict(_v4h_panels) else: nice_videos = create_save_visuals_a4d( 'train', train_step, data_batch, train_dict, self.vae, self.config, kendall_loss, metrics_dict, f'{local_path}/visuals/train', detail=tiered_detail(train_step, max_detail, interval), quality_pix=self.config.visuals_quality, quality_lat=self.config.visuals_quality + 1, gc_collect=True) nice_videos.update(_v4h_panels) # NOTE(yams_any4d) print_batch_summary(data_batch, label=f'data_batch (train s{train_step})') print_run_info('train', train_step, local_path, self._s3_path) # Gets uploaded by trainer.py later: wandb_logs = self.prepare_wandb_logs_step( 'train', train_step, data_batch, train_dict, metrics_dict, nice_videos) train_dict['wandb_logs'] = wandb_logs # BVH shape notes (vanilla): # see Predict2Video2WorldModel.training_step() return (train_dict, kendall_loss) def on_validation_start( self, iteration: int = -1, ) -> None: ''' :param iteration (int): Batch index, follows training step. ''' self.all_val_metrics = [] # list of metrics_dict # Overridden by Any4DModel @torch.no_grad() def validation_step( self, data_batch: dict[str, torch.Tensor], iteration: int = -1, dataloader_key: str = None, local_path: str = None, directives: dict = None, val_iter: int = 0, seed: int = 0, ) -> tuple[dict[str, torch.Tensor], torch.Tensor]: ''' TODO expand docstring :param iteration (int): Batch index, follows training step. :param local_path (str): Directory to save results to. :param val_iter (int): Validation step, changes while iteration remains constant. :param seed (int): Random seed for diffusion sampling (increment for diversity across samples) ''' train_step = iteration val_step = val_iter self.pipe.device = self.device # Prepare data batch (device, transforms, Any4D entries, masks, VAE encoding, etc) data_batch = self.run_modules(data_batch, phase='val', train_step=train_step, val_step=val_step, directives=directives) data_batch['dl_key'] = dataloader_key # Wan training uses `text_encoder_path=''` and loads cached positive+negative # UMT5 embeddings from default_text_wan.pkl. Use is_negative_prompt=True there # so CFG extrapolates "positive − UMT5('')" instead of "positive − zero_tensor" # (zero tensor pulls generation toward an OOD attractor at high guidance). _is_neg = getattr(self.config, 'use_wan_backbone', False) # this calls Video2WorldPipeline.process(): y0_pred_rgb0 = self.pipe.process( data_batch, seed=seed, guidance=self.config.val_cfg_scale, num_sampling_step=self.config.val_num_steps, return_latents=True, is_negative_prompt=_is_neg, ) loss = None # NOTE(bvh): input_mask and output_mask should already be constructed by anydata_to_a4d, # but supervise_mask needs to be inferred here for visuals: data_batch['a4d_latent']['rgb0_supervise_mask'] = data_batch['a4d_latent']['rgb0_output_mask'] # Replicate other expected Any4D structure for metrics & visuals as well: x0_rgb0 = data_batch['a4d_latent']['rgb0'].detach().clone() val_dict = { 'xt_streams': { 'v0': y0_pred_rgb0, # (B, 16, T, H, W) }, 'xt_entries': { 'rgb0': y0_pred_rgb0, # (B, 16, T, H, W) }, 'y0_pred_streams': { 'v0': y0_pred_rgb0, # (B, 16, T, H, W) }, 'y0_pred_entries': { 'rgb0': y0_pred_rgb0, # (B, 16, T, H, W) }, 'x0_streams': { 'v0': x0_rgb0, # (B, 16, T, H, W) }, 'x0_entries': { 'rgb0': x0_rgb0, # (B, 16, T, H, W) }, } metrics_dict = calculate_metrics_a4d( 'val', train_step, data_batch, val_dict, self.vae, self.config) nice_videos = create_save_visuals_a4d( 'val', train_step, data_batch, val_dict, self.vae, self.config, loss, metrics_dict, f'{local_path}/visuals', detail=self.config.val_visuals_detail, quality_pix=self.config.visuals_quality, quality_lat=self.config.visuals_quality + 1, val_step=val_step, batch_idx='all', dataloader_key=dataloader_key, gc_collect=(val_step % 5 == 0)) print_run_info('val', train_step, local_path, self._s3_path) # Gets uploaded by trainer.py later: wandb_logs = self.prepare_wandb_logs_step( 'val', train_step, data_batch, val_dict, metrics_dict, nice_videos) val_dict['wandb_logs'] = wandb_logs self.all_val_metrics.append(metrics_dict) # BVH shape notes (vanilla): # data_batch: dict_keys(['dl_idx', 'anydata', 'dset_name', 'sample_id', 'fps', 'prompt', # 'scale_factor', 'image_size', 'num_frames', 'padding_mask', 'guidance', 'is_preprocessed', # 'is_anydata', 'a4d_raw', 'a4d_latent', 'meta', 't5_text_embeddings', 't5_text_mask', # 'video', 'video_latent', 'num_conditional_frames']) # metrics_dict: {'psnr': {'rgb0': [...]}, 'ssim': {'rgb0': [...]}, 'mse': {}, 'mae': {}, 'friendly': {'psnr': [...], 'ssim': [...], 'mse': [...], 'mae': [...]}} # nice_videos: {'pix': '/basile/ws/repos/cpred2-basile/logs_a4d/vidar/debug/08-14-00-34_dv2/visuals/dv2_val_s0_r0_lbmv12_i0_v0_all.mp4'} # wandb_logs: TBD return (val_dict, loss) def prepare_wandb_logs_step( self, phase: str = 'train', train_step: int = -1, data_batch: dict[str, torch.Tensor] = None, output_dict: dict[str, torch.Tensor] = None, metrics_dict: dict[str, list[float]] = None, nice_videos: dict[str, str] = None, dataloader_key: str = None, ) -> dict[str, float]: ''' NOTE(bvh): In practice, these items will only get uploaded on global rank 0, so we should check averages and S3 for the rest. :param phase (str): train / val / test. :param train_step (int): Batch index, follows training step. :param data_batch (dict): Input batch, contains vidar, fps, prompt, etc. :param output_dict (dict): train_dict / val_dict. :param nice_videos (dict): Maps video descriptor to path. :param metrics_dict (dict): Maps metric name to list of values. :return wandb_logs (dict): Maps wandb key to appropriate value. ''' wandb_logs = dict() # Prepare individual losses, overall and per entry if 'loss' in output_dict: wandb_logs[f'loss_{phase}/total'] = output_dict['loss'].mean().item() if 'edm_loss_means' in output_dict: for (k, v) in output_dict['edm_loss_means'].items(): # Any4D entries cur_loss_val = v.mean().item() if cur_loss_val > 0.0: wandb_logs[f'loss_{phase}/{k}'] = cur_loss_val # Prepare individual quantitative metrics for wandb # metrics_dict is metric -> entry -> list of float/None # becomes e.g. 'metrics_val/lbmv12/rgb1/ssim--32': 0.43393746 for (metric, entries_dict) in metrics_dict.items(): if metric in ('friendly', 'dset_name', 'dl_idx', 'dl_key', 'sample_id', 'true_val_size'): continue if 'friendly' in metric: continue for (entry, values_list) in entries_dict.items(): for i in range(len(values_list)): if values_list[i] is not None: dset_name = metrics_dict['dset_name'][i] dl_idx = metrics_dict['dl_idx'][i] if dataloader_key is not None: wandb_key = f'metrics_{phase}_each/{dataloader_key}--{dset_name}/{entry}/{metric}--{dl_idx}' else: wandb_key = f'metrics_{phase}_each/{dset_name}/{entry}/{metric}--{dl_idx}' wandb_logs[wandb_key] = values_list[i] # Prepare individual qualitative videos for wandb (via local file paths) # Keys in nice_videos are already formatted like dl_key--dset_name/A4D--dl_idx # becomes e.g. 'visuals_val/lbmv12/pix--32': for (video_key, video_path) in nice_videos.items(): # if dataloader_key is not None: # wandb_key = f'visuals_{phase}/{dataloader_key}/{video_key}' # else: wandb_key = f'visuals_{phase}/{video_key}' # NOTE(yams_any4d): static panels (V4Head viz) are pngs -> wandb.Image. # Guarded: a bad media file must never kill training. try: if str(video_path).endswith('.png'): wandb_logs[wandb_key] = wandb.Image(video_path) else: wandb_logs[wandb_key] = wandb.Video(video_path) except Exception as _e: print(f'[yams_any4d] WARNING: skipping wandb media {video_path}: {_e}') return wandb_logs def on_validation_end( self, iteration: int = -1, local_path: str = None, ) -> dict[str, float]: ''' :param iteration (int): Batch index, follows training step. :param local_path (str): Run output directory (from trainer). ''' train_step = iteration # Average & report validation set metrics wandb_logs = self.prepare_wandb_logs_end( phase='val', all_metrics=self.all_val_metrics, train_step=train_step, local_path=local_path) return wandb_logs def prepare_wandb_logs_end( self, phase: str = 'val', all_metrics: list[dict[str, dict[str, list[float]]]] = None, train_step: int = -1, local_path: str = None, ): my_rank_td = torch.distributed.get_rank() wandb_logs = dict() # First collect all lists across all processes / ranks world_size = torch.distributed.get_world_size() world_metrics = [None] * world_size torch.distributed.all_gather_object(world_metrics, all_metrics) world_metrics_flat = list(chain(*world_metrics)) # Results here should only be uploaded on global rank 0, so return empty dict if my_rank_td != 0: return wandb_logs # Then convert long list of metrics_dict into a single metrics_dict # (values will become long lists of tuples) final_agg_metrics = aggregate_all_metrics(world_metrics_flat) # Convert back to average metrics avg_dset_metrics = average_metrics_per_dataset(final_agg_metrics) # Prepare averaged quantitative metrics for wandb # avg_dset_metrics is metric -> entry -> dset -> (avg: float, count: int) # becomes e.g. 'metrics_val/lbmv12/rgb1/SSIM--avg': 0.43393746 for (metric, entries_dict) in avg_dset_metrics.items(): if 'friendly' in metric or 'dset_name' in metric or 'dl_idx' in metric: continue for (entry, dset_dict) in entries_dict.items(): for (dset_name, (avg, count)) in dset_dict.items(): if count > 0: wandb_key = f'metrics_{phase}/{dset_name}/{entry}/{metric.upper()}' wandb_logs[wandb_key + '--avg'] = avg wandb_logs[wandb_key + '--count'] = count # Print and save aggregated metrics locally (rank 0 only) log.info('=== Validation metrics summary ===') local_summary = {} for (metric, entries_dict) in avg_dset_metrics.items(): if 'friendly' in metric or 'dset_name' in metric or 'dl_idx' in metric: continue for (entry, dset_dict) in entries_dict.items(): for (dset_name, (avg, count)) in dset_dict.items(): if count > 0: key = f'{dset_name}/{entry}/{metric.upper()}' local_summary[key] = {'avg': float(avg), 'count': int(count)} log.info(f' {key}: {avg:.4f} (n={count})') # Save JSON iter_tag = f'_iter{train_step:06d}' if train_step >= 0 else '' metrics_path = os.path.join(local_path, f'val_metrics_summary_basile{iter_tag}.json') try: save_json(local_summary, metrics_path, indent=2) log.info(f'Saved metrics JSON to {metrics_path}') except Exception as e: log.warning(f'Could not save metrics JSON: {e}') return wandb_logs class Any4DModel(VidarModel): def __init__( self, config ): super().__init__(config) assert self.config.any4d_active == True # Initialize autoregressive handler for multi-segment inference self.autoreg_handler = AutoregressiveHandler(config) # NaN loss tracking for safeguard in compute_loss self._nan_loss_count = 0 self._total_loss_count = 0 def _reduce_entries_and_filter_active(self, edm_loss_entries, train_step): '''Per-entry reduction (tensor -> scalar via mean) + active-entry filtering. Does not reduce across entries yet (see compute_loss). 1) Mean per entry (supervise mask is already baked into edm_loss_entries). 2) Filter NaN entries (should not happen after input sanitization, but log loudly if it does). 3) Filter inactive entries (all-zero, e.g. padded/stub views). Returns (edm_loss_means, active_edm_loss_means). NOTE / WARNING(bvh): entries with fewer active tokens are upweighted relative to larger entries.''' # Step 1: mean per entry edm_loss_means = {k: v.mean() for k, v in edm_loss_entries.items()} # Step 2: separate into active / nan / inactive active = {} nan_entries = [] for k, v in edm_loss_means.items(): if v.isnan(): nan_entries.append(k) elif not v.any(): continue # inactive (zero supervise mask) else: active[k] = v # Step 3: NaN safeguard self._total_loss_count += 1 if nan_entries: self._nan_loss_count += 1 my_rank = torch.distributed.get_rank() log.error( f'[RANK{my_rank}] NaN in EDM loss at iter {train_step}! ' f'entries: {nan_entries}, ' f'NaN count: {self._nan_loss_count}/{self._total_loss_count}', rank0_only=False) if self._total_loss_count >= 100 and self._nan_loss_count / self._total_loss_count >= 0.10: raise RuntimeError( f'NaN loss rate {self._nan_loss_count}/{self._total_loss_count} ' f'({100*self._nan_loss_count/self._total_loss_count:.1f}%) exceeds 10% threshold') if len(active) == 0: raise RuntimeError( f'No active entries for EDM loss averaging! ' f'all means: {edm_loss_means}, nan: {nan_entries}') return (edm_loss_means, active) def _build_stub_output_dict(self, x0_streams, masks): '''Data-only debug helper: shape-correct dict with y0_pred = x0 / 2, zero loss. Skips conditioner / denoise / compute_loss so dataloader + visuals still run. Output streams = all x0 streams minus xattn / adaln (matches compute_loss filter). ''' device = next(iter(x0_streams.values())).device output_keys = [k for k in x0_streams.keys() if k not in ('xattn', 'adaln')] y0_pred_streams = {k: x0_streams[k] / 2.0 for k in output_keys} xt_streams = {k: v.clone() for k, v in x0_streams.items()} x0_entries = unpack_entries_from_streams(x0_streams, self.config, detach=True, ignore_masks=False) y0_pred_entries = unpack_entries_from_streams(y0_pred_streams, self.config, detach=True, ignore_masks=True) xt_entries = unpack_entries_from_streams(xt_streams, self.config, detach=True, ignore_masks=True) # Zero-filled sigmas so a4d_visuals.create_save_visuals_a4d (line 397) can index # `output_dict['sigmas'][k][batch_idx].item()` and produce a filename suffix. B = next(iter(x0_streams.values())).shape[0] sigmas = {k: torch.zeros(B, device=device) for k in output_keys} return { 'x0_streams': x0_streams, 'x0_entries': x0_entries, 'masks': masks, 'xt_streams': xt_streams, 'xt_entries': xt_entries, 'y0_pred_streams': y0_pred_streams, 'y0_pred_entries': y0_pred_entries, 'edm_loss_means': {}, # so prepare_wandb_logs_step's loop is a no-op. 'loss': torch.zeros((), device=device, dtype=torch.float32), 'sigmas': sigmas, } def _apply_action_override( self, data_batch: dict[str, torch.Tensor], method: str, seed: int = 0, ) -> None: ''' Apply an action perturbation directly on data_batch (after run_modules), preserving struct keys and shapes. Only modifies existing values of 'action' in a4d_latent and a4d_raw. :param data_batch (dict): Input batch, contains a4d_latent and a4d_raw. :param method (str): Perturbation method to use ('basile1', 'basile2', or None). :param seed (int): Random seed for perturbation. ''' if method is None: return # a4d_latent action (preferred) act = data_batch['a4d_latent']['action'] act_raw = data_batch['a4d_raw']['action'] # Use cond_frames_raw from meta cond_frames_raw = int(data_batch['meta']['cond_frames_raw']) if method == 'basile2': data_batch['a4d_latent']['action'] = perturb_action_basile2(act, cond_frames_raw, seed=seed) data_batch['a4d_raw']['action'] = perturb_action_basile2(act_raw, cond_frames_raw, seed=seed) elif method == 'basile1': data_batch['a4d_latent']['action'] = perturb_action_basile1(act, cond_frames_raw, seed=seed) data_batch['a4d_raw']['action'] = perturb_action_basile1(act_raw, cond_frames_raw, seed=seed) # Overrides Predict2Video2WorldModel # used @ train only def training_step( self, data_batch: dict[str, torch.Tensor], iteration: int = -1, local_path: str = None, directives: dict = None, ) -> tuple[dict, torch.Tensor]: ''' :param data_batch (dict): Input batch, contains vidar, fps, prompt, etc. :param iteration (int): Batch index, follows training step. :param local_path (str): Directory to save results to. :return: train_dict (dict): Output samples batch. kendall_loss (tensor): Single float to backprop with. ''' train_step = iteration my_rank_td = torch.distributed.get_rank() self.pipe.device = self.device # Reset peak-memory tracker so the max_memory_allocated() read below reflects this # step's forward+loss (backward happens after this function returns). if torch.cuda.is_available(): torch.cuda.reset_peak_memory_stats() # BVH shape notes (any4d): # data_batch (before run_modules): dict_keys(['dl_idx', 'anydata', 'dset_name', 'sample_id', # 'fps', 'prompt', 'scale_factor', 'image_size', 'num_frames', 'padding_mask', 'guidance', # 'is_preprocessed', 'is_anydata']) # local_path: '/basile/ws/repos/cpred2-basile/logs_a4d/a4d2/debug/08-13-16-18_b2' # Prepare data batch (device, transforms, Any4D entries, masks, VAE encoding, etc) data_batch = self.run_modules(data_batch, phase='train', train_step=train_step, directives=directives) # Process the data batch to assemble all inputs, outputs, masks, and conditions. (x0_streams, masks) = pack_streams_from_entries(data_batch['a4d_latent'], self.config) # Periodically verify masks if train_step % 111 == 1: validate_masks(masks) if getattr(self.config, 'data_only_disable_model', False): # Data-only debug: skip conditioner / denoise / loss; build stub train_dict. # _update_train_stats is also skipped because it touches self.pipe.dit, which is # None when data_only_disable_model implies remove_dit=True (see video2world.py). train_dict = self._build_stub_output_dict(x0_streams, masks) kendall_loss = train_dict['loss'] else: # Update video counter self._update_train_stats(data_batch) # Apply conditioner (now only used for text & random dropout thereof). condition = self.pipe.apply_conditioner(data_batch) # Sample pertubation noise levels and N(0, 1) noise vectors. # NOTE on naming: The main difference between x and y dicts is that x contains all streams # (so both input and output/GT), but y only contains output/GT streams. # x0_shapes = {k: v.shape for k, v in x0_streams.items()} y0_shapes = {k: v.shape for k, v in x0_streams.items() if not k in ['xattn', 'adaln']} (sigmas, epsilons) = self.draw_training_noise( y0_shapes, harmonize_streams=self.config.harmonize_streams, harmonize_frames=self.config.harmonize_frames) # assert self.pipe.get_context_parallel_group() is None, \ # 'Please see broadcast_split_for_model_parallelsim() [sic]' self.pipe.broadcast_split_for_model_parallelism() # Compute loss with epsilons and sigmas. sample_info = list(zip(data_batch.get('dset_name', []), data_batch.get('sample_id', []))) train_dict = self.compute_loss( x0_streams, masks, condition, epsilons, sigmas, train_step=train_step, sample_info=sample_info) # NOTE(yams_any4d): V4Head aux losses on the stashed DiT features. if getattr(self.config, 'v4head_enabled', False): train_dict = self._v4head_compute(data_batch, train_dict, train_step) # Reduce & scale loss as desired assert self.loss_reduce == 'mean', 'Not yet supported in Any4D context' kendall_loss = train_dict['loss'] # single float. kendall_loss *= self.loss_scale assert kendall_loss.dtype == torch.float32, \ f'kendall_loss should be float32, got {kendall_loss.dtype}' # In the beginning, we want to save visuals twice as frequently. max_detail = self.config.train_visuals_detail interval = self.config.train_visuals_interval want_visuals = (interval >= 1 and max_detail >= 1 and (train_step % interval == 1 or train_step <= 2 or (train_step <= interval * 10 and train_step % (interval // 2) == 1))) if want_visuals: metrics_dict = calculate_metrics_a4d( 'train', train_step, data_batch, train_dict, self.vae, self.config) nice_videos = create_save_visuals_a4d( 'train', train_step, data_batch, train_dict, self.vae, self.config, kendall_loss, metrics_dict, f'{local_path}/visuals/train', detail=tiered_detail(train_step, max_detail, interval), quality_pix=self.config.visuals_quality, quality_lat=self.config.visuals_quality + 1, gc_collect=True) # NOTE(yams_any4d): V4Head panels (PCA / keypoints / bin grids), v4-trainer style. if getattr(self.config, 'v4head_enabled', False) and 'v4head_out' in train_dict: from custom.any4d.v4_head_viz import make_v4head_visuals _dit = self.pipe.dit _head = getattr(_dit, 'v4head', None) or getattr(getattr(_dit, 'module', _dit), 'v4head', None) if _head is not None: nice_videos.update(make_v4head_visuals( _head, train_dict['v4head_out'], data_batch, f'{local_path}/visuals/train/v4head', train_step)) elif getattr(self.config, 'v4head_enabled', False): print(f'[v4head_viz] NOTE: v4head_out missing from train_dict at step {train_step}') print_batch_summary(data_batch, label=f'data_batch (train s{train_step})') print_run_info('train', train_step, local_path, self._s3_path) # Gets uploaded by trainer.py later: wandb_logs = self.prepare_wandb_logs_step( 'train', train_step, data_batch, train_dict, metrics_dict, nice_videos) train_dict['wandb_logs'] = wandb_logs # BVH shape notes (any4d): # train_dict: see compute_loss() # kendall_loss: tensor grad MeanBackward0 cuda:0 239.801 # metrics_dict: {'psnr': {'rgb0': [...], 'rgb1': [...]}, 'ssim': {'rgb0': [...], 'rgb1': [...]}, 'mse': {'depth0': [...], 'depth1': [...], 'action': [...]}, 'mae': {'depth0': [...], 'depth1': [...], 'action': [...]}, 'friendly': {'psnr': [...], 'ssim': [...], 'mse': [...], 'mae': [...]}} # nice_videos: {'pix': '/basile/ws/repos/cpred2-basile/logs_a4d/a4d2/debug/08-13-16-18_b2/visuals/b2_val_s1_r0_lbmv12_d885044_v1,0_all.mp4'} # Peak VRAM for this step (forward + loss; backward happens after return) if torch.cuda.is_available(): peak_gb = torch.cuda.max_memory_allocated() / 1e9 log.info(f'[step {train_step} rank{my_rank_td}] peak VRAM: {peak_gb:.1f} GB', rank0_only=(train_step > 5)) return (train_dict, kendall_loss) # Replaces Predict2Video2WorldModel.draw_training_sigma_and_epsilon() # used @ train only def draw_training_noise( self, y0_shapes: dict[str, torch.Tensor], harmonize_streams: bool = True, harmonize_frames: bool = True, ): ''' :param y0_shapes (dict): Maps stream name to (B, C, T, H, W) or (B, T, C). :param harmonize_streams (bool): Whether to use equal noise levels (sigma) across streams. :param harmonize_frames (bool): Whether to use equal noise levels (sigma) across frames. # :return sigmas (dict): Maps stream name to (B, 1) tensor of float. :return sigmas (dict): Maps stream name to (B, 1/T, 1) or (B, 1, 1/T, 1, 1) tensor of float. :return epsilons (dict): Maps stream name to (B, C, T, H, W) or (B, T, C) tensor of float. ''' sigmas = dict() epsilons = dict() assert harmonize_frames, 'Cosmos does not support harmonize_frames = False by default, to be implemented in Any4D if deemed necessary' # Loop over streams and call existing method # NOTE(bvh): sigma values depend on state_t via video_noise_multiplier! for (k, y0_shape) in y0_shapes.items(): (sigmas[k], epsilons[k]) = self.draw_training_sigma_and_epsilon(y0_shape, None) # ^ sigma returned as (B, T) but with T = 1 for now # NOTE(yams_any4d 2026-07-13): for the live frozen-head run, override the diffusion # timestep to a small LOW band (the "last generation timestep") instead of the full EDM # schedule. The backbone is frozen so the RGB EDM loss is inert — there is no reason to # ever feed the head high-sigma (near-pure-noise) features. Clamping to a tight low band # gives the head consistently informative features (like the extraction cache's fixed # sigma=0.05) and removes the huge gradient variance that destabilized the full-schedule # run. Log-uniform in [lo, hi]. Applied BEFORE harmonize so all streams share one value. band = getattr(self.config, 'v4head_train_sigma_band', None) if band is not None: (lo, hi) = (float(band[0]), float(band[1])) for k in sigmas.keys(): s = sigmas[k] log_lo = torch.log(torch.as_tensor(lo, device=s.device, dtype=s.dtype)) log_hi = torch.log(torch.as_tensor(hi, device=s.device, dtype=s.dtype)) sigmas[k] = torch.exp(torch.rand_like(s) * (log_hi - log_lo) + log_lo) # Harmonize noise levels if desired, keeping noise vectors intact if harmonize_streams: for k in sigmas.keys(): sigmas[k] = sigmas['v0'] # Expand noise level shapes to match corresponding stream for (k, y0_shape) in y0_shapes.items(): if len(y0_shape) == 3: # (B, T, C) sigmas[k] = sigmas[k][:, :, None] # (B, 1/T, 1) for low-dim elif len(y0_shape) == 5: # (B, C, T, H, W) sigmas[k] = sigmas[k][:, None, :, None, None] # (B, 1, 1/T, 1, 1) for high-dim else: raise ValueError(f'Unexpected shape: {y0_shape}') # BVH shape notes (any4d): # sigmas: { # 'v0': tensor[2, 1] μ=7.094 σ=2.201 cuda:0 [[5.538], [8.651]], # 'v1': tensor[2, 1] μ=5.920 σ=0.617 cuda:0 [[6.357], [5.484]], # 'sattn': tensor[2, 1] μ=5.582 σ=2.179 cuda:0 [[4.041], [7.123]], # 'xattn': tensor[2, 1] μ=10.524 σ=8.962 cuda:0 [[4.187], [16.862]], # 'adaln': tensor[2, 1] μ=1.192e+04 σ=1.685e+04 cuda:0 [[2.383e+04], [3.642]] # } # epsilons: { # 'v0': tensor[2, 70, 7, 16, 20] n=313600 (1.2Mb) x∈[-4.727, 4.856] μ=0.002 σ=1.001 cuda:0, # 'v1': tensor[2, 70, 7, 16, 20] n=313600 (1.2Mb) x∈[-4.719, 4.347] μ=-0.001 σ=1.001 cuda:0, # 'sattn': tensor[2, 25, 22] n=1100 (4.3Kb) x∈[-3.169, 2.900] μ=-0.021 σ=1.003 cuda:0, # 'xattn': tensor[2, 30, 14] n=840 (3.3Kb) x∈[-2.847, 3.418] μ=-0.005 σ=0.997 cuda:0, # 'adaln': tensor[2, 7, 21] n=294 (1.1Kb) x∈[-2.712, 2.113] μ=0.013 σ=0.951 cuda:0 # } return (sigmas, epsilons) # NOTE(yams_any4d): V4Head — run the volumetric action head on the DiT features # stashed by a4d_network this forward, add weighted CE losses to train_dict['loss'], # and expose per-loss means so prepare_wandb_logs_step logs them for free. def _v4head_compute(self, data_batch, train_dict, train_step): dit = self.pipe.dit head = getattr(dit, 'v4head', None) or getattr(getattr(dit, 'module', dit), 'v4head', None) feats = getattr(dit, '_v4head_feats', None) # GT MUST come from the unmasked copy: a4d_raw['action'] is zeroed under # pure-forecast task masks (found 2026-07-10, cost one invalid training run). action_raw = data_batch.get('a4d_raw', {}).get('action_gt_full', None) if action_raw is None: action_raw = data_batch.get('a4d_raw', {}).get('action', None) if action_raw is not None and train_step <= 2: print('[v4head] WARNING: action_gt_full missing — falling back to masked ' 'a4d_raw[action]; targets are ZERO under forecast tasks!') if head is None or feats is None or 0 not in feats or action_raw is None: return train_dict # NOTE(yams_any4d 2026-07-12): optionally detach features so the (freshly-init) head # trains on the video model AS A FROZEN BACKBONE and its large early gradients do NOT # flow back into the LoRA/DiT. Without this, joint training DIVERGES (loss 8->1e15 # after warmup): the random head destabilizes the well-trained video model. The video # LoRA still adapts on its own RGB loss. if getattr(self.config, 'v4head_detach_features', False): feats = {k: v.detach() for (k, v) in feats.items()} (Hp, Wp) = feats[0].shape[2], feats[0].shape[3] pixel_hw = (Hp * 16, Wp * 16) # tokens -> pixels: patch_spatial(2) x VAE spatial(8) head_out = head(feats, pixel_hw, action_raw.to(feats[0].device)) if train_step <= 2 or train_step % 500 == 0: # sanity: EE z in metres, tensor non-degenerate z = action_raw[:, :, 2].float() if not (0.05 < z.median().item() < 0.6) or action_raw.float().abs().sum() < 1e-3: print(f'[v4head] WARNING: GT EE z median {z.median().item():.3f} / abs-sum ' f'{action_raw.float().abs().sum().item():.3f} — targets look degenerate!') total = None for (k, lv) in head_out['losses'].items(): w = self.config.v4head_loss_weights.get(k, 0.05) total = (w * lv) if total is None else (total + w * lv) train_dict['edm_loss_means'][k] = lv.detach().float() # wandb: loss_train/ train_dict['edm_loss_means']['v4h_total_weighted'] = total.detach().float() train_dict['loss'] = train_dict['loss'] + total.float() train_dict['v4head_out'] = head_out return train_dict # Replaces Predict2Video2WorldModel.compute_loss_with_epsilon_and_sigma() # used @ train only def compute_loss( self, x0_streams: dict[str, torch.Tensor], masks: dict[str, torch.Tensor], condition: dict[str, torch.Tensor], epsilons: dict[str, torch.Tensor], sigmas: dict[str, torch.Tensor], train_step: int = -1, sample_info: list = None, ): ''' EDM training loss: noise the GT streams with epsilons * sigmas, run the network, compute sigma-weighted MSE under supervise masks. Also applies cond_aug noise to conditioning latents and sanitizes NaN streams. Returns the loss dict + entry tensors. ''' streams = sigmas.keys() # Output ones only. # Get the mean and std of the marginal probability distribution means = x0_streams stds = sigmas # Generate noisy observations (corrupted GT data) yt_streams = {k: means[k] + stds[k] * epsilons[k] for k in streams} # Conditioning augmentation (train): add Gaussian noise to conditioning latents. (x0_streams_noised, cond_aug_sigmas) = apply_cond_aug( x0_streams, self.config.train_cond_aug_sigma_range) # Replace NaN in streams with zeros + zero entire affected supervise masks # (corrupted data safety/fallback to prevent NaNs in gradients). sanitize_nan_inf_streams( [x0_streams, x0_streams_noised, yt_streams], masks=masks, iteration=train_step, sample_info=sample_info) # Make prediction using the network with the constructed input tensors. # NOTE(bvh): x0_streams_noised for network input, clean x0_streams kept for loss targets below. (xt_streams, y0_pred_streams, eps_pred_streams, logvar) = self.pipe.denoise( x0_streams_noised, yt_streams, masks, sigmas, condition, iteration=train_step, phase='train') # Loss weights for different noise levels weights_per_sigma = {k: self.get_per_sigma_loss_weights(sigma=sigma) for (k, sigma) in sigmas.items()} # Calculate MSE losses unmasked_mse_losses = {k: (y0_pred_streams[k] - x0_streams[k]) ** 2.0 for k in streams} # Apply supervision masks onto loss tensors mse_loss_streams = multiply_dicts(unmasked_mse_losses, masks['supervise']) # Multiply loss streams by noise-dependent weight factors edm_loss_streams = multiply_dicts(mse_loss_streams, weights_per_sigma) # Convert streams to entries at this point for finer-grained manipulation, # discarding masks where irrelevant x0_entries = unpack_entries_from_streams(x0_streams, self.config, detach=True, ignore_masks=False) yt_entries = unpack_entries_from_streams(yt_streams, self.config, detach=True, ignore_masks=True) xt_entries = unpack_entries_from_streams(xt_streams, self.config, detach=True, ignore_masks=True) y0_pred_entries = unpack_entries_from_streams(y0_pred_streams, self.config, detach=True, ignore_masks=True) edm_loss_entries = unpack_entries_from_streams(edm_loss_streams, self.config, detach=False, ignore_masks=True) # Multiple loss entries by their customized weights for (k, lw) in self.config.loss_weights.items(): if k in edm_loss_entries: edm_loss_entries[k] = edm_loss_entries[k] * lw # Average across all dimensions (batch, sequence, etc) within entries (= per-entry reduce), # then filter out NaN / all-zero entries to obtain the active subset. (edm_loss_means, active_edm_loss_means) = self._reduce_entries_and_filter_active( edm_loss_entries, train_step) # Then reduce across remaining active entries. # NOTE(bvh): batches with fewer active entries used to be upweighted because of mean(), # but changed to sum() on 04/21/2026 (and added / 4.0 for balancing) because of: # http://10.110.170.251:8082/browse/s3://tri-ml-sandbox-16011-us-west-2-datasets/any4d/a4d2/sm/04-21-17-24_s32c_rwm3_robot33m_b4/visuals/DROID?depth=1&flat=1&theater=2&select=&keyword=gal&sort_by=name_asc # kendall_loss = torch.mean(torch.stack(list(active_edm_loss_means.values()))) # single float. kendall_loss = torch.sum(torch.stack(list(active_edm_loss_means.values()))) # single float. kendall_loss /= 4.0 # ^ together with LR / 2.0, corresponds to averaging over 8 (not necessarily all active) # entries in the "old" Any4D2 train_dict = { 'x0_streams': x0_streams, 'x0_streams_noised': x0_streams_noised, 'x0_entries': x0_entries, 'masks': masks, 'yt_streams': yt_streams, 'yt_entries': yt_entries, 'xt_streams': xt_streams, 'xt_entries': xt_entries, 'y0_pred_streams': y0_pred_streams, 'y0_pred_entries': y0_pred_entries, 'eps_pred_streams': eps_pred_streams, #### 'sigmas': sigmas, 'cond_aug_sigmas': cond_aug_sigmas, # (B,) or None 'weights_per_sigma': weights_per_sigma, 'condition': condition, 'logvar': logvar, #### 'mse_loss_streams': mse_loss_streams, 'edm_loss_streams': edm_loss_streams, 'edm_loss_entries': edm_loss_entries, 'edm_loss_means': edm_loss_means, 'active_edm_loss_means': active_edm_loss_means, 'loss': kendall_loss, } # BVH shape notes (any4d): # train_dict: { # 'x0_streams': {'v0': tensor[2, 70, 7, 16, 20] bf16 n=313600 (0.6Mb) x∈[-3.609, 3.422] μ=-0.109 σ=0.910 cuda:0, 'v1': tensor[2, 70, 7, 16, 20] bf16 n=313600 (0.6Mb) x∈[-3.578, 3.188] μ=-0.105 σ=0.941 cuda:0, 'sattn': tensor[2, 25, 22] bf16 n=1100 (2.1Kb) x∈[-0.957, 0.984] μ=0.076 σ=0.475 cuda:0, 'xattn': tensor[2, 30, 14] bf16 n=840 (1.6Kb) all_zeros cuda:0, 'adaln': tensor[2, 25, 21] bf16 n=1050 (2.1Kb) x∈[-0.957, 1.000] μ=0.124 σ=0.516 cuda:0}, # 'x0_entries': {'rgb0': tensor[2, 16, 7, 16, 20] bf16 n=71680 (0.1Mb) x∈[-2.578, 2.891] μ=-0.006 σ=0.609 cuda:0, 'rgb0_input_mask': tensor[2, 1, 7, 16, 20] bf16 n=4480 (8.8Kb) x∈[1.000, 1.000] μ=1.000 σ=0. cuda:0, 'rgb0_output_mask': tensor[2, 1, 7, 16, 20] bf16 n=4480 (8.8Kb) all_zeros cuda:0, 'depth0': tensor[2, 16, 7, 16, 20] bf16 n=71680 (0.1Mb) all_zeros cuda:0, 'depth0_input_mask': tensor[2, 1, 7, 16, 20] bf16 n=4480 (8.8Kb) all_zeros cuda:0, 'depth0_output_mask': tensor[2, 1, 7, 16, 20] bf16 n=4480 (8.8Kb) all_zeros cuda:0, 'cams0': tensor[2, 32, 7, 16, 20] bf16 n=143360 (0.3Mb) x∈[-3.609, 3.422] μ=-0.299 σ=1.227 cuda:0, 'cams0_input_mask': tensor[2, 1, 7, 16, 20] bf16 n=4480 (8.8Kb) x∈[1.000, 1.000] μ=1.000 σ=0. cuda:0, 'cams0_output_mask': tensor[2, 1, 7, 16, 20] bf16 n=4480 (8.8Kb) all_zeros cuda:0, 'rgb1': tensor[2, 16, 7, 16, 20] bf16 n=71680 (0.1Mb) x∈[-2.906, 2.859] μ=0.002 σ=0.656 cuda:0, 'rgb1_input_mask': tensor[2, 1, 7, 16, 20] bf16 n=4480 (8.8Kb) all_zeros cuda:0, 'rgb1_output_mask': tensor[2, 1, 7, 16, 20] bf16 n=4480 (8.8Kb) x∈[1.000, 1.000] μ=1.000 σ=0. cuda:0, 'depth1': tensor[2, 16, 7, 16, 20] bf16 n=71680 (0.1Mb) all_zeros cuda:0, 'depth1_input_mask': tensor[2, 1, 7, 16, 20] bf16 n=4480 (8.8Kb) all_zeros cuda:0, 'depth1_output_mask': tensor[2, 1, 7, 16, 20] bf16 n=4480 (8.8Kb) all_zeros cuda:0, 'cams1': tensor[2, 32, 7, 16, 20] bf16 n=143360 (0.3Mb) x∈[-3.578, 3.188] μ=-0.293 σ=1.266 cuda:0, 'cams1_input_mask': tensor[2, 1, 7, 16, 20] bf16 n=4480 (8.8Kb) x∈[1.000, 1.000] μ=1.000 σ=0. cuda:0, 'cams1_output_mask': tensor[2, 1, 7, 16, 20] bf16 n=4480 (8.8Kb) all_zeros cuda:0, 'action': tensor[2, 25, 20] bf16 n=1000 (2.0Kb) x∈[-0.957, 0.984] μ=0.083 σ=0.496 cuda:0, ...}, # 'masks': {'is_mask': {...}, 'input': {...}, 'output': {...}, 'supervise': {...}}, # 'yt_streams': {'v0': tensor[2, 70, 7, 16, 20] n=313600 (1.2Mb) x∈[-38.076, 38.914] μ=-0.123 σ=7.317 cuda:0, 'v1': tensor[2, 70, 7, 16, 20] n=313600 (1.2Mb) x∈[-27.928, 29.208] μ=-0.113 σ=6.024 cuda:0, 'sattn': tensor[2, 25, 22] n=1100 (4.3Kb) x∈[-21.224, 21.280] μ=0.201 σ=5.677 cuda:0}, # 'yt_entries': {'rgb0': tensor[2, 16, 7, 16, 20] n=71680 (0.3Mb) x∈[-36.712, 37.794] μ=-0.035 σ=7.280 cuda:0, 'depth0': tensor[2, 16, 7, 16, 20] n=71680 (0.3Mb) x∈[-33.218, 38.914] μ=0.036 σ=7.250 cuda:0, 'cams0': tensor[2, 32, 7, 16, 20] n=143360 (0.5Mb) x∈[-35.416, 37.147] μ=-0.327 σ=7.373 cuda:0, 'rgb1': tensor[2, 16, 7, 16, 20] n=71680 (0.3Mb) x∈[-26.312, 27.286] μ=-0.024 σ=5.961 cuda:0, 'depth1': tensor[2, 16, 7, 16, 20] n=71680 (0.3Mb) x∈[-26.021, 23.999] μ=-0.019 σ=5.965 cuda:0, 'cams1': tensor[2, 32, 7, 16, 20] n=143360 (0.5Mb) x∈[-27.928, 29.208] μ=-0.285 σ=6.092 cuda:0, 'action': tensor[2, 25, 20] n=1000 (3.9Kb) x∈[-21.224, 21.280] μ=0.075 σ=5.702 cuda:0}, 'xt_streams': {'adaln': tensor[2, 25, 21] bf16 n=1050 (2.1Kb) x∈[-0.957, 1.000] μ=0.124 σ=0.516 cuda:0, 'sattn': tensor[2, 25, 22] bf16 n=1100 (2.1Kb) all_zeros cuda:0, 'v0': tensor[2, 70, 7, 16, 20] bf16 n=313600 (0.6Mb) x∈[-3.609, 3.422] μ=-0.109 σ=0.910 cuda:0, 'v1': tensor[2, 70, 7, 16, 20] bf16 n=313600 (0.6Mb) x∈[-3.578, 3.703] μ=-0.106 σ=0.980 cuda:0, 'xattn': tensor[2, 30, 14] bf16 n=840 (1.6Kb) all_zeros cuda:0}, 'xt_entries': {'rgb0': tensor[2, 16, 7, 16, 20] bf16 n=71680 (0.1Mb) x∈[-2.578, 2.891] μ=-0.006 σ=0.609 cuda:0, 'depth0': tensor[2, 16, 7, 16, 20] bf16 n=71680 (0.1Mb) all_zeros cuda:0, 'cams0': tensor[2, 32, 7, 16, 20] bf16 n=143360 (0.3Mb) x∈[-3.609, 3.422] μ=-0.299 σ=1.227 cuda:0, 'rgb1': tensor[2, 16, 7, 16, 20] bf16 n=71680 (0.1Mb) x∈[-3.578, 3.703] μ=-0.003 σ=0.859 cuda:0, 'depth1': tensor[2, 16, 7, 16, 20] bf16 n=71680 (0.1Mb) all_zeros cuda:0, 'cams1': tensor[2, 32, 7, 16, 20] bf16 n=143360 (0.3Mb) x∈[-3.578, 3.188] μ=-0.293 σ=1.266 cuda:0, 'action': tensor[2, 25, 20] bf16 n=1000 (2.0Kb) all_zeros cuda:0, 'bbox': tensor[2, 30, 13] bf16 n=780 (1.5Kb) all_zeros cuda:0, 'proprio': tensor[2, 25, 20] bf16 n=1000 (2.0Kb) x∈[-0.957, 0.984] μ=0.080 σ=0.490 cuda:0}, # 'y0_pred_streams': {'sattn': tensor[2, 25, 22] n=1100 (4.3Kb) all_zeros cuda:0, 'v0': tensor[2, 70, 7, 16, 20] n=313600 (1.2Mb) x∈[-3.609, 3.422] μ=-0.109 σ=0.910 cuda:0, 'v1': tensor[2, 70, 7, 16, 20] n=313600 (1.2Mb) x∈[-345.020, 288.295] μ=-0.593 σ=16.492 cuda:0}, # 'y0_pred_entries': {'rgb0': tensor[2, 16, 7, 16, 20] n=71680 (0.3Mb) x∈[-2.578, 2.891] μ=-0.006 σ=0.611 cuda:0, 'depth0': tensor[2, 16, 7, 16, 20] n=71680 (0.3Mb) all_zeros cuda:0, 'cams0': tensor[2, 32, 7, 16, 20] n=143360 (0.5Mb) x∈[-3.609, 3.422] μ=-0.299 σ=1.225 cuda:0, 'rgb1': tensor[2, 16, 7, 16, 20] n=71680 (0.3Mb) x∈[-345.020, 288.295] μ=-2.134 σ=34.401 cuda:0, 'depth1': tensor[2, 16, 7, 16, 20] n=71680 (0.3Mb) all_zeros cuda:0, 'cams1': tensor[2, 32, 7, 16, 20] n=143360 (0.5Mb) x∈[-3.578, 3.188] μ=-0.293 σ=1.267 cuda:0, 'action': tensor[2, 25, 20] n=1000 (3.9Kb) all_zeros cuda:0}, # 'eps_pred_streams': {'sattn': tensor[2, 25, 22] n=1100 (4.3Kb) x∈[-2.613, 2.620] μ=0.025 σ=0.831 cuda:0, 'v0': tensor[2, 70, 7, 16, 20] n=313600 (1.2Mb) x∈[-37.650, 44.647] μ=-0.080 σ=3.987 grad DivBackward0 cuda:0, 'v1': tensor[2, 70, 7, 16, 20] n=313600 (1.2Mb) x∈[-54.355, 64.312] μ=0.350 σ=5.747 grad DivBackward0 cuda:0}, # 'sigmas': {'v0': tensor[2, 1, 1, 1, 1] μ=7.094 σ=2.201 cuda:0 [[[[[5.538]]]], [[[[8.651]]]]], 'v1': tensor[2, 1, 1, 1, 1] μ=5.920 σ=0.617 cuda:0 [[[[[6.357]]]], [[[[5.484]]]]], 'sattn': tensor[2, 1, 1] μ=5.582 σ=2.179 cuda:0 [[[4.041]], [[7.123]]]}, # 'weights_per_sigma': {'v0': tensor[2, 1, 1, 1, 1] μ=1.023 σ=0.014 cuda:0 [[[[[1.033]]]], [[[[1.013]]]]], 'v1': tensor[2, 1, 1, 1, 1] μ=1.029 σ=0.006 cuda:0 [[[[[1.025]]]], [[[[1.033]]]]], 'sattn': tensor[2, 1, 1] μ=1.040 σ=0.029 cuda:0 [[[1.061]], [[1.020]]]}, # 'condition': Vid2VidCondition(_is_broadcasted=False, crossattn_emb=tensor[2, 512, 1024] bf16 n=1048576 (2Mb) x∈[-0.648, 0.578] μ=1.317e-05 σ=0.018 cuda:0, data_type=, padding_mask=tensor[2, 1, 128, 160] bf16 n=40960 (80Kb) all_zeros cuda:0, fps=tensor[2] bf16 μ=10.000 σ=0. cuda:0 [10.000, 10.000], use_video_condition=tensor[1] bool cuda:0 [True], gt_frames=None, condition_video_input_mask_B_C_T_H_W=None), # 'logvar': None, # 'mse_loss_streams': {'sattn': tensor[2, 25, 22] n=1100 (4.3Kb) all_zeros cuda:0, 'v0': tensor[2, 70, 7, 16, 20] n=313600 (1.2Mb) all_zeros cuda:0, 'v1': tensor[2, 70, 7, 16, 20] n=313600 (1.2Mb) x∈[0., 1.192e+05] μ=271.219 σ=2.509e+03 cuda:0}, # 'edm_loss_streams': {'sattn': tensor[2, 25, 22] n=1100 (4.3Kb) all_zeros cuda:0, 'v0': tensor[2, 70, 7, 16, 20] n=313600 (1.2Mb) all_zeros cuda:0, 'v1': tensor[2, 70, 7, 16, 20] n=313600 (1.2Mb) x∈[0., 1.232e+05] μ=279.096 σ=2.582e+03 cuda:0}, # 'edm_loss_entries': {'rgb0': tensor[2, 16, 7, 16, 20] n=71680 (0.3Mb) all_zeros cuda:0, 'depth0': tensor[2, 16, 7, 16, 20] n=71680 (0.3Mb) all_zeros cuda:0, 'cams0': tensor[2, 32, 7, 16, 20] n=143360 (0.5Mb) all_zeros cuda:0, 'rgb1': tensor[2, 16, 7, 16, 20] n=71680 (0.3Mb) x∈[1.333e-09, 1.232e+05] μ=1.221e+03 σ=5.293e+03 cuda:0, 'depth1': tensor[2, 16, 7, 16, 20] n=71680 (0.3Mb) all_zeros cuda:0, 'cams1': tensor[2, 32, 7, 16, 20] n=143360 (0.5Mb) all_zeros cuda:0, 'action': tensor[2, 25, 20] n=1000 (3.9Kb) all_zeros cuda:0}, # 'edm_loss_means': {'rgb0': tensor cuda:0 0., 'depth0': tensor cuda:0 0., 'cams0': tensor cuda:0 0., 'rgb1': tensor cuda:0 1.221e+03, 'depth1': tensor cuda:0 0., 'cams1': tensor cuda:0 0., 'action': tensor cuda:0 0.}, # 'loss': tensor grad MeanBackward0 cuda:0 239.801} return train_dict # Overrides VidarModel @torch.no_grad() def validation_step( self, data_batch: dict[str, torch.Tensor], iteration: int = -1, dataloader_key: str = None, local_path: str = None, directives: dict = None, val_iter: int = 0, seed: int = 0, is_padding: bool = False, ) -> tuple[dict[str, torch.Tensor], torch.Tensor]: ''' :param data_batch (dict): Input batch, contains vidar, fps, prompt, etc. :param iteration (int): Batch index, follows training step. :param local_path (str): Directory to save results to. :param val_iter (int): Validation step, changes while iteration remains constant. :param seed (int): Random seed for diffusion sampling (increment for diversity across samples) :param is_padding (bool): If True, forward pass runs but metrics/visuals are skipped (FSDP sync). :return val_dict (dict): Output samples batch. ''' train_step = iteration val_step = val_iter self.pipe.device = self.device data_batch['dl_key'] = dataloader_key val_dict = dict() loss = None if getattr(self.config, 'data_only_disable_model', False): # Data-only debug: skip pipe.generate / autoregressive; build stub val_dict. # Entries are re-unpacked below by the shared post-processing block. data_batch = self.run_modules( data_batch, phase='val', train_step=train_step, val_step=val_step, directives=directives) (x0_streams, masks) = pack_streams_from_entries(data_batch['a4d_latent'], self.config) val_dict = self._build_stub_output_dict(x0_streams, masks) elif not directives.get('run_autoregressive', False): # Single-pass inference (original behavior) # Prepare data batch using all directives (device, transforms, Any4D entries, masks, VAE encoding, etc) data_batch = self.run_modules(data_batch, phase='val', train_step=train_step, val_step=val_step, directives=directives) # Wan path uses cached UMT5("") negative prompt (see Any4DPipeline.generate); # Cosmos path keeps the legacy zero-tensor uncond. _is_neg = getattr(self.config, 'use_wan_backbone', False) with misc.timer('pipe.generate()'): val_dict = self.pipe.generate( data_batch, seed=seed, num_sampling_steps=self.config.val_num_steps, guidance=self.config.val_cfg_scale, sigma_max=self.config.val_sigma_max, sigma_min=self.config.val_sigma_min, cfg_zero_entries=self.config.val_cfg_zero_entries, is_negative_prompt=_is_neg, phase='val', cond_aug_sigma=self.config.val_cond_aug_sigma, ) else: # Autoregressive inference mode # Strip action perturbation config from directives before run_modules directives_no_action = {k: v for k, v in directives.items() if k != 'perturb_action'} # Prepare data batch (device, transforms, Any4D entries, masks, VAE encoding, etc) data_batch = self.run_modules(data_batch, phase='val', train_step=train_step, val_step=val_step, directives=directives_no_action) # Store original action for reference and GT visuals data_batch['meta']['orig_action'] = data_batch['a4d_raw']['action'].clone() # Apply action override outside run_modules (only values, same structs) if 'perturb_action' in directives and directives['perturb_action']: self._apply_action_override(data_batch, directives['perturb_action'], seed=seed) # Run autoregressive inference, which iteratively updates data_batch # and produces val_dict with combined streams from each step val_dict, data_batch = self.autoreg_handler.run_autoregressive_inference( data_batch, self.pipe, self.config, directives, seed=seed, ) data_batch = self.autoreg_handler.update_masks_for_visualization(data_batch, val_dict) # Padding steps: forward pass ran for FSDP sync, skip everything else. if is_padding: return ({'wandb_logs': {}}, loss) # Convert streams to entries for metrics & visualization val_dict['x0_entries'] = unpack_entries_from_streams( val_dict['x0_streams'], self.config, detach=True, ignore_masks=False) val_dict['xt_entries'] = unpack_entries_from_streams( val_dict['xt_streams'], self.config, detach=True, ignore_masks=True) val_dict['y0_pred_entries'] = unpack_entries_from_streams( val_dict['y0_pred_streams'], self.config, detach=True, ignore_masks=True) # For autoregressive mode, override action entry with full-length adaln stream. # unpack_entries_from_streams() truncates to config-specified temporal size, but # autoregressive mode has extended action sequences that we want to preserve. if directives.get('run_autoregressive', False): action_cfg = self.config.lowdim_adaln_entries['action'] c_start, c_end = action_cfg[2], action_cfg[3] for entries_key, streams_key in [ ('x0_entries', 'x0_streams'), ('xt_entries', 'xt_streams'), ]: if 'adaln' in val_dict[streams_key]: adaln_full = val_dict[streams_key]['adaln'] # (B, T_full, C) val_dict[entries_key]['action'] = adaln_full[:, :, c_start:c_end].detach() metrics_dict = calculate_metrics_a4d( 'val', train_step, data_batch, val_dict, self.vae, self.config) # For inference mode, save directly to local_path without /visuals subdirectory if dataloader_key in ['infer', 'gradio']: visuals_path = local_path else: visuals_path = f'{local_path}/visuals' if dataloader_key is not None: visuals_path += f'/{dataloader_key}' nice_videos = create_save_visuals_a4d( 'val', train_step, data_batch, val_dict, self.vae, self.config, loss, metrics_dict, visuals_path, detail=self.config.val_visuals_detail, quality_pix=self.config.visuals_quality, quality_lat=self.config.visuals_quality + 1, val_step=val_step, batch_idx='all', dataloader_key=dataloader_key, gc_collect=(val_step % 5 == 0)) print_run_info('val', train_step, local_path, self._s3_path) remove_cond = directives.get('remove_cond', True) if getattr(self, 'unroll_val', False): val_dict['unrolled'] = unroll_a4d( data_batch, val_dict, remove_cond=remove_cond) # Gets uploaded by trainer.py later: wandb_logs = self.prepare_wandb_logs_step( 'val', train_step, data_batch, val_dict, metrics_dict, nice_videos, dataloader_key) val_dict['wandb_logs'] = wandb_logs self.all_val_metrics.append(metrics_dict) # BVH shape notes (vanilla): # data_batch: dict_keys(['dl_idx', 'anydata', 'dset_name', 'sample_id', 'fps', 'prompt', # 'scale_factor', 'image_size', 'num_frames', 'padding_mask', 'guidance', 'is_preprocessed', # 'is_anydata', 'a4d_raw', 'a4d_latent', 'meta', 't5_text_embeddings', 't5_text_mask']) # val_dict: dict_keys(['x0_streams', 'masks', 'xt_streams', 'y0_pred_streams', # 'sampling_params', 'x0_entries', 'xt_entries', 'y0_pred_entries', 'wandb_logs']) # metrics_dict: {'psnr': {'rgb0': [...], 'rgb1': [...]}, 'ssim': {'rgb0': [...], 'rgb1': [...]}, 'mse': {'depth0': [...], 'depth1': [...], 'action': [...]}, 'mae': {'depth0': [...], 'depth1': [...], 'action': [...]}, 'friendly': {'psnr': [...], 'ssim': [...], 'mse': [...], 'mae': [...]}} # nice_videos: {'pix_b0': '/basile/ws/repos/cpred2-basile/logs_a4d/a4d2/debug/08-15-02-28_da2/visuals/da2_val_s0_r0_lbmv12_i0_b0_v0,2_all.mp4', 'pix_b1': '/basile/ws/repos/cpred2-basile/logs_a4d/a4d2/debug/08-15-02-28_da2/visuals/da2_val_s0_r0_lbmv12_i0_b1_v0,2_all.mp4'} # wandb_logs: TBD return (val_dict, loss)