import os, sys, re import torch import torch.nn as nn import logging logger = logging.getLogger(__name__) _rank_env = os.environ.get("RANK") or os.environ.get("LOCAL_RANK") or os.environ.get("ACCELERATE_PROCESS_INDEX") or "0" _rank = int(str(_rank_env)) _level = logging.INFO if _rank == 0 else logging.WARNING logging.basicConfig( level=_level, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", datefmt="%Y-%m-%d %H:%M:%S", force=True, ) logger.setLevel(_level) current_file_abs = os.path.abspath(__file__) project_root = os.path.dirname(os.path.dirname(os.path.dirname(current_file_abs))) if project_root not in sys.path: sys.path.insert(0, project_root) modules_to_clear = [ 'diffsynth.models.memory.framepack_length', 'diffsynth.models.memory.framepack_weight', 'diffsynth.models.memory.spatial_grid_memory', 'diffsynth.models.memory.videossm_hybrid', 'diffsynth.models.memory.block_wise_ssm', 'diffsynth.models.memory', 'diffsynth.pipelines.wan_video_new', 'diffsynth.trainers.utils', 'diffsynth.models.wan_video_dit', 'diffsynth.lora.flux_lora', 'diffsynth.lora', 'diffsynth.configs.model_config', 'diffsynth.configs', 'diffsynth.pipelines', 'diffsynth.trainers', 'diffsynth.models', 'diffsynth', ] for mod in modules_to_clear: if mod in sys.modules: del sys.modules[mod] import importlib importlib.invalidate_caches() from diffsynth.pipelines.wan_video_new import WanVideoPipeline, ModelConfig try: import transformers if not hasattr(transformers, "HybridCache") and hasattr(transformers, "DynamicCache"): transformers.HybridCache = transformers.DynamicCache except Exception: pass from diffsynth.trainers.utils import DiffusionTrainingModule, ModelLogger as BaseModelLogger, VideoDataset, CamVideoDataset, wan_parser from diffsynth.models.wan_video_dit import SelfAttention, CrossAttention, GateModule, modulate from diffsynth.models.memory.videossm_hybrid import HybridStateSpaceMemory from diffsynth.models.memory.block_wise_ssm import BlockWiseStateSpaceMemory from diffsynth.models.memory.u_vit_cgla_blocks import CGLATransformerBlock, remap_wan_to_cgla try: import diffsynth.trainers.utils as utils_module utils_file = utils_module.__file__ if hasattr(utils_module, '__file__') else 'unknown' is_local = 'site-packages' not in utils_file if is_local: logger.info(f"[VERIFIED] Using LOCAL diffsynth code from: {utils_file}") else: logger.warning(f"Using INSTALLED diffsynth package from: {utils_file}") except Exception as e: logger.error(f"Failed to verify code location: {e}") import random import numpy as np os.environ["TOKENIZERS_PARALLELISM"] = "false" from safetensors.torch import load_file as safe_load_file from src.model_training.fov_retrieval import setup_fov_retriever_for_training from src.model_training.training_modules import DiTBlock_w_Action, WanTrainingModule def set_seed(seed=42): """Set random seeds for reproducible training.""" random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed(seed) torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False os.environ['PYTHONHASHSEED'] = str(seed) os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8' logger.info(f"Random seed set to {seed}") def _log_dit_freeze_summary(dit: torch.nn.Module) -> None: by_module: dict[str, tuple[int, bool]] = {} for name, p in dit.named_parameters(): numel = p.numel() trainable = p.requires_grad parts = name.split(".") prefix = ".".join(parts[:-1]) if len(parts) > 1 else name if prefix not in by_module: by_module[prefix] = (0, False) prev_numel, prev_trainable = by_module[prefix] by_module[prefix] = (prev_numel + numel, prev_trainable or trainable) trainable_list = [(k, v[0]) for k, v in by_module.items() if v[1]] frozen_list = [(k, v[0]) for k, v in by_module.items() if not v[1]] trainable_list.sort(key=lambda x: x[0]) frozen_list.sort(key=lambda x: x[0]) total_trainable = sum(n for _, n in trainable_list) total_frozen = sum(n for _, n in frozen_list) examples = ", ".join(name for name, _ in trainable_list[:8]) logger.info( f"[DiT freeze] trainable={total_trainable:,} ({len(trainable_list)} groups), " f"frozen={total_frozen:,} ({len(frozen_list)} groups), examples=[{examples}]" ) set_seed(42) from src.model_training.training_modules.model_logger import ModelLogger from src.model_training.training_modules.training_loop import launch_training_task if __name__ == "__main__": parser = wan_parser() def _add_arg_if_missing(*args, **kwargs): if args and args[0] in parser._option_string_actions: return parser.add_argument(*args, **kwargs) for name, kwargs in [ ("--tokenizer_path", dict(type=str, default=None, help="Local tokenizer path.")), ("--wandb_run_name", dict(type=str, default=None)), ("--ckpt_interval", dict(type=int, default=None)), ("--trainable_dit_modules", dict(type=str, default=None, help="Comma-separated DiT modules to unfreeze.")), ("--num_workers", dict(type=int, default=0, help="DataLoader workers.")), ("--max_train_steps", dict(type=int, default=0, help="Stop after N optimizer steps.")), ("--progress_total_steps", dict(type=int, default=0, help="tqdm total steps override.")), ("--log_interval", dict(type=int, default=20, help="Log loss/grad_norm/lr every N optimizer steps.")), ("--max_grad_norm", dict(type=float, default=1.0, help="Grad-norm clip (default 1.0; 0 = no clip, norm still logged).")), ("--resume_from_checkpoint", dict(type=str, default=None)), ("--context_memory_frames", dict(type=int, default=8)), ("--training_mode", dict(type=str, default="predict", choices=["predict", "context", "condition"])), ("--context_drop_prob", dict(type=float, default=0.0)), ("--retrieval_method", dict(type=str, default="fov", choices=["fov", "latent_sim"])), ("--latent_retrieval_dir", dict(type=str, default=None)), ("--fov_top_k", dict(type=int, default=4)), ("--context_attention_weight", dict(type=float, default=1.0)), ("--context_temporal_decay", dict(type=float, default=1.0)), ("--spike_threshold", dict(type=float, default=5.0)), ("--spatial_memory_tokens", dict(type=int, default=64)), ("--spatial_memory_grid", dict(type=int, default=8)), ("--spatial_memory_inject_mode", dict(type=str, default="concat_text", choices=["concat_text", "none", "cross_attn_readout"])), ("--geometry_memory_column", dict(type=str, default="geometry_memory")), ("--geometry_memory_root", dict(type=str, default=None)), ("--geometry_spatial_memory_tokens", dict(type=int, default=64)), ("--geometry_spatial_memory_grid", dict(type=int, default=8)), ("--geometry_spatial_memory_temporal_bins", dict(type=int, default=4)), ("--geometry_spatial_memory_inject_mode", dict(type=str, default="concat_text", choices=["concat_text", "none", "cross_attn_readout"])), ("--framepack_ratio", dict(type=int, default=2)), ("--framepack_length_strategy", dict(type=str, default="distance_merge", choices=["distance_merge", "mean", "uniform", "recent_weighted", "weighted_recent", "packed_multiscale"])), ("--framepack_recent_keep_ratio", dict(type=float, default=0.5)), ("--framepack_multiscale_w2", dict(type=float, default=0.25)), ("--framepack_multiscale_w4", dict(type=float, default=0.15)), ("--context_source", dict(type=str, default="fov", choices=["fov", "replay", "prev_chunk_tail"])), ("--ssm_num_blocks_hint", dict(type=int, default=21)), ("--ssm_every_n_blocks", dict(type=int, default=4)), ("--videossm_kernel_size", dict(type=int, default=3)), ("--videossm_expand", dict(type=int, default=2)), ("--videossm_every_n_blocks", dict(type=int, default=4)), # Camera-Guided Linear Attention (CGLA) memory row. ("--cgla_every_n_blocks", dict(type=int, default=4)), ("--cgla_num_heads", dict(type=int, default=0)), # 0 -> auto (dim // head_dim) ("--cgla_head_dim", dict(type=int, default=128)), ("--cgla_num_sparse_partition", dict(type=int, default=4)), ("--cgla_num_writer", dict(type=int, default=1)), ("--cgla_num_reader", dict(type=int, default=1)), ("--cgla_gate_logit_normalizer", dict(type=int, default=16)), ("--cgla_gate_low_rank_dim", dict(type=int, default=16)), ("--cgla_pose_dim", dict(type=int, default=12)), ("--cgla_pose_bottleneck", dict(type=int, default=64)), ("--cgla_aux_loss_weight", dict(type=float, default=0.0)), ("--cgla_mechanism", dict(type=str, default="cgla", choices=["cgla", "prope", "ucpe"])), ("--cgla_emb_dim", dict(type=int, default=1024)), ("--cgla_num_patches", dict(type=int, default=880)), ("--cgla_temporal_length", dict(type=int, default=21)), ("--cgla_ffn_dim", dict(type=int, default=0)), # 0 -> use Wan DiT ffn_dim ("--sampling_interval_steps", dict(type=int, default=0)), ("--sampling_negative_prompt", dict(type=str, default="oversaturated colors, overexposed, static, blurry details")), ("--sampling_height", dict(type=int, default=352)), ("--sampling_width", dict(type=int, default=640)), ("--sampling_num_frames", dict(type=int, default=81)), ("--sampling_num_inference_steps", dict(type=int, default=50)), ("--sampling_action_path", dict(type=str, default=None)), ("--sampling_two_chunk_action_path", dict(type=str, default=None)), ("--sampling_eval_dataset_base", dict(type=str, default=None)), ("--sampling_eval_metadata_path", dict(type=str, default=None)), ("--samples_per_epoch", dict(type=int, default=0)), ("--camera_encoder_scale", dict(type=float, default=1.0)), ("--camera_inject_mode", dict(type=str, default="post", choices=["post", "pre_norm", "pre_qkv", "pre_qkv_post", "pre_modulate", "pre_qkv_gated"])), ]: _add_arg_if_missing(name, **kwargs) for name in [ "--save_full_model", "--add_action_attn", "--action_use_temporal_attention", "--action_inject_after_spatial_attn", "--use_camera_encoder", "--camera_encoder_shallow", "--camera_encoder_separate_t_r", "--camera_encoder_explicit_yaw", "--yaw_flip_aug", "--camera_encoder_sincos_yaw", "--camera_encoder_r_mlp_no_layernorm", "--add_camera_outside_gate", "--no_camera_encoder_zero_init", "--camera_encoder_full_zero_init", "--enable_context_memory", "--context_per_frame_vae", "--cfg_target_only", "--enable_fov_retrieval", "--use_rt_relative", "--strict_overlap_context", "--use_anchor_frame", "--use_spatial_memory", "--use_geometry_spatial_memory", "--use_spatial_memory_legacy", "--use_framepack_memory", "--use_framepack_length_compress", "--use_block_wise_ssm", "--use_videossm_hybrid", "--sampling_two_chunk_memory", "--use_cgla_memory", "--cgla_use_pose_rope", "--cgla_use_pose_gate_mod", ]: _add_arg_if_missing(name, action="store_true") for name, kwargs in [ ("--per_device_train_batch_size", dict(type=int, default=None)), ("--timestep_shift", dict(type=float, default=1.0)), ("--action_base_path", dict(type=str, default=None)), ("--ckpt_path", dict(type=str, default=None)), ("--cam_position_scale", dict(type=float, default=0.01)), ("--resume_from", dict(type=str, default=None)), ("--verify_ckpt_step", dict(type=int, default=0)), ("--verify_high_noise_first_steps", dict(type=int, default=0)), ("--moc_temperature", dict(type=float, default=1.0)), ("--moc_top_k", dict(type=int, default=0)), ("--prev_chunk_frames", dict(type=int, default=81)), ("--implicit_type", dict(type=str, default="summary")), ("--context_compressor_ratio", dict(type=int, default=2)), ("--episodic_buffer_size", dict(type=int, default=0)), ("--episodic_replay_interval", dict(type=int, default=0)), ("--episodic_replay_weight", dict(type=float, default=0.0)), ]: _add_arg_if_missing(name, **kwargs) for name in [ "--enable_video_sampling", "--sampling_atomic_left_right", "--sampling_four_prompts", "--sampling_two_prompts", "--train_action_module", "--train_cam_pose", "--action_module_only", "--use_moc", "--unified_implicit", "--use_implicit_memory", "--use_memory_v2v_compressor", "--use_slow_fast_memory", "--use_entity_memory", "--use_episodic_memory", ]: _add_arg_if_missing(name, action="store_true") args = parser.parse_args() def _arg(name, default=None): return getattr(args, name, default) def _normalize_and_validate_args(): # Backward-compat mappings if _arg("per_device_train_batch_size", None) is None: args.per_device_train_batch_size = int(_arg("batch_size", 1)) if _arg("sampling_atomic_left_right", False) and not _arg("sampling_two_chunk_memory", False): # Legacy monitor intent maps to current two-chunk monitor. args.sampling_two_chunk_memory = True if _arg("enable_video_sampling", False) and int(_arg("sampling_interval_steps", 0)) <= 0: args.sampling_interval_steps = 1000 # Keep paper-style block-wise SSM and legacy VideoSSM hybrid explicitly separated. if _arg("use_block_wise_ssm", False) and _arg("use_videossm_hybrid", False): raise ValueError( "--use_block_wise_ssm and --use_videossm_hybrid are mutually exclusive; " "use block-wise SSM for paper-aligned runs or VideoSSM hybrid for legacy baselines." ) # CGLA is another per-block temporal-memory module on the same hook as # block-wise SSM / VideoSSM hybrid; only one may be active at a time. _per_block_memory_active = ( bool(_arg("use_block_wise_ssm", False)) or bool(_arg("use_videossm_hybrid", False)) ) if _arg("use_cgla_memory", False) and _per_block_memory_active: raise ValueError( "--use_cgla_memory is mutually exclusive with --use_block_wise_ssm / " "--use_videossm_hybrid (two per-block temporal-memory modules on the same hook)." ) if _arg("use_cgla_memory", False) and not _arg("train_cam_pose", False): logger.warning( "--use_cgla_memory expects per-frame camera pose; enabling without " "--train_cam_pose means no RT reaches the module (it will run as vanilla GLA)." ) # Explicit retrieval strategy visibility: default fov, latent_sim degrades to fov when cache dir is absent. if _arg("retrieval_method", "fov") == "latent_sim": if not _arg("latent_retrieval_dir", None): logger.warning("retrieval_method=latent_sim but latent_retrieval_dir is empty; runtime will fallback to fov retrieval.") else: logger.info(f"retrieval_method=latent_sim latent_retrieval_dir={args.latent_retrieval_dir}") else: logger.info("retrieval_method=fov") # 2-chunk sampling defaults: keep left/right_45 semantics compatible with existing shell wrappers. if _arg("sampling_two_chunk_action_path", None) in (None, ""): args.sampling_two_chunk_action_path = _arg("sampling_action_path", None) _normalize_and_validate_args() resume_step_count = 0 if args.resume_from_checkpoint is not None: if (_arg('trainable_dit_modules', None) or "").strip() or _arg('resume_weights_only', False): logger.info("resume_from_checkpoint used for weights only (trainable_dit_modules set or resume_weights_only), step count starts from 0, no skip data") resume_step_count = 0 else: checkpoint_filename = os.path.basename(args.resume_from_checkpoint) step_match = re.search(r'Step-(\d+)', checkpoint_filename) epoch_match = re.search(r'epoch-(\d+)', checkpoint_filename) if step_match: resume_step_count = int(step_match.group(1)) logger.info(f"Resuming from step {resume_step_count} (extracted from checkpoint filename)") elif epoch_match: logger.info(f"Resuming from epoch checkpoint (epoch-{epoch_match.group(1)}), step count will start from 0") resume_step_count = 0 else: logger.warning("Could not extract step count from checkpoint filename, starting from step 0") set_seed(42) args.enable_icl = False args.icl_num_examples = 2 args.icl_context_frames = 8 if _arg('train_cam_pose', False): dataset = CamVideoDataset(args=args) else: dataset = VideoDataset(args=args, action_base_path=args.action_base_path) def _log_dataset_validation(ds): ds_size = len(ds) ds_repeat = _arg('dataset_repeat', 1) logger.info( f"[Dataset] size={ds_size}, repeat={ds_repeat}, " f"epochs={args.num_epochs}, total_samples={ds_size * ds_repeat * args.num_epochs}" ) _log_dataset_validation(dataset) model = WanTrainingModule( model_paths=args.model_paths, model_id_with_origin_paths=args.model_id_with_origin_paths, tokenizer_path=_arg('tokenizer_path', None), trainable_models=_arg('trainable_models', None), lora_base_model=args.lora_base_model, lora_target_modules=args.lora_target_modules, lora_rank=args.lora_rank, use_gradient_checkpointing_offload=args.use_gradient_checkpointing_offload, extra_inputs=args.extra_inputs, resume_from_checkpoint=args.resume_from_checkpoint, dataset_base_path=_arg('dataset_base_path', None), enable_context_memory=_arg('enable_context_memory', False), context_drop_prob=_arg('context_drop_prob', 0.0), context_drop_seed=42, omit_context_actions=_arg('omit_context_actions', False) or (_arg('context_memory_frames', 8) == 1), # ctx=1: no context action injection context_noise_prob=_arg('context_noise_prob', 0.0), context_noise_std=_arg('context_noise_std', 0.02), context_fixed_noise_std=_arg('context_fixed_noise_std', None), context_memory_frames=_arg('context_memory_frames', 8), context_per_frame_vae=_arg('context_per_frame_vae', False), training_mode=_arg('training_mode', 'predict'), teacher_forcing_prob=_arg('teacher_forcing_prob', 0.0), yaw_flip_aug=_arg('yaw_flip_aug', False), context_source=_arg('context_source', 'fov'), use_framepack_memory=_arg('use_framepack_memory', False), context_temporal_decay=_arg('context_temporal_decay', 1.0), context_attention_weight=_arg('context_attention_weight', 1.0), use_framepack_length_compress=_arg('use_framepack_length_compress', False), framepack_ratio=_arg('framepack_ratio', 2), framepack_length_strategy=_arg('framepack_length_strategy', 'distance_merge'), framepack_recent_keep_ratio=_arg('framepack_recent_keep_ratio', 0.5), framepack_multiscale_w2=_arg('framepack_multiscale_w2', 0.25), framepack_multiscale_w4=_arg('framepack_multiscale_w4', 0.15), use_spatial_memory=_arg('use_spatial_memory', False), use_spatial_memory_legacy=_arg('use_spatial_memory_legacy', False), spatial_memory_tokens=_arg('spatial_memory_tokens', 64), spatial_memory_grid=_arg('spatial_memory_grid', 8), spatial_memory_inject_mode=_arg('spatial_memory_inject_mode', 'concat_text'), use_geometry_spatial_memory=_arg('use_geometry_spatial_memory', False), geometry_spatial_memory_tokens=_arg('geometry_spatial_memory_tokens', 64), geometry_spatial_memory_grid=_arg('geometry_spatial_memory_grid', 8), geometry_spatial_memory_temporal_bins=_arg('geometry_spatial_memory_temporal_bins', 4), geometry_spatial_memory_inject_mode=_arg( 'geometry_spatial_memory_inject_mode', 'concat_text', ), use_moc=_arg('use_moc', False), moc_temperature=_arg('moc_temperature', 1.0), moc_top_k=_arg('moc_top_k', 0), timestep_shift=float(_arg('timestep_shift', 1.0)), ) if _arg('use_moc', False): logger.info( f"[MoC] enabled with temperature={float(_arg('moc_temperature', 1.0))}, " f"top_k={int(_arg('moc_top_k', 0) or 0)}" ) if _arg('use_geometry_spatial_memory', False): logger.info( "[Geometry Spatial Memory] enabled; expects TSDF/point-cloud rendered condition " f"from metadata column '{_arg('geometry_memory_column', 'geometry_memory')}'" ) # ── VWM-style: Replace DiT blocks with DiTBlock_w_Action ── _use_cam_pose = bool(_arg('train_cam_pose', False)) use_cgla_memory = bool(_arg('use_cgla_memory', False)) if _arg('train_action_module', False) or _use_cam_pose or use_cgla_memory: dit = model.pipe.dit old_blocks = dit.blocks has_image_input = dit.has_image_input dim = dit.dim num_heads = dit.num_heads ffn_dim = dit.ffn_dim eps = 1e-6 block_dtype = next(old_blocks[0].parameters()).dtype use_block_wise_ssm = bool(_arg('use_block_wise_ssm', False)) use_videossm_hybrid = bool(_arg('use_videossm_hybrid', False)) ssm_every_n = max(int(_arg('ssm_every_n_blocks', 4)), 1) videossm_every_n = max(int(_arg('videossm_every_n_blocks', 4)), 1) # CGLA: which DiT blocks become CGLATransformerBlock (the rest stay # DiTBlock_w_Action = Wan softmax + cam-pose). cgla_every_n_blocks=1 => # all blocks (matches tests/test_cgla_wan.py); =4 => every 4th block, # the same attach cadence as the SSM/VideoSSM rows (controlled ablation). cgla_every_n = max(int(_arg('cgla_every_n_blocks', 4) or 4), 1) _cgla_head_dim = int(_arg('cgla_head_dim', 128) or 128) _cgla_head_dim = (_cgla_head_dim if _cgla_head_dim > 0 else (dim // num_heads)) new_blocks = nn.ModuleList() for block_id, old_block in enumerate(old_blocks): attach_block_ssm = use_block_wise_ssm and (block_id % ssm_every_n == 0) attach_videossm = use_videossm_hybrid and (block_id % videossm_every_n == 0) attach_cgla = use_cgla_memory and (block_id % cgla_every_n == 0) if attach_cgla: # CGLA row: the DiT block IS CGLATransformerBlock (Wan DiTBlock # with self-attn swapped for the SSE-GLA). Wan's self_attn.{q,k,v,o} # initialise the SSE-GLA {q,k,v,o}_proj (same shapes, via remap); # cross_attn/norm/ffn/modulation load directly. new_block = CGLATransformerBlock( has_image_input=has_image_input, dim=dim, num_heads=num_heads, ffn_dim=ffn_dim, eps=eps, head_dim=_cgla_head_dim, num_sparse_partition=int(_arg('cgla_num_sparse_partition', 4) or 4), num_writer=int(_arg('cgla_num_writer', 1) or 1), num_reader=int(_arg('cgla_num_reader', 1) or 1), pose_dim=int(_arg('cgla_pose_dim', 12) or 12), pose_bottleneck=int(_arg('cgla_pose_bottleneck', 64) or 64), gate_logit_normalizer=int(_arg('cgla_gate_logit_normalizer', 16) or 16), gate_low_rank_dim=int(_arg('cgla_gate_low_rank_dim', 16) or 16), use_pose_gate_mod=bool(_arg('cgla_use_pose_gate_mod', False)), layer_idx=block_id, mechanism=str(_arg('cgla_mechanism', 'cgla') or 'cgla'), bidirectional=True, add_action_attn=_arg('add_action_attn', False), action_use_temporal_attention=_arg('action_use_temporal_attention', False), use_cam_pose=_use_cam_pose, ) new_block = new_block.to(dtype=block_dtype, device=next(old_block.parameters()).device) # Copy the Wan submodules that match by name (cross_attn/norm/ffn). for attr in ("cross_attn", "norm1", "norm2", "norm3", "ffn"): if hasattr(old_block, attr) and hasattr(new_block, attr): getattr(new_block, attr).load_state_dict(getattr(old_block, attr).state_dict()) if hasattr(old_block, "modulation") and hasattr(new_block, "modulation"): with torch.no_grad(): new_block.modulation.copy_(old_block.modulation.to(dtype=block_dtype)) # # self_attn: SSE-GLA structure differs from Wan SelfAttention, # # so remap Wan's {q,k,v,o} -> {q,k,v,o}_proj and load (norm_q/ # # norm_k have no SSE-GLA counterpart -> dropped). # _sa_sd = {} # for _k, _v in old_block.self_attn.state_dict().items(): # if _k.startswith("q."): # _sa_sd["q_proj." + _k[2:]] = _v # elif _k.startswith("k."): # _sa_sd["k_proj." + _k[2:]] = _v # elif _k.startswith("v."): # _sa_sd["v_proj." + _k[2:]] = _v # elif _k.startswith("o."): # _sa_sd["o_proj." + _k[2:]] = _v # new_block.self_attn.load_state_dict(_sa_sd, strict=False) else: new_block = DiTBlock_w_Action( has_image_input=has_image_input, dim=dim, num_heads=num_heads, ffn_dim=ffn_dim, eps=eps, add_action_attn=_arg('add_action_attn', False), action_use_temporal_attention=_arg('action_use_temporal_attention', False), use_cam_pose=_use_cam_pose, use_block_wise_ssm=attach_block_ssm, use_videossm_hybrid=attach_videossm, videossm_kernel_size=int(_arg('videossm_kernel_size', 3) or 3), videossm_expand=int(_arg('videossm_expand', 2) or 2), ) new_block = new_block.to(dtype=block_dtype, device=next(old_block.parameters()).device) for attr in ("self_attn", "cross_attn", "norm1", "norm2", "norm3", "ffn"): if hasattr(old_block, attr) and hasattr(new_block, attr): getattr(new_block, attr).load_state_dict(getattr(old_block, attr).state_dict()) if hasattr(old_block, "modulation") and hasattr(new_block, "modulation"): with torch.no_grad(): new_block.modulation.copy_(old_block.modulation.to(dtype=block_dtype)) new_blocks.append(new_block) dit.blocks = new_blocks _n_cgla = sum(1 for b in new_blocks if isinstance(b, CGLATransformerBlock)) if use_cgla_memory: _mech = str(_arg('cgla_mechanism', 'cgla') or 'cgla') logger.info(f"[CGLA] Replaced {_n_cgla}/{len(new_blocks)} DiT blocks with " f"CGLATransformerBlock (every_n={cgla_every_n}, mechanism={_mech}, " f"head_dim={_cgla_head_dim}, bidirectional); " f"Wan self_attn q/k/v/o -> SSE-GLA q/k/v/o_proj") _mlp_type = "MLP_CamPose" if _use_cam_pose else "MLP_Action" logger.info(f"[VWM-style] Replaced {len(new_blocks)} DiT blocks with DiTBlock_w_Action ({_mlp_type}, zero-init)") if use_block_wise_ssm: logger.info(f"[Block-wise SSM] attached to every {ssm_every_n} DiT block(s)") if use_videossm_hybrid: logger.info(f"[VideoSSM hybrid] attached to every {videossm_every_n} DiT block(s)") device = next(dit.parameters()).device _ckpt_path = _arg('ckpt_path', None) or _arg('resume_from_checkpoint', None) if _ckpt_path is not None: ckpt = safe_load_file(_ckpt_path) if use_cgla_memory: # Remap Wan self_attn.{q,k,v,o} -> SSE-GLA {q,k,v,o}_proj so a Wan # checkpoint loads the linear-attention projections. No-op on an # already-CGLA checkpoint (keys are already *_proj). ckpt = remap_wan_to_cgla(ckpt) missing, unexpected = dit.load_state_dict(ckpt, strict=False) dit.to(device=device) logger.info(f"[VWM-style] Loaded ckpt: {len(ckpt)} keys, missing={len(missing)}, unexpected={len(unexpected)}") if use_cgla_memory: # Block-aware freeze (supports cgla_every_n_blocks > 1, where some # blocks stay DiTBlock_w_Action = Wan softmax). CGLA blocks: train # the SSE-GLA self_attn + cam_encoder + cgla_gate + noise_write_gate, # freeze the Wan backbone (cross_attn/norm/ffn/modulation). Non-CGLA # blocks: VWM pattern (train action_mlp / self_attn_with_action / # SSM, freeze Wan softmax self_attn + backbone). for block in dit.blocks: if isinstance(block, CGLATransformerBlock): for name, param in block.named_parameters(): if ("action_mlp" in name) or ("self_attn_with_action" in name) or ("self_attn" in name) \ ("cam_encoder" in name) or ("noise_write_gate" in name): param.requires_grad = True else: param.requires_grad = False else: for name, param in block.named_parameters(): if ("action_mlp" in name) or ("self_attn_with_action" in name) or \ ("block_wise_ssm" in name) or ("videossm_hybrid" in name): param.requires_grad = True else: param.requires_grad = False elif _arg('action_module_only', False): if _arg('add_action_attn', False): for block in dit.blocks: for name, param in block.named_parameters(): if ("action_mlp" in name) or ("self_attn_with_action" in name) or ("block_wise_ssm" in name) or ("videossm_hybrid" in name): param.requires_grad = True else: param.requires_grad = False else: for block in dit.blocks: for name, param in block.named_parameters(): if "action_mlp" in name or "self_attn" in name or "block_wise_ssm" in name or "videossm_hybrid" in name: param.requires_grad = True else: param.requires_grad = False else: for block in dit.blocks: for name, param in block.named_parameters(): if "action_mlp" in name or "self_attn_with_action" in name or "block_wise_ssm" in name or "videossm_hybrid" in name: param.requires_grad = True else: param.requires_grad = False _log_dit_freeze_summary(dit) _resume_from = _arg('resume_from', None) if _resume_from: logger.info(f"Loading full resume checkpoint: {_resume_from}") ckpt = safe_load_file(_resume_from) if use_cgla_memory: # No-op on an already-CGLA ckpt (keys are *_proj); remaps a Wan ckpt. ckpt = remap_wan_to_cgla(ckpt) model.pipe.dit.load_state_dict(ckpt, strict=False) logger.info(f"Checkpoint loaded, resuming from step {resume_step_count}") model_logger = ModelLogger( args.output_path, remove_prefix_in_ckpt=args.remove_prefix_in_ckpt, wandb_run_name=args.wandb_run_name, ckpt_interval=args.ckpt_interval, resume_step_count=resume_step_count, save_full_model=_arg('save_full_model', False), context_drop_prob=float(_arg("context_drop_prob", 0.0)), enable_video_sampling=_arg("enable_video_sampling", False), sampling_interval_steps=int(_arg("sampling_interval_steps", 0)), sampling_two_chunk_memory=_arg("sampling_two_chunk_memory", False), sampling_action_path=_arg("sampling_action_path", None), sampling_two_chunk_action_path=_arg("sampling_two_chunk_action_path", None), sampling_negative_prompt=_arg("sampling_negative_prompt", ""), sampling_height=int(_arg("sampling_height", 352)), sampling_width=int(_arg("sampling_width", 640)), sampling_num_frames=int(_arg("sampling_num_frames", 81)), sampling_num_inference_steps=int(_arg("sampling_num_inference_steps", 50)), context_memory_frames=int(_arg("context_memory_frames", 1)), context_source=_arg("context_source", "replay"), context_per_frame_vae=_arg("context_per_frame_vae", False), # Monitor samples with the SAME noise-schedule shift training uses # (--timestep_shift), so the in-training video reflects the trained # schedule (15 for two-chunk rows, 5 for legacy ctx rows). sampling_sigma_shift=float(_arg("timestep_shift", 1.0) or 1.0), ) optimizer = torch.optim.AdamW(model.trainable_modules(), lr=args.learning_rate) scheduler = torch.optim.lr_scheduler.ConstantLR(optimizer) # Setup FOV retriever for context-based memory training (also for ModelLogger sampling) enable_fov_retrieval = _arg('enable_fov_retrieval', False) fov_retriever = None dataset_base_path = _arg('dataset_base_path', None) if enable_fov_retrieval: fov_retriever = setup_fov_retriever_for_training( dataset_base_path=dataset_base_path, enable_fov_retrieval=True ) launch_training_task( dataset, model, model_logger, optimizer, scheduler, num_epochs=args.num_epochs, gradient_accumulation_steps=args.gradient_accumulation_steps, per_device_train_batch_size=int(_arg("per_device_train_batch_size", 1)), spike_threshold=_arg('spike_threshold', 5.0), resume_step_count=resume_step_count, enable_fov_retrieval=enable_fov_retrieval, retrieval_method=_arg('retrieval_method', 'fov'), latent_retrieval_dir=_arg('latent_retrieval_dir', None), dataset_base_path=_arg('dataset_base_path', None), fov_retriever=fov_retriever, context_memory_frames=_arg('context_memory_frames', 8), prev_chunk_frames=int(_arg('prev_chunk_frames', 81)), fov_top_k=_arg('fov_top_k', 4), # Number of overlap frames (4), GT frame 0 added automatically use_rt_relative=_arg('use_rt_relative', False), # Experiment 1_4_2: RT relative conversion strict_overlap_context=_arg('strict_overlap_context', False), dataset_repeat=_arg('dataset_repeat', 1), # Pass dataset_repeat for step calculation use_camera_encoder=_arg('use_camera_encoder', False), # exp1_4_3: DDP find_unused_parameters num_workers=_arg('num_workers', 0), context_source=_arg('context_source', 'fov'), max_train_steps=int(_arg('max_train_steps', 0)), progress_total_steps=int(_arg('progress_total_steps', 0)), log_interval=max(1, int(_arg('log_interval', 20) or 20)), max_grad_norm=float(_arg('max_grad_norm', 1.0)), ) model_logger.finish()