# Trains Wan 2.1 Fun-V1.1-1.3B-Control with structured noise for style transfer # text (per-scene VLM caption) + canny edge map -> image, single-view (M=1) at # T=1 (image mode). Random camera per __getitem__ from all 11 gaia cameras. import os from cosmos_predict2.callbacks.grad_clip import GradClipCallback from custom.experiment.template import template_any4d_wan_control_1_3b from hydra.core.config_store import ConfigStore from imaginaire.lazy_config import LazyCall as L # Registers the 'adverse_weather_canny' data_library (AnyData Unified loader + # precomputed UMT5 caption embeddings). Import for its side-effect BEFORE the # template() call below resolves data_library via the registry. import custom.dataset.adverse_weather_canny_dataset # noqa: F401 job = dict( project='a4d2', group='canny_control', name='aw_canny_1frame', prepend_datetime=True, # Local 8-GPU training; the in-container launcher (run_gaia_train.sh) overrides # local_root -> /workspace/Any4D/output_a4d and s3_root -> '' on the CLI. local_root='/workspace/Any4D/output_a4d', s3_root='', ) wandb = dict( enabled=True, entity=os.environ.get('WANDB_ENTITY', 'yuzeng-tri'), project=os.environ.get('WANDB_PROJECT') or 'a4d2', num_validation_logs=8, ) any4d_config = dict( transforms='default', # AnyData "Unified" loader + precomputed caption embeddings (captions_t5_all), # registered in custom/dataset/adverse_weather_canny_dataset.py. data_library='adverse_weather_canny', dataloader='unified_anydrive_vid2vid', vae='a4d_vae_wan', use_wan_backbone=True, wan_variant='control', # PPD structured noise (https://arxiv.org/abs/2512.05106): replace ε with # phase-preserved noise so the clean latent's geometry is baked into the # diffusion forward process. cutoff_scale = mean of exponential dist; # PPD reference is 10.0. use_structured_noise=True, structured_noise_cutoff_scale=10.0, # Exp(scale) for per-iter cutoff_radius (PPD ref) # Inject the clean latent's low-freq phase into only a fixed random K=4-D channel # subspace (orthonormal P; complement stays pure Gaussian). Caps structural locking at # "4 channels' worth" so the 16-ch Wan latent doesn't over-anchor the clear scene at # val and suppress weather transfer. Applies to both training and validation. structured_noise_channels=4, structured_noise_seed=0, # Validation weather-transfer: seed sampling noise with the clear-scene latent's # low-freq phase. With K=4 (4x less locking) a moderate fixed cutoff is fine. val_structured_noise=True, val_structured_noise_cutoff=10.0, # Single view (M=1). Stream layout per view is 50 channels: # [0:16] noisy RGB latent (xt / target) # [16:17] input_mask (Any4D bookkeeping) # [17:18] output_mask (Any4D bookkeeping) # [18:34] reference latent (zeros for T2V mode; forced zero in DiT) # [34:50] edge_control latent (canny VAE-encoded by a4d_vae_wan) # Reshaped to Wan's 48-channel DiT input by Any4DWanControlDiT._to_wan_48ch: # [0:16]/[16:32]/[32:48] = noisy / control / zeros. # Future modalities (``rgb0_depth_control`` / ``rgb0_seg_control``) would # replace ``rgb0_edge_control`` at the same channel range — no DiT change. video_entries=[ dict( rgb0=(0, 16, 'load', 'load'), rgb0_input_mask=(16, 17, 'load', None), rgb0_output_mask=(17, 18, 'load', None), rgb0_ref=(18, 34, 'load', None), rgb0_edge_control=(34, 50, 'load', None), ), ], num_views=1, video_concat_mode='view', video_proj_mode='per_view', view_timestep_mode='per_view', view_emb_std=0.06, block_gate_fix=True, loss_weights=dict(rgb0=1.0), harmonize_streams=True, harmonize_frames=True, load_modals=['rgb'], task_probs=dict( cross_modal=0.0, dyn_view_synth=0.0, forecast=1.0, pose_estimation=0.0, inverse_dynamics=0.0, policy=0.0, world_model=0.0, ), use_views='all', shuffle_cams2views=False, normalizer_stats='', relative_actions=False, action_multiplier=2.0, # T=1 (single-frame): unified_anydrive's forecast task tries to randomly pick # cond_frames from range(1, (T*2)//3, 4) which is empty for T=1. Force # cond_frames_raw=0 so the whole frame is a generation target (text+canny only). train_directives=dict(cond_frames_raw=0), val_directives=dict( tasks='forecast', cond_frames_raw=0, remove_cond=True, ), data_train_overrides=dict(frame_stride=1, resize=[704, 1280]), # 704x1280 (16:9, /16) # subsample=64 so the val set has enough samples to fill max_val_iter=8 across # dp_world=8 (8 samples shown, not 1). single_sample='first' keeps it deterministic. data_val_overrides=dict(frame_stride=1, subsample=64, resize=[704, 1280]), track_metrics=dict( rgb0=['psnr', 'ssim'], ), # Train-visual saving runs on all ranks and is memory/IO-heavy; firing it every 99 # steps stalled the box for ~2.3h once (host OOM-killed a dataloader worker -> crash). # 500 matches the validation cadence (validation_iter=500) -- 5x fewer visual passes. train_visuals_interval=500, train_visuals_detail=1, val_visuals_detail=1, visuals_quality=8, viz_input_blacklist=[], viz_mask_border_width=2, val_num_steps=35, ) WAN_CTRL_CKPT_DIR = os.environ.get( 'WAN_1_3B_CONTROL_CKPT_DIR', '/workspace/Any4D/checkpoints/wan/Wan2.1-Fun-V1.1-1.3B-Control', ) WAN_VAE_PATH = os.environ.get( 'WAN_VAE_PATH', f'{WAN_CTRL_CKPT_DIR}/Wan2.1_VAE.pth', ) model = dict( dit_path=WAN_CTRL_CKPT_DIR, # Empty -> no UMT5 text encoder is loaded; the dataset (adverse_weather_canny) # injects precomputed UMT5 caption embeddings per sample, which # a4d_vae_wan.encode_text_as_needed consumes via its fast-path. Saves ~11GB # VRAM/GPU (enabling a larger batch). default_text_wan.pkl covers any # caption missing from the t5 store. text_encoder_path='', vae_path=WAN_VAE_PATH, fsdp_shard_size=8, # Karras EDM2 power-function EMA (rate=σ_rel=0.1 default). Builds a frozen dit_ema that # the training loop averages into; validation runs under ema_scope() so it evaluates the # EMA weights. EMA is seeded from the loaded init checkpoint by the checkpointer. ema_enabled=True, run_validation=True, # 'delay3' => first validation runs at step 3 (sanity check that EMA + validation work), # instead of the usual pre-training pass. Subsequent validations follow validation_iter. skip_first_validation='delay3', validation_iter=500, max_val_iter=8, # weather-transfer: log ~8 clear-scene x rain/snow-prompt examples max_iter=800000, grad_accum_iter=1, context_parallel_size=1, device_monitor=0, manual_gc_iter=288, manual_gc_warm_up=-1, state_t=1, # T=1: single-frame image-mode training (Tl=1) tokenizer_chunk_duration=1, # raw frames per VAE chunk = 1 ) checkpoint = dict( save_iter=1000, early_sanity_check=5, ) optimizer = dict( lr=3.0e-5, ) scheduler = dict( warm_up_steps=[500], cycle_lengths=[model['max_iter']], f_start=[0.01], f_max=[1.0], f_min=[0.05], ) _AW_CANNY_YAML = 'custom/config/anydata/local/adverse_weather_canny.yaml' # Validation uses a DIFFERENT split: clear (non-rain/snow) scenes + rain/snow prompt # override, for weather-transfer eval. See the val yaml header. _AW_CANNY_VAL_YAML = 'custom/config/anydata/local/adverse_weather_canny_val_transfer.yaml' dataset_train = dict( config=_AW_CANNY_YAML, # 8 ranks x 8 workers x prefetch held a lot of host RAM with no swap to absorb a # spike; T=1 samples are tiny so 4 workers keep the GPUs fed at far lower peak RAM. num_workers=4, prefetch_factor=2, batch_size=1, # overridden by the launcher's BATCH env (probed for largest fit) ) dataset_val = dict( config=dict( AdverseWeatherCanny=_AW_CANNY_VAL_YAML, ), num_workers=0, batch_size=1, ) metrics = None cs = ConfigStore.instance() this_config = template_any4d_wan_control_1_3b( job, wandb, any4d_config, model, checkpoint, optimizer, scheduler, metrics, dataset_train, dataset_val) this_config['trainer']['callbacks']['grad_clip'] = L(GradClipCallback)( max_norm=1.0, log_every_n=10, ) experiment_name = 'any4d_' + os.path.splitext(os.path.basename(__file__))[0] cs.store( group='experiment', package='_global_', name=experiment_name, node=this_config, )