# Created by BVH, Jun - Aug 2025. # Manages the network, loss, flexible masking, and diffusion sampling. from typing import Any, Callable, Dict, List, Optional, Tuple, Union import torch from torch import Tensor from cosmos_predict2.configs.base.config_video2world import Video2WorldPipelineConfig from cosmos_predict2.functional.batch_ops import batch_mul from cosmos_predict2.pipelines.video2world import Video2WorldPipeline from custom.any4d.a4d_logistics import ( add_dicts, assemble_clean_predictions, assemble_network_inputs, dict_to_cuda_recursive, impute_data_batch_masks, multiply_dicts, pack_streams_from_entries, unpack_entries_from_streams, validate_masks, ) from imaginaire.utils import log, misc # from custom.any4d import a4d_surgery # avoid due to circular import # from custom.any4d.a4d_config import Any4DConfig # avoid due to circular import class Any4DPipeline(Video2WorldPipeline): def __init__( self, *args, **kwargs, ): super().__init__(*args, **kwargs) self.any4d_active = True # since not accessible from pipeline config self.any4d_config = None def is_image_batch( self, data_batch: dict[str, torch.Tensor] ) -> bool: return False # Overrides Video2WorldPipeline.load_checkpoint() # used @ train & inference def load_checkpoint( self, ckpt: dict, any4d_config=None, ): ''' Performs the necessary tensor logistics to map checkpoints from disk to actual architectures. :param ckpt (dict): state_dict_dit_compatible, as loaded from dit_path, without net. prefix. ''' # Load available pretrained weights, copy around and/or initialize new ones as needed from custom.any4d import a4d_surgery self.dit = a4d_surgery.load_checkpoint(self.dit, ckpt, any4d_config) # Replaces Video2WorldPipeline.broadcast_split_for_model_parallelsim() [sic] # used @ train & inference (but NOP for 2B) def broadcast_split_for_model_parallelism( self, ) -> None: """ Broadcast and split the input data and condition for model parallelism. Currently, we only support context parallelism. """ cp_group = self.get_context_parallel_group() cp_size = 1 if cp_group is None else cp_group.size() assert cp_size == 1, 'Context parallelism is not yet supported in Any4D context.' self.dit.disable_context_parallel() # Replaces Video2WorldPipeline.get_data_and_condition() # used @ train & inference def apply_conditioner( self, data_batch: dict[str, torch.Tensor] ): ''' NOTE: Any4D replaces the Cosmos conditioner mechanism almost entirely with dataloader-controlled flexible masking, so this is now only used for text & random dropout thereof. ''' condition = self.conditioner(data_batch) return condition # Overrides Video2WorldPipeline # used @ train & inference def denoise( self, x0_streams: dict[str, torch.Tensor], yt_streams: dict[str, torch.Tensor], masks: dict[str, dict[str, torch.Tensor]], sigmas: dict[str, torch.Tensor], condition: Any, iteration: int = -1, phase: str = 'train', ): ''' :param x0_streams (dict): Maps stream name to (B, C, T, H, W) or (B, T, C) tensor of float. :param yt_streams (dict): Maps stream name to (B, C, T, H, W) or (B, T, C) tensor of float. :param masks (dict): Maps stream name to (B, C, T, H, W) or (B, T, C) tensor of float. :param sigmas (dict): Maps stream name to (B, T, 1) or (B, 1, T, 1, 1) tensor of float. :param condition (Vid2VidCondition). :param iteration (int): Current data batch index. :param phase (str). :return xt_assemb_streams (dict): Maps stream name to (B, C, T, H, W) or (B, T, C) tensor of float. :return y0_pred_streams (dict): Maps stream name to (B, C, T, H, W) or (B, T, C) tensor of float. :return eps_pred_streams (dict): Maps stream name to (B, C, T, H, W) or (B, T, C) tensor of float. :return logvar (dict): Unused. ''' # Calculate preconditions for the network # NOTE(bvh): New vs old Cosmos repo work very differently in this regard! # Notably, there is no explicit conditioning augmentation anymore: # https://github.com/nvidia-cosmos/cosmos-predict2/issues/128 (cs_skip, cs_out, cs_in, cs_noise_pred) = (dict(), dict(), dict(), dict()) (cs_noise_cond, cs_noise_assemb) = (dict(), dict()) for (k, v) in sigmas.items(): sigma = v # This only happens at inference right now: if sigma.ndim == 1: # (B) => (B, 1, 1) or (B, 1, 1, 1, 1) while sigma.ndim < yt_streams[k].ndim: sigma = sigma[:, None] (cs_skip[k], cs_out[k], cs_in[k], cs_noise_pred[k]) = self.scaling(sigma=sigma) # ^ values = (B, T=1, 1) or (B, 1, T=1, 1, 1), see plot for insight (_, _, _, c_noise_cond) = self.scaling(sigma=self.config.sigma_conditional) cs_noise_cond[k] = torch.ones_like(cs_noise_pred[k]) * c_noise_cond # ^ values = (B, T=1, 1) or (B, 1, T=1, 1, 1), all with ~1e-4 # OLD: # Figure out final c_noise for each timestep, as conditional frames # are now treated / marked as having very low noise. # NOTE(bvh): This might feel a bit weird in case of spatial masking, # but I think it makes the most sense to match the pretrained model. # NOTE(bvh): We do not use input masks here since they include mask # values themselves, which would make this part even weirder. # if sigma.ndim == 3: # (B, T, C) => (B, T, 1) # output_pct = masks['output'][k].mean(dim=[2], keepdim=True) # elif sigma.ndim == 5: # (B, C, T, H, W) => (B, 1, T, 1, 1) # output_pct = masks['output'][k].mean(dim=[1, 3, 4], keepdim=True) # cs_noise_assemb[k] = cs_noise_pred[k] * output_pct + cs_noise_cond[k] * (1.0 - output_pct) if not self.any4d_config.force_cond_noise_level: # NEW 1: # The above makes no sense because we very often have mixed inputs and outputs per frame # and per stream, so we just have to teach the model that the preconditions always refer # to outputs only, and unteach the weird conditional noise thing. cs_noise_assemb[k] = cs_noise_pred[k] else: # DEBUG / NEW 2: # To investigate equivalence between vanilla vs any4d # Similar to my initial strategy, except we are doing either one or the other # with binary flags if sigma.ndim == 3: # (B, T, C) => (B, T, 1) output_flag = (masks['output'][k].mean(dim=[2], keepdim=True) > 0.1).to(torch.float32) elif sigma.ndim == 5: # (B, C, T, H, W) => (B, 1, T, 1, 1) output_flag = (masks['output'][k].mean(dim=[1, 3, 4], keepdim=True) > 0.1).to(torch.float32) cs_noise_assemb[k] = cs_noise_pred[k] * output_flag + cs_noise_cond[k] * (1.0 - output_flag) # Scale the noisy observations yt_scaled_streams = {k: cs_in[k] * v for (k, v) in yt_streams.items()} # Create the masked network input tensors # by combining augmented conditioning inputs and scaled noisy target values. xt_assemb_streams = assemble_network_inputs(x0_streams, yt_scaled_streams, masks) # Ensure the whole network operates in bfloat16 precision xt_assemb_streams = {k: v.to(**self.tensor_kwargs) for (k, v) in xt_assemb_streams.items()} cs_noise_assemb = {k: v.to(**self.tensor_kwargs) for (k, v) in cs_noise_assemb.items()} # BVH shape notes: # xt_assemb_streams: { # 'sattn': tensor[2, 25, 22] bf16 n=1100 (2.1Kb) 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, # 'xattn': tensor[2, 30, 14] bf16 n=840 (1.6Kb) all_zeros 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, # 'v0': tensor[2, 70, 7, 16, 20] bf16 n=313600 (0.6Mb) x∈[-3.609, 3.422] μ=-0.109 σ=0.910 cuda:0} # DiT model forward call - Eq. 7 of https://arxiv.org/pdf/2206.00364.pdf # NOTE(bvh): Despite divergent random initializations of new weights across ranks, # at this point they appear to have been synchronized automatically. net_output = self.dit( x_in_streams=xt_assemb_streams, timesteps=cs_noise_assemb, **condition.to_dict(), iteration=iteration) # Calculate predicted clean samples based on SDE y0_raw_pred_streams = dict() eps_pred_streams = dict() for k in net_output.keys(): net_output[k] = net_output[k].to(torch.float32) # match float() in vanilla pipe y0_raw_pred = cs_skip[k] * yt_streams[k] + cs_out[k] * net_output[k] eps_pred = (yt_streams[k] - y0_raw_pred) / sigmas[k] # NOTE(bvh): We could also multiply by output_mask here already, but we leave that to # compute_loss() and assemble_clean_predictions() instead, such that we can visualize # y0_raw_pred_streams directly. y0_raw_pred_streams[k] = y0_raw_pred eps_pred_streams[k] = eps_pred # Create the clean prediction tensors by combining non-augmented conditioning inputs and # predicted outputs. y0_pred_streams = assemble_clean_predictions(x0_streams, y0_raw_pred_streams, masks, config=self.any4d_config) # BVH shape notes (before sattn was properly implemented): # net_output: { # 'v0': tensor[2, 70, 7, 16, 20] bf16 n=313600 (0.6Mb) x∈[-314.000, 276.000] μ=-0.428 σ=29.750 grad UnsafeViewBackward0 cuda:0, # 'v1': tensor[2, 70, 7, 16, 20] bf16 n=313600 (0.6Mb) x∈[-336.000, 410.000] μ=2.500 σ=38.750 grad UnsafeViewBackward0 cuda:0, # 'sattn': tensor[2, 25, 22] bf16 n=1100 (2.1Kb) all_zeros cuda:0} # y0_pred_streams: { # 'sattn': tensor[2, 25, 22] n=1100 (4.3Kb) all_zeros cuda:0, # 'v1': tensor[2, 70, 7, 16, 20] n=313600 (1.2Mb) x∈[-346.712, 284.912] μ=-0.596 σ=16.493 cuda:0, # 'v0': tensor[2, 70, 7, 16, 20] n=313600 (1.2Mb) x∈[-3.609, 3.422] μ=-0.109 σ=0.910 cuda:0} # eps_pred_streams: { # 'v0': tensor[2, 70, 7, 16, 20] n=313600 (1.2Mb) x∈[-37.803, 44.341] μ=-0.078 σ=3.989 grad DivBackward0 cuda:0, # 'v1': tensor[2, 70, 7, 16, 20] n=313600 (1.2Mb) x∈[-53.738, 64.620] μ=0.352 σ=5.747 grad DivBackward0 cuda:0, # 'sattn': tensor[2, 25, 22] n=1100 (4.3Kb) x∈[-2.613, 2.620] μ=0.025 σ=0.831 cuda:0} logvar = None return (xt_assemb_streams, y0_pred_streams, eps_pred_streams, logvar) def _customize_initial_noise(self, yt_noise_start_streams, x0_streams, state_shapes, device, num_sampling_steps, **kwargs): """Extension point for subclasses to customize the initial sampling noise / start step. Returns ``(yt_noise_start_streams, i_start)``. Base is a NO-OP: noise unchanged, i_start=0 (sample from pure noise, all steps). Overridden by the Wan vid2vid pipeline to inject PPD structured noise and/or SDEdit, keeping that logic out of this shared base loop. """ return yt_noise_start_streams, 0 # Replaces Video2WorldPipeline.process() # used @ inference only def generate( self, data_batch: Dict, seed: int = 0, num_sampling_steps: int = 20, guidance: float = 0.0, # 7.0, sigma_max: Optional[float] = None, sigma_min: Optional[float] = None, cfg_zero_entries: list[str] = [], is_negative_prompt: bool = False, phase: str = 'val', cond_aug_sigma: float = 0.0, **kwargs, ) -> Tensor: ''' X ''' assert not self.dit.is_context_parallel_enabled, \ 'Context parallelism is not yet supported in Any4D context.' if sigma_max is not None and sigma_max > 0.0: self.scheduler.config.sigma_max = sigma_max if sigma_min is not None and sigma_min > 0.0: self.scheduler.config.sigma_min = sigma_min # Process the data batch to assemble all inputs, outputs, masks, and conditions. (x0_streams, masks) = pack_streams_from_entries(data_batch['a4d_latent'], self.any4d_config) # Conditioning augmentation: add fixed noise to conditioning latents at val/inference. # NOTE(bvh): keep consistent with a4d_model._apply_cond_aug(); # it's fine to not do per batch index here since val batch size = 1. # TODO(bvh): keep mask channels clean if cond_aug_sigma > 0.0: # TODO(bvh): debug/verify cond aug rigorously; # also, should we noise differently for every sampling step or not? x0_streams_noised = { k: v + cond_aug_sigma * torch.randn_like(v) for k, v in x0_streams.items() } else: x0_streams_noised = x0_streams B = x0_streams['v0'].shape[0] device = x0_streams['v0'].device state_shapes = {k: v.shape[1:] for (k, v) in x0_streams.items() if not k in ['xattn', 'adaln']} # Obtain sampling iteration function (which includes CFG). # NOTE(bvh): noised x0 for network input; clean x0 kept for assemble_clean_predictions below. y0_pred_fn = self.create_y0_pred_fn_from_batch( data_batch, x0_streams_noised, masks, guidance, cfg_zero_entries=cfg_zero_entries, is_negative_prompt=is_negative_prompt, phase=phase, ) # ys_noise_max = {k: torch.randn((B, *state_shapes[k]), **self.tensor_kwargs, # generator=self._noise_generator) * self.sde.sigma_max # for k in state_shapes} # # NOTE(bvh): ^ this is simply epsilon / starting noise # # with sde max sigma (= 80) as scaling factor. yt_noise_start_streams = {k: misc.arch_invariant_rand( (B,) + tuple(state_shapes[k]), torch.float32, self.tensor_kwargs['device'], seed * 100 + v, ) * self.scheduler.config.sigma_max for (v, k) in enumerate(state_shapes.keys())} # NOTE(bvh): ^ 5/6 fixed accidental seed overlap across streams (thanks Yu!) # yt_noise_start_streams: { # 'v0': tensor[1, 18, 11, 26, 48] n=247104 (0.9Mb) x∈[-388.169, 342.868] μ=0.101 σ=79.858 cuda:3, # 'v1': tensor[1, 18, 11, 26, 48] n=247104 (0.9Mb) x∈[-346.151, 362.472] μ=0.267 σ=79.942 cuda:3} # Since the diffusion mathematics in res_sampler.py operate with a single tensor which # we wish to avoid messing with, we need to combine our multiple streams into one. # ys_noise_max_flat = self.streams_dict_to_flat(ys_noise_max) # ------------------------------------------------------------------ # # Sampling loop driven by `RectifiedFlowAB2Scheduler` # ------------------------------------------------------------------ # # Construct sigma schedule (L + 1 entries including simga_min) and timesteps self.scheduler.set_timesteps(num_sampling_steps, device=device) # Extension point (see _customize_initial_noise): subclasses may replace the per-stream # initial noise (e.g. structured noise) and/or return a start-step index (e.g. SDEdit) via # kwargs. Base is a no-op -> plain sampling from pure noise at step 0. yt_noise_start_streams, i_start = self._customize_initial_noise( yt_noise_start_streams, x0_streams, state_shapes, device, num_sampling_steps, **kwargs) # Bring the initial latent into the precision expected by the scheduler sample_streams = {k: yt_noise_start_streams[k].to(dtype=torch.float32) for k in state_shapes} first_xt_streams = None y0_prev_streams = None for (i, _) in enumerate(self.scheduler.timesteps): if i < i_start: # SDEdit: skip the high-noise steps above s_start continue # Current noise level (sigma_t). sigma_t = self.scheduler.sigmas[i].to(device, dtype=torch.float32) # single float. # sigma_in = {k: sigma_t.repeat(B) for k in state_shapes} # ^ dict mapping stream name -> (B) tensor of float. sigma_in = {k: sigma_t.repeat(B, *([1] * len(state_shapes[k]))) for k in state_shapes} # ^ dict mapping stream name -> (B, 1, 1) or (B, 1, 1, 1, 1) tensor of float. # y0 prediction with conditional and unconditional branches (xt_streams, y0_pred_streams) = y0_pred_fn(sample_streams, sigma_in) if first_xt_streams is None: first_xt_streams = xt_streams # Scheduler step updates the noisy sample and returns the cached y0 (sample_streams, y0_prev_streams) = self.scheduler.step_streams( x0_pred_streams=y0_pred_streams, i=i, sample_streams=sample_streams, x0_prev_streams=y0_prev_streams) # Final clean pass at sigma_min sigma_min = self.scheduler.sigmas[-1].to(device, dtype=torch.float32) # single float. sigma_in = {k: sigma_min.repeat(B, *([1] * len(state_shapes[k]))) for k in state_shapes} (xt_streams, sample_streams) = y0_pred_fn(sample_streams, sigma_in) # Create the clean prediction tensors by combining non-augmented conditioning inputs and # predicted outputs. y0_pred_streams = assemble_clean_predictions(x0_streams, sample_streams, masks, config=self.any4d_config) samples_dict = { 'x0_streams': x0_streams, # ground truth (without cond aug) 'masks': masks, 'xt_streams': first_xt_streams, # input (with full / starting noise) 'y0_pred_streams': y0_pred_streams, # prediction #### 'sampling_params': { 'num_steps': num_sampling_steps, 'guidance': guidance, 'cfg_zero': cfg_zero_entries, 'sigma_max': self.scheduler.config.sigma_max, 'sigma_min': self.scheduler.config.sigma_min, 'timesteps': self.scheduler.timesteps, 'sigmas': self.scheduler.sigmas, }, } return samples_dict # Replaces Video2WorldPipeline.get_x0_fn_from_batch() # used @ inference only def create_y0_pred_fn_from_batch( self, data_batch: Dict, x0_streams: Dict[str, torch.Tensor], masks: Dict[str, Dict[str, torch.Tensor]], guidance: float = 0.0, # 1.5, cfg_zero_entries: list[str] = [], is_negative_prompt: bool = False, phase: str = 'val', ) -> Callable: ''' X ''' if is_negative_prompt: (condition, uncondition) = \ self.conditioner.get_condition_with_negative_prompt(data_batch) else: (condition, uncondition) = \ self.conditioner.get_condition_uncondition(data_batch) self.broadcast_split_for_model_parallelism() (cond_x0_streams, cond_masks) = (x0_streams, masks) if guidance > 0.0: (uncond_x0_streams, uncond_masks) = self.create_unconditional_x0_streams_masks( x0_streams, masks, cfg_zero_entries) def y0_pred_fn(yt_streams: Dict[str, torch.Tensor], sigmas: Dict[str, torch.Tensor]) -> torch.Tensor: (cond_xt_streams, cond_y0_pred_streams, _, _) = self.denoise( cond_x0_streams, yt_streams, cond_masks, sigmas, condition, phase=phase) if guidance > 0.0: (uncond_xt_streams, uncond_y0_pred_streams, _, _) = self.denoise( uncond_x0_streams, yt_streams, uncond_masks, sigmas, uncondition, phase=phase) cfg_y0_pred_streams = {k: cond_y0_pred_streams[k] + guidance * (cond_y0_pred_streams[k] - uncond_y0_pred_streams[k]) for k in cond_y0_pred_streams.keys()} else: cfg_y0_pred_streams = cond_y0_pred_streams return (cond_xt_streams, cfg_y0_pred_streams) return y0_pred_fn def create_unconditional_x0_streams_masks( self, x0_streams: Dict[str, torch.Tensor], masks: Dict[str, Dict[str, torch.Tensor]], cfg_zero_entries: list[str] ) -> Dict[str, torch.Tensor]: ''' Modifies the conditional x0 streams to become (partially) unconditional by setting the contents to zero to hide them, and reflecting that in their corresponding input masks. ''' uncond_x0_streams = {k: x0_streams[k].detach().clone() for k in x0_streams.keys()} uncond_masks = {k1: {k2: masks[k1][k2].detach().clone() for k2 in masks[k1].keys()} for k1 in masks.keys()} V_max = self.any4d_config.num_views for v in range(V_max): if f'v{v}' not in x0_streams: continue for (k, c) in self.any4d_config.video_entries[v].items(): if k in cfg_zero_entries: uncond_x0_streams[f'v{v}'][:, c[0]:c[1]] = 0.0 uncond_masks['input'][f'v{v}'][:, c[0]:c[1]] = 0.0 c_mask = self.any4d_config.video_entries[v][k + '_input_mask'] uncond_x0_streams[f'v{v}'][:, c_mask[0]:c_mask[1]] = 0.0 if self.any4d_config.has_sattn_stream: for (k, c) in self.any4d_config.lowdim_sattn_entries.items(): if k in cfg_zero_entries: uncond_x0_streams['sattn'][:, c[0]:c[1], c[2]:c[3]] = 0.0 uncond_masks['input']['sattn'][:, c[0]:c[1], c[2]:c[3]] = 0.0 c_mask = self.any4d_config.lowdim_sattn_entries[k + '_input_mask'] uncond_x0_streams['sattn'][:, c_mask[0]:c_mask[1], c_mask[2]:c_mask[3]] = 0.0 if self.any4d_config.has_xattn_stream: for (k, c) in self.any4d_config.lowdim_xattn_entries.items(): if k in cfg_zero_entries: uncond_x0_streams['xattn'][:, c[0]:c[1], c[2]:c[3]] = 0.0 uncond_masks['input']['xattn'][:, c[0]:c[1], c[2]:c[3]] = 0.0 c_mask = self.any4d_config.lowdim_xattn_entries[k + '_input_mask'] uncond_x0_streams['xattn'][:, c_mask[0]:c_mask[1], c_mask[2]:c_mask[3]] = 0.0 if self.any4d_config.has_adaln_stream: for (k, c) in self.any4d_config.lowdim_adaln_entries.items(): if k in cfg_zero_entries: uncond_x0_streams['adaln'][:, c[0]:c[1], c[2]:c[3]] = 0.0 uncond_masks['input']['adaln'][:, c[0]:c[1], c[2]:c[3]] = 0.0 c_mask = self.any4d_config.lowdim_adaln_entries[k + '_input_mask'] uncond_x0_streams['adaln'][:, c_mask[0]:c_mask[1], c_mask[2]:c_mask[3]] = 0.0 return (uncond_x0_streams, uncond_masks)