# Created by BVH & VG, Jul 2025. import os from hydra.core.config_store import ConfigStore from custom.experiment.template import template_any4d_2b ################################################################# job = dict( project='a4d2', group='cvpr', name='mixed', prepend_datetime=True, ) wandb = dict( enabled=True, # enabled=False, entity='tri', project='a4d2_cvpr', num_validation_logs=5, ) any4d_config = dict( transforms='default_zo1', vidar2a4d='basile', vae='a4d_vae', video_entries=[ dict( rgb0=(0, 16, 'load', 'load'), rgb0_input_mask=(16, 17, 'load', None), rgb0_output_mask=(17, 18, 'zero', None), cams0=(18, 50, 'rand/load', 'zero/load'), cams0_input_mask=(50, 51, 'zero/load', None), cams0_output_mask=(51, 52, 'zero', None), ), dict( rgb1=(0, 16, 'copy:rgb0/load', 'copy:rgb0/load'), rgb1_input_mask=(16, 17, 'copy:rgb0_input_mask/load', None), rgb1_output_mask=(17, 18, 'zero', None), cams1=(18, 50, 'zero/load', 'zero/load'), cams1_input_mask=(50, 51, 'zero/load', None), cams1_output_mask=(51, 52, 'zero', None), ), dict( rgb2=(0, 16, 'copy:rgb0/load', 'copy:rgb0/load'), rgb2_input_mask=(16, 17, 'copy:rgb0_input_mask/load', None), rgb2_output_mask=(17, 18, 'zero', None), cams2=(18, 50, 'zero/load', 'zero/load'), cams2_input_mask=(50, 51, 'zero/load', None), cams2_output_mask=(51, 52, 'zero', None), ), dict( rgb3=(0, 16, 'copy:rgb0/load', 'copy:rgb0/load'), rgb3_input_mask=(16, 17, 'copy:rgb0_input_mask/load', None), rgb3_output_mask=(17, 18, 'zero', None), cams3=(18, 50, 'zero/load', 'zero/load'), cams3_input_mask=(50, 51, 'zero/load', None), cams3_output_mask=(51, 52, 'zero', None), ), ], num_views=4, video_concat_mode='view', video_proj_mode='per_view', loss_weights=dict( rgb0=1.0, rgb1=1.0, rgb2=1.0, rgb3=1.0, ), harmonize_streams=True, harmonize_frames=True, load_modals=['rgb','cams'], 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='rand2', track_metrics=dict( rgb0=['psnr','ssim'], rgb1=['psnr','ssim'], rgb2=['psnr','ssim'], rgb3=['psnr','ssim'], ), train_visuals_interval=0, train_visuals_detail=0, val_visuals_detail=1, train_directives=dict( num_pred_views=1, ), val_directives=dict( tasks='dyn_view_synth', num_pred_views=1, remove_cond=True, ), ) ################################################################# S3_PRETRAINED_PREFIX = r's3://tri-ml-sandbox-16011-us-west-2-datasets/cosmos-predict-2/checkpoints' model = dict( # dit_path=f'{S3_PRETRAINED_PREFIX}/nvidia/Cosmos-Predict2-2B-Video2World/model-480p-10fps.pt', # dit_path='s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/a4d2/cvpr/10-01-10-22_testFix/model/iter_000034000_008704000.pt', # dit_path=f's3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/a4d2/debug/08-27-03-19_s2b_dvs2_pd/model/iter_000038000_002432000.pt', # dit_path=f's3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/a4d2/cvpr/09-19-00-59_multiv1/model/iter_000001000_000192000.pt', # dit_path='s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/a4d2/cvpr/09-20-18-35_2cams/model/iter_000008000_000768000.pt', # dit_path='s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/a4d2/cvpr/10-01-10-22_testFix/model/iter_000035000_008960000.pt', dit_path='s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/a4d2/cvpr/10-01-10-22_testFix/model/iter_000038000_009728000.pt', text_encoder_path=None, # f'{S3_PRETRAINED_PREFIX}/google-t5/t5-11b', vae_path=f'{S3_PRETRAINED_PREFIX}/nvidia/Cosmos-Predict2-2B-Video2World/tokenizer/tokenizer.pth', fsdp_shard_size=8, ############ # run_validation=False, run_validation=True, validation_iter=1000, max_iter=100000, grad_accum_iter=1, context_parallel_size=1, device_monitor=0, manual_gc=100, disable_auto_gc=-1, ############ state_t=11, tokenizer_chunk_duration=41, ) checkpoint = dict( save_iter=1000, s3_folder='s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/', early_sanity_check=-1, # pretrained='s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/a4d2/cvpr/09-23-23-51_dvs1/model/iter_000006000_001536000.pt', # pretrained='s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/a4d2/cvpr/10-01-10-22_testFix/model/iter_000015000_003840000.pt', # pretrained='s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/a4d2/cvpr/10-01-10-22_testFix/model/iter_000031000_007936000.pt', ) optimizer = dict( lr=5.0e-5, ) scheduler = dict( 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_4v.yaml', num_workers=2, batch_size=1, # 4, ) dataset_val = dict( config=dict( ### DRIVING ArgoVerse2sync = 'custom/config/cvpr/mixed/val_4v/argoverse2sync.yaml', DDAD = 'custom/config/cvpr/mixed/val_4v/ddad.yaml', Waymo = 'custom/config/cvpr/mixed/val_4v/waymo.yaml', LyftL5 = 'custom/config/cvpr/mixed/val_4v/lyftl5.yaml', ### ROBOTICS DROIDCalibID = 'custom/config/cvpr/mixed/val_4v/droidcalib_id.yaml', DROIDCalibOOD = 'custom/config/cvpr/mixed/val_4v/droidcalib_ood.yaml', LBMv12 = 'custom/config/cvpr/mixed/val_4v/lbmv12.yaml', ### 3D DL3DV1 = 'custom/config/cvpr/mixed/val_4v/dl3dv1.yaml', # DL3DV2 = 'custom/config/cvpr/mixed/val_4v/dl3dv2.yaml', MVImgNet1 = 'custom/config/cvpr/mixed/val_4v/mvimgnet1.yaml', RE10K1 = 'custom/config/cvpr/mixed/val_4v/re10k1.yaml', # RE10K2 = 'custom/config/cvpr/mixed/val_4v/re10k2.yaml', WildRGBD1 = 'custom/config/cvpr/mixed/val_4v/wildrgbd1.yaml', # WildRGBD2 = 'custom/config/cvpr/mixed/val_4v/wildrgbd2.yaml', ### 4D Kubric4D = 'custom/config/cvpr/mixed/val_4v/kubric4d.yaml', PD4D = 'custom/config/cvpr/mixed/val_4v/pd4d.yaml', GCDKubric4Dgrad = 'custom/config/cvpr/mixed/val_4v/gcd_kubric4d_grad.yaml', GCDPD4Dgrad = 'custom/config/cvpr/mixed/val_4v/gcd_pd4d_grad.yaml', ### TODO EgoExo4D = 'custom/config/cvpr/mixed/val_4v/egoexo4d.yaml', # Kubric5D = 'custom/config/cvpr/mixed/val_4v/kubric5d.yaml', ), num_workers=1, batch_size=1, ) metrics = dict( gt_source='gt', modes=['r','rpv','apv','ipv'], rgb=dict(), ) ################################################################# cs = ConfigStore.instance() this_config = template_any4d_2b( job, wandb, any4d_config, model, checkpoint, optimizer, scheduler, metrics, dataset_train, dataset_val) experiment_name = 'cvpr_' + os.path.splitext(os.path.basename(__file__))[0] cs.store( group='experiment', package='_global_', name=experiment_name, node=this_config, )