Added audex_causal_speech_decoder/modeling_audex_causal_speech_decoder.py

audex_causal_speech_decoder/modeling_audex_causal_speech_decoder.py

@@ -0,0 +1,482 @@
+# coding=utf-8
+# Copyright (c) 2026, NVIDIA CORPORATION.  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 __future__ import annotations
+
+from collections.abc import Iterator, Sequence
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+from transformers import PreTrainedModel
+
+from .configuration_audex_causal_speech_decoder import AudexCausalSpeechDecoderConfig
+from .streaming_utils import load_audex_causal_speech_decoder as _load_decoder_for_remote_code
+
+REMOTE_CODE_IMPORTS = (_load_decoder_for_remote_code,)
+
+
+class RotaryPositionalEmbeddings(nn.Module):
+    def __init__(self, dim: int, max_seq_len: int = 4096, base: int = 10_000) -> None:
+        super().__init__()
+        self.dim = dim
+        self.base = base
+        self.max_seq_len = max_seq_len
+        self.rope_init()
+        self._rope_ready = False
+
+    def rope_init(self, device: "torch.device | None" = None) -> None:
+        theta = 1.0 / (
+            self.base
+            ** (torch.arange(0, self.dim, 2, device=device)[: (self.dim // 2)].float() / self.dim)
+        )
+        self.register_buffer("theta", theta, persistent=False)
+        self.build_rope_cache(self.max_seq_len)
+
+    def build_rope_cache(self, max_seq_len: int = 4096) -> None:
+        self.max_seq_len = max_seq_len
+        seq_idx = torch.arange(max_seq_len, dtype=self.theta.dtype, device=self.theta.device)
+        idx_theta = torch.einsum("i, j -> ij", seq_idx, self.theta).float()
+        cache = torch.stack([torch.cos(idx_theta), torch.sin(idx_theta)], dim=-1)
+        self.register_buffer("cache", cache, persistent=False)
+
+    def forward(self, x: torch.Tensor, *, input_pos: torch.Tensor | None = None) -> torch.Tensor:
+        seq_len = x.size(1)
+        needed_seq_len = seq_len if input_pos is None else int(input_pos.max().item()) + 1
+        if (
+            not getattr(self, "_rope_ready", False)
+            or self.theta.device != x.device
+            or needed_seq_len > self.cache.size(0)
+        ):
+            self.rope_init(device=x.device)
+            if needed_seq_len > self.cache.size(0):
+                self.build_rope_cache(max(needed_seq_len, self.cache.size(0) * 2))
+            self._rope_ready = True
+
+        rope_cache = self.cache[:seq_len] if input_pos is None else self.cache[input_pos]
+        xshaped = x.float().reshape(*x.shape[:-1], -1, 2)
+        rope_cache = rope_cache.view(-1, xshaped.size(1), 1, xshaped.size(3), 2)
+        x_out = torch.stack(
+            [
+                xshaped[..., 0] * rope_cache[..., 0] - xshaped[..., 1] * rope_cache[..., 1],
+                xshaped[..., 1] * rope_cache[..., 0] + xshaped[..., 0] * rope_cache[..., 1],
+            ],
+            -1,
+        )
+        return x_out.flatten(3).type_as(x)
+
+
+class CausalCodecDecoderCache:
+    def __init__(self) -> None:
+        self.key_values: dict[int, tuple[Tensor, Tensor]] = {}
+        self.position = 0
+
+    def input_positions(self, length: int, device: torch.device) -> Tensor:
+        return torch.arange(self.position, self.position + length, device=device).unsqueeze(0)
+
+    def update(self, layer_idx: int, key: Tensor, value: Tensor) -> tuple[Tensor, Tensor]:
+        if layer_idx in self.key_values:
+            prev_key, prev_value = self.key_values[layer_idx]
+            key = torch.cat([prev_key, key], dim=2)
+            value = torch.cat([prev_value, value], dim=2)
+        self.key_values[layer_idx] = (key, value)
+        return key, value
+
+    def advance(self, length: int) -> None:
+        self.position += length
+
+    def reset(self) -> None:
+        self.key_values.clear()
+        self.position = 0
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        norm_x = torch.mean(x**2, dim=-1, keepdim=True)
+        return x * torch.rsqrt(norm_x + self.eps) * self.weight
+
+
+class MLP(nn.Module):
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.fc1 = nn.Linear(dim, 4 * dim, bias=False)
+        self.silu = nn.SiLU()
+        self.fc2 = nn.Linear(4 * dim, dim, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.fc2(self.silu(self.fc1(x)))
+
+
+class Attention(nn.Module):
+    def __init__(self, dim: int, n_heads: int, rotary_embed: RotaryPositionalEmbeddings, layer_idx: int):
+        super().__init__()
+        if dim % n_heads != 0:
+            raise ValueError(f"dim must be divisible by n_heads, got dim={dim}, n_heads={n_heads}")
+        self.n_heads = n_heads
+        self.layer_idx = layer_idx
+        self.rotary_embed = rotary_embed
+        self.c_attn = nn.Linear(dim, 3 * dim, bias=False)
+        self.c_proj = nn.Linear(dim, dim, bias=False)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cache: CausalCodecDecoderCache | None = None,
+        input_pos: Tensor | None = None,
+    ) -> torch.Tensor:
+        batch_size, seq_len, _ = x.shape
+        qkv = self.c_attn(x)
+        head_dim = qkv.size(-1) // (3 * self.n_heads)
+        qkv = qkv.view(batch_size, seq_len, 3, self.n_heads, head_dim).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv.unbind(0)
+
+        q = self.rotary_embed(q.transpose(1, 2), input_pos=input_pos).transpose(1, 2)
+        k = self.rotary_embed(k.transpose(1, 2), input_pos=input_pos).transpose(1, 2)
+        if cache is None:
+            y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        else:
+            if input_pos is None:
+                raise ValueError("input_pos is required when cache is set")
+            k, v = cache.update(self.layer_idx, k, v)
+            key_pos = torch.arange(k.size(2), device=x.device).view(1, 1, 1, -1)
+            attn_mask = key_pos <= input_pos.view(input_pos.size(0), 1, -1, 1)
+            y = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
+        return self.c_proj(y.transpose(1, 2).contiguous().view(batch_size, seq_len, -1))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, dim: int, n_heads: int, rotary_embed: RotaryPositionalEmbeddings, layer_idx: int):
+        super().__init__()
+        self.att_norm = RMSNorm(dim)
+        self.ffn_norm = RMSNorm(dim)
+        self.att = Attention(dim=dim, n_heads=n_heads, rotary_embed=rotary_embed, layer_idx=layer_idx)
+        self.mlp = MLP(dim=dim)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cache: CausalCodecDecoderCache | None = None,
+        input_pos: Tensor | None = None,
+    ) -> torch.Tensor:
+        x = x + self.att(self.att_norm(x), cache=cache, input_pos=input_pos)
+        return x + self.mlp(self.ffn_norm(x))
+
+
+class PatchHead(nn.Module):
+    def __init__(self, dim: int, hop_length: int = 320):
+        super().__init__()
+        self.proj = nn.Linear(dim, hop_length, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = torch.tanh(self.proj(x))
+        return x.reshape(x.size(0), 1, -1)
+
+
+class CausalVocosBackbone(nn.Module):
+    def __init__(
+        self,
+        hidden_dim: int = 2048,
+        depth: int = 12,
+        heads: int = 32,
+        pos_meb_dim: int = 64,
+    ):
+        super().__init__()
+        rotary_embed = RotaryPositionalEmbeddings(dim=pos_meb_dim)
+        self.transformers = nn.ModuleList(
+            [
+                TransformerBlock(dim=hidden_dim, n_heads=heads, rotary_embed=rotary_embed, layer_idx=idx)
+                for idx in range(depth)
+            ]
+        )
+        self.final_layer_norm = RMSNorm(hidden_dim)
+
+    def forward(self, x: torch.Tensor, cache: CausalCodecDecoderCache | None = None) -> torch.Tensor:
+        input_pos = None
+        if cache is not None:
+            input_pos = cache.input_positions(x.size(1), x.device).expand(x.size(0), -1)
+        for block in self.transformers:
+            x = block(x, cache=cache, input_pos=input_pos)
+        if cache is not None:
+            cache.advance(x.size(1))
+        return self.final_layer_norm(x)
+
+
+class CausalCodecDecoderVocos(nn.Module):
+    def __init__(
+        self,
+        hidden_dim: int = 2048,
+        depth: int = 12,
+        heads: int = 32,
+        pos_meb_dim: int = 64,
+        hop_length: int = 320,
+        vq_dim: int = 2048,
+        lookahead_steps: int = 0,
+    ):
+        super().__init__()
+        if lookahead_steps < 0:
+            raise ValueError(f"lookahead_steps must be >= 0, got {lookahead_steps}")
+
+        self.wav_proj = nn.Linear(hop_length, hidden_dim, bias=False)
+        self.fc_post_a = nn.Linear(vq_dim, hidden_dim, bias=False)
+        self.lookahead_steps = lookahead_steps
+        if lookahead_steps > 0:
+            self.lookahead_conv = nn.Conv1d(
+                hidden_dim,
+                hidden_dim,
+                kernel_size=lookahead_steps + 1,
+                padding=0,
+                groups=hidden_dim,
+                bias=False,
+            )
+            self.lookahead_act = nn.SiLU()
+            self.lookahead_proj = nn.Conv1d(hidden_dim, hidden_dim, kernel_size=1, bias=False)
+            nn.init.zeros_(self.lookahead_proj.weight)
+        else:
+            self.lookahead_conv = None
+            self.lookahead_act = None
+            self.lookahead_proj = None
+        self.backbone = CausalVocosBackbone(hidden_dim, depth, heads, pos_meb_dim)
+        self.head = PatchHead(hidden_dim, hop_length)
+
+    def _project_tokens(self, vq_emb: torch.Tensor) -> torch.Tensor:
+        return self.fc_post_a(vq_emb)
+
+    def _apply_lookahead(self, x: torch.Tensor) -> torch.Tensor:
+        if self.lookahead_conv is None:
+            return x
+        if self.lookahead_act is None or self.lookahead_proj is None:
+            raise RuntimeError("lookahead modules are not initialized")
+        h = F.pad(x.transpose(1, 2), (0, self.lookahead_steps))
+        h = self.lookahead_proj(self.lookahead_act(self.lookahead_conv(h)))
+        return x + h.transpose(1, 2)
+
+    def _apply_lookahead_window(self, x: torch.Tensor) -> torch.Tensor:
+        if self.lookahead_conv is None:
+            return x
+        if self.lookahead_act is None or self.lookahead_proj is None:
+            raise RuntimeError("lookahead modules are not initialized")
+        if x.size(1) <= self.lookahead_steps:
+            raise ValueError(f"lookahead window needs more than {self.lookahead_steps} frames, got {x.size(1)}")
+        h = self.lookahead_proj(self.lookahead_act(self.lookahead_conv(x.transpose(1, 2))))
+        return x[:, : h.size(2)] + h.transpose(1, 2)
+
+    def decode_cached(
+        self,
+        vq_emb: torch.Tensor,
+        cache: CausalCodecDecoderCache,
+        lookahead_vq_emb: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        x = self._project_tokens(vq_emb)
+        if self.lookahead_steps > 0:
+            if lookahead_vq_emb is None:
+                lookahead_vq_emb = vq_emb.new_zeros(vq_emb.size(0), self.lookahead_steps, vq_emb.size(-1))
+            if lookahead_vq_emb.size(1) != self.lookahead_steps:
+                raise ValueError(
+                    f"lookahead_vq_emb must have {self.lookahead_steps} frames, got {lookahead_vq_emb.size(1)}"
+                )
+            lookahead_x = self._project_tokens(lookahead_vq_emb)
+            x = self._apply_lookahead_window(torch.cat([x, lookahead_x], dim=1))
+        x = self.backbone(x, cache=cache)
+        return self.head(x)
+
+    def forward(
+        self,
+        vq_emb: torch.Tensor,
+        patched_wav: torch.Tensor | None = None,
+        alpha: float = 0.0,
+    ) -> torch.Tensor:
+        x = self._project_tokens(vq_emb)
+        x = self._apply_lookahead(x)
+        if patched_wav is not None:
+            h = self.wav_proj(patched_wav)
+            mask = torch.bernoulli(
+                torch.full(
+                    (x.size(0), x.size(1), 1),
+                    min(max(alpha, 0.0), 1.0),
+                    device=x.device,
+                    dtype=x.dtype,
+                )
+            )
+            x = x + h * mask
+        return self.head(self.backbone(x))
+
+
+class AudexSpeechTokenEmbedder(nn.Module):
+    def __init__(
+        self,
+        output_dim: int,
+        token_embed_dim: int,
+        codebook_levels: Sequence[int],
+    ) -> None:
+        super().__init__()
+        if len(codebook_levels) != token_embed_dim:
+            raise ValueError(
+                f"token_embed_dim={token_embed_dim} must match codebook_levels length={len(codebook_levels)}"
+            )
+        self.codebook_levels = tuple(int(level) for level in codebook_levels)
+        self.project_out = nn.Linear(token_embed_dim, output_dim)
+
+    def forward(self, indices: torch.Tensor) -> torch.Tensor:
+        if indices.size(-1) != 1:
+            raise ValueError(f"indices last dimension must be 1, got {indices.size(-1)}")
+        levels = torch.tensor(self.codebook_levels, dtype=torch.long, device=indices.device)
+        basis = torch.cumprod(torch.cat([levels.new_ones(1), levels[:-1]]), dim=0)
+        level_indices = (indices.long() // basis) % levels
+        dtype = self.project_out.weight.dtype
+        codes = level_indices.to(dtype=dtype)
+        levels = levels.to(dtype=dtype)
+        codes = codes * (2.0 / (levels - 1.0)) - 1.0
+        return self.project_out(codes)
+
+    def get_output_from_indices(self, indices: torch.Tensor) -> torch.Tensor:
+        return self(indices)
+
+
+class AudexCausalSpeechDecoderModel(PreTrainedModel):
+    config_class = AudexCausalSpeechDecoderConfig
+    base_model_prefix = "module"
+    all_tied_weights_keys: dict[str, Any] = {}
+    Cache = CausalCodecDecoderCache
+
+    def __init__(self, config: AudexCausalSpeechDecoderConfig):
+        super().__init__(config)
+        self.audex_speech_token_embedder = AudexSpeechTokenEmbedder(
+            output_dim=config.vq_dim,
+            token_embed_dim=config.token_embed_dim,
+            codebook_levels=config.codebook_levels,
+        )
+        self.module = CausalCodecDecoderVocos(
+            hidden_dim=config.hidden_dim,
+            depth=config.depth,
+            heads=config.heads,
+            pos_meb_dim=config.pos_meb_dim,
+            hop_length=config.hop_length,
+            vq_dim=config.vq_dim,
+            lookahead_steps=config.lookahead_steps,
+        )
+
+    @property
+    def lookahead_steps(self) -> int:
+        return self.module.lookahead_steps
+
+    def create_cache(self) -> CausalCodecDecoderCache:
+        return CausalCodecDecoderCache()
+
+    def decode_cached(
+        self,
+        vq_emb: torch.Tensor,
+        cache: CausalCodecDecoderCache,
+        lookahead_vq_emb: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        return self.module.decode_cached(vq_emb, cache, lookahead_vq_emb=lookahead_vq_emb)
+
+    def create_session(
+        self,
+        *,
+        chunk_frames: int = 1,
+        sample_rate: int | None = None,
+        return_numpy: bool = True,
+    ) -> "AudexCausalSpeechDecoderSession":
+        return AudexCausalSpeechDecoderSession(
+            decoder=self,
+            chunk_frames=chunk_frames,
+            sample_rate=sample_rate or self.config.sample_rate,
+            return_numpy=return_numpy,
+        )
+
+    def forward(
+        self,
+        vq_emb: torch.Tensor,
+        patched_wav: torch.Tensor | None = None,
+        alpha: float = 0.0,
+    ) -> torch.Tensor:
+        return self.module(vq_emb, patched_wav=patched_wav, alpha=alpha)
+
+
+class AudexCausalSpeechDecoderSession:
+    def __init__(
+        self,
+        decoder: AudexCausalSpeechDecoderModel,
+        *,
+        chunk_frames: int,
+        sample_rate: int,
+        return_numpy: bool,
+    ):
+        if chunk_frames <= 0:
+            raise ValueError(f"chunk_frames must be positive, got {chunk_frames}")
+        self.decoder = decoder
+        self.chunk_frames = chunk_frames
+        self.sample_rate = sample_rate
+        self.return_numpy = return_numpy
+        self.cache = decoder.create_cache()
+        self.buffer: list[list[int]] = []
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.decoder.parameters()).device
+
+    def reset(self) -> None:
+        self.cache = self.decoder.create_cache()
+        self.buffer.clear()
+
+    def push(self, token_frames: Sequence[Sequence[int]]) -> Iterator[tuple[int, Any]]:
+        self.buffer.extend(list(frame) for frame in token_frames)
+        yield from self._drain(flush=False)
+
+    def flush(self) -> Iterator[tuple[int, Any]]:
+        yield from self._drain(flush=True)
+
+    def _drain(self, *, flush: bool) -> Iterator[tuple[int, Any]]:
+        ready_frames = len(self.buffer) - self.decoder.lookahead_steps
+        while self.buffer and (flush or ready_frames >= self.chunk_frames):
+            emit_frames = min(self.chunk_frames, len(self.buffer)) if flush else self.chunk_frames
+            wav = self._decode_buffered_frames(emit_frames, flush=flush)
+            del self.buffer[:emit_frames]
+            ready_frames = len(self.buffer) - self.decoder.lookahead_steps
+            yield self.sample_rate, self._format_chunk(wav)
+
+    def _embed_speech_token_frames(self, token_frames: Sequence[Sequence[int]]) -> torch.Tensor:
+        indices = torch.tensor(token_frames, dtype=torch.long, device=self.device).unsqueeze(0)
+        return self.decoder.audex_speech_token_embedder.get_output_from_indices(indices)
+
+    def _decode_buffered_frames(self, emit_frames: int, *, flush: bool) -> torch.Tensor:
+        with torch.inference_mode():
+            vq_emb = self._embed_speech_token_frames(self.buffer[:emit_frames])
+            lookahead_vq_emb = None
+            if self.decoder.lookahead_steps > 0:
+                future_frames = self.buffer[emit_frames : emit_frames + self.decoder.lookahead_steps]
+                future_parts = []
+                if future_frames:
+                    future_parts.append(self._embed_speech_token_frames(future_frames))
+                missing_frames = self.decoder.lookahead_steps - len(future_frames) if flush else 0
+                if missing_frames > 0:
+                    future_parts.append(vq_emb.new_zeros(vq_emb.size(0), missing_frames, vq_emb.size(-1)))
+                lookahead_vq_emb = torch.cat(future_parts, dim=1) if future_parts else None
+            return self.decoder.decode_cached(vq_emb, self.cache, lookahead_vq_emb=lookahead_vq_emb)
+
+    def _format_chunk(self, wav: torch.Tensor) -> Any:
+        chunk = wav.squeeze().float().detach().cpu()
+        if not self.return_numpy:
+            return chunk
+
+        import numpy as np
+
+        return chunk.numpy().astype(np.float32, copy=False)