# Created by BVH, Jul 2025. import os from hydra.core.config_store import ConfigStore from custom.experiment.template import template_any4d_2b ################################################################# job = dict( # this becomes JobConfig project='a4d2', group='cvpr_eval', name='dvs2', prepend_datetime=True, # becomes {now:%m-%d-%H-%M}_{config.job.name} ) wandb = dict( enabled=True, entity='tri', project='a4d2_cvpr_eval', num_validation_logs=5, ) any4d_config = dict( # this is part of Any4DConfig, which inherits from Predict2Video2WorldModelConfig transforms='default', # reference_camera now separate per dataset dataloader='basile', vae='a4d_vae', video_entries=[ # first entry = base / pretrained 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 cams0=(18, 50, 'load', 'load'), cams0_input_mask=(50, 51, 'load', None), cams0_output_mask=(51, 52, 'load', None), ), # later entries = new viewpoints dict( # key: (channel_start, channel_end, in_proj_init, out_proj_init) rgb1=(0, 16, 'load', 'load'), rgb1_input_mask=(16, 17, 'load', None), rgb1_output_mask=(17, 18, 'load', None), cams1=(18, 50, 'load', 'load'), cams1_input_mask=(50, 51, 'load', None), cams1_output_mask=(51, 52, 'load', None), ), ], num_views=2, video_concat_mode='view', video_proj_mode='per_view', view_timestep_mode='per_view', loss_weights=dict( rgb0=1.0, ), harmonize_streams=True, harmonize_frames=True, load_modals=['rgb', 'cams'], # 'language', 'action', ], task_probs=dict( cross_modal=0.0, dyn_view_synth=1.0, forecast=0.0, pose_est=0.0, inv_dyn=0.0, policy=0.0, world_model=0.0, ), use_views='all', shuffle_cams2views=False, train_directives=dict( num_pred_views=1, ), val_directives=dict( tasks='dyn_view_synth', num_pred_views=1, remove_cond=True, ), # NOTE(bvh): pre-refactor flat fields are deprecated; semantics now in data_*_overrides. # override_zero_origin=True, # makes more sense for 1->1 DVS setting? # override_temporal_stride_val=1, data_train_overrides=dict(zero_origin=True), data_val_overrides=dict(frame_stride=1, zero_origin=True), track_metrics=dict( rgb0=['psnr', 'ssim', 'lpips'], ), train_visuals_interval=99, train_visuals_detail=0, val_visuals_detail=3, visuals_quality=8, viz_input_blacklist=[], # overridden to show cams viz_mask_border_width=0, # overridden because DVS is temporally invariant # val_num_steps=15, # val_num_steps=25, val_num_steps=35, ) ################################################################# S3_PRETRAINED_PREFIX = r's3://tri-ml-sandbox-16011-us-west-2-datasets/cosmos-predict-2/checkpoints' model = dict( # this makes up parts of Predict2Video2WorldModelConfig and Predict2ModelManagerConfig # dit_path='s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/a4d2/cvpr/10-24-01-02_s25_rdvs2_mix_b16/model/iter_000001000_001024000.pt', dit_path='s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/a4d2/cvpr/anyview_v1_30k.pt', # dit_path='s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/a4d2/cvpr/10-31-19-46_officialHR2v/model/iter_000008700_001670400.pt', # text_encoder_path=f'{S3_PRETRAINED_PREFIX}/google-t5/t5-11b', text_encoder_path='', vae_path=f'{S3_PRETRAINED_PREFIX}/nvidia/Cosmos-Predict2-2B-Video2World/tokenizer/tokenizer.pth', # fsdp_shard_size=8, fsdp_shard_size=2, # fsdp_shard_size=1, ############ run_validation=True, validation_iter=100, # frequency / interval of validation runs max_iter=1, # eval-only mode grad_accum_iter=1, context_parallel_size=1, device_monitor=0, manual_gc_iter=288, manual_gc_warm_up=-1, # never disable automatic GC to be safe (weird VRAM issue) ############ state_t=11, # for noise level; = latent # frames for now tokenizer_chunk_duration=41, # for VAE; = num_frames for now ) checkpoint = dict( save_iter=1000, s3_folder='s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/', early_sanity_check=5, ) optimizer = dict( lr=5.0e-5, ) scheduler = dict( # warm_up_steps = [200], warm_up_steps = [1000], cycle_lengths = [model['max_iter']], f_start=[0.01], f_max=[1.0], f_min=[0.05], ) dataset_train = dict( config='custom/config/cvpr/mixed/train_2v_32gpu.yaml', num_workers=1, batch_size=1, ) dataset_val = dict( config=dict( ### 4D AssemblyHandsQNT = 'custom/config/cvpr/quant_v2_LR/assemblyhands.yaml', # 64 # DyCheckM5QNT = 'custom/config/cvpr/quant_v2_LR/dycheckm_5seq.yaml', # 76 # EgoExo4DID01QNT = 'custom/config/cvpr/quant_v2_LR/egoexo4d_ID_01.yaml', # 8 # EgoExo4DID12QNT = 'custom/config/cvpr/quant_v2_LR/egoexo4d_ID_12.yaml', # 8 # EgoExo4DID23QNT = 'custom/config/cvpr/quant_v2_LR/egoexo4d_ID_23.yaml', # 8 # EgoExo4DID30QNT = 'custom/config/cvpr/quant_v2_LR/egoexo4d_ID_30.yaml', # 8 # EgoExo4DID10QNT = 'custom/config/cvpr/quant_v2_LR/egoexo4d_ID_10.yaml', # 8 # EgoExo4DID21QNT = 'custom/config/cvpr/quant_v2_LR/egoexo4d_ID_21.yaml', # 8 # EgoExo4DID32QNT = 'custom/config/cvpr/quant_v2_LR/egoexo4d_ID_32.yaml', # 8 # EgoExo4DID03QNT = 'custom/config/cvpr/quant_v2_LR/egoexo4d_ID_03.yaml', # 8 # EgoExo4DOOD01QNT = 'custom/config/cvpr/quant_v2_LR/egoexo4d_OOD_01.yaml', # 8 # EgoExo4DOOD12QNT = 'custom/config/cvpr/quant_v2_LR/egoexo4d_OOD_12.yaml', # 8 # EgoExo4DOOD23QNT = 'custom/config/cvpr/quant_v2_LR/egoexo4d_OOD_23.yaml', # 8 # EgoExo4DOOD30QNT = 'custom/config/cvpr/quant_v2_LR/egoexo4d_OOD_30.yaml', # 8 # EgoExo4DOOD10QNT = 'custom/config/cvpr/quant_v2_LR/egoexo4d_OOD_10.yaml', # 8 # EgoExo4DOOD21QNT = 'custom/config/cvpr/quant_v2_LR/egoexo4d_OOD_21.yaml', # 8 # EgoExo4DOOD32QNT = 'custom/config/cvpr/quant_v2_LR/egoexo4d_OOD_32.yaml', # 8 # EgoExo4DOOD03QNT = 'custom/config/cvpr/quant_v2_LR/egoexo4d_OOD_03.yaml', # 8 # Kubric4DGradQNT = 'custom/config/cvpr/quant_v2_LR/kubric4d_grad.yaml', # 40 # Kubric4DDirQNT = 'custom/config/cvpr/quant_v2_LR/kubric4d_dir.yaml', # 40 Kubric5DQNT = 'custom/config/cvpr/quant_v2_LR/kubric5d.yaml', # 64 ### DRIVING Argoverse2SyncFLQNT = 'custom/config/cvpr/quant_v2_LR/argoverse_FL.yaml', # 16 Argoverse2SyncFRQNT = 'custom/config/cvpr/quant_v2_LR/argoverse_FR.yaml', # 16 Argoverse2SyncLFQNT = 'custom/config/cvpr/quant_v2_LR/argoverse_LF.yaml', # 16 Argoverse2SyncRFQNT = 'custom/config/cvpr/quant_v2_LR/argoverse_RF.yaml', # 16 DDADFLQNT = 'custom/config/cvpr/quant_v2_LR/ddad_FL.yaml', # 16 DDADFRQNT = 'custom/config/cvpr/quant_v2_LR/ddad_FR.yaml', # 16 DDADLFQNT = 'custom/config/cvpr/quant_v2_LR/ddad_LF.yaml', # 16 DDADRFQNT = 'custom/config/cvpr/quant_v2_LR/ddad_RF.yaml', # 16 LyftL5FLQNT = 'custom/config/cvpr/quant_v2_LR/lyftl5_FL.yaml', # 16 LyftL5FRQNT = 'custom/config/cvpr/quant_v2_LR/lyftl5_FR.yaml', # 16 LyftL5LFQNT = 'custom/config/cvpr/quant_v2_LR/lyftl5_LF.yaml', # 16 LyftL5RFQNT = 'custom/config/cvpr/quant_v2_LR/lyftl5_RF.yaml', # 16 PD4DQNT = 'custom/config/cvpr/quant_v2_LR/pd4d.yaml', # 64 # PD4DDirQNT = 'custom/config/cvpr/quant_v2_LR/pd4d_dir.yaml', # 20 # PD4DGradQNT = 'custom/config/cvpr/quant_v2_LR/pd4d_grad.yaml', # 20 WaymoFLQNT = 'custom/config/cvpr/quant_v2_LR/waymo_FL.yaml', # 16 WaymoFRQNT = 'custom/config/cvpr/quant_v2_LR/waymo_FR.yaml', # 16 WaymoLFQNT = 'custom/config/cvpr/quant_v2_LR/waymo_LF.yaml', # 16 WaymoRFQNT = 'custom/config/cvpr/quant_v2_LR/waymo_RF.yaml', # 16 ### ROBOTICS DROIDIDLRQNT = 'custom/config/cvpr/quant_v2_LR/droid_ID_LR.yaml', # 32 DROIDIDRLQNT = 'custom/config/cvpr/quant_v2_LR/droid_ID_RL.yaml', # 32 DROIDOODLRQNT = 'custom/config/cvpr/quant_v2_LR/droid_OOD_LR.yaml', # 32 DROIDOODRLQNT = 'custom/config/cvpr/quant_v2_LR/droid_OOD_RL.yaml', # 32 LBMv12LRQNT = 'custom/config/cvpr/quant_v2_LR/lbmv12_LR.yaml', # 32 LBMv12RLQNT = 'custom/config/cvpr/quant_v2_LR/lbmv12_RL.yaml', # 32 ), num_workers=0, # can crash debugger otherwise batch_size=1, ) metrics = dict( # gt_source='rec', gt_source='gt', modes=['r','rpv','apv'], # ,'ipv'] unroll=True, rgb=dict(), ) ################################################################# cs = ConfigStore.instance() this_config = template_any4d_2b( job, wandb, any4d_config, model, checkpoint, optimizer, scheduler, metrics, dataset_train, dataset_val) # Use the filename (without extension) as the experiment name experiment_name = 'any4d_' + os.path.splitext(os.path.basename(__file__))[0] cs.store( group='experiment', package='_global_', name=experiment_name, node=this_config, )