Upload 5 files

meanvc2/dit_kvcache.py

@@ -0,0 +1,332 @@
+"""
+ein notation:
+b - batch
+n - sequence
+nt - text sequence
+nw - raw wave length
+d - dimension
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+from einops import rearrange
+
+from x_transformers.x_transformers import RotaryEmbedding
+
+# from src.model.prompt_vp import MRTE
+from .model_modules import (
+    TimestepEmbedding,
+    ConvNeXtV2Block,
+    ConvPositionEmbedding,
+    AdaLayerNorm_Final,
+    precompute_freqs_cis,
+    get_pos_embed_indices,
+)
+
+from .modules_kvcache import (
+    DiTBlock,
+    ChunkDiTBlock,
+)
+
+class GlobalTimbreMemory(nn.Module):
+    """Global Timbre Memory (GTM)
+    Decomposes the global speaker embedding into K reusable timbre prototype slots.
+    A speaker-specific 2-layer MLP generates speaker-specific KV, fused with a universal prior.
+    """
+    def __init__(self, spk_dim=256, memory_slots=8, hidden_dim=256):
+        super().__init__()
+        self.memory_slots = memory_slots
+        self.hidden_dim = hidden_dim
+
+        # Speaker-specific mapping -- 2-layer MLP for increased expressiveness
+        self.mlp_k = nn.Sequential(
+            nn.Linear(spk_dim, spk_dim),
+            nn.SiLU(),
+            nn.Linear(spk_dim, memory_slots * hidden_dim),
+        )
+        self.mlp_v = nn.Sequential(
+            nn.Linear(spk_dim, spk_dim),
+            nn.SiLU(),
+            nn.Linear(spk_dim, memory_slots * hidden_dim),
+        )
+
+        # Universal speaker-agnostic prototypes (learns common pronunciation patterns), small-variance init
+        self.k_prior = nn.Parameter(torch.zeros(memory_slots, hidden_dim))
+        self.v_prior = nn.Parameter(torch.zeros(memory_slots, hidden_dim))
+        nn.init.normal_(self.k_prior, std=0.02)
+        nn.init.normal_(self.v_prior, std=0.02)
+
+        # LayerNorm after fusion for stable training
+        self.norm_k = nn.LayerNorm(hidden_dim)
+        self.norm_v = nn.LayerNorm(hidden_dim)
+
+    def forward(self, spks):
+        # spks: [B, spk_dim] static global speaker embedding
+        B = spks.shape[0]
+        # Generate speaker-specific key-value pairs
+        k_spk = self.mlp_k(spks).reshape(B, self.memory_slots, self.hidden_dim)
+        v_spk = self.mlp_v(spks).reshape(B, self.memory_slots, self.hidden_dim)
+        # Fuse with universal prototypes + LayerNorm
+        k = self.norm_k(k_spk + torch.tanh(self.k_prior).unsqueeze(0))  # [B, K, D]
+        v = self.norm_v(v_spk + torch.tanh(self.v_prior).unsqueeze(0))  # [B, K, D]
+        return k, v
+    
+
+class TemporalTimbreEncoder(nn.Module):
+    """Temporal Timbre Encoder (TVT processing block)
+    Frame-level content vectors serve as Query in Multi-Head Cross-Attention with GTM.
+    The resulting time-varying timbre features are fused with the global speaker embedding
+    via gated Slerp, preserving the unit hypersphere geometry of the speaker embedding.
+    """
+    def __init__(self, content_dim=256, hidden_dim=256, attn_dim=128, num_heads=4):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = attn_dim // num_heads
+        assert attn_dim % num_heads == 0
+
+        # Multi-Head Cross-Attention projections
+        self.q_proj = nn.Linear(content_dim, attn_dim)
+        self.k_proj = nn.Linear(hidden_dim, attn_dim)   # Input from GTM hidden_dim
+        self.v_proj = nn.Linear(hidden_dim, attn_dim)   # Independent V projection
+        self.out_proj = nn.Linear(attn_dim, content_dim)    # Output projection back to content_dim
+
+        self.attn_scale = self.head_dim ** -0.5
+        self.attn_norm = nn.LayerNorm(content_dim)
+
+
+    def forward(self, bn, k_mem, v_mem):
+        """Args:
+            bn:     [B, T, content_dim]  frame-level content features
+            k_mem:  [B, K, hidden_dim]   GTM keys
+            v_mem:  [B, K, hidden_dim]   GTM values
+        Returns:
+            timbre_cond: [B, T, content_dim] time-varying timbre conditioning
+        """
+        B, T, _ = bn.shape
+        K = k_mem.shape[1]
+
+        # ---- 1. Multi-Head Cross-Attention: content queries × GTM ----
+        q = self.q_proj(bn)     # [B, T, attn_dim]
+        k = self.k_proj(k_mem)  # [B, K, attn_dim]
+        v = self.v_proj(v_mem)  # [B, K, attn_dim]
+
+        # reshape → [B, num_heads, seq_len, head_dim]
+        q = q.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+        k = k.view(B, K, self.num_heads, self.head_dim).transpose(1, 2)
+        v = v.view(B, K, self.num_heads, self.head_dim).transpose(1, 2)
+
+        attn = torch.matmul(q, k.transpose(-2, -1)) * self.attn_scale  # [B, H, T, K]
+        attn = torch.softmax(attn, dim=-1)
+        v_t = torch.matmul(attn, v)  # [B, H, T, head_dim]
+
+        # merge heads -> output projection
+        v_t = v_t.transpose(1, 2).contiguous().view(B, T, -1)  # [B, T, attn_dim]
+        timbre_cond = self.attn_norm(self.out_proj(v_t))                # [B, T, content_dim]
+
+        return timbre_cond
+
+class InputEmbedding(nn.Module):
+    def __init__(self, mel_dim, cond_dim, out_dim):
+        super().__init__()
+        self.proj = nn.Linear(mel_dim + cond_dim * 2, out_dim)
+        # self.conv_pos_embed = ConvPositionEmbedding(dim=out_dim)
+
+    def forward(self, x: float["b n d"], cond: float["b n d"], spks: float["b n d"], drop_audio_cond=False):  # noqa: F722
+    # def forward(self, x: float["b n d"], cond: float["b n d"], timbre_cond: float["b n d"], drop_audio_cond=False):  # noqa: F722
+        if drop_audio_cond:  # cfg for cond audio
+            cond = torch.zeros_like(cond)
+            spks = torch.zeros_like(spks)
+
+        x = self.proj(torch.cat((x, cond, spks), dim=-1))
+        # x = self.conv_pos_embed(x) + x
+        return x
+
+
+
+# Transformer backbone using DiT blocks
+
+
+class DiT(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        depth=8,
+        heads=8,
+        dim_head=64,
+        dropout=0.1,
+        ff_mult=4,
+        mel_dim=80,
+        bn_dim=256,
+        qk_norm=None,
+        conv_layers=0,
+        chunk_size=8,
+        block_size=4,
+        pe_attn_head=None,
+        long_skip_connection=False,
+        checkpoint_activations=False,
+        forward_layers=[0],    # Layer 0 allowed to look ahead
+        backward_layers=[0,1,2,3],   # Layer 3 allowed to look behind
+        t_f_num=[1,0,0,0], 
+        t_p_num=[2,2,1,1],
+    ):
+        super().__init__()
+
+        self.t_time_embed = TimestepEmbedding(dim)
+        self.r_time_embed = TimestepEmbedding(dim)
+        self.input_embed = InputEmbedding(mel_dim, bn_dim, dim)
+        self.rotary_embed = RotaryEmbedding(dim_head)
+
+        self.dim = dim
+        self.depth = depth
+
+        # GTM + TVT time-varying timbre module (replaces MRTE)
+        self.gtm = GlobalTimbreMemory(spk_dim=bn_dim, memory_slots=32, hidden_dim=bn_dim)
+        self.temporal_timbre = TemporalTimbreEncoder(
+            content_dim=bn_dim, hidden_dim=bn_dim,
+            attn_dim=128, num_heads=4,
+        )
+
+        forward_layers = set(forward_layers) if forward_layers else set()
+        backward_layers = set(backward_layers) if backward_layers else set()
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                ChunkDiTBlock(
+                    dim=dim,
+                    heads=heads,
+                    dim_head=dim_head,
+                    ff_mult=ff_mult,
+                    dropout=dropout,
+                    qk_norm=qk_norm,
+                    chunk_size=chunk_size,
+                    block_size=block_size,
+                    pe_attn_head=pe_attn_head,
+                    t_p=t_p_num[i] if i in backward_layers else 0,  # backward
+                    t_f=t_f_num[i] if i in forward_layers else 0,    # forward
+                )
+                for i in range(depth)
+            ]
+        )
+        self.long_skip_connection = nn.Linear(dim * 2, dim, bias=False) if long_skip_connection else None
+
+        self.norm_out = AdaLayerNorm_Final(dim)  # final modulation
+        self.proj_out = nn.Linear(dim, mel_dim)
+
+        self.checkpoint_activations = checkpoint_activations
+
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        # Zero-out AdaLN layers in DiT blocks:
+        for block in self.transformer_blocks:
+            nn.init.constant_(block.attn_norm.linear.weight, 0)
+            nn.init.constant_(block.attn_norm.linear.bias, 0)
+
+        # Zero-out output layers:
+        nn.init.constant_(self.norm_out.linear.weight, 0)
+        nn.init.constant_(self.norm_out.linear.bias, 0)
+        nn.init.constant_(self.proj_out.weight, 0)
+        nn.init.constant_(self.proj_out.bias, 0)
+
+    def ckpt_wrapper(self, module):
+        # https://github.com/chuanyangjin/fast-DiT/blob/main/models.py
+        def ckpt_forward(*inputs):
+            outputs = module(*inputs)
+            return outputs
+
+        return ckpt_forward
+
+
+    def forward(
+        self,
+        x: float["b n d"],  # nosied input audio  # noqa: F722           B, T, 80
+        t: float["b"] | float[""],  # time step  # noqa: F821 F722
+        r: float["b"] | float[""],  # time step  # noqa: F821 F722
+        cache: float["b n d"],
+        cond: float["b n d"],  # bn  # noqa: F722         B, T, 256
+        spks: float["b d"],  # spks  # noqa: F722       B, 256
+        offset=0,
+        mask: bool["b n"] | None = None,  # noqa: F722
+        is_inference: bool = False,
+        is_uncondition: bool = False,
+        cfg_mask: bool["b"] | None = None,  # noqa: F722
+        kv_cache=None,
+    ):
+        
+        batch, seq_len = x.shape[0], x.shape[1]
+
+        # ---- timestep embedding ----
+        t = self.t_time_embed(t)
+        r = self.r_time_embed(r)
+        t = t + r
+
+        # ---- GTM: global speaker embedding -> timbre memory key-value pairs ----
+        k_mem, v_mem = self.gtm(spks)  # spks: [B, 256] -> k,v: [B, K, 256]
+
+        # ---- TVT: frame-level BN x GTM -> timbre-enhanced timbre_cond ----
+        timbre_cond = self.temporal_timbre(cond, k_mem, v_mem)  # [B, T, bn_dim]
+
+        # Expand spks_global to frame level, as pure global identity condition (independent of timbre_cond)
+        spks_expanded = spks.unsqueeze(1).expand(-1, cond.shape[1], -1)  # [B, T, spk_dim]
+
+        # ---- CFG masking ----
+        if cfg_mask is not None:
+            cfg_mask_ = rearrange(cfg_mask, "b -> b 1 1")
+            timbre_cond = torch.where(cfg_mask_, torch.zeros_like(timbre_cond), timbre_cond)
+            spks_expanded = torch.where(cfg_mask_, torch.zeros_like(spks_expanded), spks_expanded)
+
+        # Dual-path input: timbre_cond (content+timbre) + spks_expanded (pure global identity)
+        x = self.input_embed(x, timbre_cond, spks_expanded, drop_audio_cond=is_uncondition)
+        
+        # train
+        if not is_inference:
+
+            rope = self.rotary_embed.forward_from_seq_len(seq_len)
+        # infer
+        else:
+            if cache != None:
+                cache = self.cache_embed(cache)
+                x = torch.concat((cache, x), dim=1)   # [b, cache_len + seq_len, dim]
+                
+                # inference does not need to consider mask
+                cache_len = cache.shape[1]
+                rope_cache = self.rotary_embed.forward_from_seq_len(cache_len)
+                rope_x = self.rotary_embed.forward_from_seq_len(offset + seq_len)
+                rope = (torch.concat((rope_cache[0], rope_x[0][:, -seq_len:, :]), dim=1), rope_cache[1])
+            else:
+                rope = self.rotary_embed.forward_from_seq_len(offset + seq_len)
+                
+
+        if self.long_skip_connection is not None:
+            residual = x
+
+        new_kv_cache = []
+        # inner_hidden_states = []
+        for index_block, block in enumerate(self.transformer_blocks):
+            if kv_cache is not None:
+                block_kv_cache = kv_cache[index_block]
+            else:
+                block_kv_cache = None
+            if self.checkpoint_activations:
+                # https://pytorch.org/docs/stable/checkpoint.html#torch.utils.checkpoint.checkpoint
+                x, new_block_kv_cache = torch.utils.checkpoint.checkpoint(self.ckpt_wrapper(block), x, t, mask, rope, block_kv_cache, use_reentrant=False)
+            else:
+                x, new_block_kv_cache = block(x, t, mask=mask, rope=rope, is_inference=is_inference, kv_cache=block_kv_cache)
+            new_kv_cache.append(new_block_kv_cache)
+        if self.long_skip_connection is not None:
+            x = self.long_skip_connection(torch.cat((x, residual), dim=-1))
+        
+        
+        # x = x[:, -seq_len:, :]
+        x = self.norm_out(x, t)
+        
+        output = self.proj_out(x)
+        
+        return output, new_kv_cache
+
+
+

meanvc2/model_modules.py

@@ -0,0 +1,1040 @@
+"""
+ein notation:
+b - batch
+n - sequence
+nt - text sequence
+nw - raw wave length
+d - dimension
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+import torchaudio
+from librosa.filters import mel as librosa_mel_fn
+from torch import nn
+from x_transformers.x_transformers import apply_rotary_pos_emb
+
+
+# raw wav to mel spec
+
+
+mel_basis_cache = {}
+hann_window_cache = {}
+
+
+def get_bigvgan_mel_spectrogram(
+    waveform,
+    n_fft=1024,
+    n_mel_channels=100,
+    target_sample_rate=24000,
+    hop_length=256,
+    win_length=1024,
+    fmin=0,
+    fmax=None,
+    center=False,
+):  # Copy from https://github.com/NVIDIA/BigVGAN/tree/main
+    device = waveform.device
+    key = f"{n_fft}_{n_mel_channels}_{target_sample_rate}_{hop_length}_{win_length}_{fmin}_{fmax}_{device}"
+
+    if key not in mel_basis_cache:
+        mel = librosa_mel_fn(sr=target_sample_rate, n_fft=n_fft, n_mels=n_mel_channels, fmin=fmin, fmax=fmax)
+        mel_basis_cache[key] = torch.from_numpy(mel).float().to(device)  # TODO: why they need .float()?
+        hann_window_cache[key] = torch.hann_window(win_length).to(device)
+
+    mel_basis = mel_basis_cache[key]
+    hann_window = hann_window_cache[key]
+
+    padding = (n_fft - hop_length) // 2
+    waveform = torch.nn.functional.pad(waveform.unsqueeze(1), (padding, padding), mode="reflect").squeeze(1)
+
+    spec = torch.stft(
+        waveform,
+        n_fft,
+        hop_length=hop_length,
+        win_length=win_length,
+        window=hann_window,
+        center=center,
+        pad_mode="reflect",
+        normalized=False,
+        onesided=True,
+        return_complex=True,
+    )
+    spec = torch.sqrt(torch.view_as_real(spec).pow(2).sum(-1) + 1e-9)
+
+    mel_spec = torch.matmul(mel_basis, spec)
+    mel_spec = torch.log(torch.clamp(mel_spec, min=1e-5))
+
+    return mel_spec
+
+
+def get_vocos_mel_spectrogram(
+    waveform,
+    n_fft=1024,
+    n_mel_channels=100,
+    target_sample_rate=24000,
+    hop_length=256,
+    win_length=1024,
+):
+    mel_stft = torchaudio.transforms.MelSpectrogram(
+        sample_rate=target_sample_rate,
+        n_fft=n_fft,
+        win_length=win_length,
+        hop_length=hop_length,
+        n_mels=n_mel_channels,
+        power=1,
+        center=True,
+        normalized=False,
+        norm=None,
+    ).to(waveform.device)
+    if len(waveform.shape) == 3:
+        waveform = waveform.squeeze(1)  # 'b 1 nw -> b nw'
+
+    assert len(waveform.shape) == 2
+
+    mel = mel_stft(waveform)
+    mel = mel.clamp(min=1e-5).log()
+    return mel
+
+
+class MelSpec(nn.Module):
+    def __init__(
+        self,
+        n_fft=1024,
+        hop_length=256,
+        win_length=1024,
+        n_mel_channels=100,
+        target_sample_rate=24_000,
+        mel_spec_type="vocos",
+    ):
+        super().__init__()
+        assert mel_spec_type in ["vocos", "bigvgan"], print("We only support two extract mel backend: vocos or bigvgan")
+
+        self.n_fft = n_fft
+        self.hop_length = hop_length
+        self.win_length = win_length
+        self.n_mel_channels = n_mel_channels
+        self.target_sample_rate = target_sample_rate
+
+        if mel_spec_type == "vocos":
+            self.extractor = get_vocos_mel_spectrogram
+        elif mel_spec_type == "bigvgan":
+            self.extractor = get_bigvgan_mel_spectrogram
+
+        self.register_buffer("dummy", torch.tensor(0), persistent=False)
+
+    def forward(self, wav):
+        if self.dummy.device != wav.device:
+            self.to(wav.device)
+
+        mel = self.extractor(
+            waveform=wav,
+            n_fft=self.n_fft,
+            n_mel_channels=self.n_mel_channels,
+            target_sample_rate=self.target_sample_rate,
+            hop_length=self.hop_length,
+            win_length=self.win_length,
+        )
+
+        return mel
+
+
+# sinusoidal position embedding
+
+
+class SinusPositionEmbedding(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x, scale=1000):
+        device = x.device
+        half_dim = self.dim // 2
+        emb = math.log(10000) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb)
+        emb = scale * x.unsqueeze(1) * emb.unsqueeze(0)
+        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
+        return emb
+
+
+# convolutional position embedding
+
+
+class ConvPositionEmbedding(nn.Module):
+    def __init__(self, dim, kernel_size=31, groups=16):
+        super().__init__()
+        assert kernel_size % 2 != 0
+        self.conv1d = nn.Sequential(
+            nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=kernel_size // 2),
+            nn.Mish(),
+            nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=kernel_size // 2),
+            nn.Mish(),
+        )
+
+    def forward(self, x: float["b n d"], mask: bool["b n"] | None = None):  # noqa: F722
+        if mask is not None:
+            mask = mask[..., None]
+            x = x.masked_fill(~mask, 0.0)
+
+        x = x.permute(0, 2, 1)
+        x = self.conv1d(x)
+        out = x.permute(0, 2, 1)
+
+        if mask is not None:
+            out = out.masked_fill(~mask, 0.0)
+
+        return out
+
+
+# rotary positional embedding related
+
+
+def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, theta_rescale_factor=1.0):
+    # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning
+    # has some connection to NTK literature
+    # https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/
+    # https://github.com/lucidrains/rotary-embedding-torch/blob/main/rotary_embedding_torch/rotary_embedding_torch.py
+    theta *= theta_rescale_factor ** (dim / (dim - 2))
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
+    t = torch.arange(end, device=freqs.device)  # type: ignore
+    freqs = torch.outer(t, freqs).float()  # type: ignore
+    freqs_cos = torch.cos(freqs)  # real part
+    freqs_sin = torch.sin(freqs)  # imaginary part
+    return torch.cat([freqs_cos, freqs_sin], dim=-1)
+
+
+def get_pos_embed_indices(start, length, max_pos, scale=1.0):
+    # length = length if isinstance(length, int) else length.max()
+    scale = scale * torch.ones_like(start, dtype=torch.float32)  # in case scale is a scalar
+    pos = (
+        start.unsqueeze(1)
+        + (torch.arange(length, device=start.device, dtype=torch.float32).unsqueeze(0) * scale.unsqueeze(1)).long()
+    )
+    # avoid extra long error.
+    pos = torch.where(pos < max_pos, pos, max_pos - 1)
+    return pos
+
+
+# Global Response Normalization layer (Instance Normalization ?)
+
+
+class GRN(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.zeros(1, 1, dim))
+        self.beta = nn.Parameter(torch.zeros(1, 1, dim))
+
+    def forward(self, x):
+        Gx = torch.norm(x, p=2, dim=1, keepdim=True)
+        Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6)
+        return self.gamma * (x * Nx) + self.beta + x
+
+
+# ConvNeXt-V2 Block https://github.com/facebookresearch/ConvNeXt-V2/blob/main/models/convnextv2.py
+# ref: https://github.com/bfs18/e2_tts/blob/main/rfwave/modules.py#L108
+
+
+class ConvNeXtV2Block(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        intermediate_dim: int,
+        dilation: int = 1,
+    ):
+        super().__init__()
+        padding = (dilation * (7 - 1)) // 2
+        self.dwconv = nn.Conv1d(
+            dim, dim, kernel_size=7, padding=padding, groups=dim, dilation=dilation
+        )  # depthwise conv
+        self.norm = nn.LayerNorm(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(dim, intermediate_dim)  # pointwise/1x1 convs, implemented with linear layers
+        self.act = nn.GELU()
+        self.grn = GRN(intermediate_dim)
+        self.pwconv2 = nn.Linear(intermediate_dim, dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x
+        x = x.transpose(1, 2)  # b n d -> b d n
+        x = self.dwconv(x)
+        x = x.transpose(1, 2)  # b d n -> b n d
+        x = self.norm(x)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.grn(x)
+        x = self.pwconv2(x)
+        return residual + x
+
+
+# RMSNorm
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.native_rms_norm = float(torch.__version__[:3]) >= 2.4
+
+    def forward(self, x):
+        if self.native_rms_norm:
+            if self.weight.dtype in [torch.float16, torch.bfloat16]:
+                x = x.to(self.weight.dtype)
+            x = F.rms_norm(x, normalized_shape=(x.shape[-1],), weight=self.weight, eps=self.eps)
+        else:
+            variance = x.to(torch.float32).pow(2).mean(-1, keepdim=True)
+            x = x * torch.rsqrt(variance + self.eps)
+            if self.weight.dtype in [torch.float16, torch.bfloat16]:
+                x = x.to(self.weight.dtype)
+            x = x * self.weight
+
+        return x
+
+
+# AdaLayerNorm
+# return with modulated x for attn input, and params for later mlp modulation
+
+
+class AdaLayerNorm(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+
+        self.silu = nn.SiLU()
+        self.linear = nn.Linear(dim, dim * 6)
+
+        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+
+    def forward(self, x, emb=None):
+        emb = self.linear(self.silu(emb))
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(emb, 6, dim=1)
+
+        x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
+        return x, gate_msa, shift_mlp, scale_mlp, gate_mlp
+
+
+# AdaLayerNorm for final layer
+# return only with modulated x for attn input, cuz no more mlp modulation
+
+
+class AdaLayerNorm_Final(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+
+        self.silu = nn.SiLU()
+        self.linear = nn.Linear(dim, dim * 2)
+
+        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+
+    def forward(self, x, emb):
+        emb = self.linear(self.silu(emb))
+        scale, shift = torch.chunk(emb, 2, dim=1)
+
+        x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
+        return x
+
+
+# FeedForward
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, dim_out=None, mult=4, dropout=0.0, approximate: str = "none"):
+        super().__init__()
+        inner_dim = int(dim * mult)
+        dim_out = dim_out if dim_out is not None else dim
+
+        activation = nn.GELU(approximate=approximate)
+        project_in = nn.Sequential(nn.Linear(dim, inner_dim), activation)
+        self.ff = nn.Sequential(project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out))
+
+    def forward(self, x):
+        return self.ff(x)
+
+
+# Attention with possible joint part
+# modified from diffusers/src/diffusers/models/attention_processor.py
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        processor: JointAttnProcessor | AttnProcessor | ChunkAttnProcessor | BlockAttnProcessor,
+        dim: int,
+        heads: int = 8,
+        dim_head: int = 64,
+        dropout: float = 0.0,
+        context_dim: Optional[int] = None,  # if not None -> joint attention
+        context_pre_only: bool = False,
+        qk_norm: Optional[str] = None,
+    ):
+        super().__init__()
+
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError("Attention equires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+
+        self.processor = processor
+
+        self.dim = dim
+        self.heads = heads
+        self.inner_dim = dim_head * heads
+        self.dropout = dropout
+
+        self.context_dim = context_dim
+        self.context_pre_only = context_pre_only
+
+        self.to_q = nn.Linear(dim, self.inner_dim)
+        self.to_k = nn.Linear(dim, self.inner_dim)
+        self.to_v = nn.Linear(dim, self.inner_dim)
+
+        if qk_norm is None:
+            self.q_norm = None
+            self.k_norm = None
+        elif qk_norm == "rms_norm":
+            self.q_norm = RMSNorm(dim_head, eps=1e-6)
+            self.k_norm = RMSNorm(dim_head, eps=1e-6)
+        else:
+            raise ValueError(f"Unimplemented qk_norm: {qk_norm}")
+
+        if self.context_dim is not None:
+            self.to_q_c = nn.Linear(context_dim, self.inner_dim)
+            self.to_k_c = nn.Linear(context_dim, self.inner_dim)
+            self.to_v_c = nn.Linear(context_dim, self.inner_dim)
+            if qk_norm is None:
+                self.c_q_norm = None
+                self.c_k_norm = None
+            elif qk_norm == "rms_norm":
+                self.c_q_norm = RMSNorm(dim_head, eps=1e-6)
+                self.c_k_norm = RMSNorm(dim_head, eps=1e-6)
+
+        self.to_out = nn.ModuleList([])
+        self.to_out.append(nn.Linear(self.inner_dim, dim))
+        self.to_out.append(nn.Dropout(dropout))
+
+        if self.context_dim is not None and not self.context_pre_only:
+            self.to_out_c = nn.Linear(self.inner_dim, context_dim)
+
+    def forward(
+        self,
+        x: float["b n d"],  # noised input x  # noqa: F722
+        c: float["b n d"] = None,  # context c  # noqa: F722
+        mask: bool["b n"] | None = None,  # noqa: F722
+        rope=None,  # rotary position embedding for x
+        c_rope=None,  # rotary position embedding for c
+        is_inference=False,
+    ) -> torch.Tensor:
+        if c is not None:
+            return self.processor(self, x, c=c, mask=mask, rope=rope, c_rope=c_rope, is_inference=is_inference)
+        else:
+            return self.processor(self, x, mask=mask, rope=rope, is_inference=is_inference)
+
+
+# Attention processor
+
+
+class AttnProcessor:
+    def __init__(
+        self,
+        pe_attn_head: int | None = None,  # number of attention head to apply rope, None for all
+    ):
+        
+        self.pe_attn_head = pe_attn_head
+
+    def __call__(
+        self,
+        attn: Attention,
+        x: float["b n d"],  # noised input x  # noqa: F722
+        mask: bool["b n"] | None = None,  # noqa: F722
+        rope=None,  # rotary position embedding
+    ) -> torch.FloatTensor:
+        batch_size = x.shape[0]
+
+        # `sample` projections
+        query = attn.to_q(x)
+        key = attn.to_k(x)
+        value = attn.to_v(x)
+
+        # attention
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        # qk norm
+        if attn.q_norm is not None:
+            query = attn.q_norm(query)
+        if attn.k_norm is not None:
+            key = attn.k_norm(key)
+
+        # apply rotary position embedding
+        if rope is not None:
+            freqs, xpos_scale = rope
+            q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)
+
+            if self.pe_attn_head is not None:
+                pn = self.pe_attn_head
+                query[:, :pn, :, :] = apply_rotary_pos_emb(query[:, :pn, :, :], freqs, q_xpos_scale)
+                key[:, :pn, :, :] = apply_rotary_pos_emb(key[:, :pn, :, :], freqs, k_xpos_scale)
+            else:
+                query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)
+                key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)
+
+        # mask. e.g. inference got a batch with different target durations, mask out the padding
+        if mask is not None:
+            attn_mask = mask
+            attn_mask = attn_mask.unsqueeze(1).unsqueeze(1)  # 'b n -> b 1 1 n'
+            attn_mask = attn_mask.expand(batch_size, attn.heads, query.shape[-2], key.shape[-2])
+        else:
+            attn_mask = None
+
+        x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False)
+        # x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=True)
+        x = x.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        x = x.to(query.dtype)
+
+        # linear proj
+        x = attn.to_out[0](x)
+        # dropout
+        x = attn.to_out[1](x)
+
+        if mask is not None:
+            mask = mask.unsqueeze(-1)
+            x = x.masked_fill(~mask, 0.0)
+
+        return x
+
+def scaled_dot_product_attention_only(query, key, value, attn_mask=None, dropout_p=0.0,
+        is_causal=False, scale=None, enable_gqa=False) -> torch.Tensor:
+    
+    L, S = query.size(-2), key.size(-2)
+    B = query.size(0)
+    scale_factor = 1 / math.sqrt(query.size(-1)) if scale is None else scale
+    attn_bias = torch.zeros(B, 1, L, S, dtype=query.dtype, device=query.device)
+    if is_causal:
+        assert attn_mask is None
+        temp_mask = torch.ones(B, 1, L, S, dtype=torch.bool).tril(diagonal=0)
+        attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
+        attn_bias.to(query.dtype)
+
+    if attn_mask is not None:
+        if attn_mask.dtype == torch.bool:
+            attn_bias.masked_fill_(attn_mask.logical_not(), float("-inf"))
+        else:
+            attn_bias = attn_mask + attn_bias
+
+    if enable_gqa:
+        key = key.repeat_interleave(query.size(-3)//key.size(-3), -3)
+        value = value.repeat_interleave(query.size(-3)//value.size(-3), -3)
+
+    attn_weight = query @ key.transpose(-2, -1) * scale_factor
+    attn_weight += attn_bias
+    attn_weight = torch.softmax(attn_weight, dim=-1)
+    attn_weight = torch.dropout(attn_weight, dropout_p, train=True)
+    return attn_weight @ value
+
+class BlockAttnProcessor:
+    def __init__(self, chunk_size: int, block_size: int, t_p: int, t_f: int, pe_attn_head: int | None = None,):
+        """
+        Args:
+            chunk_size (int): Number of tokens per chunk
+            block_size (int): Number of tokens per block.
+            t_p (int): Number of past chunks to attend to
+            t_f (int): Number of leading blocks in the future chunk to attend to
+        """
+        self.pe_attn_head = pe_attn_head
+        self.chunk_size = chunk_size
+        self.block_size = block_size
+        self.t_p = t_p
+        self.t_f = t_f
+        
+    def try_cached_mask(self, seq_len, device):
+        idx = torch.arange(seq_len, device=device)
+        ci = idx // self.chunk_size
+
+        qi = ci[:, None]
+        kj = ci[None, :]
+
+        # Within the same chunk
+        same_chunk = (qi == kj)
+
+        # Previous t_p chunks
+        # prev_chunk = (kj == (qi - self.t_p))
+        prev_chunk = (kj >= (qi - self.t_p)) & (kj < qi)
+
+        # First t_f blocks of the next chunk
+        # First compute key(j)'s offset within its own chunk; this offset is bounded by block_size * t_f
+        offset_in_chunk = (idx % self.chunk_size)[None, :]
+        next_chunk_first_block = (kj == (qi + 1)) & (offset_in_chunk < self.block_size * self.t_f)
+
+        computed_mask = same_chunk | prev_chunk | next_chunk_first_block
+
+        return computed_mask
+
+    def __call__(
+        self,
+        attn: Attention,
+        x: float["b n d"],  # noised input x  # noqa: F722
+        mask: bool["b n"] | None = None,  # noqa: F722
+        rope=None,  # rotary position embedding
+        is_inference=False,
+    ) -> torch.FloatTensor:
+
+        batch_size, seq_len, _ = x.shape
+        device = x.device
+
+        # 1. Compute query, key, value projections
+        query = attn.to_q(x)    # Linear layer expands dims [b, n, d * heads]
+        key = attn.to_k(x)
+        value = attn.to_v(x) #torch.Size([batch, seq, 1024])
+        
+        ## 3. Reshape query, key, value into multi-head format: [batch, heads, seq_len, head_dim] attention
+        inner_dim = key.shape[-1]   # d * heads
+        head_dim = inner_dim // attn.heads
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)        # [b, heads, n, d]
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        # qk norm
+        if attn.q_norm is not None:
+            query = attn.q_norm(query)
+        if attn.k_norm is not None:
+            key = attn.k_norm(key)
+
+
+        # apply rotary position embedding
+        # Apply rotary position encoding to q and k
+        if rope is not None:
+            freqs, xpos_scale = rope
+            q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)
+
+            if self.pe_attn_head is not None:
+                pn = self.pe_attn_head
+                query[:, :pn, :, :] = apply_rotary_pos_emb(query[:, :pn, :, :], freqs, q_xpos_scale)
+                key[:, :pn, :, :] = apply_rotary_pos_emb(key[:, :pn, :, :], freqs, k_xpos_scale)
+            else:
+                query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)
+                key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)
+
+        # [b, i, j] = True means token i can attend to token j,
+        # [b, i, j] = False means token i cannot attend to token j,
+        computed_mask = self.try_cached_mask(seq_len,device).unsqueeze(0).expand(batch_size, -1, -1)    # [b, seq_len, seq_len]
+
+        # 4.5 Expand the final mask to multi-head dimensions; shape becomes [batch, heads, seq_len, seq_len] #torch.Size([2, 16, 636, 636])
+        attn_mask = computed_mask.unsqueeze(1).expand(batch_size, 1, seq_len, seq_len)
+
+        # 5. Call PyTorch 2.0 scaled_dot_product_attention
+        attn_output = scaled_dot_product_attention_only(
+            query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False
+        )#attn_mask.to(query.dtype)
+        # attn_output shape: [batch, heads, seq_len, head_dim]
+
+        # 6. Restore shape; concatenate multi-head back to original dims [batch, seq_len, inner_dim]
+        attn_output = attn_output.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        attn_output = attn_output.to(query.dtype)
+
+        # 7. Pass through output projection and dropout
+        attn_output = attn.to_out[0](attn_output)
+        attn_output = attn.to_out[1](attn_output)
+
+        # Expand mask to [batch, seq_len, 1] and zero out the output accordingly
+        if mask is not None:
+            mask = mask.unsqueeze(-1)
+            attn_output = attn_output.masked_fill(~mask, 0.0)
+
+        return attn_output
+
+class ChunkAttnProcessor:
+    def __init__(
+        self,
+        chunk_size: int,
+        pe_attn_head=None,  # number of attention head to apply rope, None for all
+    ):
+        self.chunk_size = chunk_size
+        self.pe_attn_head = pe_attn_head
+
+    def __call__(
+        self,
+        attn: Attention,
+        x: float["b 2*N*chunk_size d"],  # noised input x  # noqa: F722
+        mask: bool["b n"] | None = None,  # noqa: F722
+        rope=None,  # rotary position embedding
+        is_inference=False,
+    ) -> torch.FloatTensor:
+        batch_size, seq_len, _ = x.shape
+
+        # `sample` projections
+        query = attn.to_q(x)
+        key = attn.to_k(x)
+        value = attn.to_v(x)
+
+        # attention
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        # qk norm
+        if attn.q_norm is not None:
+            query = attn.q_norm(query)
+        if attn.k_norm is not None:
+            key = attn.k_norm(key)
+
+        # apply rotary position embedding
+        if rope is not None:
+            freqs, xpos_scale = rope
+            q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)
+
+            if self.pe_attn_head is not None:
+                pn = self.pe_attn_head
+                query[:, :pn, :, :] = apply_rotary_pos_emb(query[:, :pn, :, :], freqs, q_xpos_scale)
+                key[:, :pn, :, :] = apply_rotary_pos_emb(key[:, :pn, :, :], freqs, k_xpos_scale)
+            else:
+                query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)
+                key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)
+
+        indices = torch.arange(seq_len, device=x.device)
+        chunk_indices = indices // self.chunk_size
+        N = int(seq_len / 2 / self.chunk_size)
+    
+        # attn_mask_1 = chunk_indices.unsqueeze(0) <= chunk_indices.unsqueeze(1)
+        # attn_mask_2 = (chunk_indices.unsqueeze(0) + N < chunk_indices.unsqueeze(1)) | (chunk_indices.unsqueeze(0) == chunk_indices.unsqueeze(1))
+        # attn_mask = attn_mask_1 & attn_mask_2
+        
+        # Generate left/right side identifiers (left side = first N*chunk_size frames)
+        is_right_side = indices >= (N * self.chunk_size)
+
+        # Left blocks (M_i) can attend to <= current block's left blocks
+        # left_mask = chunk_indices.unsqueeze(0) <= chunk_indices.unsqueeze(1)
+
+        # Right blocks (M'_i) can only attend to left clean blocks (all M_j, j < i) and itself
+        # right_mask = (
+        #     (chunk_indices.unsqueeze(0) < (chunk_indices.unsqueeze(1) - N)) |  # Access left clean blocks
+        #     (chunk_indices.unsqueeze(0) == chunk_indices.unsqueeze(1))          # Access itself
+        # )
+        
+        max_lookback = 5
+        num_cache_blocks = N           # N
+
+        # 3. Expand dims for broadcasting
+        ci = chunk_indices.unsqueeze(0)   # [L,1], row: chunk the query position belongs to ??? shouldn't it be [1, L]?
+        cj = chunk_indices.unsqueeze(1)   # [1,L], col: chunk the key position belongs to ??? shouldn't it be [L, 1]?
+
+        # 4. Compute relative new block index: for block j, rel_j = cj - N; only rel_j >= 0 is a new block
+        rel_j = cj - num_cache_blocks       # [1,L]
+
+        # 5. Self-attention: token can always attend to itself
+        mask_self = ci == cj                # [L,L]
+
+        mask_cache = (
+            (rel_j >= 0) &
+            (ci < num_cache_blocks) &
+            (ci <  rel_j) &
+            (ci >= rel_j - max_lookback)
+        )
+
+        right_mask = mask_self | mask_cache   # [L,L] boolean matrix
+
+
+        lookback_k = 5  # Look back at most 5 previous blocks + self = 6 blocks total
+        block_diff = cj - ci
+        left_mask = (block_diff >= 0) & (block_diff <= lookback_k)
+
+        # Combine masks
+        if not is_inference:
+            # attn_mask = torch.where(
+            #     is_right_side.unsqueeze(1),  # Apply right_mask to right-side blocks
+            #     right_mask,
+            #     left_mask,                   # Apply left_mask to left-side blocks
+            # )
+            attn_mask = right_mask
+        else:
+            attn_mask = left_mask
+
+
+        if mask is not None:
+            pad_mask = mask.unsqueeze(1) & mask.unsqueeze(2) | torch.eye(seq_len, device=x.device).unsqueeze(0).bool()
+            attn_mask = attn_mask.unsqueeze(0).expand(batch_size, -1, -1) & pad_mask
+        else:
+            attn_mask = attn_mask.unsqueeze(0).expand(batch_size, -1, -1)
+        
+        attn_mask = attn_mask.unsqueeze(1).expand(batch_size, 1, seq_len, seq_len)
+
+        # x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False)
+        x = scaled_dot_product_attention_only(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False)
+        # x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=True)
+        x = x.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        x = x.to(query.dtype)
+
+        # linear proj
+        x = attn.to_out[0](x)
+        # dropout
+        x = attn.to_out[1](x)
+
+        if mask is not None:
+            mask = mask.unsqueeze(-1)
+            x = x.masked_fill(~mask, 0.0)
+
+        return x
+
+# Joint Attention processor for MM-DiT
+# modified from diffusers/src/diffusers/models/attention_processor.py
+
+
+class JointAttnProcessor:
+    def __init__(self):
+        pass
+
+    def __call__(
+        self,
+        attn: Attention,
+        x: float["b n d"],  # noised input x  # noqa: F722
+        c: float["b nt d"] = None,  # context c, here text # noqa: F722
+        mask: bool["b n"] | None = None,  # noqa: F722
+        rope=None,  # rotary position embedding for x
+        c_rope=None,  # rotary position embedding for c
+    ) -> torch.FloatTensor:
+        residual = x
+
+        batch_size = c.shape[0]
+
+        # `sample` projections
+        query = attn.to_q(x)
+        key = attn.to_k(x)
+        value = attn.to_v(x)
+
+        # `context` projections
+        c_query = attn.to_q_c(c)
+        c_key = attn.to_k_c(c)
+        c_value = attn.to_v_c(c)
+
+        # attention
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        c_query = c_query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        c_key = c_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        c_value = c_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        # qk norm
+        if attn.q_norm is not None:
+            query = attn.q_norm(query)
+        if attn.k_norm is not None:
+            key = attn.k_norm(key)
+        if attn.c_q_norm is not None:
+            c_query = attn.c_q_norm(c_query)
+        if attn.c_k_norm is not None:
+            c_key = attn.c_k_norm(c_key)
+
+        # apply rope for context and noised input independently
+        if rope is not None:
+            freqs, xpos_scale = rope
+            q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)
+            query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)
+            key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)
+        if c_rope is not None:
+            freqs, xpos_scale = c_rope
+            q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)
+            c_query = apply_rotary_pos_emb(c_query, freqs, q_xpos_scale)
+            c_key = apply_rotary_pos_emb(c_key, freqs, k_xpos_scale)
+
+        # joint attention
+        query = torch.cat([query, c_query], dim=2)
+        key = torch.cat([key, c_key], dim=2)
+        value = torch.cat([value, c_value], dim=2)
+
+        # mask. e.g. inference got a batch with different target durations, mask out the padding
+        if mask is not None:
+            attn_mask = F.pad(mask, (0, c.shape[1]), value=True)  # no mask for c (text)
+            attn_mask = attn_mask.unsqueeze(1).unsqueeze(1)  # 'b n -> b 1 1 n'
+            attn_mask = attn_mask.expand(batch_size, attn.heads, query.shape[-2], key.shape[-2])
+        else:
+            attn_mask = None
+
+        x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False)
+        x = x.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        x = x.to(query.dtype)
+
+        # Split the attention outputs.
+        x, c = (
+            x[:, : residual.shape[1]],
+            x[:, residual.shape[1] :],
+        )
+
+        # linear proj
+        x = attn.to_out[0](x)
+        # dropout
+        x = attn.to_out[1](x)
+        if not attn.context_pre_only:
+            c = attn.to_out_c(c)
+
+        if mask is not None:
+            mask = mask.unsqueeze(-1)
+            x = x.masked_fill(~mask, 0.0)
+            # c = c.masked_fill(~mask, 0.)  # no mask for c (text)
+
+        return x, c
+
+
+# DiT Block
+
+
+class DiTBlock(nn.Module):
+    def __init__(self, dim, heads, dim_head, ff_mult=4, dropout=0.1, qk_norm=None, pe_attn_head=None):
+        super().__init__()
+
+        self.attn_norm = AdaLayerNorm(dim)
+        self.attn = Attention(
+            processor=AttnProcessor(pe_attn_head=pe_attn_head),
+            dim=dim,
+            heads=heads,
+            dim_head=dim_head,
+            dropout=dropout,
+            qk_norm=qk_norm,
+        )
+
+        self.ff_norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ff = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh")
+
+    def forward(self, x, t, mask=None, rope=None):  # x: noised input, t: time embedding
+        # pre-norm & modulation for attention input
+        norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, emb=t)
+
+        # attention
+        attn_output = self.attn(x=norm, mask=mask, rope=rope)
+
+        # process attention output for input x
+        x = x + gate_msa.unsqueeze(1) * attn_output
+
+        norm = self.ff_norm(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+        ff_output = self.ff(norm)
+        x = x + gate_mlp.unsqueeze(1) * ff_output
+
+        return x
+
+
+class ChunkDiTBlock(nn.Module):
+    def __init__(self, dim, heads, dim_head, ff_mult=4, dropout=0.1, qk_norm=None, chunk_size=16, block_size=8, t_p=0, t_f=0, pe_attn_head=None):
+        super().__init__()
+
+        self.attn_norm = AdaLayerNorm(dim)
+        self.attn = Attention(
+            processor=BlockAttnProcessor(chunk_size=chunk_size, block_size=block_size, t_p=t_p, t_f=t_f, pe_attn_head=pe_attn_head),
+            dim=dim,
+            heads=heads,
+            dim_head=dim_head,
+            dropout=dropout,
+            qk_norm=qk_norm,
+        )
+
+        self.ff_norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ff = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh")
+
+    def forward(self, x, t, mask=None, rope=None, is_inference=False):  # x: noised input, t: time embedding
+        # pre-norm & modulation for attention input
+        norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, emb=t)
+
+        # attention
+        attn_output = self.attn(x=norm, mask=mask, rope=rope, is_inference=is_inference)
+
+        # process attention output for input x
+        x = x + gate_msa.unsqueeze(1) * attn_output
+
+        norm = self.ff_norm(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+        ff_output = self.ff(norm)
+        x = x + gate_mlp.unsqueeze(1) * ff_output
+
+        return x
+
+# MMDiT Block https://arxiv.org/abs/2403.03206
+
+
+class MMDiTBlock(nn.Module):
+    r"""
+    modified from diffusers/src/diffusers/models/attention.py
+
+    notes.
+    _c: context related. text, cond, etc. (left part in sd3 fig2.b)
+    _x: noised input related. (right part)
+    context_pre_only: last layer only do prenorm + modulation cuz no more ffn
+    """
+
+    def __init__(
+        self, dim, heads, dim_head, ff_mult=4, dropout=0.1, context_dim=None, context_pre_only=False, qk_norm=None
+    ):
+        super().__init__()
+        if context_dim is None:
+            context_dim = dim
+        self.context_pre_only = context_pre_only
+
+        self.attn_norm_c = AdaLayerNorm_Final(context_dim) if context_pre_only else AdaLayerNorm(context_dim)
+        self.attn_norm_x = AdaLayerNorm(dim)
+        self.attn = Attention(
+            processor=JointAttnProcessor(),
+            dim=dim,
+            heads=heads,
+            dim_head=dim_head,
+            dropout=dropout,
+            context_dim=context_dim,
+            context_pre_only=context_pre_only,
+            qk_norm=qk_norm,
+        )
+
+        if not context_pre_only:
+            self.ff_norm_c = nn.LayerNorm(context_dim, elementwise_affine=False, eps=1e-6)
+            self.ff_c = FeedForward(dim=context_dim, mult=ff_mult, dropout=dropout, approximate="tanh")
+        else:
+            self.ff_norm_c = None
+            self.ff_c = None
+        self.ff_norm_x = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ff_x = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh")
+
+    def forward(self, x, c, t, mask=None, rope=None, c_rope=None):  # x: noised input, c: context, t: time embedding
+        # pre-norm & modulation for attention input
+        if self.context_pre_only:
+            norm_c = self.attn_norm_c(c, t)
+        else:
+            norm_c, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.attn_norm_c(c, emb=t)
+        norm_x, x_gate_msa, x_shift_mlp, x_scale_mlp, x_gate_mlp = self.attn_norm_x(x, emb=t)
+
+        # attention
+        x_attn_output, c_attn_output = self.attn(x=norm_x, c=norm_c, mask=mask, rope=rope, c_rope=c_rope)
+
+        # process attention output for context c
+        if self.context_pre_only:
+            c = None
+        else:  # if not last layer
+            c = c + c_gate_msa.unsqueeze(1) * c_attn_output
+
+            norm_c = self.ff_norm_c(c) * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
+            c_ff_output = self.ff_c(norm_c)
+            c = c + c_gate_mlp.unsqueeze(1) * c_ff_output
+
+        # process attention output for input x
+        x = x + x_gate_msa.unsqueeze(1) * x_attn_output
+
+        norm_x = self.ff_norm_x(x) * (1 + x_scale_mlp[:, None]) + x_shift_mlp[:, None]
+        x_ff_output = self.ff_x(norm_x)
+        x = x + x_gate_mlp.unsqueeze(1) * x_ff_output
+
+        return c, x
+
+
+# time step conditioning embedding
+
+
+class TimestepEmbedding(nn.Module):
+    def __init__(self, dim, freq_embed_dim=256):
+        super().__init__()
+        self.time_embed = SinusPositionEmbedding(freq_embed_dim)
+        self.time_mlp = nn.Sequential(nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim))
+
+    def forward(self, timestep: float["b"]):  # noqa: F821
+        time_hidden = self.time_embed(timestep)
+        time_hidden = time_hidden.to(timestep.dtype)
+        time = self.time_mlp(time_hidden)  # b d
+        return time

meanvc2/modules_kvcache.py

@@ -0,0 +1,1058 @@
+"""
+ein notation:
+b - batch
+n - sequence
+nt - text sequence
+nw - raw wave length
+d - dimension
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+import torchaudio
+from librosa.filters import mel as librosa_mel_fn
+from torch import nn
+from x_transformers.x_transformers import apply_rotary_pos_emb
+
+
+# raw wav to mel spec
+
+
+mel_basis_cache = {}
+hann_window_cache = {}
+
+
+def get_bigvgan_mel_spectrogram(
+    waveform,
+    n_fft=1024,
+    n_mel_channels=100,
+    target_sample_rate=24000,
+    hop_length=256,
+    win_length=1024,
+    fmin=0,
+    fmax=None,
+    center=False,
+):  # Copy from https://github.com/NVIDIA/BigVGAN/tree/main
+    device = waveform.device
+    key = f"{n_fft}_{n_mel_channels}_{target_sample_rate}_{hop_length}_{win_length}_{fmin}_{fmax}_{device}"
+
+    if key not in mel_basis_cache:
+        mel = librosa_mel_fn(sr=target_sample_rate, n_fft=n_fft, n_mels=n_mel_channels, fmin=fmin, fmax=fmax)
+        mel_basis_cache[key] = torch.from_numpy(mel).float().to(device)  # TODO: why they need .float()?
+        hann_window_cache[key] = torch.hann_window(win_length).to(device)
+
+    mel_basis = mel_basis_cache[key]
+    hann_window = hann_window_cache[key]
+
+    padding = (n_fft - hop_length) // 2
+    waveform = torch.nn.functional.pad(waveform.unsqueeze(1), (padding, padding), mode="reflect").squeeze(1)
+
+    spec = torch.stft(
+        waveform,
+        n_fft,
+        hop_length=hop_length,
+        win_length=win_length,
+        window=hann_window,
+        center=center,
+        pad_mode="reflect",
+        normalized=False,
+        onesided=True,
+        return_complex=True,
+    )
+    spec = torch.sqrt(torch.view_as_real(spec).pow(2).sum(-1) + 1e-9)
+
+    mel_spec = torch.matmul(mel_basis, spec)
+    mel_spec = torch.log(torch.clamp(mel_spec, min=1e-5))
+
+    return mel_spec
+
+
+def get_vocos_mel_spectrogram(
+    waveform,
+    n_fft=1024,
+    n_mel_channels=100,
+    target_sample_rate=24000,
+    hop_length=256,
+    win_length=1024,
+):
+    mel_stft = torchaudio.transforms.MelSpectrogram(
+        sample_rate=target_sample_rate,
+        n_fft=n_fft,
+        win_length=win_length,
+        hop_length=hop_length,
+        n_mels=n_mel_channels,
+        power=1,
+        center=True,
+        normalized=False,
+        norm=None,
+    ).to(waveform.device)
+    if len(waveform.shape) == 3:
+        waveform = waveform.squeeze(1)  # 'b 1 nw -> b nw'
+
+    assert len(waveform.shape) == 2
+
+    mel = mel_stft(waveform)
+    mel = mel.clamp(min=1e-5).log()
+    return mel
+
+
+class MelSpec(nn.Module):
+    def __init__(
+        self,
+        n_fft=1024,
+        hop_length=256,
+        win_length=1024,
+        n_mel_channels=100,
+        target_sample_rate=24_000,
+        mel_spec_type="vocos",
+    ):
+        super().__init__()
+        assert mel_spec_type in ["vocos", "bigvgan"], print("We only support two extract mel backend: vocos or bigvgan")
+
+        self.n_fft = n_fft
+        self.hop_length = hop_length
+        self.win_length = win_length
+        self.n_mel_channels = n_mel_channels
+        self.target_sample_rate = target_sample_rate
+
+        if mel_spec_type == "vocos":
+            self.extractor = get_vocos_mel_spectrogram
+        elif mel_spec_type == "bigvgan":
+            self.extractor = get_bigvgan_mel_spectrogram
+
+        self.register_buffer("dummy", torch.tensor(0), persistent=False)
+
+    def forward(self, wav):
+        if self.dummy.device != wav.device:
+            self.to(wav.device)
+
+        mel = self.extractor(
+            waveform=wav,
+            n_fft=self.n_fft,
+            n_mel_channels=self.n_mel_channels,
+            target_sample_rate=self.target_sample_rate,
+            hop_length=self.hop_length,
+            win_length=self.win_length,
+        )
+
+        return mel
+
+
+# sinusoidal position embedding
+
+
+class SinusPositionEmbedding(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x, scale=1000):
+        device = x.device
+        half_dim = self.dim // 2
+        emb = math.log(10000) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb)
+        emb = scale * x.unsqueeze(1) * emb.unsqueeze(0)
+        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
+        return emb
+
+
+# convolutional position embedding
+
+
+class ConvPositionEmbedding(nn.Module):
+    def __init__(self, dim, kernel_size=31, groups=16):
+        super().__init__()
+        assert kernel_size % 2 != 0
+        self.conv1d = nn.Sequential(
+            nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=kernel_size // 2),
+            nn.Mish(),
+            nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=kernel_size // 2),
+            nn.Mish(),
+        )
+
+    def forward(self, x: float["b n d"], mask: bool["b n"] | None = None):  # noqa: F722
+        if mask is not None:
+            mask = mask[..., None]
+            x = x.masked_fill(~mask, 0.0)
+
+        x = x.permute(0, 2, 1)
+        x = self.conv1d(x)
+        out = x.permute(0, 2, 1)
+
+        if mask is not None:
+            out = out.masked_fill(~mask, 0.0)
+
+        return out
+
+
+# rotary positional embedding related
+
+
+def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, theta_rescale_factor=1.0):
+    # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning
+    # has some connection to NTK literature
+    # https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/
+    # https://github.com/lucidrains/rotary-embedding-torch/blob/main/rotary_embedding_torch/rotary_embedding_torch.py
+    theta *= theta_rescale_factor ** (dim / (dim - 2))
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
+    t = torch.arange(end, device=freqs.device)  # type: ignore
+    freqs = torch.outer(t, freqs).float()  # type: ignore
+    freqs_cos = torch.cos(freqs)  # real part
+    freqs_sin = torch.sin(freqs)  # imaginary part
+    return torch.cat([freqs_cos, freqs_sin], dim=-1)
+
+
+def get_pos_embed_indices(start, length, max_pos, scale=1.0):
+    # length = length if isinstance(length, int) else length.max()
+    scale = scale * torch.ones_like(start, dtype=torch.float32)  # in case scale is a scalar
+    pos = (
+        start.unsqueeze(1)
+        + (torch.arange(length, device=start.device, dtype=torch.float32).unsqueeze(0) * scale.unsqueeze(1)).long()
+    )
+    # avoid extra long error.
+    pos = torch.where(pos < max_pos, pos, max_pos - 1)
+    return pos
+
+
+# Global Response Normalization layer (Instance Normalization ?)
+
+
+class GRN(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.zeros(1, 1, dim))
+        self.beta = nn.Parameter(torch.zeros(1, 1, dim))
+
+    def forward(self, x):
+        Gx = torch.norm(x, p=2, dim=1, keepdim=True)
+        Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6)
+        return self.gamma * (x * Nx) + self.beta + x
+
+
+# ConvNeXt-V2 Block https://github.com/facebookresearch/ConvNeXt-V2/blob/main/models/convnextv2.py
+# ref: https://github.com/bfs18/e2_tts/blob/main/rfwave/modules.py#L108
+
+
+class ConvNeXtV2Block(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        intermediate_dim: int,
+        dilation: int = 1,
+    ):
+        super().__init__()
+        padding = (dilation * (7 - 1)) // 2
+        self.dwconv = nn.Conv1d(
+            dim, dim, kernel_size=7, padding=padding, groups=dim, dilation=dilation
+        )  # depthwise conv
+        self.norm = nn.LayerNorm(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(dim, intermediate_dim)  # pointwise/1x1 convs, implemented with linear layers
+        self.act = nn.GELU()
+        self.grn = GRN(intermediate_dim)
+        self.pwconv2 = nn.Linear(intermediate_dim, dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x
+        x = x.transpose(1, 2)  # b n d -> b d n
+        x = self.dwconv(x)
+        x = x.transpose(1, 2)  # b d n -> b n d
+        x = self.norm(x)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.grn(x)
+        x = self.pwconv2(x)
+        return residual + x
+
+
+# RMSNorm
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.native_rms_norm = float(torch.__version__[:3]) >= 2.4
+
+    def forward(self, x):
+        if self.native_rms_norm:
+            if self.weight.dtype in [torch.float16, torch.bfloat16]:
+                x = x.to(self.weight.dtype)
+            x = F.rms_norm(x, normalized_shape=(x.shape[-1],), weight=self.weight, eps=self.eps)
+        else:
+            variance = x.to(torch.float32).pow(2).mean(-1, keepdim=True)
+            x = x * torch.rsqrt(variance + self.eps)
+            if self.weight.dtype in [torch.float16, torch.bfloat16]:
+                x = x.to(self.weight.dtype)
+            x = x * self.weight
+
+        return x
+
+
+# AdaLayerNorm
+# return with modulated x for attn input, and params for later mlp modulation
+
+
+class AdaLayerNorm(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+
+        self.silu = nn.SiLU()
+        self.linear = nn.Linear(dim, dim * 6)
+
+        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+
+    def forward(self, x, emb=None):
+        emb = self.linear(self.silu(emb))
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(emb, 6, dim=1)
+
+        x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
+        return x, gate_msa, shift_mlp, scale_mlp, gate_mlp
+
+
+# AdaLayerNorm for final layer
+# return only with modulated x for attn input, cuz no more mlp modulation
+
+
+class AdaLayerNorm_Final(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+
+        self.silu = nn.SiLU()
+        self.linear = nn.Linear(dim, dim * 2)
+
+        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+
+    def forward(self, x, emb):
+        emb = self.linear(self.silu(emb))
+        scale, shift = torch.chunk(emb, 2, dim=1)
+
+        x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
+        return x
+
+
+# FeedForward
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, dim_out=None, mult=4, dropout=0.0, approximate: str = "none"):
+        super().__init__()
+        inner_dim = int(dim * mult)
+        dim_out = dim_out if dim_out is not None else dim
+
+        activation = nn.GELU(approximate=approximate)
+        project_in = nn.Sequential(nn.Linear(dim, inner_dim), activation)
+        self.ff = nn.Sequential(project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out))
+
+    def forward(self, x):
+        return self.ff(x)
+
+
+# Attention with possible joint part
+# modified from diffusers/src/diffusers/models/attention_processor.py
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        processor: JointAttnProcessor | AttnProcessor | ChunkAttnProcessor | BlockAttnProcessor,
+        dim: int,
+        heads: int = 8,
+        dim_head: int = 64,
+        dropout: float = 0.0,
+        context_dim: Optional[int] = None,  # if not None -> joint attention
+        context_pre_only: bool = False,
+        qk_norm: Optional[str] = None,
+    ):
+        super().__init__()
+
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError("Attention equires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+
+        self.processor = processor
+
+        self.dim = dim
+        self.heads = heads
+        self.inner_dim = dim_head * heads
+        self.dropout = dropout
+
+        self.context_dim = context_dim
+        self.context_pre_only = context_pre_only
+
+        self.to_q = nn.Linear(dim, self.inner_dim)
+        self.to_k = nn.Linear(dim, self.inner_dim)
+        self.to_v = nn.Linear(dim, self.inner_dim)
+
+        if qk_norm is None:
+            self.q_norm = None
+            self.k_norm = None
+        elif qk_norm == "rms_norm":
+            self.q_norm = RMSNorm(dim_head, eps=1e-6)
+            self.k_norm = RMSNorm(dim_head, eps=1e-6)
+        else:
+            raise ValueError(f"Unimplemented qk_norm: {qk_norm}")
+
+        if self.context_dim is not None:
+            self.to_q_c = nn.Linear(context_dim, self.inner_dim)
+            self.to_k_c = nn.Linear(context_dim, self.inner_dim)
+            self.to_v_c = nn.Linear(context_dim, self.inner_dim)
+            if qk_norm is None:
+                self.c_q_norm = None
+                self.c_k_norm = None
+            elif qk_norm == "rms_norm":
+                self.c_q_norm = RMSNorm(dim_head, eps=1e-6)
+                self.c_k_norm = RMSNorm(dim_head, eps=1e-6)
+
+        self.to_out = nn.ModuleList([])
+        self.to_out.append(nn.Linear(self.inner_dim, dim))
+        self.to_out.append(nn.Dropout(dropout))
+
+        if self.context_dim is not None and not self.context_pre_only:
+            self.to_out_c = nn.Linear(self.inner_dim, context_dim)
+
+    def forward(
+        self,
+        x: float["b n d"],  # noised input x  # noqa: F722
+        c: float["b n d"] = None,  # context c  # noqa: F722
+        mask: bool["b n"] | None = None,  # noqa: F722
+        rope=None,  # rotary position embedding for x
+        c_rope=None,  # rotary position embedding for c
+        is_inference=False,
+        kv_cache=None,
+    ) -> torch.Tensor:
+        if c is not None:
+            return self.processor(self, x, c=c, mask=mask, rope=rope, c_rope=c_rope, is_inference=is_inference, kv_cache=kv_cache)
+        else:
+            return self.processor(self, x, mask=mask, rope=rope, is_inference=is_inference, kv_cache=kv_cache)
+
+
+# Attention processor
+
+
+class AttnProcessor:
+    def __init__(
+        self,
+        pe_attn_head: int | None = None,  # number of attention head to apply rope, None for all
+    ):
+        
+        self.pe_attn_head = pe_attn_head
+
+    def __call__(
+        self,
+        attn: Attention,
+        x: float["b n d"],  # noised input x  # noqa: F722
+        mask: bool["b n"] | None = None,  # noqa: F722
+        rope=None,  # rotary position embedding
+    ) -> torch.FloatTensor:
+        batch_size = x.shape[0]
+
+        # `sample` projections
+        query = attn.to_q(x)
+        key = attn.to_k(x)
+        value = attn.to_v(x)
+
+        # attention
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        # qk norm
+        if attn.q_norm is not None:
+            query = attn.q_norm(query)
+        if attn.k_norm is not None:
+            key = attn.k_norm(key)
+
+        # apply rotary position embedding
+        if rope is not None:
+            freqs, xpos_scale = rope
+            q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)
+
+            if self.pe_attn_head is not None:
+                pn = self.pe_attn_head
+                query[:, :pn, :, :] = apply_rotary_pos_emb(query[:, :pn, :, :], freqs, q_xpos_scale)
+                key[:, :pn, :, :] = apply_rotary_pos_emb(key[:, :pn, :, :], freqs, k_xpos_scale)
+            else:
+                query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)
+                key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)
+
+        # mask. e.g. inference got a batch with different target durations, mask out the padding
+        if mask is not None:
+            attn_mask = mask
+            attn_mask = attn_mask.unsqueeze(1).unsqueeze(1)  # 'b n -> b 1 1 n'
+            attn_mask = attn_mask.expand(batch_size, attn.heads, query.shape[-2], key.shape[-2])
+        else:
+            attn_mask = None
+
+        x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False)
+        # x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=True)
+        x = x.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        x = x.to(query.dtype)
+
+        # linear proj
+        x = attn.to_out[0](x)
+        # dropout
+        x = attn.to_out[1](x)
+
+        if mask is not None:
+            mask = mask.unsqueeze(-1)
+            x = x.masked_fill(~mask, 0.0)
+
+        return x
+
+def scaled_dot_product_attention_only(query, key, value, attn_mask=None, dropout_p=0.0,
+        is_causal=False, scale=None, enable_gqa=False) -> torch.Tensor:
+    
+    L, S = query.size(-2), key.size(-2)
+    B = query.size(0)
+    scale_factor = 1 / math.sqrt(query.size(-1)) if scale is None else scale
+    attn_bias = torch.zeros(B, 1, L, S, dtype=query.dtype, device=query.device)
+    if is_causal:
+        assert attn_mask is None
+        temp_mask = torch.ones(B, 1, L, S, dtype=torch.bool).tril(diagonal=0)
+        attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
+        attn_bias.to(query.dtype)
+
+    if attn_mask is not None:
+        if attn_mask.dtype == torch.bool:
+            attn_bias.masked_fill_(attn_mask[:,:,-attn_bias.shape[2]:,:].logical_not(), float("-inf"))
+        else:
+            attn_bias = attn_mask + attn_bias
+
+    if enable_gqa:
+        key = key.repeat_interleave(query.size(-3)//key.size(-3), -3)
+        value = value.repeat_interleave(query.size(-3)//value.size(-3), -3)
+
+    attn_weight = query @ key.transpose(-2, -1) * scale_factor
+    attn_weight += attn_bias
+    attn_weight = torch.softmax(attn_weight, dim=-1)
+    attn_weight = torch.dropout(attn_weight, dropout_p, train=True)
+    return attn_weight @ value
+
+class BlockAttnProcessor:
+    def __init__(self, chunk_size: int, block_size: int, t_p: int, t_f: int, pe_attn_head: int | None = None,):
+        """
+        Args:
+            chunk_size (int): Number of tokens per chunk
+            block_size (int): Number of tokens per block.
+            t_p (int): Number of past chunks to attend to
+            t_f (int): Number of leading blocks in the future chunk to attend to
+        """
+        self.pe_attn_head = pe_attn_head
+        self.chunk_size = chunk_size
+        self.block_size = block_size
+        self.t_p = t_p
+        self.t_f = t_f
+        
+    def try_cached_mask(self, seq_len, device):
+        idx = torch.arange(seq_len, device=device)
+        ci = idx // self.chunk_size
+
+        qi = ci[:, None]
+        kj = ci[None, :]
+
+        # Within the same chunk
+        same_chunk = (qi == kj)
+
+        # Previous t_p chunks
+        # prev_chunk = (kj == (qi - self.t_p))
+        prev_chunk = (kj >= (qi - self.t_p)) & (kj < qi)
+
+        # First t_f blocks of the next chunk
+        # First compute key(j)'s offset within its own chunk; this offset is bounded by block_size * t_f
+        offset_in_chunk = (idx % self.chunk_size)[None, :]
+        next_chunk_first_block = (kj == (qi + 1)) & (offset_in_chunk < self.block_size * self.t_f)
+
+        computed_mask = same_chunk | prev_chunk | next_chunk_first_block
+
+        return computed_mask
+
+    def __call__(
+        self,
+        attn: Attention,
+        x: float["b n d"],  # noised input x  # noqa: F722
+        mask: bool["b n"] | None = None,  # noqa: F722
+        rope=None,  # rotary position embedding
+        is_inference=False,
+        kv_cache=None,
+    ) -> torch.FloatTensor:
+
+        # batch_size, seq_len, _ = x.shape
+        batch_size = x.shape[0]
+        device = x.device
+
+        # 1. Compute query, key, value projections
+        query = attn.to_q(x)    # Linear layer expands dims [b, n, d * heads]
+        key = attn.to_k(x)
+        value = attn.to_v(x) #torch.Size([batch, seq, 1024])
+        
+        ## 3. Reshape query, key, value into multi-head format: [batch, heads, seq_len, head_dim] attention
+        inner_dim = key.shape[-1]   # d * heads
+        head_dim = inner_dim // attn.heads
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)        # [b, heads, n, d]
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        # qk norm
+        if attn.q_norm is not None:
+            query = attn.q_norm(query)
+        if attn.k_norm is not None:
+            key = attn.k_norm(key)
+
+        # kvcache
+        if kv_cache is None:
+            key_cache = None
+            value_cache = None
+        else:
+            key_cache, value_cache = kv_cache
+        
+        if kv_cache is not None:
+            key = torch.cat([key_cache, key], dim=2)
+        if value_cache is not None:
+            value = torch.cat([value_cache, value], dim=2)
+        
+        new_kv_cache = (key[:, :, :-self.block_size, :], value[:, :, :-self.block_size, :])  # Subsequent blocks are future info; do not cache them
+            
+        batch_size, _, seq_len, _ = key.shape
+
+        # apply rotary position embedding
+        # Apply rotary position encoding to q and k
+        if rope is not None:
+            freqs, xpos_scale = rope
+            q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)
+
+            if self.pe_attn_head is not None:
+                pn = self.pe_attn_head
+                query[:, :pn, :, :] = apply_rotary_pos_emb(query[:, :pn, :, :], freqs, q_xpos_scale)
+                key[:, :pn, :, :] = apply_rotary_pos_emb(key[:, :pn, :, :], freqs, k_xpos_scale)
+            else:
+                query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)
+                key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)
+
+        # [b, i, j] = True means token i can attend to token j,
+        # [b, i, j] = False means token i cannot attend to token j,
+        computed_mask = self.try_cached_mask(seq_len,device).unsqueeze(0).expand(batch_size, -1, -1)    # [b, seq_len, seq_len]
+
+        # 4.5 Expand the final mask to multi-head dimensions; shape becomes [batch, heads, seq_len, seq_len] #torch.Size([2, 16, 636, 636])
+        attn_mask = computed_mask.unsqueeze(1).expand(batch_size, 1, seq_len, seq_len)
+
+        # 5. Call PyTorch 2.0 scaled_dot_product_attention
+        attn_output = scaled_dot_product_attention_only(
+            query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False
+        )#attn_mask.to(query.dtype)
+        # attn_output shape: [batch, heads, seq_len, head_dim]
+
+        # 6. Restore shape; concatenate multi-head back to original dims [batch, seq_len, inner_dim]
+        attn_output = attn_output.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        attn_output = attn_output.to(query.dtype)
+
+        # 7. Pass through output projection and dropout
+        attn_output = attn.to_out[0](attn_output)
+        attn_output = attn.to_out[1](attn_output)
+
+        # Expand mask to [batch, seq_len, 1] and zero out the output accordingly
+        if mask is not None:
+            mask = mask.unsqueeze(-1)
+            attn_output = attn_output.masked_fill(~mask, 0.0)
+
+        return attn_output, new_kv_cache
+
+class ChunkAttnProcessor:
+    def __init__(
+        self,
+        chunk_size: int,
+        pe_attn_head=None,  # number of attention head to apply rope, None for all
+    ):
+        self.chunk_size = chunk_size
+        self.pe_attn_head = pe_attn_head
+
+    def __call__(
+        self,
+        attn: Attention,
+        x: float["b 2*N*chunk_size d"],  # noised input x  # noqa: F722
+        mask: bool["b n"] | None = None,  # noqa: F722
+        rope=None,  # rotary position embedding
+        is_inference=False,
+    ) -> torch.FloatTensor:
+        batch_size, seq_len, _ = x.shape
+
+        # `sample` projections
+        query = attn.to_q(x)
+        key = attn.to_k(x)
+        value = attn.to_v(x)
+
+        # attention
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        # qk norm
+        if attn.q_norm is not None:
+            query = attn.q_norm(query)
+        if attn.k_norm is not None:
+            key = attn.k_norm(key)
+
+        # apply rotary position embedding
+        if rope is not None:
+            freqs, xpos_scale = rope
+            q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)
+
+            if self.pe_attn_head is not None:
+                pn = self.pe_attn_head
+                query[:, :pn, :, :] = apply_rotary_pos_emb(query[:, :pn, :, :], freqs, q_xpos_scale)
+                key[:, :pn, :, :] = apply_rotary_pos_emb(key[:, :pn, :, :], freqs, k_xpos_scale)
+            else:
+                query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)
+                key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)
+
+        indices = torch.arange(seq_len, device=x.device)
+        chunk_indices = indices // self.chunk_size
+        N = int(seq_len / 2 / self.chunk_size)
+    
+        # attn_mask_1 = chunk_indices.unsqueeze(0) <= chunk_indices.unsqueeze(1)
+        # attn_mask_2 = (chunk_indices.unsqueeze(0) + N < chunk_indices.unsqueeze(1)) | (chunk_indices.unsqueeze(0) == chunk_indices.unsqueeze(1))
+        # attn_mask = attn_mask_1 & attn_mask_2
+        
+        # Generate left/right side identifiers (left side = first N*chunk_size frames)
+        is_right_side = indices >= (N * self.chunk_size)
+
+        # Left blocks (M_i) can attend to <= current block's left blocks
+        # left_mask = chunk_indices.unsqueeze(0) <= chunk_indices.unsqueeze(1)
+
+        # Right blocks (M'_i) can only attend to left clean blocks (all M_j, j < i) and itself
+        # right_mask = (
+        #     (chunk_indices.unsqueeze(0) < (chunk_indices.unsqueeze(1) - N)) |  # Access left clean blocks
+        #     (chunk_indices.unsqueeze(0) == chunk_indices.unsqueeze(1))          # Access itself
+        # )
+        
+        max_lookback = 5
+        num_cache_blocks = N           # N
+
+        # 3. Expand dims for broadcasting
+        ci = chunk_indices.unsqueeze(0)   # [L,1], row: chunk the query position belongs to ??? shouldn't it be [1, L]?
+        cj = chunk_indices.unsqueeze(1)   # [1,L], col: chunk the key position belongs to ??? shouldn't it be [L, 1]?
+
+        # 4. Compute relative new block index: for block j, rel_j = cj - N; only rel_j >= 0 is a new block
+        rel_j = cj - num_cache_blocks       # [1,L]
+
+        # 5. Self-attention: token can always attend to itself
+        mask_self = ci == cj                # [L,L]
+
+        mask_cache = (
+            (rel_j >= 0) &
+            (ci < num_cache_blocks) &
+            (ci <  rel_j) &
+            (ci >= rel_j - max_lookback)
+        )
+
+        right_mask = mask_self | mask_cache   # [L,L] boolean matrix
+
+
+        lookback_k = 5  # Look back at most 5 previous blocks + self = 6 blocks total
+        block_diff = cj - ci
+        left_mask = (block_diff >= 0) & (block_diff <= lookback_k)
+
+        # Combine masks
+        if not is_inference:
+            # attn_mask = torch.where(
+            #     is_right_side.unsqueeze(1),  # Apply right_mask to right-side blocks
+            #     right_mask,
+            #     left_mask,                   # Apply left_mask to left-side blocks
+            # )
+            attn_mask = right_mask
+        else:
+            attn_mask = left_mask
+
+
+        if mask is not None:
+            pad_mask = mask.unsqueeze(1) & mask.unsqueeze(2) | torch.eye(seq_len, device=x.device).unsqueeze(0).bool()
+            attn_mask = attn_mask.unsqueeze(0).expand(batch_size, -1, -1) & pad_mask
+        else:
+            attn_mask = attn_mask.unsqueeze(0).expand(batch_size, -1, -1)
+        
+        attn_mask = attn_mask.unsqueeze(1).expand(batch_size, 1, seq_len, seq_len)
+
+        # x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False)
+        x = scaled_dot_product_attention_only(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False)
+        # x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=True)
+        x = x.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        x = x.to(query.dtype)
+
+        # linear proj
+        x = attn.to_out[0](x)
+        # dropout
+        x = attn.to_out[1](x)
+
+        if mask is not None:
+            mask = mask.unsqueeze(-1)
+            x = x.masked_fill(~mask, 0.0)
+
+        return x
+
+# Joint Attention processor for MM-DiT
+# modified from diffusers/src/diffusers/models/attention_processor.py
+
+
+class JointAttnProcessor:
+    def __init__(self):
+        pass
+
+    def __call__(
+        self,
+        attn: Attention,
+        x: float["b n d"],  # noised input x  # noqa: F722
+        c: float["b nt d"] = None,  # context c, here text # noqa: F722
+        mask: bool["b n"] | None = None,  # noqa: F722
+        rope=None,  # rotary position embedding for x
+        c_rope=None,  # rotary position embedding for c
+    ) -> torch.FloatTensor:
+        residual = x
+
+        batch_size = c.shape[0]
+
+        # `sample` projections
+        query = attn.to_q(x)
+        key = attn.to_k(x)
+        value = attn.to_v(x)
+
+        # `context` projections
+        c_query = attn.to_q_c(c)
+        c_key = attn.to_k_c(c)
+        c_value = attn.to_v_c(c)
+
+        # attention
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        c_query = c_query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        c_key = c_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        c_value = c_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        # qk norm
+        if attn.q_norm is not None:
+            query = attn.q_norm(query)
+        if attn.k_norm is not None:
+            key = attn.k_norm(key)
+        if attn.c_q_norm is not None:
+            c_query = attn.c_q_norm(c_query)
+        if attn.c_k_norm is not None:
+            c_key = attn.c_k_norm(c_key)
+
+        # apply rope for context and noised input independently
+        if rope is not None:
+            freqs, xpos_scale = rope
+            q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)
+            query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)
+            key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)
+        if c_rope is not None:
+            freqs, xpos_scale = c_rope
+            q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)
+            c_query = apply_rotary_pos_emb(c_query, freqs, q_xpos_scale)
+            c_key = apply_rotary_pos_emb(c_key, freqs, k_xpos_scale)
+
+        # joint attention
+        query = torch.cat([query, c_query], dim=2)
+        key = torch.cat([key, c_key], dim=2)
+        value = torch.cat([value, c_value], dim=2)
+
+        # mask. e.g. inference got a batch with different target durations, mask out the padding
+        if mask is not None:
+            attn_mask = F.pad(mask, (0, c.shape[1]), value=True)  # no mask for c (text)
+            attn_mask = attn_mask.unsqueeze(1).unsqueeze(1)  # 'b n -> b 1 1 n'
+            attn_mask = attn_mask.expand(batch_size, attn.heads, query.shape[-2], key.shape[-2])
+        else:
+            attn_mask = None
+
+        x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False)
+        x = x.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        x = x.to(query.dtype)
+
+        # Split the attention outputs.
+        x, c = (
+            x[:, : residual.shape[1]],
+            x[:, residual.shape[1] :],
+        )
+
+        # linear proj
+        x = attn.to_out[0](x)
+        # dropout
+        x = attn.to_out[1](x)
+        if not attn.context_pre_only:
+            c = attn.to_out_c(c)
+
+        if mask is not None:
+            mask = mask.unsqueeze(-1)
+            x = x.masked_fill(~mask, 0.0)
+            # c = c.masked_fill(~mask, 0.)  # no mask for c (text)
+
+        return x, c
+
+
+# DiT Block
+
+
+class DiTBlock(nn.Module):
+    def __init__(self, dim, heads, dim_head, ff_mult=4, dropout=0.1, qk_norm=None, pe_attn_head=None):
+        super().__init__()
+
+        self.attn_norm = AdaLayerNorm(dim)
+        self.attn = Attention(
+            processor=AttnProcessor(pe_attn_head=pe_attn_head),
+            dim=dim,
+            heads=heads,
+            dim_head=dim_head,
+            dropout=dropout,
+            qk_norm=qk_norm,
+        )
+
+        self.ff_norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ff = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh")
+
+    def forward(self, x, t, mask=None, rope=None):  # x: noised input, t: time embedding
+        # pre-norm & modulation for attention input
+        norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, emb=t)
+
+        # attention
+        attn_output = self.attn(x=norm, mask=mask, rope=rope)
+
+        # process attention output for input x
+        x = x + gate_msa.unsqueeze(1) * attn_output
+
+        norm = self.ff_norm(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+        ff_output = self.ff(norm)
+        x = x + gate_mlp.unsqueeze(1) * ff_output
+
+        return x
+
+
+class ChunkDiTBlock(nn.Module):
+    def __init__(self, dim, heads, dim_head, ff_mult=4, dropout=0.1, qk_norm=None, chunk_size=16, block_size=8, t_p=0, t_f=0, pe_attn_head=None):
+        super().__init__()
+
+        self.attn_norm = AdaLayerNorm(dim)
+        self.attn = Attention(
+            processor=BlockAttnProcessor(chunk_size=chunk_size, block_size=block_size, t_p=t_p, t_f=t_f, pe_attn_head=pe_attn_head),
+            dim=dim,
+            heads=heads,
+            dim_head=dim_head,
+            dropout=dropout,
+            qk_norm=qk_norm,
+        )
+
+        self.ff_norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ff = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh")
+
+    def forward(self, x, t, mask=None, rope=None, is_inference=False, kv_cache=None):  # x: noised input, t: time embedding
+        # pre-norm & modulation for attention input
+        norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, emb=t)
+
+        # attention
+        attn_output, new_kv_cache = self.attn(x=norm, mask=mask, rope=rope, is_inference=is_inference, kv_cache=kv_cache)
+
+        # process attention output for input x
+        x = x + gate_msa.unsqueeze(1) * attn_output
+
+        norm = self.ff_norm(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+        ff_output = self.ff(norm)
+        x = x + gate_mlp.unsqueeze(1) * ff_output
+
+        return x, new_kv_cache
+
+# MMDiT Block https://arxiv.org/abs/2403.03206
+
+
+class MMDiTBlock(nn.Module):
+    r"""
+    modified from diffusers/src/diffusers/models/attention.py
+
+    notes.
+    _c: context related. text, cond, etc. (left part in sd3 fig2.b)
+    _x: noised input related. (right part)
+    context_pre_only: last layer only do prenorm + modulation cuz no more ffn
+    """
+
+    def __init__(
+        self, dim, heads, dim_head, ff_mult=4, dropout=0.1, context_dim=None, context_pre_only=False, qk_norm=None
+    ):
+        super().__init__()
+        if context_dim is None:
+            context_dim = dim
+        self.context_pre_only = context_pre_only
+
+        self.attn_norm_c = AdaLayerNorm_Final(context_dim) if context_pre_only else AdaLayerNorm(context_dim)
+        self.attn_norm_x = AdaLayerNorm(dim)
+        self.attn = Attention(
+            processor=JointAttnProcessor(),
+            dim=dim,
+            heads=heads,
+            dim_head=dim_head,
+            dropout=dropout,
+            context_dim=context_dim,
+            context_pre_only=context_pre_only,
+            qk_norm=qk_norm,
+        )
+
+        if not context_pre_only:
+            self.ff_norm_c = nn.LayerNorm(context_dim, elementwise_affine=False, eps=1e-6)
+            self.ff_c = FeedForward(dim=context_dim, mult=ff_mult, dropout=dropout, approximate="tanh")
+        else:
+            self.ff_norm_c = None
+            self.ff_c = None
+        self.ff_norm_x = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ff_x = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh")
+
+    def forward(self, x, c, t, mask=None, rope=None, c_rope=None):  # x: noised input, c: context, t: time embedding
+        # pre-norm & modulation for attention input
+        if self.context_pre_only:
+            norm_c = self.attn_norm_c(c, t)
+        else:
+            norm_c, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.attn_norm_c(c, emb=t)
+        norm_x, x_gate_msa, x_shift_mlp, x_scale_mlp, x_gate_mlp = self.attn_norm_x(x, emb=t)
+
+        # attention
+        x_attn_output, c_attn_output = self.attn(x=norm_x, c=norm_c, mask=mask, rope=rope, c_rope=c_rope)
+
+        # process attention output for context c
+        if self.context_pre_only:
+            c = None
+        else:  # if not last layer
+            c = c + c_gate_msa.unsqueeze(1) * c_attn_output
+
+            norm_c = self.ff_norm_c(c) * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
+            c_ff_output = self.ff_c(norm_c)
+            c = c + c_gate_mlp.unsqueeze(1) * c_ff_output
+
+        # process attention output for input x
+        x = x + x_gate_msa.unsqueeze(1) * x_attn_output
+
+        norm_x = self.ff_norm_x(x) * (1 + x_scale_mlp[:, None]) + x_shift_mlp[:, None]
+        x_ff_output = self.ff_x(norm_x)
+        x = x + x_gate_mlp.unsqueeze(1) * x_ff_output
+
+        return c, x
+
+
+# time step conditioning embedding
+
+
+class TimestepEmbedding(nn.Module):
+    def __init__(self, dim, freq_embed_dim=256):
+        super().__init__()
+        self.time_embed = SinusPositionEmbedding(freq_embed_dim)
+        self.time_mlp = nn.Sequential(nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim))
+
+    def forward(self, timestep: float["b"]):  # noqa: F821
+        time_hidden = self.time_embed(timestep)
+        time_hidden = time_hidden.to(timestep.dtype)
+        time = self.time_mlp(time_hidden)  # b d
+        return time

meanvc2/speaker.py

@@ -0,0 +1,385 @@
+"""
+Speaker embedding extraction using WavLM Large + ECAPA-TDNN.
+
+Loads a fine-tuned speaker verification model and extracts 256-dim speaker
+embeddings from reference audio for use in the VC pipeline.
+
+Source: meanvc_run/speaker_verification/ (ecapa_tdnn.py + verification.py)
+
+Dependencies: torch, torchaudio, soundfile, numpy
+Optional: s3prl (for WavLM feature extraction; install from source if needed)
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchaudio.transforms as trans
+from torchaudio.transforms import Resample
+
+
+# ---------------------------------------------------------------------------
+# Lightweight ECAPA-TDNN components (self-contained, no s3prl for fbank mode)
+# ---------------------------------------------------------------------------
+
+class Conv1dReluBn(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size=1, stride=1,
+                 padding=0, dilation=1, bias=True):
+        super().__init__()
+        self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, stride,
+                              padding, dilation, bias=bias)
+        self.bn = nn.BatchNorm1d(out_channels)
+
+    def forward(self, x):
+        return self.bn(F.relu(self.conv(x)))
+
+
+class Res2Conv1dReluBn(nn.Module):
+    def __init__(self, channels, kernel_size=1, stride=1, padding=0,
+                 dilation=1, bias=True, scale=4):
+        super().__init__()
+        assert channels % scale == 0
+        self.scale = scale
+        self.width = channels // scale
+        self.nums = scale if scale == 1 else scale - 1
+        self.convs = nn.ModuleList([
+            nn.Conv1d(self.width, self.width, kernel_size, stride, padding, dilation, bias=bias)
+            for _ in range(self.nums)
+        ])
+        self.bns = nn.ModuleList([
+            nn.BatchNorm1d(self.width) for _ in range(self.nums)
+        ])
+
+    def forward(self, x):
+        out = []
+        spx = torch.split(x, self.width, 1)
+        sp = None
+        for i in range(self.nums):
+            sp = spx[i] if sp is None else sp + spx[i]
+            sp = self.bns[i](F.relu(self.convs[i](sp)))
+            out.append(sp)
+        if self.scale != 1:
+            out.append(spx[self.nums])
+        return torch.cat(out, dim=1)
+
+
+class SE_Connect(nn.Module):
+    def __init__(self, channels, se_bottleneck_dim=128):
+        super().__init__()
+        self.linear1 = nn.Linear(channels, se_bottleneck_dim)
+        self.linear2 = nn.Linear(se_bottleneck_dim, channels)
+
+    def forward(self, x):
+        out = x.mean(dim=2)
+        out = F.relu(self.linear1(out))
+        out = torch.sigmoid(self.linear2(out))
+        return x * out.unsqueeze(2)
+
+
+class SE_Res2Block(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride, padding,
+                 dilation, scale, se_bottleneck_dim):
+        super().__init__()
+        # Submodule names must match the reference models/ecapa_tdnn.py
+        # so that the fine-tuned checkpoint weights load correctly.
+        self.Conv1dReluBn1 = Conv1dReluBn(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+        self.Res2Conv1dReluBn = Res2Conv1dReluBn(out_channels, kernel_size, stride, padding, dilation, scale=scale)
+        self.Conv1dReluBn2 = Conv1dReluBn(out_channels, out_channels, kernel_size=1, stride=1, padding=0)
+        self.SE_Connect = SE_Connect(out_channels, se_bottleneck_dim)
+        self.shortcut = None
+        if in_channels != out_channels:
+            self.shortcut = nn.Conv1d(in_channels, out_channels, kernel_size=1)
+
+    def forward(self, x):
+        residual = self.shortcut(x) if self.shortcut else x
+        x = self.Conv1dReluBn1(x)
+        x = self.Res2Conv1dReluBn(x)
+        x = self.Conv1dReluBn2(x)
+        x = self.SE_Connect(x)
+        return x + residual
+
+
+class AttentiveStatsPool(nn.Module):
+    def __init__(self, in_dim, attention_channels=128, global_context_att=False):
+        super().__init__()
+        self.global_context_att = global_context_att
+        if global_context_att:
+            self.linear1 = nn.Conv1d(in_dim * 3, attention_channels, kernel_size=1)
+        else:
+            self.linear1 = nn.Conv1d(in_dim, attention_channels, kernel_size=1)
+        self.linear2 = nn.Conv1d(attention_channels, in_dim, kernel_size=1)
+
+    def forward(self, x):
+        if self.global_context_att:
+            context_mean = torch.mean(x, dim=-1, keepdim=True).expand_as(x)
+            context_std = torch.sqrt(torch.var(x, dim=-1, keepdim=True) + 1e-10).expand_as(x)
+            x_in = torch.cat((x, context_mean, context_std), dim=1)
+        else:
+            x_in = x
+        alpha = torch.softmax(self.linear2(torch.tanh(self.linear1(x_in))), dim=2)
+        mean = torch.sum(alpha * x, dim=2)
+        residuals = torch.sum(alpha * (x ** 2), dim=2) - mean ** 2
+        std = torch.sqrt(residuals.clamp(min=1e-9))
+        return torch.cat([mean, std], dim=1)
+
+
+# ---------------------------------------------------------------------------
+# ECAPA-TDNN with WavLM feature extractor
+# ---------------------------------------------------------------------------
+
+class ECAPA_TDNN(nn.Module):
+    """ECAPA-TDNN speaker embedding model with WavLM Large backbone."""
+
+    def __init__(self, feat_dim=1024, channels=512, emb_dim=256,
+                 feat_type='wavlm_large', sr=16000,
+                 feature_selection="hidden_states", update_extract=False,
+                 config_path=None):
+        super().__init__()
+
+        self.feat_type = feat_type
+        self.feature_selection = feature_selection
+        self.update_extract = update_extract
+        self.sr = sr
+
+        if feat_type == "fbank" or feat_type == "mfcc":
+            self.update_extract = False
+            win_len = int(sr * 0.025)
+            hop_len = int(sr * 0.01)
+            if feat_type == 'fbank':
+                self.feature_extract = trans.MelSpectrogram(
+                    sample_rate=sr, n_fft=512, win_length=win_len,
+                    hop_length=hop_len, f_min=0.0, f_max=sr // 2,
+                    pad=0, n_mels=feat_dim,
+                )
+            else:
+                melkwargs = {'n_fft': 512, 'win_length': win_len, 'hop_length': hop_len,
+                             'f_min': 0.0, 'f_max': sr // 2, 'pad': 0}
+                self.feature_extract = trans.MFCC(
+                    sample_rate=sr, n_mfcc=feat_dim, log_mels=False, melkwargs=melkwargs,
+                )
+        else:
+            if config_path is not None:
+                # Build UpstreamExpert from tiny config (~10 KB), skipping the
+                # 1.2 GB wavlm_large.pt.  Replicates UpstreamExpert.__init__
+                # (expert.py:34-54) but without torch.load(ckpt) and
+                # load_state_dict — the fine-tuned ckpt provides all weights.
+                from s3prl_wavlm.WavLM import WavLM, WavLMConfig
+                from s3prl_wavlm.expert import UpstreamExpert
+                from s3prl_wavlm.interfaces import UpstreamBase
+
+                if isinstance(config_path, dict):
+                    cfg_dict = config_path
+                else:
+                    cfg_dict = torch.load(config_path, map_location='cpu')
+                cfg = WavLMConfig(cfg_dict)
+                wavlm = WavLM(cfg)
+                wavlm.feature_grad_mult = 0.0
+                wavlm.encoder.layerdrop = 0.0
+
+                expert = UpstreamExpert.__new__(UpstreamExpert)
+                UpstreamBase.__init__(expert)
+                expert.cfg = cfg
+                expert.model = wavlm
+                expert.model.feature_grad_mult = 0.0
+                expert.model.encoder.layerdrop = 0.0
+
+                if len(expert.hooks) == 0:
+                    for module_id in range(len(wavlm.encoder.layers)):
+                        expert.add_hook(
+                            f"self.model.encoder.layers[{module_id}]",
+                            lambda input, output: input[0].transpose(0, 1),
+                        )
+                    expert.add_hook("self.model.encoder",
+                                    lambda input, output: output[0])
+                expert._init_layerdrop = wavlm.encoder.layerdrop
+
+                self.feature_extract = expert
+            else:
+                raise ValueError("config_path (WavLM cfg dict) is required")
+
+            # Disable fp32_attention for layers that have it (compatibility)
+            if len(self.feature_extract.model.encoder.layers) == 24:
+                for layer_idx in [11, 23]:
+                    layer = self.feature_extract.model.encoder.layers[layer_idx]
+                    if hasattr(layer.self_attn, "fp32_attention"):
+                        layer.self_attn.fp32_attention = False
+
+            self.feat_num = self._get_feat_num()
+            self.feature_weight = nn.Parameter(torch.zeros(self.feat_num))
+
+        if feat_type != 'fbank' and feat_type != 'mfcc':
+            freeze_list = ['final_proj', 'label_embs_concat', 'mask_emb', 'project_q', 'quantizer']
+            for name, param in self.feature_extract.named_parameters():
+                for freeze_val in freeze_list:
+                    if freeze_val in name:
+                        param.requires_grad = False
+                        break
+
+        if not self.update_extract:
+            for param in self.feature_extract.parameters():
+                param.requires_grad = False
+
+        self.instance_norm = nn.InstanceNorm1d(feat_dim)
+        self.channels = [channels] * 4 + [1536]
+
+        self.layer1 = Conv1dReluBn(feat_dim, self.channels[0], kernel_size=5, padding=2)
+        self.layer2 = SE_Res2Block(self.channels[0], self.channels[1],
+                                   kernel_size=3, stride=1, padding=2, dilation=2,
+                                   scale=8, se_bottleneck_dim=128)
+        self.layer3 = SE_Res2Block(self.channels[1], self.channels[2],
+                                   kernel_size=3, stride=1, padding=3, dilation=3,
+                                   scale=8, se_bottleneck_dim=128)
+        self.layer4 = SE_Res2Block(self.channels[2], self.channels[3],
+                                   kernel_size=3, stride=1, padding=4, dilation=4,
+                                   scale=8, se_bottleneck_dim=128)
+
+        cat_channels = channels * 3
+        self.conv = nn.Conv1d(cat_channels, self.channels[-1], kernel_size=1)
+        self.pooling = AttentiveStatsPool(self.channels[-1], attention_channels=128,
+                                          global_context_att=False)
+        self.bn = nn.BatchNorm1d(self.channels[-1] * 2)
+        self.linear = nn.Linear(self.channels[-1] * 2, emb_dim)
+
+    def _get_feat_num(self):
+        self.feature_extract.eval()
+        wav = [torch.randn(self.sr).to(next(self.feature_extract.parameters()).device)]
+        with torch.no_grad():
+            features = self.feature_extract(wav)
+        select_feature = features[self.feature_selection]
+        if isinstance(select_feature, (list, tuple)):
+            return len(select_feature)
+        return 1
+
+    def _get_feat(self, x):
+        if self.update_extract:
+            x = self.feature_extract([sample for sample in x])
+        else:
+            with torch.no_grad():
+                if self.feat_type == 'fbank' or self.feat_type == 'mfcc':
+                    x = self.feature_extract(x) + 1e-6
+                else:
+                    x = self.feature_extract([sample for sample in x])
+
+        if self.feat_type == 'fbank':
+            x = x.log()
+
+        if self.feat_type != "fbank" and self.feat_type != "mfcc":
+            x = x[self.feature_selection]
+            if isinstance(x, (list, tuple)):
+                x = torch.stack(x, dim=0)
+            else:
+                x = x.unsqueeze(0)
+            norm_weights = F.softmax(self.feature_weight, dim=-1).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
+            x = (norm_weights * x).sum(dim=0)
+            x = torch.transpose(x, 1, 2) + 1e-6
+
+        x = self.instance_norm(x)
+        return x
+
+    def forward(self, x):
+        x = self._get_feat(x)
+        out1 = self.layer1(x)
+        out2 = self.layer2(out1)
+        out3 = self.layer3(out2)
+        out4 = self.layer4(out3)
+        out = torch.cat([out2, out3, out4], dim=1)
+        out = F.relu(self.conv(out))
+        out = self.bn(self.pooling(out))
+        out = self.linear(out)
+        return out
+
+
+def ECAPA_TDNN_SMALL(feat_dim, emb_dim=256, feat_type='fbank', sr=16000,
+                     feature_selection="hidden_states", update_extract=False,
+                     config_path=None):
+    return ECAPA_TDNN(
+        feat_dim=feat_dim, channels=512, emb_dim=emb_dim,
+        feat_type=feat_type, sr=sr,
+        feature_selection=feature_selection,
+        update_extract=update_extract,
+        config_path=config_path,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+
+def init_speaker_model(ckpt_path=None, device='cpu', wavlm_config=None):
+    """
+    Load the WavLM Large + ECAPA-TDNN speaker verification model.
+
+    Args:
+        ckpt_path: path to fine-tuned checkpoint (wavlm_large_finetune.pth).
+            Contains ALL backbone + ECAPA-TDNN weights.
+        device: 'cpu' or 'cuda'
+        wavlm_config: path to WavLM config (wavlm_large_cfg.pt, ~10 KB).
+            Extracted from wavlm_large.pt via extract_wavlm_config.py.
+            When provided, skips loading the 1.2 GB base checkpoint — the
+            WavLM backbone is built from config and all weights come from
+            ckpt_path via load_state_dict().
+    Returns:
+        model: ECAPA_TDNN model in eval mode, on the specified device
+    """
+    model = ECAPA_TDNN_SMALL(
+        feat_dim=1024, emb_dim=256,
+        feat_type='wavlm_large',
+        feature_selection="hidden_states",
+        update_extract=False,
+        config_path=wavlm_config,
+    )
+    if ckpt_path is not None:
+        state_dict = torch.load(ckpt_path, map_location='cpu', weights_only=True)
+        if 'model' in state_dict:
+            state_dict = state_dict['model']
+        missing, unexpected = model.load_state_dict(state_dict, strict=False)
+        missing, unexpected = list(missing), list(unexpected)
+        print(f"[speaker] loaded ckpt: {len(state_dict)} tensors | "
+              f"missing={len(missing)} unexpected={len(unexpected)}", flush=True)
+        if missing:
+            print(f"[speaker] missing (first 10): {missing[:10]}", flush=True)
+        if unexpected:
+            print(f"[speaker] unexpected (first 10): {unexpected[:10]}", flush=True)
+    model.eval()
+    model.to(device)
+    return model
+
+
+def extract_embedding(model, wav, sample_rate=16000, device='cpu'):
+    """
+    Extract 256-dim speaker embedding from audio.
+
+    Args:
+        model: ECAPA_TDNN model from init_speaker_model()
+        wav: either a file path (str) or a torch.Tensor [1, samples] or numpy array
+        sample_rate: target sample rate (used if loading from file)
+        device: compute device
+
+    Returns:
+        torch.Tensor [1, 256] — L2-normalized speaker embedding
+    """
+    import soundfile as sf
+
+    if isinstance(wav, str):
+        data, sr = sf.read(wav)
+        if data.ndim == 2:
+            data = np.mean(data, axis=1)
+        wav = torch.from_numpy(data).unsqueeze(0).float().to(device)
+        if sr != sample_rate:
+            resample = Resample(orig_freq=sr, new_freq=sample_rate).to(device)
+            wav = resample(wav)
+    elif isinstance(wav, np.ndarray):
+        wav = torch.from_numpy(wav).unsqueeze(0).float().to(device)
+    elif isinstance(wav, torch.Tensor):
+        wav = wav.to(device)
+        if wav.ndim == 1:
+            wav = wav.unsqueeze(0)
+
+    with torch.no_grad():
+        emb = model(wav)
+
+    # NOTE: No L2 normalization — aligned with extract_spk_emb_wavlm_multi_mp3.py
+    # which outputs raw (unnormalized) embeddings.
+    return emb