Mirror modeling_moss_tts.py

modeling_moss_tts.py

@@ -0,0 +1,623 @@
+# coding=utf-8
+"""Modeling code for the MOSS-TTS-Local-Transformer-v1.5 HuggingFace release."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Optional, Union
+
+import torch
+import torch.nn as nn
+from transformers.modeling_outputs import BaseModelOutputWithPast
+from transformers.modeling_utils import PreTrainedModel
+from transformers.models.gpt2.configuration_gpt2 import GPT2Config
+from transformers.utils import ModelOutput
+
+from .configuration_moss_tts import MossTTSLocalConfig
+from .gpt2_decoder import MossTTSNanoGPT2Model
+from .qwen3_decoder import MossQwen3Model
+
+
+@dataclass
+class MossTTSLocalOutput(ModelOutput):
+    last_hidden_state: Optional[torch.FloatTensor] = None
+    past_key_values: Optional[tuple[tuple[torch.Tensor, torch.Tensor], ...]] = None
+    hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[tuple[torch.FloatTensor, ...]] = None
+
+
+def _find_last_equal(input_ids: torch.LongTensor, value: int) -> torch.LongTensor:
+    matches = input_ids.eq(int(value))
+    if not bool(matches.any(dim=1).all().item()):
+        raise ValueError(f"Every sample must contain token id {int(value)}.")
+    positions = torch.arange(input_ids.shape[1], device=input_ids.device, dtype=torch.long)
+    masked_positions = positions.unsqueeze(0).masked_fill(~matches, -1)
+    return masked_positions.max(dim=1).values
+
+
+class MossTTSLocalPreTrainedModel(PreTrainedModel):
+    config_class = MossTTSLocalConfig
+    base_model_prefix = "transformer"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["MossTTSNanoGPT2Block", "MossQwen3DecoderLayer"]
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _supports_cache_class = True
+
+    def _set_gradient_checkpointing(self, module: nn.Module, value: bool = False) -> None:
+        if isinstance(module, MossTTSNanoGPT2Model) or isinstance(module, MossQwen3Model):
+            module.gradient_checkpointing = value
+
+
+class MossTTSLocalModel(MossTTSLocalPreTrainedModel):
+    _tied_weights_keys = None
+
+    def __init__(self, config: MossTTSLocalConfig) -> None:
+        super().__init__(config)
+        self._tied_weights_keys = self._build_tied_weights_keys(config)
+
+        config.qwen3_config.pad_token_id = config.pad_token_id
+        config.qwen3_config._attn_implementation = config.attn_implementation
+        local_gpt2_config = config.gpt2_config.to_dict()
+        local_gpt2_config["n_layer"] = int(getattr(config, "local_transformer_layers", config.gpt2_config.n_layer))
+        local_gpt2_config["n_positions"] = int(config.n_vq) + 1
+        local_gpt2_config["n_ctx"] = int(config.n_vq) + 1
+        local_gpt2_config = GPT2Config(**local_gpt2_config)
+        local_gpt2_config.pad_token_id = config.pad_token_id
+        local_gpt2_config._attn_implementation = config.local_transformer_attn_implementation
+
+        self.transformer = MossQwen3Model(config.qwen3_config)
+        self.local_transformer = MossTTSNanoGPT2Model(
+            local_gpt2_config,
+            attn_implementation=config.local_transformer_attn_implementation,
+        )
+        self.local_transformer.wte = nn.Identity()
+
+        hidden_size = int(config.hidden_size)
+        self.audio_embeddings = nn.ModuleList(
+            [
+                nn.Embedding(int(config.audio_codebook_sizes[index]), hidden_size)
+                for index in range(config.n_vq)
+            ]
+        )
+        self.text_lm_head = nn.Linear(hidden_size, int(config.vocab_size), bias=False)
+        self.audio_lm_heads = nn.ModuleList(
+            [
+                nn.Linear(hidden_size, int(config.audio_codebook_sizes[index]), bias=False)
+                for index in range(config.n_vq)
+            ]
+        )
+        self.local_text_lm_head = (
+            nn.Linear(hidden_size, 2, bias=False)
+            if self._use_binary_local_text_head()
+            else None
+        )
+
+        self.post_init()
+        self.tie_weights()
+        self.initialize_local_text_lm_head_from_text_lm_head()
+
+    def can_generate(self) -> bool:
+        return True
+
+    @staticmethod
+    def _build_tied_weights_keys(config: MossTTSLocalConfig) -> dict[str, str]:
+        tied_weights = {"text_lm_head.weight": "transformer.embed_tokens.weight"}
+        tied_weights.update(
+            {
+                f"audio_lm_heads.{index}.weight": f"audio_embeddings.{index}.weight"
+                for index in range(config.n_vq)
+            }
+        )
+        return tied_weights
+
+    def tie_weights(self, *args, **kwargs) -> None:
+        del args, kwargs
+        self.text_lm_head.weight = self.transformer.embed_tokens.weight
+        for embedding, head in zip(self.audio_embeddings, self.audio_lm_heads):
+            head.weight = embedding.weight
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.transformer.embed_tokens
+
+    def set_input_embeddings(self, value: nn.Embedding) -> None:
+        self.transformer.embed_tokens = value
+        self.tie_weights()
+        self.initialize_local_text_lm_head_from_text_lm_head()
+
+    def get_output_embeddings(self) -> nn.Linear:
+        return self.text_lm_head
+
+    def set_output_embeddings(self, new_embeddings: nn.Linear) -> None:
+        self.text_lm_head = new_embeddings
+        self.tie_weights()
+        self.initialize_local_text_lm_head_from_text_lm_head()
+
+    def _use_binary_local_text_head(self) -> bool:
+        return str(getattr(self.config, "local_text_head_mode", "full_vocab")).strip().lower() == "binary"
+
+    def _local_text_candidate_ids(self, device: torch.device) -> torch.LongTensor:
+        return torch.tensor(
+            [
+                int(self.config.audio_assistant_slot_token_id),
+                int(self.config.audio_end_token_id),
+            ],
+            dtype=torch.long,
+            device=device,
+        )
+
+    def initialize_local_text_lm_head_from_text_lm_head(self) -> None:
+        if not self._use_binary_local_text_head() or self.local_text_lm_head is None:
+            return
+        candidate_ids = self._local_text_candidate_ids(self.text_lm_head.weight.device)
+        with torch.no_grad():
+            source_weight = self.text_lm_head.weight.index_select(0, candidate_ids)
+            if tuple(source_weight.shape) == tuple(self.local_text_lm_head.weight.shape):
+                self.local_text_lm_head.weight.copy_(
+                    source_weight.to(
+                        device=self.local_text_lm_head.weight.device,
+                        dtype=self.local_text_lm_head.weight.dtype,
+                    )
+                )
+
+    def _resolve_fixed_nq(
+        self,
+        n_vq_for_inference: Optional[int] = None,
+        nq: Optional[int] = None,
+    ) -> int:
+        requested = n_vq_for_inference if n_vq_for_inference is not None else nq
+        config_nq = int(self.config.n_vq)
+        if requested is not None and int(requested) != config_nq:
+            raise ValueError(
+                "This MOSS-TTS-Local-Transformer-v1.5 release is trained with a fixed RVQ depth. "
+                f"Expected n_vq={config_nq}, got {int(requested)}."
+            )
+        return config_nq
+
+    def _build_inputs_embeds(self, input_ids: torch.LongTensor) -> torch.FloatTensor:
+        if input_ids.ndim != 3 or input_ids.shape[-1] != self.config.n_vq + 1:
+            raise ValueError(
+                f"Expected input_ids shape [batch, seq, {self.config.n_vq + 1}], "
+                f"got {tuple(input_ids.shape)}."
+            )
+        text_ids = input_ids[..., 0]
+        inputs_embeds = self.transformer.embed_tokens(text_ids)
+        for channel_index, embedding in enumerate(self.audio_embeddings):
+            channel_ids = input_ids[..., channel_index + 1]
+            valid_mask = channel_ids.ne(self.config.audio_pad_token_id)
+            safe_ids = channel_ids.masked_fill(~valid_mask, 0)
+            audio_embeds = embedding(safe_ids) * valid_mask.unsqueeze(-1)
+            inputs_embeds = inputs_embeds + audio_embeds
+        return inputs_embeds
+
+    def _global_hidden_to_local(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:
+        return hidden_states
+
+    @staticmethod
+    def _local_past_length(past_key_values: Optional[tuple[Any, ...]]) -> int:
+        if past_key_values is None or len(past_key_values) == 0:
+            return 0
+        first_layer_past = past_key_values[0]
+        if isinstance(first_layer_past, dict) and bool(first_layer_past.get("static_kv_cache", False)):
+            return int(first_layer_past.get("length", 0))
+        return int(first_layer_past[0].shape[1])
+
+    def _new_static_local_past_key_values(
+        self,
+        batch_size: int,
+        max_length: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> tuple[dict[str, Any], ...]:
+        layers = []
+        for block in self.local_transformer.h:
+            attn = block.attn
+            cache_shape = (
+                int(batch_size),
+                int(max_length),
+                int(attn.num_heads),
+                int(attn.head_dim),
+            )
+            layers.append(
+                {
+                    "static_kv_cache": True,
+                    "key": torch.empty(cache_shape, device=device, dtype=dtype),
+                    "value": torch.empty(cache_shape, device=device, dtype=dtype),
+                    "length": 0,
+                }
+            )
+        return tuple(layers)
+
+    def _decode_local_hidden_states_with_cache(
+        self,
+        local_inputs_embeds: torch.FloatTensor,
+        past_key_values: Optional[tuple[Any, ...]] = None,
+    ) -> tuple[torch.FloatTensor, Optional[tuple[Any, ...]]]:
+        if (
+            past_key_values is None
+            and not self.training
+            and bool(getattr(self.config, "use_static_local_kv_cache", True))
+        ):
+            max_length = max(int(getattr(self.config, "n_vq", 0)) + 1, int(local_inputs_embeds.shape[1]))
+            past_key_values = self._new_static_local_past_key_values(
+                batch_size=int(local_inputs_embeds.shape[0]),
+                max_length=max_length,
+                device=local_inputs_embeds.device,
+                dtype=local_inputs_embeds.dtype,
+            )
+        past_length = self._local_past_length(past_key_values)
+        local_seq_len = int(local_inputs_embeds.shape[1])
+        local_position_ids = torch.arange(
+            past_length,
+            past_length + local_seq_len,
+            device=local_inputs_embeds.device,
+            dtype=torch.long,
+        ).unsqueeze(0)
+        if int(local_inputs_embeds.shape[0]) != 1:
+            local_position_ids = local_position_ids.expand(int(local_inputs_embeds.shape[0]), -1)
+        local_outputs = self.local_transformer(
+            input_ids=None,
+            past_key_values=past_key_values,
+            attention_mask=None,
+            position_ids=local_position_ids,
+            inputs_embeds=local_inputs_embeds,
+            use_cache=True,
+            output_attentions=False,
+            output_hidden_states=False,
+            return_dict=True,
+            cu_seqlens=None,
+            num_sequences=None,
+        )
+        return local_outputs.last_hidden_state, local_outputs.past_key_values
+
+    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[tuple[tuple[torch.Tensor, torch.Tensor], ...]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = True,
+        **kwargs,
+    ) -> Union[tuple, MossTTSLocalOutput]:
+        del kwargs
+        if inputs_embeds is None:
+            if input_ids is None:
+                raise ValueError("Either input_ids or inputs_embeds must be provided.")
+            inputs_embeds = self._build_inputs_embeds(input_ids)
+        outputs = self.transformer(
+            input_ids=None,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            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,
+            cu_seqlens=None,
+            num_sequences=None,
+        )
+        if not return_dict:
+            return (
+                outputs.last_hidden_state,
+                outputs.past_key_values,
+                outputs.hidden_states,
+                outputs.attentions,
+            )
+        return MossTTSLocalOutput(
+            last_hidden_state=outputs.last_hidden_state,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def _decode_local_last_hidden_state(
+        self,
+        local_inputs_embeds: torch.FloatTensor,
+    ) -> torch.FloatTensor:
+        local_seq_len = int(local_inputs_embeds.shape[1])
+        local_position_ids = torch.arange(
+            0,
+            local_seq_len,
+            device=local_inputs_embeds.device,
+            dtype=torch.long,
+        ).unsqueeze(0)
+        if int(local_inputs_embeds.shape[0]) != 1:
+            local_position_ids = local_position_ids.expand(int(local_inputs_embeds.shape[0]), -1)
+        local_outputs = self.local_transformer(
+            input_ids=None,
+            attention_mask=None,
+            position_ids=local_position_ids,
+            inputs_embeds=local_inputs_embeds,
+            use_cache=False,
+            output_attentions=False,
+            output_hidden_states=False,
+            return_dict=True,
+            cu_seqlens=None,
+            num_sequences=None,
+        )
+        return local_outputs.last_hidden_state[:, -1, :]
+
+    def _filter_logits(
+        self,
+        logits: torch.FloatTensor,
+        top_k: Optional[int],
+        top_p: Optional[float],
+    ) -> torch.FloatTensor:
+        scores = logits
+        if top_k is not None and int(top_k) > 0 and int(top_k) < scores.shape[-1]:
+            kth = torch.topk(scores, int(top_k), dim=-1).values[..., -1, None]
+            scores = scores.masked_fill(scores < kth, -torch.inf)
+        if top_p is not None and 0.0 < float(top_p) < 1.0:
+            sorted_scores, sorted_indices = torch.sort(scores, descending=True, dim=-1)
+            sorted_probs = torch.softmax(sorted_scores, dim=-1)
+            cumulative_probs = sorted_probs.cumsum(dim=-1)
+            sorted_mask = cumulative_probs > float(top_p)
+            sorted_mask[..., 1:] = sorted_mask[..., :-1].clone()
+            sorted_mask[..., 0] = False
+            remove_mask = torch.zeros_like(scores, dtype=torch.bool)
+            remove_mask.scatter_(dim=-1, index=sorted_indices, src=sorted_mask)
+            scores = scores.masked_fill(remove_mask, -torch.inf)
+        return scores
+
+    def _apply_repetition_penalty(
+        self,
+        scores: torch.FloatTensor,
+        previous_token_ids: Optional[torch.LongTensor],
+        penalty: float,
+    ) -> torch.FloatTensor:
+        if previous_token_ids is None or float(penalty) == 1.0:
+            return scores
+        if previous_token_ids.ndim == 1:
+            previous_token_ids = previous_token_ids.unsqueeze(0)
+        updated = scores.clone()
+        for batch_index in range(updated.shape[0]):
+            unique_token_ids = torch.unique(previous_token_ids[batch_index])
+            unique_token_ids = unique_token_ids[
+                (unique_token_ids >= 0) & (unique_token_ids < updated.shape[-1])
+            ]
+            if unique_token_ids.numel() == 0:
+                continue
+            token_scores = updated[batch_index].index_select(0, unique_token_ids)
+            token_scores = torch.where(
+                token_scores < 0,
+                token_scores * float(penalty),
+                token_scores / float(penalty),
+            )
+            updated[batch_index].scatter_(0, unique_token_ids, token_scores)
+        return updated
+
+    def _sample_next_token(
+        self,
+        logits: torch.FloatTensor,
+        do_sample: bool,
+        temperature: float,
+        top_k: Optional[int],
+        top_p: Optional[float],
+        previous_token_ids: Optional[torch.LongTensor] = None,
+        repetition_penalty: float = 1.0,
+    ) -> torch.LongTensor:
+        scores = logits.float()
+        scores = self._apply_repetition_penalty(scores, previous_token_ids, repetition_penalty)
+        if not do_sample:
+            return torch.argmax(scores, dim=-1)
+        if float(temperature) <= 0:
+            raise ValueError("temperature must be positive when do_sample=True.")
+        scores = scores / float(temperature)
+        scores = self._filter_logits(scores, top_k=top_k, top_p=top_p)
+        probs = torch.softmax(scores, dim=-1)
+        return torch.multinomial(probs, num_samples=1).squeeze(-1)
+
+    def _sample_next_assistant_text_token(
+        self,
+        local_hidden_states: torch.FloatTensor,
+        do_sample: bool,
+        temperature: float,
+        top_k: Optional[int],
+        top_p: Optional[float],
+    ) -> torch.LongTensor:
+        if self._use_binary_local_text_head() and self.local_text_lm_head is not None:
+            logits = self.local_text_lm_head(local_hidden_states)
+            sampled_indices = self._sample_next_token(
+                logits=logits,
+                do_sample=do_sample,
+                temperature=temperature,
+                top_k=top_k,
+                top_p=top_p,
+            )
+            candidate_ids = self._local_text_candidate_ids(logits.device)
+            return candidate_ids[sampled_indices]
+
+        candidate_ids = self._local_text_candidate_ids(local_hidden_states.device)
+        logits = self.text_lm_head(local_hidden_states).index_select(dim=-1, index=candidate_ids)
+        sampled_indices = self._sample_next_token(
+            logits=logits,
+            do_sample=do_sample,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+        )
+        return candidate_ids[sampled_indices]
+
+    def _build_generation_row(
+        self,
+        batch_size: int,
+        device: torch.device,
+        audio_token_ids: torch.LongTensor,
+    ) -> torch.LongTensor:
+        row = torch.full(
+            (batch_size, 1, self.config.n_vq + 1),
+            int(self.config.audio_pad_token_id),
+            dtype=torch.long,
+            device=device,
+        )
+        row[:, :, 0] = int(self.config.audio_assistant_slot_token_id)
+        row[:, :, 1:] = audio_token_ids.unsqueeze(1)
+        return row
+
+    @torch.inference_mode()
+    def generate(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        max_new_tokens: Optional[int] = None,
+        max_new_frames: Optional[int] = None,
+        do_sample: bool = True,
+        text_temperature: float = 1.0,
+        text_top_p: float = 1.0,
+        text_top_k: int = 50,
+        audio_temperature: Optional[float] = None,
+        audio_top_p: Optional[float] = None,
+        audio_top_k: Optional[int] = None,
+        audio_repetition_penalty: Optional[float] = None,
+        temperature: float = 1.0,
+        top_p: float = 0.95,
+        top_k: int = 50,
+        repetition_penalty: float = 1.0,
+        use_kv_cache: bool = True,
+        n_vq_for_inference: Optional[int] = None,
+        nq: Optional[int] = None,
+        **kwargs,
+    ) -> list[tuple[int, torch.LongTensor]]:
+        del kwargs
+        self._resolve_fixed_nq(n_vq_for_inference=n_vq_for_inference, nq=nq)
+
+        if input_ids.ndim == 2:
+            input_ids = input_ids.unsqueeze(0)
+        if input_ids.ndim != 3:
+            raise ValueError(f"Expected input_ids with 3 dims, got {tuple(input_ids.shape)}.")
+        if input_ids.shape[-1] != self.config.n_vq + 1:
+            raise ValueError(
+                f"Expected {self.config.n_vq + 1} channels from config.n_vq, got {input_ids.shape[-1]}."
+            )
+        if attention_mask is None:
+            attention_mask = torch.ones(input_ids.shape[:2], dtype=torch.bool, device=input_ids.device)
+        elif attention_mask.ndim == 1:
+            attention_mask = attention_mask.unsqueeze(0)
+        attention_mask = attention_mask.to(device=input_ids.device, dtype=torch.bool)
+
+        frame_budget = max_new_frames if max_new_frames is not None else max_new_tokens
+        if frame_budget is None:
+            frame_budget = 4096
+        frame_budget = int(frame_budget)
+
+        audio_temperature = float(temperature if audio_temperature is None else audio_temperature)
+        audio_top_p = float(top_p if audio_top_p is None else audio_top_p)
+        audio_top_k = int(top_k if audio_top_k is None else audio_top_k)
+        audio_repetition_penalty = float(
+            repetition_penalty if audio_repetition_penalty is None else audio_repetition_penalty
+        )
+
+        batch_size = input_ids.shape[0]
+        input_ids_length = input_ids.shape[1]
+        current_input_ids = input_ids
+        current_attention_mask = attention_mask
+        current_model_input_ids = current_input_ids
+        generated_frames: list[torch.LongTensor] = []
+        finished = torch.zeros(batch_size, dtype=torch.bool, device=input_ids.device)
+        past_key_values = None
+        local_dtype = self.local_transformer.ln_f.weight.dtype
+
+        for _ in range(frame_budget):
+            generated_audio_history = torch.stack(generated_frames, dim=1) if generated_frames else None
+            global_inputs_embeds = self._build_inputs_embeds(current_model_input_ids)
+            global_outputs = self.transformer(
+                input_ids=None,
+                past_key_values=past_key_values,
+                attention_mask=current_attention_mask,
+                position_ids=None,
+                inputs_embeds=global_inputs_embeds,
+                use_cache=use_kv_cache,
+                output_attentions=False,
+                output_hidden_states=False,
+                return_dict=True,
+                cu_seqlens=None,
+                num_sequences=None,
+            )
+            global_hidden_states = global_outputs.last_hidden_state[:, -1, :]
+            local_global_hidden_states = self._global_hidden_to_local(global_hidden_states).to(dtype=local_dtype)
+
+            local_prefix_hidden_states, local_prefix_past_key_values = self._decode_local_hidden_states_with_cache(
+                local_global_hidden_states.unsqueeze(1)
+            )
+            local_hidden_states = local_prefix_hidden_states[:, -1, :]
+            next_text_tokens = self._sample_next_assistant_text_token(
+                local_hidden_states=local_hidden_states,
+                do_sample=do_sample,
+                temperature=text_temperature,
+                top_k=text_top_k,
+                top_p=text_top_p,
+            )
+            should_continue = next_text_tokens.eq(int(self.config.audio_assistant_slot_token_id)) & ~finished
+            finished = finished | next_text_tokens.eq(int(self.config.audio_end_token_id))
+            if not bool(should_continue.any().item()):
+                break
+
+            next_frame_tokens = []
+            for channel_index in range(int(self.config.n_vq)):
+                channel_logits = self.audio_lm_heads[channel_index](local_hidden_states)
+                channel_token = self._sample_next_token(
+                    logits=channel_logits,
+                    do_sample=do_sample,
+                    temperature=audio_temperature,
+                    top_k=audio_top_k,
+                    top_p=audio_top_p,
+                    previous_token_ids=(
+                        None
+                        if generated_audio_history is None
+                        else generated_audio_history[:, :, channel_index]
+                    ),
+                    repetition_penalty=audio_repetition_penalty,
+                )
+                next_frame_tokens.append(channel_token)
+                if channel_index + 1 < int(self.config.n_vq):
+                    current_local_input = self.audio_embeddings[channel_index](channel_token).to(dtype=local_dtype)
+                    local_token_hidden_states, local_prefix_past_key_values = (
+                        self._decode_local_hidden_states_with_cache(
+                            current_local_input.unsqueeze(1),
+                            past_key_values=local_prefix_past_key_values,
+                        )
+                    )
+                    local_hidden_states = local_token_hidden_states[:, -1, :]
+
+            next_frame = torch.stack(next_frame_tokens, dim=-1)
+            next_frame = next_frame.masked_fill(
+                ~should_continue.unsqueeze(-1),
+                int(self.config.audio_pad_token_id),
+            )
+            generated_frames.append(next_frame)
+
+            next_row = self._build_generation_row(
+                batch_size=batch_size,
+                device=input_ids.device,
+                audio_token_ids=next_frame,
+            )
+            if bool((~should_continue).any().item()):
+                next_row[~should_continue, 0, 0] = int(self.config.pad_token_id)
+                next_row[~should_continue, 0, 1:] = int(self.config.audio_pad_token_id)
+
+            current_input_ids = torch.cat([current_input_ids, next_row], dim=1)
+            current_attention_mask = torch.cat(
+                [current_attention_mask, should_continue.unsqueeze(1)],
+                dim=1,
+            )
+            if use_kv_cache:
+                current_model_input_ids = next_row
+                past_key_values = global_outputs.past_key_values
+            else:
+                current_model_input_ids = current_input_ids
+
+        start_indices = _find_last_equal(input_ids[..., 0], int(self.config.audio_start_token_id))
+        start_lengths = input_ids_length - start_indices - 1
+        outputs: list[tuple[int, torch.LongTensor]] = []
+        for start_index, start_length, generation_ids in zip(
+            start_indices.tolist(),
+            start_lengths.tolist(),
+            current_input_ids,
+        ):
+            outputs.append((int(start_length), generation_ids[int(start_index):].detach().cpu()))
+        return outputs