Repair PersonaPlex for Blackwell ZeroGPU

PERSONAPLEX_LICENSE-MIT

@@ -0,0 +1,23 @@
+Permission is hereby granted, free of charge, to any
+person obtaining a copy of this software and associated
+documentation files (the "Software"), to deal in the
+Software without restriction, including without
+limitation the rights to use, copy, modify, merge,
+publish, distribute, sublicense, and/or sell copies of
+the Software, and to permit persons to whom the Software
+is furnished to do so, subject to the following
+conditions:
+
+The above copyright notice and this permission notice
+shall be included in all copies or substantial portions
+of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF
+ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED
+TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
+PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
+SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
+CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
+OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR
+IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+DEALINGS IN THE SOFTWARE.

PERSONAPLEX_LICENSE.audiocraft

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) Meta Platforms, Inc. and affiliates.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

PERSONAPLEX_LICENSE.moshi

@@ -0,0 +1,27 @@
+MIT License
+
+Copyright (c) Kyutai
+
+Permission is hereby granted, free of charge, to any
+person obtaining a copy of this software and associated
+documentation files (the "Software"), to deal in the
+Software without restriction, including without
+limitation the rights to use, copy, modify, merge,
+publish, distribute, sublicense, and/or sell copies of
+the Software, and to permit persons to whom the Software
+is furnished to do so, subject to the following
+conditions:
+
+The above copyright notice and this permission notice
+shall be included in all copies or substantial portions
+of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF
+ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED
+TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
+PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
+SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
+CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
+OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR
+IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+DEALINGS IN THE SOFTWARE.

README.md

@@ -61,6 +61,10 @@ ZeroGPU. PersonaPlex is the practical feasibility fallback.
 - Model: `nvidia/personaplex-7b-v1`
 - Model revision: `fdaf4090a61cb315c138a1faee287ffd6c716309`
 - NVIDIA PersonaPlex source revision: `3428dfd95309a7f3c84fd93259ded0f810d1ff91`
+- The required `moshi` Python package is vendored unchanged from that revision so RoleForge can use the
+  ZeroGPU-supported PyTorch 2.8 / CUDA 12.8 runtime. The upstream package metadata restricts PyTorch to `<2.5`,
+  which cannot execute on ZeroGPU's current Blackwell `sm_120` workers. The upstream MIT/Moshi/AudioCraft
+  license files are preserved at the repository root.
 - Base ZeroGPU implementation reviewed from `Javedalam/PersonaPlex` at
   `1b0acf6f195ec8d33e48907531909f1409c17d34`
 

VENDORED_PERSONAPLEX.md

@@ -0,0 +1,18 @@
+# Vendored PersonaPlex runtime
+
+RoleForge vendors the `moshi/moshi` Python source directory unchanged from NVIDIA PersonaPlex commit
+`3428dfd95309a7f3c84fd93259ded0f810d1ff91`.
+
+Upstream: <https://github.com/NVIDIA/personaplex>
+
+Reason: the upstream package metadata requires `torch >=2.2,<2.5`. Hugging Face ZeroGPU now uses NVIDIA RTX Pro
+6000 Blackwell workers (`sm_120`) and officially supports PyTorch 2.8 or newer. Installing the upstream package
+normally therefore selected a PyTorch build with kernels only through `sm_90`, causing every real GPU turn to
+fail with `no kernel image is available for execution on the device`.
+
+The vendored Python files are unchanged. Only dependency resolution moved into RoleForge's `requirements.txt`,
+which pins PyTorch 2.8.0 with CUDA 12.8. The corresponding upstream license files are preserved as:
+
+- `PERSONAPLEX_LICENSE-MIT`
+- `PERSONAPLEX_LICENSE.moshi`
+- `PERSONAPLEX_LICENSE.audiocraft`

app.py

@@ -16,6 +16,9 @@ import torch
 from huggingface_hub import hf_hub_download
 
 
+print(f"RoleForge runtime: torch={torch.__version__}, cuda={torch.version.cuda}")
+
+
 MODEL_REPO = "nvidia/personaplex-7b-v1"
 MODEL_REVISION = "fdaf4090a61cb315c138a1faee287ffd6c716309"
 DEVICE = "cuda"

moshi/__init__.py

@@ -0,0 +1,18 @@
+# Copyright (c) Kyutai, all rights reserved.
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+moshi is the inference codebase for Kyutai audio generation models.
+
+The code has been adapted from Audiocraft, see LICENSE.audiocraft
+  Copyright (c) Meta Platforms, Inc. and affiliates.
+"""
+
+# flake8: noqa
+from . import utils
+from . import modules
+from . import models
+from . import quantization
+
+__version__ = "0.1.0"

moshi/client_utils.py

@@ -0,0 +1,213 @@
+# Copyright (c) Kyutai, all rights reserved.
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+"""Utilities for the command line client, in particular for handling interactions with the terminal.
+"""
+
+from dataclasses import dataclass
+import sys
+
+
+def colorize(text, color):
+    code = f"\033[{color}m"
+    restore = "\033[0m"
+    return "".join([code, text, restore])
+
+
+def make_log(level: str, msg: str) -> str:
+    if level == "warning":
+        prefix = colorize("[Warn]", "1;31")
+    elif level == "info":
+        prefix = colorize("[Info]", "1;34")
+    elif level == "error":
+        prefix = colorize("[Err ]", "1;31")
+    else:
+        raise ValueError(f"Unknown level {level}")
+    return prefix + " " + msg
+
+
+class RawPrinter:
+    def __init__(self, stream=sys.stdout, err_stream=sys.stderr):
+        self.stream = stream
+        self.err_stream = err_stream
+
+    def print_header(self):
+        pass
+
+    def print_token(self, token: str):
+        self.stream.write(token)
+        self.stream.flush()
+
+    def log(self, level: str, msg: str):
+        print(f"{level.capitalize()}: {msg}", file=self.err_stream)
+
+    def print_lag(self):
+        self.err_stream.write(colorize(" [LAG]", "31"))
+        self.err_stream.flush()
+
+    def print_pending(self):
+        pass
+
+
+@dataclass
+class LineEntry:
+    msg: str
+    color: str | None = None
+
+    def render(self):
+        if self.color is None:
+            return self.msg
+        else:
+            return colorize(self.msg, self.color)
+
+    def __len__(self):
+        return len(self.msg)
+
+
+class Line:
+    def __init__(self, stream):
+        self.stream = stream
+        self._line: list[LineEntry] = []
+        self._has_padding: bool = False
+        self._max_line_length = 0
+
+    def __bool__(self):
+        return bool(self._line)
+
+    def __len__(self):
+        return sum(len(entry) for entry in self._line)
+
+    def add(self, msg: str, color: str | None = None) -> int:
+        entry = LineEntry(msg, color)
+        return self._add(entry)
+
+    def _add(self, entry: LineEntry) -> int:
+        if self._has_padding:
+            self.erase(count=0)
+        self._line.append(entry)
+        self.stream.write(entry.render())
+        self._max_line_length = max(self._max_line_length, len(self))
+        return len(entry)
+
+    def erase(self, count: int = 1):
+        if count:
+            entries = list(self._line[:-count])
+        else:
+            entries = list(self._line)
+        self._line.clear()
+        self.stream.write("\r")
+        for entry in entries:
+            self._line.append(entry)
+            self.stream.write(entry.render())
+
+        self._has_padding = False
+
+    def newline(self):
+        missing = self._max_line_length - len(self)
+        if missing > 0:
+            self.stream.write(" " * missing)
+        self.stream.write("\n")
+        self._line.clear()
+        self._max_line_length = 0
+        self._has_padding = False
+
+    def flush(self):
+        missing = self._max_line_length - len(self)
+        if missing > 0:
+            self.stream.write(" " * missing)
+            self._has_padding = True
+        self.stream.flush()
+
+
+class Printer:
+    def __init__(self, max_cols: int = 80, stream=sys.stdout, err_stream=sys.stderr):
+        self.max_cols = max_cols
+        self.line = Line(stream)
+        self.stream = stream
+        self.err_stream = err_stream
+        self._pending_count = 0
+        self._pending_printed = False
+
+    def print_header(self):
+        self.line.add(" " + "-" * (self.max_cols) + " ")
+        self.line.newline()
+        self.line.flush()
+        self.line.add("| ")
+
+    def _remove_pending(self) -> bool:
+        if self._pending_printed:
+            self._pending_printed = False
+            self.line.erase(1)
+            return True
+        return False
+
+    def print_token(self, token: str, color: str | None = None):
+        self._remove_pending()
+        remaining = self.max_cols - len(self.line)
+        if len(token) <= remaining:
+            self.line.add(token, color)
+        else:
+            end = " " * remaining + " |"
+            if token.startswith(" "):
+                token = token.lstrip()
+                self.line.add(end)
+                self.line.newline()
+                self.line.add("| ")
+                self.line.add(token, color)
+            else:
+                assert color is None
+                erase_count = None
+                cumulated = ""
+                for idx, entry in enumerate(self.line._line[::-1]):
+                    if entry.color:
+                        # probably a LAG message
+                        erase_count = idx
+                        break
+                    if entry.msg.startswith(" "):
+                        erase_count = idx + 1
+                        cumulated = entry.msg + cumulated
+                        break
+                if erase_count is not None:
+                    if erase_count > 0:
+                        self.line.erase(erase_count)
+                    remaining = self.max_cols - len(self.line)
+                    end = " " * remaining + " |"
+                    self.line.add(end)
+                    self.line.newline()
+                    self.line.add("| ")
+                    token = cumulated.lstrip() + token
+                    self.line.add(token)
+                else:
+                    self.line.add(token[:remaining])
+                    self.line.add(" |")
+                    self.line.newline()
+                    self.line.add("| ")
+                    self.line.add(token[remaining:])
+        self.line.flush()
+
+    def log(self, level: str, msg: str):
+        msg = make_log(level, msg)
+        self._remove_pending()
+        if self.line:
+            self.line.newline()
+        self.line.flush()
+        print(msg, file=self.err_stream)
+        self.err_stream.flush()
+
+    def print_lag(self):
+        self.print_token(" [LAG]", "31")
+
+    def print_pending(self):
+        chars = ["|", "/", "-", "\\"]
+        count = int(self._pending_count / 5)
+        char = chars[count % len(chars)]
+        colors = ["32", "33", "31"]
+        self._remove_pending()
+        self.line.add(char, colors[count % len(colors)])
+        self._pending_printed = True
+        self._pending_count += 1
+
+
+AnyPrinter = Printer | RawPrinter
+
+

moshi/models/__init__.py

@@ -0,0 +1,14 @@
+# Copyright (c) Kyutai, all rights reserved.
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+"""
+Models for the compression model Moshi,
+"""
+
+# flake8: noqa
+from .compression import (
+    CompressionModel,
+    MimiModel,
+)
+from .lm import LMModel, LMGen
+from .loaders import get_mimi, get_moshi_lm

moshi/models/compression.py

@@ -0,0 +1,474 @@
+# Copyright (c) Kyutai, all rights reserved.
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Part of this file is adapted from encodec.py in https://github.com/facebookresearch/audiocraft
+# released under the following license.
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+"""Compression models or wrapper around existing models. In particular, provides the implementation
+for Mimi. Also defines the main interface that a model must follow to be usable as an audio tokenizer.
+"""
+
+from abc import abstractmethod
+from contextlib import nullcontext
+from dataclasses import dataclass
+import logging
+import typing as tp
+
+import torch
+from torch import nn
+
+
+from ..quantization import (
+    QuantizedResult,
+    BaseQuantizer,
+    SplitResidualVectorQuantizer,
+    ResidualVectorQuantizer,
+)
+from ..modules.resample import ConvDownsample1d, ConvTrUpsample1d
+from ..modules.streaming import StreamingModule, State
+from ..utils.compile import no_compile, CUDAGraphed
+
+
+logger = logging.getLogger()
+
+
+class CompressionModel(StreamingModule[State]):
+    """Base API for all compression model that aim at being used as audio tokenizers
+    with a language model.
+    """
+
+    @abstractmethod
+    def forward(self, x: torch.Tensor) -> QuantizedResult: ...
+
+    @abstractmethod
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """See `MimiModel.encode`."""
+        ...
+
+    @abstractmethod
+    def decode(self, codes: torch.Tensor) -> torch.Tensor:
+        """See `MimiModel.decode`."""
+        ...
+
+    @abstractmethod
+    def decode_latent(self, codes: torch.Tensor) -> torch.Tensor:
+        """Decode from the discrete codes to continuous latent space."""
+        ...
+
+    @property
+    @abstractmethod
+    def channels(self) -> int: ...
+
+    @property
+    @abstractmethod
+    def frame_rate(self) -> float: ...
+
+    @property
+    @abstractmethod
+    def sample_rate(self) -> int: ...
+
+    @property
+    @abstractmethod
+    def cardinality(self) -> int: ...
+
+    @property
+    @abstractmethod
+    def num_codebooks(self) -> int: ...
+
+    @property
+    @abstractmethod
+    def total_codebooks(self) -> int: ...
+
+    @abstractmethod
+    def set_num_codebooks(self, n: int):
+        """Set the active number of codebooks used by the quantizer."""
+        ...
+
+
+@dataclass
+class _MimiState:
+    graphed_tr_enc: CUDAGraphed | None
+    graphed_tr_dec: CUDAGraphed | None
+
+    def reset(self):
+        pass
+
+
+class MimiModel(CompressionModel[_MimiState]):
+    """Mimi model operating on the raw waveform.
+
+    Args:
+        encoder (nn.Module): Encoder network.
+        decoder (nn.Module): Decoder network.
+        quantizer (qt.BaseQuantizer): Quantizer network.
+        frame_rate (float): Final frame rate of the quantized representatiopn.
+        encoder_frame_rate (float): frame rate of the encoder model. Note that if `frame_rate != encopder_frame_rate`,
+            the latent will be resampled linearly to match the desired `frame_rate` before and after quantization.
+        sample_rate (int): Audio sample rate.
+        channels (int): Number of audio channels.
+        causal (bool): Whether to use a causal version of the model.
+        encoder_transformer (nn.Module or None): optional transformer for the encoder.
+        decoder_transformer (nn.Module or None): optional transformer for the decoder.
+        resample_method (str): method to use for resampling the latent space before the quantizer.
+        upsample_channel_wise_bug (bool): controls whether the upsampling is channel wise.
+            Defaults to true to reproduce bug in original implementation.
+        freeze_encoder: whether to freeze the encoder weights.
+        freeze_quantizer: whether to freeze the quantizer weights.
+        freeze_quantizer_level: If positive, freeze the quantizer up to this level.
+        torch_compile_encoder_decoder (bool): if True, uses torch.compile on the encoder / decoder.
+            Deactivated by default for training as this is incompatible at the moment with weight norm.
+            See https://github.com/pytorch/pytorch/issues/121902
+            Also this seems to work well with 2.2.0, but completely fail with 2.4.0.
+    """
+
+    def __init__(
+        self,
+        encoder: nn.Module,
+        decoder: nn.Module,
+        quantizer: BaseQuantizer,
+        frame_rate: float,
+        encoder_frame_rate: float,
+        sample_rate: int,
+        channels: int,
+        causal: bool = False,
+        encoder_transformer: tp.Optional[nn.Module] = None,
+        decoder_transformer: tp.Optional[nn.Module] = None,
+        resample_method: str = "interpolate",
+        upsample_channel_wise_bug: bool = True,
+        freeze_encoder: bool = False,
+        freeze_quantizer: bool = False,
+        freeze_quantizer_level: int = -1,
+        torch_compile_encoder_decoder: bool = False,
+    ):
+        super().__init__()
+        self.encoder = encoder
+        self.decoder = decoder
+        self.encoder_transformer = encoder_transformer
+        self.decoder_transformer = decoder_transformer
+        self.quantizer = quantizer
+        self._frame_rate = frame_rate
+        self._sample_rate = sample_rate
+        self._channels = channels
+        self.encoder_frame_rate = encoder_frame_rate
+        self.torch_compile_encoder_decoder = torch_compile_encoder_decoder
+
+        if freeze_encoder:
+            for p in self.encoder.parameters():
+                p.requires_grad = False
+            if self.encoder_transformer is not None:
+                for p in self.encoder_transformer.parameters():
+                    p.requires_grad = False
+            for name, p in self.quantizer.named_parameters():
+                if name.endswith("input_proj.weight"):
+                    p.requires_grad = False
+        if freeze_quantizer:
+            self.quantizer.ema_frozen_(True)
+        self.freeze_quantizer = freeze_quantizer
+        self.freeze_quantizer_level = (
+            freeze_quantizer_level
+            if freeze_quantizer_level > 0
+            else self.quantizer.num_codebooks
+        )
+
+        # We will need the dimension for the resampling. In general the encoder will be a SeanetEncoder
+        # which exposes a `dimension` attribute.
+        dimension = encoder.dimension
+        assert isinstance(
+            dimension, int
+        ), f"Dimension should be int, got {dimension} of type {type(dimension)}."
+        self.dimension = dimension
+
+        assert resample_method in [
+            "interpolate",
+            "conv",
+            "avg_pool",
+        ], f"Invalid resample_method {resample_method}"
+        self.resample_method = resample_method
+        if encoder_frame_rate != frame_rate:
+            assert not (
+                causal and resample_method == "interpolate"
+            ), "Cannot interpolate with causal model."
+            if resample_method in ["conv", "avg_pool"]:
+                assert (
+                    self.encoder_frame_rate > self.frame_rate
+                ), "Cannot upsample with conv."
+                downsample_stride = self.encoder_frame_rate / self.frame_rate
+                assert downsample_stride == int(
+                    downsample_stride
+                ), f"Only integer strides are supported, got {downsample_stride}"
+                learnt = resample_method == "conv"
+                self.downsample = ConvDownsample1d(
+                    int(downsample_stride),
+                    dimension=dimension,
+                    learnt=learnt,
+                    causal=causal,
+                )
+                if freeze_encoder:
+                    for p in self.downsample.parameters():
+                        p.requires_grad = False
+                self.upsample = ConvTrUpsample1d(
+                    int(downsample_stride),
+                    dimension=dimension,
+                    learnt=learnt,
+                    causal=causal,
+                    channel_wise=upsample_channel_wise_bug,
+                )
+
+    def _init_streaming_state(self, batch_size: int) -> _MimiState:
+        device = next(self.parameters()).device
+        disable = device.type != 'cuda'
+        graphed_tr_dec = None
+        graphed_tr_enc = None
+        if self.encoder_transformer is not None:
+            graphed_tr_enc = CUDAGraphed(self.encoder_transformer, disable=disable)
+        if self.decoder_transformer is not None:
+            graphed_tr_dec = CUDAGraphed(self.decoder_transformer, disable=disable)
+        return _MimiState(graphed_tr_enc, graphed_tr_dec)
+
+    @property
+    def channels(self) -> int:
+        return self._channels
+
+    @property
+    def frame_rate(self) -> float:
+        return self._frame_rate
+
+    @property
+    def sample_rate(self) -> int:
+        return self._sample_rate
+
+    @property
+    def total_codebooks(self):
+        """Total number of quantizer codebooks available."""
+        return self.quantizer.total_codebooks
+
+    @property
+    def num_codebooks(self):
+        """Active number of codebooks used by the quantizer."""
+        return self.quantizer.num_codebooks
+
+    def set_num_codebooks(self, n: int):
+        """Set the active number of codebooks used by the quantizer."""
+        self.quantizer.set_num_codebooks(n)
+
+    @property
+    def cardinality(self):
+        """Cardinality of each codebook."""
+        return self.quantizer.cardinality
+
+    def _to_framerate(self, x: torch.Tensor):
+        # Convert from the encoder frame rate to the overall framerate.
+        _, _, length = x.shape
+        frame_rate = self.encoder_frame_rate
+        new_frame_rate = self.frame_rate
+        if frame_rate == new_frame_rate:
+            return x
+        if self.resample_method == "interpolate":
+            target_length = int(length * new_frame_rate / frame_rate)
+            return nn.functional.interpolate(x, size=target_length, mode="linear")
+        else:
+            return self.downsample(x)
+
+    def _to_encoder_framerate(self, x: torch.Tensor):
+        # Convert from overall framerate to the encoder frame rate.
+        _, _, length = x.shape
+        frame_rate = self.encoder_frame_rate
+        new_frame_rate = self.frame_rate
+        if frame_rate == new_frame_rate:
+            return x
+        if self.resample_method == "interpolate":
+            target_length = int(length * new_frame_rate / frame_rate)
+            return nn.functional.interpolate(x, size=target_length, mode="linear")
+        else:
+            return self.upsample(x)
+
+    @property
+    def _context_for_encoder_decoder(self):
+        if self.torch_compile_encoder_decoder:
+            return nullcontext()
+        else:
+            return no_compile()
+
+    def forward(self, x: torch.Tensor) -> QuantizedResult:
+        assert x.dim() == 3
+        length = x.shape[-1]
+        extra_metrics: tp.Dict[str, torch.Tensor] = {}
+
+        if self.freeze_quantizer:
+            if isinstance(self.quantizer, SplitResidualVectorQuantizer):
+                self.quantizer.rvq_first.eval()
+                for i in range(
+                    self.freeze_quantizer_level - self.quantizer.n_q_semantic
+                ):
+                    self.quantizer.rvq_rest.vq.layers[i].eval()
+            elif isinstance(self.quantizer, ResidualVectorQuantizer):
+                for i in range(self.freeze_quantizer_level):
+                    self.quantizer.vq.layers[i].eval()
+            else:
+                raise ValueError(f"Unsupported quantizer type {type(self.quantizer)}")
+
+        with self._context_for_encoder_decoder:
+            emb = self.encoder(x)
+        if self.encoder_transformer is not None:
+            (emb,) = self.encoder_transformer(emb)
+        emb = self._to_framerate(emb)
+        expected_length = self.frame_rate * length / self.sample_rate
+        # Checking that we have the proper length given the advertised frame rate.
+        assert abs(emb.shape[-1] - expected_length) < 1, (
+            emb.shape[-1],
+            expected_length,
+        )
+
+        q_res = self.quantizer(emb, self.frame_rate)
+        emb = q_res.x
+        emb = self._to_encoder_framerate(emb)
+        if self.decoder_transformer is not None:
+            (emb,) = self.decoder_transformer(emb)
+
+        with self._context_for_encoder_decoder:
+            out = self.decoder(emb)
+
+        # remove extra padding added by the encoder and decoder
+        assert out.shape[-1] >= length, (out.shape[-1], length)
+        out = out[..., :length]
+
+        q_res.x = out
+        q_res.metrics.update(extra_metrics)
+        return q_res
+
+    def _encode_to_unquantized_latent(self, x: torch.Tensor) -> torch.Tensor:
+        """Projects a batch of waveforms to unquantized latent space.
+
+        Args:
+            x (torch.Tensor): Float tensor of shape [B, C, T].
+
+        Returns:
+            Unquantized embeddings.
+        """
+        assert (
+            x.dim() == 3
+        ), f"CompressionModel._encode_to_unquantized_latent expects audio of shape [B, C, T] but got {x.shape}"
+        state = self._streaming_state
+        with self._context_for_encoder_decoder:
+            emb = self.encoder(x)
+        if self.encoder_transformer is not None:
+            if state is None:
+                (emb,) = self.encoder_transformer(emb)
+            else:
+                assert state.graphed_tr_enc is not None
+                (emb,) = state.graphed_tr_enc(emb)
+        emb = self._to_framerate(emb)
+        return emb
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """Encode the given input tensor to quantized representation.
+
+        Args:
+            x (torch.Tensor): Float tensor of shape [B, C, T]
+
+        Returns:
+            codes (torch.Tensor): an int tensor of shape [B, K, T]
+                with K the number of codebooks used and T the timestep.
+        """
+        emb = self._encode_to_unquantized_latent(x)
+        codes = self.quantizer.encode(emb)
+        return codes
+
+    def encode_to_latent(self, x: torch.Tensor, quantize: bool = True) -> torch.Tensor:
+        """Projects a batch of waveforms to latent space.
+
+        Args:
+            x (torch.Tensor): Float tensor of shape [B, C, T].
+
+        Returns:
+            Embeddings, either quantized or not.
+        """
+        emb = self._encode_to_unquantized_latent(x)
+        if not quantize:
+            return emb
+        else:
+            codes = self.quantizer.encode(emb)
+            return self.decode_latent(codes)
+
+    def decode(self, codes: torch.Tensor):
+        """Decode the given codes to a reconstructed representation.
+
+        Args:
+            codes (torch.Tensor): Int tensor of shape [B, K, T]
+
+        Returns:
+            out (torch.Tensor): Float tensor of shape [B, C, T], the reconstructed audio.
+        """
+        state = self._streaming_state
+        emb = self.decode_latent(codes)
+        emb = self._to_encoder_framerate(emb)
+        if self.decoder_transformer is not None:
+            if state is None:
+                (emb,) = self.decoder_transformer(emb)
+            else:
+                assert state.graphed_tr_dec is not None
+                (emb,) = state.graphed_tr_dec(emb)
+        with self._context_for_encoder_decoder:
+            out = self.decoder(emb)
+        # out contains extra padding added by the encoder and decoder
+        return out
+
+    def decode_latent(self, codes: torch.Tensor) -> torch.Tensor:
+        """Decode from the discrete codes to continuous latent space."""
+        return self.quantizer.decode(codes)
+
+
+class WrapperCompressionModel(CompressionModel[State]):
+    """Base API for CompressionModel wrappers that do not depend on external frameworks."""
+
+    def __init__(self, model: CompressionModel):
+        super().__init__()
+        self.model = model
+
+    def forward(self, x: torch.Tensor) -> QuantizedResult:
+        return self.model.forward(x)
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        return self.model.encode(x)
+
+    def decode(self, codes: torch.Tensor) -> torch.Tensor:
+        return self.model.decode(codes)
+
+    def decode_latent(self, codes: torch.Tensor) -> torch.Tensor:
+        return self.model.decode_latent(codes)
+
+    def set_num_codebooks(self, n: int):
+        self.model.set_num_codebooks(n)
+
+    @property
+    def quantizer(self):
+        return self.model.quantizer
+
+    @property
+    def channels(self) -> int:
+        return self.model.channels
+
+    @property
+    def frame_rate(self) -> float:
+        return self.model.frame_rate
+
+    @property
+    def sample_rate(self) -> int:
+        return self.model.sample_rate
+
+    @property
+    def cardinality(self) -> int:
+        return self.model.cardinality
+
+    @property
+    def num_codebooks(self) -> int:
+        return self.model.num_codebooks
+
+    @property
+    def total_codebooks(self) -> int:
+        return self.model.total_codebooks

moshi/models/lm.py

@@ -0,0 +1,1178 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: MIT
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+
+# Copyright (c) Kyutai, all rights reserved.
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from functools import partial
+from os.path import splitext
+import logging
+import numpy as np
+import sys
+from typing import Optional, Union, List, Tuple, Callable, Iterator
+import sphn
+import torch
+from tqdm.auto import tqdm
+
+from ..utils.sampling import sample_token
+from ..utils.compile import CUDAGraphed
+from ..modules.streaming import StreamingStateDict, StreamingContainer, StreamingModule, load_streaming_state
+from ..modules.transformer import (
+    StreamingTransformer,
+    create_norm_fn,
+)
+
+logger = logging.getLogger(__name__)
+
+AUDIO_TOKENS_PER_STREAM = 8
+FRAME_RATE_HZ = 12.5
+SILENCE_TOKENS = np.array([948, 243, 1178, 546, 1736, 1030, 1978, 2008], dtype=np.int64)
+SINE_TOKENS    = np.array([430, 1268, 381, 1611, 1095, 1495, 56, 472], dtype=np.int64)
+
+
+@dataclass
+class LMOutput:
+    # The logits are already re-aligned with the input codes
+    # hence no extra shift is required, e.g. when computing CE
+    logits: torch.Tensor  # [B, K, T, card]
+    mask: torch.Tensor  # [B, K, T]
+    text_logits: torch.Tensor  # [B, 1, T, text_card]
+    text_mask: torch.Tensor  # [B, 1, T]
+
+
+def _delay_sequence(delays: List[int], tensor: torch.Tensor, padding: torch.Tensor) -> torch.Tensor:
+    B, K, T = tensor.shape
+    assert len(delays) == K, (len(delays), K)
+    outs = []
+
+    for k, delay in enumerate(delays):
+        assert delay >= 0
+        line = tensor[:, k].roll(delay, dims=1)
+        if delay > 0:
+            line[:, :delay] = padding[:, k]
+        outs.append(line)
+    return torch.stack(outs, dim=1)
+
+
+def _undelay_sequence(delays: List[int], tensor: torch.Tensor,
+                      fill_value: Union[int, float] = float('NaN')) -> Tuple[torch.Tensor, torch.Tensor]:
+    B, K, T, *_ = tensor.shape
+    assert len(delays) == K
+    mask = torch.ones(B, K, T, dtype=torch.bool, device=tensor.device)
+    outs = []
+    if all([delay == 0 for delay in delays]):
+        return tensor, mask
+    for k, delay in enumerate(delays):
+        assert delay >= 0
+        line = tensor[:, k].roll(-delay, dims=1)
+        if delay > 0:
+            line[:, -delay:] = fill_value
+            mask[:, k, -delay:] = 0
+        outs.append(line)
+    return torch.stack(outs, dim=1), mask
+
+
+def create_sinewave(duration: float, sample_rate: int) -> np.ndarray:
+    """Return a 440 Hz 'silent' sinewave of the given duration."""
+    t = np.linspace(0.0, duration, int(sample_rate * duration), endpoint=False)
+    amplitude = 0.5
+    return amplitude * np.sin(2 * np.pi * 440.0 * t).astype(np.float32)
+
+
+def normalize_audio(wav: np.ndarray, sr: int, target_lufs: float) -> np.ndarray:
+    """Normalize **mono** audio to a target LUFS level."""
+    import pyloudnorm as pyln
+    # Ensure shape is (T,)
+    if wav.ndim == 2 and wav.shape[0] == 1:
+        wav = wav[0]
+
+    meter = pyln.Meter(sr)
+    loudness = meter.integrated_loudness(wav)
+    return pyln.normalize.loudness(wav, loudness, target_lufs)
+
+
+def load_audio(
+    filepath: str, sample_rate: int, 
+):
+    """Yields audio samples in intervals of sample_interval_size"""
+    sample_pcm, sample_sr = sphn.read(filepath)
+    sample_pcm = sphn.resample(
+        sample_pcm, src_sample_rate=sample_sr, dst_sample_rate=sample_rate
+    )  # shape: (C, T)
+    return sample_pcm
+
+def _iterate_audio(sample_pcm, sample_interval_size, max_len=sys.maxsize, pad=True):
+    cnt = 0
+    while sample_pcm.shape[-1] > 0 and cnt < max_len:
+        sample = sample_pcm[:, :sample_interval_size]
+        sample_pcm = sample_pcm[:, sample_interval_size:]
+        if sample_pcm.shape[-1] == 0 and pad:
+            sample = np.concatenate(
+                [
+                    sample,
+                    np.zeros(
+                        (
+                            sample.shape[0],
+                            sample_interval_size - sample.shape[-1],
+                        )
+                    ),
+                ],
+                axis=1,
+            )
+        cnt += 1
+        yield sample[0:1]  # shape: (1, T)
+
+
+def encode_from_sphn(mimi, samples, max_batch=sys.maxsize):
+    """
+    Takes an iterator of samples, batches them, encodes them;
+    and yields the encoded samples one sample at a time in the same order.
+    """
+    device = next(mimi.parameters()).device
+    current_batch = []
+    done_flag = False
+    # TO-DO: Fix the batching bug
+    max_batch = 1
+
+    while True:
+        try:
+            sample = next(samples)
+            tensor = torch.tensor(sample, dtype=torch.float32, device=device)
+            tensor = tensor.unsqueeze(0)  # shape: (1, C, T)                                                                                                      
+            current_batch.append(tensor)
+        except StopIteration:
+            done_flag = True
+
+        if (not done_flag) and len(current_batch) < max_batch:
+            continue
+        if not current_batch:
+            break
+
+        batch = torch.cat(current_batch, dim=0)  # shape: (B, C, T)
+        encoded = mimi.encode(batch)  # shape: (B, K, F)
+        separated = torch.unbind(encoded, dim=0)  # shape: (K, F)
+        reshaped = [x.unsqueeze(0) for x in separated]  # shape: (1, K, F)
+        detached = [x.detach().clone() for x in reshaped]
+
+        current_batch = []
+        yield from detached  # shape: (1, K, F)
+
+        if done_flag:
+            break
+
+
+class ScaledEmbedding(torch.nn.Embedding):
+    """Boost learning rate for embeddings (with `scale`).
+
+    Args:
+        norm (bool): if True, uses a layer norm after the embedding.
+        zero_idx (int): special value indicating that the output should be exactly 0.
+    """
+
+    def __init__(self, *args, norm: bool = False, zero_idx: int = -1, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.norm = None
+        if norm:
+            self.norm = create_norm_fn("layer_norm", self.embedding_dim)
+        assert zero_idx < 0, "Please use negative values for the zero_idx."
+        self.zero_idx = zero_idx
+
+    def forward(self, input, *args, **kwargs):
+        is_zero = input == self.zero_idx
+        zero = torch.zeros(1, dtype=input.dtype, device=input.device)
+        input = input.clamp(min=0)
+        y = super().forward(input, *args, **kwargs)
+        if self.norm is not None:
+            y = self.norm(y)
+        y = torch.where(is_zero[..., None], zero, y)
+        return y
+
+
+class LMModel(StreamingContainer):
+    """Transformer-based language model on multiple streams of codes.
+
+    Args:
+        n_q (int): Number of parallel streams to model as input.
+        dep_q (int): Number of parallel streams to model in the depformer.
+        card (int): Cardinality, vocabulary size.
+        text_card (int): Cardinality of the text vocabulary.
+        dim (int): Dimension of the transformer encoder.
+        num_heads (int): Number of heads for the transformer encoder.
+        hidden_scale (int): Scale for hidden feed forward dimension of the transformer encoder.
+        norm (str): Normalization method.
+        norm_emb (bool): Whether to normalize embeddings.
+        bias_proj (bool): Use bias for output projections.
+        depformer_*: params used for the Depformer Transformer, all the other will be shared.
+        depformer_multi_linear (bool): if True, uses one linear layer per codebook to project the
+            output of the main transformer to the Depformer latent space.
+        depformer_dim_feedforward (int| list[int]| None): If None, defaults to hidden_scale * depformer_dim.
+        existing_text_padding_id (bool): if True, will use a different token for the initial text token, and
+            the text padding token.
+        same_initial (bool): if True, uses the same initial tokens for both text and audio mode.
+        **kwargs: Additional parameters for the transformer encoder.
+    """
+
+    def __init__(
+        self,
+        delays: List[int] = [0],
+        n_q: int = 8,
+        dep_q: int = 8,
+        card: int = 1024,
+        text_card: int = 32000,
+        dim: int = 128,
+        num_heads: int = 8,
+        hidden_scale: int = 4,
+        norm: str = "layer_norm",
+        norm_emb: bool = False,
+        bias_proj: bool = False,
+        depformer_dim: int = 256,
+        depformer_dim_feedforward: int | list[int] | None = None,
+        depformer_multi_linear: bool = False,
+        depformer_weights_per_step: bool = False,
+        depformer_weights_per_step_schedule: list[int] | None = None,
+        depformer_pos_emb: str = "sin",
+        existing_text_padding_id: Optional[int] = None,
+        context: Optional[int] = None,
+        device=None,
+        dtype=None,
+        **kwargs,
+    ):
+        super().__init__()
+        self.n_q = n_q
+        self.dep_q = dep_q
+        self.card = card
+        self.text_card = text_card
+        assert len(delays) == self.num_codebooks, "unexpected number of delays"
+        self.delays = delays
+        self.dim = dim
+        self.existing_text_padding_id = existing_text_padding_id
+        self.context = context
+        self.depformer_weights_per_step_schedule = depformer_weights_per_step_schedule
+        if depformer_weights_per_step_schedule is not None:
+            assert len(depformer_weights_per_step_schedule) == dep_q
+        kwargs["context"] = context
+        EmbeddingFactory = partial(
+            ScaledEmbedding,
+            norm=norm_emb,
+            device=device,
+            dtype=dtype,
+            zero_idx=self.zero_token_id,
+        )
+        self.EmbeddingFactory = EmbeddingFactory
+        self.emb = torch.nn.ModuleList(
+            [EmbeddingFactory(self.card + 1, dim) for _ in range(n_q)]
+        )
+        # Text card + padding token (if not in the original tokenizer)
+        extra_text = self.existing_text_padding_id is None
+        # Unlike for audio, here we authorize the model to output the special token.
+        self.text_emb = EmbeddingFactory(text_card + 1, dim)
+        self.text_linear = torch.nn.Linear(dim, text_card + extra_text, bias=bias_proj)
+        depformer_prefix = "depformer_"
+        main_kwargs = {
+            k: v for k, v in kwargs.items() if not k.startswith(depformer_prefix)
+        }
+        self.transformer = StreamingTransformer(
+            d_model=dim,
+            num_heads=num_heads,
+            dim_feedforward=int(hidden_scale * dim),
+            norm=norm,
+            device=device,
+            dtype=dtype,
+            **main_kwargs,
+        )
+        self.out_norm = create_norm_fn(norm, dim)
+        self.depformer_multi_linear = depformer_multi_linear
+        kwargs_dep = main_kwargs.copy()
+        kwargs_dep.update(
+            {
+                k.removeprefix(depformer_prefix): v
+                for k, v in kwargs.items()
+                if k.startswith(depformer_prefix)
+            }
+        )
+        kwargs_dep["positional_embedding"] = depformer_pos_emb
+        kwargs_dep["context"] = None
+        if depformer_weights_per_step:
+            kwargs_dep["weights_per_step"] = dep_q
+        if depformer_multi_linear:
+            # One linear layer per codebook to project different informations from the main model.
+            self.depformer_in = torch.nn.ModuleList(
+                [torch.nn.Linear(dim, depformer_dim, bias=False) for _ in range(dep_q)]
+            )
+        else:
+            self.depformer_in = torch.nn.ModuleList(
+                [torch.nn.Linear(dim, depformer_dim, bias=False)]
+            )
+        # Only using up to dep_q - 1 because the last codebook is never an input to Depformer.
+        self.depformer_emb = torch.nn.ModuleList(
+            [EmbeddingFactory(self.card + 1, depformer_dim) for _ in range(dep_q - 1)]
+        )
+        self.depformer_text_emb = EmbeddingFactory(text_card + 1, depformer_dim)
+        if depformer_dim_feedforward is None:
+            depformer_dim_feedforward = int(hidden_scale * depformer_dim)
+        self.depformer = StreamingTransformer(
+            d_model=depformer_dim,
+            dim_feedforward=depformer_dim_feedforward,
+            norm=norm,
+            device=device,
+            dtype=dtype,
+            **kwargs_dep,
+        )
+        self.depformer.set_streaming_propagate(False)
+        dim = depformer_dim  # we will directly apply the next linears to the output of the Depformer.
+
+        self.linears = torch.nn.ModuleList(
+            [torch.nn.Linear(dim, self.card, bias=bias_proj) for _ in range(dep_q)]
+        )
+
+    @property
+    def initial_token_id(self) -> int:
+        """Token id for the start of sequence (audio)."""
+        return self.card
+
+    @property
+    def text_initial_token_id(self) -> int:
+        """Token id for the start of sequence (text)."""
+        return self.text_card
+
+    @property
+    def text_padding_token_id(self) -> int:
+        """Token id for text padding."""
+        if self.existing_text_padding_id is None:
+            return self.text_card
+        else:
+            return self.existing_text_padding_id
+
+    @property
+    def end_of_text_padding_id(self) -> int:
+        """Token id for optionally marking the last padding step for a word."""
+        return 0
+
+    @property
+    def zero_token_id(self) -> int:
+        """Special value in the input tokens, indicating that no sampling should
+        happen for that value, and no input should be given to the model."""
+        return -1
+
+    @property
+    def ungenerated_token_id(self) -> int:
+        """Special value that can be provided in the prompt to indicate that this specific
+        value should be predicted and sampled. This allows for partial teacher forcing, by generating
+        one modality, with the other one fixed.
+        """
+        return -2
+
+    @property
+    def device(self):
+        first_param = next(iter(self.parameters()))
+        return first_param.device
+
+    @property
+    def num_codebooks(self) -> int:
+        return self.n_q + 1
+
+    @property
+    def num_audio_codebooks(self) -> int:
+        return self.n_q
+
+    @property
+    def audio_offset(self) -> int:
+        return 1
+
+    def _get_initial_token(self) -> torch.Tensor:
+        # Returns the initial token that will be fed to the model to predict the very first timestep.
+        # The output shape will be [B, K, 1].
+        device = next(iter(self.parameters())).device
+        zero = torch.full(
+            [1, 1, 1], self.zero_token_id, device=device, dtype=torch.long
+        )
+        special = torch.full_like(zero, self.initial_token_id)
+
+        text_special = torch.full_like(zero, self.text_initial_token_id)
+        audio_token = special
+        text_token = text_special
+        audio_token = audio_token.expand(-1, self.num_audio_codebooks, -1)
+        token = torch.cat([text_token, audio_token], dim=1)
+        return token
+    
+    def embed_codes(self, sequence: torch.Tensor) -> torch.Tensor:
+        B, K, S = sequence.shape
+        assert (
+            K == self.num_codebooks
+        ), f"Sequence shape {sequence.shape} must match the number of codebooks."
+        input_sequence = sequence
+        input_ = None
+        for cb_index in range(self.num_audio_codebooks):
+            audio_emb = self.emb[cb_index](
+                input_sequence[:, cb_index + self.audio_offset]
+            )
+            input_ = audio_emb if input_ is None else input_ + audio_emb
+        text_emb = self.text_emb(input_sequence[:, 0])
+        input_ = text_emb if input_ is None else input_ + text_emb
+        return input_
+
+    def forward_codes(
+        self,
+        sequence: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return self.forward_embeddings(self.embed_codes(sequence))
+    
+    def forward_embeddings(self, input: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        # print("EMBED:", input[0, 0, :10].float().cpu().tolist()) # DEBUG
+        transformer_out = self.transformer(input)
+        if self.out_norm:
+            transformer_out = self.out_norm(transformer_out)
+        assert isinstance(transformer_out, torch.Tensor)
+        text_logits = self.text_linear(transformer_out)
+        text_logits = text_logits[:, None]
+        return transformer_out, text_logits
+
+    def forward_depformer(
+        self,
+        depformer_cb_index: int,
+        sequence: torch.Tensor,
+        transformer_out: torch.Tensor,
+    ) -> torch.Tensor:
+        B, K, S = sequence.shape
+        assert (
+            K == 1
+        ), f"Codebooks for Depformer streaming should be passed 1 by 1, got {K}."
+        assert (
+            S == 1
+        ), f"Steps for Depformer streaming should be passed 1 by 1, got {S}."
+        assert (
+            transformer_out.shape[1] == 1
+        ), "Transformer out should be a for a single step."
+        last_token_input: Optional[torch.Tensor] = None
+        depformer_input = transformer_out
+        if self.depformer_multi_linear:
+            depformer_input = self.depformer_in[depformer_cb_index](depformer_input)
+        else:
+            depformer_input = self.depformer_in[0](depformer_input)
+        if depformer_cb_index == 0:
+            last_token_input = self.depformer_text_emb(sequence[:, 0])
+        else:
+            last_token_input = self.depformer_emb[depformer_cb_index - 1](
+                sequence[:, 0]
+            )
+        depformer_input = depformer_input + last_token_input
+        assert depformer_input.shape[1] == 1
+        # depformer_input is [B, 1, depformer_dim].
+        # The streaming state of the depformer ensures that the proper layer is run.
+        dep_output = self.depformer(depformer_input)
+        logits = self.linears[depformer_cb_index](dep_output)
+        logits = logits[:, None]
+        assert logits.dim() == 4, logits.shape  # [B, Ka, S, card]
+        return logits
+
+    def forward_depformer_training(
+        self,
+        sequence: torch.Tensor,
+        transformer_out: torch.Tensor,
+    ) -> torch.Tensor:
+        B, K, T = sequence.shape
+        Ka = self.dep_q
+        assert (
+            K == self.num_codebooks
+        ), f"Codebooks for Depformer training should be passed all at once, got {K,}."
+        depformer_inputs = []
+        for cb_index in range(Ka):
+            if self.depformer_multi_linear:
+                linear_index = cb_index
+                if self.depformer_weights_per_step_schedule is not None:
+                    linear_index = self.depformer_weights_per_step_schedule[cb_index]
+                transformer_in = self.depformer_in[linear_index](transformer_out)
+            else:
+                transformer_in = self.depformer_in[0](transformer_out)
+            if cb_index == 0:
+                token_in = self.depformer_text_emb(sequence[:, 0])
+            else:
+                token_in = self.depformer_emb[cb_index - 1](sequence[:, cb_index + self.audio_offset - 1])
+            depformer_inputs.append(token_in + transformer_in)
+        depformer_input = torch.stack(depformer_inputs, 2)
+        # depformer_input is [B, T, K, depformer_dim], reshaping to [B * T, K, D]
+        depformer_input = depformer_input.view(B * T, Ka, -1)
+        depformer_output = self.depformer(depformer_input)
+        all_logits = []
+        for cb_index in range(Ka):
+            logits = self.linears[cb_index](depformer_output[:, cb_index])
+            all_logits.append(logits.view(B, T, -1))
+        logits = torch.stack(all_logits, 1)
+        assert logits.dim() == 4, logits.shape  # [B, Ka, T, card]
+        return logits
+
+    def forward_train(self, codes: torch.Tensor):
+        B, K, T = codes.shape
+        # Delaying codes and removing the last time step that will never be an input.
+        initial = self._get_initial_token().expand(B, -1, -1)
+        delayed_codes = _delay_sequence(self.delays, codes, initial)
+        # Inserting the empty tokens for the first time step.
+        delayed_codes = torch.cat([initial, delayed_codes], dim=2)
+
+        # LLM Backbone
+        transformer_out, text_logits = self.forward_codes(delayed_codes[:, :, :-1])
+        logits = self.forward_depformer_training(delayed_codes[:, :, 1:], transformer_out)
+
+        # map back the logits on pattern sequence to logits on original codes: [B, K, S, card] -> [B, K, T, card]
+        # and provide the corresponding mask over invalid positions of tokens. We will with NaN values invalid positions
+        # to ensure they properly handled.
+        logits, logits_mask = _undelay_sequence(
+            self.delays[self.audio_offset:self.audio_offset + self.dep_q],
+            logits, fill_value=float('NaN'))
+        logits_mask &= (codes[:, self.audio_offset: self.audio_offset + self.dep_q] != self.zero_token_id)
+        text_logits, text_logits_mask = _undelay_sequence(self.delays[:1], text_logits, fill_value=float('NaN'))
+        text_logits_mask &= (codes[:, :1] != self.zero_token_id)
+        return LMOutput(logits, logits_mask, text_logits, text_logits_mask)
+
+
+@dataclass
+class _LMGenState:
+    cache: torch.Tensor
+    provided: torch.Tensor
+    initial: torch.Tensor
+    graphed_main: CUDAGraphed
+    graphed_embeddings: CUDAGraphed
+    graphed_depth: CUDAGraphed
+    offset: int = 0
+
+    def reset(self):
+        self.offset = 0
+        self.provided[:] = False
+
+
+@torch.no_grad()
+def create_loss_report(
+    state_cache: torch.Tensor,
+    lm_model: LMModel,
+    text_logits: torch.Tensor,
+    audio_logits: torch.Tensor,
+    target: torch.Tensor,
+    sampled_text_token: torch.Tensor,
+    sampled_audio_tokens: torch.Tensor,
+    target_position: int,
+) -> dict[str, torch.Tensor]:
+    report = {}
+    B = state_cache.shape[0]
+    # model_tokens is the sampled output from model_logits
+    model_tokens = torch.zeros_like(state_cache[:, :, target_position])
+    model_tokens[:, 0] = sampled_text_token
+    model_tokens[:, 1 : lm_model.dep_q + 1] = sampled_audio_tokens
+
+    report.update(
+        {
+            "forced_tokens": torch.zeros((B, lm_model.dep_q + 1)),
+            "model_tokens": torch.zeros((B, lm_model.dep_q + 1)),
+            "ranks_of_forced": torch.zeros((B, lm_model.dep_q + 1)),
+            "losses": torch.zeros((B, lm_model.dep_q+1)),
+        }
+    )
+    report["model_tokens"] = model_tokens.clone()
+    report["forced_tokens"] = target.clone()
+
+    # Text Channel
+    text_logits = text_logits.squeeze(dim=1).squeeze(dim=1)
+    target = target[:, 0].squeeze(1).clone()
+
+    text_probs = torch.softmax(text_logits, dim=-1)
+    text_ranks = torch.argsort(text_probs, dim=-1, descending=True)
+    for b in range(B):
+        forced_token = target[b].item()
+        try:
+            rank = (text_ranks[b] == forced_token).nonzero().item()
+        except RuntimeError:
+            rank = lm_model.zero_token_id
+        report["ranks_of_forced"][b, 0] = rank
+
+    target[target == lm_model.text_initial_token_id] = -100
+    text_loss = torch.nn.functional.cross_entropy(
+        text_logits,
+        target,
+        ignore_index=-100,
+        )
+    report["losses"][:, 0] = text_loss
+
+    # Audio Channels
+    for k in range(lm_model.dep_q):
+        target = target[:, k+1].squeeze(1).clone()
+        channel_logits = audio_logits[:, k, :]
+
+        audio_probs = torch.softmax(channel_logits, dim=-1)
+        audio_ranks = torch.argsort(audio_probs, dim=-1, descending=True)
+        for b in range(B):
+            forced_token = target[b].item()
+            try:
+                rank = (audio_ranks[b] == forced_token).nonzero().item()
+            except RuntimeError:
+                rank = lm_model.zero_token_id
+            report["ranks_of_forced"][b, k + 1] = rank
+
+        target[target == lm_model.initial_token_id] = -100
+        audio_loss = torch.nn.functional.cross_entropy(
+            channel_logits,
+            target,
+            ignore_index=-100,
+        )
+        report["losses"][:, k + 1] = audio_loss
+    return report
+
+
+class LMGen(StreamingModule[_LMGenState]):
+    def __init__(
+        self,
+        lm_model: LMModel,
+        device: str | torch.device,
+        use_sampling: bool = True,
+        temp: float = 0.8,
+        temp_text: float = 0.7,
+        top_k: int = 250,
+        top_k_text: int = 25,
+        check: bool = False,
+        report_loss: bool = False,
+        return_logits: bool = False,
+        audio_silence_frame_cnt: int = 1,
+        text_prompt_tokens: Optional[list[int]] = None,
+        save_voice_prompt_embeddings: bool = False,
+        sample_rate: int = 32000,
+        frame_rate: int = FRAME_RATE_HZ,
+    ):
+        assert not lm_model.training, "generation shouldn't be used in training mode."
+        super().__init__()
+
+        self.lm_model = lm_model
+        self.use_sampling = use_sampling
+        self.temp = temp
+        self.temp_text = temp_text
+        self.top_k = top_k
+        self.top_k_text = top_k_text
+        self.text_prompt_tokens = text_prompt_tokens
+        self.audio_silence_frame_cnt = audio_silence_frame_cnt
+        self.voice_prompt = None
+        self.zero_text_code = 3
+        self._frame_rate = frame_rate
+        self._sample_rate = sample_rate
+        self._frame_size = int(self._sample_rate / self._frame_rate)
+        self._zero_frame = torch.zeros(1, 1, self._frame_size, device=device)
+        duration = self._frame_size / self._sample_rate
+        sine = create_sinewave(duration, self._sample_rate)
+        self._sine_frame = torch.tensor(sine, device=device).unsqueeze(0).unsqueeze(0)  # (1,1,T)
+        self.check = check
+        self.report_loss = report_loss
+        if report_loss:
+            return_logits = True
+        self.return_logits = return_logits
+        self.max_delay = max(
+            lm_model.delays
+        )  # with delays, we need to generate a few more time steps.
+        self.delays_cuda = torch.tensor(
+            lm_model.delays, device=lm_model.device, dtype=torch.long
+        )
+        self.save_voice_prompt_embeddings = save_voice_prompt_embeddings
+        self.voice_prompt_audio: Optional[torch.Tensor] = None
+        self.voice_prompt_cache: Optional[torch.Tensor] = None
+        self.voice_prompt_embeddings: Optional[torch.Tensor] = None
+        #self.voice_prompt_mimi_streaming_state: Optional[StreamingStateDict] = None
+
+    def _init_streaming_state(self, batch_size: int) -> _LMGenState:
+        lm_model = self.lm_model
+        initial = lm_model._get_initial_token()
+        cache = torch.full(
+            (batch_size, self.lm_model.num_codebooks, self.max_delay + 3),
+            lm_model.ungenerated_token_id,
+            device=lm_model.device,
+            dtype=torch.long,
+        )
+        provided = torch.full(
+            (batch_size, self.lm_model.num_codebooks, self.max_delay + 3),
+            False,
+            device=lm_model.device,
+            dtype=torch.bool
+        )
+
+        disable = lm_model.device.type != 'cuda'
+        # disable = True # DEBUG
+        graphed_main = CUDAGraphed(lm_model.forward_codes, disable=disable)
+        graphed_embeddings = CUDAGraphed(lm_model.forward_embeddings, disable=disable)
+        graphed_depth = CUDAGraphed(self.depformer_step, disable=disable)
+
+        return _LMGenState(cache, provided, initial, graphed_main, graphed_embeddings, graphed_depth)
+    
+    @torch.no_grad()
+    def prepare_step_input(self,
+                           input_tokens: torch.Tensor=None,
+                           moshi_tokens:torch.Tensor=None,
+                           text_token:torch.Tensor=None,
+                           ):
+        state = self._streaming_state
+        if state is None:
+            raise RuntimeError(
+                "You should wrap those calls with a `with lm_gen.streaming(): ...`."
+            )
+        lm_model = self.lm_model
+
+        # audio_tokens_per_stream = lm_model.dep_q//2
+        needed_tokens = lm_model.num_codebooks - AUDIO_TOKENS_PER_STREAM - 1
+        CT = state.cache.shape[2]
+
+        ####
+        # Fill Cache with provided tokens at state.offset (target) + delays
+
+        if input_tokens is not None:
+            assert input_tokens.dim() == 3, "Shape should be [B, K, T]."
+            B, Ki, S = input_tokens.shape
+            assert S == 1, "Only support being given steps one by one."
+            assert (
+                Ki == needed_tokens
+            ), f"We expect {needed_tokens} tokens from the user stream, got {Ki}."
+
+            for q_other in range(input_tokens.shape[1]):
+                k = AUDIO_TOKENS_PER_STREAM + 1 + q_other
+                delay = lm_model.delays[k]
+                write_position = (state.offset + delay) % CT
+                state.cache[:, k, write_position : write_position + 1] = input_tokens[:, q_other]
+                state.provided[:, k, write_position : write_position + 1] = True
+
+        if moshi_tokens is not None:
+            assert moshi_tokens.dim() == 3, "Shape should be [B, K, T]."
+            B, Ki, S = moshi_tokens.shape
+            assert S == 1, "Only support being given steps one by one."
+            assert (
+                Ki == needed_tokens
+            ), f"We expect {needed_tokens} tokens from the moshi stream, got {Ki}."
+
+            for q_moshi in range(moshi_tokens.shape[1]):
+                k = 1 + q_moshi
+                delay = lm_model.delays[k]
+                write_position = (state.offset + delay) % CT
+                state.cache[:, k, write_position : write_position + 1] = moshi_tokens[:, q_moshi]
+                state.provided[:, k, write_position : write_position + 1] = True
+
+        if text_token is not None:
+            write_position = (state.offset + lm_model.delays[0]) % CT
+            state.cache[:, 0, write_position] = text_token
+            state.provided[:, 0, write_position] = True
+
+        for k, delay in enumerate(lm_model.delays):
+            # Only for the very beginning, we extend the initial token for the acoustic
+            # token that are delayed, and thus have no good value to take.
+            if state.offset <= delay:
+                state.cache[:, k, state.offset % CT] = state.initial[:, k, 0]
+                state.provided[:, k, state.offset % CT] = True
+
+        ####
+        # Perform inference at state.offset - 1 (model_input); forcing with tokens at state.offset (target) when provided
+
+        if state.offset == 0:
+            # We can't report loss or force depth tranformer tokens until we're at step 2
+            # And we need to initialize the delay-0 cache where it's not provided for step 2
+            state.cache[:, :, 0] = state.initial[:, :, 0] # torch.where(state.provided[:, :, 0], state.cache[:, :, 0], state.initial[:, :, 0])
+            state.offset += 1
+            return None
+
+        model_input_position = (state.offset-1) % CT
+        target_position = state.offset % CT
+        input_ = state.cache[:, :, model_input_position : model_input_position + 1]
+        target_ = state.cache[:, :, target_position : target_position + 1]
+        provided_ = state.provided[:, :, target_position : target_position + 1]
+
+        if self.check:
+            # Check that we are not feeding in any value that is not generated yet.
+            assert not (input_ == lm_model.ungenerated_token_id).any(), (
+                state.offset,
+                input_,
+            )
+            assert (input_[:, lm_model.audio_offset :] <= lm_model.card).all(), input_
+            assert (input_[:, :1] <= lm_model.text_card).all()
+        return input_, provided_, target_, model_input_position, target_position
+
+    @torch.no_grad()
+    def step(self, input_tokens: torch.Tensor=None, moshi_tokens:torch.Tensor=None, text_token:torch.Tensor=None,
+             return_embeddings: bool=False) \
+        -> torch.Tensor | tuple[torch.Tensor, torch.Tensor] | tuple[torch.Tensor, dict[str, torch.Tensor]]:
+        state = self._streaming_state
+        lm_model = self.lm_model
+        prepared_inputs = self.prepare_step_input(
+            input_tokens, moshi_tokens, text_token,
+        )
+        # print("INPUT:", None if input_tokens is None else input_tokens.squeeze().cpu().tolist()) # DEBUG
+        # print("MOSHI:", None if moshi_tokens is None else moshi_tokens.squeeze().cpu().tolist()) # DEBUG
+        if prepared_inputs is None:
+            return (None, None) if self.report_loss or self.return_logits else None
+        input_, provided_, target_, model_input_position, target_position = prepared_inputs
+        if self.check:
+            # Check that we are not feeding in any value that is not generated yet.
+            assert not (input_ == lm_model.ungenerated_token_id).any(), (
+                state.offset,
+                input_,
+            )
+            assert (input_[:, lm_model.audio_offset :] <= lm_model.card).all(), input_
+            assert (input_[:, :1] <= lm_model.text_card).all()
+        embeddings = None
+        if return_embeddings:
+            embeddings = self.lm_model.embed_codes(input_)
+        transformer_out, text_logits = state.graphed_main(input_)
+        output = self.process_transformer_output(
+            transformer_out,
+            text_logits,
+            provided_,
+            target_,
+            model_input_position,
+            target_position,
+        )
+        if return_embeddings:
+            return output, embeddings
+        return output
+    
+    @torch.no_grad()
+    def step_embeddings(self, embeddings: torch.Tensor):
+        state = self._streaming_state
+        lm_model = self.lm_model
+        needed_input_tokens = lm_model.num_codebooks - AUDIO_TOKENS_PER_STREAM - 1
+        _dummy_audio_token = lm_model._get_initial_token()
+        while True:
+            prepared_inputs = self.prepare_step_input(
+                input_tokens=_dummy_audio_token[:, 1:1+needed_input_tokens], moshi_tokens=_dummy_audio_token[:, 1+needed_input_tokens:], text_token=self.zero_text_code,
+            )
+            if prepared_inputs is not None:
+                break
+        _, provided_, target_, model_input_position, target_position = prepared_inputs
+        transformer_out, text_logits = state.graphed_embeddings(embeddings)
+        return self.process_transformer_output(
+            transformer_out,
+            text_logits,
+            provided_,
+            target_,
+            model_input_position,
+            target_position,
+        )
+
+    @torch.no_grad()
+    def process_transformer_output(self, transformer_out, text_logits, provided_, target_, model_input_position, target_position):
+        state = self._streaming_state
+        lm_model = self.lm_model
+
+        # Shape of text_logits should be [B, K_text=1, T=1, Card_text]
+        sampled_text_token = sample_token(
+            text_logits.float(),
+            self.use_sampling,
+            self.temp_text,
+            self.top_k_text,
+        )
+        assert sampled_text_token.dim() == 3, sampled_text_token.shape
+        assert sampled_text_token.shape[2] == 1
+        assert sampled_text_token.shape[1] == 1, "Only one text stream supported."
+        sampled_text_token = sampled_text_token[:, 0, 0]  # shape is [B]
+
+        next_text_token = torch.where(provided_[:, 0, 0], target_[:, 0, 0], sampled_text_token)
+
+        if self.return_logits:
+            sampled_audio_tokens, audio_logits = state.graphed_depth(next_text_token, transformer_out, target_[:,lm_model.audio_offset:,0], provided_[:,lm_model.audio_offset:,0]) # [B, K_audio, Card_audio]
+        else:
+            sampled_audio_tokens = state.graphed_depth(next_text_token, transformer_out, target_[:,lm_model.audio_offset:,0], provided_[:,lm_model.audio_offset:,0])
+
+        state.provided[:, :, model_input_position] = False
+        ####
+        # Fill cache with generated tokens at state.offset (where not provided)
+
+        state.cache[:, 0, target_position] = torch.where(
+            ~state.provided[:, 0, target_position],
+            sampled_text_token,
+            state.cache[:, 0, target_position],
+        )
+        state.cache[:, 1 : lm_model.dep_q + 1, target_position] = torch.where(
+            ~state.provided[:, 1 : lm_model.dep_q + 1, target_position],
+            sampled_audio_tokens,
+            state.cache[:, 1 : lm_model.dep_q + 1, target_position],
+        )
+
+        ####
+        # Calculate loss of model logits (based on state.offset - 1) compared to target (state.offset)
+
+        report = {}
+        if self.report_loss:
+            report = create_loss_report(
+                state_cache=state.cache,
+                lm_model=lm_model,
+                text_logits=text_logits,
+                audio_logits=audio_logits,
+                target=target_,
+                sampled_text_token=sampled_text_token,
+                sampled_audio_tokens=sampled_audio_tokens,
+                target_position=target_position,
+            )
+
+        ####
+        # Collect outputs for state.offset - max_delay
+
+        if state.offset <= self.max_delay:
+            state.offset += 1
+            if self.report_loss:
+                return None, report
+            if self.return_logits:
+                return None, None
+            else:
+                return None
+        
+        B = state.cache.shape[0]
+        CT = state.cache.shape[2]
+        gen_delays_cuda = self.delays_cuda[: lm_model.dep_q + 1]
+        index = (
+            ((state.offset - self.max_delay + gen_delays_cuda) % CT)
+            .view(1, -1, 1)
+            .expand(B, -1, 1)
+        )
+        out = state.cache.gather(dim=2, index=index)
+
+        state.offset += 1
+        if self.report_loss:
+            return out, report
+        elif self.return_logits and not self.report_loss:
+            return out, (text_logits.clone(), audio_logits.clone())
+        else:
+            return out
+
+    def load_voice_prompt(self, voice_prompt: str):
+        self.voice_prompt = voice_prompt
+        raw_audio = load_audio(
+            voice_prompt, self._sample_rate,
+        )  # shape: (1, T) for mono
+
+        # Normalize to -24 LUFS (mono-safe)
+        raw_audio = normalize_audio(raw_audio, self._sample_rate, -24.0)
+
+        # Keep shape (1, T) because your encoder expects channels-first
+        if raw_audio.ndim == 1:
+            raw_audio = raw_audio[None, :]
+
+        self.voice_prompt_audio = raw_audio
+        self.voice_prompt_cache: Optional[torch.Tensor] = None
+        self.voice_prompt_embeddings: Optional[torch.Tensor] = None
+
+    def load_voice_prompt_embeddings(self, path: str):
+        self.voice_prompt = path
+        state = torch.load(path)
+
+        self.voice_prompt_audio = None
+        self.voice_prompt_embeddings = state["embeddings"].to(self.lm_model.device)
+        self.voice_prompt_cache = state["cache"].to(self.lm_model.device)
+
+    def _encode_zero_frame(self) -> torch.Tensor:
+        return torch.as_tensor(
+            SILENCE_TOKENS,
+            dtype=torch.long,
+            device=self.lm_model.device,
+        ).view(1, 8, 1)
+
+    def _encode_sine_frame(self) -> torch.Tensor:
+        return torch.as_tensor(
+            SINE_TOKENS,
+            dtype=torch.long,
+            device=self.lm_model.device,
+        ).view(1, 8, 1)
+
+    def _encode_voice_prompt_frames(self, mimi):
+        return encode_from_sphn(
+            mimi,
+            _iterate_audio(
+                self.voice_prompt_audio,
+                sample_interval_size=self._frame_size,
+                pad=True,
+            ),
+            max_batch=1,
+        )
+
+    def _step_voice_prompt_frame(self,
+                                 voice_prompt_frame_tokens: torch.Tensor,
+                                 saved_embeddings: Optional[list[torch.Tensor]]=None,
+                                 ):
+        # Always use zero_text_code during voice prompt
+        out = self.step(
+            moshi_tokens=voice_prompt_frame_tokens,
+            text_token=self.zero_text_code,
+            input_tokens=self._encode_sine_frame(),
+            return_embeddings=self.save_voice_prompt_embeddings,
+        )
+        if out is not None and self.save_voice_prompt_embeddings:
+            _, embeddings = out
+            saved_embeddings.append(embeddings)
+
+    def _step_voice_prompt_core(self, mimi) -> Iterator[None]:
+        """Shared core for stepping through the voice prompt.
+
+        This generator yields at each *checkpoint* where the async wrapper may want to
+        consult `is_alive`. The core itself is intentionally unaware of connection state.
+        """
+        if self.voice_prompt_embeddings is not None:
+            # Replay stored voice prompt embeddings
+            for next_embed in self.voice_prompt_embeddings:
+                yield
+                self.step_embeddings(next_embed)
+
+            state = self._streaming_state
+            state.cache.copy_(self.voice_prompt_cache)
+            return
+
+        elif self.voice_prompt_audio is not None:
+            saved_embeddings = []
+            for voice_prompt_frame_tokens in self._encode_voice_prompt_frames(mimi):
+                yield
+                self._step_voice_prompt_frame(
+                    voice_prompt_frame_tokens,
+                    saved_embeddings
+                )
+            # One last checkpoint before any optional save (nice-to-have for async disconnect)
+            yield
+
+            if self.save_voice_prompt_embeddings:
+                # Offset int(self._streaming_state.offset) is not needed since calling step() for len(voice_prompt_frame_tokens)
+                # and calling step_embeddings() for len(voice_prompt_embeddings) will increment offset by the same amount
+                torch.save(
+                    {
+                        "embeddings": torch.stack(saved_embeddings, dim=0).detach().cpu(),
+                        "cache": self._streaming_state.cache
+                    },
+                    splitext(self.voice_prompt)[0] + ".pt",
+                )
+        print('Done loading voice prompt.')
+
+    def _step_voice_prompt(self, mimi):
+        # Sync path intentionally does not support `is_alive` / disconnect checks.
+        for _ in self._step_voice_prompt_core(mimi):
+            pass
+
+    async def _step_voice_prompt_async(self, mimi, is_alive: Optional[Callable]=None):
+        for _ in self._step_voice_prompt_core(mimi):
+            if is_alive is not None and not await is_alive():
+                break
+
+    def _step_audio_silence_core(self) -> Iterator[None]:
+        # For slots of silence (default 0.5s) after voice/text prompts
+        # (agent text, user audio, agent audio) : (PADs, silence, sine)
+        for _ in range(self.audio_silence_frame_cnt):
+            yield
+            self.step(
+                moshi_tokens=self._encode_zero_frame(),
+                text_token=self.zero_text_code,
+                input_tokens=self._encode_sine_frame(),
+            )
+        print('Done loading audio silence.')
+
+    def _step_audio_silence(self):
+        # Sync path intentionally does not support `is_alive` / disconnect checks.
+        for _ in self._step_audio_silence_core():
+            pass
+
+    async def _step_audio_silence_async(self, is_alive: Optional[Callable]=None):
+        for _ in self._step_audio_silence_core():
+            if is_alive is not None and not await is_alive():
+                break
+
+    def _step_text_prompt_core(self) -> Iterator[None]:
+        for text_prompt_token in self.text_prompt_tokens:
+            yield
+            self.step(
+                moshi_tokens=self._encode_zero_frame(),
+                text_token=text_prompt_token,
+                input_tokens=self._encode_sine_frame(),
+            )
+        print('Done loading text prompt.')
+
+
+    def _step_text_prompt(self):
+        # Sync path intentionally does not support `is_alive` / disconnect checks.
+        for _ in self._step_text_prompt_core():
+            pass
+
+    async def _step_text_prompt_async(self, is_alive: Optional[Callable]=None):
+        for _ in self._step_text_prompt_core():
+            if is_alive is not None and not await is_alive():
+                break
+
+    async def step_system_prompts_async(self, mimi, is_alive: Optional[Callable]=None):
+        await self._step_voice_prompt_async(mimi, is_alive)
+        await self._step_audio_silence_async(is_alive)
+        await self._step_text_prompt_async(is_alive)
+        await self._step_audio_silence_async(is_alive)
+
+    def step_system_prompts(self, mimi):
+        self._step_voice_prompt(mimi)
+        self._step_audio_silence()
+        self._step_text_prompt()
+        self._step_audio_silence()
+
+    def depformer_step(
+        self,
+        text_token: torch.Tensor,
+        transformer_out: torch.Tensor,
+        audio_tokens: torch.Tensor,
+        audio_provided: torch.Tensor
+    ) -> torch.Tensor:
+        (B,) = text_token.shape
+        prev_token = text_token
+        lm_model = self.lm_model
+        depformer_tokens: list[torch.Tensor] = []
+        depformer_logits: list[torch.Tensor] = []
+        assert not lm_model.depformer.is_streaming
+        with lm_model.depformer.streaming(B):
+            for cb_index in range(lm_model.dep_q):
+                input_ = prev_token[:, None, None]
+                logits = lm_model.forward_depformer(cb_index, input_, transformer_out)
+                if self.return_logits:
+                    assert logits.shape == (B, 1, 1, lm_model.card), logits.shape
+                    ret_logits = logits.squeeze(dim=1).squeeze(dim=1)
+                    assert ret_logits.shape == (B, lm_model.card), ret_logits.shape
+                    depformer_logits.append(ret_logits.float())
+                next_token = sample_token(
+                    logits.float(),
+                    self.use_sampling,
+                    self.temp,
+                    self.top_k,
+                )
+                assert next_token.shape == (B, 1, 1)
+                next_token = next_token[:, 0, 0]  # shape is B
+                prev_token = torch.where(
+                    audio_provided[:, cb_index],
+                    audio_tokens[:, cb_index],
+                    next_token,
+                )
+                depformer_tokens.append(next_token)
+
+        assert len(depformer_tokens) == lm_model.dep_q, (
+            len(depformer_tokens),
+            lm_model.dep_q,
+        )
+        tokens = torch.stack(depformer_tokens, dim=1)
+        assert tokens.shape == (B, lm_model.dep_q), tokens.shape
+        if self.return_logits:
+            all_logits = torch.stack(depformer_logits, dim=1)
+            assert all_logits.shape == (B, lm_model.dep_q, lm_model.card), all_logits.shape
+            return tokens, all_logits
+        else:
+            return tokens
+

moshi/models/loaders.py

@@ -0,0 +1,364 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: MIT
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+
+# Copyright (c) Kyutai, all rights reserved.
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+"""Retrieves the pretrained models for Moshi and Mimi."""
+from pathlib import Path
+import logging
+
+from safetensors.torch import load_model, load_file
+import torch
+
+logger = logging.getLogger(__name__)
+
+from .compression import MimiModel
+from .lm import LMModel
+from ..modules import SEANetEncoder, SEANetDecoder, transformer
+from ..quantization import SplitResidualVectorQuantizer
+
+SAMPLE_RATE = 24000
+FRAME_RATE = 12.5
+
+TEXT_TOKENIZER_NAME = 'tokenizer_spm_32k_3.model'
+MOSHI_NAME = 'model.safetensors'
+MIMI_NAME = 'tokenizer-e351c8d8-checkpoint125.safetensors'
+DEFAULT_REPO = 'nvidia/personaplex-7b-v1'
+
+
+_seanet_kwargs = {
+    "channels": 1,
+    "dimension": 512,
+    "causal": True,
+    "n_filters": 64,
+    "n_residual_layers": 1,
+    "activation": "ELU",
+    "compress": 2,
+    "dilation_base": 2,
+    "disable_norm_outer_blocks": 0,
+    "kernel_size": 7,
+    "residual_kernel_size": 3,
+    "last_kernel_size": 3,
+    # We train using weight_norm but then the weights are pre-processed for inference so
+    # that we can use a normal convolution.
+    "norm": "none",
+    "pad_mode": "constant",
+    "ratios": [8, 6, 5, 4],
+    "true_skip": True,
+}
+_quantizer_kwargs = {
+    "dimension": 256,
+    "n_q": 32,
+    "bins": 2048,
+    "input_dimension": _seanet_kwargs["dimension"],
+    "output_dimension": _seanet_kwargs["dimension"],
+}
+_transformer_kwargs = {
+    "d_model": _seanet_kwargs["dimension"],
+    "num_heads": 8,
+    "num_layers": 8,
+    "causal": True,
+    "layer_scale": 0.01,
+    "context": 250,
+    "conv_layout": True,
+    "max_period": 10000,
+    "gating": "none",
+    "norm": "layer_norm",
+    "positional_embedding": "rope",
+    "dim_feedforward": 2048,
+    "input_dimension": _seanet_kwargs["dimension"],
+    "output_dimensions": [_seanet_kwargs["dimension"]],
+}
+
+_lm_kwargs = {
+    "dim": 4096,
+    "text_card": 32000,
+    "existing_text_padding_id": 3,
+    "n_q": 16,
+    "dep_q": 8,
+    "card": _quantizer_kwargs["bins"],
+    "num_heads": 32,
+    "num_layers": 32,
+    "hidden_scale": 4.125,
+    "causal": True,
+    "layer_scale": None,
+    "context": 3000,
+    "max_period": 10000,
+    "gating": "silu",
+    "norm": "rms_norm_f32",
+    "positional_embedding": "rope",
+    "depformer_dim": 1024,
+    "depformer_dim_feedforward": int(4.125 * 1024),
+    "depformer_num_heads": 16,
+    "depformer_num_layers": 6,
+    "depformer_causal": True,
+    "depformer_layer_scale": None,
+    "depformer_multi_linear": True,
+    "depformer_context": 8,
+    "depformer_max_period": 10000,
+    "depformer_gating": "silu",
+    "depformer_pos_emb": "none",
+    "depformer_weights_per_step": True,
+    "delays": [0, 0, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1],
+}
+
+
+def _is_safetensors(path: Path | str) -> bool:
+    return Path(path).suffix in (".safetensors", ".sft", ".sfts")
+
+
+def get_mimi(filename: str | Path,
+             device: torch.device | str = 'cpu') -> MimiModel:
+    """Return a pretrained Mimi model."""
+    encoder = SEANetEncoder(**_seanet_kwargs)
+    decoder = SEANetDecoder(**_seanet_kwargs)
+    encoder_transformer = transformer.ProjectedTransformer(
+        device=device, **_transformer_kwargs
+    )
+    decoder_transformer = transformer.ProjectedTransformer(
+        device=device, **_transformer_kwargs
+    )
+    quantizer = SplitResidualVectorQuantizer(
+        **_quantizer_kwargs,
+    )
+    model = MimiModel(
+        encoder,
+        decoder,
+        quantizer,
+        channels=1,
+        sample_rate=SAMPLE_RATE,
+        frame_rate=FRAME_RATE,
+        encoder_frame_rate=SAMPLE_RATE / encoder.hop_length,
+        causal=True,
+        resample_method="conv",
+        encoder_transformer=encoder_transformer,
+        decoder_transformer=decoder_transformer,
+    ).to(device=device)
+    model.eval()
+    if _is_safetensors(filename):
+        load_model(model, filename)
+    else:
+        pkg = torch.load(filename, "cpu")
+        model.load_state_dict(pkg["model"])
+    model.set_num_codebooks(8)
+    return model
+
+
+def get_moshi_lm(
+    filename: str | Path | None,
+    copy_missing_weights: bool = True,
+    device: torch.device | str = "cpu",
+    dtype: torch.dtype = torch.bfloat16,
+    delays=None,
+    cpu_offload: bool = False,
+) -> LMModel:
+    """Return a pretrained Moshi LM model.
+
+    Args:
+        filename: Path to model weights.
+        copy_missing_weights: Whether to copy missing weights from existing layers.
+        device: Target device for the model.
+        dtype: Data type for model weights.
+        delays: Optional custom delays configuration.
+        cpu_offload: If True, offload model layers to CPU when GPU memory is
+                     insufficient. Uses accelerate's device_map="auto".
+    """
+    # Copy to avoid mutating a shared/global dict
+    lm_kwargs = dict(_lm_kwargs)
+    lm_kwargs["dep_q"] = 16
+    if delays is not None:
+        lm_kwargs["delays"] = delays
+
+    if cpu_offload and filename is not None:
+        return _get_moshi_lm_with_offload(
+            filename, copy_missing_weights, device, dtype, lm_kwargs
+        )
+
+    # Init with meta device to avoid init dummy memory
+    init_device = "meta" if filename is not None else device
+    model = LMModel(device=init_device, dtype=dtype, **lm_kwargs)
+    if filename is None:
+        model.to(device=device, dtype=dtype)
+        model.eval()
+        return model
+
+    filename = str(filename)
+
+    # Load state_dict
+    if filename.endswith(".safetensors"):
+        # safetensors does not support mps directly
+        dev = torch.device(device) if isinstance(device, str) else device
+        if dev.type == "mps":
+            state_dict = load_file(filename, device="cpu")
+        else:
+            state_dict = load_file(filename, device=dev.type)
+    else:
+        # torch checkpoint
+        with open(filename, "rb") as f:
+            state_dict = torch.load(f, map_location="cpu")
+    # Patch 1: expand depformer self_attn weights if needed
+    model_sd = model.state_dict()
+    for name, tensor in list(state_dict.items()):
+        if "depformer" in name and "self_attn" in name and name in model_sd:
+            if tensor.shape != model_sd[name].shape:
+                print("Expanding %s", name)
+                missing = (
+                    tensor
+                    if copy_missing_weights
+                    else model_sd[name][tensor.shape[0] :]
+                )
+                state_dict[name] = torch.concat([tensor, missing], dim=0)
+
+    # Patch 2: fill missing keys by copying 0..7 -> 8..15 for certain groups
+    if copy_missing_weights:
+        to_replace = ["gating", "linears", "depformer_in", "depformer_emb"]
+        for name in model_sd.keys():
+            if name in state_dict:
+                continue
+            replaced = False
+            for old, new in zip(range(8), range(8, 16)):
+                for rep in to_replace:
+                    needle = f"{rep}.{new}."
+                    if needle in name:
+                        src = name.replace(needle, f"{rep}.{old}.")
+                        if src in state_dict:
+                            print("Replacing %s <- %s", name, src)
+                            state_dict[name] = state_dict[src]
+                            replaced = True
+                        break
+                if replaced:
+                    break
+            if not replaced:
+                print("Missing %s", name)
+
+    # Assign weights to target device
+    dev = torch.device(device) if isinstance(device, str) else device
+    for key in state_dict:
+        state_dict[key] = state_dict[key].to(device=dev, dtype=dtype)
+    
+    model.load_state_dict(state_dict, strict=False, assign=True)
+    model.eval()
+    return model.to(device=device, dtype=dtype)
+
+
+def _get_moshi_lm_with_offload(
+    filename: str | Path,
+    copy_missing_weights: bool,
+    device: torch.device | str,
+    dtype: torch.dtype,
+    lm_kwargs: dict,
+) -> LMModel:
+    """Load Moshi LM with CPU offloading using accelerate.
+
+    This function distributes model layers across GPU and CPU based on
+    available GPU memory. Layers that don't fit on GPU are kept on CPU
+    and moved to GPU only during forward pass.
+    """
+    try:
+        from accelerate import infer_auto_device_map, dispatch_model
+    except ImportError:
+        raise ImportError(
+            "CPU offloading requires the 'accelerate' package. "
+            "Install it with: pip install accelerate"
+        )
+
+    filename = str(filename)
+    logger.info("Loading model with CPU offloading enabled")
+
+    # First, create model on CPU to get the architecture
+    model = LMModel(device="cpu", dtype=dtype, **lm_kwargs)
+
+    # Load state_dict to CPU
+    if filename.endswith(".safetensors"):
+        state_dict = load_file(filename, device="cpu")
+    else:
+        with open(filename, "rb") as f:
+            state_dict = torch.load(f, map_location="cpu")
+
+    # Apply weight patches (same as non-offload path)
+    model_sd = model.state_dict()
+    for name, tensor in list(state_dict.items()):
+        if "depformer" in name and "self_attn" in name and name in model_sd:
+            if tensor.shape != model_sd[name].shape:
+                logger.info(f"Expanding {name}")
+                missing = (
+                    tensor
+                    if copy_missing_weights
+                    else model_sd[name][tensor.shape[0]:]
+                )
+                state_dict[name] = torch.concat([tensor, missing], dim=0)
+
+    if copy_missing_weights:
+        to_replace = ["gating", "linears", "depformer_in", "depformer_emb"]
+        for name in model_sd.keys():
+            if name in state_dict:
+                continue
+            replaced = False
+            for old, new in zip(range(8), range(8, 16)):
+                for rep in to_replace:
+                    needle = f"{rep}.{new}."
+                    if needle in name:
+                        src = name.replace(needle, f"{rep}.{old}.")
+                        if src in state_dict:
+                            logger.info(f"Replacing {name} <- {src}")
+                            state_dict[name] = state_dict[src]
+                            replaced = True
+                        break
+                if replaced:
+                    break
+            if not replaced:
+                logger.warning(f"Missing {name}")
+
+    model.load_state_dict(state_dict, strict=False, assign=True)
+
+    # Determine target device
+    dev = torch.device(device) if isinstance(device, str) else device
+
+    if dev.type != "cuda":
+        # If not using CUDA, just move to the target device without offloading
+        logger.info(f"CPU offload requested but device is {dev}, skipping offload")
+        model.to(dev)
+        model.eval()
+        return model
+
+    # Infer device map based on available GPU memory
+    device_map = infer_auto_device_map(
+        model,
+        max_memory=None,  # Let accelerate auto-detect available memory
+        no_split_module_classes=["StreamingTransformerLayer"],
+        dtype=dtype,
+    )
+
+    # Log the device distribution
+    gpu_layers = sum(1 for v in device_map.values() if v == 0 or v == "cuda:0")
+    cpu_layers = sum(1 for v in device_map.values() if v == "cpu")
+    logger.info(f"Device map: {gpu_layers} modules on GPU, {cpu_layers} modules on CPU")
+
+    # Dispatch model across devices
+    model = dispatch_model(
+        model,
+        device_map=device_map,
+        offload_dir="offload_weights",  # Directory for disk offload if needed
+    )
+
+    model.eval()
+    return model

moshi/modules/__init__.py

@@ -0,0 +1,23 @@
+# Copyright (c) Kyutai, all rights reserved.
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+"""Modules used for building the models."""
+
+# flake8: noqa
+from .conv import (
+    NormConv1d,
+    NormConvTranspose1d,
+    StreamingConv1d,
+    StreamingConvTranspose1d,
+    pad_for_conv1d,
+    pad1d,
+    unpad1d,
+)
+from .seanet import SEANetEncoder, SEANetDecoder
+from .transformer import StreamingTransformer

moshi/modules/conv.py

@@ -0,0 +1,350 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: MIT
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+
+# Copyright (c) Kyutai, all rights reserved.
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from dataclasses import dataclass
+import math
+import typing as tp
+import warnings
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch.nn.utils import weight_norm
+
+from .streaming import RawStreamingConv1d, RawStreamingConvTranspose1d, StreamingModule
+
+
+CONV_NORMALIZATIONS = frozenset(["none", "weight_norm"])
+
+
+class TransposedLayerNorm(nn.Module):
+    """LayerNorm for [B, C, T] inputs."""
+
+    def __init__(self, **kwargs):
+        super().__init__()
+        self.layer_norm = nn.LayerNorm(**kwargs)
+
+    def forward(self, x):
+        x = x.transpose(1, 2)
+        x = self.layer_norm(x)
+        return x.transpose(1, 2)
+
+
+def apply_parametrization_norm(module: nn.Module, norm: str = "none"):
+    assert norm in CONV_NORMALIZATIONS
+    if norm == "weight_norm":
+        return weight_norm(module)
+    else:
+        # We already check was in CONV_NORMALIZATION, so any other choice
+        # doesn't need reparametrization.
+        return module
+
+
+def get_extra_padding_for_conv1d(
+    x: torch.Tensor, kernel_size: int, stride: int, padding_total: int = 0
+) -> int:
+    """See `pad_for_conv1d`."""
+    length = x.shape[-1]
+    n_frames = (length - kernel_size + padding_total) / stride + 1
+    ideal_length = (math.ceil(n_frames) - 1) * stride + (kernel_size - padding_total)
+    return ideal_length - length
+
+
+def pad_for_conv1d(
+    x: torch.Tensor, kernel_size: int, stride: int, padding_total: int = 0
+):
+    """Pad for a convolution to make sure that the last window is full.
+    Extra padding is added at the end. This is required to ensure that we can rebuild
+    an output of the same length, as otherwise, even with padding, some time steps
+    might get removed.
+    For instance, with total padding = 4, kernel size = 4, stride = 2:
+        0 0 1 2 3 4 5 0 0   # (0s are padding)
+        1   2   3           # (output frames of a convolution, last 0 is never used)
+        0 0 1 2 3 4 5 0     # (output of tr. conv., but pos. 5 is going to get removed as padding)
+            1 2 3 4         # once you removed padding, we are missing one time step !
+    """
+    extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total)
+    return F.pad(x, (0, extra_padding))
+
+
+def pad1d(
+    x: torch.Tensor,
+    paddings: tp.Tuple[int, int],
+    mode: str = "constant",
+    value: float = 0.0,
+):
+    """Tiny wrapper around F.pad, just to allow for reflect padding on small input.
+    If this is the case, we insert extra 0 padding to the right before the reflection happen.
+    """
+    length = x.shape[-1]
+    padding_left, padding_right = paddings
+    assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)
+    if mode == "reflect":
+        max_pad = max(padding_left, padding_right)
+        extra_pad = 0
+        if length <= max_pad:
+            extra_pad = max_pad - length + 1
+            x = F.pad(x, (0, extra_pad))
+        padded = F.pad(x, paddings, mode, value)
+        end = padded.shape[-1] - extra_pad
+        return padded[..., :end]
+    else:
+        return F.pad(x, paddings, mode, value)
+
+
+def unpad1d(x: torch.Tensor, paddings: tp.Tuple[int, int]):
+    """Remove padding from x, handling properly zero padding. Only for 1d!"""
+    padding_left, padding_right = paddings
+    assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)
+    assert (padding_left + padding_right) <= x.shape[-1]
+    end = x.shape[-1] - padding_right
+    return x[..., padding_left:end]
+
+
+class NormConv1d(nn.Module):
+    """Wrapper around Conv1d and normalization applied to this conv
+    to provide a uniform interface across normalization approaches.
+    """
+
+    def __init__(
+        self,
+        *args,
+        causal: bool = False,
+        norm: str = "none",
+        norm_kwargs: tp.Dict[str, tp.Any] = {},
+        **kwargs,
+    ):
+        super().__init__()
+        self.conv = apply_parametrization_norm(
+            RawStreamingConv1d(*args, **kwargs), norm
+        )
+        self.norm_type = norm
+
+    def forward(self, x):
+        x = self.conv(x)
+        return x
+
+
+class NormConvTranspose1d(nn.Module):
+    """Wrapper around ConvTranspose1d and normalization applied to this conv
+    to provide a uniform interface across normalization approaches.
+    """
+
+    def __init__(
+        self,
+        *args,
+        causal: bool = False,
+        norm: str = "none",
+        norm_kwargs: tp.Dict[str, tp.Any] = {},
+        **kwargs,
+    ):
+        super().__init__()
+        self.convtr = apply_parametrization_norm(
+            RawStreamingConvTranspose1d(*args, **kwargs), norm
+        )
+        self.norm_type = norm
+
+    def forward(self, x):
+        x = self.convtr(x)
+        return x
+
+
+@dataclass
+class _StreamingConv1dState:
+    padding_to_add: int
+    original_padding_to_add: int
+
+    def reset(self):
+        self.padding_to_add = self.original_padding_to_add
+
+
+class StreamingConv1d(StreamingModule[_StreamingConv1dState]):
+    """Conv1d with some builtin handling of asymmetric or causal padding
+    and normalization.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int = 1,
+        dilation: int = 1,
+        groups: int = 1,
+        bias: bool = True,
+        causal: bool = False,
+        norm: str = "none",
+        norm_kwargs: tp.Dict[str, tp.Any] = {},
+        pad_mode: str = "reflect",
+    ):
+        super().__init__()
+        # warn user on unusual setup between dilation and stride
+        if stride > 1 and dilation > 1:
+            warnings.warn(
+                "StreamingConv1d has been initialized with stride > 1 and dilation > 1"
+                f" (kernel_size={kernel_size} stride={stride}, dilation={dilation})."
+            )
+        self.conv = NormConv1d(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride,
+            dilation=dilation,
+            groups=groups,
+            bias=bias,
+            causal=causal,
+            norm=norm,
+            norm_kwargs=norm_kwargs,
+        )
+        self.causal = causal
+        self.pad_mode = pad_mode
+
+    @property
+    def _stride(self) -> int:
+        return self.conv.conv.stride[0]
+
+    @property
+    def _kernel_size(self) -> int:
+        return self.conv.conv.kernel_size[0]
+
+    @property
+    def _effective_kernel_size(self) -> int:
+        dilation = self.conv.conv.dilation[0]
+        return (
+            self._kernel_size - 1
+        ) * dilation + 1  # effective kernel size with dilations
+
+    @property
+    def _padding_total(self) -> int:
+        return self._effective_kernel_size - self._stride
+
+    def _init_streaming_state(self, batch_size: int) -> _StreamingConv1dState:
+        assert self.causal, "streaming is only supported for causal convs"
+        return _StreamingConv1dState(self._padding_total, self._padding_total)
+
+    def forward(self, x):
+        B, C, T = x.shape
+        padding_total = self._padding_total
+        extra_padding = get_extra_padding_for_conv1d(
+            x, self._effective_kernel_size, self._stride, padding_total
+        )
+        state = self._streaming_state
+        if state is None:
+            if self.causal:
+                # Left padding for causal
+                x = pad1d(x, (padding_total, extra_padding), mode=self.pad_mode)
+            else:
+                # Asymmetric padding required for odd strides
+                padding_right = padding_total // 2
+                padding_left = padding_total - padding_right
+                x = pad1d(
+                    x, (padding_left, padding_right + extra_padding), mode=self.pad_mode
+                )
+        else:
+            if state.padding_to_add > 0 and x.shape[-1] > 0:
+                x = pad1d(x, (state.padding_to_add, 0), mode=self.pad_mode)
+                state.padding_to_add = 0
+        return self.conv(x)
+
+
+@dataclass
+class _StreamingConvTr1dState:
+    pass
+
+    def reset(self):
+        pass
+
+
+class StreamingConvTranspose1d(StreamingModule[_StreamingConvTr1dState]):
+    """ConvTranspose1d with some builtin handling of asymmetric or causal padding
+    and normalization.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int = 1,
+        groups: int = 1,
+        bias: bool = True,
+        causal: bool = False,
+        norm: str = "none",
+        trim_right_ratio: float = 1.0,
+        norm_kwargs: tp.Dict[str, tp.Any] = {},
+    ):
+        super().__init__()
+        self.convtr = NormConvTranspose1d(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride,
+            groups=groups,
+            bias=bias,
+            causal=causal,
+            norm=norm,
+            norm_kwargs=norm_kwargs,
+        )
+        self.causal = causal
+        self.trim_right_ratio = trim_right_ratio
+        assert (
+            self.causal or self.trim_right_ratio == 1.0
+        ), "`trim_right_ratio` != 1.0 only makes sense for causal convolutions"
+        assert self.trim_right_ratio >= 0.0 and self.trim_right_ratio <= 1.0
+
+    def _init_streaming_state(self, batch_size: int) -> _StreamingConvTr1dState:
+        assert self.causal, "streaming is only supported for causal convtrs"
+        return _StreamingConvTr1dState()
+
+    def forward(self, x):
+        kernel_size = self.convtr.convtr.kernel_size[0]
+        stride = self.convtr.convtr.stride[0]
+        padding_total = kernel_size - stride
+
+        y = self.convtr(x)
+
+        if not self.is_streaming:
+            # We will only trim fixed padding. Extra padding from `pad_for_conv1d` would be
+            # removed at the very end, when keeping only the right length for the output,
+            # as removing it here would require also passing the length at the matching layer
+            # in the encoder.
+            if self.causal:
+                # Trim the padding on the right according to the specified ratio
+                # if trim_right_ratio = 1.0, trim everything from right
+                padding_right = math.ceil(padding_total * self.trim_right_ratio)
+                padding_left = padding_total - padding_right
+                y = unpad1d(y, (padding_left, padding_right))
+            else:
+                # Asymmetric padding required for odd strides
+                padding_right = padding_total // 2
+                padding_left = padding_total - padding_right
+                y = unpad1d(y, (padding_left, padding_right))
+        return y

moshi/modules/gating.py

@@ -0,0 +1,103 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: MIT
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+
+# Copyright (c) Kyutai, all rights reserved.
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from ..utils.compile import torch_compile_lazy
+
+
+@torch_compile_lazy
+def gating_forward_kernel(
+    weight_in: torch.Tensor, weight_out: torch.Tensor, activation, x: torch.Tensor
+):
+    x = F.linear(x, weight_in)
+    B, T, _ = x.shape
+    x = x.view(B, T, 2, -1)
+    x = activation(x[..., 0, :]) * x[..., 1, :]
+    x = F.linear(x, weight_out)
+    return x
+
+
+class ActivationGating(nn.Module):
+    """
+    Gating FFN layer, using the given activation.
+    Args:
+        dim (int): dimension of the input and output of the transformer.
+        activation (any callable Tensor to Tensor): activation function to use.
+        **factory_kwargs: other kwargs passed to the linear layer, in particular device and dtype.
+    """
+
+    _fsdp_final = True
+
+    def __init__(self, dim: int, dim_feedforward: int, activation, **factory_kwargs):
+        super().__init__()
+        # We should have 8 d^2 param, instead we will have
+        # 2 * h * d + h * d = 3 h * d = 8 d^2
+        # so h = 8 d / 3 but following Hervé's advice we use 21 / 8 as an approx.
+        if dim_feedforward == 4 * dim:
+            hidden = (21 * dim) // 8
+        else:
+            hidden = (2 * dim_feedforward) // 3
+        self.linear_in = nn.Linear(dim, 2 * hidden, bias=False, **factory_kwargs)
+        self.linear_out = nn.Linear(hidden, dim, bias=False, **factory_kwargs)
+        self.activation = activation
+
+    def forward(self, x: torch.Tensor):
+        return gating_forward_kernel(
+            self.linear_in.weight, self.linear_out.weight, self.activation, x
+        )
+
+
+def _get_activation(name: str):
+    if name in ["sigmoid", "tanh", "relu"]:
+        return getattr(torch, name)
+    elif name in ["leaky_relu", "elu", "gelu", "silu", "mish", "softsign"]:
+        return getattr(torch.nn.functional, name)
+    elif name == "identity":
+        return torch.nn.Identity()
+    else:
+        raise ValueError(f"Unknown activation {name}")
+
+
+def _make_gating(
+    name: str, dim: int, dim_feedforward: int, **factory_kwargs
+) -> nn.Module:
+    return ActivationGating(
+        dim, dim_feedforward, _get_activation(name), **factory_kwargs
+    )
+
+
+def make_gating(
+    name: str, dim: int, dim_feedforward: int, **factory_kwargs
+) -> nn.Module:
+    gating = _make_gating(name, dim, dim_feedforward, **factory_kwargs)
+    max_params = 2 * dim * dim_feedforward
+    params = sum(p.numel() for p in gating.parameters())
+    assert (
+        params <= max_params
+    ), f"{name} gating has {params} params, max is {max_params}"
+    return gating

moshi/modules/resample.py

@@ -0,0 +1,140 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: MIT
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+
+# Copyright (c) Kyutai, all rights reserved.
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import typing as tp
+
+from einops import rearrange
+import torch
+from torch import nn
+
+from .conv import StreamingConv1d, StreamingConvTranspose1d
+
+
+class ConvDownsample1d(nn.Module):
+    """
+    Downsampling by some integer amount `stride` using convolutions
+    with a kernel size of twice the stride.
+    If `causal` is True, the output uses a causal convolution.
+    """
+
+    def __init__(
+        self,
+        stride: int,
+        dimension: tp.Optional[int] = None,
+        causal: bool = False,
+        learnt: bool = False,
+        channel_wise: bool = False,
+    ):
+        super().__init__()
+        self.learnt = learnt
+        self.channel_wise = channel_wise
+        groups = 1
+        if learnt:
+            assert dimension is not None, "Dimension required for learnt convolutions."
+            in_channels = dimension
+            out_channels = dimension
+            if channel_wise:
+                groups = dimension
+        else:
+            in_channels = 1
+            out_channels = 1
+
+        self.conv = StreamingConv1d(
+            in_channels,
+            out_channels,
+            kernel_size=2 * stride,
+            stride=stride,
+            causal=causal,
+            groups=groups,
+            bias=False,
+            pad_mode="replicate",
+        )
+        if not learnt:
+            actual_conv = self.conv.conv.conv
+            actual_conv.weight.requires_grad_(False)
+            actual_conv.weight.data.fill_(1.0 / (2 * stride))
+
+    def forward(self, x: torch.Tensor):
+        batch_size = len(x)
+        if not self.learnt:
+            x = rearrange(x, "b c t -> (b c) () t")
+        y = self.conv(x)
+        if not self.learnt:
+            y = rearrange(y, "(b c) () t -> b c t", b=batch_size)
+        return y
+
+
+class ConvTrUpsample1d(nn.Module):
+    """
+    Upsample by some integer amount `stride` using transposed convolutions.
+    """
+
+    def __init__(
+        self,
+        stride: int,
+        dimension: tp.Optional[int] = None,
+        causal: bool = False,
+        learnt: bool = False,
+        channel_wise: bool = False,
+    ):
+        super().__init__()
+        self.learnt = learnt
+        self.channel_wise = channel_wise
+        groups = 1
+        if learnt:
+            assert dimension is not None, "Dimension required for learnt convolutions."
+            in_channels = dimension
+            out_channels = dimension
+            if channel_wise:
+                groups = dimension
+        else:
+            in_channels = 1
+            out_channels = 1
+
+        self.convtr = StreamingConvTranspose1d(
+            in_channels,
+            out_channels,
+            kernel_size=2 * stride,
+            stride=stride,
+            causal=causal,
+            groups=groups,
+            bias=False,
+        )
+        if not learnt:
+            actual_convtr = self.convtr.convtr.convtr
+            actual_convtr.weight.requires_grad_(False)
+            actual_convtr.weight.data.fill_(1.0)
+
+    def forward(self, x: torch.Tensor):
+        batch_size = len(x)
+        if not self.learnt:
+            x = rearrange(x, "b c t -> (b c) () t")
+        y = self.convtr(x)
+        if not self.learnt:
+            x_for_normalization = torch.ones_like(x[:1])
+            normalization = self.convtr(x_for_normalization)
+            y = y / normalization
+            y = rearrange(y, "(b c) () t -> b c t", b=batch_size)
+        return y

moshi/modules/rope.py

@@ -0,0 +1,111 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: MIT
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+
+# Copyright (c) Kyutai, all rights reserved.
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from torch import nn
+import math
+import torch
+from ..utils.compile import torch_compile_lazy
+
+
+@torch_compile_lazy
+def apply_rope(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    offset: torch.Tensor,
+    max_period: float = 10_000,
+    time_before_heads: bool = False,
+):
+    """
+    Args:
+        q (torch.Tensor): queries, shape `[B, T, H, D]`.
+        k (torch.Tensor): keys, shape `[B, T, H, D]`.
+        offset (int): current offset, e.g. when streaming.
+        max_period (float): maximum period for the cos and sin.
+        time_before_heads (bool):  if True, expected [B, T, H, D], else [B, H, T ,D]
+    """
+
+    if time_before_heads:
+        B, T, H, D = q.shape
+    else:
+        B, H, T, D = q.shape
+    assert k.shape == q.shape
+    assert D > 0
+    assert D % 2 == 0
+    assert max_period > 0
+
+    ds = torch.arange(D // 2, device=q.device, dtype=torch.float32)
+    freqs = torch.exp(ds * (-math.log(max_period) * 2 / D))
+    ts = offset.float() + torch.arange(T, device=q.device, dtype=torch.float32)
+    if time_before_heads:
+        ts = ts.view(-1, 1, 1)
+    else:
+        ts = ts.view(1, -1, 1)
+
+    dims = q.shape[:-1]
+    q = q.view(*dims, D // 2, 2)
+    k = k.view(*dims, D // 2, 2)
+
+    # convention is `r` suffix is real part, `i` is imaginary.
+    qr = q[..., 0].float()
+    qi = q[..., 1].float()
+
+    kr = k[..., 0].float()
+    ki = k[..., 1].float()
+
+    rotr = torch.cos(freqs * ts)
+    roti = torch.sin(freqs * ts)
+    qor = qr * rotr - qi * roti
+    qoi = qr * roti + qi * rotr
+
+    kor = kr * rotr - ki * roti
+    koi = kr * roti + ki * rotr
+
+    dtype = q.dtype
+    qo = torch.stack([qor.to(dtype), qoi.to(dtype)], dim=-1)
+    ko = torch.stack([kor.to(dtype), koi.to(dtype)], dim=-1)
+
+    return qo.view(*dims, D), ko.view(*dims, D)
+
+
+class RotaryEmbedding(nn.Module):
+    """Rotary positional embedding (RoPE) from [Su et al 2022](https://arxiv.org/abs/2104.09864).
+
+    Args:
+        max_period (float): Maximum period of the rotation frequencies.
+    """
+
+    def __init__(self, max_period: float = 10000.0):
+        super().__init__()
+        self.max_period = max_period
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        offset: torch.Tensor,
+        time_before_heads: bool = False,
+    ):
+        """Apply rope rotation to query or key tensor."""
+        return apply_rope(q, k, offset, self.max_period, time_before_heads)

moshi/modules/seanet.py

@@ -0,0 +1,416 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: MIT
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+
+# Copyright (c) Kyutai, all rights reserved.
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import typing as tp
+
+import numpy as np
+import torch.nn as nn
+
+from .conv import StreamingConv1d, StreamingConvTranspose1d
+from .streaming import StreamingContainer, StreamingAdd
+from ..utils.compile import torch_compile_lazy
+
+
+class SEANetResnetBlock(StreamingContainer):
+    """Residual block from SEANet model.
+
+    Args:
+        dim (int): Dimension of the input/output.
+        kernel_sizes (list): List of kernel sizes for the convolutions.
+        dilations (list): List of dilations for the convolutions.
+        activation (str): Activation function.
+        activation_params (dict): Parameters to provide to the activation function.
+        norm (str): Normalization method.
+        norm_params (dict): Parameters to provide to the underlying normalization used along with the convolution.
+        causal (bool): Whether to use fully causal convolution.
+        pad_mode (str): Padding mode for the convolutions.
+        compress (int): Reduced dimensionality in residual branches (from Demucs v3).
+        true_skip (bool): Whether to use true skip connection or a simple
+            (streamable) convolution as the skip connection.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        kernel_sizes: tp.List[int] = [3, 1],
+        dilations: tp.List[int] = [1, 1],
+        activation: str = "ELU",
+        activation_params: dict = {"alpha": 1.0},
+        norm: str = "none",
+        norm_params: tp.Dict[str, tp.Any] = {},
+        causal: bool = False,
+        pad_mode: str = "reflect",
+        compress: int = 2,
+        true_skip: bool = True,
+    ):
+        super().__init__()
+        assert len(kernel_sizes) == len(
+            dilations
+        ), "Number of kernel sizes should match number of dilations"
+        act = getattr(nn, activation)
+        hidden = dim // compress
+        block = []
+        for i, (kernel_size, dilation) in enumerate(zip(kernel_sizes, dilations)):
+            in_chs = dim if i == 0 else hidden
+            out_chs = dim if i == len(kernel_sizes) - 1 else hidden
+            block += [
+                act(**activation_params),
+                StreamingConv1d(
+                    in_chs,
+                    out_chs,
+                    kernel_size=kernel_size,
+                    dilation=dilation,
+                    norm=norm,
+                    norm_kwargs=norm_params,
+                    causal=causal,
+                    pad_mode=pad_mode,
+                ),
+            ]
+        self.block = nn.Sequential(*block)
+        self.add = StreamingAdd()
+        self.shortcut: nn.Module
+        if true_skip:
+            self.shortcut = nn.Identity()
+        else:
+            self.shortcut = StreamingConv1d(
+                dim,
+                dim,
+                kernel_size=1,
+                norm=norm,
+                norm_kwargs=norm_params,
+                causal=causal,
+                pad_mode=pad_mode,
+            )
+
+    def forward(self, x):
+        u, v = self.shortcut(x), self.block(x)
+        return self.add(u, v)
+
+
+class SEANetEncoder(StreamingContainer):
+    """SEANet encoder.
+
+    Args:
+        channels (int): Audio channels.
+        dimension (int): Intermediate representation dimension.
+        n_filters (int): Base width for the model.
+        n_residual_layers (int): nb of residual layers.
+        ratios (Sequence[int]): kernel size and stride ratios. The encoder uses downsampling ratios instead of
+            upsampling ratios, hence it will use the ratios in the reverse order to the ones specified here
+            that must match the decoder order. We use the decoder order as some models may only employ the decoder.
+        activation (str): Activation function.
+        activation_params (dict): Parameters to provide to the activation function.
+        norm (str): Normalization method.
+        norm_params (dict): Parameters to provide to the underlying normalization used along with the convolution.
+        kernel_size (int): Kernel size for the initial convolution.
+        last_kernel_size (int): Kernel size for the initial convolution.
+        residual_kernel_size (int): Kernel size for the residual layers.
+        dilation_base (int): How much to increase the dilation with each layer.
+        causal (bool): Whether to use fully causal convolution.
+        pad_mode (str): Padding mode for the convolutions.
+        true_skip (bool): Whether to use true skip connection or a simple
+            (streamable) convolution as the skip connection in the residual network blocks.
+        compress (int): Reduced dimensionality in residual branches (from Demucs v3).
+        disable_norm_outer_blocks (int): Number of blocks for which we don't apply norm.
+            For the encoder, it corresponds to the N first blocks.
+        mask_fn (nn.Module): Optional mask function to apply after convolution layers.
+        mask_position (int): Position of the mask function, with mask_position == 0 for the first convolution layer,
+            mask_position == 1 for the first conv block, etc.
+    """
+
+    def __init__(
+        self,
+        channels: int = 1,
+        dimension: int = 128,
+        n_filters: int = 32,
+        n_residual_layers: int = 3,
+        ratios: tp.List[int] = [8, 5, 4, 2],
+        activation: str = "ELU",
+        activation_params: dict = {"alpha": 1.0},
+        norm: str = "none",
+        norm_params: tp.Dict[str, tp.Any] = {},
+        kernel_size: int = 7,
+        last_kernel_size: int = 7,
+        residual_kernel_size: int = 3,
+        dilation_base: int = 2,
+        causal: bool = False,
+        pad_mode: str = "reflect",
+        true_skip: bool = True,
+        compress: int = 2,
+        disable_norm_outer_blocks: int = 0,
+        mask_fn: tp.Optional[nn.Module] = None,
+        mask_position: tp.Optional[int] = None,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.dimension = dimension
+        self.n_filters = n_filters
+        self.ratios = list(reversed(ratios))
+        del ratios
+        self.n_residual_layers = n_residual_layers
+        self.hop_length = int(np.prod(self.ratios))
+        self.n_blocks = len(self.ratios) + 2  # first and last conv + residual blocks
+        self.disable_norm_outer_blocks = disable_norm_outer_blocks
+        assert (
+            self.disable_norm_outer_blocks >= 0 and self.disable_norm_outer_blocks <= self.n_blocks
+        ), (
+            "Number of blocks for which to disable norm is invalid."
+            "It should be lower or equal to the actual number of blocks in the network and greater or equal to 0."
+        )
+
+        act = getattr(nn, activation)
+        mult = 1
+        model: tp.List[nn.Module] = [
+            StreamingConv1d(
+                channels,
+                mult * n_filters,
+                kernel_size,
+                norm="none" if self.disable_norm_outer_blocks >= 1 else norm,
+                norm_kwargs=norm_params,
+                causal=causal,
+                pad_mode=pad_mode,
+            )
+        ]
+        if mask_fn is not None and mask_position == 0:
+            model += [mask_fn]
+        # Downsample to raw audio scale
+        for i, ratio in enumerate(self.ratios):
+            block_norm = "none" if self.disable_norm_outer_blocks >= i + 2 else norm
+            # Add residual layers
+            for j in range(n_residual_layers):
+                model += [
+                    SEANetResnetBlock(
+                        mult * n_filters,
+                        kernel_sizes=[residual_kernel_size, 1],
+                        dilations=[dilation_base**j, 1],
+                        norm=block_norm,
+                        norm_params=norm_params,
+                        activation=activation,
+                        activation_params=activation_params,
+                        causal=causal,
+                        pad_mode=pad_mode,
+                        compress=compress,
+                        true_skip=true_skip,
+                    )
+                ]
+
+            # Add downsampling layers
+            model += [
+                act(**activation_params),
+                StreamingConv1d(
+                    mult * n_filters,
+                    mult * n_filters * 2,
+                    kernel_size=ratio * 2,
+                    stride=ratio,
+                    norm=block_norm,
+                    norm_kwargs=norm_params,
+                    causal=causal,
+                    pad_mode=pad_mode,
+                ),
+            ]
+            mult *= 2
+            if mask_fn is not None and mask_position == i + 1:
+                model += [mask_fn]
+
+        model += [
+            act(**activation_params),
+            StreamingConv1d(
+                mult * n_filters,
+                dimension,
+                last_kernel_size,
+                norm=(
+                    "none" if self.disable_norm_outer_blocks == self.n_blocks else norm
+                ),
+                norm_kwargs=norm_params,
+                causal=causal,
+                pad_mode=pad_mode,
+            ),
+        ]
+
+        self.model = nn.Sequential(*model)
+
+    @torch_compile_lazy
+    def forward(self, x):
+        return self.model(x)
+
+
+class SEANetDecoder(StreamingContainer):
+    """SEANet decoder.
+
+    Args:
+        channels (int): Audio channels.
+        dimension (int): Intermediate representation dimension.
+        n_filters (int): Base width for the model.
+        n_residual_layers (int): nb of residual layers.
+        ratios (Sequence[int]): kernel size and stride ratios.
+        activation (str): Activation function.
+        activation_params (dict): Parameters to provide to the activation function.
+        final_activation (str): Final activation function after all convolutions.
+        final_activation_params (dict): Parameters to provide to the activation function.
+        norm (str): Normalization method.
+        norm_params (dict): Parameters to provide to the underlying normalization used along with the convolution.
+        kernel_size (int): Kernel size for the initial convolution.
+        last_kernel_size (int): Kernel size for the initial convolution.
+        residual_kernel_size (int): Kernel size for the residual layers.
+        dilation_base (int): How much to increase the dilation with each layer.
+        causal (bool): Whether to use fully causal convolution.
+        pad_mode (str): Padding mode for the convolutions.
+        true_skip (bool): Whether to use true skip connection or a simple.
+            (streamable) convolution as the skip connection in the residual network blocks.
+        compress (int): Reduced dimensionality in residual branches (from Demucs v3).
+        disable_norm_outer_blocks (int): Number of blocks for which we don't apply norm.
+            For the decoder, it corresponds to the N last blocks.
+        trim_right_ratio (float): Ratio for trimming at the right of the transposed convolution under the causal setup.
+            If equal to 1.0, it means that all the trimming is done at the right.
+    """
+
+    def __init__(
+        self,
+        channels: int = 1,
+        dimension: int = 128,
+        n_filters: int = 32,
+        n_residual_layers: int = 3,
+        ratios: tp.List[int] = [8, 5, 4, 2],
+        activation: str = "ELU",
+        activation_params: dict = {"alpha": 1.0},
+        final_activation: tp.Optional[str] = None,
+        final_activation_params: tp.Optional[dict] = None,
+        norm: str = "none",
+        norm_params: tp.Dict[str, tp.Any] = {},
+        kernel_size: int = 7,
+        last_kernel_size: int = 7,
+        residual_kernel_size: int = 3,
+        dilation_base: int = 2,
+        causal: bool = False,
+        pad_mode: str = "reflect",
+        true_skip: bool = True,
+        compress: int = 2,
+        disable_norm_outer_blocks: int = 0,
+        trim_right_ratio: float = 1.0,
+    ):
+        super().__init__()
+        self.dimension = dimension
+        self.channels = channels
+        self.n_filters = n_filters
+        self.ratios = ratios
+        del ratios
+        self.n_residual_layers = n_residual_layers
+        self.hop_length = int(np.prod(self.ratios))
+        self.n_blocks = len(self.ratios) + 2  # first and last conv + residual blocks
+        self.disable_norm_outer_blocks = disable_norm_outer_blocks
+        assert (
+            self.disable_norm_outer_blocks >= 0 and self.disable_norm_outer_blocks <= self.n_blocks
+        ), (
+            "Number of blocks for which to disable norm is invalid."
+            "It should be lower or equal to the actual number of blocks in the network and greater or equal to 0."
+        )
+
+        act = getattr(nn, activation)
+        mult = int(2 ** len(self.ratios))
+        model: tp.List[nn.Module] = [
+            StreamingConv1d(
+                dimension,
+                mult * n_filters,
+                kernel_size,
+                norm=(
+                    "none" if self.disable_norm_outer_blocks == self.n_blocks else norm
+                ),
+                norm_kwargs=norm_params,
+                causal=causal,
+                pad_mode=pad_mode,
+            )
+        ]
+
+        # Upsample to raw audio scale
+        for i, ratio in enumerate(self.ratios):
+            block_norm = (
+                "none"
+                if self.disable_norm_outer_blocks >= self.n_blocks - (i + 1)
+                else norm
+            )
+            # Add upsampling layers
+            model += [
+                act(**activation_params),
+                StreamingConvTranspose1d(
+                    mult * n_filters,
+                    mult * n_filters // 2,
+                    kernel_size=ratio * 2,
+                    stride=ratio,
+                    norm=block_norm,
+                    norm_kwargs=norm_params,
+                    causal=causal,
+                    trim_right_ratio=trim_right_ratio,
+                ),
+            ]
+            # Add residual layers
+            for j in range(n_residual_layers):
+                model += [
+                    SEANetResnetBlock(
+                        mult * n_filters // 2,
+                        kernel_sizes=[residual_kernel_size, 1],
+                        dilations=[dilation_base**j, 1],
+                        activation=activation,
+                        activation_params=activation_params,
+                        norm=block_norm,
+                        norm_params=norm_params,
+                        causal=causal,
+                        pad_mode=pad_mode,
+                        compress=compress,
+                        true_skip=true_skip,
+                    )
+                ]
+
+            mult //= 2
+
+        # Add final layers
+        model += [
+            act(**activation_params),
+            StreamingConv1d(
+                n_filters,
+                channels,
+                last_kernel_size,
+                norm="none" if self.disable_norm_outer_blocks >= 1 else norm,
+                norm_kwargs=norm_params,
+                causal=causal,
+                pad_mode=pad_mode,
+            ),
+        ]
+        # Add optional final activation to decoder (eg. tanh)
+        if final_activation is not None:
+            final_act = getattr(nn, final_activation)
+            final_activation_params = final_activation_params or {}
+            model += [final_act(**final_activation_params)]
+        self.model = nn.Sequential(*model)
+
+    @torch_compile_lazy
+    def forward(self, z):
+        y = self.model(z)
+        return y

moshi/modules/streaming.py

@@ -0,0 +1,630 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: MIT
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+
+# Copyright (c) Kyutai, all rights reserved.
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Streaming module API that should be implemented by all Streaming components,
+"""
+
+import abc
+from contextlib import contextmanager
+from dataclasses import dataclass, fields, is_dataclass
+import itertools
+import math
+import json
+from typing import List, Union, Protocol, TypeVar, Generic, Any, Optional
+import torch
+from safetensors.torch import save_file, load_file
+
+
+class Resetable(Protocol):
+    def reset(self) -> None:
+        pass
+
+
+State = TypeVar("State", bound=Resetable)
+StreamingStateDict = dict[str, Union[torch.Tensor, int, float, str, bool, None]]
+
+
+def is_dataclass_instance(obj):
+    """Check if obj is an instance of a dataclass (not the class itself).
+    
+    Parameters
+    ----------
+    obj : Any
+        Object to check.
+        
+    Returns
+    -------
+    bool
+        True if obj is an instance of a dataclass, False otherwise.
+    """
+    return is_dataclass(obj) and not isinstance(obj, type)
+
+
+def _restore_streaming_state_pt(value: torch.Tensor,
+                                name: str,
+                                state_dict: dict[str, torch.Tensor],
+                                ):
+    """Restore the streaming state from the given pt_state dict
+    
+    Parameters
+    ----------
+    value : torch.Tensor
+        Specific streaming state tensor that needs to be set.
+    name : str
+        Name of the tensor in the state dict.
+    state_dict: StreamingStateDict
+        Flattened state dict containing the values to set.
+    """
+    if name in state_dict:
+        value.copy_(state_dict[name].to(value.device))
+        state_dict.pop(name)
+    else:
+        raise KeyError(f"Expected to find a streaming state for {name}.")
+
+    
+def _set_streaming_state_inplace(streaming_state: State,
+                                 state_dict: StreamingStateDict,
+                                 prefix: str,
+                                 device: torch.device,
+                                 ):
+    """Set the streaming state in-place from the given `state_dict` dict.
+
+    Parameters
+    ----------
+    streaming_state : State
+        Specific streaming state object that needs to be set.
+    state_dict: StreamingStateDict
+        Flattened state dict containing the values to set.
+    prefix : str
+        Prefix to add to each key when looking up in `state_dict`.
+    device : torch.device
+        Device to move tensors to if needed.
+    """
+    if isinstance(streaming_state, torch.Tensor):
+        _restore_streaming_state_pt(streaming_state, prefix, state_dict)
+    elif is_dataclass_instance(streaming_state):
+        _restore_streaming_state_from_keys(streaming_state, state_dict, prefix, [field.name for field in fields(streaming_state)], device)
+    elif hasattr(streaming_state, "asdict"):
+        _restore_streaming_state_from_keys(streaming_state, state_dict, prefix, list(streaming_state.asdict().keys()), device)
+    else:
+        raise TypeError(f"Unsupported type {type(streaming_state)} for streaming state with prefix {prefix}.")
+
+
+def _restore_streaming_state_from_keys(streaming_state: State,
+                                       state_dict: StreamingStateDict,
+                                       prefix: str,
+                                       keys: List[str],
+                                       device: torch.device,
+                                       ):
+    """Restores the streaming state from the given `state_dict` dict
+    looking up fields by adding `prefix` to each key in `keys` to look
+    up values.
+    `torch.Tensor` are copied to `device` if no `torch.Tensor` is already present
+    otherwise, the data is copied to the device of the existing `torch.Tensor`.
+    
+    Parameters
+    ----------
+    streaming_state : State
+        Specific streaming state object that needs to be set.
+    state_dict: StreamingStateDict
+        Flattened state dict containing the values to set.
+    prefix : str
+        Prefix to add to each key when looking up in `state_dict`.
+    keys : List[str]
+        List of keys to look up in `state_dict`.
+    device : torch.device
+        Device to move tensors to if needed.
+    """
+    for key in keys:
+        full_key = f"{prefix}.{key}"
+        existing_value = getattr(streaming_state, key)
+        if isinstance(existing_value, torch.Tensor):
+            _restore_streaming_state_pt(existing_value, full_key, state_dict)
+        elif isinstance(existing_value, (int, float, str, bool, type(None))):
+            if full_key in state_dict:
+                restored_value = state_dict[full_key]
+                if isinstance(restored_value, torch.Tensor):
+                    restored_value = restored_value.to(device)
+                setattr(streaming_state, key, restored_value)
+                
+                state_dict.pop(full_key)
+            else:
+                raise RuntimeError(f"Key {full_key} not found in state_dict.")
+        else:
+            _set_streaming_state_inplace(existing_value, state_dict, full_key, device)
+
+
+def safe_asdict(dataclass_obj):
+    """
+    safe_asdict(dataclass_obj)
+
+    Converts a dataclass object to a dict, skipping empty nested
+    dataclasses without requiring values to be pickleable.
+    
+    Parameters
+    ----------
+    dataclass_obj : Any
+        Dataclass object to convert.
+        
+    Returns
+    -------
+    dict
+        Dictionary representation of the dataclass object.
+    """
+    out = {}
+    for field in fields(dataclass_obj):
+        value = getattr(dataclass_obj, field.name)
+        if is_dataclass_instance(value):
+            subvalue = safe_asdict(value)
+            if len(subvalue) > 0:
+                out[field.name] = subvalue
+        else:
+            out[field.name] = value
+    return out
+        
+
+def _flatten_streaming_state(state_dict: dict[str, torch.Tensor],
+                             state_dict_metadata: dict[str, Union[int, float, str, None]],
+                             state: dict[str, State],
+                             prefix: str,
+                             ):
+    """
+    _flatten_streaming_state(state_dict, state_dict_metadata, state, prefix)
+
+    Helper function for recursively flattening the streaming state into a dict of tensors
+    and a dict of metadata (non-tensor values).
+
+    Parameters
+    ----------
+    state_dict : dict[str, torch.Tensor]
+        Dictionary to store the flattened tensor states.
+    state_dict_metadata : dict[str, Union[int, float, str, None]]
+        Dictionary to store the flattened non-tensor states.
+    state : dict[str, State]
+        The streaming state to flatten.
+    prefix : str
+        Prefix to add to each key in the flattened state.
+    """
+    for key, value in state.items():
+        if isinstance(value, torch.Tensor):
+            state_dict[f"{prefix}{key}"] = value.contiguous()
+        elif is_dataclass_instance(value):
+            _flatten_streaming_state(state_dict, state_dict_metadata, safe_asdict(value), prefix=f"{prefix}{key}.")
+        elif isinstance(value, dict):
+            _flatten_streaming_state(state_dict, state_dict_metadata, value, prefix=f"{prefix}{key}.")
+        elif isinstance(value, (str, int, float, bool, type(None))):
+            state_dict_metadata[f"{prefix}{key}"] = value
+        elif hasattr(value, "asdict"):
+            _flatten_streaming_state(state_dict, state_dict_metadata, value.asdict(), prefix=f"{prefix}{key}.")
+        else:
+            raise TypeError(f"Unsupported type {type(value)} for key {key} (prefix={prefix}) in streaming state.")
+
+
+def load_streaming_state(path: str,
+                         metadata_path: str,
+                         device: Union[str, int] = 'cpu',
+                         ) -> StreamingStateDict:
+    """
+    load_streaming_state(path, metadata_path)
+
+    Loads a streaming state from a safetensors file and its associated metadata json file.
+
+    Parameters
+    ----------
+    str : path
+        Path to the safetensors file.
+    str : metadata_path
+        Path to the metadata json file.
+    device : Union[str, int], optional
+        Device to load the tensors onto, by default 'cpu'.
+
+    Returns
+    -------
+    dict
+        The loaded streaming state flattened as a dictionary.
+    """
+    state_dict = load_file(path, device=device)
+    with open(metadata_path, "rt", encoding="utf-8") as fin:
+        state_dict_metadata = json.load(fin)
+    state_dict.update(state_dict_metadata)
+    return state_dict
+
+
+class StreamingModule(abc.ABC, torch.nn.Module, Generic[State]):
+    """Common API for streaming components.
+
+    Each streaming component has a streaming state, which is just a dict[str, Tensor].
+    By convention, the first dim of each tensor must be the batch size.
+    Don't use dots in the key names, as this would clash with submodules
+    (like in state_dict).
+
+    If `self._is_streaming` is True, the component should use and remember
+    the proper state inside `self._streaming_state`.
+
+    To set a streaming component in streaming state, use
+
+        with module.streaming():
+            ...
+
+    This will automatically reset the streaming state when exiting the context manager.
+    This also automatically propagates to all streaming children module.
+
+    Some module might also implement the `StreamingModule.flush` method, although
+    this one is trickier, as all parents module must be StreamingModule and implement
+    it as well for it to work properly. See `StreamingSequential` after.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._streaming_state: State | None = None
+        self._streaming_propagate: bool = True
+
+    @property
+    def is_streaming(self):
+        return self._streaming_state is not None
+
+    def set_streaming_propagate(self, streaming_propagate: bool):
+        self._streaming_propagate = streaming_propagate
+
+    def _apply_named_streaming(self, fn: Any):
+        def _handle_module(prefix: str, module: torch.nn.Module, recurse: bool = True):
+            propagate = True
+            if isinstance(module, StreamingModule):
+                if module._streaming_propagate:
+                    fn(prefix, module)
+                else:
+                    propagate = False
+            if not recurse:
+                return
+            if propagate:
+                for name, child in module.named_children():
+                    _handle_module(prefix + "." + name, child)
+
+        _handle_module("", self, recurse=False)
+        for name, child in self.named_children():
+            _handle_module(name, child)
+
+    def _start_streaming(self, batch_size: int):
+        def _start_streaming(name: str, module: StreamingModule):
+            module._streaming_state = module._init_streaming_state(batch_size)
+
+        self._apply_named_streaming(_start_streaming)
+
+    def _stop_streaming(self):
+        def _stop_streaming(name: str, module: StreamingModule):
+            module._streaming_state = None
+
+        self._apply_named_streaming(_stop_streaming)
+
+    @abc.abstractmethod
+    def _init_streaming_state(self, batch_size: int) -> State: ...
+
+    def streaming_forever(self, batch_size: int):
+        self._start_streaming(batch_size)
+
+    @contextmanager
+    def streaming(self, batch_size: int):
+        """Context manager to enter streaming mode. Reset streaming state on exit."""
+
+        self._start_streaming(batch_size)
+        try:
+            yield
+        finally:
+            self._stop_streaming()
+
+    def reset_streaming(self):
+        """Reset the streaming state."""
+
+        def _reset(name: str, module: StreamingModule):
+            state = module._streaming_state
+            if state is None:
+                raise ValueError(
+                    f"Trying to reset streaming, but {name} wasn't streaming."
+                )
+            state.reset()
+
+        self._apply_named_streaming(_reset)
+
+    def get_streaming_state(self) -> dict[str, Any]:
+        """Return the complete streaming state, including that of sub-modules."""
+        state: dict[str, Any] = {}
+
+        def _add(name: str, module: StreamingModule):
+            state[name] = module._streaming_state
+
+        self._apply_named_streaming(_add)
+        return state
+
+    def save_streaming_state(self,
+                             save_path: str,
+                             metadata_save_path: str,
+                             extra_state_dict: Optional[dict[str, torch.Tensor]] = None,
+                             ):
+        """Save the streaming state, including that of sub-modules, to the given paths.
+        
+        Parameters
+        ----------
+        save_path : str
+            Path to save the streaming state tensors (safetensors format).
+        metadata_save_path : str
+            Path to save the streaming state metadata (json format).
+        extra_state_dict : Optional[dict[str, torch.Tensor]], optional
+            Extra state dict to include in the saved streaming state tensors, by default None.
+        """
+        state_dict = {}
+        if extra_state_dict is not None:
+            state_dict.update(extra_state_dict)
+        state_dict_metadata = {}
+        state = self.get_streaming_state()
+        _flatten_streaming_state(state_dict, state_dict_metadata, state, prefix="")
+        save_file(state_dict, save_path)
+        with open(metadata_save_path, "wt", encoding="utf-8") as fout:
+            json.dump(state_dict_metadata, fout)
+
+    def set_streaming_state_inplace(self, state: StreamingStateDict):
+        """
+        Set the streaming state in-place, including that of
+        sub-modules using a flattened-state dict.
+        """
+        device = next(self.parameters()).device
+        def _set(name: str, module: StreamingModule):
+            _set_streaming_state_inplace(module._streaming_state, state, prefix=name, device=device)
+        self._apply_named_streaming(_set)
+        if state:
+            raise RuntimeError(f"Some states were not consumed: {list(state.keys())}")
+
+    def set_streaming_state(self, state: dict[str, Any]):
+        """Set the streaming state, including that of sub-modules."""
+        state = dict(state)
+
+        def _set(name: str, module: StreamingModule):
+            if name in state:
+                module._streaming_state = state[name]
+                state.pop(name)
+            else:
+                raise RuntimeError(f"Expected to find a streaming state for {name}.")
+
+        self._apply_named_streaming(_set)
+        if state:
+            raise RuntimeError(f"Some states were not consumed: {list(state.keys())}")
+
+
+@dataclass
+class _NullState:
+    pass
+
+    def reset(self) -> None:
+        pass
+
+
+class StreamingContainer(StreamingModule[_NullState]):
+    def _init_streaming_state(self, batch_size: int) -> _NullState:
+        return _NullState()
+
+
+@dataclass
+class _StreamingAddState:
+    previous_x: torch.Tensor | None = None
+    previous_y: torch.Tensor | None = None
+
+    def reset(self):
+        self.previous_x = None
+        self.previous_y = None
+
+
+class StreamingAdd(StreamingModule[_StreamingAddState]):
+    def _init_streaming_state(self, batch_size: int) -> _StreamingAddState:
+        return _StreamingAddState()
+
+    def forward(self, x: torch.Tensor, y: torch.Tensor):
+        if self._streaming_state is None:
+            return x + y
+        else:
+            prev_x = self._streaming_state.previous_x
+            prev_y = self._streaming_state.previous_y
+            if prev_x is not None:
+                x = torch.cat([prev_x, x], dim=-1)
+            if prev_y is not None:
+                y = torch.cat([prev_y, y], dim=-1)
+            m_l = min(x.shape[-1], y.shape[-1])
+            self._streaming_state.previous_x = x[..., m_l:]
+            self._streaming_state.previous_y = y[..., m_l:]
+            return x[..., :m_l] + y[..., :m_l]
+
+
+@dataclass
+class _StreamingConvState:
+    previous: torch.Tensor | None = None
+
+    def reset(self):
+        self.previous = None
+
+
+class RawStreamingConv1d(torch.nn.Conv1d, StreamingModule[_StreamingConvState]):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        assert self.padding[0] == 0, "Padding should be handled outside."
+        assert (
+            self.stride[0] <= self.kernel_size[0]
+        ), "stride must be less than kernel_size."
+
+    def _init_streaming_state(self, batch_size: int) -> _StreamingConvState:
+        return _StreamingConvState()
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        stride = self.stride[0]
+        # Effective kernel size accounting for dilation.
+        kernel = (self.kernel_size[0] - 1) * self.dilation[0] + 1
+        if self._streaming_state is None:
+            return super().forward(input)
+        else:
+            # Due to the potential overlap, we might have some cache of the previous time steps.
+            previous = self._streaming_state.previous
+            if previous is not None:
+                input = torch.cat([previous, input], dim=-1)
+            B, C, T = input.shape
+            # We now compute the number of full convolution frames, i.e. the frames
+            # that are ready to be computed.
+            num_frames = max(0, int(math.floor((T - kernel) / stride) + 1))
+            offset = num_frames * stride
+            # We will compute `num_frames` outputs, and we are advancing by `stride`
+            # for each of the frame, so we know the data before `stride * num_frames`
+            # will never be used again.
+            self._streaming_state.previous = input[..., offset:]
+            if num_frames > 0:
+                input_length = (num_frames - 1) * stride + kernel
+                out = super().forward(input[..., :input_length])
+            else:
+                # Not enough data as this point to output some new frames.
+                out = torch.empty(
+                    B, self.out_channels, 0, device=input.device, dtype=input.dtype
+                )
+            return out
+
+
+@dataclass
+class _StreamingConvTrState:
+    partial: torch.Tensor | None = None
+
+    def reset(self):
+        self.partial = None
+
+
+class RawStreamingConvTranspose1d(
+    torch.nn.ConvTranspose1d, StreamingModule[_StreamingConvTrState]
+):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        assert self.padding[0] == 0, "Padding should be handled outside."
+        assert self.dilation[0] == 1, "No dilation for now"
+        assert (
+            self.stride[0] <= self.kernel_size[0]
+        ), "stride must be less than kernel_size."
+        assert self.output_padding[0] == 0, "Output padding not supported."
+
+    def _init_streaming_state(self, batch_size: int) -> _StreamingConvTrState:
+        return _StreamingConvTrState()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:  # type: ignore
+        B, C, T = x.shape
+        stride = self.stride[0]
+        kernel = self.kernel_size[0]
+        if self._streaming_state is None:
+            return super().forward(x)
+        else:
+            if T == 0:
+                return torch.empty(
+                    B, self.out_channels, 0, device=x.device, dtype=x.dtype
+                )
+            out = super().forward(x)
+            OT = out.shape[-1]
+            partial = self._streaming_state.partial
+            if partial is not None:
+                # Due to the potential overlap, the rightmost output of the conv transpose is not
+                # ready to be output, as it will receive contributions from the next input frames.
+                # Here we recover those `partial` output frames. We know that the first time step
+                # of the `partial` tensor corresponds to the first time step of `out` as anything
+                # coming before the first time step of `out` would have been already flushed.
+                PT = partial.shape[-1]
+                if self.bias is not None:
+                    out[..., :PT] += partial - self.bias[:, None]
+                else:
+                    out[..., :PT] += partial
+            # The input is T, the output is S * (T - 1) + K.
+            # The offset of the left of the next frame will be S * T
+            # so everything between 0 and S * T is ready to be output, and we need
+            # to keep in the internal state everything beyond that, i.e. S (T - 1) + K - S T = K - S
+            invalid_steps = kernel - stride
+            partial = out[..., OT - invalid_steps :]
+            out = out[..., : OT - invalid_steps]
+            self._streaming_state.partial = partial
+            return out
+
+
+def test():
+    torch.manual_seed(1234)
+    device = "cpu"
+    if torch.cuda.is_available():
+        # Avoid the cuda optimizations that would take place on single precision
+        # floats for convolutions.
+        torch.backends.cudnn.enabled = True
+        torch.backends.cudnn.benchmark = False
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cuda.matmul.allow_tf32 = False
+        torch.backends.cudnn.allow_tf32 = False
+        device = "cuda:0"
+
+    kernel_sizes = [1, 3, 4, 8, 15, 16]
+    strides = [1, 2, 3, 4, 5, 6, 7, 8, 9]
+    chin = 6
+    chout = 12
+
+    for kernel, stride in itertools.product(kernel_sizes, strides):
+        if stride > kernel:
+            continue
+        conv = RawStreamingConv1d(chin, chout, kernel, stride).to(device)
+        convtr = RawStreamingConvTranspose1d(chout, chin, kernel, stride).to(device)
+
+        for length in [4, 8, 32, 54, 65, 128, 1043]:
+            print(f"ksize {kernel} strides {stride} len {length}")
+            if length < kernel:
+                continue
+            batch_size = 3
+            x = torch.randn(batch_size, chin, length).to(device)
+            y = conv(x)
+            z = convtr(y)
+            for chunk_size in [1, 3, 5, 8]:
+                ys = []
+                zs = []
+                with conv.streaming(batch_size), convtr.streaming(batch_size):
+                    for offset in range(0, length, chunk_size):
+                        chunk = x[..., offset : offset + chunk_size]
+                        ys.append(conv(chunk))
+                        zs.append(convtr(ys[-1]))
+                y_stream = torch.cat(ys, dim=-1)
+                z_stream = torch.cat(zs, dim=-1)
+                y = y[..., : y_stream.shape[-1]]
+                z = z[..., : z_stream.shape[-1]]
+                assert y.shape == y_stream.shape, (y.shape, y_stream.shape)
+                delta = (y_stream - y).norm() / y.norm()
+                assert delta <= 1e-6, delta
+                num_frames = int((length - kernel) / stride) + 1
+                assert num_frames == y_stream.shape[-1]
+
+                assert z.shape == z_stream.shape, (z.shape, z_stream.shape)
+                delta = (z_stream - z).norm() / z.norm()
+                assert delta <= 1e-6, (delta, (z_stream - z).abs().mean(dim=(0, 1)))
+
+
+if __name__ == "__main__":
+    with torch.no_grad():
+        test()

moshi/modules/transformer.py

@@ -0,0 +1,774 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: MIT
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+
+# Copyright (c) Kyutai, all rights reserved.
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Transformer model, with streaming support, + CUDA Graphable.
+Optimized for inference.
+
+See `StreamingTransformer` for more information.
+"""
+
+from contextlib import ExitStack
+from dataclasses import dataclass
+import typing as tp
+
+from einops import rearrange
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+from ..utils.compile import no_compile
+from .gating import make_gating
+from .rope import RotaryEmbedding
+from .streaming import StreamingModule, StreamingContainer
+
+
+class LayerNormF32(nn.LayerNorm):
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        x_f32 = input.float()
+        out_f32 = super().forward(x_f32)
+        return out_f32.to(input.dtype)
+
+
+def _rms_norm(
+    x: torch.Tensor,
+    alpha: torch.Tensor,
+    dtype: tp.Optional[torch.dtype],
+    eps: float,
+):
+    assert x.dim() == 3, f"RMSNorm expects 3D inputs but got {x.shape}"
+    x_dtype = x.dtype
+    if dtype is not None:
+        x = x.to(dtype)
+    var = eps + torch.mean(x**2, dim=2, keepdim=True)
+    y = (x * (alpha.to(var) * torch.rsqrt(var))).to(x_dtype)
+    return y
+
+
+class RMSNorm(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        eps: float = 1e-5,
+        dtype: tp.Optional[torch.dtype] = None,
+        device=None,
+    ):
+        super().__init__()
+        self.eps = eps
+        self.dtype = dtype
+        self.alpha = nn.Parameter(
+            torch.full((1, 1, dim), 1.0, requires_grad=True, device=device, dtype=dtype)
+        )
+
+    def forward(self, x: torch.Tensor):
+        return _rms_norm(x, self.alpha, self.dtype, self.eps)
+
+
+class LayerScale(nn.Module):
+    """Layer scale from [Touvron et al 2021] (https://arxiv.org/pdf/2103.17239.pdf).
+    This rescales diagonally the residual outputs close to 0, with a learnt scale.
+
+    Args:
+        channels (int): Number of channels.
+        init (float): Initial scale.
+        channel_last (bool): If True, expect `[*, C]` shaped tensors, otherwise, `[*, C, T]`.
+        device (torch.device or str, optional): Device on which to initialize the module.
+        dtype (torch.dtype, optional): dtype to use to initialize the module.
+    """
+
+    def __init__(
+        self,
+        channels: int,
+        init: float = 1e-4,
+        channel_last: bool = True,
+        device=None,
+        dtype=None,
+    ):
+        super().__init__()
+        self.channel_last = channel_last
+        self.scale = nn.Parameter(
+            torch.full(
+                (channels,), init, requires_grad=True, device=device, dtype=dtype
+            )
+        )
+
+    def forward(self, x: torch.Tensor):
+        if self.channel_last:
+            return self.scale * x
+        else:
+            return self.scale[:, None] * x
+
+
+def create_norm_fn(norm_type: str, dim: int, **kwargs) -> nn.Module:
+    """Create normalization module for transformer encoder layer.
+
+    Args:
+        norm_type (str): Normalization method.
+        dim (int): Dimension of the normalized layer.
+        **kwargs (dict): Additional parameters for normalization layer.
+    Returns:
+        nn.Module: Normalization module.
+    """
+    if norm_type == "layer_norm":
+        return nn.LayerNorm(dim, eps=1e-5, **kwargs)
+    elif norm_type == "layer_norm_f32":
+        kwargs.pop("dtype", None)
+        return LayerNormF32(dim, eps=1e-8, **kwargs)
+    elif norm_type in {"rms_norm"}:
+        return RMSNorm(dim, eps=1e-5, **kwargs)
+    elif norm_type in {"rms_norm_f32"}:
+        kwargs.pop("dtype", None)
+        return RMSNorm(dim, eps=1e-8, dtype=torch.float, **kwargs)
+    else:
+        raise ValueError(f"Unknown norm type: {norm_type}")
+
+
+def create_sin_embedding(
+    positions: torch.Tensor,
+    dim: int,
+    max_period: float = 10000,
+    dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
+    """Create sinusoidal positional embedding, with shape `[B, T, C]`.
+
+    Args:
+        positions (torch.Tensor): LongTensor of positions.
+        dim (int): Dimension of the embedding.
+        max_period (float): Maximum period of the cosine/sine functions.
+        dtype (torch.dtype or str): dtype to use to generate the embedding.
+    Returns:
+        torch.Tensor: Sinusoidal positional embedding.
+    """
+    # We aim for BTC format
+    assert dim % 2 == 0
+    half_dim = dim // 2
+    positions = positions.to(dtype)
+    adim = torch.arange(half_dim, device=positions.device, dtype=dtype).view(1, 1, -1)
+    max_period_tensor = torch.full(
+        [], max_period, device=positions.device, dtype=dtype
+    )  # avoid sync point
+    phase = positions / (max_period_tensor ** (adim / (half_dim - 1)))
+    return torch.cat([torch.cos(phase), torch.sin(phase)], dim=-1)
+
+
+def multi_linear(
+    num_linear: int,
+    weight: torch.Tensor,
+    x: torch.Tensor,
+    offset: int,
+):
+    """Utility to apply a multi linear layer to the given input. A multi linear layer
+    applies a different set of weight for each time step.
+
+    Args:
+        num_linear (int): Number of possible time steps and so number of linears.
+        weight (torch.Tensor): Weight tensor, with shape `[num_linear * chout, chin]`.
+        x (torch.Tensor): Input tensor, with shape `[B, T, C]`.
+        offset (int): offset for the current time step, in particular for decoding, with
+            time steps provided one by one.
+    """
+    B, T, C = x.shape
+    ys = []
+    chout, chin = weight.shape
+    weight = weight.view(num_linear, -1, chin)
+    for t in range(T):
+        y = F.linear(x[:, t], weight[t + offset])
+        ys.append(y)
+    out = torch.stack(ys, 1)
+    return out
+
+
+def set_attention_context(model: nn.Module, context: tp.Optional[int] = None) -> None:
+    """Deactivates or changes the context span (in time steps) in a model.
+    Args:
+        model (nn.Module): model over which to look for attentions.
+        context (int or None): new temporary context value.
+
+    ..Note:: this is not a context manager but a plain function changing the context forever.
+        Initially, it was a context manager, but that led to interesting bugs when using
+        activation checkpointing, with the context being inconsistent between the forward
+        and backward.
+    """
+    for module in model.modules():
+        if isinstance(module, StreamingMultiheadAttention):
+            module.context = context
+
+
+class KVCacheResult(tp.NamedTuple):
+    keys: torch.Tensor
+    values: torch.Tensor
+    positions: torch.Tensor
+
+    @staticmethod
+    def from_kv(keys: torch.Tensor, values: torch.Tensor) -> "KVCacheResult":
+        B, H, T, D = keys.shape
+        assert tuple(values.shape[:-1]) == (B, H, T)
+        positions = torch.arange(T, device=keys.device, dtype=torch.long)
+        return KVCacheResult(keys, values, positions)
+
+
+class RingKVCache:
+    """Efficient streaming KVCache to be compatible with Cuda Graph.
+
+    Args:
+        batch_size (int): Batch size.
+        num_heads (int): Number of heads in the attention.
+        dim_per_head (int): Dimension per head.
+        device (torch.device): Device on which to initialize the cache.
+        dtype (torch.dtype): dtype to use for the cache.
+    """
+
+    def __init__(
+        self,
+        batch_size: int,
+        num_heads: int,
+        dim_per_head: int,
+        capacity: int,
+        device: torch.device = torch.device("cuda"),
+        dtype: torch.dtype = torch.bfloat16,
+    ):
+        self.capacity = capacity
+        self.cache = torch.zeros(
+            (2, batch_size, num_heads, capacity, dim_per_head),
+            device=device,
+            dtype=dtype,
+        )
+        self.end_offset = torch.zeros(1, device=device, dtype=torch.long)
+
+    def reset(self):
+        self.end_offset.zero_()
+
+    def complete(self, k: torch.Tensor, v: torch.Tensor) -> KVCacheResult:
+        assert k.shape[:-1] == v.shape[:-1], (k.shape, v.shape)
+        B, H, T, D = k.shape
+        indexes = torch.arange(T, device=self.end_offset.device, dtype=self.end_offset.dtype) + self.end_offset
+        indexes = indexes % self.capacity
+        self.cache[0].index_copy_(2, indexes, k)
+        self.cache[1].index_copy_(2, indexes, v)
+        self.end_offset.add_(T)
+
+        keys = self.cache[0]
+        values = self.cache[1]
+
+        indexes = torch.arange(
+            self.capacity, device=self.end_offset.device, dtype=torch.long
+        )
+        invalid = indexes >= self.end_offset
+
+        end_index = self.end_offset % self.capacity
+        delta = indexes - end_index
+
+        # If last key is for step S, and capacity is C, last key was written at index S % C.
+        # then end_offset = S + 1, and end_index = (S + 1) % C.
+        # Then for index = (S % C), delta = -1, and the next code gives us:
+        # position(index) = (S + 1) - 1 = S, all good.
+        # Now the time step at end_offset is actually the oldest in the KVCache, e.g., its
+        # position should be (S - self.capacity + 1).
+        # The following code gives us:
+        # position(index + 1) = S + 1 + 0 - self.capacity.
+
+        positions = torch.where(
+            delta <= 0,
+            self.end_offset + delta,
+            self.end_offset + delta - self.capacity,
+        )
+        positions = torch.where(invalid, torch.full_like(positions, -1), positions)
+
+        return KVCacheResult(keys, values, positions)
+
+    def asdict(self):
+        return {"cache": self.cache, "end_offset": self.end_offset}
+
+
+@dataclass
+class _MHAState:
+    kv_cache: RingKVCache
+    offset: torch.Tensor
+    offset_cpu: int
+
+    def reset(self):
+        self.kv_cache.reset()
+        self.offset.zero_()
+        self.offset_cpu = 0
+
+
+class StreamingMultiheadAttention(StreamingModule[_MHAState]):
+    """Similar to `nn.MultiheadAttention` but with support for streaming, causal evaluation.
+
+    Args:
+        embed_dim (int): Dimension to project to.
+        num_heads (int): Number of heads.
+        causal (bool): Causal mask applied automatically.
+        context (int, optional): Number of time steps the attention can access to.
+            When causal, can access `context` time steps into the past, and when non causal,
+            can access `context // 2` steps in the past, and the same in the future.
+        rope (`RotaryEmbedding`, optional): Rope embedding to use.
+        weights_per_step (int): use different weights per time step. If non zero, should correspond to the
+            number of possible time steps.
+        device (torch.device, optional): Device on which to initialize.
+        dtype (torch.dtype, optional): dtype to use.
+    """
+
+    _fsdp_final = True
+
+    def __init__(
+        self,
+        embed_dim: int,
+        num_heads: int,
+        causal: bool = False,
+        context: tp.Optional[int] = None,
+        rope: tp.Optional[RotaryEmbedding] = None,
+        weights_per_step: int = 0,
+        device=None,
+        dtype=None,
+    ):
+        super().__init__()
+        factory_kwargs = {"device": device, "dtype": dtype}
+
+        self.embed_dim = embed_dim
+        self.causal = causal
+        self.context = context
+        self.rope = rope
+        self.num_heads = num_heads
+
+        out_dim = embed_dim
+        out_dim = 3 * embed_dim
+        mult = 1
+        self.weights_per_step = weights_per_step
+        if weights_per_step:
+            mult = weights_per_step
+        in_proj = nn.Linear(embed_dim, mult * out_dim, bias=False, **factory_kwargs)
+        # We try to follow the default PyTorch MHA convention, to easily compare results.
+        self.in_proj_weight = in_proj.weight
+        self.in_proj_bias = in_proj.bias
+        self.out_proj = nn.Linear(
+            embed_dim, mult * embed_dim, bias=False, **factory_kwargs
+        )
+
+    def _init_streaming_state(self, batch_size: int) -> _MHAState:
+        if self.context is None:
+            if self.weights_per_step:
+                capacity = self.weights_per_step
+            else:
+                raise RuntimeError(
+                    "Cannot create a streaming KVCache without a context to estimate capacity."
+                )
+        else:
+            capacity = self.context
+        device = self.in_proj_weight.device
+        # TODO: the following estimation will not work great with FSDP.
+        dtype = self.in_proj_weight.dtype
+        dim_per_head = self.embed_dim // self.num_heads
+        kv_cache = RingKVCache(
+            batch_size, self.num_heads, dim_per_head, capacity, device, dtype
+        )
+        return _MHAState(
+            kv_cache,
+            offset=torch.zeros(1, device=device, dtype=torch.long),
+            offset_cpu=0,
+        )
+
+    def _complete_kv(self, k, v) -> KVCacheResult:
+        state = self._streaming_state
+        if state is None:
+            return KVCacheResult.from_kv(k, v)
+        else:
+            return state.kv_cache.complete(k, v)
+
+    def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor):
+        state = self._streaming_state
+        T = query.shape[1]
+
+        if state is None:
+            offset = torch.zeros(1, device=query.device, dtype=torch.long)
+            offset_cpu = 0
+        else:
+            assert self.causal, "Streaming only available for causal"
+            offset = state.offset
+            offset_cpu = state.offset_cpu
+
+        if self.weights_per_step:
+            projected = multi_linear(
+                self.weights_per_step, self.in_proj_weight, query, offset_cpu
+            )
+        else:
+            projected = nn.functional.linear(query, self.in_proj_weight)
+        q, k, v = rearrange(
+            projected, "b t (p h d) -> p b h t d", p=3, h=self.num_heads
+        )
+
+        if self.rope:
+            q, k = self.rope(q, k, offset, time_before_heads=False)
+
+        k, v, pos_k = self._complete_kv(k, v)
+        if self.causal:
+            pos_k = pos_k.view(1, -1)
+            pos_q = offset + torch.arange(T, device=q.device, dtype=torch.long).view(
+                -1, 1
+            )
+            delta = pos_q - pos_k
+            attn_bias = (pos_k >= 0) & (delta >= 0)
+            if self.context is not None:
+                attn_bias = attn_bias & (delta < self.context)
+        else:
+            attn_bias = None
+        x = F.scaled_dot_product_attention(q, k, v, attn_bias, dropout_p=0.0)
+
+        x = rearrange(x, "b h t d -> b t (h d)")
+        if self.weights_per_step:
+            x = multi_linear(self.weights_per_step, self.out_proj.weight, x, offset_cpu)
+        else:
+            x = self.out_proj(x)
+        if state is not None:
+            state.offset.add_(T)
+            state.offset_cpu += T
+        return x
+
+
+@dataclass
+class _LayerState:
+    offset_cpu: int
+
+    def reset(self):
+        self.offset_cpu = 0
+
+
+class StreamingTransformerLayer(StreamingModule[_LayerState]):
+    """TransformerLayer with Streaming / Causal support.
+
+    Args:
+        d_model (int): Dimension of the data.
+        num_heads (int): Number of heads.
+        dim_feedforward (int): Intermediate dimension of FF module.
+        causal (bool): Causal mask applied automatically.
+        context (int, optional): Receptive field for the causal mask, infinite if None.
+        custom (bool): Use custom MHA implementation, for testing / benchmarking.
+        rope (`RotaryEmbedding`, optional): Rope embedding to use.
+        norm (str): Normalization to use. Currently, only 'layer_norm' is supported.
+        layer_scale (float, optional): If not None, LayerScale will be used with the given value as initial scale.
+        gating (str): if provided, replaces FFN with special gating, like GLU, GSiGLU etc.
+        weights_per_step (int): use different weights per time step. If non zero, should correspond to the
+            number of possible time steps.
+        skip_self_attn: If true, skips the self attention module and the norm
+        device (torch.device, optional): Device on which to initialize.
+        dtype (torch.dtype, optional): dtype to use.
+    """
+
+    _fsdp_final = True
+
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        dim_feedforward: int | list[int] = 2048,
+        causal: bool = False,
+        context: tp.Optional[int] = None,
+        rope: tp.Optional[RotaryEmbedding] = None,
+        norm: str = "layer_norm",
+        layer_scale: tp.Optional[float] = None,
+        gating: str = "none",
+        weights_per_step: int = 0,
+        activation=F.gelu,
+        skip_self_attn: bool = False,
+        device=None,
+        dtype=None,
+    ):
+        super().__init__()
+        factory_kwargs = {"device": device, "dtype": dtype}
+        # Redefine self_attn to our streaming multi-head attention
+        attn_kwargs: tp.Dict[str, tp.Any] = {
+            "embed_dim": d_model,
+            "num_heads": num_heads,
+        }
+        if not skip_self_attn:
+            self.self_attn: StreamingMultiheadAttention = StreamingMultiheadAttention(
+                causal=causal,
+                context=context,
+                rope=rope,
+                weights_per_step=weights_per_step,
+                **attn_kwargs,  # type: ignore
+                **factory_kwargs,  # type: ignore
+            )  # type: ignore
+            self.norm1 = create_norm_fn(norm, d_model, **factory_kwargs)
+        self.norm2 = create_norm_fn(norm, d_model, **factory_kwargs)
+        # Redefine feedforward layers to expose bias parameter
+        self.weights_per_step = weights_per_step
+        self.gating: tp.Optional[nn.Module] = None
+        self.linear1: tp.Optional[nn.Module] = None
+        self.linear2: tp.Optional[nn.Module] = None
+        self.activation = activation
+        self.skip_self_attn = skip_self_attn
+
+        if isinstance(dim_feedforward, list):
+            assert dim_feedforward
+            assert len(dim_feedforward) == weights_per_step, (
+                "Length of dim_feedforward must match weights_per_step,"
+                f" got {len(dim_feedforward)} != {weights_per_step}"
+            )
+        if gating == "none":
+            assert (
+                not weights_per_step
+            ), "weights_per_step without gating not supported for now."
+            assert not isinstance(
+                dim_feedforward, list
+            ), "List dim_feedforward without gating not supported for now."
+            self.linear1 = nn.Linear(
+                d_model, dim_feedforward, bias=False, **factory_kwargs
+            )
+            self.linear2 = nn.Linear(
+                dim_feedforward, d_model, bias=False, **factory_kwargs
+            )
+        else:
+            self.linear1 = None
+            self.linear2 = None
+            if weights_per_step:
+                if isinstance(dim_feedforward, int):
+                    dim_feedforward = [dim_feedforward] * weights_per_step
+                assert isinstance(dim_feedforward, list), dim_feedforward
+                self.gating = nn.ModuleList(
+                    [
+                        make_gating(gating, d_model, dim, **factory_kwargs)
+                        for dim in dim_feedforward
+                    ]
+                )
+            else:
+                assert isinstance(dim_feedforward, int)
+                self.gating = make_gating(
+                    gating, d_model, dim_feedforward, **factory_kwargs
+                )
+
+        self.layer_scale_1: nn.Module
+        self.layer_scale_2: nn.Module
+        if layer_scale is None:
+            self.layer_scale_1 = nn.Identity()
+            self.layer_scale_2 = nn.Identity()
+        else:
+            self.layer_scale_1 = LayerScale(d_model, layer_scale, **factory_kwargs)  # type: ignore
+            self.layer_scale_2 = LayerScale(d_model, layer_scale, **factory_kwargs)  # type: ignore
+
+    def _init_streaming_state(self, batch_size: int) -> _LayerState:
+        return _LayerState(offset_cpu=0)
+
+    # feed forward block
+    def _ff_block(self, x: torch.Tensor) -> torch.Tensor:
+        state = self._streaming_state
+        offset = 0
+        if state is not None:
+            offset = state.offset_cpu
+        x_orig = x
+        x = self.norm2(x)
+        if self.gating is None:
+            assert self.linear1 is not None
+            assert self.linear2 is not None
+            update = self.linear2(self.activation(self.linear1(x)))
+        else:
+            if self.weights_per_step:
+                assert isinstance(self.gating, nn.ModuleList)
+                B, T, D = x.shape
+                ys = []
+                for t in range(T):
+                    y = self.gating[offset + t](x[:, t : t + 1])
+                    ys.append(y)
+                update = torch.cat(ys, dim=1)
+            else:
+                update = self.gating(x)
+        return x_orig + self.layer_scale_2(update)
+
+    def _sa_block(self, x: torch.Tensor):
+        if self.skip_self_attn:
+            return x
+        x_orig = x
+        x = self.norm1(x)
+        update = self.self_attn(x, x, x)
+        return x_orig + self.layer_scale_1(update)
+
+    def forward(self, x: torch.Tensor):
+        with ExitStack() as stack:
+            if x.device.type != 'cuda':
+                stack.enter_context(no_compile())
+            x = self._sa_block(x)
+            x = self._ff_block(x)
+            state = self._streaming_state
+            if state:
+                state.offset_cpu += x.shape[1]
+            return x
+
+
+@dataclass
+class _TransformerState:
+    offset: torch.Tensor
+
+    def reset(self):
+        self.offset.zero_()
+
+
+class StreamingTransformer(StreamingModule[_TransformerState]):
+    """Transformer with Streaming / Causal support.
+
+    Args:
+        d_model (int): Dimension of the data.
+        num_heads (int): Number of heads.
+        dim_feedforward (int): Intermediate dimension of FF module.
+        causal (bool): Causal mask applied automatically.
+        context (int, optional): Receptive field for the causal mask, infinite if None.
+        layer_scale (float, optional): If not None, LayerScale will be used
+            with the given value as initial scale.
+        positional_embedding (str): Positional embedding strategy (sin, rope, sin_rope, or none).
+        max_period (float): Maximum period of the time embedding.
+        positional_scale (float): Scale of positional embedding, set to 0 to deactivate.
+        layer_class: (subclass of `StreamingTransformerLayer): class to use
+            to initialize the layers, allowing further customization outside of AudioCraft.
+        device (torch.device, optional): Device on which to initialize.
+        dtype (torch.dtype, optional): dtype to use.
+        **kwargs: See `StreamingTransformerLayer`.
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        num_layers: int,
+        dim_feedforward: int | list[int] = 2048,
+        causal: bool = False,
+        context: tp.Optional[int] = None,
+        positional_embedding: str = "sin",
+        max_period: float = 10_000,
+        positional_scale: float = 1.0,
+        betas: tp.Optional[tp.Tuple[float, float]] = None,
+        layer_class: tp.Type[StreamingTransformerLayer] = StreamingTransformerLayer,
+        device=None,
+        dtype=None,
+        **kwargs,
+    ):
+        super().__init__()
+        assert d_model % num_heads == 0
+
+        self.positional_embedding = positional_embedding
+        self.max_period = max_period
+        self.positional_scale = positional_scale
+        self.betas = betas
+
+        assert positional_embedding in {"sin", "rope", "sin_rope", "none"}
+        self.rope: tp.Optional[RotaryEmbedding] = None
+        if self.positional_embedding in {"rope", "sin_rope"}:
+            self.rope = RotaryEmbedding(max_period=max_period)
+
+        self.layers = nn.ModuleList()
+        for _ in range(num_layers):
+            self.layers.append(
+                layer_class(
+                    d_model=d_model,
+                    num_heads=num_heads,
+                    dim_feedforward=dim_feedforward,
+                    causal=causal,
+                    context=context,
+                    rope=self.rope,
+                    device=device,
+                    dtype=dtype,
+                    **kwargs,
+                )
+            )
+
+    def _init_streaming_state(self, batch_size: int) -> _TransformerState:
+        device = next(self.parameters()).device
+        return _TransformerState(offset=torch.zeros(1, device=device, dtype=torch.long))
+
+    def forward(self, x: torch.Tensor, *args, **kwargs):
+        B, T, C = x.shape
+
+        state = self._streaming_state
+        if state is None:
+            offset = torch.zeros(1, dtype=torch.long, device=x.device)
+        else:
+            offset = state.offset
+
+        if self.positional_embedding in {"sin", "sin_rope"}:
+            positions = torch.arange(T, device=x.device).view(1, -1, 1)
+            positions = positions + offset.view(-1, 1, 1)
+            pos_emb = create_sin_embedding(
+                positions, C, max_period=self.max_period, dtype=x.dtype
+            )
+            x = x + self.positional_scale * pos_emb
+
+        for layer in self.layers:
+            x = layer(x, *args, **kwargs)
+
+        if state is not None:
+            state.offset.add_(T)
+        return x
+
+
+class ProjectedTransformer(StreamingContainer):
+    """Transformer with optional projections of the input and output to different dimensions when needed.
+    Supports multiple outputs.
+
+    Args:
+        input_dimension (int): dimension of the input.
+        output_dimensions (tuple[int]): dimensions of the outputs.
+        d_model (int): inner dimension of the Transformer.
+        conv_layout (bool): If True, expects `[B, C, T]` shaped tensors, otherwise, `[B, T, C]`.
+            Similarly, the output will have the same layout.
+    """
+
+    def __init__(
+        self,
+        input_dimension: int,
+        output_dimensions: tp.Tuple[int, ...],
+        d_model: int,
+        *,
+        conv_layout: bool = False,
+        **kwargs,
+    ):
+        super().__init__()
+        self.transformer = StreamingTransformer(d_model=d_model, **kwargs)
+        self.input_dimension = input_dimension
+        self.output_dimensions = output_dimensions
+        self.conv_layout = conv_layout
+        self.input_proj = None
+        if d_model != input_dimension:
+            self.input_proj = nn.Linear(input_dimension, d_model, bias=False)
+
+        self.output_projs = nn.ModuleList()
+        for output_dimension in output_dimensions:
+            if d_model == output_dimension:
+                self.output_projs.append(nn.Identity())
+            else:
+                self.output_projs.append(
+                    nn.Linear(d_model, output_dimension, bias=False)
+                )
+
+    def forward(self, x, *args, **kwargs):
+        if self.conv_layout:
+            x = x.transpose(1, 2)
+        if self.input_proj is not None:
+            x = self.input_proj(x)
+        z = self.transformer(x, *args, **kwargs)
+        ys = []
+        for output_proj in self.output_projs:
+            y = output_proj(z)
+            if self.conv_layout:
+                y = y.transpose(1, 2)
+            ys.append(y)
+        return ys

moshi/offline.py

@@ -0,0 +1,431 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: MIT
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+
+
+# Copyright (c) Kyutai, all rights reserved.
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Offline inference entrypoint for PersonaPlex that mirrors server.py behavior without a WebSocket server.
+
+High-level flow:
+- Load Mimi encoders/decoders, Moshi LM, and tokenizer (same as server.py)
+- Warmup to initialize CUDA graphs and streaming state
+- Prompt phase: load system text tokens and a voice prompt WAV (agent side)
+- Streaming-like phase: feed user audio frames from a WAV file into the "input" channels,
+  autoregressively sample text + agent audio channels each step, and decode audio frames
+- Concatenate generated frames and write an output WAV matching the input duration
+
+This script reuses helpers from lm.py (load_audio, _iterate_audio, encode_from_sphn) to
+keep parity with voice-prompt feeding logic in the server.
+"""
+
+import argparse
+import os
+import tarfile
+from pathlib import Path
+import json
+from typing import Optional, List
+
+import numpy as np
+import torch
+import sentencepiece
+import sphn
+from huggingface_hub import hf_hub_download
+
+from .client_utils import make_log
+from .models import loaders, LMGen, MimiModel
+from .models.lm import load_audio as lm_load_audio
+from .models.lm import _iterate_audio as lm_iterate_audio
+from .models.lm import encode_from_sphn as lm_encode_from_sphn
+
+
+def log(level: str, msg: str):
+    print(make_log(level, msg))
+
+
+def seed_all(seed: int):
+    """Seed torch, CUDA, numpy, and Python RNG for reproducible runs.
+
+    Matches the seeding strategy in server.py.
+    """
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+    import random
+    import numpy as _np
+    random.seed(seed)
+    _np.random.seed(seed)
+    torch.backends.cudnn.deterministic = False
+    torch.backends.cudnn.benchmark = False
+
+
+def wrap_with_system_tags(text: str) -> str:
+    """Add system tags as the model expects if they are missing.
+    Example: "<system> You enjoy having a good conversation. Have a deep conversation about technology. Your name is Jane. <system>"
+    """
+    cleaned = text.strip()
+    if cleaned.startswith("<system>") and cleaned.endswith("<system>"):
+        return cleaned
+    return f"<system> {cleaned} <system>"
+
+
+def warmup(mimi: MimiModel, other_mimi: MimiModel, lm_gen: LMGen, device: str, frame_size: int):
+    """Run a short warmup loop to initialize CUDA graphs and streaming state.
+
+    Replicates the same warmup behavior as server.py: zeros → encode → LMGen.step → decode.
+    """
+    for _ in range(4):
+        chunk = torch.zeros(1, 1, frame_size, dtype=torch.float32, device=device)
+        codes = mimi.encode(chunk)
+        _ = other_mimi.encode(chunk)
+        for c in range(codes.shape[-1]):
+            tokens = lm_gen.step(codes[:, :, c : c + 1])
+            if tokens is None:
+                continue
+            # Decode agent audio channels to ensure decode graphs/states are primed
+            _ = mimi.decode(tokens[:, 1:9])
+            _ = other_mimi.decode(tokens[:, 1:9])
+    if torch.cuda.is_available():
+        torch.cuda.synchronize()
+
+
+def decode_tokens_to_pcm(mimi: MimiModel, other_mimi: MimiModel, lm_gen: LMGen, tokens: torch.Tensor) -> np.ndarray:
+    """Decode a single step of model tokens to PCM using Mimi.
+
+    tokens is shaped [B, dep_q+1, 1]; channels 1..dep_q are the agent audio codebooks.
+    Returns a 1D float32 numpy array (mono) for the current frame.
+    """
+    pcm = mimi.decode(tokens[:, 1:9])
+    _ = other_mimi.decode(tokens[:, 1:9])
+    pcm = pcm.detach().cpu().numpy()[0, 0]
+    return pcm
+
+
+def _get_voice_prompt_dir(voice_prompt_dir: Optional[str], hf_repo: str) -> Optional[str]:
+    """
+    If voice_prompt_dir is None:
+      - download voices.tgz from HF
+      - extract it once
+      - return extracted directory
+    If voice_prompt_dir is provided:
+      - just return it
+    """
+    if voice_prompt_dir is not None:
+        return voice_prompt_dir
+
+    log("info", "retrieving voice prompts")
+    voices_tgz = hf_hub_download(hf_repo, "voices.tgz")
+    voices_tgz = Path(voices_tgz)
+    voices_dir = voices_tgz.parent / "voices"
+
+    if not voices_dir.exists():
+        log("info", f"extracting {voices_tgz} to {voices_dir}")
+        with tarfile.open(voices_tgz, "r:gz") as tar:
+            tar.extractall(path=voices_tgz.parent)
+
+    if not voices_dir.exists():
+        raise RuntimeError("voices.tgz did not contain a 'voices/' directory")
+
+    return str(voices_dir)
+
+
+def run_inference(
+    input_wav: str,
+    output_wav: str,
+    output_text: str,
+    text_prompt: str,
+    voice_prompt_path: str,
+    tokenizer_path: Optional[str],
+    moshi_weight: Optional[str],
+    mimi_weight: Optional[str],
+    hf_repo: str,
+    device: str,
+    seed: Optional[int],
+    temp_audio: float,
+    temp_text: float,
+    topk_audio: int,
+    topk_text: int,
+    greedy: bool,
+    save_voice_prompt_embeddings: bool,
+    cpu_offload: bool = False,
+):
+    """Run offline inference using an input WAV as the user-side stream.
+
+    - Loads/initializes models and tokenizer
+    - Warms up execution
+    - Loads system text tokens and voice prompt
+    - Runs prompt phases (text + voice + silences) via LMGen.step_system_prompts
+    - Streams the user WAV frames into the input channels and samples model outputs
+    - Decodes and writes an output WAV of the same duration
+    """
+    if seed is not None and seed != -1:
+        seed_all(seed)
+
+    # Download config.json to increment download counter
+    # No worries about double-counting since config.json will be cached the second time
+    hf_hub_download(hf_repo, "config.json")
+
+    # 1) Load Mimi encoders/decoders (same as server.py)
+    log("info", "loading mimi")
+    if mimi_weight is None:
+        mimi_weight = hf_hub_download(hf_repo, loaders.MIMI_NAME)  # type: ignore
+    mimi = loaders.get_mimi(mimi_weight, device)
+    other_mimi = loaders.get_mimi(mimi_weight, device)
+    log("info", "mimi loaded")
+
+    # 2) Load tokenizer
+    if tokenizer_path is None:
+        tokenizer_path = hf_hub_download(hf_repo, loaders.TEXT_TOKENIZER_NAME)  # type: ignore
+    text_tokenizer = sentencepiece.SentencePieceProcessor(tokenizer_path)  # type: ignore
+
+    # 3) Load Moshi LM and eval mode
+    log("info", "loading moshi")
+    if moshi_weight is None:
+        moshi_weight = hf_hub_download(hf_repo, loaders.MOSHI_NAME)  # type: ignore
+    lm = loaders.get_moshi_lm(moshi_weight, device=device, cpu_offload=cpu_offload)
+    lm.eval()
+    log("info", "moshi loaded")
+
+    # 4) Construct LMGen like server.py's ServerState does
+    frame_size = int(mimi.sample_rate / mimi.frame_rate)
+    lm_gen = LMGen(
+        lm,
+        audio_silence_frame_cnt=int(0.5 * mimi.frame_rate),  # spacer after prompts
+        sample_rate=mimi.sample_rate,
+        device=device,
+        frame_rate=mimi.frame_rate,
+        save_voice_prompt_embeddings=save_voice_prompt_embeddings,
+        use_sampling=not greedy,
+        temp=temp_audio,
+        temp_text=temp_text,
+        top_k=topk_audio,
+        top_k_text=topk_text,
+    )
+    # Keep models in streaming mode similar to the server
+    mimi.streaming_forever(1)
+    other_mimi.streaming_forever(1)
+    lm_gen.streaming_forever(1)
+
+    # 5) Warmup
+    log("info", "warming up the model")
+    warmup(mimi, other_mimi, lm_gen, device, frame_size)
+
+    # 6) Prompt configuration (text + voice)
+    # System text tokens (k=0) and agent voice-prompt audio (k=1..dep_q) are forced
+    if voice_prompt_path.endswith('.pt'):
+        # Load pre-saved voice prompt embeddings
+        lm_gen.load_voice_prompt_embeddings(voice_prompt_path)
+    else:
+        lm_gen.load_voice_prompt(voice_prompt_path)
+    lm_gen.text_prompt_tokens = (
+        text_tokenizer.encode(wrap_with_system_tags(text_prompt)) if len(text_prompt) > 0 else None
+    )
+
+    # 7) Reset streaming and run initial prompt phases
+    #    - Voice prompt injection
+    #    - Audio silence
+    #    - Text prompt injection
+    #    - Final audio silence
+    mimi.reset_streaming()
+    other_mimi.reset_streaming()
+    lm_gen.reset_streaming()
+    lm_gen.step_system_prompts(mimi)
+    # Reset mimi streaming after voice prompt encoding
+    mimi.reset_streaming()
+
+    # 8) Load and iterate user audio frames for feeding into the input channels
+    sample_rate = mimi.sample_rate
+    user_audio = lm_load_audio(input_wav, sample_rate)  # (C, T) at model SR
+
+    # 9) Encode user audio with Mimi (same iterator logic used for voice prompts),
+    #    and step the model one frame at a time, collecting decoded PCM frames
+    generated_frames: List[np.ndarray] = []
+    generated_text_tokens: List[str] = []
+    total_target_samples = user_audio.shape[-1]
+
+    for user_encoded in lm_encode_from_sphn(
+        mimi,
+        lm_iterate_audio(
+            user_audio, sample_interval_size=lm_gen._frame_size, pad=True
+        ),
+        max_batch=1,
+    ):
+        # user_encoded: [1, K, T]. Feed one step at a time (usually T==1)
+        steps = user_encoded.shape[-1]
+        for c in range(steps):
+            step_in = user_encoded[:, :, c : c + 1]
+            # Feed user-side input channels; text + agent audio are sampled
+            tokens = lm_gen.step(step_in)
+            if tokens is None:
+                continue
+            # Decode current sampled agent frame to PCM
+            pcm = decode_tokens_to_pcm(mimi, other_mimi, lm_gen, tokens)
+            generated_frames.append(pcm)
+            # Decode text token
+            text_token = tokens[0, 0, 0].item()
+            if text_token not in (0, 3):
+                _text = text_tokenizer.id_to_piece(text_token)  # type: ignore
+                _text = _text.replace("▁", " ")
+                log("info", f"text token '{_text}'")
+                generated_text_tokens.append(_text)
+            else:
+                text_token_map = ['EPAD', 'BOS', 'EOS', 'PAD']
+                log("info", f"text token '{text_token_map[text_token]}'")
+                generated_text_tokens.append(text_token_map[text_token])
+
+    if len(generated_frames) == 0:
+        log("error", "No audio frames were generated. Check input file and configuration.")
+        return
+
+    # 10) Concatenate frames and trim/pad to match input duration
+    output_pcm = np.concatenate(generated_frames, axis=-1)
+    if output_pcm.shape[-1] > total_target_samples:
+        output_pcm = output_pcm[:total_target_samples]
+    elif output_pcm.shape[-1] < total_target_samples:
+        pad_len = total_target_samples - output_pcm.shape[-1]
+        output_pcm = np.concatenate(
+            [output_pcm, np.zeros(pad_len, dtype=output_pcm.dtype)], axis=-1
+        )
+
+    # 11) Write mono WAV at model sample rate
+    sphn.write_wav(output_wav, output_pcm, sample_rate)
+    log("info", f"Wrote output audio to {output_wav}")
+
+    # 12) Write text tokens
+    with open(output_text, "w") as file:
+        json.dump(generated_text_tokens, file, ensure_ascii=False)
+    log("info", f"Wrote output text to {output_text}")    
+
+
+def main():
+    """Parse CLI args and run offline inference."""
+    parser = argparse.ArgumentParser(
+        description="Offline inference from WAV input using Moshi server components."
+    )
+    parser.add_argument(
+        "--input-wav", required=True, type=str, help="Path to input WAV file (user audio)"
+    )
+    parser.add_argument(
+        "--output-wav", required=True, type=str, help="Path to output WAV file of agent audio to write"
+    )
+    parser.add_argument(
+        "--output-text", required=True, type=str, help="Path to output JSON file of agent text to write"
+    )
+    parser.add_argument("--text-prompt", default="You are a wise and friendly teacher. Answer questions or provide advice in a clear and engaging way.", type=str, help="Text prompt")
+
+    parser.add_argument(
+        "--voice-prompt", required=True, type=str, help="Voice prompt filename (basename) inside --voice-prompt-dir (e.g. 'NATM1.pt')."
+    )
+    parser.add_argument(
+        "--voice-prompt-dir",
+        type=str,
+        help=(
+            "Directory containing voice prompt files. "
+            "If omitted, voices.tgz is downloaded from HF and extracted."
+            "Voice prompt filenames from -voice-prompt arg will be joined with this directory path."
+        )
+    )
+
+    # Model assets
+    parser.add_argument("--tokenizer", type=str, help="Path to a local tokenizer file.")
+    parser.add_argument("--moshi-weight", type=str, help="Path to a local checkpoint file for Moshi.")
+    parser.add_argument("--mimi-weight", type=str, help="Path to a local checkpoint file for Mimi.")
+    parser.add_argument(
+        "--hf-repo",
+        type=str,
+        default=loaders.DEFAULT_REPO,
+        help="HF repo to look into (defaults to pre-trained model repo)",
+    )
+
+    # Runtime / sampling controls (mirror UI semantics)
+    parser.add_argument(
+        "--temp-audio", type=float, default=0.8, help="Audio sampling temperature (default: 0.8)"
+    )
+    parser.add_argument(
+        "--temp-text", type=float, default=0.7, help="Text sampling temperature (default: 0.7)"
+    )
+    parser.add_argument(
+        "--topk-audio", type=int, default=250, help="Audio top-k sampling (default: 250)"
+    )
+    parser.add_argument(
+        "--topk-text", type=int, default=25, help="Text top-k sampling (default: 25)"
+    )
+    parser.add_argument(
+        "--greedy", action="store_true", help="Disable sampling (greedy decoding)"
+    )
+    parser.add_argument(
+        "--device", type=str, default="cuda", help="Device on which to run, defaults to 'cuda'."
+    )
+    parser.add_argument("--cpu-offload", action="store_true",
+                        help="Offload LM model layers to CPU when GPU memory is insufficient. "
+                             "Requires 'accelerate' package.")
+    parser.add_argument("--seed", type=int, default=-1, help="Seed for reproducibility (-1 disables)")
+
+    args = parser.parse_args()
+
+    # If --voice-prompt-dir is omitted, voices.tgz is downloaded from HF and extracted.
+    voice_prompt_dir = _get_voice_prompt_dir(
+        args.voice_prompt_dir,
+        args.hf_repo,
+    )
+    if not os.path.exists(voice_prompt_dir):
+        raise FileNotFoundError(f"voice_prompt_dir does not exist: {voice_prompt_dir}")
+    log("info", f"voice_prompt_dir = {voice_prompt_dir}")
+
+    # Join basename with directory (DO NOT mutate args.voice_prompt)
+    voice_prompt_path = os.path.join(voice_prompt_dir, args.voice_prompt)
+    if not os.path.exists(voice_prompt_path):
+        raise FileNotFoundError(
+            f"Voice prompt '{args.voice_prompt}' not found in "
+            f"'{voice_prompt_dir}' (resolved: {voice_prompt_path})"
+        )
+
+    # Normalize greedy flag behavior (True if present, False otherwise)
+    greedy = bool(args.greedy)
+
+    with torch.no_grad():
+        run_inference(
+            input_wav=args.input_wav,
+            output_wav=args.output_wav,
+            output_text=args.output_text,
+            text_prompt=args.text_prompt,
+            voice_prompt_path=voice_prompt_path,
+            tokenizer_path=args.tokenizer,
+            moshi_weight=args.moshi_weight,
+            mimi_weight=args.mimi_weight,
+            hf_repo=args.hf_repo,
+            device=args.device,
+            seed=args.seed,
+            temp_audio=args.temp_audio,
+            temp_text=args.temp_text,
+            topk_audio=args.topk_audio,
+            topk_text=args.topk_text,
+            greedy=greedy,
+            save_voice_prompt_embeddings=False,
+            cpu_offload=args.cpu_offload,
+        )
+
+
+if __name__ == "__main__":
+    main()

moshi/quantization/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) Kyutai, all rights reserved.
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+"""RVQ."""
+# flake8: noqa
+from .vq import ResidualVectorQuantizer, SplitResidualVectorQuantizer
+from .base import BaseQuantizer, DummyQuantizer, QuantizedResult

moshi/quantization/base.py

@@ -0,0 +1,170 @@
+# Copyright (c) Kyutai, all rights reserved.
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Base class for all quantizers.
+"""
+
+from dataclasses import dataclass, field
+import typing as tp
+
+import torch
+from torch import nn
+
+
+@dataclass
+class QuantizedResult:
+    x: torch.Tensor
+    codes: torch.Tensor
+    bandwidth: torch.Tensor  # bandwidth in kb/s used, per batch item.
+    penalty: tp.Optional[torch.Tensor] = None
+    metrics: dict = field(default_factory=dict)
+
+
+class BaseQuantizer(nn.Module):
+    """Base class for quantizers."""
+
+    def __init__(self):
+        super().__init__()
+        self._ema_frozen = False
+
+    def forward(self, x: torch.Tensor, frame_rate: int) -> QuantizedResult:
+        """
+        Given input tensor x, returns first the quantized (or approximately quantized)
+        representation along with quantized codes, bandwidth, and any penalty term for the loss.
+        Finally, this returns a dict of metrics to update logging etc.
+        Frame rate must be passed so that the bandwidth is properly computed.
+        """
+        raise NotImplementedError()
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """Encode a given input tensor with the specified sample rate at the given bandwidth."""
+        raise NotImplementedError()
+
+    def decode(self, codes: torch.Tensor) -> torch.Tensor:
+        """Decode the given codes to the quantized representation."""
+        raise NotImplementedError()
+
+    @property
+    def cardinality(self) -> int:
+        """Cardinality of each codebook."""
+        raise NotImplementedError()
+
+    @property
+    def total_codebooks(self) -> int:
+        """Total number of codebooks."""
+        raise NotImplementedError()
+
+    @property
+    def num_codebooks(self) -> int:
+        """Number of active codebooks."""
+        raise NotImplementedError()
+
+    @property
+    def semantic_quantizer(self) -> 'BaseQuantizer':
+        """This returns the quantizer that models the first level of the hierarchy (typically semantic).
+
+        In this case, it's the quantizer itself.
+        """
+        return self
+
+    @property
+    def acoustic_quantizer(self) -> 'BaseQuantizer':
+        """This returns the quantizer that models the higher levels of the hierarchy (typically acoustic).
+
+        In this case, it's the quantizer itself.
+        """
+        return self
+
+    def set_num_codebooks(self, n: int) -> None:
+        """Set the number of active codebooks."""
+        raise NotImplementedError()
+
+    @property
+    def ema_frozen(self) -> bool:
+        """Whether to apply ema to the codebooks."""
+        return self._ema_frozen
+
+    def ema_frozen_(self, ema_frozen: bool) -> None:
+        """Set whether ema should be applied to the codebooks."""
+        self._ema_frozen = ema_frozen
+
+
+class DummyQuantizer(BaseQuantizer):
+    """Fake quantizer that actually does not perform any quantization."""
+
+    def __init__(
+        self,
+        dimension: int,
+        input_dimension: tp.Optional[int] = None,
+        output_dimension: tp.Optional[int] = None,
+    ):
+        super().__init__()
+        self.dimension = dimension
+        self.input_dimension = input_dimension or dimension
+        self.output_dimension = output_dimension or dimension
+        self.input_proj: torch.nn.Module
+        self.output_proj: torch.nn.Module
+        if self.input_dimension == self.dimension:
+            self.input_proj = torch.nn.Identity()
+        else:
+            self.input_proj = torch.nn.Conv1d(
+                self.input_dimension, self.dimension, 1, bias=False
+            )
+        if self.input_dimension == self.dimension:
+            self.output_proj = torch.nn.Identity()
+        else:
+            self.output_proj = torch.nn.Conv1d(
+                self.dimension, self.output_dimension, 1, bias=False
+            )
+
+    def forward(self, x: torch.Tensor, frame_rate: int):
+        q = x.unsqueeze(1)
+        x = self.output_proj(self.input_proj(x))
+        return QuantizedResult(
+            x, q, torch.tensor(q.numel() * 32 * frame_rate / 1000 / len(x)).to(x)
+        )
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """Encode a given input tensor with the specified sample rate at the given bandwidth.
+        In the case of the DummyQuantizer, the codes are actually identical
+        to the input and resulting quantized representation as no quantization is done.
+        """
+        x = self.input_proj(x)
+        return x.unsqueeze(1)
+
+    def decode(self, codes: torch.Tensor) -> torch.Tensor:
+        """Decode the given codes to the quantized representation.
+        In the case of the DummyQuantizer, the codes are actually identical
+        to the input and resulting quantized representation as no quantization is done.
+        """
+        y = codes.squeeze(1)
+        return self.output_proj(y)
+
+    @property
+    def total_codebooks(self):
+        """Total number of codebooks."""
+        return 1
+
+    @property
+    def num_codebooks(self):
+        """Total number of codebooks."""
+        return self.total_codebooks
+
+    def set_num_codebooks(self, n: int):
+        """Set the number of active codebooks."""
+        raise AttributeError(
+            "Cannot override the number of codebooks for the dummy quantizer"
+        )
+
+    @property
+    def cardinality(self) -> int:
+        """Cardinality of each codebook."""
+        return 1

moshi/quantization/core_vq.py

@@ -0,0 +1,384 @@
+# Copyright (c) Kyutai, all rights reserved.
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import typing as tp
+
+from einops import rearrange
+import torch
+from torch import nn
+from torch import distributed
+import torch.nn.functional as F
+
+
+class _CodebookForwardResult(tp.NamedTuple):
+    quantized: torch.Tensor
+    codes: torch.Tensor
+    metrics: tp.Dict[str, torch.Tensor]
+
+
+class _VQForwardResult(tp.NamedTuple):
+    quantized: torch.Tensor
+    codes: torch.Tensor
+    loss: torch.Tensor
+    metrics: tp.Dict[str, torch.Tensor]
+
+
+def _ema_inplace(moving_avg: torch.Tensor, new: torch.Tensor, decay: float) -> None:
+    moving_avg.data.mul_(decay).add_(new, alpha=(1 - decay))
+
+
+def _uniform_init(*shape: int) -> torch.Tensor:
+    t = torch.empty(shape)
+    nn.init.kaiming_uniform_(t)
+    return t
+
+
+def _sample_vectors(samples: torch.Tensor, num: int) -> torch.Tensor:
+    num_samples, device = samples.shape[0], samples.device
+
+    if num_samples >= num:
+        indices = torch.randperm(num_samples, device=device)[:num]
+    else:
+        indices = torch.randint(0, num_samples, (num,), device=device)
+
+    return samples[indices]
+
+
+def _compute_entropy(usage: torch.Tensor) -> torch.Tensor:
+    # Usage is some unnormalized distribution.
+    proba = usage / usage.sum()
+    p_log_p = torch.where(
+        proba == 0, zero_scalar(usage.device), proba * torch.log(proba)
+    )
+    return -p_log_p.sum()
+
+
+def _is_distributed() -> bool:
+    # Checks if we need to use distributed routines.
+    return distributed.is_initialized() and distributed.get_world_size() > 1
+
+
+def zero_scalar(device) -> torch.Tensor:
+    """Returns a 0. value on the given device without introducing a synchronization point."""
+    return torch.zeros([1], device=device)[0]
+
+
+class EuclideanCodebook(nn.Module):
+    """Codebook with Euclidean distance.
+
+    Args:
+        dim (int): Dimension.
+        codebook_size (int): Codebook size.
+        decay (float): Decay for exponential moving average over the codebooks.
+        epsilon (float): Epsilon value for numerical stability.
+        threshold_usage_ratio (float): Defines the threshold for the cluster usage under which a centroid
+            is replaced. This is expressed as a fraction of the usage a centroid would get under
+            a uniform distribution, so that it doesn't depend on the batch size etc.
+        replaced_usage_ratio (float): When replacing a centroid, use this as an initial centroid usage,
+            to avoid the centroid getting replaced too quickly.
+        check_unused_every (int): Check for unused centroids every `check_unused_every` iterations.
+            This is to avoid too many synchronization points.
+
+    Buffers:
+        cluster_usage (torch.Tensor): EMA of the cluster usage per batch, e.g. this will
+            be dependent on the batch size etc.
+        embedding_sum (torch.Tensor): EMA of the sum of the assigned points to each cluster.
+            In particular, this can be normalized by `cluster_usage` to obtain the
+            actual cluster centroids.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        codebook_size: int,
+        decay: float = 0.99,
+        epsilon: float = 1e-5,
+        threshold_usage_ratio: float = 0.1,
+        replaced_usage_ratio: float = 1.0,
+        check_unused_every: int = 5,
+    ):
+        super().__init__()
+        self.decay = decay
+        embedding = torch.zeros(codebook_size, dim)
+
+        self.dim = dim
+        self.codebook_size = codebook_size
+
+        self.epsilon = epsilon
+        self.threshold_usage_ratio = threshold_usage_ratio
+        self.replaced_usage_ratio = replaced_usage_ratio
+        self.check_unused_every = check_unused_every
+        self._next_unused_check = check_unused_every
+
+        self.register_buffer("_initialized", torch.tensor([False], dtype=torch.float))
+        self.register_buffer("cluster_usage", torch.ones(codebook_size))
+        self.register_buffer("embedding_sum", embedding)
+        self.register_buffer("_embedding", None, persistent=False)
+        self._cached_initialized = False
+
+    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs) -> None:
+        # Mapping old names to new names
+        mappings = {
+            "inited": "_initialized",
+            "cluster_size": "cluster_usage",
+            "embed_avg": "embedding_sum",
+            "embed_sum": "embedding_sum",
+        }
+        for old_name, new_name in mappings.items():
+            old_name = prefix + old_name
+            if old_name in state_dict:
+                value = state_dict.pop(old_name)
+                if new_name is not None:
+                    state_dict[prefix + new_name] = value
+        super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
+
+    @property
+    def embedding(self) -> torch.Tensor:
+        if self._embedding is None:
+            embedding = (
+                self.embedding_sum / self.cluster_usage.clamp(min=self.epsilon)[:, None]
+            )
+            self.register_buffer("_embedding", embedding, persistent=False)
+            return embedding
+        return self._embedding
+
+    def _broadcast_buffers(self) -> None:
+        if _is_distributed():
+            for buffer in self.buffers():
+                distributed.broadcast(buffer, 0)
+
+    def _replace_expired_codes(self, samples: torch.Tensor, mask: torch.Tensor) -> None:
+        # Replaces expired centroids, as indicated by `mask` (a true value indicate the code needs to be replaced).
+        # The new codes are sampled from the batch `samples`.
+        new_vectors = _sample_vectors(samples, self.codebook_size)
+        replace_cluster_usage = (
+            self.replaced_usage_ratio * self.cluster_usage.sum() / self.codebook_size
+        )
+        self.embedding_sum[:] = torch.where(
+            mask[:, None], replace_cluster_usage * new_vectors, self.embedding_sum
+        )
+        self.cluster_usage[:] = torch.where(
+            mask, replace_cluster_usage, self.cluster_usage
+        )
+
+    def _reshape_input(self, x: torch.Tensor) -> torch.Tensor:
+        # Flattens all the dimensions but the last one, e.g. return a vector of shape `[N, D]`.
+        x = rearrange(x, "... d -> (...) d")
+        return x
+
+    def _reshape_codes(self, codes: torch.Tensor, shape: torch.Size) -> torch.Tensor:
+        return codes.view(*shape[:-1])
+
+    def _quantize(self, x: torch.Tensor) -> torch.Tensor:
+        # Projects each vector in `x` over the nearest centroid and return its index.
+        # `x` should be `[N, D]` with `N` the number of input vectors and `D` the dimension.
+        assert x.dim() == 2
+        dists = torch.cdist(x[None], self.embedding[None], p=2)[0]
+        codes = dists.argmin(dim=-1)
+        return codes
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """Given a tensor `x` of shape `[*, D]`, returns a tensor of integer codes of shape `[*]`.
+        The codes are defined as the indexes of the centroids nearest to each vector in `x`.
+        """
+        assert x.dtype.is_floating_point, f"Input should be floats, got {x.dtype}"
+        shape = x.shape
+        x = self._reshape_input(x)
+        codes = self._quantize(x)
+        codes = self._reshape_codes(codes, shape)
+        return codes
+
+    def decode(self, codes: torch.Tensor) -> torch.Tensor:
+        """Given a tensor of codes of shape `[*]`, returns a tensor of shape `[*, D]`,
+        corresponding to the centroids associated to each code index.
+        """
+        assert (
+            not codes.dtype.is_floating_point
+        ), f"Codes should be integers, got {codes.dtype}"
+        quantized = F.embedding(codes, self.embedding)
+        return quantized
+
+    def forward(
+        self, x: torch.Tensor, initialize: bool = True
+    ) -> _CodebookForwardResult:
+        shape = x.shape
+        x = self._reshape_input(x)
+
+        flat_codes = self._quantize(x)
+        codes = self._reshape_codes(flat_codes, shape)
+        quantized = self.decode(codes)
+        metrics: tp.Dict[str, torch.Tensor] = {}
+
+        return _CodebookForwardResult(quantized, codes, metrics)
+
+
+class VectorQuantization(nn.Module):
+    """Vector quantization implementation.
+    Currently supports only euclidean distance.
+
+    Args:
+        dim (int): Dimension
+        codebook_size (int): Codebook size
+        codebook_dim (int): Codebook dimension. If not defined, uses the specified dimension in dim.
+        decay (float): Decay for exponential moving average over the codebooks.
+        epsilon (float): Epsilon value for numerical stability.
+        threshold_usage_ratio (float): Defines the threshold for the cluster usage under which a centroid
+            is replaced. This is expressed as a fraction of the usage a centroid would get under
+            a uniform distribution, so that it doesn't depend on the batch size etc.
+        replaced_usage_ratio (float): When replacing a centroid, use this as an initial centroid usage,
+            to avoid the centroid getting replaced too quickly.
+        check_unused_every (int): Check for unused centroids every `check_unused_every` iterations.
+            This is to avoid too many synchronization points.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        codebook_size: int,
+        codebook_dim: tp.Optional[int] = None,
+        decay: float = 0.99,
+        epsilon: float = 1e-5,
+        threshold_usage_ratio: float = 0.1,
+        **kwargs,
+    ):
+        super().__init__()
+        if codebook_dim is None:
+            codebook_dim = dim
+
+        requires_projection = codebook_dim != dim
+        self.project_in = (
+            nn.Linear(dim, codebook_dim) if requires_projection else nn.Identity()
+        )
+        self.project_out = (
+            nn.Linear(codebook_dim, dim) if requires_projection else nn.Identity()
+        )
+        self.epsilon = epsilon
+        self._codebook = EuclideanCodebook(
+            dim=codebook_dim,
+            codebook_size=codebook_size,
+            decay=decay,
+            epsilon=epsilon,
+            threshold_usage_ratio=threshold_usage_ratio,
+            **kwargs,
+        )
+        self.codebook_size = codebook_size
+
+    @property
+    def embedding(self):
+        return self._codebook.embedding
+
+    def _rearrange_input(self, x):
+        x = rearrange(x, "b d n -> b n d")
+        return x
+
+    def _rearrange_output(self, quantized):
+        quantized = rearrange(quantized, "b n d -> b d n")
+        return quantized
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """Encodes `x` into discrete integer codes."""
+        x = self._rearrange_input(x)
+        x = self.project_in(x)
+        codes = self._codebook.encode(x)
+        return codes
+
+    def decode(self, codes: torch.Tensor) -> torch.Tensor:
+        """Converts integer codes into quantized vectors."""
+        quantized = self._codebook.decode(codes)
+        quantized = self.project_out(quantized)
+        quantized = self._rearrange_output(quantized)
+        return quantized
+
+    def forward(self, x: torch.Tensor, initialize: bool = True) -> _VQForwardResult:
+        x = self._rearrange_input(x)
+        quantized, codes, metrics = self._codebook(x, initialize=initialize)
+
+        loss = zero_scalar(x.device)
+
+        quantized = self.project_out(quantized)
+        quantized = self._rearrange_output(quantized)
+
+        return _VQForwardResult(quantized, codes, loss, metrics)
+
+
+class ResidualVectorQuantization(nn.Module):
+    """Residual vector quantization implementation.
+
+    Follows Algorithm 1. in https://arxiv.org/pdf/2107.03312.pdf
+    """
+
+    def __init__(self, *, num_quantizers: int, codebook_offset: int, **kwargs):
+        super().__init__()
+        self.layers = nn.ModuleList(
+            [VectorQuantization(**kwargs) for _ in range(num_quantizers)]
+        )
+        self.codebook_offset = codebook_offset
+
+    def forward(
+        self, x: torch.Tensor, n_q: tp.Optional[int] = None
+    ) -> _VQForwardResult:
+        """
+        Args:
+            x (torch.Tensor): input tensor to quantize, of shape `[B, C, T]`.
+            n_q (int or None): if provided, number of codebook levels to use in RVQ.
+        """
+
+        quantized_out = zero_scalar(x.device)
+        residual = x
+
+        all_losses = []
+        all_codes = []
+        all_metrics: tp.Dict[str, torch.Tensor] = {}
+
+        n_q = n_q or len(self.layers)
+        previous_layer_is_initialized = True
+
+        for i, layer in enumerate(self.layers[:n_q]):  # type: ignore
+            quantized, codes, loss, metrics = layer(
+                residual, initialize=previous_layer_is_initialized
+            )
+
+            quantized = quantized.detach()
+            residual = residual - quantized
+            quantized_out = quantized_out + quantized
+
+            all_codes.append(codes)
+            all_losses.append(loss)
+
+            for key, value in metrics.items():
+                if key in all_metrics:
+                    all_metrics[key] += value / n_q
+                else:
+                    all_metrics[key] = value / n_q
+                all_metrics[key + f"_{i + self.codebook_offset}"] = value
+
+        out_losses, out_codes = map(torch.stack, (all_losses, all_codes))
+        return _VQForwardResult(quantized_out, out_codes, out_losses, all_metrics)
+
+    def encode(self, x: torch.Tensor, n_q: tp.Optional[int] = None) -> torch.Tensor:
+        """Encodes `x` into discrete integer codes. If `n_q` is provided, only uses the first `n_q` codebook levels."""
+        residual = x
+        all_indices = []
+        n_q = n_q or len(self.layers)
+        for layer in self.layers[:n_q]:  # type: ignore
+            indices = layer.encode(residual)
+            quantized = layer.decode(indices)
+            residual = residual - quantized
+            all_indices.append(indices)
+        out_indices = torch.stack(all_indices)
+        return out_indices
+
+    def decode(self, codes: torch.Tensor) -> torch.Tensor:
+        """Converts the integer codes into quantized vectors."""
+        quantized = zero_scalar(codes.device)
+        for idx, layer_codes in enumerate(codes):
+            layer = self.layers[idx]
+            quantized = quantized + layer.decode(layer_codes)
+        return quantized

moshi/quantization/vq.py

@@ -0,0 +1,361 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: MIT
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+
+# Copyright (c) Kyutai, all rights reserved.
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+import typing as tp
+
+import torch
+
+from .base import BaseQuantizer, QuantizedResult
+from .core_vq import ResidualVectorQuantization
+
+
+class ResidualVectorQuantizer(BaseQuantizer):
+    """Residual Vector Quantizer.
+
+    Args:
+        dimension (int): Dimension of the codebooks.
+        input_dimension (None or int): dimension of the input, defaults to `dimension` if not provided.
+        output_dimension (None or int): dimension of the output, defaults to `dimension` if not provided.
+        n_q (int): Number of vector quantizers used.
+        q_dropout (bool): Random quantizer drop out at train time.
+        no_quantization_rate (float): Gives the probability of applying no quantization at all
+            at train time. The RVQ codebooks will still get the input value to learn the proper codebook.
+        bins (int): Codebook size.
+        decay (float): Decay for exponential moving average over the codebooks.
+        threshold_usage_ratio (float): Defines the threshold for the cluster usage under which a centroid
+            is replaced. This is expressed as a fraction of the usage a centroid would get under
+            a uniform distribution, so that it doesn't depend on the batch size etc.
+        replaced_usage_ratio (float): When replacing a centroid, use this as an initial centroid usage,
+            to avoid the centroid getting replaced too quickly.
+        codebook_offset (int): Offset to use for the codebook indices. This is useful when using multiple quantizers
+            such as in SplitResidualVectorQuantizer.
+        force_projection (bool): Whether to force input and output projections even when dimension is constant.
+        generator_seed (int or None): seed used to initialize the RNG used for no quantization.
+    """
+
+    def __init__(
+        self,
+        dimension: int = 128,
+        input_dimension: tp.Optional[int] = None,
+        output_dimension: tp.Optional[int] = None,
+        n_q: int = 8,
+        q_dropout: bool = False,
+        q_first_only_proba: float = 0.0,
+        no_quantization_rate: float = 0.0,
+        bins: int = 1024,
+        decay: float = 0.99,
+        threshold_usage_ratio: float = 0.1,
+        replaced_usage_ratio: float = 1.0,
+        codebook_offset: int = 0,
+        force_projection: bool = False,
+        generator_seed: tp.Optional[int] = None,
+    ):
+        super().__init__()
+        self.max_n_q = n_q
+        self.n_q = n_q
+        self.q_dropout = q_dropout
+        self.no_quantization_rate = no_quantization_rate
+        self.q_first_only_proba = q_first_only_proba
+        self.dimension = dimension
+        self.input_dimension = input_dimension or dimension
+        self.output_dimension = output_dimension or dimension
+        self.bins = bins
+        self.decay = decay
+        self.input_proj: torch.nn.Module
+        self.output_proj: torch.nn.Module
+        self.generator = None
+        if generator_seed is not None:
+            self.generator = torch.Generator(
+                device="cuda" if torch.cuda.is_available() else "cpu"
+            )
+            self.generator.manual_seed(generator_seed)
+        if self.input_dimension == self.dimension and not force_projection:
+            self.input_proj = torch.nn.Identity()
+        else:
+            self.input_proj = torch.nn.Conv1d(
+                self.input_dimension, self.dimension, 1, bias=False
+            )
+        if self.output_dimension == self.dimension and not force_projection:
+            self.output_proj = torch.nn.Identity()
+        else:
+            self.output_proj = torch.nn.Conv1d(
+                self.dimension, self.output_dimension, 1, bias=False
+            )
+        self.vq = ResidualVectorQuantization(
+            dim=self.dimension,
+            codebook_size=self.bins,
+            num_quantizers=self.n_q,
+            decay=self.decay,
+            threshold_usage_ratio=threshold_usage_ratio,
+            replaced_usage_ratio=replaced_usage_ratio,
+            codebook_offset=codebook_offset,
+        )
+
+    def forward(self, x: torch.Tensor, frame_rate: int):
+        """
+        Args:
+            x (torch.Tensor): Input tensor of shape [B, C, T] with `C` number of channels.
+            frame_rate (int): frame rate of the input (e.g `T = frame_rate * duration`), used to compute
+                the bandwidth.
+
+        Returns:
+            QuantizedResult: Quantized result with the following attributes:
+                - `x` (torch.Tensor): Quantized tensor of shape [B, C, T].
+                - `codes` (torch.Tensor): Quantized codes of shape [B, K, T] with `K` number of codebooks.
+                - `bw` (torch.Tensor): Bandwidth of the quantized tensor in kbits per second.
+                - `penalty` (torch.Tensor): Commitment loss.
+                - `metrics` (dict): RVQ metrics, in particular rate of dead code replacement, and entropy.
+        """
+        n_q = self.n_q
+        x = self.input_proj(x)
+
+        bw_per_q = math.log2(self.bins) * frame_rate / 1000
+        quantized, codes, commit_loss, metrics = self.vq(x, n_q=n_q)
+        B, _, _ = quantized.shape
+        quantized = self.output_proj(quantized)
+        codes = codes.transpose(0, 1)
+        # codes is [B, K, T], with T frames, K nb of codebooks.
+        bw = torch.tensor(n_q * bw_per_q).to(x)
+        return QuantizedResult(
+            quantized, codes, bw, penalty=torch.mean(commit_loss), metrics=metrics
+        )
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """Encode a given input tensor with the specified frame rate at the given bandwidth.
+        The RVQ encode method sets the appropriate number of quantizer to use
+        and returns indices for each quantizer.
+        """
+        n_q = self.n_q
+        if x.shape[-1] == 0:
+            return torch.empty((x.shape[0], n_q, 0), device=x.device, dtype=torch.int64)
+
+        x = self.input_proj(x)
+        codes = self.vq.encode(x, n_q=n_q)
+        codes = codes.transpose(0, 1)
+        # codes is [B, K, T], with T frames, K nb of codebooks.
+        return codes
+
+    def decode(self, codes: torch.Tensor) -> torch.Tensor:
+        """Decode the given codes to the quantized representation."""
+        # codes is [B, K, T], with T frames, K nb of codebooks, vq.decode expects [K, B, T].
+        codes = codes.transpose(0, 1)
+        quantized = self.vq.decode(codes)
+        quantized = self.output_proj(quantized)
+        return quantized
+
+    @property
+    def total_codebooks(self):
+        return self.max_n_q
+
+    @property
+    def num_codebooks(self):
+        return self.n_q
+
+    def set_num_codebooks(self, n: int):
+        assert n >= 0 and n <= self.max_n_q
+        self.n_q = n
+
+    @property
+    def cardinality(self) -> int:
+        return self.bins
+
+
+class SplitResidualVectorQuantizer(BaseQuantizer):
+    """Residual Vector Quantizer with separate projections for the first quantizer and the rest.
+
+    Args:
+        n_q (int): Number of residual vector quantizers used.
+        n_semantic_q (int): Number of residual vector quantizers used for the semantic quantizer.
+        no_quantization_mode (str): if 'true_skip', when doing no quantization, the input will not go
+            through the sub quantizers. If `independent`, independent decisions are taken by
+            the semantic and acoustic quantizers. If `same` (the default), the same decision is taken by both.
+        **kwargs: Arguments to the constructor of `ResidualVectorQuantizer` that are shared between both.
+    """
+
+    def __init__(
+        self,
+        *,
+        n_q: int = 8,
+        no_quantization_rate: float = 0.0,
+        no_quantization_mode: str = "same",
+        n_q_semantic: int = 1,
+        **kwargs,
+    ):
+        super().__init__()
+        assert n_q > n_q_semantic, (
+            f"Number of quantizers {n_q} must be larger "
+            f"than the number of semantic quantizers {n_q_semantic}."
+        )
+        self.max_n_q = n_q
+        self.n_q_semantic = n_q_semantic
+        self.n_q_acoustic = n_q - n_q_semantic
+        if no_quantization_mode == "true_skip":
+            self.no_quantization_rate = no_quantization_rate
+            # Setting to zero for the underlying RVQ.
+            no_quantization_rate = 0.0
+        else:
+            self.no_quantization_rate = 0.0
+        if no_quantization_mode == "same":
+            kwargs["generator_seed"] = 1234
+        kwargs["no_quantization_rate"] = no_quantization_rate
+        q_dropout = kwargs.pop("q_dropout", False)
+        self.rvq_first = ResidualVectorQuantizer(
+            n_q=n_q_semantic, force_projection=True, q_dropout=False, **kwargs
+        )
+        self.rvq_rest = ResidualVectorQuantizer(
+            n_q=n_q - n_q_semantic,
+            codebook_offset=1,
+            force_projection=True,
+            q_dropout=q_dropout,
+            **kwargs,
+        )
+        if no_quantization_mode == "true_skip":
+            assert self.rvq_first.input_dimension == self.rvq_first.output_dimension
+            assert self.rvq_rest.input_dimension == self.rvq_rest.output_dimension
+
+    def _renorm_and_add(
+        self,
+        first_val: torch.Tensor,
+        rest_val: torch.Tensor,
+        n_q_semantic: int,
+        n_q_acoustic: int,
+    ):
+        """Renormalizes values from `rvq_first` and `rvq_rest` and adds them.
+
+        This allows correcting statistics that are normalized by the number of quantizers. To renormalize, we use the
+        number of quantizers that are actually used, e.g. taking into account quantizer dropout.
+        """
+        n_q = n_q_semantic + n_q_acoustic
+        renorm_first_val = first_val * n_q_semantic / n_q
+        renorm_rest_val = rest_val * n_q_acoustic / n_q
+        return renorm_first_val + renorm_rest_val
+
+    def forward(self, x: torch.Tensor, frame_rate: int):
+        """
+        Args:
+            x (torch.Tensor): Input tensor of shape [B, C, T] with `C` number of channels.
+            frame_rate (int): frame rate of the input (e.g `T = frame_rate * duration`), used to compute
+                the bandwidth.
+
+        Returns:
+            QuantizedResult: Quantized result with the following attributes:
+                - `x` (torch.Tensor): Quantized tensor of shape [B, C, T].
+                - `codes` (torch.Tensor): Quantized codes of shape [B, K, T] with `K` number of codebooks.
+                - `bw` (torch.Tensor): Bandwidth of the quantized tensor in kbits per second.
+                - `penalty` (torch.Tensor): Commitment loss.
+                - `metrics` (dict): RVQ metrics, in particular rate of dead code replacement, and entropy.
+        """
+        semantic_result = self.rvq_first(x, frame_rate)
+        if self.n_q == self.n_q_semantic:
+            return semantic_result
+        acoustic_result = self.rvq_rest(x, frame_rate)
+        full_quantized_emb = semantic_result.x + acoustic_result.x
+        full_quantized_codes = torch.cat(
+            [semantic_result.codes, acoustic_result.codes], dim=1
+        )
+        # This is the actual number of quantizers used,  e.g. taking into account quantizer dropout.
+        n_q_semantic = semantic_result.codes.shape[1]
+        n_q_acoustic = acoustic_result.codes.shape[1]
+        full_quantized_bandwidth = semantic_result.bandwidth + acoustic_result.bandwidth
+        full_quantized_penalty = self._renorm_and_add(
+            semantic_result.penalty, acoustic_result.penalty, n_q_semantic, n_q_acoustic
+        )
+        full_quantized_metrics = semantic_result.metrics
+        for key, value in acoustic_result.metrics.items():
+            if key in full_quantized_metrics:
+                full_quantized_metrics[key] = self._renorm_and_add(
+                    full_quantized_metrics[key], value, n_q_semantic, n_q_acoustic
+                )
+            else:
+                full_quantized_metrics[key] = value
+        return QuantizedResult(
+            full_quantized_emb,
+            full_quantized_codes,
+            full_quantized_bandwidth,
+            penalty=full_quantized_penalty,
+            metrics=full_quantized_metrics,
+        )
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """Encode a given input tensor with the specified frame rate at the given bandwidth.
+        The RVQ encode method sets the appropriate number of quantizer to use
+        and returns indices for each quantizer.
+        """
+        codes = self.rvq_first.encode(x)
+        if self.n_q > self.n_q_semantic:
+            acoustic_codes = self.rvq_rest.encode(x)
+            codes = torch.cat([codes, acoustic_codes], dim=1)
+        # codes is [B, K, T], with T frames, K nb of codebooks.
+        return codes
+
+    def decode(self, codes: torch.Tensor) -> torch.Tensor:
+        """Decode the given codes to the quantized representation."""
+        # codes is [B, K, T], with T frames, K nb of codebooks.
+        quantized = self.rvq_first.decode(codes[:, : self.n_q_semantic])
+        if codes.shape[1] > self.n_q_semantic:
+            quantized += self.rvq_rest.decode(codes[:, self.n_q_semantic :])
+        return quantized
+
+    @property
+    def total_codebooks(self):
+        return self.rvq_first.max_n_q + self.rvq_rest.max_n_q
+
+    @property
+    def num_codebooks(self):
+        return self.rvq_first.num_codebooks + self.rvq_rest.num_codebooks
+
+    @property
+    def n_q(self):
+        return self.rvq_first.n_q + self.rvq_rest.n_q
+
+    @property
+    def dimension(self):
+        return self.rvq_first.dimension
+
+    @property
+    def semantic_quantizer(self) -> ResidualVectorQuantizer:
+        """This returns the quantizer that models the first level of the hierarchy (typically semantic)."""
+        return self.rvq_first
+
+    @property
+    def acoustic_quantizer(self) -> ResidualVectorQuantizer:
+        """This returns the quantizer that models the higher levels of the hierarchy (typically acoustic)."""
+        return self.rvq_rest
+
+    def set_num_codebooks(self, n: int):
+        assert n >= self.n_q_semantic and n <= self.total_codebooks
+        self.rvq_rest.set_num_codebooks(n - self.n_q_semantic)
+
+    @property
+    def cardinality(self) -> int:
+        assert self.rvq_rest.cardinality == self.rvq_first.cardinality
+        return self.rvq_first.cardinality

moshi/server.py

@@ -0,0 +1,483 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: MIT
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+
+
+# Copyright (c) Kyutai, all rights reserved.
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import argparse
+import asyncio
+from dataclasses import dataclass
+import random
+import os
+from pathlib import Path
+import tarfile
+import time
+import secrets
+import sys
+from typing import Literal, Optional
+
+import aiohttp
+from aiohttp import web
+from huggingface_hub import hf_hub_download
+import numpy as np
+import sentencepiece
+import sphn
+import torch
+import random
+
+from .client_utils import make_log, colorize
+from .models import loaders, MimiModel, LMModel, LMGen
+from .utils.connection import create_ssl_context, get_lan_ip
+from .utils.logging import setup_logger, ColorizedLog
+
+
+logger = setup_logger(__name__)
+DeviceString = Literal["cuda"] | Literal["cpu"] #| Literal["mps"]
+
+def torch_auto_device(requested: Optional[DeviceString] = None) -> torch.device:
+    """Return a torch.device based on the requested string or availability."""
+    if requested is not None:
+        return torch.device(requested)
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    #elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+    #    return torch.device("mps")
+    return torch.device("cpu")
+
+
+def seed_all(seed):
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)  # for multi-GPU setups
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.backends.cudnn.deterministic = False
+    torch.backends.cudnn.benchmark = False
+
+
+def wrap_with_system_tags(text: str) -> str:
+    """Add system tags as the model expects if they are missing.
+    Example: "<system> You enjoy having a good conversation. Have a deep conversation about technology. Your name is Jane. <system>"
+    """
+    cleaned = text.strip()
+    if cleaned.startswith("<system>") and cleaned.endswith("<system>"):
+        return cleaned
+    return f"<system> {cleaned} <system>"
+
+
+@dataclass
+class ServerState:
+    mimi: MimiModel
+    other_mimi: MimiModel
+    text_tokenizer: sentencepiece.SentencePieceProcessor
+    lm_gen: LMGen
+    lock: asyncio.Lock
+
+    def __init__(self, mimi: MimiModel, other_mimi: MimiModel, text_tokenizer: sentencepiece.SentencePieceProcessor,
+                 lm: LMModel, device: str | torch.device, voice_prompt_dir: str | None = None,
+                 save_voice_prompt_embeddings: bool = False):
+        self.mimi = mimi
+        self.other_mimi = other_mimi
+        self.text_tokenizer = text_tokenizer
+        self.device = device
+        self.voice_prompt_dir = voice_prompt_dir
+        self.frame_size = int(self.mimi.sample_rate / self.mimi.frame_rate)
+        self.lm_gen = LMGen(lm,
+                            audio_silence_frame_cnt=int(0.5 * self.mimi.frame_rate),
+                            sample_rate=self.mimi.sample_rate,
+                            device=device,
+                            frame_rate=self.mimi.frame_rate,
+                            save_voice_prompt_embeddings=save_voice_prompt_embeddings,
+        )
+        
+        self.lock = asyncio.Lock()
+        self.mimi.streaming_forever(1)
+        self.other_mimi.streaming_forever(1)
+        self.lm_gen.streaming_forever(1)
+    
+    def warmup(self):
+        for _ in range(4):
+            chunk = torch.zeros(1, 1, self.frame_size, dtype=torch.float32, device=self.device)
+            codes = self.mimi.encode(chunk)
+            _ = self.other_mimi.encode(chunk)
+            for c in range(codes.shape[-1]):
+                tokens = self.lm_gen.step(codes[:, :, c: c + 1])
+                if tokens is None:
+                    continue
+                _ = self.mimi.decode(tokens[:, 1:9])
+                _ = self.other_mimi.decode(tokens[:, 1:9])
+
+        if self.device.type == 'cuda':
+            torch.cuda.synchronize()
+
+
+    async def handle_chat(self, request):
+        ws = web.WebSocketResponse()
+        await ws.prepare(request)
+        clog = ColorizedLog.randomize()
+        peer = request.remote  # IP
+        peer_port = request.transport.get_extra_info("peername")[1]  # Port
+        clog.log("info", f"Incoming connection from {peer}:{peer_port}")
+
+        # self.lm_gen.temp = float(request.query["audio_temperature"])
+        # self.lm_gen.temp_text = float(request.query["text_temperature"])
+        # self.lm_gen.top_k_text = max(1, int(request.query["text_topk"]))
+        # self.lm_gen.top_k = max(1, int(request.query["audio_topk"]))
+        
+        # Construct full voice prompt path
+        requested_voice_prompt_path = None
+        voice_prompt_path = None
+        if self.voice_prompt_dir is not None:
+            voice_prompt_filename = request.query["voice_prompt"]
+            requested_voice_prompt_path = None
+            if voice_prompt_filename is not None:
+                requested_voice_prompt_path = os.path.join(self.voice_prompt_dir, voice_prompt_filename)
+            # If the voice prompt file does not exist, find a valid (s0) voiceprompt file in the directory
+            if requested_voice_prompt_path is None or not os.path.exists(requested_voice_prompt_path):
+                raise FileNotFoundError(
+                    f"Requested voice prompt '{voice_prompt_filename}' not found in '{self.voice_prompt_dir}'"
+                )
+            else:
+                voice_prompt_path = requested_voice_prompt_path
+                
+        if self.lm_gen.voice_prompt != voice_prompt_path:
+            if voice_prompt_path.endswith('.pt'):
+                # Load pre-saved voice prompt embeddings
+                self.lm_gen.load_voice_prompt_embeddings(voice_prompt_path)
+            else:
+                self.lm_gen.load_voice_prompt(voice_prompt_path)
+        self.lm_gen.text_prompt_tokens = self.text_tokenizer.encode(wrap_with_system_tags(request.query["text_prompt"])) if len(request.query["text_prompt"]) > 0 else None
+        seed = int(request["seed"]) if "seed" in request.query else None
+
+        async def recv_loop():
+            nonlocal close
+            try:
+                async for message in ws:
+                    if message.type == aiohttp.WSMsgType.ERROR:
+                        clog.log("error", f"{ws.exception()}")
+                        break
+                    elif message.type == aiohttp.WSMsgType.CLOSED:
+                        break
+                    elif message.type == aiohttp.WSMsgType.CLOSE:
+                        break
+                    elif message.type != aiohttp.WSMsgType.BINARY:
+                        clog.log("error", f"unexpected message type {message.type}")
+                        continue
+                    message = message.data
+                    if not isinstance(message, bytes):
+                        clog.log("error", f"unsupported message type {type(message)}")
+                        continue
+                    if len(message) == 0:
+                        clog.log("warning", "empty message")
+                        continue
+                    kind = message[0]
+                    if kind == 1:  # audio
+                        payload = message[1:]
+                        opus_reader.append_bytes(payload)
+                    else:
+                        clog.log("warning", f"unknown message kind {kind}")
+            finally:
+                close = True
+                clog.log("info", "connection closed")
+
+        async def opus_loop():
+            all_pcm_data = None
+
+            while True:
+                if close:
+                    return
+                await asyncio.sleep(0.001)
+                pcm = opus_reader.read_pcm()
+                if pcm.shape[-1] == 0:
+                    continue
+                if all_pcm_data is None:
+                    all_pcm_data = pcm
+                else:
+                    all_pcm_data = np.concatenate((all_pcm_data, pcm))
+                while all_pcm_data.shape[-1] >= self.frame_size:
+                    be = time.time()
+                    chunk = all_pcm_data[: self.frame_size]
+                    all_pcm_data = all_pcm_data[self.frame_size:]
+                    chunk = torch.from_numpy(chunk)
+                    chunk = chunk.to(device=self.device)[None, None]
+                    codes = self.mimi.encode(chunk)
+                    _ = self.other_mimi.encode(chunk)
+                    for c in range(codes.shape[-1]):
+                        tokens = self.lm_gen.step(codes[:, :, c: c + 1])
+                        if tokens is None:
+                            continue
+                        assert tokens.shape[1] == self.lm_gen.lm_model.dep_q + 1
+                        main_pcm = self.mimi.decode(tokens[:, 1:9])
+                        _ = self.other_mimi.decode(tokens[:, 1:9])
+                        main_pcm = main_pcm.cpu()
+                        opus_writer.append_pcm(main_pcm[0, 0].numpy())
+                        text_token = tokens[0, 0, 0].item()
+                        if text_token not in (0, 3):
+                            _text = self.text_tokenizer.id_to_piece(text_token)  # type: ignore
+                            _text = _text.replace("▁", " ")
+                            msg = b"\x02" + bytes(_text, encoding="utf8")
+                            await ws.send_bytes(msg)
+                        else:
+                            text_token_map = ['EPAD', 'BOS', 'EOS', 'PAD']
+
+        async def send_loop():
+            while True:
+                if close:
+                    return
+                await asyncio.sleep(0.001)
+                msg = opus_writer.read_bytes()
+                if len(msg) > 0:
+                    await ws.send_bytes(b"\x01" + msg)
+
+        clog.log("info", "accepted connection")
+        if len(request.query["text_prompt"]) > 0:
+            clog.log("info", f"text prompt: {request.query['text_prompt']}")
+        if len(request.query["voice_prompt"]) > 0:
+            clog.log("info", f"voice prompt: {voice_prompt_path} (requested: {requested_voice_prompt_path})")
+        close = False
+        async with self.lock:
+            if seed is not None and seed != -1:
+                seed_all(seed)
+
+            opus_writer = sphn.OpusStreamWriter(self.mimi.sample_rate)
+            opus_reader = sphn.OpusStreamReader(self.mimi.sample_rate)
+            self.mimi.reset_streaming()
+            self.other_mimi.reset_streaming()
+            self.lm_gen.reset_streaming()
+            async def is_alive():
+                if close or ws.closed:
+                    return False
+                try:
+                    # Check for disconnect without waiting too long
+                    msg = await asyncio.wait_for(ws.receive(), timeout=0.01)
+                    if msg.type in (aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.CLOSED, aiohttp.WSMsgType.ERROR):
+                        return False
+                except asyncio.TimeoutError:
+                    # No messages → client probably still alive
+                    return True
+                except aiohttp.ClientConnectionError:
+                    return False
+                return True
+            # Reuse mimi for encoding voice prompt and then reset it before conversation starts
+            await self.lm_gen.step_system_prompts_async(self.mimi, is_alive=is_alive)
+            self.mimi.reset_streaming()
+            clog.log("info", "done with system prompts")
+            # Send the handshake.
+            if await is_alive():
+                await ws.send_bytes(b"\x00")
+                clog.log("info", "sent handshake bytes")
+                # Clean cancellation manager
+                tasks = [
+                    asyncio.create_task(recv_loop()),
+                    asyncio.create_task(opus_loop()),
+                    asyncio.create_task(send_loop()),
+                ]
+
+                done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
+                # Force-kill remaining tasks
+                for task in pending:
+                    task.cancel()
+                    try:
+                        await task
+                    except asyncio.CancelledError:
+                        pass
+                await ws.close()
+                clog.log("info", "session closed")
+                # await asyncio.gather(opus_loop(), recv_loop(), send_loop())
+        clog.log("info", "done with connection")
+        return ws
+
+
+def _get_voice_prompt_dir(voice_prompt_dir: Optional[str], hf_repo: str) -> Optional[str]:
+    """
+    If voice_prompt_dir is None:
+      - download voices.tgz from HF
+      - extract it once
+      - return extracted directory
+    If voice_prompt_dir is provided:
+      - just return it
+    """
+    if voice_prompt_dir is not None:
+        return voice_prompt_dir
+
+    logger.info("retrieving voice prompts")
+
+    voices_tgz = hf_hub_download(hf_repo, "voices.tgz")
+    voices_tgz = Path(voices_tgz)
+    voices_dir = voices_tgz.parent / "voices"
+
+    if not voices_dir.exists():
+        logger.info(f"extracting {voices_tgz} to {voices_dir}")
+        with tarfile.open(voices_tgz, "r:gz") as tar:
+            tar.extractall(path=voices_tgz.parent)
+
+    if not voices_dir.exists():
+        raise RuntimeError("voices.tgz did not contain a 'voices/' directory")
+
+    return str(voices_dir)
+
+
+def _get_static_path(static: Optional[str]) -> Optional[str]:
+    if static is None:
+        logger.info("retrieving the static content")
+        dist_tgz = hf_hub_download("nvidia/personaplex-7b-v1", "dist.tgz")
+        dist_tgz = Path(dist_tgz)
+        dist = dist_tgz.parent / "dist"
+        if not dist.exists():
+            with tarfile.open(dist_tgz, "r:gz") as tar:
+                tar.extractall(path=dist_tgz.parent)
+        return str(dist)
+    elif static != "none":
+        # When set to the "none" string, we don't serve any static content.
+        return static
+    return None
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--host", default="localhost", type=str)
+    parser.add_argument("--port", default=8998, type=int)
+    parser.add_argument("--static", type=str)
+    parser.add_argument("--gradio-tunnel", action='store_true', help='Activate a gradio tunnel.')
+    parser.add_argument("--gradio-tunnel-token",
+                        help='Provide a custom (secret) token here to keep getting the same URL.')
+
+    parser.add_argument("--tokenizer", type=str, help="Path to a local tokenizer file.")
+    parser.add_argument("--moshi-weight", type=str, help="Path to a local checkpoint file for Moshi.")
+    parser.add_argument("--mimi-weight", type=str, help="Path to a local checkpoint file for Mimi.")
+    parser.add_argument("--hf-repo", type=str, default=loaders.DEFAULT_REPO,
+                        help="HF repo to look into, defaults PersonaPlex. "
+                             "Use this to select a different pre-trained model.")
+    parser.add_argument("--device", type=str, default="cuda", help="Device on which to run, defaults to 'cuda'.")
+    parser.add_argument("--cpu-offload", action="store_true",
+                        help="Offload LM model layers to CPU when GPU memory is insufficient. "
+                             "Requires 'accelerate' package.")
+    parser.add_argument(
+        "--voice-prompt-dir",
+        type=str,
+        help=(
+            "Directory containing voice prompt files. "
+            "If omitted, voices.tgz is downloaded from HF and extracted."
+            "Voice prompt filenames from client requests will be joined with this directory path."
+        )
+    )
+    parser.add_argument(
+        "--ssl",
+        type=str,
+        help=(
+            "use https instead of http, this flag should point to a directory "
+            "that contains valid key.pem and cert.pem files"
+        )
+    )
+
+    args = parser.parse_args()
+    args.voice_prompt_dir = _get_voice_prompt_dir(
+        args.voice_prompt_dir,
+        args.hf_repo,
+    )
+    if args.voice_prompt_dir is not None:
+        assert os.path.exists(args.voice_prompt_dir), \
+            f"Directory missing: {args.voice_prompt_dir}"
+    logger.info(f"voice_prompt_dir = {args.voice_prompt_dir}")
+
+    static_path: None | str = _get_static_path(args.static)
+    assert static_path is None or os.path.exists(static_path), \
+        f"Static path does not exist: {static_path}."
+    logger.info(f"static_path = {static_path}")
+    args.device = torch_auto_device(args.device)
+
+    seed_all(42424242)
+
+    setup_tunnel = None
+    tunnel_token = ''
+    if args.gradio_tunnel:
+        try:
+            from gradio import networking  # type: ignore
+        except ImportError:
+            logger.error("Cannot find gradio which is required to activate a tunnel. "
+                         "Please install with `pip install gradio`.")
+            sys.exit(1)
+        setup_tunnel = networking.setup_tunnel
+        if args.gradio_tunnel_token is None:
+            tunnel_token = secrets.token_urlsafe(32)
+        else:
+            tunnel_token = args.gradio_tunnel_token
+
+    # Download config.json to increment download counter
+    # No worries about double-counting since config.json will be cached the second time
+    hf_hub_download(args.hf_repo, "config.json")
+
+    logger.info("loading mimi")
+    if args.mimi_weight is None:
+        args.mimi_weight = hf_hub_download(args.hf_repo, loaders.MIMI_NAME)
+    mimi = loaders.get_mimi(args.mimi_weight, args.device)
+    other_mimi = loaders.get_mimi(args.mimi_weight, args.device)
+    logger.info("mimi loaded")
+
+    if args.tokenizer is None:
+        args.tokenizer = hf_hub_download(args.hf_repo, loaders.TEXT_TOKENIZER_NAME)
+    text_tokenizer = sentencepiece.SentencePieceProcessor(args.tokenizer)  # type: ignore
+
+    logger.info("loading moshi")
+    if args.moshi_weight is None:
+        args.moshi_weight = hf_hub_download(args.hf_repo, loaders.MOSHI_NAME)
+    lm = loaders.get_moshi_lm(args.moshi_weight, device=args.device, cpu_offload=args.cpu_offload)
+    lm.eval()
+    logger.info("moshi loaded")
+    state = ServerState(
+        mimi=mimi,
+        other_mimi=other_mimi,
+        text_tokenizer=text_tokenizer,
+        lm=lm,
+        device=args.device,
+        voice_prompt_dir=args.voice_prompt_dir,
+        save_voice_prompt_embeddings=False,
+    )
+    logger.info("warming up the model")
+    state.warmup()
+    app = web.Application()
+    app.router.add_get("/api/chat", state.handle_chat)
+    if static_path is not None:
+        async def handle_root(_):
+            return web.FileResponse(os.path.join(static_path, "index.html"))
+
+        logger.info(f"serving static content from {static_path}")
+        app.router.add_get("/", handle_root)
+        app.router.add_static(
+            "/", path=static_path, follow_symlinks=True, name="static"
+        )
+    protocol = "http"
+    ssl_context = None
+    if args.ssl is not None:
+        ssl_context, protocol = create_ssl_context(args.ssl)
+    host_ip = args.host if args.host not in ("0.0.0.0", "::", "localhost") else get_lan_ip()
+    logger.info(f"Access the Web UI directly at {protocol}://{host_ip}:{args.port}")
+    if setup_tunnel is not None:
+        tunnel = setup_tunnel('localhost', args.port, tunnel_token, None)
+        logger.info(f"Tunnel started, if executing on a remote GPU, you can use {tunnel}.")
+    web.run_app(app, port=args.port, ssl_context=ssl_context)
+
+
+with torch.no_grad():
+    main()

moshi/utils/__init__.py

@@ -0,0 +1,10 @@
+# Copyright (c) Kyutai, all rights reserved.
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+"""Utilities."""

moshi/utils/autocast.py

@@ -0,0 +1,45 @@
+# Copyright (c) Kyutai, all rights reserved.
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+
+
+class TorchAutocast:
+    """TorchAutocast utility class.
+    Allows you to enable and disable autocast. This is specially useful
+    when dealing with different architectures and clusters with different
+    levels of support.
+
+    Args:
+        enabled (bool): Whether to enable torch.autocast or not.
+        args: Additional args for torch.autocast.
+        kwargs: Additional kwargs for torch.autocast
+    """
+
+    def __init__(self, enabled: bool, *args, **kwargs):
+        self.autocast = torch.autocast(*args, **kwargs) if enabled else None
+
+    def __enter__(self):
+        if self.autocast is None:
+            return
+        try:
+            self.autocast.__enter__()
+        except RuntimeError:
+            device = self.autocast.device
+            dtype = self.autocast.fast_dtype
+            raise RuntimeError(
+                f"There was an error autocasting with dtype={dtype} device={device}\n"
+                "If you are on the FAIR Cluster, you might need to use autocast_dtype=float16"
+            )
+
+    def __exit__(self, *args, **kwargs):
+        if self.autocast is None:
+            return
+        self.autocast.__exit__(*args, **kwargs)

moshi/utils/compile.py

@@ -0,0 +1,308 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: MIT
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+
+# Copyright (c) Kyutai, all rights reserved.
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Provides some extra utilities around torch compile, in particular with a way
+to fully deactivate it easily with a context manager.
+Provides a simple activation checkpointing that is compatible with FSDP and torch compile.
+Finally, provides some utilities for CUDA graphing functions.
+"""
+from contextlib import contextmanager
+from functools import wraps
+import inspect
+import os
+import typing as tp
+
+import torch
+from torch import cuda
+
+
+_compile_disabled: bool = False
+
+
+@contextmanager
+def no_compile():
+    """Disable torch.compile locally. Now Pytorch 2.4 provides a function to do that."""
+    global _compile_disabled
+
+    prev_disabled = _compile_disabled
+    _compile_disabled = True
+    try:
+        yield
+    finally:
+        _compile_disabled = prev_disabled
+
+
+def torch_compile_lazy(fun):
+    """torch.compile creates a huge pool of processes, even when not using the function at all,
+    e.g. with Dora. This can polute stderr when doing CTRL+C. So we do it in a lazy way.
+    """
+    if os.environ.get("NO_TORCH_COMPILE"):
+        return fun
+    fun_compiled = None
+
+    @wraps(fun)
+    def _wrapped(*args, **kwargs):
+        nonlocal fun_compiled
+        if _compile_disabled:
+            return fun(*args, **kwargs)
+        if fun_compiled is None:
+            fun_compiled = torch.compile(fun)
+        return fun_compiled(*args, **kwargs)
+
+    return _wrapped
+
+
+class Checkpoint(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, function, *args) -> tp.Any:
+        to_save = []
+        ctx.others = []
+        ctx.function = function
+        # Sources will indicate whether the arg in position N is
+        # a tensor stored in ctx.save_for_backward, or inside ctx.others.
+        ctx.sources = []
+        new_args = []
+        for arg in args:
+            if isinstance(arg, torch.Tensor):
+                to_save.append(arg)
+                ctx.sources.append("tensor")
+                new_args.append(arg.detach())
+            else:
+                ctx.sources.append("other")
+                ctx.others.append(arg)
+                new_args.append(arg)
+        ctx.save_for_backward(*to_save)
+        # During the forward, we just make a pass with no gradient computed.
+        with torch.no_grad():
+            res = function(*new_args)
+        return res
+
+    @staticmethod
+    def backward(ctx, *grads) -> tp.Tuple[tp.Optional[torch.Tensor], ...]:
+        pseudo_tensors = []
+        with torch.set_grad_enabled(True):
+            # We create leaf tensors to collect the output gradients.
+            # We call them pseudo_tensors because they are pretending to be the input
+            # to `function` but are not directly
+            for tensor in ctx.saved_tensors:
+                pseudo_tensor = tensor.detach()
+                pseudo_tensor.requires_grad_(True)
+                pseudo_tensors.append(pseudo_tensor)
+            pseudo_tensors_copy = list(pseudo_tensors)
+            args = []
+            for source in ctx.sources:
+                if source == "other":
+                    args.append(ctx.others.pop(0))
+                else:
+                    assert source == "tensor"
+                    args.append(pseudo_tensors_copy.pop(0))
+            res = ctx.function(*args)
+            # The second forward with grad computation allows us to connect the input leaf tensors
+            # inside pseudo_tensors, to the outputs of the function called.
+        if not isinstance(res, tuple):
+            res = (res,)
+        # Now we just ask Torch to compute the derivative of `res` given the gradient coming from above
+        # `grads`. The computed gradient will end up into the `pseudo_tensors` grad attributes.
+        torch.autograd.backward(res, grads)
+        out: tp.List[tp.Optional[torch.Tensor]] = [None]
+        for source in ctx.sources:
+            # We still need to output `None` values for non tensor parameters.
+            if source == "other":
+                out.append(None)
+            else:
+                assert source == "tensor"
+                out.append(pseudo_tensors.pop(0).grad)
+        return tuple(out)
+
+
+def simple_checkpoint(module: torch.nn.Module, *args, **kwargs):
+    """Custom implementation of checkpointing in PyTorch as the builtin implementation is broken
+    when using torch compile. Only supports wrapping a `nn.Module` with a forward with no `*args` or `**kwargs`.
+
+    https://github.com/pytorch/pytorch/issues/97436.
+    Should be resolved in nightlies, but it is quite fun and simple to code it ourselves.
+    """
+    if hasattr(module, "_fsdp_wrapped_module"):
+        module_for_sig = module._fsdp_wrapped_module
+    else:
+        module_for_sig = module
+    sig = inspect.signature(module_for_sig.forward)
+    # We first flatten all arguments to use only *args, to make things easier and because
+    # torch.autograd.Function has weird support for kwargs.
+    bounded = sig.bind(*args, **kwargs)
+    new_args = []
+    for name, param in sig.parameters.items():
+        if param.kind in {
+            inspect.Parameter.VAR_POSITIONAL,
+            inspect.Parameter.VAR_KEYWORD,
+        }:
+            raise RuntimeError("simple_checkpoint doesn't support var args.")
+        if name not in bounded.arguments:
+            break
+        new_args.append(bounded.arguments[name])
+    return Checkpoint.apply(module, *new_args)
+
+
+_in_cuda_graph = False
+_disable_cuda_graph = False
+
+
+def in_cuda_graph() -> bool:
+    """Indicate whether we are in a function that is CUDA Graphed (or will be soon)."""
+    return _in_cuda_graph
+
+
+@contextmanager
+def _set_in_cuda_graph():
+    global _in_cuda_graph
+    assert not _in_cuda_graph
+    _in_cuda_graph = True
+    try:
+        yield
+    finally:
+        _in_cuda_graph = False
+
+
+def _is_cuda_graph_enabled() -> bool:
+    if _disable_cuda_graph:
+        return False
+    no_cuda_graph = os.environ.get("NO_CUDA_GRAPH", "")
+    if no_cuda_graph.lower() not in {"0", "no", "n", ""}:
+        return False
+    return True
+
+
+@contextmanager
+def no_cuda_graph():
+    """Deactivate CUDA Graphing for all the calls in this context manager."""
+    global _disable_cuda_graph
+    old_value = _disable_cuda_graph
+    _disable_cuda_graph = True
+    try:
+        yield
+    finally:
+        _disable_cuda_graph = old_value
+
+
+class CUDAGraphed:
+    """Allow simple CUDA Graphing of a function.
+
+    Args:
+        func: callable, taking any number of arguments. Its tensors arguments should
+            be top level args, not nested in structures (tuples, dicts, etc). Keyword
+            arguments are NOT supported for simplicity.
+        warmup_steps: how many call to make normally before CUDA Graphing. In particular, this
+            allows torch.compiled functions to get properly compiled.
+        disabled: if True, just call the func directly, useful to quickly deactivate on CPU.
+    """
+
+    def __init__(self, func: tp.Callable, warmup_steps: int = 1, disable: bool = False):
+        self.func = func
+        self.warmup_steps = warmup_steps
+        self.disable = disable
+        self._graph: cuda.CUDAGraph | None = None
+        self._output: tuple | None = None
+        self._args: tuple | None = None
+
+    def reset(self, warmup_steps: int = 0) -> None:
+        """Reset the state, meaning the next call we get CUDA Graphed again. Useful if some
+        shapes have changed, or external state (e.g. KVCache) has changed."""
+        self.warmup_steps = warmup_steps
+        self._graph = None
+        self._output = None
+        self._args = None
+
+    def asdict(self):
+        return {}
+
+    def __call__(self, *args, **kwargs) -> tp.Any:
+        if kwargs:
+            raise RuntimeError("Named arguments not supported for now.")
+        if self.disable or not _is_cuda_graph_enabled() or in_cuda_graph():
+            return self.func(*args, **kwargs)
+
+        def _clone_tensors(args: tuple) -> tuple:
+            out: list = []
+            for arg in args:
+                if isinstance(arg, torch.Tensor):
+                    arg = arg.clone()
+                out.append(arg)
+            return tuple(out)
+
+        def _match_values_copy_tensors(args: tuple, target_args: tuple) -> None:
+            if len(args) != len(target_args):
+                raise ValueError(
+                    f"Expected {len(target_args)}, but got {args} for CUDA Graphed function."
+                )
+            for idx, (source, target) in enumerate(zip(args, target_args)):
+                if isinstance(target, torch.Tensor):
+                    if not isinstance(source, torch.Tensor):
+                        raise ValueError(
+                            f"Argument #{idx} was a tensor, and is no longer (now {source})."
+                        )
+                    if source.shape != target.shape:
+                        raise ValueError(
+                            f"Argument #{idx} had shape {target.shape}, but got shae {source.shape}"
+                        )
+                    target.copy_(source)
+                else:
+                    if isinstance(source, torch.Tensor):
+                        raise ValueError(
+                            f"Argument #{idx} was not a tensor {target}, but is now one."
+                        )
+                    if source is not target and source != target:
+                        raise ValueError(
+                            f"Argument #{idx} changed value from {target} to {source}."
+                        )
+
+        with _set_in_cuda_graph():
+            # Prevent any one under us to try and CUDA Graph things.
+            if self._graph is None:
+                if self.warmup_steps <= 0:
+                    self._graph = cuda.CUDAGraph()
+                    # Making a copy just to ensure those are not used else where.
+                    self._args = _clone_tensors(args)
+                    with cuda.graph(self._graph):
+                        self._output = self.func(*self._args)
+                    # At this point nothing really happened, so we have to make it run for real.
+                    self._graph.replay()
+                    return self._output
+                else:
+                    self.warmup_steps -= 1
+                    return self.func(*args)
+            else:
+                assert self._args is not None
+                assert self._output is not None
+                _match_values_copy_tensors(args, self._args)
+                self._graph.replay()
+                return self._output
+
+
+def cuda_graph(func: tp.Callable, warmup_steps: int = 1):
+    """Just calls `CUDAGraphed` on the given function."""
+    if not _is_cuda_graph_enabled():
+        return func
+    return CUDAGraphed(func, warmup_steps)

moshi/utils/connection.py

@@ -0,0 +1,228 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: MIT
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+
+import socket
+import subprocess
+import platform
+from shutil import which
+import urllib.request
+from pathlib import Path
+from .logging import setup_logger
+
+logger = setup_logger(__name__)
+
+def get_lan_ip():
+    """
+    get_lan_ip()
+    
+    Get the LAN IP address of the current machine.
+
+    Returns
+    -------
+    str
+        LAN IP address as a string.
+    """
+    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+    try:
+        s.connect(("8.8.8.8", 80))
+        return s.getsockname()[0]
+    finally:
+        s.close()
+
+
+def ensure_mkcert():
+    """
+    Ensure mkcert exists; if not, download the correct binary.
+    
+    Returns
+    -------
+    str | None
+        Path to mkcert binary, or None if installation failed.
+    """
+
+    # mkcert already installed?
+    mkcert_bin_path = which("mkcert")
+    if mkcert_bin_path:
+        return mkcert_bin_path
+
+    logger.info("[auto-cert] mkcert not found, attempting to download...")
+
+    # Determine OS + arch
+    system = platform.system().lower()   # "linux", "darwin", "windows"
+    machine = platform.machine().lower() # "x86_64", "amd64", "arm64", etc.
+
+    # Normalize architecture
+    if machine in ("x86_64", "amd64"):
+        arch = "amd64"
+    elif machine in ("aarch64", "arm64"):
+        arch = "arm64"
+    elif machine.startswith("arm"):
+        arch = "arm"
+    else:
+        logger.error(f"[auto-cert] Unsupported architecture: {machine}")
+        return None
+
+    # Determine filename and URL
+    if system == "linux":
+        filename = f"mkcert-v1.4.4-linux-{arch}"
+    elif system == "darwin":
+        filename = f"mkcert-v1.4.4-darwin-{arch}"
+    elif system == "windows":
+        filename = f"mkcert-v1.4.4-windows-{arch}.exe"
+    else:
+        logger.error(f"[auto-cert] Unsupported OS: {system}")
+        return None
+
+    url = f"https://github.com/FiloSottile/mkcert/releases/download/v1.4.4/{filename}"
+    dest_dir = Path.home() / ".local" / "bin"
+    dest_dir.mkdir(parents=True, exist_ok=True)
+    dest = dest_dir / ("mkcert.exe" if system == "windows" else "mkcert")
+
+    logger.info(f"[auto-cert] Downloading: {url}")
+    try:
+        urllib.request.urlretrieve(url, dest)
+    except Exception as e:
+        logger.error(f"[auto-cert] Failed to download mkcert: {e}")
+        return None
+
+    # Make executable if Unix
+    if system != "windows":
+        dest.chmod(0o755)
+
+    logger.info(f"[auto-cert] mkcert installed at {dest}")
+
+    # Verify it works
+    try:
+        subprocess.check_call([str(dest), "-help"],
+                              stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
+    except Exception:
+        logger.error("[auto-cert] mkcert failed to run after install.")
+        return None
+
+    return str(dest)
+
+
+def _run_command(cmd):
+    """Run command, return True on success."""
+    try:
+        subprocess.check_call(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
+        return True
+    except subprocess.CalledProcessError:
+        return False
+
+
+def ensure_mkcert_ca(mkcert_bin: str):
+    """
+    Install mkcert CA if not installed yet.
+
+    Parameters
+    ----------
+    mkcert_bin : str
+        Path to mkcert binary.
+    """
+    _run_command([mkcert_bin, "-install"])
+
+
+def create_cert_if_needed(cert_dir: str):
+    """
+    Create cert.pem and key.pem using mkcert if they don't already exist.
+    Returns (cert_file, key_file) or (None, None) if mkcert unavailable.
+
+    Parameters
+    ----------
+    cert_dir : str
+        Directory to store or find cert.pem and key.pem.
+    
+    Returns
+    -------
+    pathlib.Path | None
+        Path to cert.pem, or None if not created.
+    pathlib.Path | None
+        Path to key.pem, or None if not created.
+    """
+    cert_dir = Path(cert_dir)
+    cert_dir.mkdir(parents=True, exist_ok=True)
+
+    cert_file = cert_dir / "cert.pem"
+    key_file = cert_dir / "key.pem"
+
+    # Already exists → nothing to do
+    if cert_file.exists() and key_file.exists():
+        return cert_file, key_file
+
+    mkcert_bin = ensure_mkcert()
+
+    if not mkcert_bin:
+        logger.warning("[auto-cert] mkcert not installed; falling back to HTTP.")
+        return None, None
+
+    logger.info("[auto-cert] mkcert detected. Ensuring local CA installed...")
+    ensure_mkcert_ca(mkcert_bin)
+
+    # Create cert for localhost, loopback, and LAN IP
+    lan_ip = get_lan_ip()
+    logger.info(f"[auto-cert] Generating certificate for localhost and {lan_ip}...")
+
+    success = _run_command([
+        mkcert_bin,
+        "-cert-file", str(cert_file),
+        "-key-file", str(key_file),
+        "localhost",
+        "127.0.0.1",
+        "::1",
+        lan_ip
+    ])
+
+    if not success:
+        logger.warning("[auto-cert] mkcert failed. Using HTTP.")
+        return None, None
+
+    logger.info("[auto-cert] Certificate generated.")
+    return cert_file, key_file
+
+
+def create_ssl_context(cert_dir: str):
+    """
+    Main entry: create SSL context if certificates can be created.
+    Returns (ssl_context, protocol_str) where protocol_str is "http" or "https".
+
+    Parameters
+    ----------
+    cert_dir : str
+        Directory to store or find cert.pem and key.pem.
+
+    Returns
+    -------
+    ssl.SSLContext | None
+        SSL context if HTTPS is available, else None.
+    str
+        "https" if SSL context created, else "http".
+    """
+    cert_file, key_file = create_cert_if_needed(cert_dir)
+
+    if cert_file is None:
+        # mkcert missing → use HTTP
+        return None, "http"
+
+    import ssl
+    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
+    ctx.load_cert_chain(certfile=str(cert_file), keyfile=str(key_file))
+    return ctx, "https"

moshi/utils/logging.py

@@ -0,0 +1,75 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: MIT
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+
+import logging
+import sys
+import random
+import string
+from typing import Optional
+from ..client_utils import make_log, colorize
+
+
+def random_id(n=4):
+    return "".join(random.choices(string.ascii_uppercase + string.digits, k=n))
+
+
+def setup_logger(name: str, log_file=None, level=logging.INFO):
+    logger = logging.getLogger(name)
+    logger.setLevel(level)
+
+    formatter = logging.Formatter(
+        "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+    )
+
+    console_handler = logging.StreamHandler(sys.stdout)
+    console_handler.setFormatter(formatter)
+    logger.addHandler(console_handler)
+
+    if log_file:
+        file_handler = logging.FileHandler(log_file)
+        file_handler.setFormatter(formatter)
+        logger.addHandler(file_handler)
+
+    return logger
+
+
+def print_log(level: str, msg: str, prefix: Optional[str] = None, info_color: Optional[str] = None):
+    colorized_msg = make_log(level, msg) if info_color is None or level != "info" else colorize(msg, info_color)
+    if prefix is None:
+        print(colorized_msg)
+    else:
+        print(prefix + colorized_msg)
+
+
+class ColorizedLog(object):
+    def __init__(self, prefix: str, info_color: str):
+        self.prefix = prefix
+        self.info_color = info_color
+
+    def log(self, level: str, msg: str):
+        print_log(level, msg, prefix=self.prefix, info_color=self.info_color)
+
+    @classmethod
+    def randomize(cls):
+        cid = random_id()
+        color = random.choice(["91", "92", "93", "94", "95", "96", "97"])
+        prefix = colorize(f"[{cid}] ", color)
+        return cls(prefix=prefix, info_color=color)

moshi/utils/sampling.py

@@ -0,0 +1,147 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: MIT
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+
+# Copyright (c) Kyutai, all rights reserved.
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+import torch
+
+
+def multinomial(
+    input: torch.Tensor, num_samples: int, replacement=False, *, generator=None
+):
+    """torch.multinomial with arbitrary number of dimensions, and number of candidates on the last dimension.
+
+    Args:
+        input (torch.Tensor): The input tensor containing probabilities.
+        num_samples (int): Number of samples to draw.
+        replacement (bool): Whether to draw with replacement or not.
+    Keywords args:
+        generator (torch.Generator): A pseudorandom number generator for sampling.
+    Returns:
+        torch.Tensor: Last dimension contains num_samples indices
+            sampled from the multinomial probability distribution
+            located in the last dimension of tensor input.
+    """
+    input_ = input.reshape(-1, input.shape[-1])
+    # We should probably be able to remove this once the following PR has landed:
+    # https://github.com/pytorch/pytorch/pull/134818/files
+    # In the meantime, we specialize the case no-replacement, nsamples=1 so as not
+    # to have a synchronization point.
+    if replacement or num_samples != 1:
+        output_ = torch.multinomial(
+            input_,
+            num_samples=num_samples,
+            replacement=replacement,
+            generator=generator,
+        )
+    else:
+        q = torch.empty_like(input_).exponential_(1, generator=generator)
+        q = input_ / q
+        output_ = q.argmax(dim=-1, keepdim=True)
+    output = output_.reshape(*list(input.shape[:-1]), -1)
+    return output
+
+
+def sample_top_k(probs: torch.Tensor, k: int) -> torch.Tensor:
+    """Sample next token from top K values along the last dimension of the input probs tensor.
+
+    Args:
+        probs (torch.Tensor): Input probabilities with token candidates on the last dimension.
+        k (int): The k in “top-k”.
+    Returns:
+        torch.Tensor: Sampled tokens.
+    """
+    probs, indices = torch.topk(probs, k, dim=-1)
+    next_token = multinomial(probs, num_samples=1)
+    next_token = indices.gather(-1, next_token)
+    return next_token
+
+
+def sample_top_p(probs: torch.Tensor, p: float) -> torch.Tensor:
+    """Sample next token from top P probabilities along the last dimension of the input probs tensor.
+
+    Args:
+        probs (torch.Tensor): Input probabilities with token candidates on the last dimension.
+        p (int): The p in “top-p”.
+    Returns:
+        torch.Tensor: Sampled tokens.
+    """
+    probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)
+    probs_sum = torch.cumsum(probs_sort, dim=-1)
+    mask = probs_sum - probs_sort > p
+    probs_sort *= (~mask).float()
+    probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True))
+    next_token = multinomial(probs_sort, num_samples=1)
+    next_token = torch.gather(probs_idx, -1, next_token)
+    return next_token
+
+
+def sample_token(
+    logits: torch.Tensor,
+    use_sampling: bool = False,
+    temp: float = 1.0,
+    top_k: int = 0,
+    top_p: float = 0.0,
+) -> torch.Tensor:
+    """Given logits of shape [*, Card], returns a LongTensor of shape [*]."""
+    # Apply softmax for sampling if temp > 0. Else, do greedy sampling to avoid zero division error.
+    if use_sampling and temp > 0.0:
+        probs = torch.softmax(logits / temp, dim=-1)
+        if top_p > 0.0:
+            next_token = sample_top_p(probs, p=top_p)
+        elif top_k > 0:
+            next_token = sample_top_k(probs, k=top_k)
+        else:
+            next_token = multinomial(probs, num_samples=1)
+    else:
+        next_token = torch.argmax(logits, dim=-1, keepdim=True)
+    assert next_token.shape[-1] == 1
+    return next_token[..., 0]
+
+
+if __name__ == "__main__":
+    torch.manual_seed(1234)
+    device = "cpu"
+    if torch.cuda.is_available():
+        torch.backends.cuda.matmul.allow_tf32 = False
+        torch.backends.cudnn.allow_tf32 = False
+        device = "cuda:0"
+
+    ps = torch.tensor([5.0, 2.0, 12.0, 6.0, 8.0, 1.0, 0.0, 4.0], device=device)
+    cnts = torch.zeros(ps.shape, dtype=torch.long, device=device)
+    total_samples = 1000
+    for _ in range(total_samples):
+        vs = multinomial(ps, num_samples=1, replacement=False)
+        cnts[vs] += 1
+    diff = cnts / cnts.sum() - ps / ps.sum()
+    max_diff = diff.abs().max().cpu().item()
+    print(ps / ps.sum())
+    print(cnts / cnts.sum())
+    assert max_diff < 1.5e-2

requirements.txt

@@ -1,11 +1,15 @@
+--extra-index-url https://download.pytorch.org/whl/cu128
 gradio==4.44.1
 fastapi==0.112.4
 starlette==0.38.6
-torch
-numpy
+torch==2.8.0+cu128
+numpy>=1.26,<2.2
 pydantic==2.10.6
-huggingface_hub
-sentencepiece
-sphn
-safetensors
-git+https://github.com/NVIDIA/personaplex.git@3428dfd95309a7f3c84fd93259ded0f810d1ff91#subdirectory=moshi
+huggingface_hub>=0.24,<0.25
+sentencepiece==0.2
+sphn>=0.1.4,<0.2
+safetensors>=0.4,<0.5
+einops==0.7
+tqdm
+sounddevice==0.5
+aiohttp>=3.10.5,<3.11