src/model_training/train.py

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)

from src.model_training.transformers_compat import patch_transformers_hybrid_cache

patch_transformers_hybrid_cache()

from diffsynth.trainers.utils import VideoDataset, CamVideoDataset, wan_parser

import diffsynth.trainers.utils as utils_module

utils_file = utils_module.__file__
if 'site-packages' in utils_file:
    logger.warning(f"Using INSTALLED diffsynth package from: {utils_file}")
else:
    logger.info(f"[VERIFIED] Using LOCAL diffsynth code from: {utils_file}")

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.")),
        ("--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"])),
        ("--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)),
        ("--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_spatial_memory_legacy", "--use_framepack_memory", "--use_framepack_length_compress",
        "--use_block_wise_ssm", "--use_videossm_hybrid", "--sampling_two_chunk_memory",
    ]:
        _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."
            )

        # 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'),
        timestep_shift=float(_arg('timestep_shift', 1.0)),
    )

    # ── VWM-style: Replace DiT blocks with DiTBlock_w_Action ──
    _use_cam_pose = bool(_arg('train_cam_pose', False))
    if _arg('train_action_module', False) or _use_cam_pose:
        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)

        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)
            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)),
                videossm_expand=int(_arg('videossm_expand', 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"):
                getattr(new_block, attr).load_state_dict(getattr(old_block, attr).state_dict())
            with torch.no_grad():
                new_block.modulation.copy_(old_block.modulation.to(dtype=block_dtype))
            new_blocks.append(new_block)

        dit.blocks = new_blocks
        _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)
            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 _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)
        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),
    )
    
    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)),
    )

    model_logger.finish()