Add mossttsrealtime model

mossttsrealtime/__init__.py

@@ -0,0 +1,31 @@
+# Copyright 2026 OpenMOSS and the HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import TYPE_CHECKING
+
+from transformers.utils import _LazyModule
+from transformers.utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+    from .configuration_mossttsrealtime import *
+    from .modeling_mossttsrealtime import *
+    from .modeling_mossttsrealtime_local import *
+    from .processing_mossttsrealtime import *
+    from .streaming_mossttsrealtime import *
+else:
+    import sys
+
+    _file = globals()["__file__"]
+    sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)

mossttsrealtime/configuration_mossttsrealtime.py

@@ -0,0 +1,115 @@
+# Copyright 2026 OpenMOSS and the HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""MossTTSRealtimeModel configuration."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.models.qwen3 import Qwen3Config
+
+
+def _ensure_config(cfg: Any, cls: type[PretrainedConfig]) -> PretrainedConfig:
+    if isinstance(cfg, cls):
+        return cfg
+    if cfg is None:
+        return cls()
+    if isinstance(cfg, dict):
+        return cls(**cfg)
+    raise TypeError(f"Unsupported config type for {cls.__name__}: {type(cfg)}")
+
+
+class MossTTSRealtimeLocalTransformerConfig(PretrainedConfig):
+    model_type = "moss_tts_realtime_local_transformer"
+
+    def __init__(
+        self,
+        head_dim: int = 128,
+        use_cache: bool = True,
+        hidden_size: int = 2048,
+        rms_norm_eps: float = 1e-6,
+        num_hidden_layers: int = 4,
+        intermediate_size: int = 6144,
+        num_attention_heads: int = 16,
+        initializer_range: float = 0.02,
+        attention_bias: bool = False,
+        attention_dropout: float = 0.0,
+        max_position_embeddings: int = 33,
+        num_key_value_heads: int = 8,
+        hidden_act: str = "silu",
+        rope_theta: int = 1000000,
+        rope_type: str = "linear",
+        pad_token_id: int = 1024,
+        rope_parameters: dict | None = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.head_dim = head_dim
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.hidden_act = hidden_act
+        self.rope_theta = rope_theta
+        self.rope_type = rope_type
+        if rope_parameters is None:
+            rope_parameters = {"rope_type": rope_type, "rope_theta": rope_theta, "factor": 1.0}
+        self.rope_parameters = rope_parameters
+        self.attention_bias = attention_bias
+        self.attention_dropout = attention_dropout
+        self.num_key_value_heads = num_key_value_heads
+        self.max_position_embeddings = max_position_embeddings
+        self.pad_token_id = pad_token_id
+
+        self.audio_pad_token = 1024
+        self.audio_vocab_size = 1027
+        self.rvq = 16
+
+
+class MossTTSRealtimeConfig(PretrainedConfig):
+    model_type = "moss_tts_realtime"
+
+    def __init__(
+        self,
+        language_config: Qwen3Config | dict | None = None,
+        local_config: MossTTSRealtimeLocalTransformerConfig | dict | None = None,
+        rvq: int = 16,
+        audio_pad_token: int = 1024,
+        audio_vocab_size: int = 1027,
+        reference_audio_pad: int = 151654,
+        text_pad: int = 151655,
+        initializer_range: float = 0.02,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.rvq = rvq
+        self.initializer_range = initializer_range
+        self.audio_pad_token = audio_pad_token
+        self.audio_vocab_size = audio_vocab_size
+        self.reference_audio_pad = reference_audio_pad
+        self.text_pad = text_pad
+        self.language_config = _ensure_config(language_config, Qwen3Config)
+        self.local_config = _ensure_config(local_config, MossTTSRealtimeLocalTransformerConfig)
+
+        attn_impl = self._attn_implementation
+        self.language_config._attn_implementation = attn_impl
+        self.local_config._attn_implementation = attn_impl
+
+
+__all__ = ["MossTTSRealtimeConfig", "MossTTSRealtimeLocalTransformerConfig"]

mossttsrealtime/modeling_mossttsrealtime.py

@@ -0,0 +1,190 @@
+# Copyright 2026 OpenMOSS and the HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""MossTTSRealtime model."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers import initialization as init
+from transformers.cache_utils import Cache
+from transformers.modeling_outputs import ModelOutput
+from transformers.modeling_utils import PreTrainedModel
+from transformers.models.qwen3 import Qwen3Model
+from transformers.models.qwen3.modeling_qwen3 import Qwen3Attention, Qwen3DecoderLayer
+from .configuration_mossttsrealtime import MossTTSRealtimeConfig
+from .modeling_mossttsrealtime_local import MossTTSRealtimeLocalTransformerForCausalLM
+
+
+class MossTTSRealtimePretrainedModel(PreTrainedModel):
+    config_class = MossTTSRealtimeConfig
+    config: MossTTSRealtimeConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["Qwen3DecoderLayer"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_sdpa = True
+    _supports_flex_attn = True
+    _supports_flash_attn = True
+    _can_compile_fullgraph = True
+    _supports_attention_backend = True
+    _can_record_outputs = {
+        "hidden_states": Qwen3DecoderLayer,
+        "attentions": Qwen3Attention,
+    }
+
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear):
+            init.normal_(module.weight, mean=0.0, std=std)
+            if module.bias is not None:
+                init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            init.normal_(module.weight, mean=0.0, std=std)
+            if module.padding_idx is not None:
+                init.zeros_(module.weight[module.padding_idx])
+
+
+@dataclass
+class MossTTSRealtimeOutputWithPast(ModelOutput):
+    loss: Optional[torch.FloatTensor] = None
+    logits: Optional[torch.FloatTensor] = None
+    past_key_values: Optional[Cache] = None
+    last_hidden_state: Optional[torch.FloatTensor] = None
+    hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[tuple[torch.FloatTensor]] = None
+    local_loss: Optional[torch.FloatTensor] = None
+    local_logits: Optional[torch.FloatTensor] = None
+    local_past_key_values: Optional[tuple[tuple[torch.FloatTensor]]] = None
+    local_hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None
+    local_attentions: Optional[tuple[torch.FloatTensor, ...]] = None
+    backbone_loss: Optional[torch.FloatTensor] = None
+
+
+class MossTTSRealtime(MossTTSRealtimePretrainedModel):
+    def __init__(self, config: MossTTSRealtimeConfig):
+        super().__init__(config)
+        self.config = config
+        self.embed_tokens = nn.ModuleList([])
+        self.embed_tokens.append(
+            nn.Embedding(
+                config.language_config.vocab_size,
+                config.language_config.hidden_size,
+                config.language_config.pad_token_id,
+            )
+        )
+        self.audio_vocab_size = self.config.audio_vocab_size
+        for _ in range(self.config.rvq):
+            self.embed_tokens.append(
+                nn.Embedding(self.audio_vocab_size, config.language_config.hidden_size, self.config.audio_pad_token)
+            )
+        self.language_model = Qwen3Model._from_config(config.language_config)
+        self.local_transformer = MossTTSRealtimeLocalTransformerForCausalLM._from_config(config.local_config)
+        self.post_init()
+
+    def get_input_embeddings(self, input_ids):
+        if input_ids.device != self.embed_tokens[0].weight.device:
+            input_ids = input_ids.to(self.embed_tokens[0].weight.device)
+        inputs_embeds = self.embed_tokens[0](input_ids[..., 0])
+        for i, embed in enumerate(self.embed_tokens):
+            if i == 0:
+                continue
+            inputs_embeds = inputs_embeds + embed(input_ids[..., i])
+        return inputs_embeds
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[list[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = False,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        hidden_out_layers: Optional[list] = None,
+        **kwargs,
+    ):
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if inputs_embeds is None:
+            inputs_embeds = self.get_input_embeddings(input_ids)
+
+        outputs = self.language_model(
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        loss = None
+        local_outputs = None
+        if labels is not None:
+            audio_labels = labels[:, :, 1:]
+            train_mask = ~(audio_labels == -100).all(dim=-1)
+            local_input_ids = audio_labels[train_mask][..., : self.config.rvq - 1]
+            local_input_ids[local_input_ids == -100] = self.config.audio_pad_token
+            local_input_ids = F.pad(local_input_ids, (1, 0), value=0)
+
+            train_idx = train_mask.nonzero(as_tuple=True)
+            hidden_positions = torch.clamp(train_idx[1] - 1, min=0)
+            local_hidden_states = outputs.last_hidden_state[train_idx[0], hidden_positions, :].reshape(
+                -1, 1, self.config.local_config.hidden_size
+            )
+            local_labels = audio_labels[train_mask]
+
+            local_outputs = self.local_transformer(
+                input_ids=local_input_ids,
+                backbone_last_hidden_state=local_hidden_states,
+                use_cache=use_cache,
+                return_dict=True,
+                labels=local_labels,
+                **kwargs,
+            )
+            loss = local_outputs.loss
+
+        output = MossTTSRealtimeOutputWithPast(
+            loss=loss,
+            logits=None,
+            past_key_values=outputs.past_key_values,
+            last_hidden_state=outputs.last_hidden_state,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            local_logits=local_outputs.logits if local_outputs is not None else None,
+            local_past_key_values=local_outputs.past_key_values if local_outputs is not None else None,
+            local_hidden_states=local_outputs.hidden_states if local_outputs is not None else None,
+            local_attentions=local_outputs.attentions if local_outputs is not None else None,
+        )
+        if not return_dict:
+            return output.to_tuple()
+        return output
+
+
+__all__ = ["MossTTSRealtime", "MossTTSRealtimeConfig", "MossTTSRealtimeOutputWithPast", "MossTTSRealtimePretrainedModel"]

mossttsrealtime/modeling_mossttsrealtime_local.py

@@ -0,0 +1,471 @@
+# Copyright 2026 OpenMOSS and the HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Local transformer used by MossTTSRealtime for RVQ codebook decoding."""
+
+from __future__ import annotations
+
+from typing import Optional, Union
+
+import torch
+import torch.nn as nn
+
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache, StaticCache
+from transformers.generation import GenerationMixin
+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+from transformers.modeling_layers import GradientCheckpointingLayer
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from transformers.masking_utils import create_causal_mask
+from transformers.processing_utils import Unpack
+from transformers.loss.loss_utils import ForCausalLMLoss
+from transformers.utils import TransformersKwargs, logging
+from .configuration_mossttsrealtime import MossTTSRealtimeLocalTransformerConfig
+
+logger = logging.get_logger(__name__)
+
+
+class MossTTSRealtimeLocalTransformerRMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+    def extra_repr(self):
+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+class MossTTSRealtimeLocalTransformerMLP(nn.Module):
+    def __init__(self, config: MossTTSRealtimeLocalTransformerConfig):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x):
+        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+        return down_proj
+
+
+def rotate_half(x):
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: Optional[torch.Tensor],
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs: Unpack[TransformersKwargs],
+):
+    key_states = repeat_kv(key, module.num_key_value_groups)
+    value_states = repeat_kv(value, module.num_key_value_groups)
+    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+        attn_weights = attn_weights + causal_mask
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+    return attn_output, attn_weights
+
+
+class MossTTSRealtimeLocalTransformerAttention(nn.Module):
+    def __init__(self, config: MossTTSRealtimeLocalTransformerConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = config.attention_dropout
+        self.is_causal = True
+
+        self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias)
+        self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias)
+        self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias)
+        self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias)
+        self.q_norm = MossTTSRealtimeLocalTransformerRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.k_norm = MossTTSRealtimeLocalTransformerRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.sliding_window = None
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+        past_key_values: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+        query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+        key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        cos, sin = position_embeddings
+
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        if past_key_values is not None:
+            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+        attention_interface = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            sliding_window=self.sliding_window,
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+
+class MossTTSRealtimeLocalTransformerDecoderLayer(GradientCheckpointingLayer):
+    def __init__(self, config: MossTTSRealtimeLocalTransformerConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.self_attn = MossTTSRealtimeLocalTransformerAttention(config=config, layer_idx=layer_idx)
+        self.mlp = MossTTSRealtimeLocalTransformerMLP(config)
+        self.input_layernorm = MossTTSRealtimeLocalTransformerRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = MossTTSRealtimeLocalTransformerRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.attention_type = "full_attention"
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        use_cache: Optional[bool] = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states, _ = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )
+        hidden_states = residual + hidden_states
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+
+class MossTTSRealtimeLocalTransformerPreTrainedModel(PreTrainedModel):
+
+    config_class = MossTTSRealtimeLocalTransformerConfig
+    config: MossTTSRealtimeLocalTransformerConfig
+
+    base_model_prefix = "local_transformer"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["MossTTSRealtimeLocalTransformerDecoderLayer"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_sdpa = True
+    _supports_flex_attn = True
+    _supports_flash_attn = True
+    _can_compile_fullgraph = True
+    _supports_attention_backend = True
+
+    _can_record_outputs = {
+        "hidden_states": MossTTSRealtimeLocalTransformerDecoderLayer,
+        "attentions": MossTTSRealtimeLocalTransformerAttention,
+    }
+
+
+class MossTTSRealtimeLocalTransformerRotaryEmbedding(nn.Module):
+    inv_freq: torch.Tensor
+
+    def __init__(self, config: MossTTSRealtimeLocalTransformerConfig, device=None):
+        super().__init__()
+        self.config = config
+        self.rope_type = getattr(config, "rope_type", "linear")
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.original_inv_freq = self.inv_freq
+
+    @torch.no_grad()
+    @dynamic_rope_update
+    def forward(self, x, position_ids):
+        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+        position_ids_expanded = position_ids[:, None, :].float()
+        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+        with torch.autocast(device_type=device_type, enabled=False):
+            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos() * self.attention_scaling
+            sin = emb.sin() * self.attention_scaling
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+class MossTTSRealtimeLocalTransformer(MossTTSRealtimeLocalTransformerPreTrainedModel):
+    def __init__(self, config: MossTTSRealtimeLocalTransformerConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.embed_tokens = nn.ModuleList(
+            [nn.Embedding(config.audio_vocab_size, config.hidden_size, config.audio_pad_token) for _ in range(config.rvq - 1)]
+        )
+        self.layers = nn.ModuleList(
+            [MossTTSRealtimeLocalTransformerDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.norm = MossTTSRealtimeLocalTransformerRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = MossTTSRealtimeLocalTransformerRotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+        self.has_sliding_layers = None
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        backbone_last_hidden_state: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        codebook_idx: Optional[int] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutputWithPast:
+        if position_ids is not None and not torch.compiler.is_compiling():
+            position_ids = None
+
+        if (input_ids is None) == (inputs_embeds is None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds.")
+
+        if use_cache and past_key_values is None:
+            device = inputs_embeds.device if inputs_embeds is not None else input_ids.device
+            past_key_values = StaticCache(config=self.config, max_cache_len=self.config.rvq, device=device)
+
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            inputs_seq_length = inputs_embeds.shape[1] if inputs_embeds is not None else input_ids.shape[1]
+            device = inputs_embeds.device if inputs_embeds is not None else input_ids.device
+            cache_position = torch.arange(past_seen_tokens, past_seen_tokens + inputs_seq_length, device=device)
+
+        if inputs_embeds is None:
+            if codebook_idx is not None:
+                if codebook_idx <= 0:
+                    raise ValueError(f"`codebook_idx` must be in [1, {len(self.embed_tokens)}], got {codebook_idx}.")
+                if codebook_idx > len(self.embed_tokens):
+                    raise ValueError(f"`codebook_idx` must be in [1, {len(self.embed_tokens)}], got {codebook_idx}.")
+                if input_ids.ndim == 1:
+                    input_ids = input_ids.unsqueeze(1)
+                token_emb = self.embed_tokens[codebook_idx - 1](input_ids[:, 0]).unsqueeze(1)  # [B,1,H]
+                inputs_embeds = token_emb
+            else:
+                if input_ids.shape[1] != cache_position.shape[0]:
+                    raise ValueError(
+                        "`input_ids` and `cache_position` must align in sequence length: "
+                        f"got {input_ids.shape[1]} and {cache_position.shape[0]}."
+                    )
+                codebook_idxs = torch.clamp(cache_position - 1, min=0, max=len(self.embed_tokens) - 1)
+                inputs_embeds = torch.stack(
+                    [
+                        self.embed_tokens[codebook_idx](input_ids[:, seq_idx])
+                        for seq_idx, codebook_idx in enumerate(codebook_idxs.tolist())
+                    ],
+                    dim=1,
+                )
+
+                input_ids_are_first_codebook = bool(cache_position[0] == 0)
+                if backbone_last_hidden_state is not None:
+                    inputs_embeds[:, 0, :] = backbone_last_hidden_state[:, 0, :]
+                else:
+                    if not torch.compiler.is_compiling() and input_ids_are_first_codebook:
+                        logger.warning(
+                            "When the first codebook token is provided, `backbone_last_hidden_state` should also be provided for correct inference."
+                        )
+
+        causal_mask = create_causal_mask(
+            config=self.config,
+            input_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            cache_position=cache_position,
+            past_key_values=past_key_values,
+            position_ids=position_ids,
+        )
+
+        hidden_states = inputs_embeds
+        position_ids = cache_position.unsqueeze(0)
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=causal_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                position_embeddings=position_embeddings,
+                **kwargs,
+            )
+        hidden_states = self.norm(hidden_states)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values if use_cache else None,
+        )
+
+
+class MossTTSRealtimeLocalTransformerForCausalLM(MossTTSRealtimeLocalTransformerPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = None
+    _tp_plan = None
+    _pp_plan = None
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = MossTTSRealtimeLocalTransformer(config)
+        self.audio_vocab_size = self.config.audio_vocab_size
+
+        self.local_lm_heads = nn.ModuleList(
+            [nn.Linear(config.hidden_size, config.audio_vocab_size, bias=False) for _ in range(config.rvq)]
+        )
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        backbone_last_hidden_state: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Union[Cache, list[torch.FloatTensor]]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        codebook_idx: Optional[int] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> Union[tuple, CausalLMOutputWithPast]:
+        outputs = self.model(
+            input_ids=input_ids,
+            backbone_last_hidden_state=backbone_last_hidden_state,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            codebook_idx=codebook_idx,
+            **kwargs,
+        )
+
+        hidden_states = outputs.last_hidden_state
+
+        if isinstance(logits_to_keep, int):
+            if logits_to_keep == 0:
+                slice_indices = slice(0, None)
+            else:
+                slice_indices = slice(-logits_to_keep, None)
+        else:
+            slice_indices = logits_to_keep
+        hs = hidden_states[:, slice_indices, :]
+
+        if cache_position is not None:
+            if codebook_idx is None:
+                raise ValueError("`codebook_idx` must be provided when `cache_position` is provided.")
+            logits = self.local_lm_heads[codebook_idx](hs[:, 0, :]).unsqueeze(1)
+        else:
+            if hs.shape[1] > len(self.local_lm_heads):
+                raise ValueError(
+                    f"Cannot project {hs.shape[1]} codebooks with only {len(self.local_lm_heads)} LM heads."
+                )
+            logits_list = []
+            for i in range(hs.shape[1]):
+                logits_list.append(self.local_lm_heads[i](hs[:, i, :]))
+            logits = torch.stack(logits_list, dim=1)
+
+        logits = logits.contiguous()
+        loss = None
+        if labels is not None:
+            loss = ForCausalLMLoss(logits, None, self.audio_vocab_size, shift_labels=labels.contiguous())
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+__all__ = [
+    "MossTTSRealtimeLocalTransformer",
+    "MossTTSRealtimeLocalTransformerAttention",
+    "MossTTSRealtimeLocalTransformerConfig",
+    "MossTTSRealtimeLocalTransformerDecoderLayer",
+    "MossTTSRealtimeLocalTransformerForCausalLM",
+    "MossTTSRealtimeLocalTransformerPreTrainedModel",
+    "MossTTSRealtimeLocalTransformerRMSNorm",
+    "MossTTSRealtimeLocalTransformerRotaryEmbedding",
+]
+

mossttsrealtime/processing_mossttsrealtime.py

@@ -0,0 +1,178 @@
+# Copyright 2026 OpenMOSS and the HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Processing utilities for MossTTSRealtime."""
+
+from __future__ import annotations
+
+from typing import Iterable, Optional
+
+import numpy as np
+
+from transformers.processing_utils import ProcessorMixin
+
+
+class MossTTSRealtimeProcessor(ProcessorMixin):
+    """Builds MossTTSRealtime prompt inputs with text and audio codebooks.
+    This processor focuses on preparing the mixed text/audio token layout expected by MossTTSRealtime.
+    It does not perform audio encoding/decoding by itself.
+    """
+
+    attributes = ["tokenizer"]
+    tokenizer_class = "AutoTokenizer"
+
+    def __init__(
+        self,
+        tokenizer,
+        audio_pad_token: str = "<|audio_pad|>",
+        text_pad_token: str = "<|text_pad|>",
+        tts_system_prompt: Optional[str] = None,
+        channels: int = 16,
+        audio_channel_pad: int = 1024,
+        audio_bos_token: int = 1025,
+        audio_eos_token: int = 1026,
+        delay_tokens_len: int = 12,
+    ):
+        super().__init__(tokenizer=tokenizer)
+        self.audio_pad_token = audio_pad_token
+        self.text_pad_token = text_pad_token
+        self.channels = channels
+        self.audio_channel_pad = audio_channel_pad
+        self.audio_bos_token = audio_bos_token
+        self.audio_eos_token = audio_eos_token
+        self.delay_tokens_len = delay_tokens_len
+
+        self.audio_pad_token_id = self._convert_token_to_id(audio_pad_token)
+        self.text_pad_token_id = self._convert_token_to_id(text_pad_token)
+
+        if tts_system_prompt is None:
+            tts_system_prompt = (
+                "<|im_start|>system\n"
+                "You are a highly expressive text-to-speech (TTS) engine developed by Mosi Intelligence. \n"
+                "You possess natural language understanding, emotional modeling, and multi-style speech generation "
+                "capabilities, allowing you to generate the corresponding speech based on the text given in the assistant."
+                "<|im_end|>\n"
+            )
+        self.tts_system_prompt = tts_system_prompt
+
+    def _convert_token_to_id(self, token: str) -> int:
+        if hasattr(self.tokenizer, "convert_tokens_to_ids"):
+            token_id = self.tokenizer.convert_tokens_to_ids(token)
+            if token_id is not None and token_id != self.tokenizer.unk_token_id:
+                return int(token_id)
+        token_ids = self.tokenizer.encode(token, add_special_tokens=False)
+        if not token_ids:
+            raise ValueError(f"Token '{token}' could not be converted to an id.")
+        if len(token_ids) != 1:
+            raise ValueError(f"Token '{token}' maps to multiple ids: {token_ids}")
+        return int(token_ids[0])
+
+    def make_voice_clone_prompt(self, prompt_audio_tokens_len: int) -> str:
+        padded_audio_prompt = f"{self.audio_pad_token * prompt_audio_tokens_len}"
+        voice_clone = (
+            "<|im_start|>context\n"
+            "The assistant section should be synthesized using the following voice timbre:"
+            f"{padded_audio_prompt}"
+        )
+        return voice_clone
+
+    def _normalize_audio_tokens(self, audio_tokens: np.ndarray | Iterable) -> np.ndarray:
+        tokens = np.array(audio_tokens)
+        if tokens.ndim != 2:
+            raise ValueError(f"Expected 2D audio tokens, got shape {tokens.shape}")
+        # Accept [channels, T] or [T, channels], and slice to expected channels if needed.
+        if tokens.shape[0] == self.channels:
+            tokens = tokens.T
+        elif tokens.shape[1] == self.channels:
+            tokens = tokens
+        elif tokens.shape[0] > self.channels and tokens.shape[1] != self.channels:
+            tokens = tokens[: self.channels, :].T
+        elif tokens.shape[1] > self.channels and tokens.shape[0] != self.channels:
+            tokens = tokens[:, : self.channels]
+        if tokens.shape[1] != self.channels:
+            raise ValueError(f"Expected {self.channels} channels, got shape {tokens.shape}")
+        return tokens
+
+    def make_ensemble(self, prompt_audio_tokens: Optional[np.ndarray] = None) -> np.ndarray:
+        if prompt_audio_tokens is not None:
+            prompt_audio_tokens = self._normalize_audio_tokens(prompt_audio_tokens)
+            prompt_audio_tokens = prompt_audio_tokens[:, : self.channels]
+            system_prompt_text = f"{self.tts_system_prompt}" + f"{self.make_voice_clone_prompt(prompt_audio_tokens.shape[0])}"
+        else:
+            system_prompt_text = f"{self.tts_system_prompt}"
+
+        system_prompt_tokens = self.tokenizer(system_prompt_text)["input_ids"]
+        system_prompt_tokens_full = np.full(
+            shape=(len(system_prompt_tokens), self.channels + 1), fill_value=self.audio_channel_pad, dtype=np.int64
+        )
+        system_prompt_tokens_full[:, 0] = system_prompt_tokens
+
+        if prompt_audio_tokens is not None:
+            system_prompt_tokens = np.array(system_prompt_tokens)
+            indices = np.where(system_prompt_tokens == self.audio_pad_token_id)[0]
+            if indices.size == 0:
+                raise ValueError("No <|audio_pad|> tokens found in the system prompt.")
+            prompt_audio_start_pos, prompt_audio_end_pos = indices[0], indices[-1]
+            system_prompt_tokens_full[prompt_audio_start_pos : prompt_audio_end_pos + 1, 1:] = prompt_audio_tokens
+
+        return system_prompt_tokens_full
+
+    def make_user_prompt(self, text: str, audio_tokens: np.ndarray) -> np.ndarray:
+        prefill_temp = "<|im_end|>\n<|im_start|>user\n"
+        text_tokens = self.tokenizer(text)["input_ids"]
+        text_start_pos = len(self.tokenizer.encode(prefill_temp))
+        token = self._normalize_audio_tokens(audio_tokens)
+
+        text_len = len(text_tokens)
+        audio_len = token.shape[0]
+
+        if text_len >= self.delay_tokens_len:
+            padded_text_len = audio_len + self.delay_tokens_len - text_len + 1
+            cur_input_id_ch1 = prefill_temp + text + "<|text_pad|>" * padded_text_len
+            assistant_tokens_ch1 = self.tokenizer(cur_input_id_ch1)["input_ids"]
+            cur_input_id = np.full(
+                shape=(len(assistant_tokens_ch1), self.channels + 1),
+                fill_value=self.audio_channel_pad,
+                dtype=np.int64,
+            )
+            cur_input_id[:, 0] = assistant_tokens_ch1
+            cur_input_id[
+                text_start_pos + self.delay_tokens_len : text_start_pos + self.delay_tokens_len + audio_len, 1:
+            ] = token
+            cur_input_id[text_start_pos + self.delay_tokens_len - 1, 1] = self.audio_bos_token
+            cur_input_id[text_start_pos + self.delay_tokens_len + audio_len, 1] = self.audio_eos_token
+        else:
+            padded_text_len = audio_len + 1
+            cur_input_id_ch1 = prefill_temp + text + "<|text_pad|>" * padded_text_len
+            assistant_tokens_ch1 = self.tokenizer(cur_input_id_ch1)["input_ids"]
+            cur_input_id = np.full(
+                shape=(len(assistant_tokens_ch1), self.channels + 1),
+                fill_value=self.audio_channel_pad,
+                dtype=np.int64,
+            )
+            cur_input_id[:, 0] = assistant_tokens_ch1
+            cur_input_id[-(audio_len + 1) : -1, 1:] = token
+            cur_input_id[-(audio_len + 2), 1] = self.audio_bos_token
+            cur_input_id[-1, 1] = self.audio_eos_token
+
+        begin_of_response = self.tokenizer.encode("<|im_end|>\n<|im_start|>assistant\n")
+        begin_of_response_full = np.full(
+            shape=(len(begin_of_response), self.channels + 1), fill_value=self.audio_channel_pad, dtype=np.int64
+        )
+        begin_of_response_full[:, 0] = begin_of_response
+
+        input_ids = np.concatenate([cur_input_id, begin_of_response_full], axis=0)
+        return input_ids
+
+
+__all__ = ["MossTTSRealtimeProcessor"]

mossttsrealtime/streaming_mossttsrealtime.py

@@ -0,0 +1,1051 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Streaming inference utilities for MossTTSRealtime."""
+
+from __future__ import annotations
+
+import contextlib
+import re
+from typing import Iterable, Iterator, List, Optional, Sequence
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from transformers.cache_utils import DynamicCache, StaticCache
+from transformers.utils import is_torchaudio_available, requires_backends
+from transformers.utils.import_utils import requires
+
+if is_torchaudio_available():
+    import torchaudio
+
+
+@requires(backends=("torch",))
+class MossTTSRealtimeInference:
+    """Step-wise inference wrapper for MossTTSRealtime.
+    This class mirrors the non-streaming inference logic but exposes a
+    prefill/step/finish API for streaming usage.
+    """
+
+    def __init__(
+        self,
+        model,
+        tokenizer,
+        max_length: int = 1000,
+        channels: int = 16,
+        audio_channel_pad: int = 1024,
+        audio_bos_token: int = 1025,
+        audio_eos_token: int = 1026,
+        text_pad_id: int = 151655,
+        aud_pad_id: int = 151654,
+    ):
+        self.model = model
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.channels = channels
+        self.audio_channel_pad = audio_channel_pad
+        self.audio_bos_token = audio_bos_token
+        self.audio_eos_token = audio_eos_token
+        self.text_pad_id = text_pad_id
+        self.aud_pad_id = aud_pad_id
+
+        self.past_key_values = None
+        self.attention_mask = None
+        self._generated_tokens: List[torch.Tensor] = []
+        self._is_stopping = None
+        self._last_audio_tokens = None
+        self._step_idx = 0
+        attn_impl = ""
+        for cfg in (
+            getattr(getattr(self.model, "local_transformer", None), "config", None),
+            getattr(getattr(self.model, "config", None), "local_config", None),
+            getattr(self.model, "config", None),
+        ):
+            if cfg is None:
+                continue
+            for name in ("_attn_implementation", "attn_implementation"):
+                candidate = getattr(cfg, name, None)
+                if isinstance(candidate, str) and candidate.strip():
+                    attn_impl = candidate.strip().lower()
+                    break
+            if attn_impl:
+                break
+        self._use_dynamic_local_cache = attn_impl == "flash_attention_2"
+        self._should_compile_local_transformer = not self._use_dynamic_local_cache
+        self._compiled_local_transformer = None
+
+    @property
+    def device(self):
+        return next(self.model.parameters()).device
+
+    @property
+    def is_finished(self) -> bool:
+        return self._is_stopping is not None and bool(self._is_stopping.all())
+
+    def _build_local_past_key_values(self):
+        if self._use_dynamic_local_cache:
+            return DynamicCache()
+        return StaticCache(config=self.model.local_transformer.config, max_cache_len=self.channels)
+
+    def _get_local_transformer_runner(self):
+        if not self._should_compile_local_transformer:
+            return self._generate_local_transformer_impl
+        if self._compiled_local_transformer is None:
+            self._compiled_local_transformer = torch.compile(self._generate_local_transformer_impl, fullgraph=True)
+        return self._compiled_local_transformer
+
+    def reset_generation_state(self, keep_cache: bool = True):
+        if not keep_cache:
+            self.past_key_values = None
+            self.attention_mask = None
+        # Keep the mask when reusing cache so it stays aligned with past_key_values.
+        # This allows concatenation with the next turn prefill mask.
+        self._generated_tokens = []
+        self._is_stopping = None
+        self._last_audio_tokens = None
+        self._step_idx = 0
+
+    def _normalize_input_ids(self, input_ids):
+        if isinstance(input_ids, torch.Tensor):
+            input_ids = input_ids.detach().cpu().numpy()
+        if isinstance(input_ids, np.ndarray):
+            if input_ids.ndim == 2:
+                return [input_ids]
+            if input_ids.ndim == 3:
+                return [input_ids[i] for i in range(input_ids.shape[0])]
+        if isinstance(input_ids, (list, tuple)):
+            return [np.array(item) for item in input_ids]
+        raise ValueError("input_ids must be a list/array/tensor of shape [T, C] or [B, T, C].")
+
+    def _normalize_text_prefix(self, text_prefix_ids, batch_size: int) -> list[list[int]]:
+        if text_prefix_ids is None:
+            raise ValueError("text_prefix_ids must be provided for prefill.")
+        if isinstance(text_prefix_ids, torch.Tensor):
+            text_prefix_ids = text_prefix_ids.detach().cpu().tolist()
+        if isinstance(text_prefix_ids, np.ndarray):
+            text_prefix_ids = text_prefix_ids.tolist()
+        if isinstance(text_prefix_ids, list):
+            if len(text_prefix_ids) == 0:
+                return [[] for _ in range(batch_size)]
+            if isinstance(text_prefix_ids[0], (int, np.integer)):
+                return [list(text_prefix_ids)]
+            if len(text_prefix_ids) == 1 and batch_size > 1:
+                return [list(text_prefix_ids[0]) for _ in range(batch_size)]
+            if len(text_prefix_ids) != batch_size:
+                raise ValueError(
+                    f"text_prefix_ids batch size mismatch: got {len(text_prefix_ids)}, expected {batch_size}."
+                )
+            return [list(item) for item in text_prefix_ids]
+        raise ValueError("text_prefix_ids must be list-like or tensor-like.")
+
+    @torch.inference_mode()
+    def prefill(
+        self,
+        input_ids,
+        text_prefix_ids,
+        max_prefill_len: Optional[int] = None,
+        past_key_values=None,
+        device: Optional[torch.device] = None,
+        temperature: float = 0.8,
+        top_p: float = 0.6,
+        top_k: int = 30,
+        do_sample: bool = True,
+        repetition_penalty: Optional[float] = 1.1,
+        repetition_window: Optional[int] = 50,
+    ) -> torch.Tensor:
+        if device is None:
+            device = self.device
+
+        if past_key_values is not None:
+            self.past_key_values = past_key_values
+
+        input_ids_list = self._normalize_input_ids(input_ids)
+        batch_size = len(input_ids_list)
+        text_prefix_list = self._normalize_text_prefix(text_prefix_ids, batch_size)
+
+        concat_inputs_id_list = []
+        for i in range(batch_size):
+            prefix = text_prefix_list[i]
+            if max_prefill_len is not None:
+                prefix = prefix[:max_prefill_len]
+            if len(prefix) == 0:
+                raise ValueError("Prefill requires at least one text token.")
+
+            text_seg = np.full((len(prefix), self.channels + 1), self.audio_channel_pad, dtype=np.int64)
+            text_seg[:, 0] = np.array(prefix, dtype=np.int64)
+            text_seg[len(prefix) - 1, 1] = self.audio_bos_token
+            concat_inputs_id = np.concatenate([input_ids_list[i], text_seg], axis=0)
+            concat_inputs_id_list.append(concat_inputs_id)
+
+        attention_masks = [np.ones(ids.shape[0], dtype=np.bool_) for ids in concat_inputs_id_list]
+        max_len = max(ids.shape[0] for ids in concat_inputs_id_list)
+        padded_input_ids, padded_attns = [], []
+        pad_token_id = self.tokenizer.pad_token_id if self.tokenizer.pad_token_id is not None else self.text_pad_id
+
+        for ids, attn in zip(concat_inputs_id_list, attention_masks):
+            pad_len = max_len - ids.shape[0]
+            input_pad = np.full((pad_len, self.channels + 1), self.audio_channel_pad, dtype=np.int64)
+            input_pad[:, 0] = pad_token_id
+            padded_input_ids.append(np.concatenate([input_pad, ids]))
+            attn_pad = np.zeros(pad_len, dtype=np.bool_)
+            padded_attns.append(np.concatenate([attn_pad, attn]))
+
+        current_input_ids = torch.from_numpy(np.stack(padded_input_ids)).to(device)
+        current_attention_mask = torch.from_numpy(np.stack(padded_attns)).to(device)
+
+        # For multi-turn continuation, concatenate the cached mask and the current prefill mask.
+        if self.attention_mask is not None and self.past_key_values is not None:
+            current_attention_mask = torch.cat([self.attention_mask, current_attention_mask], dim=-1)
+
+        outputs = self.model(
+            input_ids=current_input_ids,
+            attention_mask=current_attention_mask,
+            past_key_values=self.past_key_values,
+            use_cache=True,
+            return_dict=True,
+        )
+        self.past_key_values = outputs.past_key_values
+        self.attention_mask = current_attention_mask
+
+        backbone_hidden_states = outputs.last_hidden_state[:, -1:, :]
+        audio_tokens = self.generate_local_transformer(
+            hidden_states=backbone_hidden_states,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            do_sample=do_sample,
+            repetition_penalty=repetition_penalty,
+            repetition_window=repetition_window,
+            generated_tokens=None,
+            gen_step=0,
+        )
+
+        self._generated_tokens = [audio_tokens]
+        self._last_audio_tokens = audio_tokens
+        self._is_stopping = audio_tokens[:, 0] == self.audio_eos_token
+        self._step_idx = 1
+        return audio_tokens
+
+    @torch.inference_mode()
+    def step(
+        self,
+        text_token: Optional[Iterable[int] | torch.Tensor | int],
+        temperature: float = 0.8,
+        top_p: float = 0.6,
+        top_k: int = 30,
+        do_sample: bool = True,
+        repetition_penalty: Optional[float] = 1.1,
+        repetition_window: Optional[int] = 50,
+    ) -> torch.Tensor:
+        if self._last_audio_tokens is None or self.attention_mask is None:
+            raise ValueError("You must call prefill() before step().")
+        if self.is_finished:
+            return self._last_audio_tokens
+
+        batch_size = self._last_audio_tokens.shape[0]
+        if text_token is None:
+            text_tokens = [self.text_pad_id] * batch_size
+        elif isinstance(text_token, torch.Tensor):
+            text_tokens = text_token.detach().cpu().tolist()
+        elif isinstance(text_token, (list, tuple, np.ndarray)):
+            text_tokens = list(text_token)
+        else:
+            text_tokens = [int(text_token)]
+
+        if len(text_tokens) != batch_size:
+            raise ValueError(f"text_token batch size mismatch: got {len(text_tokens)}, expected {batch_size}.")
+
+        device = self._last_audio_tokens.device
+        text_t = torch.tensor(text_tokens, device=device, dtype=torch.long)
+        step_ids = torch.cat([text_t[:, None, None], self._last_audio_tokens.unsqueeze(1)], dim=2)
+        self.attention_mask = torch.cat([self.attention_mask, (~self._is_stopping).unsqueeze(-1)], dim=-1)
+
+        outputs = self.model(
+            input_ids=step_ids,
+            attention_mask=self.attention_mask,
+            past_key_values=self.past_key_values,
+            use_cache=True,
+            return_dict=True,
+        )
+        self.past_key_values = outputs.past_key_values
+        backbone_hidden_states = outputs.last_hidden_state[:, -1:, :]
+
+        history = torch.stack(self._generated_tokens, dim=1) if self._generated_tokens else None
+        audio_tokens = self.generate_local_transformer(
+            hidden_states=backbone_hidden_states,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            do_sample=do_sample,
+            repetition_penalty=repetition_penalty,
+            repetition_window=repetition_window,
+            generated_tokens=history,
+            gen_step=self._step_idx,
+        )
+
+        self._generated_tokens.append(audio_tokens)
+        self._last_audio_tokens = audio_tokens
+        self._is_stopping |= audio_tokens[:, 0] == self.audio_eos_token
+        self._step_idx += 1
+        return audio_tokens
+
+    @torch.inference_mode()
+    def finish(
+        self,
+        max_steps: Optional[int] = None,
+        temperature: float = 0.8,
+        top_p: float = 0.6,
+        top_k: int = 30,
+        do_sample: bool = True,
+        repetition_penalty: Optional[float] = 1.1,
+        repetition_window: Optional[int] = 50,
+    ) -> list[torch.Tensor]:
+        outputs = []
+        steps_left = max_steps if max_steps is not None else self.max_length
+        while steps_left > 0 and not self.is_finished:
+            outputs.append(
+                self.step(
+                    text_token=None,
+                    temperature=temperature,
+                    top_p=top_p,
+                    top_k=top_k,
+                    do_sample=do_sample,
+                    repetition_penalty=repetition_penalty,
+                    repetition_window=repetition_window,
+                )
+            )
+            steps_left -= 1
+        return outputs
+
+    def generate_local_transformer(
+        self,
+        hidden_states: torch.Tensor,
+        temperature: float,
+        top_p: float,
+        top_k: int,
+        do_sample: bool,
+        repetition_penalty: Optional[float],
+        repetition_window: Optional[int],
+        generated_tokens: Optional[torch.Tensor],
+        gen_step: int,
+    ) -> torch.Tensor:
+        runner = self._get_local_transformer_runner()
+        return runner(
+            hidden_states=hidden_states,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            do_sample=do_sample,
+            repetition_penalty=repetition_penalty,
+            repetition_window=repetition_window,
+            generated_tokens=generated_tokens,
+            gen_step=gen_step,
+        )
+
+    def _generate_local_transformer_impl(
+        self,
+        hidden_states: torch.Tensor,
+        temperature: float,
+        top_p: float,
+        top_k: int,
+        do_sample: bool,
+        repetition_penalty: Optional[float],
+        repetition_window: Optional[int],
+        generated_tokens: Optional[torch.Tensor],
+        gen_step: int,
+    ) -> torch.Tensor:
+        batch_size = hidden_states.shape[0]
+        device = hidden_states.device
+        local_inputs = hidden_states.reshape(-1, 1, self.model.config.local_config.hidden_size)
+        output_token = torch.empty(batch_size, self.channels, dtype=torch.long, device=device)
+
+        past_key_values = self._build_local_past_key_values()
+        local_token = None
+
+        cache_pos_t = torch.zeros(1, dtype=torch.long, device=device)
+
+        for i in range(self.channels):
+            cache_pos_t.fill_(i)
+
+            local_outputs = self.model.local_transformer(
+                input_ids=local_token,
+                inputs_embeds=local_inputs,
+                past_key_values=past_key_values,
+                cache_position=cache_pos_t,
+                codebook_idx=i,
+                use_cache=True,
+                logits_to_keep=1,
+            )
+            logits = local_outputs.logits
+
+            if repetition_penalty and repetition_penalty != 1.0 and generated_tokens is not None:
+                logits = self.apply_repetition_penalty(
+                    scores=logits,
+                    history_tokens=generated_tokens[:, :gen_step, i],
+                    penalty=float(repetition_penalty),
+                    repetition_window=repetition_window,
+                )
+
+            local_token = self.sample_token(
+                logits=logits,
+                temperature=temperature,
+                top_p=top_p,
+                top_k=top_k,
+                do_sample=do_sample,
+            )
+            output_token[:, i] = local_token.squeeze(-1)
+
+            if i == 0:
+                local_inputs = None
+        return output_token
+
+    def apply_repetition_penalty(
+        self,
+        scores: torch.Tensor,
+        history_tokens: torch.Tensor,
+        penalty: float = 1.1,
+        repetition_window: Optional[int] = None,
+    ):
+        scores_ = scores[:, 0, :]
+        ht = history_tokens
+
+        if repetition_window is not None and repetition_window > 0:
+            ht = ht[:, -repetition_window:]
+
+        cur = scores_.gather(1, ht)
+        new = torch.where(cur < 0, cur * penalty, cur / penalty)
+        scores_.scatter_(1, ht, new)
+        return scores_
+
+    def sample_token(self, logits, temperature, top_p=0.6, top_k=30, do_sample=True):
+        if not do_sample or temperature == 0:
+            return torch.argmax(logits, dim=-1)
+        logits = logits / temperature
+        original_shape = logits.shape
+        vocab_size = original_shape[-1]
+        reshaped_logits = logits.reshape(-1, vocab_size)
+
+        if top_k is not None:
+            reshaped_logits = self.apply_top_k(reshaped_logits, top_k)
+
+        if top_p is not None:
+            reshaped_logits = self.apply_top_p(reshaped_logits, top_p)
+
+        probs = F.softmax(reshaped_logits, dim=-1)
+        next_tokens_flat = torch.multinomial(probs, num_samples=1)
+
+        output_shape = original_shape[:-1]
+        return next_tokens_flat.view(output_shape)
+
+    def apply_top_k(self, logits, top_k, filter_value=float("-inf"), min_tokens_to_keep: int = 1):
+        if not isinstance(top_k, int) or top_k <= 0:
+            raise ValueError(f"`top_k` has to be a strictly positive integer, but is {top_k}")
+        batch_size, vocab_size = logits.shape
+        top_k = max(top_k, min_tokens_to_keep)
+        top_k = min(top_k, vocab_size)
+        indices_to_remove = torch.topk(logits, top_k, dim=-1).values[..., -1, None]
+        return logits.masked_fill(logits < indices_to_remove, filter_value)
+
+    def apply_top_p(self, logits, top_p, filter_value=float("-inf"), min_tokens_to_keep: int = 1):
+        top_p = float(top_p)
+        if top_p < 0 or top_p > 1.0:
+            raise ValueError(f"`top_p` has to be a float > 0 and < 1, but is {top_p}")
+
+        sorted_logits, sorted_indices = torch.sort(logits, descending=False)
+        cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)
+
+        sorted_indices_to_remove = cumulative_probs <= (1 - top_p)
+        sorted_indices_to_remove[..., -min_tokens_to_keep:] = 0
+        indices_to_remove = torch.zeros_like(logits, dtype=torch.bool).scatter(1, sorted_indices, sorted_indices_to_remove)
+        logits_processed = logits.masked_fill(indices_to_remove, filter_value)
+        return logits_processed
+
+
+@requires(backends=("torch",))
+class MossTTSRealtimeStreamingSession:
+    """Manage text-to-audio streaming for a single conversation."""
+
+    _split_pattern = re.compile(
+        r"[。！？!?\.\u2026]\s*"  # sentence boundaries: 。！？ ! ? . …
+        r"|[,，;；:：\u2014\u2013\-]\s*"  # short pauses: , ， ; ； : ： — – -
+        r"|\)\s*|\]\s*"  # closing brackets: ) ]
+        r"|\n"
+    )
+
+    def __init__(
+        self,
+        inferencer: MossTTSRealtimeInference,
+        processor,
+        codec=None,
+        codec_sample_rate: int = 24000,
+        codec_encode_kwargs: Optional[dict] = None,
+        prefill_text_len: int = 12,
+        text_buffer_size: int = 32,
+        min_text_chunk_chars: int = 8,
+        temperature: float = 0.8,
+        top_p: float = 0.6,
+        top_k: int = 30,
+        do_sample: bool = True,
+        repetition_penalty: Optional[float] = 1.1,
+        repetition_window: Optional[int] = 50,
+    ):
+        self.inferencer = inferencer
+        self.processor = processor
+        self.tokenizer = processor.tokenizer
+        self.codec = codec
+        self.codec_sample_rate = codec_sample_rate
+        self.codec_encode_kwargs = codec_encode_kwargs or {}
+
+        self.prefill_text_len = prefill_text_len
+        self.text_buffer_size = text_buffer_size
+        self.min_text_chunk_chars = min_text_chunk_chars
+
+        self.temperature = temperature
+        self.top_p = top_p
+        self.top_k = top_k
+        self.do_sample = do_sample
+        self.repetition_penalty = repetition_penalty
+        self.repetition_window = repetition_window
+
+        self._voice_prompt_tokens = None
+        self._turn_input_ids = None
+        self._turn_idx = 0
+
+        self._text_cache = ""
+        self._pending_tokens: list[int] = []
+        self._prefilled = False
+        self._text_ended = False
+
+    def set_voice_prompt_tokens(self, audio_tokens: np.ndarray):
+        self._voice_prompt_tokens = audio_tokens
+
+    def set_voice_prompt(self, audio, sample_rate: Optional[int] = None):
+        """Set voice prompt from either audio tokens or waveform.
+        If `audio` is a 2D array whose shape matches the codebook channels, it is
+        treated as audio tokens. Otherwise a codec is required to encode waveform
+        prompts into tokens.
+        """
+        if isinstance(audio, np.ndarray) and audio.ndim == 2:
+            if self.processor.channels in audio.shape:
+                self._voice_prompt_tokens = audio
+                return
+        if isinstance(audio, torch.Tensor) and audio.dim() == 2:
+            if self.processor.channels in audio.shape:
+                self._voice_prompt_tokens = audio.detach().cpu().numpy()
+                return
+
+        if self.codec is None:
+            raise ValueError("codec is required to encode waveform prompts.")
+
+        waveform = audio
+        if isinstance(audio, (str, bytes)):
+            requires_backends(self, ["torchaudio"])
+            wav, sr = torchaudio.load(audio)
+            if wav.shape[0] > 1:
+                wav = wav.mean(dim=0, keepdim=True)
+            waveform = wav.squeeze(0)
+            sample_rate = sr
+
+        if isinstance(waveform, np.ndarray):
+            waveform = torch.from_numpy(waveform)
+        if not isinstance(waveform, torch.Tensor):
+            raise ValueError("Unsupported audio type for voice prompt.")
+
+        if sample_rate is not None and sample_rate != self.codec_sample_rate:
+            requires_backends(self, ["torchaudio"])
+            waveform = torchaudio.functional.resample(waveform, sample_rate, self.codec_sample_rate)
+
+        waveform = waveform.to(self.inferencer.device)
+        encode_out = self.codec.encode([waveform], **self.codec_encode_kwargs)
+        if isinstance(encode_out, dict):
+            if "codes_list" in encode_out:
+                tokens = encode_out["codes_list"][0]
+            elif "audio_codes" in encode_out:
+                tokens = encode_out["audio_codes"][0]
+            else:
+                raise ValueError("codec.encode output missing audio codes.")
+        else:
+            tokens = encode_out
+        if isinstance(tokens, torch.Tensor):
+            tokens = tokens.detach().cpu().numpy()
+        self._voice_prompt_tokens = tokens
+
+    def clear_voice_prompt(self):
+        self._voice_prompt_tokens = None
+
+    def reset_turn(
+        self,
+        user_text: Optional[str] = None,
+        user_audio_tokens: Optional[np.ndarray] = None,
+        input_ids: Optional[np.ndarray] = None,
+        include_system_prompt: Optional[bool] = None,
+        reset_cache: bool = False,
+    ):
+        if include_system_prompt is None:
+            include_system_prompt = self._turn_idx == 0
+
+        if input_ids is None:
+            if user_text is None or user_audio_tokens is None:
+                raise ValueError("user_text and user_audio_tokens are required when input_ids is not provided.")
+            user_prompt = self.processor.make_user_prompt(user_text, user_audio_tokens)
+            if include_system_prompt:
+                system_prompt = self.processor.make_ensemble(self._voice_prompt_tokens)
+                input_ids = np.concatenate([system_prompt, user_prompt], axis=0)
+            else:
+                input_ids = user_prompt
+
+        self._turn_input_ids = input_ids
+        self._turn_idx += 1
+
+        self._text_cache = ""
+        self._pending_tokens = []
+        self._prefilled = False
+        self._text_ended = False
+
+        self.inferencer.reset_generation_state(keep_cache=not reset_cache)
+
+    def push_text_tokens(self, tokens: Iterable[int]) -> list[torch.Tensor]:
+        self._pending_tokens.extend([int(t) for t in tokens])
+        return self._drain_pending_tokens()
+
+    def push_text(self, text_fragment: str) -> list[torch.Tensor]:
+        self._text_cache += text_fragment
+        segments = self._extract_text_segments(force=False)
+        for segment in segments:
+            self._pending_tokens.extend(self._tokenize(segment))
+        return self._drain_pending_tokens()
+
+    def end_text(self) -> list[torch.Tensor]:
+        self._text_ended = True
+        if self._text_cache:
+            self._pending_tokens.extend(self._tokenize(self._text_cache))
+            self._text_cache = ""
+        return self._drain_pending_tokens()
+
+    def drain(self, max_steps: Optional[int] = None) -> list[torch.Tensor]:
+        if not self._prefilled:
+            return []
+        return self.inferencer.finish(
+            max_steps=max_steps,
+            temperature=self.temperature,
+            top_p=self.top_p,
+            top_k=self.top_k,
+            do_sample=self.do_sample,
+            repetition_penalty=self.repetition_penalty,
+            repetition_window=self.repetition_window,
+        )
+
+    def _tokenize(self, text: str) -> list[int]:
+        return self.tokenizer.encode(text, add_special_tokens=False)
+
+    def _extract_text_segments(self, force: bool) -> list[str]:
+        segments = []
+        if force:
+            if self._text_cache:
+                segments.append(self._text_cache)
+                self._text_cache = ""
+            return segments
+
+        while self._text_cache:
+            cut_idx = None
+            if len(self._text_cache) >= self.min_text_chunk_chars:
+                matches = list(self._split_pattern.finditer(self._text_cache))
+                for match in matches:
+                    if match.end() >= self.min_text_chunk_chars:
+                        cut_idx = match.end()
+                        break
+            if cut_idx is None and len(self._text_cache) >= self.text_buffer_size:
+                whitespace_idx = self._text_cache.rfind(" ")
+                if whitespace_idx != -1:
+                    cut_idx = whitespace_idx + 1
+            if cut_idx is None:
+                break
+            segments.append(self._text_cache[:cut_idx])
+            self._text_cache = self._text_cache[cut_idx:]
+        return segments
+
+    def _prefill_if_needed(self) -> list[torch.Tensor]:
+        if self._prefilled:
+            return []
+        if not self._pending_tokens and not self._text_ended:
+            return []
+        if len(self._pending_tokens) < self.prefill_text_len and not self._text_ended:
+            return []
+        if self._turn_input_ids is None:
+            raise ValueError("reset_turn must be called before streaming text.")
+
+        if self._text_ended:
+            prefill_len = len(self._pending_tokens)
+        else:
+            prefill_len = min(len(self._pending_tokens), self.prefill_text_len)
+
+        if prefill_len == 0:
+            return []
+
+        prefix_tokens = [self._pending_tokens.pop(0) for _ in range(prefill_len)]
+        audio_tokens = self.inferencer.prefill(
+            input_ids=[self._turn_input_ids],
+            text_prefix_ids=[prefix_tokens],
+            temperature=self.temperature,
+            top_p=self.top_p,
+            top_k=self.top_k,
+            do_sample=self.do_sample,
+            repetition_penalty=None,
+            repetition_window=self.repetition_window,
+        )
+        self._prefilled = True
+        return [audio_tokens]
+
+    def _drain_pending_tokens(self) -> list[torch.Tensor]:
+        outputs: list[torch.Tensor] = []
+        outputs.extend(self._prefill_if_needed())
+        if not self._prefilled:
+            return outputs
+
+        while self._pending_tokens and not self.inferencer.is_finished:
+            token = self._pending_tokens.pop(0)
+            outputs.append(
+                self.inferencer.step(
+                    token,
+                    temperature=self.temperature,
+                    top_p=self.top_p,
+                    top_k=self.top_k,
+                    do_sample=self.do_sample,
+                    repetition_penalty=self.repetition_penalty,
+                    repetition_window=self.repetition_window,
+                )
+            )
+        return outputs
+
+
+@requires(backends=("torch",))
+class AudioStreamDecoder:
+    """Decode audio tokens into waveform chunks with optional crossfade."""
+
+    def __init__(
+        self,
+        codec,
+        chunk_frames: int = 40,
+        overlap_frames: int = 4,
+        initial_chunk_frames: Optional[int] = None,
+        decode_chunk_duration: Optional[float] = None,
+        decode_kwargs: Optional[dict] = None,
+        device: Optional[torch.device] = None,
+    ):
+        self.codec = codec
+        self.chunk_frames = chunk_frames
+        self.overlap_frames = overlap_frames
+        self.initial_chunk_frames = initial_chunk_frames
+        self.decode_chunk_duration = decode_chunk_duration
+        self.decode_kwargs = decode_kwargs or {}
+        self.device = device
+
+        self._buffer: list[torch.Tensor] = []
+        self._buffer_len = 0
+        self._prev_tail: Optional[torch.Tensor] = None
+        self._chunks_emitted = 0
+
+    def push_tokens(self, audio_tokens: np.ndarray | torch.Tensor):
+        if isinstance(audio_tokens, np.ndarray):
+            audio_tokens = torch.from_numpy(audio_tokens)
+        if audio_tokens.dim() != 2:
+            raise ValueError(f"Expected [T, C] audio tokens, got {tuple(audio_tokens.shape)}")
+        self._buffer.append(audio_tokens)
+        self._buffer_len += audio_tokens.shape[0]
+
+    @property
+    def _active_chunk_frames(self) -> int:
+        if self.initial_chunk_frames is not None:
+            return min(self.initial_chunk_frames + self._chunks_emitted, self.chunk_frames)
+        return self.chunk_frames
+
+    def audio_chunks(self) -> Iterable[torch.Tensor]:
+        while self._buffer_len >= self._active_chunk_frames:
+            chunk_tokens = self._consume_frames(self._active_chunk_frames)
+            wav = self._decode(chunk_tokens)
+            self._chunks_emitted += 1
+            yield self._apply_crossfade(wav)
+
+    def flush(self) -> Optional[torch.Tensor]:
+        if self._buffer_len == 0:
+            return None
+        chunk_tokens = self._consume_frames(self._buffer_len)
+        wav = self._decode(chunk_tokens)
+        return self._apply_crossfade(wav, final_chunk=True)
+
+    def _consume_frames(self, num_frames: int) -> torch.Tensor:
+        frames = []
+        remaining = num_frames
+        while remaining > 0 and self._buffer:
+            head = self._buffer[0]
+            if head.shape[0] <= remaining:
+                frames.append(head)
+                remaining -= head.shape[0]
+                self._buffer.pop(0)
+            else:
+                frames.append(head[:remaining])
+                self._buffer[0] = head[remaining:]
+                remaining = 0
+        self._buffer_len -= num_frames - remaining
+        return torch.cat(frames, dim=0)
+
+    def _decode(self, tokens: torch.Tensor) -> torch.Tensor:
+        device = self.device
+        if device is None:
+            if hasattr(self.codec, "device"):
+                device = self.codec.device
+            else:
+                try:
+                    device = next(self.codec.parameters()).device
+                except Exception:
+                    device = None
+        if device is not None:
+            tokens = tokens.to(device)
+        tokens_t = tokens.permute(1, 0)
+        decode_kwargs = dict(self.decode_kwargs) if self.decode_kwargs else {}
+        decoded = self.codec.decode(tokens_t, chunk_duration=self.decode_chunk_duration, **decode_kwargs)
+        if isinstance(decoded, dict):
+            wav = decoded["audio"][0]
+        else:
+            wav = decoded
+        if isinstance(wav, np.ndarray):
+            wav = torch.from_numpy(wav)
+        if wav.dim() > 1:
+            wav = wav.squeeze(0)
+        return wav
+
+    def _apply_crossfade(self, wav: torch.Tensor, final_chunk: bool = False) -> torch.Tensor:
+        if self.overlap_frames <= 0:
+            return wav
+        if self._prev_tail is None:
+            self._prev_tail = wav[-self._overlap_samples(wav) :].clone() if not final_chunk else None
+            return wav
+
+        overlap = self._overlap_samples(wav)
+        if overlap == 0:
+            return wav
+
+        prev_tail = self._prev_tail
+        if prev_tail.numel() < overlap:
+            overlap = prev_tail.numel()
+        if overlap == 0:
+            return wav
+
+        fade_out = torch.linspace(1.0, 0.0, overlap, device=wav.device)
+        fade_in = 1.0 - fade_out
+        cross = prev_tail[-overlap:] * fade_out + wav[:overlap] * fade_in
+        merged = torch.cat([prev_tail[:-overlap], cross, wav[overlap:]], dim=-1)
+
+        self._prev_tail = None if final_chunk else wav[-overlap:].clone()
+        return merged
+
+    def _overlap_samples(self, wav: torch.Tensor) -> int:
+        if self.chunk_frames <= 0:
+            return 0
+        return int(wav.numel() * (self.overlap_frames / self.chunk_frames))
+
+
+class TextDeltaTokenizer:
+    """
+    Convert LLM streaming text (delta) into “incremental token IDs”.
+    Notes:
+    - The input is a delta that is progressively appended to the same string
+    (consistent with the common delta output behavior in vLLM).
+    - Each time, re-encode the *full text* with the tokenizer, then take only
+    the newly added token IDs.
+    - This guarantees that tokenization is consistent with the final complete
+    text, avoiding boundary mismatches caused by tokenizing partial segments.
+    """
+
+    def __init__(self, tokenizer, *, hold_back: int = 3):
+        self.tokenizer = tokenizer
+        self.hold_back = max(0, int(hold_back))
+        self._text = ""
+        self._all_ids: list[int] = []
+        self._emitted_count: int = 0
+
+    @property
+    def text(self) -> str:
+        return self._text
+
+    @property
+    def token_ids(self) -> list[int]:
+        return list(self._all_ids)
+
+    def push_delta(self, delta: str) -> list[int]:
+        """Append a text delta and return newly stable token ids (may be empty)."""
+        if not delta:
+            return []
+        self._text += str(delta)
+        self._all_ids = self.tokenizer.encode(self._text, add_special_tokens=False)
+        # Keep the tail un-emitted because the latest tokens can still change.
+        stable_count = max(self._emitted_count, len(self._all_ids) - self.hold_back)
+        new_ids = self._all_ids[self._emitted_count : stable_count]
+        self._emitted_count = stable_count
+        return new_ids
+
+    def flush(self) -> list[int]:
+        """Emit all remaining token ids at end of stream."""
+        self._all_ids = self.tokenizer.encode(self._text, add_special_tokens=False)
+        remaining = self._all_ids[self._emitted_count :]
+        self._emitted_count = len(self._all_ids)
+        return remaining
+
+
+def _sanitize_audio_tokens(
+    tokens: torch.Tensor,
+    *,
+    codebook_size: int,
+    audio_eos_token: int,
+) -> tuple[torch.Tensor, bool]:
+    """Trim rows after EOS/invalid tokens and return whether decoding should stop."""
+    if tokens.dim() == 1:
+        tokens = tokens.unsqueeze(0)
+    if tokens.numel() == 0:
+        return tokens, False
+
+    eos_rows = (tokens[:, 0] == audio_eos_token).nonzero(as_tuple=False)
+    invalid_rows = ((tokens < 0) | (tokens >= codebook_size)).any(dim=1)
+
+    stop_idx = None
+    if eos_rows.numel() > 0:
+        stop_idx = int(eos_rows[0].item())
+    if invalid_rows.any():
+        invalid_idx = int(invalid_rows.nonzero(as_tuple=False)[0].item())
+        stop_idx = invalid_idx if stop_idx is None else min(stop_idx, invalid_idx)
+
+    if stop_idx is not None:
+        return tokens[:stop_idx], True
+    return tokens, False
+
+
+def _maybe_codec_streaming(codec, *, batch_size: int):
+    if codec is None or not hasattr(codec, "streaming"):
+        return contextlib.nullcontext()
+    return codec.streaming(batch_size=batch_size)
+
+
+@requires(backends=("torch",))
+class MossTTSRealtimeTextStreamBridge:
+    """
+    Bridge: external LLM streaming text (delta) -> TTS streaming audio chunks.
+    Usage overview:
+    - First configure `MossTTSRealtimeStreamingSession` (especially `prefill_text_len=12`).
+    - Provide an `AudioStreamDecoder`, then continuously feed the LLM delta text via
+    `push_text_delta()`.
+    - Once the accumulated token count reaches `prefill_text_len`, the session will
+    start generating audio tokens; the bridge will immediately decode them into WAV
+    chunks and yield them.
+    """
+
+    def __init__(
+        self,
+        session: MossTTSRealtimeStreamingSession,
+        decoder: AudioStreamDecoder,
+        *,
+        codebook_size: Optional[int] = None,
+        audio_eos_token: Optional[int] = None,
+        batch_size: int = 1,
+    ):
+        self.session = session
+        self.decoder = decoder
+        self.batch_size = int(batch_size)
+
+        if codebook_size is None:
+            codebook_size = int(getattr(getattr(session, "codec", None), "codebook_size", 1024))
+        if audio_eos_token is None:
+            audio_eos_token = int(getattr(session.inferencer, "audio_eos_token", 1026))
+
+        self.codebook_size = int(codebook_size)
+        self.audio_eos_token = int(audio_eos_token)
+
+    def push_text_delta(self, delta: str) -> Iterator[torch.Tensor]:
+        """
+        Push a chunk of incremental text output from the LLM and return newly generated WAV chunks.
+        Internally, this directly calls `session.push_text()`, which segments the text
+        based on punctuation/length and then tokenizes the *entire segment* at once,
+        avoiding the prefix instability issues of incremental BPE tokenization.
+        """
+        audio_frames = self.session.push_text(delta)
+        yield from self._decode_audio_frames(audio_frames)
+
+    def push_text_tokens(self, token_ids: Sequence[int]) -> Iterator[torch.Tensor]:
+        """Push token ids directly (for sources that stream token ids)."""
+        if not token_ids:
+            return
+        audio_frames = self.session.push_text_tokens(token_ids)
+        yield from self._decode_audio_frames(audio_frames)
+
+    def finish(self, *, drain_step: int = 1) -> Iterator[torch.Tensor]:
+        """Mark text stream end and emit all remaining audio chunks (including flush)."""
+        audio_frames = self.session.end_text()
+        yield from self._decode_audio_frames(audio_frames)
+
+        while True:
+            more_frames = self.session.drain(max_steps=drain_step)
+            if not more_frames:
+                break
+            yield from self._decode_audio_frames(more_frames)
+            if self.session.inferencer.is_finished:
+                break
+
+        final = self.decoder.flush()
+        if final is not None and final.numel() > 0:
+            yield final.detach().cpu()
+
+    def stream_from_text_deltas(self, deltas: Iterable[str], *, drain_step: int = 1) -> Iterator[torch.Tensor]:
+        """Consume a full delta iterator and continuously yield waveform chunks."""
+        with _maybe_codec_streaming(getattr(self.session, "codec", None), batch_size=self.batch_size):
+            for delta in deltas:
+                yield from self.push_text_delta(delta)
+            yield from self.finish(drain_step=drain_step)
+
+    def _decode_audio_frames(self, audio_frames: list[torch.Tensor]) -> Iterator[torch.Tensor]:
+        for frame in audio_frames:
+            tokens = frame
+            if tokens.dim() == 3:
+                tokens = tokens[0]
+            if tokens.dim() != 2:
+                raise ValueError(f"Expected [B, C] or [1, C] audio tokens, got {tuple(tokens.shape)}")
+            if tokens.shape[0] != 1:
+                raise ValueError(
+                    f"This bridge currently supports batch_size=1 for decoding, got batch={tokens.shape[0]}."
+                )
+
+            tokens, stop = _sanitize_audio_tokens(
+                tokens,
+                codebook_size=self.codebook_size,
+                audio_eos_token=self.audio_eos_token,
+            )
+            if tokens.numel() == 0:
+                if stop:
+                    break
+                continue
+
+            self.decoder.push_tokens(tokens.detach())
+            for wav in self.decoder.audio_chunks():
+                if wav.numel() == 0:
+                    continue
+                yield wav.detach().cpu()
+            if stop:
+                break
+
+
+__all__ = [
+    "AudioStreamDecoder",
+    "MossTTSRealtimeInference",
+    "MossTTSRealtimeStreamingSession",
+    "MossTTSRealtimeTextStreamBridge",
+    "TextDeltaTokenizer",
+]