app.py

# ---------------------------------------------------------------------
# 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))