scripts/train_moe.py

#融合nuscenes和sekai数据集的MoE训练
import torch
import torch.nn as nn
import lightning as pl
import wandb
import os
import time
import copy
import json
import numpy as np
import random
import traceback
from diffsynth import WanVideoReCamMasterPipeline, ModelManager
from torchvision.transforms import v2
from einops import rearrange
from pose_classifier import PoseClassifier
import argparse
from scipy.spatial.transform import Rotation as R

def get_traj_position_change(cam_c2w, stride=1):
    positions = cam_c2w[:, :3, 3]
    
    traj_coord = []
    tarj_angle = []
    for i in range(0, len(positions) - 2 * stride):
        v1 = positions[i + stride] - positions[i]
        v2 = positions[i + 2 * stride] - positions[i + stride]

        norm1 = np.linalg.norm(v1)
        norm2 = np.linalg.norm(v2)
        if norm1 < 1e-6 or norm2 < 1e-6:
            continue

        cos_angle = np.dot(v1, v2) / (norm1 * norm2)
        angle = np.degrees(np.arccos(np.clip(cos_angle, -1.0, 1.0)))

        traj_coord.append(v1)
        tarj_angle.append(angle)
    
    return traj_coord, tarj_angle

def get_traj_rotation_change(cam_c2w, stride=1):
    rotations = cam_c2w[:, :3, :3]
    
    traj_rot_angle = []
    for i in range(0, len(rotations) - stride):
        z1 = rotations[i][:, 2]
        z2 = rotations[i + stride][:, 2]

        norm1 = np.linalg.norm(z1)
        norm2 = np.linalg.norm(z2)
        if norm1 < 1e-6 or norm2 < 1e-6:
            continue

        cos_angle = np.dot(z1, z2) / (norm1 * norm2)
        angle = np.degrees(np.arccos(np.clip(cos_angle, -1.0, 1.0)))
        traj_rot_angle.append(angle)

    return traj_rot_angle

def compute_relative_pose(pose_a, pose_b, use_torch=False):
    """计算相机B相对于相机A的相对位姿矩阵"""
    assert pose_a.shape == (4, 4), f"相机A外参矩阵形状应为(4,4)，实际为{pose_a.shape}"
    assert pose_b.shape == (4, 4), f"相机B外参矩阵形状应为(4,4)，实际为{pose_b.shape}"
    
    if use_torch:
        if not isinstance(pose_a, torch.Tensor):
            pose_a = torch.from_numpy(pose_a).float()
        if not isinstance(pose_b, torch.Tensor):
            pose_b = torch.from_numpy(pose_b).float()
        
        pose_a_inv = torch.inverse(pose_a)
        relative_pose = torch.matmul(pose_b, pose_a_inv)
    else:
        if not isinstance(pose_a, np.ndarray):
            pose_a = np.array(pose_a, dtype=np.float32)
        if not isinstance(pose_b, np.ndarray):
            pose_b = np.array(pose_b, dtype=np.float32)
        
        pose_a_inv = np.linalg.inv(pose_a)
        relative_pose = np.matmul(pose_b, pose_a_inv)
    
    return relative_pose

def compute_relative_pose_matrix(pose1, pose2):
    """
    计算相邻两帧的相对位姿，返回3×4的相机矩阵 [R_rel | t_rel]
    
    参数:
    pose1: 第i帧的相机位姿，形状为(7,)的数组 [tx1, ty1, tz1, qx1, qy1, qz1, qw1]
    pose2: 第i+1帧的相机位姿，形状为(7,)的数组 [tx2, ty2, tz2, qx2, qy2, qz2, qw2]
    
    返回:
    relative_matrix: 3×4的相对位姿矩阵，前3列是旋转矩阵R_rel，第4列是平移向量t_rel
    """
    # 分离平移向量和四元数
    t1 = pose1[:3]  # 第i帧平移 [tx1, ty1, tz1]
    q1 = pose1[3:]  # 第i帧四元数 [qx1, qy1, qz1, qw1]
    t2 = pose2[:3]  # 第i+1帧平移
    q2 = pose2[3:]  # 第i+1帧四元数
    
    # 1. 计算相对旋转矩阵 R_rel
    rot1 = R.from_quat(q1)  # 第i帧旋转
    rot2 = R.from_quat(q2)  # 第i+1帧旋转
    rot_rel = rot2 * rot1.inv()  # 相对旋转 = 后一帧旋转 × 前一帧旋转的逆
    R_rel = rot_rel.as_matrix()  # 转换为3×3矩阵
    
    # 2. 计算相对平移向量 t_rel
    R1_T = rot1.as_matrix().T  # 前一帧旋转矩阵的转置（等价于逆）
    t_rel = R1_T @ (t2 - t1)   # 相对平移 = R1^T × (t2 - t1)
    
    # 3. 组合为3×4矩阵 [R_rel | t_rel]
    relative_matrix = np.hstack([R_rel, t_rel.reshape(3, 1)])
    
    return relative_matrix

class MultiDatasetDynamicDataset(torch.utils.data.Dataset):
    """支持FramePack机制的多数据集动态历史长度数据集 - 融合nuscenes和sekai"""
    
    def __init__(self, dataset_configs, steps_per_epoch, 
                 min_condition_frames=10, max_condition_frames=40,
                 target_frames=10, height=900, width=1600):
        """
        Args:
            dataset_configs: 数据集配置列表，每个配置包含 {
                'name': 数据集名称,
                'paths': 数据集路径列表,
                'type': 数据集类型 ('sekai' 或 'nuscenes'),
                'weight': 采样权重
            }
        """
        self.dataset_configs = dataset_configs
        self.min_condition_frames = min_condition_frames
        self.max_condition_frames = max_condition_frames
        self.target_frames = target_frames
        self.height = height
        self.width = width
        self.steps_per_epoch = steps_per_epoch
        self.pose_classifier = PoseClassifier()
        
        # VAE时间压缩比例
        self.time_compression_ratio = 4
        
        # 🔧 扫描所有数据集，建立统一的场景索引
        self.scene_dirs = []
        self.dataset_info = {}  # 记录每个场景的数据集信息
        self.dataset_weights = []  # 每个场景的采样权重
        
        total_scenes = 0
        
        for config in self.dataset_configs:
            dataset_name = config['name']
            dataset_paths = config['paths'] if isinstance(config['paths'], list) else [config['paths']]
            dataset_type = config['type']
            dataset_weight = config.get('weight', 1.0)
            
            print(f"🔧 扫描数据集: {dataset_name} (类型: {dataset_type})")
            
            dataset_scenes = []
            for dataset_path in dataset_paths:
                print(f"  📁 检查路径: {dataset_path}")
                if os.path.exists(dataset_path):                    
                    if dataset_type == 'nuscenes':
                        # NuScenes使用 base_path/scenes 结构
                        scenes_path = os.path.join(dataset_path, "scenes")
                        print(f"  📂 扫描NuScenes scenes目录: {scenes_path}")
                        for item in os.listdir(scenes_path):
                            scene_dir = os.path.join(scenes_path, item)
                            if os.path.isdir(scene_dir):
                                self.scene_dirs.append(scene_dir)
                                dataset_scenes.append(scene_dir)
                                self.dataset_info[scene_dir] = {
                                    'name': dataset_name,
                                    'type': dataset_type,
                                    'weight': dataset_weight
                                }
                                self.dataset_weights.append(dataset_weight)

                    elif dataset_type in ['sekai', 'spatialvid', 'openx']:
                        # Sekai、spatialvid、OpenX等数据集直接扫描根目录
                        for item in os.listdir(dataset_path):
                            scene_dir = os.path.join(dataset_path, item)
                            if os.path.isdir(scene_dir):
                                encoded_path = os.path.join(scene_dir, "encoded_video.pth")
                                if os.path.exists(encoded_path):
                                    self.scene_dirs.append(scene_dir)
                                    dataset_scenes.append(scene_dir)
                                    self.dataset_info[scene_dir] = {
                                        'name': dataset_name,
                                        'type': dataset_type,
                                        'weight': dataset_weight
                                    }
                                    self.dataset_weights.append(dataset_weight)
                else:
                    print(f"  ❌ 路径不存在: {dataset_path}")
                
                print(f"  ✅ 找到 {len(dataset_scenes)} 个场景")
                total_scenes += len(dataset_scenes)
                    
        # 统计各数据集场景数
        dataset_counts = {}
        for scene_dir in self.scene_dirs:
            dataset_name = self.dataset_info[scene_dir]['name']
            dataset_type = self.dataset_info[scene_dir]['type']
            key = f"{dataset_name} ({dataset_type})"
            dataset_counts[key] = dataset_counts.get(key, 0) + 1
        
        for dataset_key, count in dataset_counts.items():
            print(f"  - {dataset_key}: {count} 个场景")
        
        assert len(self.scene_dirs) > 0, "No encoded scenes found!"
        
        # 🔧 计算采样概率
        total_weight = sum(self.dataset_weights)
        self.sampling_probs = [w / total_weight for w in self.dataset_weights]

    def calculate_relative_rotation(self, current_rotation, reference_rotation):
        """计算相对旋转四元数 - NuScenes专用"""
        q_current = torch.tensor(current_rotation, dtype=torch.float32)
        q_ref = torch.tensor(reference_rotation, dtype=torch.float32)

        q_ref_inv = torch.tensor([q_ref[0], -q_ref[1], -q_ref[2], -q_ref[3]])

        w1, x1, y1, z1 = q_ref_inv
        w2, x2, y2, z2 = q_current

        relative_rotation = torch.tensor([
            w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2,
            w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2,
            w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2,
            w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2
        ])

        return relative_rotation

    def prepare_framepack_inputs(self, full_latents, segment_info):
        """🔧 准备FramePack风格的多尺度输入 - 修正版，正确处理空索引"""
        # 🔧 修正：处理4维输入 [C, T, H, W]，添加batch维度
        if len(full_latents.shape) == 4:
            full_latents = full_latents.unsqueeze(0)  # [C, T, H, W] -> [1, C, T, H, W]
            B, C, T, H, W = full_latents.shape
        else:
            B, C, T, H, W = full_latents.shape
        
        # 主要latents（用于去噪预测）
        latent_indices = segment_info['latent_indices']
        main_latents = full_latents[:, :, latent_indices, :, :]  # 注意维度顺序
        
        # 🔧 1x条件帧（起始帧 + 最后1帧）
        clean_latent_indices = segment_info['clean_latent_indices']
        clean_latents = full_latents[:, :, clean_latent_indices, :, :]  # 注意维度顺序
        
        # 🔧 4x条件帧 - 总是16帧，直接用真实索引 + 0填充
        clean_latent_4x_indices = segment_info['clean_latent_4x_indices']
        
        # 创建固定长度16的latents，初始化为0
        clean_latents_4x = torch.zeros(B, C, 16, H, W, dtype=full_latents.dtype)
        clean_latent_4x_indices_final = torch.full((16,), -1, dtype=torch.long)  # -1表示padding
        
        # 🔧 修正：检查是否有有效的4x索引
        if len(clean_latent_4x_indices) > 0:
            actual_4x_frames = len(clean_latent_4x_indices)
            # 从后往前填充，确保最新的帧在最后
            start_pos = max(0, 16 - actual_4x_frames)
            end_pos = 16
            actual_start = max(0, actual_4x_frames - 16)  # 如果超过16帧，只取最后16帧
            
            clean_latents_4x[:, :, start_pos:end_pos, :, :] = full_latents[:, :, clean_latent_4x_indices[actual_start:], :, :]
            clean_latent_4x_indices_final[start_pos:end_pos] = clean_latent_4x_indices[actual_start:]
        
        # 🔧 2x条件帧 - 总是2帧，直接用真实索引 + 0填充
        clean_latent_2x_indices = segment_info['clean_latent_2x_indices']
        
        # 创建固定长度2的latents，初始化为0
        clean_latents_2x = torch.zeros(B, C, 2, H, W, dtype=full_latents.dtype)
        clean_latent_2x_indices_final = torch.full((2,), -1, dtype=torch.long)  # -1表示padding
        
        # 🔧 修正：检查是否有有效的2x索引
        if len(clean_latent_2x_indices) > 0:
            actual_2x_frames = len(clean_latent_2x_indices)
            # 从后往前填充，确保最新的帧在最后
            start_pos = max(0, 2 - actual_2x_frames)
            end_pos = 2
            actual_start = max(0, actual_2x_frames - 2)  # 如果超过2帧，只取最后2帧
            
            clean_latents_2x[:, :, start_pos:end_pos, :, :] = full_latents[:, :, clean_latent_2x_indices[actual_start:], :, :]
            clean_latent_2x_indices_final[start_pos:end_pos] = clean_latent_2x_indices[actual_start:]
        
        # 🔧 移除添加的batch维度，返回原始格式
        if B == 1:
            main_latents = main_latents.squeeze(0)  # [1, C, T, H, W] -> [C, T, H, W]
            clean_latents = clean_latents.squeeze(0)
            clean_latents_2x = clean_latents_2x.squeeze(0)
            clean_latents_4x = clean_latents_4x.squeeze(0)
        
        return {
            'latents': main_latents,
            'clean_latents': clean_latents,
            'clean_latents_2x': clean_latents_2x,
            'clean_latents_4x': clean_latents_4x,
            'latent_indices': segment_info['latent_indices'],
            'clean_latent_indices': segment_info['clean_latent_indices'],
            'clean_latent_2x_indices': clean_latent_2x_indices_final,  # 🔧 使用真实索引（含-1填充）
            'clean_latent_4x_indices': clean_latent_4x_indices_final,  # 🔧 使用真实索引（含-1填充）
        }

    def create_sekai_pose_embeddings(self, cam_data, segment_info):
        """创建Sekai风格的pose embeddings"""
        cam_data_seq = cam_data['extrinsic']
        
        # 为所有帧计算相对pose
        all_keyframe_indices = []
        for compressed_idx in range(segment_info['start_frame'], segment_info['target_range'][1]):
            all_keyframe_indices.append(compressed_idx * 4)
        
        relative_cams = []
        for idx in all_keyframe_indices:
            cam_prev = cam_data_seq[idx]
            cam_next = cam_data_seq[idx + 4]
            relative_cam = compute_relative_pose(cam_prev, cam_next)
            relative_cams.append(torch.as_tensor(relative_cam[:3, :]))
        
        pose_embedding = torch.stack(relative_cams, dim=0)
        pose_embedding = rearrange(pose_embedding, 'b c d -> b (c d)')
        pose_embedding = pose_embedding.to(torch.bfloat16)

        return pose_embedding

    def create_openx_pose_embeddings(self, cam_data, segment_info):
        """🔧 创建OpenX风格的pose embeddings - 类似sekai但处理更短的序列"""
        cam_data_seq = cam_data['extrinsic']
        
        # 为所有帧计算相对pose - OpenX使用4倍间隔
        all_keyframe_indices = []
        for compressed_idx in range(segment_info['start_frame'], segment_info['target_range'][1]):
            keyframe_idx = compressed_idx * 4
            if keyframe_idx + 4 < len(cam_data_seq):
                all_keyframe_indices.append(keyframe_idx)
        
        relative_cams = []
        for idx in all_keyframe_indices:
            if idx + 4 < len(cam_data_seq):
                cam_prev = cam_data_seq[idx]
                cam_next = cam_data_seq[idx + 4]
                relative_cam = compute_relative_pose(cam_prev, cam_next)
                relative_cams.append(torch.as_tensor(relative_cam[:3, :]))
            else:
                # 如果没有下一帧，使用单位矩阵
                identity_cam = torch.eye(3, 4)
                relative_cams.append(identity_cam)
        
        if len(relative_cams) == 0:
            return None
            
        pose_embedding = torch.stack(relative_cams, dim=0)
        pose_embedding = rearrange(pose_embedding, 'b c d -> b (c d)')
        pose_embedding = pose_embedding.to(torch.bfloat16)

        return pose_embedding
    
    def create_spatialvid_pose_embeddings(self, cam_data, segment_info):
        """🔧 创建SpatialVid风格的pose embeddings - camera间隔为1帧而非4帧"""
        cam_data_seq = cam_data['extrinsic']   # N * 4 * 4
        
        # 🔧 为所有帧（condition + target）计算camera embedding
        # SpatialVid特有：每隔1帧而不是4帧
        keyframe_original_idx = segment_info['keyframe_original_idx']
        
        relative_cams = []
        for idx in keyframe_original_idx:
            if idx + 1 < len(cam_data_seq):
                cam_prev = cam_data_seq[idx]
                cam_next = cam_data_seq[idx + 1]  # SpatialVid: 每隔1帧
                relative_cam = compute_relative_pose_matrix(cam_prev, cam_next)
                relative_cams.append(torch.as_tensor(relative_cam[:3, :]))
            else:
                # 如果没有下一帧，使用零运动
                identity_cam = torch.zeros(3, 4)
                relative_cams.append(identity_cam)
        
        if len(relative_cams) == 0:
            return None
            
        pose_embedding = torch.stack(relative_cams, dim=0)
        pose_embedding = rearrange(pose_embedding, 'b c d -> b (c d)')
        pose_embedding = pose_embedding.to(torch.bfloat16)

        return pose_embedding
                
    def create_nuscenes_pose_embeddings_framepack(self, scene_info, segment_info):
        """创建NuScenes风格的pose embeddings - 每帧都相对上一帧（7维）"""
        keyframe_poses = scene_info['keyframe_poses']
        # 生成所有需要的关键帧索引
        start_frame = segment_info['start_frame']
        total_frames = segment_info['condition_frames'] + segment_info['target_frames']
        keyframe_indices = []
        for i in range(total_frames + 1):  # +1是因为需要前后两帧
            idx = (start_frame + i) * self.time_compression_ratio
            keyframe_indices.append(idx)
        # 边界检查
        keyframe_indices = [min(idx, len(keyframe_poses)-1) for idx in keyframe_indices]

        pose_vecs = []
        for i in range(total_frames):
            pose_prev = keyframe_poses[keyframe_indices[i]]
            pose_next = keyframe_poses[keyframe_indices[i+1]]
            # 计算相对位姿
            translation = torch.tensor(
                np.array(pose_next['translation']) - np.array(pose_prev['translation']),
                dtype=torch.float32
            )
            relative_rotation = self.calculate_relative_rotation(
                pose_next['rotation'],
                pose_prev['rotation']
            )
            pose_vec = torch.cat([translation, relative_rotation], dim=0)  # [7D]
            pose_vecs.append(pose_vec)

        if not pose_vecs:
            return None
        pose_sequence = torch.stack(pose_vecs, dim=0)  # [total_frames, 7]
        return pose_sequence

    # 修改select_dynamic_segment方法
    def select_dynamic_segment(self, full_latents, dataset_type, scene_info=None):
        """🔧 根据数据集类型选择不同的段落选择策略"""
        # 原有的sekai方式
        total_lens = full_latents.shape[1]
        
        min_condition_compressed = self.min_condition_frames // self.time_compression_ratio
        max_condition_compressed = self.max_condition_frames // self.time_compression_ratio
        target_frames_compressed = self.target_frames // self.time_compression_ratio
        max_condition_compressed = min(total_lens-target_frames_compressed-1, max_condition_compressed)

        ratio = random.random()
        if ratio < 0.15:
            condition_frames_compressed = 1
        elif 0.15 <= ratio < 0.9 or total_lens <= 2*target_frames_compressed + 1:
            condition_frames_compressed = random.randint(min_condition_compressed, max_condition_compressed)
        else:
            condition_frames_compressed = target_frames_compressed
        
        # 确保有足够的帧数
        min_required_frames = condition_frames_compressed + target_frames_compressed
        if total_lens < min_required_frames:
            return None
        
        start_frame_compressed = random.randint(0, total_lens - min_required_frames - 1)
        condition_end_compressed = start_frame_compressed + condition_frames_compressed
        target_end_compressed = condition_end_compressed + target_frames_compressed

        # FramePack风格的索引处理
        latent_indices = torch.arange(condition_end_compressed, target_end_compressed)
        
        # 1x帧：起始帧 + 最后1帧
        clean_latent_indices_start = torch.tensor([start_frame_compressed])
        clean_latent_1x_indices = torch.tensor([condition_end_compressed - 1])
        clean_latent_indices = torch.cat([clean_latent_indices_start, clean_latent_1x_indices])
        
        # 🔧 2x帧：根据实际condition长度确定
        if condition_frames_compressed >= 2:
            # 取最后2帧（如果有的话）
            clean_latent_2x_start = max(start_frame_compressed, condition_end_compressed - 2-1)
            clean_latent_2x_indices = torch.arange(clean_latent_2x_start, condition_end_compressed-1)
        else:
            # 如果condition帧数不足2帧，创建空索引
            clean_latent_2x_indices = torch.tensor([], dtype=torch.long)
        
        # 🔧 4x帧：根据实际condition长度确定，最多16帧
        if condition_frames_compressed > 3:
            # 取最多16帧的历史（如果有的话）
            clean_4x_start = max(start_frame_compressed, condition_end_compressed - 16-3)
            clean_latent_4x_indices = torch.arange(clean_4x_start, condition_end_compressed-3)
        else:
            clean_latent_4x_indices = torch.tensor([], dtype=torch.long)
        
        # 对应的原始关键帧索引
        keyframe_original_idx = []
        for compressed_idx in range(start_frame_compressed, target_end_compressed):
            if dataset_type == 'spatialvid':
                keyframe_original_idx.append(compressed_idx)  # spatialvid直接使用compressed_idx
            elif dataset_type == 'openx' or 'sekai' or "nuscenes":  # 🔧 新增openx处理
                keyframe_original_idx.append(compressed_idx * 4)  # openx使用4倍间隔

        return {
            'start_frame': start_frame_compressed,
            'condition_frames': condition_frames_compressed,
            'target_frames': target_frames_compressed,
            'condition_range': (start_frame_compressed, condition_end_compressed),
            'target_range': (condition_end_compressed, target_end_compressed),
            
            # FramePack风格的索引
            'latent_indices': latent_indices,
            'clean_latent_indices': clean_latent_indices,
            'clean_latent_2x_indices': clean_latent_2x_indices,
            'clean_latent_4x_indices': clean_latent_4x_indices,
            
            'keyframe_original_idx': keyframe_original_idx,
            'original_condition_frames': condition_frames_compressed * self.time_compression_ratio,
            'original_target_frames': target_frames_compressed * self.time_compression_ratio,
        }

    # 修改create_pose_embeddings方法
    def create_pose_embeddings(self, cam_data, segment_info, dataset_type, scene_info=None):
        """🔧 根据数据集类型创建pose embeddings"""
        if dataset_type == 'nuscenes' and scene_info is not None:
            return self.create_nuscenes_pose_embeddings_framepack(scene_info, segment_info)
        elif dataset_type == 'spatialvid':  # 🔧 新增spatialvid处理
            return self.create_spatialvid_pose_embeddings(cam_data, segment_info)
        elif dataset_type == 'sekai':
            return self.create_sekai_pose_embeddings(cam_data, segment_info)
        elif dataset_type == 'openx':  # 🔧 新增openx处理
            return self.create_openx_pose_embeddings(cam_data, segment_info)
        
    def __getitem__(self, index):
        while True:
            try:
                # 根据权重随机选择场景
                scene_idx = np.random.choice(len(self.scene_dirs), p=self.sampling_probs)
                scene_dir = self.scene_dirs[scene_idx]
                dataset_info = self.dataset_info[scene_dir]
                
                dataset_name = dataset_info['name']
                dataset_type = dataset_info['type']
                
                # 🔧 根据数据集类型加载数据
                scene_info = None
                if dataset_type == 'nuscenes':
                    # NuScenes需要加载scene_info.json
                    scene_info_path = os.path.join(scene_dir, "scene_info.json")
                    if os.path.exists(scene_info_path):
                        with open(scene_info_path, 'r') as f:
                            scene_info = json.load(f)
                    
                    # NuScenes使用不同的编码文件名
                    encoded_path = os.path.join(scene_dir, "encoded_video-480p.pth")
                    if not os.path.exists(encoded_path):
                        encoded_path = os.path.join(scene_dir, "encoded_video.pth")  # fallback
                    
                    encoded_data = torch.load(encoded_path, weights_only=True, map_location="cpu")
                else:
                    # Sekai数据集
                    encoded_path = os.path.join(scene_dir, "encoded_video.pth")
                    encoded_data = torch.load(encoded_path, weights_only=False, map_location="cpu")
                
                full_latents = encoded_data['latents']
                if full_latents.shape[1] <= 10:
                    continue
                cam_data = encoded_data.get('cam_emb', encoded_data)
                
                # 🔧 验证NuScenes的latent帧数
                if dataset_type == 'nuscenes' and scene_info is not None:
                    expected_latent_frames = scene_info['total_frames'] // self.time_compression_ratio
                    actual_latent_frames = full_latents.shape[1]
                    
                    if abs(actual_latent_frames - expected_latent_frames) > 2:
                        print(f"⚠️  NuScenes Latent帧数不匹配，跳过此样本")
                        continue
                
                # 使用数据集特定的段落选择策略
                segment_info = self.select_dynamic_segment(full_latents, dataset_type, scene_info)
                if segment_info is None:
                    continue
                
                # 创建数据集特定的pose embeddings
                all_camera_embeddings = self.create_pose_embeddings(cam_data, segment_info, dataset_type, scene_info)
                if all_camera_embeddings is None:
                    continue
                
                # 准备FramePack风格的多尺度输入
                framepack_inputs = self.prepare_framepack_inputs(full_latents, segment_info)
                
                n = segment_info["condition_frames"]
                m = segment_info['target_frames']
                
                # 处理camera embedding with mask
                mask = torch.zeros(n+m, dtype=torch.float32)
                mask[:n] = 1.0
                mask = mask.view(-1, 1)
                
                # 🔧 NuScenes返回的是直接的embedding，Sekai返回的是tensor
                if isinstance(all_camera_embeddings, torch.Tensor):
                    camera_with_mask = torch.cat([all_camera_embeddings, mask], dim=1)
                else:
                    # NuScenes风格，直接就是最终的embedding
                    camera_with_mask = torch.cat([all_camera_embeddings, mask], dim=1)
                
                result = {
                    # FramePack风格的多尺度输入
                    "latents": framepack_inputs['latents'],
                    "clean_latents": framepack_inputs['clean_latents'],
                    "clean_latents_2x": framepack_inputs['clean_latents_2x'],
                    "clean_latents_4x": framepack_inputs['clean_latents_4x'],
                    "latent_indices": framepack_inputs['latent_indices'],
                    "clean_latent_indices": framepack_inputs['clean_latent_indices'],
                    "clean_latent_2x_indices": framepack_inputs['clean_latent_2x_indices'],
                    "clean_latent_4x_indices": framepack_inputs['clean_latent_4x_indices'],
                    
                    # Camera数据
                    "camera": camera_with_mask,
                    
                    # 其他数据
                    "prompt_emb": encoded_data["prompt_emb"],
                    "image_emb": encoded_data.get("image_emb", {}),
                    
                    # 元信息
                    "condition_frames": n,
                    "target_frames": m,
                    "scene_name": os.path.basename(scene_dir),
                    "dataset_name": dataset_name,
                    "dataset_type": dataset_type,
                    "original_condition_frames": segment_info['original_condition_frames'],
                    "original_target_frames": segment_info['original_target_frames'],
                }
                
                return result
                
            except Exception as e:
                print(f"Error loading sample: {e}")
                traceback.print_exc()
                continue

    def __len__(self):
        return self.steps_per_epoch

def replace_dit_model_in_manager():
    """在模型加载前替换DiT模型类为MoE版本"""
    from diffsynth.models.wan_video_dit_moe import WanModelMoe
    from diffsynth.configs.model_config import model_loader_configs
    
    # 修改model_loader_configs中的配置
    for i, config in enumerate(model_loader_configs):
        keys_hash, keys_hash_with_shape, model_names, model_classes, model_resource = config
        
        # 检查是否包含wan_video_dit模型
        if 'wan_video_dit' in model_names:
            new_model_names = []
            new_model_classes = []
            
            for name, cls in zip(model_names, model_classes):
                if name == 'wan_video_dit':
                    new_model_names.append(name)
                    new_model_classes.append(WanModelMoe)  # 🔧 使用MoE版本
                    print(f"✅ 替换了模型类: {name} -> WanModelMoe")
                else:
                    new_model_names.append(name)
                    new_model_classes.append(cls)
            
            # 更新配置
            model_loader_configs[i] = (keys_hash, keys_hash_with_shape, new_model_names, new_model_classes, model_resource)

class MultiDatasetLightningModelForTrain(pl.LightningModule):
    def __init__(
        self,
        dit_path,
        learning_rate=1e-5,
        use_gradient_checkpointing=True,
        use_gradient_checkpointing_offload=False,
        resume_ckpt_path=None,
        # 🔧 MoE参数
        use_moe=False,
        moe_config=None
    ):
        super().__init__()
        self.use_moe = use_moe
        self.moe_config = moe_config or {}
        
        replace_dit_model_in_manager()
        model_manager = ModelManager(torch_dtype=torch.bfloat16, device="cpu")
        if os.path.isfile(dit_path):
            model_manager.load_models([dit_path])
        else:
            dit_path = dit_path.split(",")
            model_manager.load_models([dit_path])
        model_manager.load_models(["models/Wan-AI/Wan2.1-T2V-1.3B/Wan2.1_VAE.pth"])
        
        self.pipe = WanVideoReCamMasterPipeline.from_model_manager(model_manager)
        self.pipe.scheduler.set_timesteps(1000, training=True)

        # 添加FramePack的clean_x_embedder
        self.add_framepack_components()
        if self.use_moe:
            self.add_moe_components()

        # 🔧 添加camera编码器（wan_video_dit_moe.py已经包含MoE逻辑）
        dim = self.pipe.dit.blocks[0].self_attn.q.weight.shape[0]
        for block in self.pipe.dit.blocks:
            # 🔧 简化：只添加传统camera编码器，MoE逻辑在wan_video_dit_moe.py中
            block.cam_encoder = nn.Linear(13, dim)
            block.projector = nn.Linear(dim, dim)
            block.cam_encoder.weight.data.zero_()
            block.cam_encoder.bias.data.zero_()
            block.projector.weight = nn.Parameter(torch.eye(dim))
            block.projector.bias = nn.Parameter(torch.zeros(dim))
        
        if resume_ckpt_path is not None:
            state_dict = torch.load(resume_ckpt_path, map_location="cpu")
            state_dict.pop("global_router.weight", None)
            state_dict.pop("global_router.bias", None)
            self.pipe.dit.load_state_dict(state_dict, strict=False)
            print('load checkpoint:', resume_ckpt_path)

        self.freeze_parameters()
        
        # 🔧 训练参数设置
        for name, module in self.pipe.denoising_model().named_modules():
            if any(keyword in name for keyword in ["cam_encoder", "projector", "self_attn", "clean_x_embedder", 
                                                "moe", "sekai_processor", "nuscenes_processor","openx_processor"]):
                for param in module.parameters():
                    param.requires_grad = True
        
        self.learning_rate = learning_rate
        self.use_gradient_checkpointing = use_gradient_checkpointing
        self.use_gradient_checkpointing_offload = use_gradient_checkpointing_offload
        
        # 创建可视化目录
        self.vis_dir = "multi_dataset_dynamic/visualizations"
        os.makedirs(self.vis_dir, exist_ok=True)

    def add_moe_components(self):
        """🔧 添加MoE相关组件 - 简化版，只为每个block添加MoE，全局processor在WanModelMoe中"""
        if not hasattr(self.pipe.dit, 'moe_config'):
            self.pipe.dit.moe_config = self.moe_config
            print("✅ 添加了MoE配置到模型")
        self.pipe.dit.top_k = self.moe_config.get("top_k", 1)
        
        # 为每个block添加MoE组件（modality processors已经在WanModelMoe中全局创建）
        dim = self.pipe.dit.blocks[0].self_attn.q.weight.shape[0]
        unified_dim = self.moe_config.get("unified_dim", 30)
        num_experts = self.moe_config.get("num_experts", 4)
        from diffsynth.models.wan_video_dit_moe import MultiModalMoE, ModalityProcessor

        self.pipe.dit.sekai_processor = ModalityProcessor("sekai", 13, unified_dim)
        self.pipe.dit.nuscenes_processor = ModalityProcessor("nuscenes", 8, unified_dim)
        self.pipe.dit.openx_processor = ModalityProcessor("openx", 13, unified_dim)  # OpenX使用13维输入，类似sekai但独立处理
        self.pipe.dit.global_router = nn.Linear(unified_dim, num_experts)

        for i, block in enumerate(self.pipe.dit.blocks):            
            # 只为每个block添加MoE网络
            block.moe = MultiModalMoE(
                unified_dim=unified_dim,
                output_dim=dim,
                num_experts=self.moe_config.get("num_experts", 4),
                top_k=self.moe_config.get("top_k", 2)
            )
            
            print(f"✅ Block {i} 添加了MoE组件 (unified_dim: {unified_dim}, experts: {self.moe_config.get('num_experts', 4)})")


    def add_framepack_components(self):
        """🔧 添加FramePack相关组件"""
        if not hasattr(self.pipe.dit, 'clean_x_embedder'):
            inner_dim = self.pipe.dit.blocks[0].self_attn.q.weight.shape[0]
            
            class CleanXEmbedder(nn.Module):
                def __init__(self, inner_dim):
                    super().__init__()
                    self.proj = nn.Conv3d(16, inner_dim, kernel_size=(1, 2, 2), stride=(1, 2, 2))
                    self.proj_2x = nn.Conv3d(16, inner_dim, kernel_size=(2, 4, 4), stride=(2, 4, 4))
                    self.proj_4x = nn.Conv3d(16, inner_dim, kernel_size=(4, 8, 8), stride=(4, 8, 8))
                
                def forward(self, x, scale="1x"):
                    if scale == "1x":
                        return self.proj(x)
                    elif scale == "2x":
                        return self.proj_2x(x)
                    elif scale == "4x":
                        return self.proj_4x(x)
                    else:
                        raise ValueError(f"Unsupported scale: {scale}")
            
            self.pipe.dit.clean_x_embedder = CleanXEmbedder(inner_dim)
            print("✅ 添加了FramePack的clean_x_embedder组件")
        
    def freeze_parameters(self):
        self.pipe.requires_grad_(False)
        self.pipe.eval()
        self.pipe.denoising_model().train()

    def training_step(self, batch, batch_idx):
        """🔧 多数据集训练步骤"""
        condition_frames = batch["condition_frames"][0].item()
        target_frames = batch["target_frames"][0].item()
        
        original_condition_frames = batch.get("original_condition_frames", [condition_frames * 4])[0]
        original_target_frames = batch.get("original_target_frames", [target_frames * 4])[0]

        dataset_name = batch.get("dataset_name", ["unknown"])[0]
        dataset_type = batch.get("dataset_type", ["sekai"])[0]
        scene_name = batch.get("scene_name", ["unknown"])[0]
        
        # 准备输入数据
        latents = batch["latents"].to(self.device)
        if len(latents.shape) == 4:
            latents = latents.unsqueeze(0)
        
        clean_latents = batch["clean_latents"].to(self.device) if batch["clean_latents"].numel() > 0 else None
        if clean_latents is not None and len(clean_latents.shape) == 4:
            clean_latents = clean_latents.unsqueeze(0)
        
        clean_latents_2x = batch["clean_latents_2x"].to(self.device) if batch["clean_latents_2x"].numel() > 0 else None
        if clean_latents_2x is not None and len(clean_latents_2x.shape) == 4:
            clean_latents_2x = clean_latents_2x.unsqueeze(0)
        
        clean_latents_4x = batch["clean_latents_4x"].to(self.device) if batch["clean_latents_4x"].numel() > 0 else None
        if clean_latents_4x is not None and len(clean_latents_4x.shape) == 4:
            clean_latents_4x = clean_latents_4x.unsqueeze(0)
        
        # 索引处理
        latent_indices = batch["latent_indices"].to(self.device)
        clean_latent_indices = batch["clean_latent_indices"].to(self.device) if batch["clean_latent_indices"].numel() > 0 else None
        clean_latent_2x_indices = batch["clean_latent_2x_indices"].to(self.device) if batch["clean_latent_2x_indices"].numel() > 0 else None
        clean_latent_4x_indices = batch["clean_latent_4x_indices"].to(self.device) if batch["clean_latent_4x_indices"].numel() > 0 else None
        
        # Camera embedding处理
        cam_emb = batch["camera"].to(self.device)
        
        # 🔧 根据数据集类型设置modality_inputs
        if dataset_type == "sekai":
            modality_inputs = {"sekai": cam_emb}
        elif dataset_type == "spatialvid":  # 🔧 spatialvid使用sekai processor
            modality_inputs = {"sekai": cam_emb}  # 注意：这里使用"sekai"键
        elif dataset_type == "nuscenes":
            modality_inputs = {"nuscenes": cam_emb}
        elif dataset_type == "openx":  # 🔧 新增：openx使用独立的processor
            modality_inputs = {"openx": cam_emb}
        else:
            modality_inputs = {"sekai": cam_emb}  # 默认
        
        camera_dropout_prob = 0.05
        if random.random() < camera_dropout_prob:
            cam_emb = torch.zeros_like(cam_emb)
            # 同时清空modality_inputs
            for key in modality_inputs:
                modality_inputs[key] = torch.zeros_like(modality_inputs[key])
            print(f"应用camera dropout for CFG training (dataset: {dataset_name}, type: {dataset_type})")
        
        prompt_emb = batch["prompt_emb"]
        prompt_emb["context"] = prompt_emb["context"][0].to(self.device)
        image_emb = batch["image_emb"]

        if "clip_feature" in image_emb:
            image_emb["clip_feature"] = image_emb["clip_feature"][0].to(self.device)
        if "y" in image_emb:
            image_emb["y"] = image_emb["y"][0].to(self.device)

        # Loss计算
        self.pipe.device = self.device
        noise = torch.randn_like(latents)
        timestep_id = torch.randint(0, self.pipe.scheduler.num_train_timesteps, (1,))
        timestep = self.pipe.scheduler.timesteps[timestep_id].to(dtype=self.pipe.torch_dtype, device=self.pipe.device)
        
        # FramePack风格的噪声处理
        noisy_condition_latents = None
        if clean_latents is not None:
            noisy_condition_latents = copy.deepcopy(clean_latents)
            is_add_noise = random.random()
            if is_add_noise > 0.2:
                noise_cond = torch.randn_like(clean_latents)
                timestep_id_cond = torch.randint(0, self.pipe.scheduler.num_train_timesteps//4*3, (1,))
                timestep_cond = self.pipe.scheduler.timesteps[timestep_id_cond].to(dtype=self.pipe.torch_dtype, device=self.pipe.device)
                noisy_condition_latents = self.pipe.scheduler.add_noise(clean_latents, noise_cond, timestep_cond)

        extra_input = self.pipe.prepare_extra_input(latents)
        origin_latents = copy.deepcopy(latents)
        noisy_latents = self.pipe.scheduler.add_noise(latents, noise, timestep)
        
        training_target = self.pipe.scheduler.training_target(latents, noise, timestep)
        
        noise_pred, specialization_loss = self.pipe.denoising_model()(
            noisy_latents, 
            timestep=timestep, 
            cam_emb=cam_emb,
            modality_inputs=modality_inputs,  # 🔧 传递多模态输入
            latent_indices=latent_indices,
            clean_latents=noisy_condition_latents if noisy_condition_latents is not None else clean_latents,
            clean_latent_indices=clean_latent_indices,
            clean_latents_2x=clean_latents_2x,
            clean_latent_2x_indices=clean_latent_2x_indices,
            clean_latents_4x=clean_latents_4x,
            clean_latent_4x_indices=clean_latent_4x_indices,
            **prompt_emb, 
            **extra_input, 
            **image_emb,
            use_gradient_checkpointing=self.use_gradient_checkpointing,
            use_gradient_checkpointing_offload=self.use_gradient_checkpointing_offload
        )
        
        # 计算loss
        # 🔧 计算总loss = 重建loss + MoE专业化loss
        reconstruction_loss = torch.nn.functional.mse_loss(noise_pred.float(), training_target.float())
        reconstruction_loss = reconstruction_loss * self.pipe.scheduler.training_weight(timestep)
        
        # 🔧 添加MoE专业化loss（交叉熵损失）
        specialization_loss_weight = self.moe_config.get("moe_loss_weight", 0.1)
        total_loss = reconstruction_loss + specialization_loss_weight * specialization_loss
        
        print(f'\n loss info (step {self.global_step}):')
        print(f'  - diff loss: {reconstruction_loss.item():.6f}')
        print(f'  - MoE specification loss: {specialization_loss.item():.6f}')
        print(f'  - Expert loss weight: {specialization_loss_weight}')
        print(f'  - Total Loss: {total_loss.item():.6f}')
        
        # 🔧 显示预期的专家映射
        modality_to_expert = {
            "sekai": 0,
            "nuscenes": 1, 
            "openx": 2
        }
        expected_expert = modality_to_expert.get(dataset_type, 0)
        print(f'  - current modality: {dataset_type} -> expected expert: {expected_expert}')

        return total_loss

    def configure_optimizers(self):
        trainable_modules = filter(lambda p: p.requires_grad, self.pipe.denoising_model().parameters())
        optimizer = torch.optim.AdamW(trainable_modules, lr=self.learning_rate)
        return optimizer
    
    def on_save_checkpoint(self, checkpoint):
        checkpoint_dir = "/share_zhuyixuan05/zhuyixuan05/ICLR2026/framepack_moe"
        os.makedirs(checkpoint_dir, exist_ok=True)
        
        current_step = self.global_step
        checkpoint.clear()
        t = time.strftime("%Y%m%d-%H%M%S")   # 20250923-153047

        
        state_dict = self.pipe.denoising_model().state_dict()
        torch.save(state_dict, os.path.join(checkpoint_dir, f"step{current_step}_nus_moe_from_{t}.ckpt"))
        print(f"Saved MoE model checkpoint: step{current_step}.ckpt")

def train_multi_dataset(args):
    """训练支持多数据集MoE的模型"""
    
    # 🔧 数据集配置
    dataset_configs = [
        {
            'name': 'sekai-drone',
            'paths': ['/share_zhuyixuan05/zhuyixuan05/sekai-game-drone'],
            'type': 'sekai',
            'weight': 0.1
        },
        {
            'name': 'sekai-walking',
            'paths': ['/share_zhuyixuan05/zhuyixuan05/sekai-game-walking'],
            'type': 'sekai',
            'weight': 0.1
        },
        {
            'name': 'spatialvid',
            'paths': ['/share_zhuyixuan05/zhuyixuan05/spatialvid'],
            'type': 'spatialvid',
            'weight': 0.1
        },
        {
            'name': 'nuscenes',
            'paths': ['/share_zhuyixuan05/zhuyixuan05/nuscenes_video_generation_dynamic'],
            'type': 'nuscenes',
            'weight': 10.0
        },
        {
            'name': 'openx-fractal',
            'paths': ['/share_zhuyixuan05/zhuyixuan05/openx-fractal-encoded'],
            'type': 'openx',
            'weight': 0.1
        }
    ]
    
    dataset = MultiDatasetDynamicDataset(
        dataset_configs,
        steps_per_epoch=args.steps_per_epoch,
        min_condition_frames=args.min_condition_frames,
        max_condition_frames=args.max_condition_frames,
        target_frames=args.target_frames,
    )
    
    dataloader = torch.utils.data.DataLoader(
        dataset,
        shuffle=True,
        batch_size=1,
        num_workers=args.dataloader_num_workers
    )
    
    # 🔧 MoE配置
    moe_config = {
        "unified_dim": args.unified_dim,  # 新增
        "num_experts": args.moe_num_experts,
        "top_k": args.moe_top_k,
        "moe_loss_weight": args.moe_loss_weight,
        "sekai_input_dim": 13,
        "nuscenes_input_dim": 8,
        "openx_input_dim": 13  
    }
    
    model = MultiDatasetLightningModelForTrain(
        dit_path=args.dit_path,
        learning_rate=args.learning_rate,
        use_gradient_checkpointing=args.use_gradient_checkpointing,
        use_gradient_checkpointing_offload=args.use_gradient_checkpointing_offload,
        resume_ckpt_path=args.resume_ckpt_path,
        use_moe=True,  # 总是使用MoE
        moe_config=moe_config
    )

    trainer = pl.Trainer(
        max_epochs=args.max_epochs,
        accelerator="gpu",
        devices="auto",
        precision="bf16",
        strategy=args.training_strategy,
        default_root_dir=args.output_path,
        accumulate_grad_batches=args.accumulate_grad_batches,
        callbacks=[],
        logger=False
    )
    trainer.fit(model, dataloader)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Train Multi-Dataset FramePack with MoE")
    parser.add_argument("--dit_path", type=str, default="models/Wan-AI/Wan2.1-T2V-1.3B/diffusion_pytorch_model.safetensors")
    parser.add_argument("--output_path", type=str, default="./")
    parser.add_argument("--learning_rate", type=float, default=1e-5)
    parser.add_argument("--steps_per_epoch", type=int, default=8000)
    parser.add_argument("--max_epochs", type=int, default=100000)
    parser.add_argument("--min_condition_frames", type=int, default=8, help="最小条件帧数")
    parser.add_argument("--max_condition_frames", type=int, default=120, help="最大条件帧数")
    parser.add_argument("--target_frames", type=int, default=32, help="目标帧数")
    parser.add_argument("--dataloader_num_workers", type=int, default=4)
    parser.add_argument("--accumulate_grad_batches", type=int, default=1)
    parser.add_argument("--training_strategy", type=str, default="deepspeed_stage_1")
    parser.add_argument("--use_gradient_checkpointing", default=False)
    parser.add_argument("--use_gradient_checkpointing_offload", action="store_true")
    parser.add_argument("--resume_ckpt_path", type=str, default="/home/zhuyixuan05/ReCamMaster/nus_dynamic/step15000_dynamic.ckpt")
    
    # 🔧 MoE参数
    parser.add_argument("--unified_dim", type=int, default=25, help="统一的中间维度")
    parser.add_argument("--moe_num_experts", type=int, default=3, help="专家数量")
    parser.add_argument("--moe_top_k", type=int, default=1, help="Top-K专家")
    parser.add_argument("--moe_loss_weight", type=float, default=0.1, help="MoE损失权重")
    
    args = parser.parse_args()
    
    print("🔧 多数据集MoE训练配置:")
    print(f"  - 使用wan_video_dit_moe.py作为模型")
    print(f"  - 统一维度: {args.unified_dim}")
    print(f"  - 专家数量: {args.moe_num_experts}")
    print(f"  - Top-K: {args.moe_top_k}")
    print(f"  - MoE损失权重: {args.moe_loss_weight}")
    print("  - 数据集:")
    print("    - sekai-game-drone (sekai模态)")
    print("    - sekai-game-walking (sekai模态)")
    print("    - spatialvid (使用sekai模态处理器)") 
    print("    - openx-fractal (使用sekai模态处理器)")
    print(f"   - nuscenes (nuscenes模态)")
    
    train_multi_dataset(args)