diffsynth/models/ace_step_tokenizer.py

# Copyright 2025 The ACESTEO Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import Optional

import torch
import torch.nn as nn
from einops import rearrange

from ..core.attention import attention_forward

from transformers.cache_utils import Cache
from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
from transformers.processing_utils import Unpack
from transformers.utils import can_return_tuple, logging
from transformers.models.qwen3.modeling_qwen3 import (
    Qwen3MLP,
    Qwen3RMSNorm,
    Qwen3RotaryEmbedding,
    apply_rotary_pos_emb,
)
from .ace_step_residual_fsq import ResidualFSQ

logger = logging.get_logger(__name__)


def create_4d_mask(
    seq_len: int,
    dtype: torch.dtype,
    device: torch.device,
    attention_mask: Optional[torch.Tensor] = None,
    sliding_window: Optional[int] = None,
    is_sliding_window: bool = False,
    is_causal: bool = True,
) -> torch.Tensor:
    indices = torch.arange(seq_len, device=device)
    diff = indices.unsqueeze(1) - indices.unsqueeze(0)
    valid_mask = torch.ones((seq_len, seq_len), device=device, dtype=torch.bool)
    if is_causal:
        valid_mask = valid_mask & (diff >= 0)
    if is_sliding_window and sliding_window is not None:
        if is_causal:
            valid_mask = valid_mask & (diff <= sliding_window)
        else:
            valid_mask = valid_mask & (torch.abs(diff) <= sliding_window)
    valid_mask = valid_mask.unsqueeze(0).unsqueeze(0)
    if attention_mask is not None:
        padding_mask_4d = attention_mask.view(attention_mask.shape[0], 1, 1, seq_len).to(torch.bool)
        valid_mask = valid_mask & padding_mask_4d
    min_dtype = torch.finfo(dtype).min
    mask_tensor = torch.full(valid_mask.shape, min_dtype, dtype=dtype, device=device)
    mask_tensor.masked_fill_(valid_mask, 0.0)
    return mask_tensor


class Lambda(nn.Module):
    def __init__(self, func):
        super().__init__()
        self.func = func

    def forward(self, x):
        return self.func(x)


class AceStepAttention(nn.Module):
    def __init__(
        self,
        hidden_size: int,
        num_attention_heads: int,
        num_key_value_heads: int,
        rms_norm_eps: float,
        attention_bias: bool,
        attention_dropout: float,
        layer_types: list,
        head_dim: Optional[int] = None,
        sliding_window: Optional[int] = None,
        layer_idx: int = 0,
        is_cross_attention: bool = False,
        is_causal: bool = False,
    ):
        super().__init__()
        self.layer_idx = layer_idx
        self.head_dim = head_dim or hidden_size // num_attention_heads
        self.num_key_value_groups = num_attention_heads // num_key_value_heads
        self.scaling = self.head_dim ** -0.5
        self.attention_dropout = attention_dropout
        if is_cross_attention:
            is_causal = False
        self.is_causal = is_causal
        self.is_cross_attention = is_cross_attention

        self.q_proj = nn.Linear(hidden_size, num_attention_heads * self.head_dim, bias=attention_bias)
        self.k_proj = nn.Linear(hidden_size, num_key_value_heads * self.head_dim, bias=attention_bias)
        self.v_proj = nn.Linear(hidden_size, num_key_value_heads * self.head_dim, bias=attention_bias)
        self.o_proj = nn.Linear(num_attention_heads * self.head_dim, hidden_size, bias=attention_bias)
        self.q_norm = Qwen3RMSNorm(self.head_dim, eps=rms_norm_eps)
        self.k_norm = Qwen3RMSNorm(self.head_dim, eps=rms_norm_eps)
        self.attention_type = layer_types[layer_idx]
        self.sliding_window = sliding_window if layer_types[layer_idx] == "sliding_attention" else None

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor],
        past_key_value: Optional[Cache] = None,
        cache_position: Optional[torch.LongTensor] = None,
        encoder_hidden_states: Optional[torch.Tensor] = None,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] = None,
        output_attentions: Optional[bool] = False,
        **kwargs: Unpack[FlashAttentionKwargs],
    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
        input_shape = hidden_states.shape[:-1]
        hidden_shape = (*input_shape, -1, self.head_dim)

        query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)

        is_cross_attention = self.is_cross_attention and encoder_hidden_states is not None

        if is_cross_attention:
            encoder_hidden_shape = (*encoder_hidden_states.shape[:-1], -1, self.head_dim)
            if past_key_value is not None:
                is_updated = past_key_value.is_updated.get(self.layer_idx)
                curr_past_key_value = past_key_value.cross_attention_cache
                if not is_updated:
                    key_states = self.k_norm(self.k_proj(encoder_hidden_states).view(encoder_hidden_shape)).transpose(1, 2)
                    value_states = self.v_proj(encoder_hidden_states).view(encoder_hidden_shape).transpose(1, 2)
                    key_states, value_states = curr_past_key_value.update(key_states, value_states, self.layer_idx)
                    past_key_value.is_updated[self.layer_idx] = True
                else:
                    key_states = curr_past_key_value.layers[self.layer_idx].keys
                    value_states = curr_past_key_value.layers[self.layer_idx].values
            else:
                key_states = self.k_norm(self.k_proj(encoder_hidden_states).view(encoder_hidden_shape)).transpose(1, 2)
                value_states = self.v_proj(encoder_hidden_states).view(encoder_hidden_shape).transpose(1, 2)

        else:
            key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
            value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
            if position_embeddings is not None:
                cos, sin = position_embeddings
                query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)

            if past_key_value is not None:
                cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
                key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)

        if self.num_key_value_groups > 1:
            key_states = key_states.unsqueeze(2).expand(-1, -1, self.num_key_value_groups, -1, -1).flatten(1, 2)
            value_states = value_states.unsqueeze(2).expand(-1, -1, self.num_key_value_groups, -1, -1).flatten(1, 2)

        attn_output = attention_forward(
            query_states, key_states, value_states,
            q_pattern="b n s d", k_pattern="b n s d", v_pattern="b n s d", out_pattern="b n s d",
            attn_mask=attention_mask,
        )
        attn_weights = None

        attn_output = attn_output.transpose(1, 2).flatten(2, 3).contiguous()
        attn_output = self.o_proj(attn_output)
        return attn_output, attn_weights


class AceStepEncoderLayer(nn.Module):
    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        num_attention_heads: int,
        num_key_value_heads: int,
        rms_norm_eps: float,
        attention_bias: bool,
        attention_dropout: float,
        layer_types: list,
        head_dim: Optional[int] = None,
        sliding_window: Optional[int] = None,
        layer_idx: int = 0,
    ):
        super().__init__()
        self.hidden_size = hidden_size
        self.layer_idx = layer_idx

        self.self_attn = AceStepAttention(
            hidden_size=hidden_size,
            num_attention_heads=num_attention_heads,
            num_key_value_heads=num_key_value_heads,
            rms_norm_eps=rms_norm_eps,
            attention_bias=attention_bias,
            attention_dropout=attention_dropout,
            layer_types=layer_types,
            head_dim=head_dim,
            sliding_window=sliding_window,
            layer_idx=layer_idx,
            is_cross_attention=False,
            is_causal=False,
        )
        self.input_layernorm = Qwen3RMSNorm(hidden_size, eps=rms_norm_eps)
        self.post_attention_layernorm = Qwen3RMSNorm(hidden_size, eps=rms_norm_eps)

        mlp_config = type('Config', (), {
            'hidden_size': hidden_size,
            'intermediate_size': intermediate_size,
            'hidden_act': 'silu',
        })()
        self.mlp = Qwen3MLP(mlp_config)
        self.attention_type = layer_types[layer_idx]

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor],
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        output_attentions: Optional[bool] = False,
        **kwargs,
    ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
        hidden_states, self_attn_weights = self.self_attn(
            hidden_states=hidden_states,
            position_embeddings=position_embeddings,
            attention_mask=attention_mask,
            position_ids=position_ids,
            output_attentions=output_attentions,
            use_cache=False,
            past_key_value=None,
            **kwargs,
        )
        hidden_states = residual + hidden_states

        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states

        outputs = (hidden_states,)
        if output_attentions:
            outputs += (self_attn_weights,)
        return outputs


class AttentionPooler(nn.Module):
    """Pools every pool_window_size frames into 1 representation via transformer + CLS token."""

    def __init__(
        self,
        hidden_size: int = 2048,
        intermediate_size: int = 6144,
        num_attention_heads: int = 16,
        num_key_value_heads: int = 8,
        rms_norm_eps: float = 1e-6,
        attention_bias: bool = False,
        attention_dropout: float = 0.0,
        layer_types: Optional[list] = None,
        head_dim: Optional[int] = None,
        sliding_window: Optional[int] = 128,
        use_sliding_window: bool = True,
        rope_theta: float = 1000000,
        max_position_embeddings: int = 32768,
        initializer_range: float = 0.02,
        num_attention_pooler_hidden_layers: int = 2,
        **kwargs,
    ):
        super().__init__()
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.num_attention_heads = num_attention_heads
        self.num_key_value_heads = num_key_value_heads
        self.rms_norm_eps = rms_norm_eps
        self.attention_bias = attention_bias
        self.attention_dropout = attention_dropout
        # Default matches target library config (24 alternating entries).
        self.layer_types = layer_types or (["sliding_attention", "full_attention"] * 12)
        self.head_dim = head_dim or hidden_size // num_attention_heads
        self.sliding_window = sliding_window
        self.use_sliding_window = use_sliding_window
        self.rope_theta = rope_theta
        self.max_position_embeddings = max_position_embeddings
        self.initializer_range = initializer_range
        self.num_attention_pooler_hidden_layers = num_attention_pooler_hidden_layers
        self._attn_implementation = kwargs.get("_attn_implementation", "sdpa")

        self.embed_tokens = nn.Linear(hidden_size, hidden_size)
        self.norm = Qwen3RMSNorm(hidden_size, eps=rms_norm_eps)
        # Slice layer_types to our own layer count
        pooler_layer_types = self.layer_types[:num_attention_pooler_hidden_layers]
        rope_config = type('RopeConfig', (), {
            'hidden_size': hidden_size,
            'num_attention_heads': num_attention_heads,
            'num_key_value_heads': num_key_value_heads,
            'head_dim': head_dim,
            'max_position_embeddings': max_position_embeddings,
            'rope_theta': rope_theta,
            'rope_parameters': {'rope_type': 'default', 'rope_theta': rope_theta},
            'rms_norm_eps': rms_norm_eps,
            'attention_bias': attention_bias,
            'attention_dropout': attention_dropout,
            'hidden_act': 'silu',
            'intermediate_size': intermediate_size,
            'layer_types': pooler_layer_types,
            'sliding_window': sliding_window,
            '_attn_implementation': self._attn_implementation,
        })()
        self.rotary_emb = Qwen3RotaryEmbedding(rope_config)
        self.gradient_checkpointing = False
        self.special_token = nn.Parameter(torch.randn(1, 1, hidden_size) * 0.02)
        self.layers = nn.ModuleList([
            AceStepEncoderLayer(
                hidden_size=hidden_size,
                intermediate_size=intermediate_size,
                num_attention_heads=num_attention_heads,
                num_key_value_heads=num_key_value_heads,
                rms_norm_eps=rms_norm_eps,
                attention_bias=attention_bias,
                attention_dropout=attention_dropout,
                layer_types=pooler_layer_types,
                head_dim=head_dim,
                sliding_window=sliding_window,
                layer_idx=layer_idx,
            )
            for layer_idx in range(num_attention_pooler_hidden_layers)
        ])

    @can_return_tuple
    def forward(
        self,
        x,
        attention_mask: Optional[torch.Tensor] = None,
        **flash_attn_kwargs: Unpack[FlashAttentionKwargs],
    ) -> torch.Tensor:
        B, T, P, D = x.shape
        x = self.embed_tokens(x)
        special_tokens = self.special_token.expand(B, T, 1, -1).to(x.device)
        x = torch.cat([special_tokens, x], dim=2)
        x = rearrange(x, "b t p c -> (b t) p c")

        cache_position = torch.arange(0, x.shape[1], device=x.device)
        position_ids = cache_position.unsqueeze(0)
        hidden_states = x
        position_embeddings = self.rotary_emb(hidden_states, position_ids)

        seq_len = x.shape[1]
        dtype = x.dtype
        device = x.device

        full_attn_mask = create_4d_mask(
            seq_len=seq_len, dtype=dtype, device=device,
            attention_mask=attention_mask, sliding_window=None,
            is_sliding_window=False, is_causal=False
        )
        sliding_attn_mask = None
        if self.use_sliding_window:
            sliding_attn_mask = create_4d_mask(
                seq_len=seq_len, dtype=dtype, device=device,
                attention_mask=attention_mask, sliding_window=self.sliding_window,
                is_sliding_window=True, is_causal=False
            )

        self_attn_mask_mapping = {
            "full_attention": full_attn_mask,
            "sliding_attention": sliding_attn_mask,
        }

        for layer_module in self.layers:
            layer_outputs = layer_module(
                hidden_states, position_embeddings,
                attention_mask=self_attn_mask_mapping[layer_module.attention_type],
                **flash_attn_kwargs,
            )
            hidden_states = layer_outputs[0]

        hidden_states = self.norm(hidden_states)
        cls_output = hidden_states[:, 0, :]
        return rearrange(cls_output, "(b t) c -> b t c", b=B)


class AceStepAudioTokenizer(nn.Module):
    """Converts continuous acoustic features (VAE latents) into discrete quantized tokens.

    Input: [B, T, 64] (VAE latent dim)
    Output: quantized [B, T/5, 2048], indices [B, T/5, 1]
    """

    def __init__(
        self,
        hidden_size: int = 2048,
        intermediate_size: int = 6144,
        num_attention_heads: int = 16,
        num_key_value_heads: int = 8,
        rms_norm_eps: float = 1e-6,
        attention_bias: bool = False,
        attention_dropout: float = 0.0,
        layer_types: Optional[list] = None,
        head_dim: Optional[int] = None,
        sliding_window: Optional[int] = 128,
        use_sliding_window: bool = True,
        rope_theta: float = 1000000,
        max_position_embeddings: int = 32768,
        initializer_range: float = 0.02,
        audio_acoustic_hidden_dim: int = 64,
        pool_window_size: int = 5,
        fsq_dim: int = 2048,
        fsq_input_levels: list = None,
        fsq_input_num_quantizers: int = 1,
        num_attention_pooler_hidden_layers: int = 2,
        **kwargs,
    ):
        super().__init__()
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.num_attention_heads = num_attention_heads
        self.num_key_value_heads = num_key_value_heads
        self.rms_norm_eps = rms_norm_eps
        self.attention_bias = attention_bias
        self.attention_dropout = attention_dropout
        # Default matches target library config (24 alternating entries).
        self.layer_types = layer_types or (["sliding_attention", "full_attention"] * 12)
        self.head_dim = head_dim or hidden_size // num_attention_heads
        self.sliding_window = sliding_window
        self.use_sliding_window = use_sliding_window
        self.rope_theta = rope_theta
        self.max_position_embeddings = max_position_embeddings
        self.initializer_range = initializer_range
        self.audio_acoustic_hidden_dim = audio_acoustic_hidden_dim
        self.pool_window_size = pool_window_size
        self.fsq_dim = fsq_dim
        self.fsq_input_levels = fsq_input_levels or [8, 8, 8, 5, 5, 5]
        self.fsq_input_num_quantizers = fsq_input_num_quantizers
        self.num_attention_pooler_hidden_layers = num_attention_pooler_hidden_layers
        self._attn_implementation = kwargs.get("_attn_implementation", "sdpa")

        self.audio_acoustic_proj = nn.Linear(audio_acoustic_hidden_dim, hidden_size)
        # Slice layer_types for the attention pooler
        pooler_layer_types = self.layer_types[:num_attention_pooler_hidden_layers]
        self.attention_pooler = AttentionPooler(
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            num_attention_heads=num_attention_heads,
            num_key_value_heads=num_key_value_heads,
            rms_norm_eps=rms_norm_eps,
            attention_bias=attention_bias,
            attention_dropout=attention_dropout,
            layer_types=pooler_layer_types,
            head_dim=head_dim,
            sliding_window=sliding_window,
            use_sliding_window=use_sliding_window,
            rope_theta=rope_theta,
            max_position_embeddings=max_position_embeddings,
            initializer_range=initializer_range,
            num_attention_pooler_hidden_layers=num_attention_pooler_hidden_layers,
        )
        self.quantizer = ResidualFSQ(
            dim=self.fsq_dim,
            levels=self.fsq_input_levels,
            num_quantizers=self.fsq_input_num_quantizers,
            force_quantization_f32=False,
        )

    @can_return_tuple
    def forward(
        self,
        hidden_states: Optional[torch.FloatTensor] = None,
        **flash_attn_kwargs: Unpack[FlashAttentionKwargs],
    ) -> tuple[torch.Tensor, torch.Tensor]:
        hidden_states = self.audio_acoustic_proj(hidden_states)
        hidden_states = self.attention_pooler(hidden_states)
        quantized, indices = self.quantizer(hidden_states)
        return quantized, indices

    def tokenize(self, x):
        """Convenience: takes [B, T, 64], rearranges to patches, runs forward."""
        x = rearrange(x, 'n (t_patch p) d -> n t_patch p d', p=self.pool_window_size)
        return self.forward(x)


class AudioTokenDetokenizer(nn.Module):
    """Converts quantized audio tokens back to continuous acoustic representations.

    Input: [B, T/5, hidden_size] (quantized vectors)
    Output: [B, T, 64] (VAE-latent-shaped continuous features)
    """

    def __init__(
        self,
        hidden_size: int = 2048,
        intermediate_size: int = 6144,
        num_attention_heads: int = 16,
        num_key_value_heads: int = 8,
        rms_norm_eps: float = 1e-6,
        attention_bias: bool = False,
        attention_dropout: float = 0.0,
        layer_types: Optional[list] = None,
        head_dim: Optional[int] = None,
        sliding_window: Optional[int] = 128,
        use_sliding_window: bool = True,
        rope_theta: float = 1000000,
        max_position_embeddings: int = 32768,
        initializer_range: float = 0.02,
        pool_window_size: int = 5,
        audio_acoustic_hidden_dim: int = 64,
        num_attention_pooler_hidden_layers: int = 2,
        **kwargs,
    ):
        super().__init__()
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.num_attention_heads = num_attention_heads
        self.num_key_value_heads = num_key_value_heads
        self.rms_norm_eps = rms_norm_eps
        self.attention_bias = attention_bias
        self.attention_dropout = attention_dropout
        # Default matches target library config (24 alternating entries).
        self.layer_types = layer_types or (["sliding_attention", "full_attention"] * 12)
        self.head_dim = head_dim or hidden_size // num_attention_heads
        self.sliding_window = sliding_window
        self.use_sliding_window = use_sliding_window
        self.rope_theta = rope_theta
        self.max_position_embeddings = max_position_embeddings
        self.initializer_range = initializer_range
        self.pool_window_size = pool_window_size
        self.audio_acoustic_hidden_dim = audio_acoustic_hidden_dim
        self.num_attention_pooler_hidden_layers = num_attention_pooler_hidden_layers
        self._attn_implementation = kwargs.get("_attn_implementation", "sdpa")

        self.embed_tokens = nn.Linear(hidden_size, hidden_size)
        self.norm = Qwen3RMSNorm(hidden_size, eps=rms_norm_eps)
        # Slice layer_types to our own layer count (use num_audio_decoder_hidden_layers)
        detok_layer_types = self.layer_types[:num_attention_pooler_hidden_layers]
        rope_config = type('RopeConfig', (), {
            'hidden_size': hidden_size,
            'num_attention_heads': num_attention_heads,
            'num_key_value_heads': num_key_value_heads,
            'head_dim': head_dim,
            'max_position_embeddings': max_position_embeddings,
            'rope_theta': rope_theta,
            'rope_parameters': {'rope_type': 'default', 'rope_theta': rope_theta},
            'rms_norm_eps': rms_norm_eps,
            'attention_bias': attention_bias,
            'attention_dropout': attention_dropout,
            'hidden_act': 'silu',
            'intermediate_size': intermediate_size,
            'layer_types': detok_layer_types,
            'sliding_window': sliding_window,
            '_attn_implementation': self._attn_implementation,
        })()
        self.rotary_emb = Qwen3RotaryEmbedding(rope_config)
        self.gradient_checkpointing = False
        self.special_tokens = nn.Parameter(torch.randn(1, pool_window_size, hidden_size) * 0.02)
        self.layers = nn.ModuleList([
            AceStepEncoderLayer(
                hidden_size=hidden_size,
                intermediate_size=intermediate_size,
                num_attention_heads=num_attention_heads,
                num_key_value_heads=num_key_value_heads,
                rms_norm_eps=rms_norm_eps,
                attention_bias=attention_bias,
                attention_dropout=attention_dropout,
                layer_types=detok_layer_types,
                head_dim=head_dim,
                sliding_window=sliding_window,
                layer_idx=layer_idx,
            )
            for layer_idx in range(num_attention_pooler_hidden_layers)
        ])
        self.proj_out = nn.Linear(hidden_size, audio_acoustic_hidden_dim)

    @can_return_tuple
    def forward(
        self,
        x,
        attention_mask: Optional[torch.Tensor] = None,
        **flash_attn_kwargs: Unpack[FlashAttentionKwargs],
    ) -> torch.Tensor:
        B, T, D = x.shape
        x = self.embed_tokens(x)
        x = x.unsqueeze(2).repeat(1, 1, self.pool_window_size, 1)
        special_tokens = self.special_tokens.expand(B, T, -1, -1)
        x = x + special_tokens.to(x.device)
        x = rearrange(x, "b t p c -> (b t) p c")

        cache_position = torch.arange(0, x.shape[1], device=x.device)
        position_ids = cache_position.unsqueeze(0)
        hidden_states = x
        position_embeddings = self.rotary_emb(hidden_states, position_ids)

        seq_len = x.shape[1]
        dtype = x.dtype
        device = x.device

        full_attn_mask = create_4d_mask(
            seq_len=seq_len, dtype=dtype, device=device,
            attention_mask=attention_mask, sliding_window=None,
            is_sliding_window=False, is_causal=False
        )
        sliding_attn_mask = None
        if self.use_sliding_window:
            sliding_attn_mask = create_4d_mask(
                seq_len=seq_len, dtype=dtype, device=device,
                attention_mask=attention_mask, sliding_window=self.sliding_window,
                is_sliding_window=True, is_causal=False
            )

        self_attn_mask_mapping = {
            "full_attention": full_attn_mask,
            "sliding_attention": sliding_attn_mask,
        }

        for layer_module in self.layers:
            layer_outputs = layer_module(
                hidden_states, position_embeddings,
                attention_mask=self_attn_mask_mapping[layer_module.attention_type],
                **flash_attn_kwargs,
            )
            hidden_states = layer_outputs[0]

        hidden_states = self.norm(hidden_states)
        hidden_states = self.proj_out(hidden_states)
        return rearrange(hidden_states, "(b t) p c -> b (t p) c", b=B, p=self.pool_window_size)


class AceStepTokenizer(nn.Module):
    """Container for AceStepAudioTokenizer + AudioTokenDetokenizer.

    Provides encode/decode convenience methods for VAE latent discretization.
    Used in cover song mode to convert source audio latents to discrete tokens
    and back to continuous conditioning hints.
    """

    def __init__(
        self,
        hidden_size: int = 2048,
        intermediate_size: int = 6144,
        num_attention_heads: int = 16,
        num_key_value_heads: int = 8,
        rms_norm_eps: float = 1e-6,
        attention_bias: bool = False,
        attention_dropout: float = 0.0,
        layer_types: Optional[list] = None,
        head_dim: Optional[int] = None,
        sliding_window: Optional[int] = 128,
        use_sliding_window: bool = True,
        rope_theta: float = 1000000,
        max_position_embeddings: int = 32768,
        initializer_range: float = 0.02,
        audio_acoustic_hidden_dim: int = 64,
        pool_window_size: int = 5,
        fsq_dim: int = 2048,
        fsq_input_levels: list = None,
        fsq_input_num_quantizers: int = 1,
        num_attention_pooler_hidden_layers: int = 2,
        num_audio_decoder_hidden_layers: int = 24,
        **kwargs,
    ):
        super().__init__()
        if layer_types is None:
            layer_types = ["sliding_attention", "full_attention"] * 12
        self.tokenizer = AceStepAudioTokenizer(
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            num_attention_heads=num_attention_heads,
            num_key_value_heads=num_key_value_heads,
            rms_norm_eps=rms_norm_eps,
            attention_bias=attention_bias,
            attention_dropout=attention_dropout,
            layer_types=layer_types,
            head_dim=head_dim,
            sliding_window=sliding_window,
            use_sliding_window=use_sliding_window,
            rope_theta=rope_theta,
            max_position_embeddings=max_position_embeddings,
            initializer_range=initializer_range,
            audio_acoustic_hidden_dim=audio_acoustic_hidden_dim,
            pool_window_size=pool_window_size,
            fsq_dim=fsq_dim,
            fsq_input_levels=fsq_input_levels,
            fsq_input_num_quantizers=fsq_input_num_quantizers,
            num_attention_pooler_hidden_layers=num_attention_pooler_hidden_layers,
            **kwargs,
        )
        self.detokenizer = AudioTokenDetokenizer(
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            num_attention_heads=num_attention_heads,
            num_key_value_heads=num_key_value_heads,
            rms_norm_eps=rms_norm_eps,
            attention_bias=attention_bias,
            attention_dropout=attention_dropout,
            layer_types=layer_types,
            head_dim=head_dim,
            sliding_window=sliding_window,
            use_sliding_window=use_sliding_window,
            rope_theta=rope_theta,
            max_position_embeddings=max_position_embeddings,
            initializer_range=initializer_range,
            pool_window_size=pool_window_size,
            audio_acoustic_hidden_dim=audio_acoustic_hidden_dim,
            num_attention_pooler_hidden_layers=num_attention_pooler_hidden_layers,
            **kwargs,
        )

    def encode(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """VAE latent [B, T, 64] → discrete tokens."""
        return self.tokenizer(hidden_states)

    def decode(self, quantized: torch.Tensor) -> torch.Tensor:
        """Discrete tokens [B, T/5, hidden_size] → continuous [B, T, 64]."""
        return self.detokenizer(quantized)

    def tokenize(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """Convenience: [B, T, 64] → quantized + indices via patch rearrangement."""
        return self.tokenizer.tokenize(x)