models/dit/mmdit_back.py

import logging
from dataclasses import dataclass
from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F

# 假设这些是你原来的导入
from .mmdit_layers import compute_rope_rotations
from .mmdit_layers import TimestepEmbedder
from .mmdit_layers import MLP, ChannelLastConv1d, ConvMLP
from .mmdit_layers import (FinalBlock, MMDitSingleBlock, JointBlock_AT)

log = logging.getLogger()


@dataclass
class PreprocessedConditions:
    text_f: torch.Tensor
    text_f_c: torch.Tensor


class MMAudio(nn.Module):
    """
    一个修改版的 MMAudio 接口尽量和LayerFusionAudioDiT一致。
    """
    def __init__(self,
                 *,
                 latent_dim: int,
                 text_dim: int,
                 hidden_dim: int,
                 depth: int,
                 fused_depth: int,
                 num_heads: int,
                 mlp_ratio: float = 4.0,
                 latent_seq_len: int,
                 text_seq_len: int = 640,
                 # --- 新增参数，对齐 LayerFusionAudioDiT ---
                 ta_context_dim: int,
                 ta_context_fusion: str = 'add', # 'add' or 'concat'
                 ta_context_norm: bool = False,
                 # --- 其他原有参数 ---
                 empty_string_feat: Optional[torch.Tensor] = None,
                 v2: bool = False) -> None:
        super().__init__()

        self.v2 = v2
        self.latent_dim = latent_dim
        self._latent_seq_len = latent_seq_len
        self._text_seq_len = text_seq_len
        self.hidden_dim = hidden_dim
        self.num_heads = num_heads
        
        # --- 1. time_aligned_context 的投影层 ---
        # 我们在这里定义一个投影层，而不是在每个 block 里都定义一个。
        # 这样更高效，也符合你代码注释中的想法：“现在是每一层proj，改为不映射”。
        # 我们的方案是：只映射一次，然后传递给所有层。
        self.ta_context_fusion = ta_context_fusion
        self.ta_context_norm_flag = ta_context_norm
        
        if self.ta_context_fusion == "add":
            # 如果是相加融合，将 ta_context 投射到和 latent 一样的维度 (hidden_dim)
            self.ta_context_projection = nn.Linear(ta_context_dim, hidden_dim, bias=False)
            self.ta_context_norm = nn.LayerNorm(ta_context_dim) if self.ta_context_norm_flag else nn.Identity()
        elif self.ta_context_fusion == "concat":
            # 如果是拼接融合，在 block 内部处理，这里不需要主投影层
            # 但你的原始代码在concat后也有一个projection，我们可以在 block 内部实现
            # 为了简化，这里先假设主要的融合逻辑在 block 内部
            self.ta_context_projection = nn.Identity()
            self.ta_context_norm = nn.Identity()
        else:
            raise ValueError(f"Unknown ta_context_fusion type: {ta_context_fusion}")


        # --- 原有的输入投影层 (基本不变) ---
        # 现在我的输入要变为editing，需要变为latent*2
        self.audio_input_proj = nn.Sequential(
            ChannelLastConv1d(latent_dim*2, hidden_dim, kernel_size=7, padding=3),
            nn.SELU(),
            ConvMLP(hidden_dim, hidden_dim * 4, kernel_size=7, padding=3),
        )
        self.text_input_proj = nn.Sequential(
            nn.Linear(text_dim, hidden_dim),
            MLP(hidden_dim, hidden_dim * 4),
        )
            
        self.text_cond_proj = nn.Linear(hidden_dim, hidden_dim)
        self.global_cond_mlp = MLP(hidden_dim, hidden_dim * 4)
        

        # 
        self.t_embed = TimestepEmbedder(hidden_dim, frequency_embedding_size=256, max_period=10000)
            
        # --- Transformer Blocks (基本不变) ---
        # **重要**: 你需要修改 JointBlock_AT 和 MMDitSingleBlock 的 forward 定义来接收 `time_aligned_context`
        self.joint_blocks = nn.ModuleList([
            JointBlock_AT(hidden_dim, num_heads, mlp_ratio=mlp_ratio, pre_only=(i == depth - fused_depth - 1))
            for i in range(depth - fused_depth)
        ])
        self.fused_blocks = nn.ModuleList([
            MMDitSingleBlock(hidden_dim, num_heads, mlp_ratio=mlp_ratio, kernel_size=3, padding=1)
            for i in range(fused_depth)
        ])
        
        # --- 输出层 (不变) ---
        self.final_layer = FinalBlock(hidden_dim, latent_dim)

        
        if empty_string_feat is None:
            empty_string_feat = torch.zeros((text_seq_len, text_dim))
            
        self.empty_string_feat = nn.Parameter(empty_string_feat, requires_grad=False)
        
        self.initialize_weights()
        self.initialize_rotations()

    def initialize_rotations(self):
        base_freq = 1.0

        # 唯一需要用到长度的
        latent_rot = compute_rope_rotations(self._latent_seq_len,
                                            self.hidden_dim // self.num_heads,
                                            10000,
                                            freq_scaling=base_freq,
                                            device="cuda" if torch.cuda.is_available() else "cpu")

        # add to model buffers
        self.register_buffer('latent_rot', latent_rot, persistent=False)
        # self.clip_rot = nn.Buffer(clip_rot, persistent=False)

    def update_seq_lengths(self, latent_seq_len: int, clip_seq_len: int, sync_seq_len: int) -> None:
        self._latent_seq_len = latent_seq_len
        self._clip_seq_len = clip_seq_len
        self._sync_seq_len = sync_seq_len
        self.initialize_rotations()

    def initialize_weights(self):

        def _basic_init(module):
            if isinstance(module, nn.Linear):
                torch.nn.init.xavier_uniform_(module.weight)
                if module.bias is not None:
                    nn.init.constant_(module.bias, 0)

        self.apply(_basic_init)

        # Initialize timestep embedding MLP:
        nn.init.normal_(self.t_embed.mlp[0].weight, std=0.02)
        nn.init.normal_(self.t_embed.mlp[2].weight, std=0.02)

        # Zero-out adaLN modulation layers in DiT blocks:兼容性保护
        for block in self.joint_blocks:
            nn.init.constant_(block.latent_block.adaLN_modulation[-1].weight, 0)
            nn.init.constant_(block.latent_block.adaLN_modulation[-1].bias, 0)
            nn.init.constant_(block.text_block.adaLN_modulation[-1].weight, 0)
            nn.init.constant_(block.text_block.adaLN_modulation[-1].bias, 0)
        for block in self.fused_blocks:
            nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
            nn.init.constant_(block.adaLN_modulation[-1].bias, 0)

        # Zero-out output layers:
        nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
        nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)
        nn.init.constant_(self.final_layer.conv.weight, 0)
        nn.init.constant_(self.final_layer.conv.bias, 0)


    
    def preprocess_conditions(self, text_f: torch.Tensor) -> PreprocessedConditions:
        # 预处理文本条件
        # assert text_f.shape[1] == self._text_seq_len, f'{text_f.shape=} {self._text_seq_len=}'
        bs = text_f.shape[0]

        # 这里固定外部的llm_embedding
        text_f = self.text_input_proj(text_f)
        # 全局的条件
        text_f_c = self.text_cond_proj(text_f.mean(dim=1))
        return PreprocessedConditions(text_f=text_f, text_f_c=text_f_c)

    def predict_flow(self, x: torch.Tensor, timesteps: torch.Tensor,
                     conditions: PreprocessedConditions, 
                     time_aligned_context: torch.Tensor) -> torch.Tensor:
        """
        核心的预测流程，现在加入了 time_aligned_context。
        """
        assert x.shape[2] == self._latent_seq_len, f'{x.shape=} {self._latent_seq_len=}'
        
        # 1. 预处理各种输入
        text_f = conditions.text_f
        text_f_c = conditions.text_f_c
        
        timesteps = timesteps.to(x.dtype)  # 保持和输入张量同 dtype

        global_c = self.global_cond_mlp(text_f_c)  # (B, D)
        
        # 2. 融合 timestep
        global_c = self.t_embed(timesteps).unsqueeze(1) + global_c.unsqueeze(1) # (B, 1, D)
        extended_c = global_c # 这个将作为 AdaLN 的条件
        """
        这里决定了x的形状,需要debug
        """
        # 3. **处理 time_aligned_context** 这里第一种方式是直接和latent进行融合，然后投影 
        # 从128->256 
        x = torch.cat([x.transpose(1, 2), time_aligned_context], dim=-1)
        latent = self.audio_input_proj(x)  # (B, N, D)

        # 4. 依次通过 Transformer Blocks
        for block in self.joint_blocks:
            # **你需要修改 JointBlock_AT.forward**
            latent, text_f = block(latent, text_f, global_c, extended_c,
                                           self.latent_rot) 

        for block in self.fused_blocks:
            # **你需要修改 MMDitSingleBlock.forward**
            latent = block(latent, extended_c, self.latent_rot)

        # 5. 通过输出层
        flow = self.final_layer(latent, global_c)
        return flow

    def forward(self, 
                x: torch.Tensor, 
                timesteps: torch.Tensor,
                context: torch.Tensor,
                time_aligned_context: torch.Tensor,
                x_mask=None,
                context_mask=None,
               ) -> torch.Tensor:
        """
        模型主入口，接口已对齐 LayerFusionAudioDiT。
        - x: 噪声 latent, shape (B, N_latent, latent_dim)
        - timesteps: 时间步, shape (B,)
        - context: 文本条件, shape (B, N_text, text_dim)
        - time_aligned_context: 时间对齐的条件, shape (B, N_ta, ta_context_dim)
        """

        if timesteps.dim() == 0:
            timesteps = timesteps.expand(x.shape[0]).to(x.device, dtype=torch.long)
        
        text_conditions = self.preprocess_conditions(context)
        
        # 调用核心预测流
        flow = self.predict_flow(x, timesteps, text_conditions, time_aligned_context)
        

        flow = flow.transpose(1, 2)




        return flow



    @property
    def latent_seq_len(self) -> int:
        return self._latent_seq_len
    

# latent(b,500，128)

def small_16k(**kwargs) -> MMAudio:
    num_heads = 16
    return MMAudio(latent_dim=128,
                   text_dim=1024,
                   hidden_dim=64 * num_heads,
                   depth=12,
                   fused_depth=8,
                   num_heads=num_heads,
                   latent_seq_len=500,
                   **kwargs)




if __name__ == '__main__':

    batch_size = 4
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")


    config = {
        "ta_context_dim": 128,
        "ta_context_fusion": "concat", 
        "ta_context_norm": False
    }


    try:
        model = small_16k(**config).to(device)
        model.eval() # 使用评估模式
        print("Model instantiated successfully!")
    except Exception as e:
        print(f"Error during model instantiation: {e}")
        exit()


    num_params = sum(p.numel() for p in model.parameters()) / 1e6
    print(f'Number of parameters: {num_params:.2f}M')


    latent_dim = 128
    latent_seq_len = 500
    text_dim = 1024
    # 
    text_seq_len = 640
    ta_context_dim = config["ta_context_dim"]

    dummy_x = torch.randn(batch_size,latent_dim, latent_seq_len, device=device)
    dummy_timesteps = torch.randint(0, 1000, (batch_size,), device=device)
    dummy_context = torch.randn(batch_size, text_seq_len, text_dim, device=device)
    
    # 这里的 time_aligned_context 形状需要和 x 一致，以便在特征维度上拼接
    dummy_ta_context = torch.randn(batch_size, latent_seq_len, ta_context_dim, device=device)

    print("\n--- Input Shapes ---")
    print(f"x (latent):           {dummy_x.shape}")
    print(f"timesteps:            {dummy_timesteps.shape}")
    print(f"context (text):       {dummy_context.shape}")
    print(f"time_aligned_context: {dummy_ta_context.shape}")
    print("--------------------\n")
    
    # 4. 执行前向传播
    try:
        with torch.no_grad(): # 在验证时不需要计算梯度
            output = model(
                x=dummy_x,
                timesteps=dummy_timesteps,
                context=dummy_context,
                time_aligned_context=dummy_ta_context
            )
        print("✅ Forward pass successful!")
        print(f"Output shape: {output.shape}")

        # 5. 验证输出形状
        expected_shape = (batch_size, latent_seq_len, latent_dim)
        assert output.shape == expected_shape, \
            f"Output shape mismatch! Expected {expected_shape}, but got {output.shape}"
        print("✅ Output shape is correct!")

    except Exception as e:
        print(f"❌ Error during forward pass: {e}")
        import traceback
        traceback.print_exc()