fount-vlm-nell-02 0.3.11__py3-none-any.whl

mlx_vlm/models/qwen3_omni_moe/code2wav.py

@@ -0,0 +1,542 @@
+import math
+from typing import Optional, Tuple
+
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+from mlx_lm.models.base import scaled_dot_product_attention
+
+from mlx_vlm.models.qwen3_omni_moe.config import Code2WavConfig
+
+
+class SnakeBeta(nn.Module):
+    def __init__(self, in_features, alpha=1.0):
+        super().__init__()
+        self.in_features = in_features
+        self.alpha = mx.zeros((in_features,)) * alpha
+        self.beta = mx.zeros((in_features,)) * alpha
+        self.no_div_by_zero = 0.000000001
+
+    def __call__(self, hidden_states):
+        alpha = mx.expand_dims(mx.expand_dims(self.alpha, axis=0), axis=-1)
+        beta = mx.expand_dims(mx.expand_dims(self.beta, axis=0), axis=-1)
+        alpha = mx.exp(alpha)
+        beta = mx.exp(beta)
+        hidden_states = hidden_states + (1.0 / (beta + self.no_div_by_zero)) * mx.power(
+            mx.sin(hidden_states * alpha), 2
+        )
+        return hidden_states
+
+
+class LayerScale(nn.Module):
+    def __init__(self, config: Code2WavConfig):
+        super().__init__()
+        channels = config.hidden_size
+        initial_scale = config.layer_scale_initial_scale
+        self.scale = mx.full((channels,), initial_scale)
+
+    def __call__(self, x: mx.array):
+        return self.scale * x
+
+
+class RoPE(nn.Module):
+    def __init__(self, config: Code2WavConfig):
+        super().__init__()
+        self.config = config
+        head_dim = config.hidden_size // config.num_attention_heads
+        dim = head_dim
+        inv_freq = 1.0 / (
+            config.rope_theta ** (np.arange(0, dim, 2, dtype=np.float32) / dim)
+        )
+        self.inv_freq = inv_freq
+        self.attention_scaling = 1.0
+
+    def __call__(self, x: mx.array, position_ids: mx.array):
+        batch_size = position_ids.shape[0]
+        inv_freq_mx = mx.array(self.inv_freq)
+        inv_freq_expanded = mx.broadcast_to(
+            inv_freq_mx[None, :, None].astype(mx.float32),
+            (batch_size, inv_freq_mx.shape[0], 1),
+        )
+        position_ids_expanded = mx.expand_dims(position_ids.astype(mx.float32), axis=1)
+        freqs = inv_freq_expanded @ position_ids_expanded
+        freqs = mx.swapaxes(freqs, 1, 2)
+        emb = mx.concatenate([freqs, freqs], axis=-1)
+        cos = mx.cos(emb) * self.attention_scaling
+        sin = mx.sin(emb) * self.attention_scaling
+        return cos.astype(x.dtype), sin.astype(x.dtype)
+
+
+class CausalConvNet(nn.Module):
+    def __init__(self, in_chn, out_chn, kernel_sz, dilation=1, stride=1, groups=1):
+        super().__init__()
+        self.conv = nn.Conv1d(
+            in_chn, out_chn, kernel_sz, stride=stride, dilation=dilation, groups=groups
+        )
+        self.stride = stride
+        self.kernel_size = (kernel_sz - 1) * dilation + 1
+        self.dilation = dilation
+        self.padding = self.kernel_size - self.stride
+
+    def _get_extra_padding_for_conv1d(self, length: int) -> int:
+        n_frames = (length - self.kernel_size + self.padding) / self.stride + 1
+        ideal_length = (math.ceil(n_frames) - 1) * self.stride + (
+            self.kernel_size - self.padding
+        )
+        return int(ideal_length - length)
+
+    def __call__(self, hidden_state: mx.array) -> mx.array:
+        length = hidden_state.shape[-1]
+        extra_padding = self._get_extra_padding_for_conv1d(length)
+        hidden_state = hidden_state.transpose(0, 2, 1)
+        pad_width = [(0, 0), (self.padding, extra_padding), (0, 0)]
+        hidden_state = mx.pad(
+            hidden_state, pad_width, mode="constant", constant_values=0
+        )
+        output = self.conv(hidden_state)
+        return output.transpose(0, 2, 1)
+
+
+class CausalTransConvNet(nn.Module):
+    def __init__(self, in_chn, out_chn, kernel_sz, stride=1):
+        super().__init__()
+        self.conv = nn.ConvTranspose1d(in_chn, out_chn, kernel_sz, stride=stride)
+        pad = kernel_sz - stride
+        self.left_pad = 0
+        self.right_pad = pad
+
+    def __call__(self, hidden_state: mx.array) -> mx.array:
+        hidden_state = hidden_state.transpose(0, 2, 1)
+        hidden_state = self.conv(hidden_state)
+        length = hidden_state.shape[-2]
+        hidden_state = hidden_state[:, self.left_pad : length - self.right_pad, :]
+        return hidden_state.transpose(0, 2, 1)
+
+
+class ConvNeXtBlock(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+
+        self.dwconv = CausalConvNet(dim, dim, kernel_sz=7, groups=dim, dilation=1)
+        self.norm = nn.LayerNorm(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(dim, 4 * dim)
+        self.pwconv2 = nn.Linear(4 * dim, dim)
+        self.gamma = mx.full((dim,), 1e-6)
+
+    def __call__(self, hidden_states: mx.array) -> mx.array:
+        input = hidden_states
+        hidden_states = self.dwconv(hidden_states)
+        hidden_states = hidden_states.transpose(0, 2, 1)
+        hidden_states = self.norm(hidden_states)
+        hidden_states = self.pwconv1(hidden_states)
+        hidden_states = nn.gelu(hidden_states)
+        hidden_states = self.pwconv2(hidden_states)
+        hidden_states = self.gamma * hidden_states
+        hidden_states = hidden_states.transpose(0, 2, 1)
+        hidden_states = input + hidden_states
+        return hidden_states
+
+
+class Code2WavDecoderResUnit(nn.Module):
+    def __init__(self, dim: int, dilation: int = 1):
+        super().__init__()
+
+        self.act1 = SnakeBeta(dim)
+        self.conv1 = CausalConvNet(dim, dim, kernel_sz=7, dilation=dilation)
+        self.act2 = SnakeBeta(dim)
+        self.conv2 = CausalConvNet(dim, dim, kernel_sz=1)
+
+    def __call__(self, hidden_state: mx.array) -> mx.array:
+        residual = hidden_state
+        hidden_state = self.act1(hidden_state)
+        hidden_state = self.conv1(hidden_state)
+        hidden_state = self.act2(hidden_state)
+        hidden_state = self.conv2(hidden_state)
+        return hidden_state + residual
+
+
+class Code2WavDecoderBlock(nn.Module):
+    def __init__(self, config: Code2WavConfig, idx: int):
+        super().__init__()
+
+        in_dim = config.decoder_dim // 2**idx
+        out_dim = config.decoder_dim // 2 ** (idx + 1)
+        upsample_rate = config.upsample_rates[idx]
+
+        self.block = [
+            SnakeBeta(in_dim),
+            CausalTransConvNet(in_dim, out_dim, 2 * upsample_rate, upsample_rate),
+        ]
+        self.block.extend(
+            [Code2WavDecoderResUnit(out_dim, dilation) for dilation in (1, 3, 9)]
+        )
+
+    def __call__(self, hidden: mx.array) -> mx.array:
+        for block in self.block:
+            hidden = block(hidden)
+        return hidden
+
+
+def rotate_half(x):
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return mx.concatenate([-x2, x1], axis=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin):
+    cos = mx.expand_dims(cos, axis=1)
+    sin = mx.expand_dims(sin, axis=1)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+class Code2WavAttention(nn.Module):
+    def __init__(self, config: Code2WavConfig, idx: int):
+        super().__init__()
+
+        self.config = config
+        self.layer_idx = idx
+        self.head_dim = getattr(
+            config, "head_dim", config.hidden_size // config.num_attention_heads
+        )
+        self.num_key_value_groups = (
+            config.num_attention_heads // config.num_key_value_heads
+        )
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = config.attention_dropout
+        self.is_causal = True
+
+        self.q_proj = nn.Linear(
+            config.hidden_size,
+            config.num_attention_heads * self.head_dim,
+            bias=config.attention_bias,
+        )
+        self.k_proj = nn.Linear(
+            config.hidden_size,
+            config.num_key_value_heads * self.head_dim,
+            bias=config.attention_bias,
+        )
+        self.v_proj = nn.Linear(
+            config.hidden_size,
+            config.num_key_value_heads * self.head_dim,
+            bias=config.attention_bias,
+        )
+        self.o_proj = nn.Linear(
+            config.num_attention_heads * self.head_dim,
+            config.hidden_size,
+            bias=config.attention_bias,
+        )
+        self.q_norm = nn.Identity()
+        self.k_norm = nn.Identity()
+        self.sliding_window = config.sliding_window
+        self.rotary_emb = RoPE(config)
+
+    def __call__(
+        self,
+        hidden_states: mx.array,
+        position_embeddings: Optional[Tuple[mx.array, mx.array]] = None,
+        attention_mask: Optional[mx.array] = None,
+        position_ids: Optional[mx.array] = None,
+    ) -> Tuple[mx.array, Optional[mx.array]]:
+        B, L, D = hidden_states.shape
+        hidden_shape = (B, L, -1, self.head_dim)
+
+        query_states = self.q_norm(
+            self.q_proj(hidden_states).reshape(*hidden_shape)
+        ).transpose(0, 2, 1, 3)
+        key_states = self.k_norm(
+            self.k_proj(hidden_states).reshape(*hidden_shape)
+        ).transpose(0, 2, 1, 3)
+        value_states = (
+            self.v_proj(hidden_states).reshape(*hidden_shape).transpose(0, 2, 1, 3)
+        )
+
+        if position_embeddings is None:
+            if position_ids is None:
+                position_ids = mx.arange(L)
+                position_ids = mx.expand_dims(position_ids, axis=0)
+            cos, sin = self.rotary_emb(hidden_states, position_ids)
+        else:
+            cos, sin = position_embeddings
+
+        query_states, key_states = apply_rotary_pos_emb(
+            query_states, key_states, cos, sin
+        )
+
+        if attention_mask is not None and isinstance(attention_mask, mx.array):
+            kv_seq_len = key_states.shape[-2]
+            if attention_mask.shape[-1] != kv_seq_len:
+                attention_mask = attention_mask[..., :kv_seq_len]
+
+        if self.is_causal and attention_mask is None:
+            attention_mask = nn.MultiHeadAttention.create_additive_causal_mask(L)
+            attention_mask = attention_mask.astype(query_states.dtype)
+
+        attn_output = scaled_dot_product_attention(
+            query_states,
+            key_states,
+            value_states,
+            None,
+            scale=self.scaling,
+            mask=attention_mask,
+        )
+
+        attn_output = attn_output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        attn_output = self.o_proj(attn_output)
+        return attn_output, None
+
+
+class Code2WavMlp(nn.Module):
+    def __init__(self, config: Code2WavConfig):
+        super().__init__()
+
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+
+        if config.hidden_act == "silu":
+            self.act_fn = nn.silu
+        elif config.hidden_act == "gelu":
+            self.act_fn = nn.gelu
+        elif config.hidden_act == "gelu_pytorch_tanh":
+            self.act_fn = nn.GELU(approx="precise")
+        else:
+            self.act_fn = nn.silu
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+
+class Code2WavTransformerLayer(nn.Module):
+    def __init__(self, config: Code2WavConfig, idx: int):
+        super().__init__()
+        self.self_attn = Code2WavAttention(config, idx)
+        self.mlp = Code2WavMlp(config)
+        self.input_layernorm = nn.RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(
+            config.hidden_size, config.rms_norm_eps
+        )
+        self.self_attn_layer_scale = LayerScale(config)
+        self.mlp_layer_scale = LayerScale(config)
+
+    def __call__(
+        self,
+        hidden_states: mx.array,
+        attention_mask: Optional[mx.array] = None,
+        position_ids: Optional[mx.array] = None,
+        position_embeddings: Optional[Tuple[mx.array, mx.array]] = None,
+    ) -> mx.array:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states, _ = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            position_embeddings=position_embeddings,
+        )
+        hidden_states = residual + self.self_attn_layer_scale(hidden_states)
+
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + self.mlp_layer_scale(hidden_states)
+
+        return hidden_states
+
+
+class Code2WavTransformerModel(nn.Module):
+    def __init__(self, config: Code2WavConfig):
+        super().__init__()
+
+        self.layers = [
+            Code2WavTransformerLayer(config, idx)
+            for idx in range(config.num_hidden_layers)
+        ]
+        self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = RoPE(config)
+
+    def __call__(
+        self,
+        inputs_embeds: mx.array,
+        attention_mask: Optional[mx.array] = None,
+        position_ids: Optional[mx.array] = None,
+    ) -> mx.array:
+        hidden_states = inputs_embeds
+
+        if position_ids is None:
+            position_ids = mx.arange(hidden_states.shape[1])
+            position_ids = mx.expand_dims(position_ids, axis=0)
+
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        for layer in self.layers:
+            hidden_states = layer(
+                hidden_states,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                position_embeddings=position_embeddings,
+            )
+
+        hidden_states = self.norm(hidden_states)
+        return hidden_states
+
+
+class Code2WavModel(nn.Module):
+    def __init__(self, config: Code2WavConfig):
+        super().__init__()
+
+        self.pre_transformer = Code2WavTransformerModel(config)
+        self.code_embedding = nn.Embedding(
+            config.codebook_size * config.num_quantizers, config.hidden_size
+        )
+        self.upsample = [
+            [
+                CausalTransConvNet(
+                    config.hidden_size, config.hidden_size, factor, factor
+                ),
+                ConvNeXtBlock(config.hidden_size),
+            ]
+            for factor in config.upsampling_ratios
+        ]
+        self.decoder = [CausalConvNet(config.hidden_size, config.decoder_dim, 7)]
+        self.decoder.extend(
+            [
+                Code2WavDecoderBlock(config, idx)
+                for idx in range(len(config.upsample_rates))
+            ]
+        )
+        output_dim = config.decoder_dim // 2 ** len(config.upsample_rates)
+        self.decoder.extend([SnakeBeta(output_dim), CausalConvNet(output_dim, 1, 7)])
+        self.config = config
+        self.code_offset = (
+            np.arange(config.num_quantizers).reshape(1, -1, 1) * config.codebook_size
+        )
+
+    def __call__(
+        self, codes: mx.array = None, input_embeds: mx.array = None
+    ) -> mx.array:
+        if input_embeds is not None:
+            hidden = input_embeds
+        elif codes is not None:
+            if codes.shape[1] != self.config.num_quantizers:
+                raise ValueError(
+                    f"Expected {self.config.num_quantizers} layer of codes, got {codes.shape[1]}"
+                )
+            hidden = self.code_embedding(codes + mx.array(self.code_offset)).mean(1)
+        else:
+            raise ValueError("Must provide codes or input_embeds")
+
+        hidden = self.pre_transformer(inputs_embeds=hidden)
+        hidden = hidden.transpose(0, 2, 1)
+        for blocks in self.upsample:
+            for block in blocks:
+                hidden = block(hidden)
+        wav = hidden
+        for block in self.decoder:
+            wav = block(wav)
+        return mx.clip(wav, -1, 1)
+
+    def chunked_decode(self, codes, chunk_size=300, left_context_size=25):
+        total_upsample_factor = 1
+        for r in self.config.upsampling_ratios:
+            total_upsample_factor *= r
+        for r in self.config.upsample_rates:
+            total_upsample_factor *= r
+
+        B, Q, L = codes.shape
+        final_wav_list = []
+
+        for start in range(0, L, chunk_size):
+            end = min(start + chunk_size, L)
+            context_start = max(0, start - left_context_size)
+            chunk_codes = codes[:, :, context_start:end]
+            wav_chunk = self(codes=chunk_codes)
+            context_len_tokens = start - context_start
+            valid_start_sample = context_len_tokens * total_upsample_factor
+            current_chunk_valid_len_tokens = end - start
+            valid_len_samples = current_chunk_valid_len_tokens * total_upsample_factor
+            chunk_valid_wav = wav_chunk[
+                :, :, valid_start_sample : valid_start_sample + valid_len_samples
+            ]
+            final_wav_list.append(chunk_valid_wav)
+
+        return mx.concatenate(final_wav_list, axis=-1)
+
+    def stream_decode(
+        self, codes_buffer, chunk_size=300, left_context_size=25, decoded_len=0
+    ):
+        total_upsample_factor = 1
+        for r in self.config.upsampling_ratios:
+            total_upsample_factor *= r
+        for r in self.config.upsample_rates:
+            total_upsample_factor *= r
+
+        L = codes_buffer.shape[2]
+        start = decoded_len
+        context_start = max(0, start - left_context_size)
+        context_len = start - context_start
+        new_tokens = chunk_size - context_len
+        if L - start < new_tokens:
+            return None, decoded_len
+
+        end = start + new_tokens
+        chunk_codes = codes_buffer[:, :, context_start:end]
+        wav_chunk = self(codes=chunk_codes)
+        context_len_tokens = start - context_start
+        valid_start_sample = context_len_tokens * total_upsample_factor
+        current_chunk_valid_len_tokens = end - start
+        valid_len_samples = current_chunk_valid_len_tokens * total_upsample_factor
+        chunk_valid_wav = wav_chunk[
+            :, :, valid_start_sample : valid_start_sample + valid_len_samples
+        ]
+        return chunk_valid_wav, end
+
+    def flush_decode(self, codes_buffer, left_context_size=25, decoded_len=0):
+        total_upsample_factor = 1
+        for r in self.config.upsampling_ratios:
+            total_upsample_factor *= r
+        for r in self.config.upsample_rates:
+            total_upsample_factor *= r
+
+        L = codes_buffer.shape[2]
+        if decoded_len >= L:
+            return None
+
+        start = decoded_len
+        context_start = max(0, start - left_context_size)
+        chunk_codes = codes_buffer[:, :, context_start:]
+        wav_chunk = self(codes=chunk_codes)
+        context_len_tokens = start - context_start
+        valid_start_sample = context_len_tokens * total_upsample_factor
+        return wav_chunk[:, :, valid_start_sample:]
+
+    def sanitize(self, weights):
+        sanitized_weights = {}
+        for k, v in weights.items():
+            if ("upsample" in k and "conv.weight" in k and "dwconv" not in k) or (
+                "decoder" in k
+                and "block" in k
+                and "conv.weight" in k
+                and "conv1" not in k
+                and "conv2" not in k
+            ):
+                sanitized_weights[k] = v.transpose(1, 2, 0)
+            elif (
+                ("dwconv.conv.weight" in k)
+                or ("decoder" in k and "conv.weight" in k and "block" not in k)
+                or (
+                    "decoder" in k
+                    and "block" in k
+                    and ("conv1.conv.weight" in k or "conv2.conv.weight" in k)
+                )
+            ):
+                sanitized_weights[k] = v.transpose(0, 2, 1)
+            else:
+                sanitized_weights[k] = v
+
+        return sanitized_weights