# coding=utf-8 from __future__ import annotations import json import logging import os import time from contextlib import nullcontext from dataclasses import dataclass from pathlib import Path from typing import Any, Iterator, Optional, Sequence, Union import numpy as np import soundfile as sf import torch import torch.nn as nn import torchaudio from transformers import AutoModel, AutoTokenizer from transformers.modeling_outputs import ModelOutput from transformers.modeling_utils import PreTrainedModel from transformers.models.gpt2.configuration_gpt2 import GPT2Config from .configuration_nanotts import NanoTTSConfig from .gpt2_decoder import NanoGPT2Block, NanoGPT2Model from .prompting import ( build_assistant_prompt_prefix, build_prompt_token_ids, build_user_prompt_after_reference, build_user_prompt_prefix, ) from .tokenization_nanotts_sentencepiece import NanoTTSSentencePieceTokenizer @dataclass class NanoTTSOutput(ModelOutput): global_hidden_states: 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 @dataclass class NanoTTSGenerationOutput(ModelOutput): audio_token_ids: torch.LongTensor prompt_input_ids: Optional[torch.LongTensor] = None MOSS_AUDIO_TOKENIZER_NANO_TYPE = "moss-audio-tokenizer-nano" DEFAULT_MOSS_AUDIO_TOKENIZER_PRETRAINED_NAME_OR_PATH = "OpenMOSS-Team/MOSS-Audio-Tokenizer-Nano" DEFAULT_VOICE_CLONE_MAX_TEXT_TOKENS = 50 DEFAULT_VOICE_CLONE_MAX_MEMORY_PER_SAMPLE_GB = 1.0 DEFAULT_VOICE_CLONE_INTER_CHUNK_PAUSE_SHORT_SECONDS = 0.40 DEFAULT_VOICE_CLONE_INTER_CHUNK_PAUSE_LONG_SECONDS = 0.24 _SENTENCE_END_PUNCTUATION = frozenset(".!?。！？；;") _CLAUSE_SPLIT_PUNCTUATION = frozenset(",，、；;：:") _CLOSING_PUNCTUATION = frozenset("\"'”’)]}）】》」』") class NanoTTSPreTrainedModel(PreTrainedModel): config_class = NanoTTSConfig base_model_prefix = "transformer" supports_gradient_checkpointing = False _no_split_modules = ["NanoGPT2Block"] _supports_flash_attn_2 = True _supports_sdpa = True class NanoTTSGlobalLocalForCausalLM(NanoTTSPreTrainedModel): _keys_to_ignore_on_load_unexpected = [r"local_transformer\.wte\.weight"] def __init__(self, config: NanoTTSConfig) -> None: super().__init__(config) config.gpt2_config.pad_token_id = config.pad_token_id config.gpt2_config._attn_implementation = config.attn_implementation self.transformer = NanoGPT2Model( config.gpt2_config, attn_implementation=config.attn_implementation, ) hidden_size = config.gpt2_config.hidden_size init_std = config.gpt2_config.initializer_range 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, config.gpt2_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) ] ) local_gpt2_config = config.gpt2_config.to_dict() local_gpt2_config["n_layer"] = int(config.local_transformer_layers) local_gpt2_config["n_positions"] = config.n_vq + 1 local_gpt2_config["n_ctx"] = config.n_vq + 1 self.local_transformer = NanoGPT2Model( GPT2Config(**local_gpt2_config), attn_implementation=str(config.local_transformer_attn_implementation), ) self.local_transformer.wte = nn.Identity() for module in list(self.audio_embeddings) + [self.text_lm_head] + list(self.audio_lm_heads): if hasattr(module, "weight") and module.weight is not None: nn.init.normal_(module.weight, mean=0.0, std=init_std) self._tied_weights_keys = tuple(self.all_tied_weights_keys.keys()) self.tie_weights() @property def all_tied_weights_keys(self) -> dict[str, str]: tied_weights = {"text_lm_head.weight": "transformer.wte.weight"} tied_weights.update( { f"audio_lm_heads.{index}.weight": f"audio_embeddings.{index}.weight" for index in range(self.config.n_vq) } ) return tied_weights def tie_weights(self, *args, **kwargs) -> None: del args, kwargs self.text_lm_head.weight = self.transformer.wte.weight for embedding, lm_head in zip(self.audio_embeddings, self.audio_lm_heads): lm_head.weight = embedding.weight def get_input_embeddings(self) -> nn.Embedding: return self.transformer.wte def set_input_embeddings(self, value: nn.Embedding) -> None: self.transformer.wte = value self.tie_weights() 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}], got {tuple(input_ids.shape)}" ) text_ids = input_ids[..., 0] inputs_embeds = self.transformer.wte(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) invalid_mask = valid_mask & ((channel_ids < 0) | (channel_ids >= embedding.num_embeddings)) if invalid_mask.any(): invalid_token_ids = channel_ids[invalid_mask] raise ValueError( "Found out-of-range audio token ids for channel " f"{channel_index}: min={int(invalid_token_ids.min().item())} " f"max={int(invalid_token_ids.max().item())} " f"codebook_size={embedding.num_embeddings} " f"audio_pad_token_id={self.config.audio_pad_token_id}" ) safe_ids = channel_ids.masked_fill(~valid_mask, 0) audio_embeds = embedding(safe_ids) audio_embeds = audio_embeds * valid_mask.unsqueeze(-1) inputs_embeds = inputs_embeds + audio_embeds return inputs_embeds def forward( self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = 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] = None, **kwargs, ): labels = kwargs.pop("labels", None) if labels is not None: raise NotImplementedError("This open-source package is inference-only and does not support training forward.") if kwargs: ignored = ", ".join(sorted(kwargs.keys())) logging.debug("ignoring unsupported forward kwargs: %s", ignored) return_dict = self.config.use_return_dict if return_dict is None else return_dict 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, past_key_values=past_key_values, attention_mask=attention_mask, position_ids=None, 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 NanoTTSOutput( global_hidden_states=outputs.last_hidden_state, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, ) def _build_text_rows( self, token_ids: list[int], device: torch.device, ) -> torch.LongTensor: rows = torch.full( (len(token_ids), self.config.n_vq + 1), self.config.audio_pad_token_id, dtype=torch.long, device=device, ) if token_ids: rows[:, 0] = torch.tensor(token_ids, dtype=torch.long, device=device) return rows def _encode_text(self, tokenizer, text: str) -> list[int]: try: return list(tokenizer.encode(text, add_special_tokens=False)) except TypeError: return list(tokenizer.encode(text)) @staticmethod def _contains_cjk(text: str) -> bool: return any( "\u4e00" <= ch <= "\u9fff" or "\u3400" <= ch <= "\u4dbf" or "\u3040" <= ch <= "\u30ff" or "\uac00" <= ch <= "\ud7af" for ch in str(text) ) @staticmethod def _prepare_text_for_sentence_chunking(text: str) -> str: normalized_text = str(text).strip() if normalized_text == "": raise ValueError("Text prompt cannot be empty.") normalized_text = normalized_text.replace("\n", " ").replace("\r", " ") while " " in normalized_text: normalized_text = normalized_text.replace(" ", " ") contains_cjk = NanoTTSGlobalLocalForCausalLM._contains_cjk(normalized_text) if contains_cjk: if normalized_text[-1] not in _SENTENCE_END_PUNCTUATION: normalized_text = normalized_text + "。" return normalized_text if not normalized_text[0].isupper(): normalized_text = normalized_text[0].upper() + normalized_text[1:] if normalized_text[-1].isalnum(): normalized_text = normalized_text + "." if len(normalized_text.split()) < 5: normalized_text = " " * 8 + normalized_text return normalized_text @staticmethod def _split_text_by_punctuation(text: str, punctuation: set[str] | frozenset[str]) -> list[str]: sentences: list[str] = [] current_chars: list[str] = [] text = str(text) index = 0 while index < len(text): char = text[index] current_chars.append(char) if char in punctuation: lookahead = index + 1 while lookahead < len(text) and text[lookahead] in _CLOSING_PUNCTUATION: current_chars.append(text[lookahead]) lookahead += 1 sentence = "".join(current_chars).strip() if sentence: sentences.append(sentence) current_chars = [] while lookahead < len(text) and text[lookahead].isspace(): lookahead += 1 index = lookahead continue index += 1 tail = "".join(current_chars).strip() if tail: sentences.append(tail) return sentences def _count_text_tokens(self, text_tokenizer, text: str) -> int: return len(self._encode_text(text_tokenizer, text)) def _split_text_by_token_budget( self, text_tokenizer, text: str, max_tokens: int, ) -> list[str]: remaining_text = str(text).strip() if remaining_text == "": return [] pieces: list[str] = [] preferred_boundary_chars = _CLAUSE_SPLIT_PUNCTUATION | _SENTENCE_END_PUNCTUATION | frozenset({" "}) while remaining_text: if self._count_text_tokens(text_tokenizer, remaining_text) <= int(max_tokens): pieces.append(remaining_text) break low = 1 high = len(remaining_text) best_prefix_length = 1 while low <= high: middle = (low + high) // 2 candidate = remaining_text[:middle].strip() if not candidate: low = middle + 1 continue if self._count_text_tokens(text_tokenizer, candidate) <= int(max_tokens): best_prefix_length = middle low = middle + 1 else: high = middle - 1 cut_index = best_prefix_length prefix = remaining_text[:best_prefix_length] preferred_index = -1 for scan_index in range(len(prefix) - 1, max(-1, len(prefix) - 25), -1): if prefix[scan_index] in preferred_boundary_chars: preferred_index = scan_index + 1 break if preferred_index > 0: cut_index = preferred_index piece = remaining_text[:cut_index].strip() if not piece: piece = remaining_text[:best_prefix_length].strip() cut_index = best_prefix_length pieces.append(piece) remaining_text = remaining_text[cut_index:].strip() return pieces @staticmethod def _join_sentence_parts(left: str, right: str) -> str: if not left: return right if not right: return left if NanoTTSGlobalLocalForCausalLM._contains_cjk(left) or NanoTTSGlobalLocalForCausalLM._contains_cjk(right): return left + right return left + " " + right def _split_text_into_best_sentences( self, text_tokenizer, text: str, max_tokens: int, ) -> list[str]: if int(max_tokens) <= 0: return [str(text)] prepared_text = self._prepare_text_for_sentence_chunking(text) sentence_candidates = self._split_text_by_punctuation(prepared_text, punctuation=_SENTENCE_END_PUNCTUATION) if not sentence_candidates: sentence_candidates = [prepared_text.strip()] sentence_slices: list[tuple[int, str]] = [] for sentence_text in sentence_candidates: normalized_sentence = sentence_text.strip() if not normalized_sentence: continue sentence_token_count = self._count_text_tokens(text_tokenizer, normalized_sentence) if sentence_token_count <= int(max_tokens): sentence_slices.append((sentence_token_count, normalized_sentence)) continue clause_candidates = self._split_text_by_punctuation( normalized_sentence, punctuation=_CLAUSE_SPLIT_PUNCTUATION, ) if len(clause_candidates) <= 1: clause_candidates = [normalized_sentence] for clause_text in clause_candidates: normalized_clause = clause_text.strip() if not normalized_clause: continue clause_token_count = self._count_text_tokens(text_tokenizer, normalized_clause) if clause_token_count <= int(max_tokens): sentence_slices.append((clause_token_count, normalized_clause)) continue for piece in self._split_text_by_token_budget( text_tokenizer=text_tokenizer, text=normalized_clause, max_tokens=max_tokens, ): normalized_piece = piece.strip() if normalized_piece: sentence_slices.append( (self._count_text_tokens(text_tokenizer, normalized_piece), normalized_piece) ) chunks: list[str] = [] current_chunk = "" current_chunk_token_count = 0 for sentence_token_count, sentence_text in sentence_slices: if current_chunk == "": current_chunk = sentence_text current_chunk_token_count = sentence_token_count continue if current_chunk_token_count + sentence_token_count > int(max_tokens): chunks.append(current_chunk.strip()) current_chunk = sentence_text current_chunk_token_count = sentence_token_count else: current_chunk = self._join_sentence_parts(current_chunk, sentence_text) current_chunk_token_count = self._count_text_tokens(text_tokenizer, current_chunk) if current_chunk: chunks.append(current_chunk.strip()) return chunks or [prepared_text.strip()] @staticmethod def _estimate_voice_clone_inter_chunk_pause_seconds(text_chunk: str) -> float: return ( DEFAULT_VOICE_CLONE_INTER_CHUNK_PAUSE_SHORT_SECONDS if len(str(text_chunk).strip().split()) <= 4 else DEFAULT_VOICE_CLONE_INTER_CHUNK_PAUSE_LONG_SECONDS ) def _concat_voice_clone_waveform_chunks( self, waveform_chunks: list[torch.FloatTensor], text_chunks: list[str], sample_rate: int, ) -> torch.FloatTensor: if not waveform_chunks: return torch.zeros((1, 0), dtype=torch.float32) if len(waveform_chunks) != len(text_chunks): raise ValueError("waveform_chunks and text_chunks must have the same length.") if len(waveform_chunks) == 1: return waveform_chunks[0] segments: list[torch.FloatTensor] = [] for chunk_index, waveform_chunk in enumerate(waveform_chunks): segments.append(waveform_chunk) if chunk_index >= len(waveform_chunks) - 1: continue pause_seconds = self._estimate_voice_clone_inter_chunk_pause_seconds(text_chunks[chunk_index]) pause_samples = max(0, int(round(float(sample_rate) * pause_seconds))) if pause_samples > 0: silence = torch.zeros((waveform_chunk.shape[0], pause_samples), dtype=waveform_chunk.dtype) segments.append(silence) return torch.cat(segments, dim=-1) @staticmethod def _resolve_inference_mode( mode: str, has_prompt_text: bool, has_prompt_audio: bool, ) -> str: normalized_mode = str(mode or "continuation").strip().lower() or "continuation" if normalized_mode not in {"continuation", "voice_clone"}: raise ValueError(f"Unsupported inference mode {mode!r}.") if normalized_mode == "voice_clone": if not has_prompt_audio: raise ValueError("voice_clone mode requires prompt_audio_path.") if has_prompt_text: raise ValueError("voice_clone mode does not accept prompt_text.") elif has_prompt_text != has_prompt_audio: raise ValueError( "continuation mode accepts either target text only, or prompt_text and prompt_audio_path together." ) return normalized_mode def _resolve_inference_nq(self, nq: Optional[int] = None) -> int: if nq is None: return int(self.config.n_vq) resolved_nq = int(nq) if resolved_nq < 1 or resolved_nq > int(self.config.n_vq): raise ValueError(f"nq must be in [1, {self.config.n_vq}], got {resolved_nq}.") return resolved_nq def _mask_unused_audio_channels( self, audio_token_ids: torch.LongTensor, nq: int, ) -> torch.LongTensor: tensor = torch.as_tensor(audio_token_ids, dtype=torch.long) if tensor.shape[-1] != self.config.n_vq: raise ValueError( f"Expected audio token ids with trailing dim {self.config.n_vq}, got {tuple(tensor.shape)}" ) if nq < self.config.n_vq: tensor = tensor.clone() tensor[..., nq:] = self.config.audio_pad_token_id return tensor def _build_audio_prefix_rows( self, prompt_audio_codes: torch.LongTensor, slot_token_id: int, device: torch.device, ) -> torch.LongTensor: rows = torch.full( (int(prompt_audio_codes.shape[0]), self.config.n_vq + 1), self.config.audio_pad_token_id, dtype=torch.long, device=device, ) if rows.shape[0] > 0: rows[:, 0] = int(slot_token_id) rows[:, 1:] = prompt_audio_codes return rows def build_inference_input_ids( self, text: str, text_tokenizer, mode: str = "continuation", prompt_text: Optional[str] = None, prompt_audio_codes: Optional[torch.LongTensor] = None, device: Optional[Union[str, torch.device]] = None, ) -> tuple[torch.LongTensor, torch.BoolTensor]: resolved_device = self._resolve_device(device) resolved_mode = self._resolve_inference_mode( mode=mode, has_prompt_text=prompt_text is not None, has_prompt_audio=prompt_audio_codes is not None, ) if resolved_mode == "voice_clone": assert prompt_audio_codes is not None text_token_ids = self._encode_text(text_tokenizer, text) prompt_token_ids = build_user_prompt_prefix(text_tokenizer, self.config) + [self.config.audio_start_token_id] suffix_token_ids = ( [self.config.audio_end_token_id] + build_user_prompt_after_reference(text_tokenizer) + text_token_ids + build_assistant_prompt_prefix(text_tokenizer, self.config) + [self.config.audio_start_token_id] ) sections = [ self._build_text_rows(prompt_token_ids, device=resolved_device), self._build_audio_prefix_rows( prompt_audio_codes=prompt_audio_codes.to(resolved_device), slot_token_id=self.config.audio_user_slot_token_id, device=resolved_device, ), self._build_text_rows(suffix_token_ids, device=resolved_device), ] input_ids = torch.cat(sections, dim=0).unsqueeze(0) attention_mask = torch.ones(input_ids.shape[:2], dtype=torch.bool, device=resolved_device) return input_ids, attention_mask effective_text = text if prompt_text is None else prompt_text + text prompt_token_ids = build_prompt_token_ids( tokenizer=text_tokenizer, config=self.config, text_token_ids=self._encode_text(text_tokenizer, effective_text), ) sections = [ self._build_text_rows(prompt_token_ids, device=resolved_device), self._build_text_rows([self.config.audio_start_token_id], device=resolved_device), ] if prompt_audio_codes is not None: sections.append( self._build_audio_prefix_rows( prompt_audio_codes=prompt_audio_codes.to(resolved_device), slot_token_id=self.config.audio_assistant_slot_token_id, device=resolved_device, ) ) input_ids = torch.cat(sections, dim=0).unsqueeze(0) attention_mask = torch.ones(input_ids.shape[:2], dtype=torch.bool, device=resolved_device) return input_ids, attention_mask def _left_pad_inference_batch( self, input_id_batches: list[torch.LongTensor], attention_mask_batches: list[torch.BoolTensor], device: torch.device, ) -> tuple[torch.LongTensor, torch.BoolTensor]: if not input_id_batches: raise ValueError("input_id_batches must not be empty.") if len(input_id_batches) != len(attention_mask_batches): raise ValueError("input_id_batches and attention_mask_batches must have the same length.") batch_size = len(input_id_batches) max_seq_len = max(int(batch.shape[1]) for batch in input_id_batches) row_width = self.config.n_vq + 1 padded_input_ids = torch.full( (batch_size, max_seq_len, row_width), self.config.audio_pad_token_id, dtype=torch.long, device=device, ) padded_input_ids[:, :, 0] = self.config.pad_token_id padded_attention_mask = torch.zeros((batch_size, max_seq_len), dtype=torch.bool, device=device) for batch_index, (input_ids, attention_mask) in enumerate(zip(input_id_batches, attention_mask_batches)): normalized_input_ids = input_ids.squeeze(0).to(device=device, dtype=torch.long) normalized_attention_mask = attention_mask.squeeze(0).to(device=device, dtype=torch.bool) seq_len = int(normalized_input_ids.shape[0]) padded_input_ids[batch_index, -seq_len:, :] = normalized_input_ids padded_attention_mask[batch_index, -seq_len:] = normalized_attention_mask return padded_input_ids, padded_attention_mask def _trim_generated_audio_token_ids( self, audio_token_ids: torch.LongTensor, effective_nq: int, ) -> torch.LongTensor: tensor = self._mask_unused_audio_channels(audio_token_ids, nq=effective_nq) if tensor.ndim != 2: raise ValueError(f"Expected a 2D audio token tensor, got {tuple(tensor.shape)}") valid_rows = tensor[:, :effective_nq].ne(self.config.audio_pad_token_id).any(dim=-1) if not bool(valid_rows.any()): return tensor[:0] last_valid_index = int(torch.nonzero(valid_rows, as_tuple=False)[-1].item()) + 1 return tensor[:last_valid_index] def _resolve_voice_clone_chunk_batch_size( self, *, resolved_device: torch.device, chunk_count: int, max_memory_per_sample_gb: float, ) -> int: if chunk_count <= 1 or max_memory_per_sample_gb <= 0 or resolved_device.type != "cuda": return 1 if not hasattr(torch.cuda, "mem_get_info"): return 1 try: free_bytes, _ = torch.cuda.mem_get_info(resolved_device) except Exception: return 1 bytes_per_sample = int(float(max_memory_per_sample_gb) * (1024**3)) if bytes_per_sample <= 0: return 1 usable_free_bytes = max(0, int(free_bytes * 0.9)) batch_size = max(1, usable_free_bytes // bytes_per_sample) resolved_batch_size = max(1, min(int(chunk_count), int(batch_size))) logging.info( "voice_clone chunk batching device=%s free_gb=%.2f max_memory_per_sample_gb=%.2f resolved_batch_size=%d chunk_count=%d", resolved_device, float(free_bytes) / float(1024**3), float(max_memory_per_sample_gb), resolved_batch_size, int(chunk_count), ) return resolved_batch_size @staticmethod def _resolve_requested_batch_size_limit(requested_batch_size: Optional[int]) -> Optional[int]: if requested_batch_size is None: return None resolved_batch_size = int(requested_batch_size) if resolved_batch_size <= 0: return None return max(1, resolved_batch_size) def _resolve_effective_voice_clone_batch_sizes( self, *, resolved_device: torch.device, chunk_count: int, max_memory_per_sample_gb: float, requested_tts_max_batch_size: Optional[int] = None, requested_codec_max_batch_size: Optional[int] = None, realtime_streaming: bool = False, ) -> tuple[int, int]: effective_tts_batch_size = self._resolve_voice_clone_chunk_batch_size( resolved_device=resolved_device, chunk_count=chunk_count, max_memory_per_sample_gb=max_memory_per_sample_gb, ) requested_tts_limit = self._resolve_requested_batch_size_limit(requested_tts_max_batch_size) requested_codec_limit = self._resolve_requested_batch_size_limit(requested_codec_max_batch_size) if requested_tts_limit is not None: effective_tts_batch_size = min(effective_tts_batch_size, requested_tts_limit) if realtime_streaming and requested_codec_limit is not None: effective_tts_batch_size = min(effective_tts_batch_size, requested_codec_limit) effective_tts_batch_size = max(1, min(int(chunk_count), int(effective_tts_batch_size))) if realtime_streaming: effective_codec_batch_size = effective_tts_batch_size elif requested_codec_limit is None: effective_codec_batch_size = 1 else: effective_codec_batch_size = max(1, min(int(requested_codec_limit), int(effective_tts_batch_size))) return int(effective_tts_batch_size), int(effective_codec_batch_size) def _generate_audio_token_ids_with_fallback( self, *, prompt_input_ids: torch.LongTensor, attention_mask: torch.BoolTensor, effective_nq: int, max_new_frames: int, do_sample: bool, text_temperature: float, text_top_p: float, text_top_k: int, audio_temperature: float, audio_top_p: float, audio_top_k: int, audio_repetition_penalty: float, use_kv_cache: bool, resolved_device: torch.device, ) -> torch.LongTensor: try: generation = self.generate( input_ids=prompt_input_ids, attention_mask=attention_mask, nq=effective_nq, max_new_frames=max_new_frames, do_sample=do_sample, text_temperature=text_temperature, text_top_p=text_top_p, text_top_k=text_top_k, audio_temperature=audio_temperature, audio_top_p=audio_top_p, audio_top_k=audio_top_k, audio_repetition_penalty=audio_repetition_penalty, use_kv_cache=use_kv_cache, return_dict_in_generate=True, ) except (RuntimeError, ValueError) as exc: if not self._is_generation_stability_error(exc): raise self._apply_inference_stability_fallback(resolved_device) generation = self.generate( input_ids=prompt_input_ids, attention_mask=attention_mask, nq=effective_nq, max_new_frames=max_new_frames, do_sample=do_sample, text_temperature=text_temperature, text_top_p=text_top_p, text_top_k=text_top_k, audio_temperature=audio_temperature, audio_top_p=audio_top_p, audio_top_k=audio_top_k, audio_repetition_penalty=audio_repetition_penalty, use_kv_cache=use_kv_cache, return_dict_in_generate=True, ) return self._mask_unused_audio_channels(generation.audio_token_ids, nq=effective_nq) def _decode_audio_token_ids_to_waveform( self, *, audio_tokenizer, audio_token_ids: torch.LongTensor, target_sample_rate: int, effective_nq: int, resolved_device: torch.device, ) -> tuple[torch.FloatTensor, int]: decoded = self._call_audio_decode( audio_tokenizer=audio_tokenizer, audio_token_ids=audio_token_ids.to(resolved_device), sample_rate=target_sample_rate, nq=effective_nq, ) return self._extract_waveform_and_sample_rate(decoded, fallback_sample_rate=target_sample_rate) def _decode_audio_token_id_batch_to_waveforms( self, *, audio_tokenizer, audio_token_id_batches: Sequence[torch.LongTensor], target_sample_rate: int, effective_nq: int, resolved_device: torch.device, ) -> tuple[list[torch.FloatTensor], int]: if not audio_token_id_batches: return [], target_sample_rate if len(audio_token_id_batches) == 1: waveform, sample_rate = self._decode_audio_token_ids_to_waveform( audio_tokenizer=audio_tokenizer, audio_token_ids=audio_token_id_batches[0], target_sample_rate=target_sample_rate, effective_nq=effective_nq, resolved_device=resolved_device, ) return [waveform], sample_rate decode_codes = [ self._prepare_audio_codes_for_decode(audio_token_ids.to(resolved_device), nq=effective_nq) for audio_token_ids in audio_token_id_batches ] batch_decode_fn = getattr(audio_tokenizer, "batch_decode", None) if batch_decode_fn is None: raise AttributeError("audio_tokenizer must provide a batch_decode method.") try: with self._audio_tokenizer_inference_context(audio_tokenizer, resolved_device): decode_output = batch_decode_fn( decode_codes, num_quantizers=effective_nq, chunk_duration=None, ) return self._extract_batch_waveforms_and_sample_rate( decode_output, fallback_sample_rate=target_sample_rate, batch_size=len(audio_token_id_batches), ) except Exception: logging.warning( "batched audio decode failed; falling back to per-chunk decode for batch_size=%d", len(audio_token_id_batches), exc_info=True, ) waveform_rows: list[torch.FloatTensor] = [] sample_rate = target_sample_rate for audio_token_ids in audio_token_id_batches: waveform_row, sample_rate = self._decode_audio_token_ids_to_waveform( audio_tokenizer=audio_tokenizer, audio_token_ids=audio_token_ids, target_sample_rate=target_sample_rate, effective_nq=effective_nq, resolved_device=resolved_device, ) waveform_rows.append(waveform_row) return waveform_rows, sample_rate 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), self.config.audio_pad_token_id, dtype=torch.long, device=device, ) row[:, :, 0] = self.config.audio_assistant_slot_token_id row[:, :, 1:] = audio_token_ids.unsqueeze(1) return row @staticmethod def _compute_stream_lead_seconds( emitted_samples_total: int, sample_rate: int, first_audio_emitted_at: Optional[float], ) -> float: if first_audio_emitted_at is None or sample_rate <= 0: return 0.0 elapsed_seconds = max(0.0, time.monotonic() - first_audio_emitted_at) emitted_seconds = float(emitted_samples_total) / float(sample_rate) return emitted_seconds - elapsed_seconds @staticmethod def _resolve_stream_decode_frame_budget( *, emitted_samples_total: int, sample_rate: int, first_audio_emitted_at: Optional[float], ) -> int: lead_seconds = NanoTTSGlobalLocalForCausalLM._compute_stream_lead_seconds( emitted_samples_total=emitted_samples_total, sample_rate=sample_rate, first_audio_emitted_at=first_audio_emitted_at, ) if first_audio_emitted_at is None or lead_seconds < 0.20: return 1 if lead_seconds < 0.55: return 2 if lead_seconds < 1.10: return 4 return 8 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 = self._apply_repetition_penalty( logits=logits, previous_token_ids=previous_token_ids, repetition_penalty=repetition_penalty, ) if not do_sample: return scores.argmax(dim=-1) if temperature <= 0: raise ValueError("temperature must be positive when do_sample=True") scores = scores / temperature if top_k is not None and top_k > 0: top_k = min(top_k, scores.shape[-1]) threshold = torch.topk(scores, top_k, dim=-1).values[..., -1, None] scores = scores.masked_fill(scores < threshold, float("-inf")) if top_p is not None and 0.0 < top_p < 1.0: sorted_scores, sorted_indices = torch.sort(scores, descending=True, dim=-1) sorted_probs = torch.softmax(sorted_scores, dim=-1) sorted_cumsum = torch.cumsum(sorted_probs, dim=-1) sorted_remove = sorted_cumsum > top_p sorted_remove[..., 1:] = sorted_remove[..., :-1].clone() sorted_remove[..., 0] = False sorted_scores = sorted_scores.masked_fill(sorted_remove, float("-inf")) scores = torch.full_like(scores, float("-inf")) scores.scatter_(dim=-1, index=sorted_indices, src=sorted_scores) probs = torch.softmax(scores, dim=-1) return torch.multinomial(probs, num_samples=1).squeeze(-1) @staticmethod def _ensure_finite_generation_logits(logits: torch.FloatTensor, name: str) -> None: if torch.isfinite(logits).all(): return finite_mask = torch.isfinite(logits) finite_logits = logits[finite_mask] min_value = float(finite_logits.min().item()) if finite_logits.numel() > 0 else float("nan") max_value = float(finite_logits.max().item()) if finite_logits.numel() > 0 else float("nan") raise RuntimeError( f"Non-finite {name} during generation: dtype={logits.dtype} " f"shape={tuple(logits.shape)} finite={int(finite_mask.sum().item())}/{int(logits.numel())} " f"min={min_value} max={max_value}" ) def _apply_repetition_penalty( self, logits: torch.FloatTensor, previous_token_ids: Optional[torch.LongTensor], repetition_penalty: float, ) -> torch.FloatTensor: if repetition_penalty <= 0: raise ValueError("repetition_penalty must be positive") if repetition_penalty == 1.0 or previous_token_ids is None: return logits token_ids = torch.as_tensor(previous_token_ids, device=logits.device, dtype=torch.long) if token_ids.ndim == 1: token_ids = token_ids.unsqueeze(0) elif token_ids.ndim > 2: token_ids = token_ids.reshape(token_ids.shape[0], -1) scores = logits.clone() vocab_size = scores.shape[-1] for batch_index in range(scores.shape[0]): valid_token_ids = token_ids[batch_index] valid_token_ids = valid_token_ids[(valid_token_ids >= 0) & (valid_token_ids < vocab_size)] if valid_token_ids.numel() == 0: continue unique_token_ids = torch.unique(valid_token_ids) token_scores = scores[batch_index].index_select(0, unique_token_ids) token_scores = torch.where( token_scores < 0, token_scores * repetition_penalty, token_scores / repetition_penalty, ) scores[batch_index].scatter_(0, unique_token_ids, token_scores) return scores def _sample_next_assistant_text_token( self, logits: torch.FloatTensor, do_sample: bool, temperature: float, top_k: Optional[int] = None, top_p: Optional[float] = None, ) -> torch.LongTensor: candidate_ids = torch.tensor( [ self.config.audio_assistant_slot_token_id, self.config.audio_end_token_id, ], dtype=torch.long, device=logits.device, ) candidate_logits = logits.index_select(dim=-1, index=candidate_ids) sampled_indices = self._sample_next_token( logits=candidate_logits, do_sample=do_sample, temperature=temperature, top_k=top_k, top_p=top_p, ) return candidate_ids[sampled_indices] def _resolve_device(self, device: Optional[Union[str, torch.device]] = None) -> torch.device: return torch.device(device) if device is not None else next(self.parameters()).device @staticmethod def _looks_like_hf_tokenizer_dir(candidate_path: Path) -> bool: if not candidate_path.is_dir(): return False if (candidate_path / "tokenizer.model").is_file(): return True if (candidate_path / "tokenizer.json").is_file(): return True if (candidate_path / "tokenizer_config.json").is_file() and ( (candidate_path / "vocab.json").is_file() or (candidate_path / "merges.txt").is_file() or (candidate_path / "special_tokens_map.json").is_file() ): return True return False def _resolve_text_tokenizer_path(self, raw_path: Union[str, Path]) -> Path: candidate_path = Path(raw_path) if candidate_path.is_file() and candidate_path.suffix == ".model": return candidate_path if not candidate_path.exists(): raise FileNotFoundError(f"Tokenizer path does not exist: {candidate_path}") if candidate_path.is_dir(): if (candidate_path / "tokenizer.model").is_file(): return candidate_path if self._looks_like_hf_tokenizer_dir(candidate_path): return candidate_path hf_dir = candidate_path / "hf_tokenizer" if self._looks_like_hf_tokenizer_dir(hf_dir): return hf_dir sentencepiece_model = candidate_path / "sentencepiece" / "nanotts_spm_bpe.model" if sentencepiece_model.is_file(): return sentencepiece_model final_summary_path = candidate_path / "final_summary.json" if final_summary_path.is_file(): final_summary = json.loads(final_summary_path.read_text(encoding="utf-8")) latest_hf_dir = final_summary.get("latest_hf_tokenizer_dir") if latest_hf_dir: latest_hf_path = Path(str(latest_hf_dir)) if self._looks_like_hf_tokenizer_dir(latest_hf_path): return latest_hf_path raise ValueError( "Could not resolve a tokenizer from the provided path. Expected a tokenizer dir, experiment dir, or SentencePiece .model file." ) def _load_resolved_text_tokenizer(self, resolved_path: Path, cache_dir: str): if resolved_path.is_file() and resolved_path.suffix == ".model": return NanoTTSSentencePieceTokenizer(vocab_file=str(resolved_path)) try: return AutoTokenizer.from_pretrained( str(resolved_path), trust_remote_code=True, use_fast=bool(self.config.tokenizer_use_fast), local_files_only=True, cache_dir=cache_dir, ) except Exception: model_path = resolved_path / "tokenizer.model" if model_path.is_file(): return NanoTTSSentencePieceTokenizer(vocab_file=str(model_path)) raise @staticmethod def _resolve_hf_cache_dir() -> str: cache_dir = Path(__file__).resolve().parent / ".cache" / "huggingface" cache_dir.mkdir(parents=True, exist_ok=True) return str(cache_dir) @staticmethod def _patch_hf_dynamic_module_cache_dir(cache_dir: str) -> None: import transformers.dynamic_module_utils as dynamic_module_utils modules_cache_dir = str(Path(cache_dir) / "modules") Path(modules_cache_dir).mkdir(parents=True, exist_ok=True) os.environ["HF_MODULES_CACHE"] = modules_cache_dir dynamic_module_utils.HF_MODULES_CACHE = modules_cache_dir def _resolve_default_text_tokenizer_path(self) -> Path: candidates: list[Path] = [] raw_name_or_path = getattr(self.config, "_name_or_path", None) if raw_name_or_path: candidates.append(Path(str(raw_name_or_path)).expanduser()) raw_model_name_or_path = getattr(self, "name_or_path", None) if raw_model_name_or_path: candidates.append(Path(str(raw_model_name_or_path)).expanduser()) candidates.append(Path(__file__).resolve().parent) checked: set[str] = set() for candidate in candidates: resolved_candidate = candidate.resolve() key = str(resolved_candidate) if key in checked: continue checked.add(key) if (resolved_candidate / "tokenizer.model").is_file(): return resolved_candidate if self._looks_like_hf_tokenizer_dir(resolved_candidate): return resolved_candidate return candidates[0].resolve() def _load_text_tokenizer(self, text_tokenizer=None, text_tokenizer_path: Optional[str] = None): if text_tokenizer is not None: return text_tokenizer resolved_path = ( self._resolve_text_tokenizer_path(text_tokenizer_path) if text_tokenizer_path is not None else self._resolve_default_text_tokenizer_path() ) normalized_path = str(resolved_path.resolve()) cached = getattr(self, "_cached_text_tokenizer", None) cached_path = getattr(self, "_cached_text_tokenizer_path", None) if cached is not None and cached_path == normalized_path: return cached cache_dir = self._resolve_hf_cache_dir() self._patch_hf_dynamic_module_cache_dir(cache_dir) tokenizer = self._load_resolved_text_tokenizer(resolved_path=resolved_path, cache_dir=cache_dir) if tokenizer.pad_token_id is None and tokenizer.eos_token is not None: tokenizer.pad_token = tokenizer.eos_token self._cached_text_tokenizer = tokenizer self._cached_text_tokenizer_path = normalized_path return tokenizer @staticmethod def _normalize_audio_tokenizer_type(audio_tokenizer_type: Optional[str]) -> Optional[str]: if audio_tokenizer_type is None: return None normalized = str(audio_tokenizer_type).strip().lower() if not normalized: return None if normalized == MOSS_AUDIO_TOKENIZER_NANO_TYPE: return MOSS_AUDIO_TOKENIZER_NANO_TYPE raise ValueError( "Unsupported audio tokenizer type. " f"The open-source package only supports '{MOSS_AUDIO_TOKENIZER_NANO_TYPE}'." ) def _resolve_audio_tokenizer_type(self, audio_tokenizer_type: Optional[str]) -> str: explicit_type = self._normalize_audio_tokenizer_type(audio_tokenizer_type) if explicit_type is not None: return explicit_type config_type = self._normalize_audio_tokenizer_type(getattr(self.config, "audio_tokenizer_type", None)) return MOSS_AUDIO_TOKENIZER_NANO_TYPE if config_type is None else config_type @staticmethod def _set_decoder_attention_implementation(decoder, attn_implementation: str) -> None: decoder.attn_implementation = str(attn_implementation) if getattr(decoder, "config", None) is not None: decoder.config._attn_implementation = str(attn_implementation) for block in getattr(decoder, "h", []): block.attn.attn_implementation = str(attn_implementation) def _set_attention_implementation( self, attn_implementation: str, local_attn_implementation: Optional[str] = None, ) -> None: resolved_global = str(attn_implementation) resolved_local = resolved_global if local_attn_implementation is None else str(local_attn_implementation) self.config.attn_implementation = resolved_global self.config.gpt2_config._attn_implementation = resolved_global self._set_decoder_attention_implementation(self.transformer, resolved_global) self.config.local_transformer_attn_implementation = resolved_local self._set_decoder_attention_implementation(self.local_transformer, resolved_local) @staticmethod def _select_fallback_attention_implementation(device: torch.device) -> str: return "sdpa" if device.type == "cuda" else "eager" @staticmethod def _is_generation_stability_error(exc: Exception) -> bool: message = str(exc) return any( marker in message for marker in ( "Non-finite", "device-side assert triggered", "probability tensor contains either", "flash_attention_2 requires fp16/bf16 tensors", ) ) def _apply_inference_stability_fallback(self, device: torch.device) -> None: fallback_attn = self._select_fallback_attention_implementation(device) if next(self.parameters()).dtype != torch.float32: self.to(device=device, dtype=torch.float32) self._set_attention_implementation(fallback_attn) logging.warning( "retrying inference with dtype=float32 attn_implementation=%s due to numerical instability", fallback_attn, ) def _load_audio_tokenizer( self, audio_tokenizer=None, audio_tokenizer_type: Optional[str] = None, audio_tokenizer_pretrained_name_or_path: Optional[str] = None, device: Optional[Union[str, torch.device]] = None, ): if audio_tokenizer is not None: return audio_tokenizer resolved_type = self._resolve_audio_tokenizer_type(audio_tokenizer_type=audio_tokenizer_type) if resolved_type != MOSS_AUDIO_TOKENIZER_NANO_TYPE: raise ValueError( f"Unsupported audio tokenizer type {resolved_type!r}; expected '{MOSS_AUDIO_TOKENIZER_NANO_TYPE}'." ) resolved_pretrained_name_or_path = ( audio_tokenizer_pretrained_name_or_path or getattr(self.config, "audio_tokenizer_pretrained_name_or_path", None) or DEFAULT_MOSS_AUDIO_TOKENIZER_PRETRAINED_NAME_OR_PATH ) candidate_path = Path(str(resolved_pretrained_name_or_path)).expanduser() if candidate_path.exists(): load_source = str(candidate_path.resolve()) load_kwargs: dict[str, object] = { "trust_remote_code": True, "local_files_only": True, "force_download": True, } cache_key = f"{resolved_type}|{load_source}" else: load_source = str(resolved_pretrained_name_or_path) load_kwargs = { "trust_remote_code": True, } cache_key = f"{resolved_type}|{load_source}" cached = getattr(self, "_cached_audio_tokenizer", None) cached_path = getattr(self, "_cached_audio_tokenizer_path", None) if cached is not None and cached_path == cache_key: tokenizer = cached else: tokenizer = AutoModel.from_pretrained(load_source, **load_kwargs) if hasattr(tokenizer, "eval"): tokenizer.eval() self._cached_audio_tokenizer = tokenizer self._cached_audio_tokenizer_path = cache_key resolved_device = self._resolve_device(device) return tokenizer.to(resolved_device) if hasattr(tokenizer, "to") else tokenizer @staticmethod def _extract_tensor_candidate(output: Any) -> Any: if torch.is_tensor(output) or isinstance(output, np.ndarray): return output for attr_name in ("audio_codes", "audio_token_ids", "codes", "tokens", "input_ids"): value = getattr(output, attr_name, None) if value is not None: return value if isinstance(output, dict): for key in ("audio_codes", "audio_token_ids", "codes", "tokens", "input_ids"): if key in output: return output[key] if len(output) == 1: return next(iter(output.values())) if isinstance(output, (list, tuple)) and output: if len(output) == 2 and isinstance(output[1], (int, float)): return output[0] return NanoTTSGlobalLocalForCausalLM._extract_tensor_candidate(output[0]) raise TypeError(f"Unsupported audio tokenizer output type: {type(output)!r}") @staticmethod def _extract_audio_code_length(output: Any) -> Optional[int]: for attr_name in ("audio_codes_lengths", "audio_token_ids_lengths", "codes_lengths", "lengths"): candidate = getattr(output, attr_name, None) if candidate is not None: lengths = torch.as_tensor(candidate).reshape(-1) if lengths.numel() > 0: return int(lengths[0].item()) if isinstance(output, dict): for key in ("audio_codes_lengths", "audio_token_ids_lengths", "codes_lengths", "lengths"): if key in output: lengths = torch.as_tensor(output[key]).reshape(-1) if lengths.numel() > 0: return int(lengths[0].item()) if isinstance(output, (list, tuple)) and len(output) >= 2: candidate = output[1] if torch.is_tensor(candidate) or isinstance(candidate, np.ndarray): lengths = torch.as_tensor(candidate).reshape(-1) if lengths.numel() > 0: return int(lengths[0].item()) if isinstance(candidate, (int, float)): return int(candidate) return None def _normalize_audio_codes(self, audio_codes: Any) -> torch.LongTensor: code_length = self._extract_audio_code_length(audio_codes) tensor = torch.as_tensor(self._extract_tensor_candidate(audio_codes)) if tensor.ndim == 1: tensor = tensor.unsqueeze(-1) if tensor.ndim == 3: if tensor.shape[1] == 1 and tensor.shape[0] >= self.config.n_vq: tensor = tensor[: self.config.n_vq, 0, :].transpose(0, 1) elif tensor.shape[0] == 1: tensor = tensor[0] elif tensor.shape[1] == self.config.n_vq: tensor = tensor.transpose(1, 2)[0] elif tensor.shape[-1] == self.config.n_vq: tensor = tensor[0] else: raise ValueError(f"Unable to normalize audio codes with shape {tuple(tensor.shape)}") if tensor.ndim != 2: raise ValueError(f"Expected audio codes with 2 dims after normalization, got {tuple(tensor.shape)}") if tensor.shape[-1] != self.config.n_vq and tensor.shape[0] == self.config.n_vq: tensor = tensor.transpose(0, 1) elif tensor.shape[-1] != self.config.n_vq and tensor.shape[0] > self.config.n_vq: tensor = tensor[: self.config.n_vq].transpose(0, 1) elif tensor.shape[-1] > self.config.n_vq: tensor = tensor[:, : self.config.n_vq] if tensor.shape[-1] != self.config.n_vq: raise ValueError( f"Expected normalized audio codes with trailing dim {self.config.n_vq}, got {tuple(tensor.shape)}" ) if code_length is not None: tensor = tensor[:code_length] return tensor.to(dtype=torch.long) def _extract_waveform_and_sample_rate( self, decode_output: Any, fallback_sample_rate: int, ) -> tuple[torch.FloatTensor, int]: sample_rate = fallback_sample_rate waveform = decode_output waveform_length = None for key in ("sample_rate", "sampling_rate"): value = getattr(decode_output, key, None) if value is not None: sample_rate = int(value) break for key in ("waveform", "audio", "wav", "samples"): value = getattr(decode_output, key, None) if value is not None: waveform = value break for key in ("audio_lengths", "waveform_lengths", "lengths"): value = getattr(decode_output, key, None) if value is not None: lengths = torch.as_tensor(value).reshape(-1) if lengths.numel() > 0: waveform_length = int(lengths[0].item()) break if isinstance(decode_output, dict): for key in ("sample_rate", "sampling_rate"): if key in decode_output: sample_rate = int(decode_output[key]) break for key in ("waveform", "audio", "wav", "samples"): if key in decode_output: waveform = decode_output[key] break for key in ("audio_lengths", "waveform_lengths", "lengths"): if key in decode_output: lengths = torch.as_tensor(decode_output[key]).reshape(-1) if lengths.numel() > 0: waveform_length = int(lengths[0].item()) break elif isinstance(decode_output, (list, tuple)) and decode_output: if len(decode_output) == 2 and isinstance(decode_output[1], (int, float)): waveform = decode_output[0] sample_rate = int(decode_output[1]) else: waveform = decode_output[0] waveform_tensor = torch.as_tensor(waveform, dtype=torch.float32) if waveform_tensor.ndim == 3 and waveform_tensor.shape[0] == 1: waveform_tensor = waveform_tensor[0] if waveform_tensor.ndim == 2 and waveform_tensor.shape[0] > waveform_tensor.shape[1]: waveform_tensor = waveform_tensor.transpose(0, 1) if waveform_tensor.ndim == 1: waveform_tensor = waveform_tensor.unsqueeze(0) if waveform_tensor.ndim != 2: raise ValueError(f"Expected decoded waveform with 2 dims, got {tuple(waveform_tensor.shape)}") if waveform_length is not None: waveform_tensor = waveform_tensor[..., : max(0, waveform_length)] return waveform_tensor.cpu(), sample_rate def _call_audio_encode( self, audio_tokenizer, waveform: torch.FloatTensor, sample_rate: int, ) -> Any: del sample_rate batch_encode_fn = getattr(audio_tokenizer, "batch_encode", None) if batch_encode_fn is None: raise AttributeError("audio_tokenizer must provide a batch_encode method.") waveform_tensor = torch.as_tensor(waveform, dtype=torch.float32, device=self._resolve_device(waveform.device)) if waveform_tensor.ndim == 1: waveform_tensor = waveform_tensor.unsqueeze(0) if waveform_tensor.ndim != 2: raise ValueError( f"MOSS audio tokenizer encode expects waveform shaped like (C, T), got {tuple(waveform_tensor.shape)}" ) with self._audio_tokenizer_inference_context(audio_tokenizer, waveform_tensor.device): return batch_encode_fn([waveform_tensor], chunk_duration=None) def _call_audio_decode( self, audio_tokenizer, audio_token_ids: torch.LongTensor, sample_rate: int, nq: Optional[int] = None, ) -> Any: del sample_rate batch_decode_fn = getattr(audio_tokenizer, "batch_decode", None) if batch_decode_fn is None: raise AttributeError("audio_tokenizer must provide a batch_decode method.") effective_nq = self._resolve_inference_nq(nq) decode_codes = self._prepare_audio_codes_for_decode(audio_token_ids, nq=effective_nq) with self._audio_tokenizer_inference_context(audio_tokenizer, decode_codes.device): return batch_decode_fn([decode_codes], num_quantizers=effective_nq, chunk_duration=None) def _extract_batch_waveforms_and_sample_rate( self, decode_output: Any, fallback_sample_rate: int, batch_size: int, ) -> tuple[list[torch.FloatTensor], int]: sample_rate = fallback_sample_rate audio = decode_output audio_lengths = None for key in ("sample_rate", "sampling_rate"): value = getattr(decode_output, key, None) if value is not None: sample_rate = int(value) break for key in ("waveform", "audio", "wav", "samples"): value = getattr(decode_output, key, None) if value is not None: audio = value break for key in ("audio_lengths", "waveform_lengths", "lengths"): value = getattr(decode_output, key, None) if value is not None: audio_lengths = value break if isinstance(decode_output, dict): for key in ("sample_rate", "sampling_rate"): if key in decode_output: sample_rate = int(decode_output[key]) break for key in ("waveform", "audio", "wav", "samples"): if key in decode_output: audio = decode_output[key] break for key in ("audio_lengths", "waveform_lengths", "lengths"): if key in decode_output: audio_lengths = decode_output[key] break audio_tensor = torch.as_tensor(audio, dtype=torch.float32) if audio_tensor.ndim == 2: audio_tensor = audio_tensor.unsqueeze(0) if audio_tensor.ndim != 3: raise ValueError(f"Expected batched decoded audio with 3 dims, got {tuple(audio_tensor.shape)}") if audio_tensor.shape[0] != int(batch_size): raise ValueError( f"Expected decoded batch size {batch_size}, got audio tensor shape {tuple(audio_tensor.shape)}" ) if audio_lengths is None: lengths_tensor = torch.full( (batch_size,), int(audio_tensor.shape[-1]), device=audio_tensor.device, dtype=torch.long, ) else: lengths_tensor = torch.as_tensor(audio_lengths, dtype=torch.long, device=audio_tensor.device).reshape(-1) if lengths_tensor.numel() != int(batch_size): raise ValueError(f"Expected {batch_size} decoded audio lengths, got {int(lengths_tensor.numel())}") waveform_rows: list[torch.FloatTensor] = [] for row_index in range(batch_size): row_length = max(0, int(lengths_tensor[row_index].item())) waveform_rows.append(audio_tensor[row_index, :, :row_length].detach().cpu()) return waveform_rows, sample_rate @staticmethod def _resolve_audio_tokenizer_downsample_rate(audio_tokenizer) -> int: for holder in (audio_tokenizer, getattr(audio_tokenizer, "config", None)): if holder is None: continue for attr_name in ("downsample_rate", "hop_length", "frame_size"): value = getattr(holder, attr_name, None) if value is not None: return int(value) sampling_rate = getattr(holder, "sampling_rate", None) frame_rate = getattr(holder, "frame_rate", None) if sampling_rate is not None and frame_rate not in (None, 0): return int(round(float(sampling_rate) / float(frame_rate))) raise ValueError("audio_tokenizer.downsample_rate is required for prompt-audio decoding.") def _resolve_audio_tokenizer_sample_rate(self, audio_tokenizer) -> int: for holder in (audio_tokenizer, getattr(audio_tokenizer, "config", None)): if holder is None: continue for attr_name in ("sampling_rate", "sample_rate"): value = getattr(holder, attr_name, None) if value is not None: return int(value) return int(self.config.audio_tokenizer_sample_rate) @staticmethod def _resolve_audio_tokenizer_channels(audio_tokenizer) -> int: for holder in (audio_tokenizer, getattr(audio_tokenizer, "config", None)): if holder is None: continue for attr_name in ("number_channels", "channels_numbers", "audio_channels", "channels", "num_channels"): value = getattr(holder, attr_name, None) if value is not None: return int(value) return 1 @staticmethod def _audio_tokenizer_inference_context(audio_tokenizer, device: Union[str, torch.device]): del audio_tokenizer, device return nullcontext() def _prepare_audio_codes_for_decode( self, audio_token_ids: torch.LongTensor, nq: Optional[int] = None, ) -> torch.LongTensor: effective_nq = self._resolve_inference_nq(nq) tensor = torch.as_tensor(audio_token_ids, dtype=torch.long) if tensor.ndim == 2: if tensor.shape[-1] == self.config.n_vq and tensor.shape[0] != self.config.n_vq: return tensor[:, :effective_nq].transpose(0, 1).contiguous() if tensor.shape[0] == self.config.n_vq: return tensor[:effective_nq].contiguous() elif tensor.ndim == 3: if tensor.shape[-1] == self.config.n_vq: return tensor[..., :effective_nq].permute(2, 0, 1).contiguous() if tensor.shape[0] == self.config.n_vq: return tensor[:effective_nq].contiguous() raise ValueError( f"Expected generated audio token ids shaped like (T, {self.config.n_vq}) or ({self.config.n_vq}, T); got {tuple(tensor.shape)}" ) def _load_reference_audio( self, reference_audio_path: Union[str, Path], target_sample_rate: int, target_channels: int, ) -> tuple[torch.FloatTensor, int]: try: waveform, sample_rate = torchaudio.load(str(reference_audio_path)) except RuntimeError as exc: logging.warning( "torchaudio.load failed for %s; falling back to soundfile", reference_audio_path, exc_info=True, ) audio_array, sample_rate = sf.read(str(reference_audio_path), dtype="float32", always_2d=True) waveform = torch.from_numpy(audio_array.T).contiguous() waveform = waveform.to(torch.float32) if sample_rate != target_sample_rate: waveform = torchaudio.functional.resample(waveform, sample_rate, target_sample_rate) sample_rate = target_sample_rate current_channels = int(waveform.shape[0]) if current_channels == target_channels: return waveform, sample_rate if current_channels == 1 and target_channels > 1: return waveform.repeat(target_channels, 1), sample_rate if current_channels > 1 and target_channels == 1: return waveform.mean(dim=0, keepdim=True), sample_rate raise ValueError(f"Unsupported reference audio channel conversion: {current_channels} -> {target_channels}") @staticmethod def _save_audio( output_path: Union[str, Path], waveform: torch.Tensor, sample_rate: int, ) -> None: path = Path(output_path) path.parent.mkdir(parents=True, exist_ok=True) try: torchaudio.save(str(path), waveform, sample_rate) except RuntimeError: logging.warning( "torchaudio.save failed for %s; falling back to soundfile", path, exc_info=True, ) waveform_np = waveform.detach().cpu().to(torch.float32).numpy().T sf.write(str(path), waveform_np, sample_rate) def _decode_local_last_hidden_state( self, local_inputs_embeds: torch.FloatTensor, ) -> torch.FloatTensor: local_attention_mask = torch.ones( local_inputs_embeds.shape[:2], dtype=torch.bool, device=local_inputs_embeds.device, ) local_outputs = self.local_transformer( input_ids=None, attention_mask=local_attention_mask, position_ids=None, 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 _iter_generation_events( self, input_ids: torch.LongTensor, attention_mask: Optional[torch.Tensor] = None, nq: Optional[int] = None, max_new_frames: int = 300, do_sample: bool = False, text_temperature: float = 1.5, text_top_p: float = 1.0, text_top_k: int = 50, audio_temperature: float = 1.7, audio_top_p: float = 0.8, audio_top_k: int = 25, audio_repetition_penalty: float = 1.0, use_kv_cache: bool = True, return_dict_in_generate: bool = True, ) -> Iterator[dict[str, Any]]: 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 shape {tuple(input_ids.shape)}") 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) effective_nq = self._resolve_inference_nq(nq) batch_size = input_ids.shape[0] current_input_ids = input_ids current_attention_mask = attention_mask.to(device=input_ids.device) current_model_input_ids = current_input_ids generated_frames = [] 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 step_index in range(max_new_frames): 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, :].to(dtype=local_dtype) local_inputs_embeds = global_hidden_states.unsqueeze(1) local_hidden_states = self._decode_local_last_hidden_state(local_inputs_embeds) text_logits = self.text_lm_head(local_hidden_states) self._ensure_finite_generation_logits(text_logits, "text logits") next_text_tokens = self._sample_next_assistant_text_token( logits=text_logits, do_sample=do_sample, temperature=text_temperature, top_k=text_top_k, top_p=text_top_p, ) should_continue = next_text_tokens.eq(self.config.audio_assistant_slot_token_id) & ~finished finished = finished | next_text_tokens.eq(self.config.audio_end_token_id) if not should_continue.any(): break next_frame_tokens = [] current_local_input = self.transformer.wte(next_text_tokens).to(dtype=local_dtype) for channel_index in range(effective_nq): local_inputs_embeds = torch.cat([local_inputs_embeds, current_local_input.unsqueeze(1)], dim=1) local_hidden_states = self._decode_local_last_hidden_state(local_inputs_embeds) channel_logits = self.audio_lm_heads[channel_index](local_hidden_states) self._ensure_finite_generation_logits(channel_logits, f"audio logits[{channel_index}]") 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) current_local_input = self.audio_embeddings[channel_index](channel_token).to(dtype=local_dtype) next_frame_prefix = torch.stack(next_frame_tokens, dim=-1) if effective_nq < self.config.n_vq: next_frame = torch.full( (batch_size, self.config.n_vq), self.config.audio_pad_token_id, dtype=next_frame_prefix.dtype, device=next_frame_prefix.device, ) next_frame[:, :effective_nq] = next_frame_prefix else: next_frame = next_frame_prefix padded_next_frame = next_frame.masked_fill(~should_continue.unsqueeze(-1), self.config.audio_pad_token_id) generated_frames.append(padded_next_frame) next_row = self._build_generation_row( batch_size=batch_size, device=input_ids.device, audio_token_ids=padded_next_frame, ) if (~should_continue).any(): next_row[~should_continue, 0, 0] = self.config.pad_token_id next_row[~should_continue, 0, 1:] = 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 yield { "type": "frame", "step_index": int(step_index), "audio_token_ids": padded_next_frame.detach().clone(), "active_mask": should_continue.detach().clone(), "finished_mask": finished.detach().clone(), } if generated_frames: audio_token_ids = torch.stack(generated_frames, dim=1) else: audio_token_ids = torch.empty((batch_size, 0, self.config.n_vq), dtype=torch.long, device=input_ids.device) if not return_dict_in_generate: yield {"type": "final", "audio_token_ids": audio_token_ids} return yield { "type": "final", "generation": NanoTTSGenerationOutput(audio_token_ids=audio_token_ids, prompt_input_ids=input_ids), } @torch.no_grad() def generate( self, input_ids: torch.LongTensor, attention_mask: Optional[torch.Tensor] = None, nq: Optional[int] = None, max_new_frames: int = 300, do_sample: bool = False, text_temperature: float = 1.5, text_top_p: float = 1.0, text_top_k: int = 50, audio_temperature: float = 1.7, audio_top_p: float = 0.8, audio_top_k: int = 25, audio_repetition_penalty: float = 1.0, use_kv_cache: bool = True, return_dict_in_generate: bool = True, ): final_output: Any = None for event in self._iter_generation_events( input_ids=input_ids, attention_mask=attention_mask, nq=nq, max_new_frames=max_new_frames, do_sample=do_sample, text_temperature=text_temperature, text_top_p=text_top_p, text_top_k=text_top_k, audio_temperature=audio_temperature, audio_top_p=audio_top_p, audio_top_k=audio_top_k, audio_repetition_penalty=audio_repetition_penalty, use_kv_cache=use_kv_cache, return_dict_in_generate=return_dict_in_generate, ): if event["type"] != "final": continue final_output = event.get("generation", event.get("audio_token_ids")) if final_output is None: raise RuntimeError("Generation finished without producing a final output.") return final_output @torch.no_grad() def generate_stream( self, input_ids: torch.LongTensor, attention_mask: Optional[torch.Tensor] = None, nq: Optional[int] = None, max_new_frames: int = 300, do_sample: bool = False, text_temperature: float = 1.5, text_top_p: float = 1.0, text_top_k: int = 50, audio_temperature: float = 1.7, audio_top_p: float = 0.8, audio_top_k: int = 25, audio_repetition_penalty: float = 1.0, use_kv_cache: bool = True, return_dict_in_generate: bool = True, ) -> Iterator[dict[str, Any]]: yield from self._iter_generation_events( input_ids=input_ids, attention_mask=attention_mask, nq=nq, max_new_frames=max_new_frames, do_sample=do_sample, text_temperature=text_temperature, text_top_p=text_top_p, text_top_k=text_top_k, audio_temperature=audio_temperature, audio_top_p=audio_top_p, audio_top_k=audio_top_k, audio_repetition_penalty=audio_repetition_penalty, use_kv_cache=use_kv_cache, return_dict_in_generate=return_dict_in_generate, ) @torch.no_grad() def inference_stream( self, text: str, output_audio_path: Union[str, Path], mode: str = "continuation", prompt_text: Optional[str] = None, prompt_audio_path: Optional[Union[str, Path]] = None, reference_audio_path: Optional[Union[str, Path]] = None, text_tokenizer=None, text_tokenizer_path: Optional[str] = None, audio_tokenizer=None, audio_tokenizer_type: Optional[str] = None, audio_tokenizer_pretrained_name_or_path: Optional[str] = None, device: Optional[Union[str, torch.device]] = None, nq: Optional[int] = None, max_new_frames: int = 300, do_sample: bool = False, text_temperature: float = 1.5, text_top_p: float = 1.0, text_top_k: int = 50, audio_temperature: float = 1.7, audio_top_p: float = 0.8, audio_top_k: int = 25, audio_repetition_penalty: float = 1.0, use_kv_cache: bool = True, voice_clone_max_text_tokens: int = DEFAULT_VOICE_CLONE_MAX_TEXT_TOKENS, voice_clone_max_memory_per_sample_gb: float = DEFAULT_VOICE_CLONE_MAX_MEMORY_PER_SAMPLE_GB, tts_max_batch_size: int = 0, codec_max_batch_size: int = 0, ) -> Iterator[dict[str, Any]]: resolved_device = self._resolve_device(device) effective_nq = self._resolve_inference_nq(nq) if next(self.parameters()).device != resolved_device: self.to(resolved_device) was_training = self.training self.eval() text_tokenizer = self._load_text_tokenizer( text_tokenizer=text_tokenizer, text_tokenizer_path=text_tokenizer_path, ) audio_tokenizer = self._load_audio_tokenizer( audio_tokenizer=audio_tokenizer, audio_tokenizer_type=audio_tokenizer_type, audio_tokenizer_pretrained_name_or_path=audio_tokenizer_pretrained_name_or_path, device=resolved_device, ) target_sample_rate = self._resolve_audio_tokenizer_sample_rate(audio_tokenizer) target_channels = self._resolve_audio_tokenizer_channels(audio_tokenizer) effective_prompt_audio_path = prompt_audio_path or reference_audio_path resolved_mode = self._resolve_inference_mode( mode=mode, has_prompt_text=prompt_text is not None, has_prompt_audio=effective_prompt_audio_path is not None, ) if reference_audio_path is not None and prompt_audio_path is None: logging.warning( "reference_audio_path=%s is treated as prompt_audio_path for backward compatibility.", reference_audio_path, ) prompt_audio_codes = None if effective_prompt_audio_path is not None: waveform, sample_rate = self._load_reference_audio( effective_prompt_audio_path, target_sample_rate, target_channels, ) encoded = self._call_audio_encode( audio_tokenizer=audio_tokenizer, waveform=waveform.to(resolved_device), sample_rate=sample_rate, ) prompt_audio_codes = self._mask_unused_audio_channels( self._normalize_audio_codes(encoded), nq=effective_nq, ).to(resolved_device) if resolved_mode == "voice_clone": split_voice_clone_text_chunks = self._split_text_into_best_sentences( text_tokenizer=text_tokenizer, text=text, max_tokens=voice_clone_max_text_tokens, ) voice_clone_text_chunks = split_voice_clone_text_chunks if len(split_voice_clone_text_chunks) > 1 else [text] else: voice_clone_text_chunks = [text] if resolved_mode == "voice_clone" and len(voice_clone_text_chunks) > 1: voice_clone_chunk_batch_size, voice_clone_codec_batch_size = self._resolve_effective_voice_clone_batch_sizes( resolved_device=resolved_device, chunk_count=len(voice_clone_text_chunks), max_memory_per_sample_gb=float(voice_clone_max_memory_per_sample_gb), requested_tts_max_batch_size=tts_max_batch_size, requested_codec_max_batch_size=codec_max_batch_size, realtime_streaming=True, ) else: voice_clone_chunk_batch_size = 1 voice_clone_codec_batch_size = 1 generated_audio_token_chunks: list[torch.LongTensor] = [] emitted_waveform_segments: list[torch.FloatTensor] = [] decoded_sample_rate: Optional[int] = None emitted_samples_total = 0 first_audio_emitted_at: Optional[float] = None streaming_reset_fn = getattr(audio_tokenizer, "_reset_batch_decode_streaming_state", None) try: for batch_start in range(0, len(voice_clone_text_chunks), voice_clone_chunk_batch_size): batch_chunks = voice_clone_text_chunks[batch_start : batch_start + voice_clone_chunk_batch_size] batch_prompt_input_ids: list[torch.LongTensor] = [] batch_attention_masks: list[torch.BoolTensor] = [] for text_chunk in batch_chunks: prompt_input_ids, attention_mask = self.build_inference_input_ids( text=text_chunk, text_tokenizer=text_tokenizer, mode=resolved_mode, prompt_text=prompt_text, prompt_audio_codes=prompt_audio_codes, device=resolved_device, ) batch_prompt_input_ids.append(prompt_input_ids) batch_attention_masks.append(attention_mask) batched_prompt_input_ids, batched_attention_mask = self._left_pad_inference_batch( input_id_batches=batch_prompt_input_ids, attention_mask_batches=batch_attention_masks, device=resolved_device, ) row_states = [ { "pending_decode_frames": [], "decoded_audio_segments": [], "generation_complete": False, "pause_emitted": False, } for _ in batch_chunks ] batch_emit_index = 0 codec_stream_started = False if resolved_mode == "continuation" and prompt_audio_codes is not None: prompt_decode_codes = self._prepare_audio_codes_for_decode(prompt_audio_codes, nq=effective_nq) _ = audio_tokenizer.batch_decode( [prompt_decode_codes], num_quantizers=effective_nq, streaming=True, max_batch_size=1, reset_stream=True, ) codec_stream_started = True def _emit_ready_segments() -> Iterator[dict[str, Any]]: nonlocal batch_emit_index, emitted_samples_total, first_audio_emitted_at active_sample_rate = decoded_sample_rate or target_sample_rate while batch_emit_index < len(batch_chunks): state = row_states[batch_emit_index] decoded_segments = state["decoded_audio_segments"] if decoded_segments: next_segment = decoded_segments.pop(0) if next_segment.numel() == 0 or int(next_segment.shape[-1]) <= 0: continue emitted_waveform_segments.append(next_segment) if first_audio_emitted_at is None: first_audio_emitted_at = time.monotonic() emitted_samples_total += int(next_segment.shape[-1]) yield { "type": "audio", "waveform": next_segment, "sample_rate": active_sample_rate, "chunk_index": batch_start + batch_emit_index, "is_pause": False, "emitted_audio_seconds": float(emitted_samples_total) / float(active_sample_rate), "lead_seconds": self._compute_stream_lead_seconds( emitted_samples_total=emitted_samples_total, sample_rate=active_sample_rate, first_audio_emitted_at=first_audio_emitted_at, ), } continue if not state["generation_complete"]: break if ( resolved_mode == "voice_clone" and len(voice_clone_text_chunks) > 1 and not state["pause_emitted"] and (batch_start + batch_emit_index) < len(voice_clone_text_chunks) - 1 ): state["pause_emitted"] = True pause_seconds = self._estimate_voice_clone_inter_chunk_pause_seconds( voice_clone_text_chunks[batch_start + batch_emit_index] ) pause_samples = max(0, int(round(float(active_sample_rate) * pause_seconds))) if pause_samples > 0: silence = torch.zeros((target_channels, pause_samples), dtype=torch.float32) emitted_waveform_segments.append(silence) if first_audio_emitted_at is None: first_audio_emitted_at = time.monotonic() emitted_samples_total += int(silence.shape[-1]) yield { "type": "audio", "waveform": silence, "sample_rate": active_sample_rate, "chunk_index": batch_start + batch_emit_index, "is_pause": True, "emitted_audio_seconds": float(emitted_samples_total) / float(active_sample_rate), "lead_seconds": self._compute_stream_lead_seconds( emitted_samples_total=emitted_samples_total, sample_rate=active_sample_rate, first_audio_emitted_at=first_audio_emitted_at, ), } batch_emit_index += 1 continue batch_emit_index += 1 def _maybe_decode_pending(force: bool) -> Iterator[dict[str, Any]]: nonlocal codec_stream_started, decoded_sample_rate pending_counts = [len(state["pending_decode_frames"]) for state in row_states] total_pending = sum(pending_counts) if total_pending <= 0: return active_sample_rate = decoded_sample_rate or target_sample_rate head_pending = pending_counts[batch_emit_index] if batch_emit_index < len(batch_chunks) else 0 lead_seconds = self._compute_stream_lead_seconds( emitted_samples_total=emitted_samples_total, sample_rate=active_sample_rate, first_audio_emitted_at=first_audio_emitted_at, ) decode_budget = self._resolve_stream_decode_frame_budget( emitted_samples_total=emitted_samples_total, sample_rate=active_sample_rate, first_audio_emitted_at=first_audio_emitted_at, ) if not force: should_decode = False if first_audio_emitted_at is None and head_pending > 0: should_decode = True elif head_pending > 0 and lead_seconds < 0.45: should_decode = True elif max(pending_counts) >= decode_budget: should_decode = True elif lead_seconds < 0.0 and total_pending > 0: should_decode = True if not should_decode: return decode_window = max(pending_counts) if force else max(1, decode_budget) empty_codes = torch.empty((effective_nq, 0), dtype=torch.long, device=resolved_device) codes_list: list[torch.Tensor] = [] for state in row_states: take_count = min(len(state["pending_decode_frames"]), decode_window) if take_count <= 0: codes_list.append(empty_codes) continue frame_rows = state["pending_decode_frames"][:take_count] del state["pending_decode_frames"][:take_count] frame_tensor = torch.cat(frame_rows, dim=0).to(device=resolved_device, dtype=torch.long) codes_list.append(frame_tensor[:, :effective_nq].transpose(0, 1).contiguous()) decode_output = audio_tokenizer.batch_decode( codes_list, num_quantizers=effective_nq, streaming=True, max_batch_size=(voice_clone_codec_batch_size if not codec_stream_started else None), reset_stream=not codec_stream_started, ) codec_stream_started = True waveform_rows, current_sample_rate = self._extract_batch_waveforms_and_sample_rate( decode_output, fallback_sample_rate=target_sample_rate, batch_size=len(batch_chunks), ) if decoded_sample_rate is None: decoded_sample_rate = current_sample_rate elif decoded_sample_rate != current_sample_rate: raise ValueError( f"Decoded sample rates differ across streaming decode calls: {decoded_sample_rate} vs {current_sample_rate}" ) for row_index, waveform_row in enumerate(waveform_rows): if waveform_row.numel() == 0 or int(waveform_row.shape[-1]) <= 0: continue row_states[row_index]["decoded_audio_segments"].append(waveform_row) yield from _emit_ready_segments() try: final_generation = None for event in self.generate_stream( input_ids=batched_prompt_input_ids, attention_mask=batched_attention_mask, nq=effective_nq, max_new_frames=max_new_frames, do_sample=do_sample, text_temperature=text_temperature, text_top_p=text_top_p, text_top_k=text_top_k, audio_temperature=audio_temperature, audio_top_p=audio_top_p, audio_top_k=audio_top_k, audio_repetition_penalty=audio_repetition_penalty, use_kv_cache=use_kv_cache, return_dict_in_generate=True, ): if event["type"] == "frame": frame_audio_token_ids = event["audio_token_ids"] active_mask = event["active_mask"] finished_mask = event["finished_mask"] for row_index in range(len(batch_chunks)): if bool(active_mask[row_index].item()): row_states[row_index]["pending_decode_frames"].append( frame_audio_token_ids[row_index : row_index + 1].detach().clone() ) if bool(finished_mask[row_index].item()): row_states[row_index]["generation_complete"] = True yield from _maybe_decode_pending(force=False) continue final_generation = event.get("generation") if final_generation is None: raise RuntimeError("Streaming generation finished without a final output.") for state in row_states: state["generation_complete"] = True yield from _maybe_decode_pending(force=True) yield from _emit_ready_segments() batched_audio_token_ids = self._mask_unused_audio_channels(final_generation.audio_token_ids, nq=effective_nq) for sample_index in range(len(batch_chunks)): generated_audio_token_chunks.append( self._trim_generated_audio_token_ids( batched_audio_token_ids[sample_index], effective_nq=effective_nq, ) ) finally: if codec_stream_started and callable(streaming_reset_fn): streaming_reset_fn() if generated_audio_token_chunks: audio_token_ids = torch.cat(generated_audio_token_chunks, dim=0) else: audio_token_ids = torch.empty((0, self.config.n_vq), dtype=torch.long, device=resolved_device) if emitted_waveform_segments: waveform = torch.cat(emitted_waveform_segments, dim=-1) else: waveform = torch.zeros((target_channels, 0), dtype=torch.float32) decoded_sample_rate = decoded_sample_rate or target_sample_rate output_path = Path(output_audio_path) self._save_audio(output_path, waveform, decoded_sample_rate) yield { "type": "result", "audio_path": str(output_path), "sample_rate": decoded_sample_rate, "audio_token_ids": audio_token_ids.detach().cpu(), "waveform": waveform, "reference_audio_token_ids": None if prompt_audio_codes is None else prompt_audio_codes.detach().cpu(), "voice_clone_text_chunks": voice_clone_text_chunks, "voice_clone_chunk_batch_size": int(voice_clone_chunk_batch_size), "voice_clone_codec_batch_size": int(voice_clone_codec_batch_size), } finally: if was_training: self.train() @torch.no_grad() def inference( self, text: str, output_audio_path: Union[str, Path], mode: str = "continuation", prompt_text: Optional[str] = None, prompt_audio_path: Optional[Union[str, Path]] = None, reference_audio_path: Optional[Union[str, Path]] = None, text_tokenizer=None, text_tokenizer_path: Optional[str] = None, audio_tokenizer=None, audio_tokenizer_type: Optional[str] = None, audio_tokenizer_pretrained_name_or_path: Optional[str] = None, device: Optional[Union[str, torch.device]] = None, nq: Optional[int] = None, max_new_frames: int = 300, do_sample: bool = False, text_temperature: float = 1.5, text_top_p: float = 1.0, text_top_k: int = 50, audio_temperature: float = 1.7, audio_top_p: float = 0.8, audio_top_k: int = 25, audio_repetition_penalty: float = 1.0, use_kv_cache: bool = True, voice_clone_max_text_tokens: int = DEFAULT_VOICE_CLONE_MAX_TEXT_TOKENS, voice_clone_max_memory_per_sample_gb: float = DEFAULT_VOICE_CLONE_MAX_MEMORY_PER_SAMPLE_GB, tts_max_batch_size: int = 0, codec_max_batch_size: int = 0, ) -> dict[str, Any]: resolved_device = self._resolve_device(device) effective_nq = self._resolve_inference_nq(nq) if next(self.parameters()).device != resolved_device: self.to(resolved_device) was_training = self.training self.eval() text_tokenizer = self._load_text_tokenizer( text_tokenizer=text_tokenizer, text_tokenizer_path=text_tokenizer_path, ) audio_tokenizer = self._load_audio_tokenizer( audio_tokenizer=audio_tokenizer, audio_tokenizer_type=audio_tokenizer_type, audio_tokenizer_pretrained_name_or_path=audio_tokenizer_pretrained_name_or_path, device=resolved_device, ) target_sample_rate = self._resolve_audio_tokenizer_sample_rate(audio_tokenizer) target_channels = self._resolve_audio_tokenizer_channels(audio_tokenizer) effective_prompt_audio_path = prompt_audio_path or reference_audio_path resolved_mode = self._resolve_inference_mode( mode=mode, has_prompt_text=prompt_text is not None, has_prompt_audio=effective_prompt_audio_path is not None, ) if reference_audio_path is not None and prompt_audio_path is None: logging.warning( "reference_audio_path=%s is treated as prompt_audio_path for backward compatibility.", reference_audio_path, ) prompt_audio_codes = None if effective_prompt_audio_path is not None: waveform, sample_rate = self._load_reference_audio( effective_prompt_audio_path, target_sample_rate, target_channels, ) encoded = self._call_audio_encode( audio_tokenizer=audio_tokenizer, waveform=waveform.to(resolved_device), sample_rate=sample_rate, ) prompt_audio_codes = self._mask_unused_audio_channels( self._normalize_audio_codes(encoded), nq=effective_nq, ).to(resolved_device) if resolved_mode == "voice_clone": split_voice_clone_text_chunks = self._split_text_into_best_sentences( text_tokenizer=text_tokenizer, text=text, max_tokens=voice_clone_max_text_tokens, ) voice_clone_text_chunks = split_voice_clone_text_chunks if len(split_voice_clone_text_chunks) > 1 else [text] else: voice_clone_text_chunks = [text] generated_audio_token_chunks: list[torch.LongTensor] = [] decoded_waveform_chunks: list[torch.FloatTensor] = [] decoded_sample_rate: Optional[int] = None if resolved_mode == "voice_clone" and len(voice_clone_text_chunks) > 1: voice_clone_chunk_batch_size, voice_clone_codec_batch_size = self._resolve_effective_voice_clone_batch_sizes( resolved_device=resolved_device, chunk_count=len(voice_clone_text_chunks), max_memory_per_sample_gb=float(voice_clone_max_memory_per_sample_gb), requested_tts_max_batch_size=tts_max_batch_size, requested_codec_max_batch_size=codec_max_batch_size, realtime_streaming=False, ) else: voice_clone_chunk_batch_size = 1 voice_clone_codec_batch_size = 1 for batch_start in range(0, len(voice_clone_text_chunks), voice_clone_chunk_batch_size): batch_chunks = voice_clone_text_chunks[batch_start : batch_start + voice_clone_chunk_batch_size] batch_prompt_input_ids: list[torch.LongTensor] = [] batch_attention_masks: list[torch.BoolTensor] = [] for text_chunk in batch_chunks: prompt_input_ids, attention_mask = self.build_inference_input_ids( text=text_chunk, text_tokenizer=text_tokenizer, mode=resolved_mode, prompt_text=prompt_text, prompt_audio_codes=prompt_audio_codes, device=resolved_device, ) batch_prompt_input_ids.append(prompt_input_ids) batch_attention_masks.append(attention_mask) batched_prompt_input_ids, batched_attention_mask = self._left_pad_inference_batch( input_id_batches=batch_prompt_input_ids, attention_mask_batches=batch_attention_masks, device=resolved_device, ) batched_audio_token_ids = self._generate_audio_token_ids_with_fallback( prompt_input_ids=batched_prompt_input_ids, attention_mask=batched_attention_mask, effective_nq=effective_nq, max_new_frames=max_new_frames, do_sample=do_sample, text_temperature=text_temperature, text_top_p=text_top_p, text_top_k=text_top_k, audio_temperature=audio_temperature, audio_top_p=audio_top_p, audio_top_k=audio_top_k, audio_repetition_penalty=audio_repetition_penalty, use_kv_cache=use_kv_cache, resolved_device=resolved_device, ) batch_audio_token_chunks: list[torch.LongTensor] = [] for sample_index in range(len(batch_chunks)): audio_token_ids = self._trim_generated_audio_token_ids( batched_audio_token_ids[sample_index], effective_nq=effective_nq, ) generated_audio_token_chunks.append(audio_token_ids) batch_audio_token_chunks.append(audio_token_ids) if resolved_mode == "voice_clone" and len(voice_clone_text_chunks) > 1: for codec_batch_start in range(0, len(batch_audio_token_chunks), voice_clone_codec_batch_size): codec_audio_token_batches = batch_audio_token_chunks[ codec_batch_start : codec_batch_start + voice_clone_codec_batch_size ] decoded_waveforms, current_sample_rate = self._decode_audio_token_id_batch_to_waveforms( audio_tokenizer=audio_tokenizer, audio_token_id_batches=codec_audio_token_batches, target_sample_rate=target_sample_rate, effective_nq=effective_nq, resolved_device=resolved_device, ) if decoded_sample_rate is None: decoded_sample_rate = current_sample_rate elif decoded_sample_rate != current_sample_rate: raise ValueError( f"Decoded sample rates differ across voice_clone chunks: {decoded_sample_rate} vs {current_sample_rate}" ) decoded_waveform_chunks.extend(decoded_waveforms) if generated_audio_token_chunks: audio_token_ids = torch.cat(generated_audio_token_chunks, dim=0) else: audio_token_ids = torch.empty((0, self.config.n_vq), dtype=torch.long, device=resolved_device) if resolved_mode == "voice_clone" and len(voice_clone_text_chunks) > 1: waveform = ( self._concat_voice_clone_waveform_chunks( waveform_chunks=decoded_waveform_chunks, text_chunks=voice_clone_text_chunks, sample_rate=decoded_sample_rate, ) if decoded_waveform_chunks else torch.zeros((target_channels, 0), dtype=torch.float32) ) else: decode_audio_token_ids = audio_token_ids prompt_waveform_prefix_samples = 0 if resolved_mode == "continuation" and prompt_audio_codes is not None: decode_audio_token_ids = torch.cat([prompt_audio_codes, audio_token_ids], dim=0) prompt_waveform_prefix_samples = ( int(prompt_audio_codes.shape[0]) * self._resolve_audio_tokenizer_downsample_rate(audio_tokenizer) ) waveform, decoded_sample_rate = self._decode_audio_token_ids_to_waveform( audio_tokenizer=audio_tokenizer, audio_token_ids=decode_audio_token_ids, target_sample_rate=target_sample_rate, effective_nq=effective_nq, resolved_device=resolved_device, ) if prompt_waveform_prefix_samples > 0: if decoded_sample_rate != target_sample_rate: prompt_waveform_prefix_samples = int( round(prompt_waveform_prefix_samples * float(decoded_sample_rate) / float(target_sample_rate)) ) prompt_waveform_prefix_samples = min(prompt_waveform_prefix_samples, int(waveform.shape[-1])) waveform = waveform[:, prompt_waveform_prefix_samples:] assert decoded_sample_rate is not None output_path = Path(output_audio_path) self._save_audio(output_path, waveform, decoded_sample_rate) if was_training: self.train() return { "audio_path": str(output_path), "sample_rate": decoded_sample_rate, "audio_token_ids": audio_token_ids.detach().cpu(), "waveform": waveform, "reference_audio_token_ids": None if prompt_audio_codes is None else prompt_audio_codes.detach().cpu(), "voice_clone_text_chunks": voice_clone_text_chunks, "voice_clone_chunk_batch_size": int(voice_clone_chunk_batch_size), "voice_clone_codec_batch_size": int(voice_clone_codec_batch_size), }