# --------------------------------------------------------------------- # Copyright (c) 2026 Qualcomm Technologies, Inc. and/or its subsidiaries. # SPDX-License-Identifier: BSD-3-Clause # --------------------------------------------------------------------- from __future__ import annotations from pathlib import Path from typing import Any import numpy as np import soundfile as sf import soxr import torch from torch import Tensor from .external_repos.neutts.neutts.neutts import NeuTTS from .model import ( BACKBONE_REPO, CODEC_REPO, CONTEXT_LENGTH, NUM_LAYERS, PREFILL_SEQ_LEN, SAMPLE_RATE, NeuTTSNano, build_attention_mask, empty_kv_cache, load_codec, ) MAX_NEW_TOKENS = 1024 MIN_NEW_TOKENS = 50 SPEECH_END_TOKEN = "<|SPEECH_GENERATION_END|>" # NeuCodec's encoder expects 16kHz mono, regardless of the reference clip's rate. CODEC_INPUT_SAMPLE_RATE = 16_000 def _right_align(history: list[Tensor], width: int) -> list[Tensor]: """Place each cache tensor at the right edge of a fixed-width zeroed buffer.""" buffers = empty_kv_cache(width) for layer in range(NUM_LAYERS): k, v = history[2 * layer], history[2 * layer + 1] n = min(k.shape[-1], width) if n: buffers[2 * layer][..., -n:] = k[..., -n:] buffers[2 * layer + 1][:, :, -n:, :] = v[:, :, -n:, :] return buffers def _append(history: list[Tensor], new: list[Tensor], skip: int) -> list[Tensor]: """Append newly computed cache entries, dropping the first ``skip`` positions.""" out: list[Tensor] = [] for layer in range(NUM_LAYERS): k_new = new[2 * layer][..., skip:] v_new = new[2 * layer + 1][:, :, skip:, :] out.append(torch.cat([history[2 * layer], k_new], dim=-1)) out.append(torch.cat([history[2 * layer + 1], v_new], dim=2)) return out def _sample(logits: Tensor, temperature: float, top_k: int) -> int: if temperature <= 0: return int(logits.argmax()) values, indices = torch.topk(logits / temperature, min(top_k, logits.shape[-1])) probs = torch.softmax(values, dim=-1) return int(indices[torch.multinomial(probs, num_samples=1)]) def generate_speech_tokens( model: NeuTTSNano, prompt_ids: list[int], temperature: float = 1.0, top_k: int = 50, max_new_tokens: int = MAX_NEW_TOKENS, ) -> list[int]: """Autoregressive decode at the prefill and decode graph shapes. On device these are two graphs of one linked binary; in torch they are the same module called at different sequence lengths. The cache is kept as a growing per-layer history and copied into a right-aligned fixed-width buffer for each call, matching the layout the graphs expect. """ backbone = model.backbone eos_id = backbone.tokenizer.convert_tokens_to_ids(SPEECH_END_TOKEN) history = empty_kv_cache(0) num_cached = 0 last_logits: Tensor | None = None # Prefill: left-pad to a whole number of chunks so real tokens stay # right-aligned. Pad rows are masked out and their cache entries dropped. pad = -len(prompt_ids) % PREFILL_SEQ_LEN padded = [0] * pad + list(prompt_ids) for start in range(0, len(padded), PREFILL_SEQ_LEN): skip = pad if start == 0 else 0 num_real = PREFILL_SEQ_LEN - skip position_ids = torch.tensor( [[0] * skip + list(range(num_cached, num_cached + num_real))] ) cos, sin = backbone.embedding.get_embedding(position_ids) with torch.no_grad(): out = backbone( torch.tensor( [padded[start : start + PREFILL_SEQ_LEN]], dtype=torch.int32 ), build_attention_mask(PREFILL_SEQ_LEN, num_cached + num_real), cos, sin, *_right_align(history, CONTEXT_LENGTH - PREFILL_SEQ_LEN), ) last_logits = out[0][0, -1] history = _append(history, out[1:], skip) num_cached += num_real assert last_logits is not None generated: list[int] = [] while len(generated) < max_new_tokens and num_cached < CONTEXT_LENGTH - 1: token = _sample(last_logits, temperature, top_k) if token == eos_id and len(generated) >= MIN_NEW_TOKENS: break generated.append(token) position_ids = torch.tensor([[num_cached]]) cos, sin = backbone.embedding.get_embedding(position_ids) with torch.no_grad(): out = backbone( torch.tensor([[token]], dtype=torch.int32), build_attention_mask(1, num_cached + 1), cos, sin, *_right_align(history, CONTEXT_LENGTH - 1), ) last_logits = out[0][0, -1] history = _append(history, out[1:], 0) num_cached += 1 return generated class NeuTTSApp: """End-to-end NeuTTS synthesis app. Generation runs at the same two fixed graph shapes that are exported to device -- a prefill pass over ``PREFILL_SEQ_LEN`` tokens and a single-token decode pass -- with the KV cache passed in and out explicitly. Everything else (phonemization, prompt assembly, sampling, codec decode) is CPU work borrowed from the upstream pipeline. Inputs ------ input_text: str The text to synthesize. ref_audio_path: str | Path Path to a 16-44kHz mono ``.wav`` clip of the target speaker (3-15s). ref_text: str Transcript of the reference audio. Output ------ ``np.ndarray`` of float32 audio samples at 24kHz. """ def __init__(self, model: NeuTTSNano, seed: int | None = None) -> None: # Upstream's infer() calls torch.manual_seed(self._call_seed()) itself, so # seeding from outside has no effect; the seed has to go in here. self.model = model self.pipeline = _GraphPipeline(model, seed=seed) def predict(self, *args: Any, **kwargs: Any) -> np.ndarray: return self.synthesize(*args, **kwargs) __call__ = predict def encode_reference(self, ref_audio_path: str | Path) -> Tensor: # soundfile + soxr rather than librosa, which pulls in numba for what is # one load and resample. soxr at HQ is librosa.load's own default, so the # reference conditioning is unchanged. wav, sample_rate = sf.read(str(ref_audio_path), dtype="float32", always_2d=True) mono = wav.mean(axis=1) if sample_rate != CODEC_INPUT_SAMPLE_RATE: mono = soxr.resample( mono, sample_rate, CODEC_INPUT_SAMPLE_RATE, quality="HQ" ) wav_tensor = torch.from_numpy(np.ascontiguousarray(mono, dtype=np.float32)) codec = load_codec() with torch.no_grad(): return ( codec.encode_code(audio_or_path=wav_tensor[None, None, :]) # type: ignore[operator] .squeeze(0) .squeeze(0) ) def generate_speech_tokens( self, prompt_ids: list[int], temperature: float = 1.0, top_k: int = 50, max_new_tokens: int = MAX_NEW_TOKENS, ) -> list[int]: return generate_speech_tokens( self.model, prompt_ids, temperature, top_k, max_new_tokens ) def synthesize( self, input_text: str, ref_audio_path: str | Path, ref_text: str, temperature: float = 1.0, top_k: int = 50, ) -> np.ndarray: ref_codes = self.encode_reference(ref_audio_path) return self.pipeline.infer( input_text, ref_codes, ref_text, temperature=temperature, top_k=top_k, ) @property def sample_rate(self) -> int: return SAMPLE_RATE class _GraphPipeline(NeuTTS): """Upstream pipeline with weight loading and generation redirected. Reuses upstream phonemization and prompt assembly, but binds the backbone and codec to the recipe's already-loaded modules (avoiding a second copy of each) and replaces ``generate()`` with the exported-graph decode loop. """ def __init__(self, model: NeuTTSNano, seed: int | None = None) -> None: self._qaihm_model = model super().__init__( backbone_repo=BACKBONE_REPO, codec_repo=CODEC_REPO, backbone_device="cpu", codec_device="cpu", seed=seed, ) def _load_backbone(self, backbone_repo: str, backbone_device: str) -> None: backbone = self._qaihm_model.backbone self.tokenizer = backbone.tokenizer self.backbone = backbone.model config = getattr(backbone.llm_config, "neuphonic", None) or {} self.input_format = config.get("input_format", "phonemes") self._supported_emotions = config.get("supported_emotions") def _load_codec(self, codec_repo: str, codec_device: str) -> None: self.codec = load_codec() def _infer_torch( self, prompt_ids: list[int], temperature: float = 1.0, top_k: int = 50 ) -> str: tokens = generate_speech_tokens( self._qaihm_model, prompt_ids, temperature, top_k ) tokenizer = self._qaihm_model.backbone.tokenizer return str(tokenizer.decode(tokens, add_special_tokens=False))