import copy from megatron.core import parallel_state from torch.utils.data import DataLoader, DistributedSampler from cosmos_predict2.callbacks.device_monitor import DeviceMonitor from cosmos_predict2.configs.action_conditioned.config_action_conditioned import ( ACTION_CONDITIONED_PREDICT2_VIDEO2WORLD_NET_2B, ) from cosmos_predict2.configs.base.config_video2world import ConditioningStrategy from cosmos_predict2.models.video2world_model import Predict2ModelManagerConfig from custom.dataloader.collate import any4d_collate from custom.dataset.anydata_dataset import AnyDataset from custom.dataset.rank_dataset import RankDataset # Lazy import: VidarDataset is only needed when data_library != 'anydata' # from custom.dataset.vidar_dataset import VidarDataset from custom.any4d.a4d_network import Any4DDiT from imaginaire.callbacks.manual_gc import ManualGarbageCollection from imaginaire.lazy_config import LazyCall as L # NOTE(bvh): DistributedSampler is needed for Unified datasets (local files, standard __getitem__). # Webbed datasets handle shard splitting by rank internally and ignore the sampler, so this has no effect on them def get_sampler(dataset, shuffle=True, drop_last=True) -> DistributedSampler: return DistributedSampler( dataset, num_replicas=parallel_state.get_data_parallel_world_size(), rank=parallel_state.get_data_parallel_rank(), shuffle=shuffle, seed=0, drop_last=drop_last, ) ############################################################ # Single source of truth for S3 bucket and path roots. # Experiment configs can override job['s3_root'] etc. ############################################################ S3_BUCKET = 's3://tri-ml-sandbox-16011-us-west-2-datasets' S3_OUTPUT_ROOT = f'{S3_BUCKET}/any4d' S3_PRETRAINED_ROOT = f'{S3_BUCKET}/any4d/pretrained' def template(master_config, job, wandb, any4d_config, model, checkpoint, optimizer, scheduler, metrics, dataset_train, dataset_val): ''' :param master_config: vidar_2b / vidar_14b / any4d_2b / any4d_14b. ''' any4d_active = 'any4d' in master_config.lower() ### Default parameters # Path roots (override in experiment config's job dict if needed) job['local_root'] = job.get('local_root', '') job['s3_root'] = job.get('s3_root', S3_OUTPUT_ROOT) wandb['dir'] = wandb.get('dir', '') wandb['tags'] = wandb.get('tags', [v for v in job.values() if isinstance(v, str) and v]) wandb['name'] = wandb.get('name', job['name']) model['model'] = model.get('model', "predict2_video2world_fsdp_2b") model['ema_enabled'] = model.get('ema_enabled', False) model['context_parallel_size'] = model.get('context_parallel_size', 1) model['manual_gc'] = model.get('manual_gc', 500) model['device_monitor'] = model.get('device_monitor', 0) model['fsdp_shard_size'] = model.get('fsdp_shard_size', 8) model['net'] = model.get('net', dict()) model['extra_nets'] = model.get('extra_nets', dict()) model['remove_dit'] = model.get('remove_dit', False) # Shorthand for sac_config mode (e.g. 'mm_only', 'none', 'predict2_2b_720', 'any4d_x', etc) if 'sac_mode' in model: model['net']['sac_config'] = dict(mode=model['sac_mode']) ### data_library = any4d_config.get('data_library', 'anydata') if data_library == 'anydata': DatasetClass = AnyDataset elif data_library == 'vidar': from custom.dataset.vidar_dataset import VidarDataset DatasetClass = VidarDataset else: # External packages can register Dataset classes via custom.dataset.registry. # See custom/dataset/registry.py for the registration API. from custom.dataset.registry import get_data_library DatasetClass = get_data_library(data_library) if DatasetClass is None: raise ValueError( f"Unknown data_library={data_library!r}. Either it must be one of " f"['anydata', 'vidar'], or an external package must call " f"custom.dataset.registry.register_data_library({data_library!r}, ...) " f"before this experiment is loaded.") # Per-sample fallback values and per-phase overrides use OmegaConf interpolation # to reference model.config, so Hydra CLI overrides flow through to the dataset. # E.g.: ++model.config.data_train_overrides.resize='[-16,480]' will reach the dataset. # dataset_train['config'] can be: # - str: single YAML path -> one AnyDataset/VidarDataset for all GPUs # - list[str]: multiple YAML paths -> per-GPU assignment via RankDataset # (each GPU loads one dataset based on data-parallel rank, round-robin) def _make_train_dataset(config_path): if isinstance(config_path, list): return L(RankDataset)( configs=config_path, phase='train', data_overrides="${model.config.data_train_overrides}", data_defaults="${model.config.data_defaults}", shuffle=True, ) return L(DatasetClass)( dataset_dir=config_path, phase='train', data_overrides="${model.config.data_train_overrides}", data_defaults="${model.config.data_defaults}", shuffle=True, ) # dataset_val['config'] can be: # - str: single YAML path -> one dataset for all GPUs # - dict[str, str]: named splits -> separate dataloaders per split for independent metrics # (e.g. {'ArgoVerse2sync': '...yaml', 'DDADfront': '...yaml', ...}) def _make_val_dataset(config_path): return L(DatasetClass)( dataset_dir=config_path, params=dataset_val.get('params', None), phase='val', data_overrides="${model.config.data_val_overrides}", data_defaults="${model.config.data_defaults}", shuffle=False, ) dataset_train_obj = _make_train_dataset(dataset_train['config']) if isinstance(dataset_val['config'], dict): dataset_val_obj = {key: _make_val_dataset(val) for key, val in dataset_val['config'].items()} else: dataset_val_obj = _make_val_dataset(dataset_val['config']) # TODO(bvh/vitor): how to efficiently set config paths from command line? # overriding dataset here does not affect sampler, # and I suspect this might cause NCCL issues. dataloader_train_obj = L(DataLoader)( dataset=dataset_train_obj, sampler=L(get_sampler)(dataset=dataset_train_obj, shuffle=True, drop_last=True), batch_size=dataset_train['batch_size'], drop_last=True, # has to be consistent with sampler num_workers=dataset_train['num_workers'], pin_memory=True, collate_fn=any4d_collate, ) # Build val dataloaders. When dataset_val_obj is a dict (from multi-dataset # validation config), we create one dataloader per split for independent metrics. def _make_val_dataloader(ds): return L(DataLoader)( dataset=ds, sampler=L(get_sampler)(dataset=ds, shuffle=False, drop_last=False), batch_size=dataset_val['batch_size'], drop_last=False, num_workers=dataset_val['num_workers'], pin_memory=False, collate_fn=any4d_collate, ) if isinstance(dataset_val_obj, dict): dataloader_val_obj = {key: _make_val_dataloader(val) for key, val in dataset_val_obj.items()} else: dataloader_val_obj = _make_val_dataloader(dataset_val_obj) wandb['dir'] = '' # Overridden by trainer.py later # params_for_wandb.any4d must mirror model.config, but CLI ++model.config.* # overrides only patch the model.config node, not this snapshot. Interpolate # each field from model.config so W&B logs the resolved (post-override) values. any4d_wandb = {k: '${model.config.' + k + '}' for k in any4d_config} # this field overrides either: # - cosmos_predict2/configs/base/defaults/model.py:VIDAR_FSDP_2B # - cosmos_predict2/configs/base/defaults/model.py:ANY4D_FSDP_2B # and becomes imaginaire/config.py:Config (many subconfigs are also in that file): retval = dict( params_for_wandb = dict( job=job, wandb=wandb, any4d=any4d_wandb, model=model, trainer=dict( max_iter=model['max_iter'], max_val_iter=model.get('max_val_iter', None), validation_iter=model['validation_iter'], run_validation=model.get('run_validation', True), skip_first_validation=model.get('skip_first_validation', False), grad_accum_iter=model.get('grad_accum_iter', 1), seed=model.get('seed', 1000), ), checkpoint=checkpoint, optimizer=optimizer, scheduler=scheduler, dataset_train=dataset_train, dataset_val=dataset_val, metrics=metrics if metrics is not None else {}, ), defaults=[ # This override /model line is important to determine the correct config & network classes etc; # see also cosmos_predict2/configs/base/defaults/model.py:register_model() # {'override /model': 'predict2_video2world_fsdp_2b'}, # {'override /model': 'vidar_2b'}, # {'override /model': 'any4d_2b'}, {'override /model': master_config}, {'override /optimizer': 'fusedadamw'}, {'override /scheduler': 'lambdalinear'}, {'override /ckpt_type': 'standard'}, '_self_', ], job=job, model=dict( # this becomes either VidarModel or Any4DModel, which both inherit from Predict2Video2WorldModel config=dict( # this becomes either Predict2Video2WorldModelConfig or Any4DConfig wandb=wandb, vidar_active=True, any4d_active=any4d_active, **any4d_config, fsdp_shard_size=model['fsdp_shard_size'], pipe_config=dict( # this becomes Video2WorldPipelineConfig (based on either PREDICT2_VIDEO2WORLD_PIPELINE_2B or ANY4D_PIPELINE_2B) ema=dict(enabled=model['ema_enabled']), guardrail_config=dict(enabled=False), prompt_refiner_config=dict(enabled=False), conditioning_strategy=str(ConditioningStrategy.FRAME_REPLACE), min_num_conditional_frames=2, max_num_conditional_frames=2, net=model['net'], extra_nets=model['extra_nets'], remove_dit=model['remove_dit'], ), model_manager_config=L(Predict2ModelManagerConfig)( **({k: v for k, v in model.items() if k in [ 'dit_path', 'dit_ema_path', 'text_encoder_path', 'vae_path', 'tokenizer_chunk_duration', ]}), ), ), ), model_parallel=dict( context_parallel_size=model['context_parallel_size'], ), dataloader_train=dataloader_train_obj, dataloader_val=dataloader_val_obj, trainer=dict( # this goes to ImaginaireTrainer distributed_parallelism='fsdp', callbacks=dict( iter_speed=dict(hit_thres=0), device_monitor=L(DeviceMonitor)( every_n=model['device_monitor'], ), manual_gc=L(ManualGarbageCollection)( every_n=model['manual_gc_iter'] if 'manual_gc_iter' in model else 288, warm_up=model['manual_gc_warm_up'] if 'manual_gc_warm_up' in model else 5, ), ), validation_iter=model['validation_iter'], run_validation=model['run_validation'] if 'run_validation' in model else True, skip_first_validation=model.get('skip_first_validation', False), max_iter=model['max_iter'], max_val_iter=model['max_val_iter'] if 'max_val_iter' in model else None, grad_accum_iter=model['grad_accum_iter'] if 'grad_accum_iter' in model else 1, **({'profiling': model['profiling']} if 'profiling' in model else {}), ), checkpoint=checkpoint, optimizer=optimizer, scheduler=scheduler, metrics=metrics, ) return retval def template_vidar_2b(*args): return template('vidar_2b', *args) def template_any4d_2b(*args): return template('any4d_2b', *args) def template_any4d_wan_1_3b(*args): """Any4D + Wan 2.1 Fun 1.3B InP experiment template. """ from cosmos_predict2.callbacks.grad_clip import GradClipCallback cfg = template('any4d_wan_1_3b', *args) cfg['trainer']['callbacks']['grad_clip'] = L(GradClipCallback)( max_norm=1.0, log_every_n=10, ) return cfg def template_any4d_wan_control_1_3b(*args): return template('any4d_wan_control_1_3b', *args) def template_any4d_wan_t2v_1_3b(*args): """Any4D + pure Wan 2.1 T2V-1.3B base (text -> image/video, in_dim=16).""" from cosmos_predict2.callbacks.grad_clip import GradClipCallback cfg = template('any4d_wan_t2v_1_3b', *args) cfg['trainer']['callbacks']['grad_clip'] = L(GradClipCallback)( max_norm=1.0, log_every_n=10, ) return cfg def template_any4d_wan_inp_canny_1_3b(*args): """Any4D + Wan 2.1 Fun-InP base with canny control extension (in_dim=52).""" return template('any4d_wan_inp_canny_1_3b', *args) def template_any4d_wan_14b(*args): return template('any4d_wan_14b', *args) def template_vidar_14b(*args): raise NotImplementedError() def template_any4d_14b(*args): raise NotImplementedError() def template_regression(*args): return template('predict2_video2world_fsdp_2b_regression', *args)