# Pure Wan 2.1 T2V-1.3B base, PURE-MIX K=4, FULL-FREQ, PER-MODALITY PRETRAIN-then-FREEZE P (variant 4). # Loads pre-optimized P (output_a4d/pretrain_p/sn_P_permod_C16_K4.pt, coherence 0.48) and FREEZES it; # denoiser trains on the fixed structural subspace, no sync loss. Decoupled from joint training. # # The "pure mix" design: a single K=4 learned subspace (sn_P_raw, no rgb anchor). Each step # the dataloader fills the phase slot from ONE source chosen uniformly from {depth, canny, # seg, rgb}; over training every one of the 4 channels is fed every modality. That shared # exposure is what pushes P toward the CROSS-MODAL CORROBORATED subspace (a direction only # helps the denoiser if the modalities' projected phases agree on it) -- the structural # subspace the paper defines. Unlike the 4+1 variant, there is NO dedicated rgb-anchor # channel: every channel is under cross-modal corroboration pressure, so no channel is free # to absorb RGB appearance. (A/B against the 4+1 run to test whether the anchor leaks.) # # FREQUENCY: full-frequency phase in BOTH train and val (cutoff=None everywhere). # VAL: weather-transfer on all modalities (96 clear scenes + rain/snow prompts; depth/seg # were generated for the clear scenes, so the val source round-robins depth/canny/seg/rgb). import os from custom.experiment.template import template_any4d_wan_t2v_1_3b from hydra.core.config_store import ConfigStore # Registers the 'adverse_weather_canny' data_library (caption attachment; modality-agnostic). import custom.dataset.adverse_weather_canny_dataset # noqa: F401 job = dict( project='a4d2', group='t2v_base', name='aw_t2i_mixphase_syncfzpm_k4', prepend_datetime=True, 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', data_library='adverse_weather_canny', dataloader='unified_anydrive_vid2vid', vae='a4d_vae_wan', use_wan_backbone=True, wan_variant='t2v', # pure Wan2.1-T2V-1.3B base (in_dim=16). NOT control, NOT InP. use_structured_noise=True, structured_noise_cutoff_scale=10.0, # unused (full-frequency overrides, both phases) # FULL-FREQUENCY phase in BOTH train and val. structured_noise_full_frequency=True, val_structured_noise=True, val_structured_noise_cutoff=10.0, # unused (full-frequency overrides) # K=4 learned orthonormal subspace, PURE MIX (no rgb anchor). structured_noise_channels=4, structured_noise_seed=0, # PRETRAIN-then-FREEZE (variant 3): load the pre-optimized P (pretrain_p.py) and freeze it; # NO sync loss in this stage (agreement_weight defaults 0). Denoiser trains on the fixed # structural subspace, decoupled from the denoising loss. structured_noise_per_modality_p=True, structured_noise_p_init_path='/workspace/Any4D/output_a4d/pretrain_p/sn_P_permod_C16_K4.pt', structured_noise_freeze_p=True, structured_noise_source_mix=['depth', 'canny', 'seg', 'rgb'], # Single view (M=1). Stream layout per view is 34 channels: # [0:16] noisy RGB latent (xt / target; the only thing the DiT consumes) # [16:17] input_mask (Any4D bookkeeping) # [17:18] output_mask (Any4D bookkeeping) # [18:34] phase-source latent (depth/canny/seg/rgb VAE latent; read-through PHASE # SOURCE -- NOT a DiT input, NOT supervised) video_entries=[ dict( rgb0=(0, 16, 'load', 'load'), rgb0_input_mask=(16, 17, 'load', None), rgb0_output_mask=(17, 18, 'load', None), rgb0_edge_control=(18, 34, '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, 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]), # subsample=32: the transfer val split has only 96 scenes; 64 would over-prune it. data_val_overrides=dict(frame_stride=1, subsample=32, resize=[704, 1280]), track_metrics=dict( rgb0=['psnr', 'ssim'], ), 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, ) # Pure Wan2.1-T2V-1.3B checkpoint (Wan-AI/Wan2.1-T2V-1.3B). WAN_T2V_CKPT_DIR = os.environ.get( 'WAN_1_3B_T2V_CKPT_DIR', '/workspace/Any4D/checkpoints/wan/Wan2.1-T2V-1.3B', ) WAN_VAE_PATH = os.environ.get( 'WAN_VAE_PATH', f'{WAN_T2V_CKPT_DIR}/Wan2.1_VAE.pth', ) model = dict( dit_path=WAN_T2V_CKPT_DIR, text_encoder_path='', vae_path=WAN_VAE_PATH, fsdp_shard_size=8, ema_enabled=True, run_validation=True, skip_first_validation='delay3', validation_iter=500, max_val_iter=8, 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_MIX_YAML = 'custom/config/anydata/local/adverse_weather_mix.yaml' # Weather-transfer val on ALL modalities (clear scenes now have generated depth/seg). _AW_TRANSFER_VAL_YAML = 'custom/config/anydata/local/adverse_weather_mix_transfer_val.yaml' dataset_train = dict( config=_AW_MIX_YAML, num_workers=4, prefetch_factor=2, batch_size=1, # overridden by the launcher's BATCH env ) dataset_val = dict( config=dict( AdverseWeatherTransfer=_AW_TRANSFER_VAL_YAML, ), num_workers=0, batch_size=1, ) metrics = None cs = ConfigStore.instance() this_config = template_any4d_wan_t2v_1_3b( job, wandb, any4d_config, model, checkpoint, optimizer, scheduler, metrics, dataset_train, dataset_val) experiment_name = 'any4d_' + os.path.splitext(os.path.basename(__file__))[0] cs.store( group='experiment', package='_global_', name=experiment_name, node=this_config, )