# Created by BVH, Jun - Aug 2025. # Defines the extra parameters that you can control in experiments. # This file mirrors parts of: # - cosmos_predict2/configs/base/defaults/model.py # - cosmos_predict2/models/video2world_model.py:Predict2Video2WorldModelConfig import copy import attrs from collections import defaultdict from typing import Any, Dict, Optional from cosmos_predict2.models.video2world_model import ( Predict2Video2WorldModelConfig, ) # This class extends cosmos_predict2/models/video2world_model.py:Predict2Video2WorldModelConfig # and is instantiated in / controlled by: # - ANY4D_FSDP_2B below # - custom/experiment/basile.py:override /model @attrs.define(slots=False) class Any4DConfig(Predict2Video2WorldModelConfig): # NOTE: Modules are still in Predict2Video2WorldModelConfig # Model surgery: video_entries: list[dict] = [ dict( # key: (channel_start, channel_end, in_proj_init, out_proj_init) # NOTE: there will be T*H*W tokens with this information rgb0=(0, 16, 'load', 'load'), # do not change rgb0_input_mask=(16, 17, 'load', None), # do not change rgb0_output_mask=(17, 18, 'load', None), # do not change ), ] lowdim_xattn_entries: dict[str, list] = dict( # key: (seq_start, seq_end, channel_start, channel_end, init) # NOTE: each token here exists only once # NOTE: language is not explicitly mentioned here (always present) ) lowdim_adaln_entries: dict[str, list] = dict( # key: (seq_start, seq_end, channel_start, channel_end, init) # NOTE: each token here exists only once, and they get flattened into one embedding ) lowdim_sattn_entries: dict[str, list] = dict( # key: (seq_start, seq_end, channel_start, channel_end, in_proj_init, out_proj_init) # NOTE: each token here exists only once ) # ^ DEPRECATED: will be removed; predictions go via high-dim streams only # Select the base video diffusion backbone. When True, the pipeline + DiT + VAE # are rebuilt on top of Wan 2.1 Fun 1.3B InP (see custom/wan/a4d_network_wan.py, # custom/wan/a4d_pipe_wan.py, custom/wan/a4d_vae_wan.py). use_wan_backbone: bool = False # ---- Wan backbone variant (only meaningful when use_wan_backbone=True) --------------------- # Which Wan DiT/pipeline to build. See custom/wan/a4d_network_wan*.py + a4d_pipe_wan*.py. # 'inp' Fun-InP base (default): in_dim=36, InP mask + ref slot. # 'control' Fun-V1.1 Control (anchor): text+canny -> video, in_dim=48, separate ref_conv. # 'inp_canny' Fun-InP + canny control extension: in_dim=52, zero-init new channels [36:52]. # 't2v' pure Wan-T2V base: text -> video, in_dim=16, no InP/ref/CLIP slots (marker only: # the network class is chosen by the experiment; routes through build_wan_pipeline). wan_variant: str = 'inp' # ---- Gradient / activation checkpointing (Wan DiT, training) ------------------------------- # Wrap the DiT block loop in torch.utils.checkpoint (non-reentrant): trade recompute for # activation memory. Needed for many-view joint cross-view training (e.g. 11 views) at full res. wan_gradient_checkpointing: bool = False # Checkpoint a block only when (block_idx % period == 0). 1 = every block (max savings/slowest); # >1 = fewer blocks checkpointed (faster, more VRAM). Trade spare VRAM for speed. wan_grad_ckpt_period: int = 1 # SAC op policy: 'mm_only' saves matmuls+attention, recomputes cheap ops (fast, higher VRAM); # 'block_wise' recomputes the whole block incl. attention (max savings, slow). wan_sac_mode: str = 'mm_only' # ---- Wan Fun-InP conditioning granularity ------------------------------------------------- # Keep only the FINAL raw frame real in an isolated trailing given latent slot (gray the other # tcf-1), matching the original Fun-InP single-end-image recipe (Wan's 4x causal temporal VAE). wan_inp_single_end_frame: bool = False # ---- Structured noise: PPD, Phase-Preserving Diffusion (arxiv 2512.05106) ------------------ # Replace the Gaussian epsilon in Any4DWanModel.compute_loss with frequency-selective # structured noise: mix a clean latent's low-freq PHASE into the noise to preserve geometry # (per-view v* streams only). Mechanics: custom/wan/structured_noise.py. Two supported setups: # (1) rgb-phase + control : source_mix=None (phase from the RGB latent), K=4, cutoff radius # (train = Exp(cutoff_scale) draw, val = fixed val_cutoff), P learned jointly with the # denoiser (shared). Paired with an edge control input (wan_variant='control'/'inp_canny'). # (2) multimodal + T2V : source_mix=[...] (phase from the per-step source the dataloader # routes into the rgb{v}_edge_control slot), full_frequency, P pretrained-then-frozen # (per-modality or shared). Base is text->video (wan_variant='t2v'), no control input. use_structured_noise: bool = False structured_noise_channels: int = 16 # K: rank of the channel subspace phase is injected into (K>=C = full PPD) structured_noise_seed: int = 0 # seed for the learnable P subspace's INIT only (re-orthonormalized via QR each step) structured_noise_cutoff_scale: float = 10.0 # setup 1: Exp scale for the per-iter TRAIN cutoff_radius (mean); PPD ref = 10.0 structured_noise_full_frequency: bool = False # setup 2: inject phase at ALL freqs (cutoff=None), TRAIN and VAL, within the K-D subspace structured_noise_source_mix: Any = None # setup 2: phase sources to mix e.g. ['depth','canny','seg','rgb']; None => phase from RGB (setup 1) structured_noise_per_modality_p: bool = False # setup 2: learn a separate projection P_m per modality (else one shared P) structured_noise_p_init_path: str = '' # setup 2: pre-optimized (C,K) P to init sn_P_raw from ('' = seeded random, learned jointly) structured_noise_freeze_p: bool = False # setup 2: freeze sn_P_raw (train the denoiser on the fixed, pre-optimized subspace) structured_noise_active_modality: str = '' # setup 2 eval: which modality's P_m the noise reads ('' = dataloader round-robin) # Validation-time structured noise (weather-transfer eval): seed the sampling noise with the # phase source's low-freq phase. Frequency follows structured_noise_full_frequency (setup 2 => # full; setup 1 => the fixed val cutoff below). val_structured_noise: bool = False val_structured_noise_cutoff: float = 30.0 # Wrap each Wan DiTBlock with PyTorch's non-reentrant checkpoint_wrapper. Trades # ~30% slower step for a large activation-memory reduction (re-computes block forward # during backward instead of storing intermediate activations). Useful for WAN14B model training wan_activation_checkpoint: bool = False num_views: int = 2 # This is technically redundant, but provide manually for sanity checking strict_populate_data: bool = False # Require data batch to contain all above entries as keys. ignore_unused_data: bool = True # Ignore data batch keys that are not claimed by any above entry. # Diffusion & weight & token management options: force_cond_noise_level: bool = False # Force input frames to follow new timestep conditioning scheme video_concat_mode: str = 'view' # view / time / horz / vert video_proj_mode: str = 'per_view' # per_view / shared view_timestep_mode: str = 'per_view' # per_view / shared view_emb_std: float = 0.06 # initial magnitude of new video embeddings lowdim_sattn_proj_mode: str = 'per_token' # per_token / shared lowdim_xattn_proj_mode: str = 'per_token' # per_token / shared lowdim_hidden_channels: int = 64 # max hidden channels for low-dim projection MLPs lowdim_timestep_mode: str = 'per_stream' # per_stream / shared # Conditioning augmentation: train_cond_aug_sigma_range: tuple = (0.001, 0.01) # Train: uniform random sigma in [lo, hi] per batch. # ^ (0, 0) = disabled. Recommended range: (0.001, 0.01). # NOTE(bvh): bottom matches Cosmos-1 default, top is ~10x that (still small relative to EDM sigma_max=80). val_cond_aug_sigma: float = 0.001 # Val/infer: fixed sigma. Defaults to 0 (override in exp cfg or CLI). # Optimization & diffusion noise level distribution parameters: loss_weights: dict[str, float] = dict( rgb0=1.0, ) harmonize_streams: bool = True harmonize_frames: bool = True # Fallback values when no other source provides one. # Priority chain: defaults < custom (yaml) < sample < overrides. data_defaults: dict = dict( frame_rate=10.0, # None = require from YAML or sample metadata frame_stride=1, frame_trim=False, # Upper bound on # frames after striding (cut at end); False/<=0 = off. frame_offset=False, # If True, randomly offset the frame_trim window start (else cut from front). prompt='do something', reference_camera=False, # Dataset camera name/index to anchor extrinsics to, or False = disabled. ref_cam_rel_per_timestep=False, # If True, anchor each frame to ref cam at that t. If False, anchor all to t=0. scale_factor=0.125, # 1.0 / 8.0 action_format_version='action_format_v1', # robot-action parser module: 'action_format_v1' (default, # current semantics) | 'action_format_v2' (no meaningful difference yet). ) # Per-phase forced overrides (top priority). Settable via CLI, e.g. # ++model.config.data_train_overrides.frame_stride=[1,2,3] (a list is randomly sampled). # frame_stride / frame_trim accept int | list (dups=weights) | dict-of-shares {1:0.5,2:0.25}. # Construction-time keys (length, resize, subsample, ...) apply in AnyDataset.__init__; # per-sample keys in supplement_sample; clip/scale_depth in a4d_model; zero_origin = vidar only. data_train_overrides: Dict[str, Any] = dict( subsample=None, single_sample='random', ) data_val_overrides: Dict[str, Any] = dict( frame_stride=1, subsample=32, single_sample='first', ) sync_temporal_transforms: bool = False # ^ If enabled, frame_stride and frame_trim are broadcasted across ranks, not just within each # batch. This is to optimize training speed by minimizing forward pass time shearing. # NOTE/WARNING(bvh): ^ this should only be done if both options are the same or similar across # all datasets (or overrides are active), because otherwise the first dataset YAML custom group # will dominate the overall random distribution. # sync_temporal_probability: float = 0.8 # Data loading & task options (used by SOME dataloaders, but not all): load_modals: list[str] = ['rgb', 'language', 'action', 'cams', 'depth'] use_views: str = 'all' # all / rand1 / rand2 shuffle_cams2views: bool = True # Randomize when subsampling dataset cameras to Any4D views. anydrive_cams2views: bool = True # Apply per-dataset camera ordering (True) or identity mapping (False). reference_view: Any = False # View index (0..V_max-1) to use as plucker reference. False/None = disabled. ref_view_rel_per_timestep: bool = True # ^ If True, anchor each frame t to ref view's camera at t. # ^ If False, anchor all frames to ref view's camera at t=0 (static multi-view). task_probs: dict[str, float] = dict( cross_modal=0.5, dyn_view_synth=0.5, forecast=0.5, pose_est=0.0, inv_dyn=0.0, policy=0.0, world_model=0.9, ) # Impose additional dataloader constraints (e.g. specific tasks, views, etc) per phase: train_directives: dict = dict() val_directives: dict = dict() # Action parameterization: normalizer_stats: str = '' # Make sure not to mix absolute or relative stats. relative_actions: bool = False # Otherwise absolute actions. action_multiplier: float = 2.0 # Multiplier for latent proprio / action values to avoid drowning in noise. traj_multiplier: float = 1.0 # Multiplier for driving ego trajectory values (2D lateral/forward offsets). # NOTE(yams_any4d): V4Head — volumetric action head on DiT features (custom/any4d/v4_head.py). v4head_enabled: bool = False v4head_prep_path: str = '' # npz with rot_centroids / grip_range / z_range / mean scene calib v4head_loss_weights: dict[str, float] = dict(v4h_vol=0.02, v4h_grip=0.05, v4h_rot=0.05) v4head_pred_size: int = 64 v4head_n_z: int = 128 # per view (production --n_height_bins 128) v4head_temporal_mode: str = 'per_step' # 'per_step' (11 volumes, 4-action chunks) or # 'stack_mlp' (stack all frames in time -> res-mlp -> 1 volume) v4head_use_wrist: bool = True # False = scene-only (no wrist sub-volume / sampling / cls) v4head_detach_features: bool = False # True = head trains on frozen backbone features (no # grad into LoRA/DiT); prevents joint-training divergence v4head_n_lat_frames: int = 11 # # latent frames the head stacks (stack_mlp input dim) — MUST # equal the video model's state_t (11 for 41-frame, 3 for 9-frame) v4head_freeze_backbone: bool = False # True = freeze base DiT + LoRA, train ONLY the v4head on # frozen video features (live head-on-frozen-model, no cache) v4head_input_channel_norm: bool = False # True = standardize incoming DiT feats per-channel # online (kills DC bias / sink channels w/o an extraction pass) v4head_train_sigma_band: Optional[tuple] = None # (lo, hi): when set, override the diffusion # timestep to a log-uniform LOW band (the "last generation # step") instead of the full EDM schedule — frozen-head runs # only want clean, informative features (e.g. (0.02, 0.12)) skip_video_visuals: bool = False # True = skip the RGB video gallery + psnr/ssim metrics in # train visuals (VAE decode/mp4 is wasted for a frozen backbone); # keep only the cheap 2D V4Head panels default_text_path: str = 'custom/vae/default_text.pkl' # cached T5 embedding used when # text_encoder_path='' (T5 not loaded). Point at a # per-prompt pkl to bake one fixed instruction. # Metrics & logging & visualization: track_metrics: dict[str, list[str]] = dict( rgb0=['psnr', 'ssim'], ) collate_quiet: bool = False # Suppress collate warnings (useful on SM to avoid flooding CloudWatch). train_visuals_interval: int = 199 # Frequency of training visuals & metrics (<= 0 to disable). train_visuals_detail: int = 2 # 1: normal, 2: many, 3: all (<= 0 to disable) val_visuals_detail: int = 2 # 1: normal, 2: many, 3: all (<= 0 to disable) visuals_quality: int = 7 # 5: low, 7: medium, 9: high cams_viz_type: str = 'comb' # How to depict Plucker video when decoding (rays / comb). depth_viz_type: str = 'viridis' # How to depict depth video when decoding (vidar / viridis / etc). viz_input_blacklist: list[str] = [] # Do not visualize these entries in the input column in galleries. viz_mask_border_width: int = 2 # Thickness of painted mask indicator around input / output / supervise frames in visualization galleries. viz_extra_modes: list[str] = [] # Additional / custom types of visualizations / galleries to export. viz_export_info_json: bool = True # Save a per-sample info.json sidecar (traj + conditioning + pruned dicts) next to visuals. # Diffusion sampling parameters: val_num_steps: int = 35 # 25 # 15 val_cfg_scale: float = 0.0 # 7.0 val_sigma_max: float = 80.0 val_sigma_min: float = 0.002 val_cfg_zero_entries: list[str] = [] # ================ Debugging or risky options (irrelevant for most Any4D users) ================ block_gate_fix: bool = True # use gate_cross_attn instead of gate_mlp (= bug discovered Dec 2025) fsdp_stub_fix: bool = True # enable isolate_stub_streams (= bug discovered May 2026) # Force the disable_risky_sharding decision in text2image_dit.MiniTrainDIT.fully_shard(). # If None (default), falls back to (legacy_network_behavior in [1, 2]). # Set True to keep final_layer / t_embedder un-sharded at any legacy value, e.g. # to mitigate multinode rdvswm4 hangs without changing legacy_network_behavior. disable_risky_sharding: Optional[bool] = None check_nan: bool = False # Check for NaN/Inf in intermediate tensors (slow). data_only_disable_model: bool = False # Data-only debug: FakeVAE + y=x/2 stub. Also forces remove_dit=True so the pipeline factory skips tokenizer / conditioner / DiT / T5 load. # Per-camera retry inside webdataset's tar_file_iterator_single streaming: tar_retry_attempts: int = 3 # = 1 disables retry (= pre patch behavior) tar_thread_join_timeout: float = 60.0 # <= 0 disables the join() safety net # For both legacy settings: # 0 = fixed modern (default); 1 = minimal inline fix on current files; # 2 = full legacy snapshot from older commits; 3 = bugged inbetween. legacy_network_behavior: int = 0 # 0 = uses modern network (fixed or bugged now depends on fsdp_stub_fix); # 1 = pins t_embedder routing to V_used[-1] and disables extra FSDP sharding; # 2 = swaps in custom/legacy/a4d_network_tembv.py and disables extra FSDP sharding; # 3 = alias for 0. legacy_logistics_behavior: int = 0 # 0 = uses ref_dims = (1, 2, 2) (efficient); # 1 = skips zero-pad stub for missing views; # 2 = dispatches to custom/legacy/a4d_logistics_vpad.py; # 3 = uses ref_dims = v0 (wasteful). # ================ Automatically assigned (do not modify) ================ validated: bool = False num_video_in_channels: list[int] = [] # per viewpoint, max across all entries num_video_out_channels: list[int] = [] # per viewpoint, max across all entries num_lowdim_sattn_tokens: int = -1 # max across all entries num_lowdim_sattn_channels: int = -1 # max across all entries num_lowdim_xattn_tokens: int = -1 # max across all entries num_lowdim_xattn_channels: int = -1 # max across all entries num_lowdim_adaln_tokens: int = -1 # max across all entries num_lowdim_adaln_channels: int = -1 # max across all entries has_sattn_stream: bool = False has_xattn_stream: bool = False has_adaln_stream: bool = False all_highdim_entries: list[str] = [] all_lowdim_entries: list[str] = [] # expect_exist_video_in_channels: int = 18 # expect_exist_video_out_channels: int = 16 # TODO(bvh): vary / dive into these aspects at some point (not all implemented yet): # - rope_enable_fps_modulation (False) # - video_proj_mode (per_view) # - view_timestep_mode (per_view) # - video_concat_mode (view) # - lowdim_sattn_proj_mode (per_token) # - lowdim_xattn_proj_mode (per_token)