tts_utils.py

from __future__ import annotations

import base64
import tempfile
import time
from dataclasses import dataclass
from pathlib import Path
from queue import Queue
from threading import Event, Thread
from typing import Any, Iterator

import numpy as np
import soundfile as sf
import torch
from transformers import AutoModel
from transformers.generation.logits_process import LogitsProcessor
from transformers.generation.stopping_criteria import StoppingCriteria
from transformers.generation.streamers import BaseStreamer


SAMPLE_RATE = 16000
TTS_PREFIX = "<|text to speech|> Generate speech for this transcription. "
DEFAULT_SYSTEM_PROMPT = (
    "You are a helpful and harmless assistant.\n\n"
    "You are not allowed to use any tools."
)


@dataclass(frozen=True)
class SpeechTokenMap:
    start: int
    end: int
    codec_start: int
    codec_end: int
    eos: int


@dataclass(frozen=True)
class TTSStreamEvent:
    pcm: np.ndarray | None
    token_count: int
    elapsed_seconds: float
    done: bool = False
    truncated: bool = False


class SpeechTokenLogitsProcessor(LogitsProcessor):
    def __init__(self, token_map: SpeechTokenMap):
        self.token_map = token_map

    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        eos_scores = scores[:, self.token_map.eos].clone()
        scores[:, : self.token_map.end] = -float("inf")
        scores[:, self.token_map.codec_end + 1 :] = -float("inf")
        scores[:, self.token_map.eos] = eos_scores
        return scores


class EventStoppingCriteria(StoppingCriteria):
    def __init__(self, event: Event):
        self.event = event

    def __call__(
        self,
        input_ids: torch.LongTensor,
        scores: torch.FloatTensor,
        **kwargs: Any,
    ) -> torch.BoolTensor:
        return torch.full(
            (input_ids.shape[0],),
            self.event.is_set(),
            dtype=torch.bool,
            device=input_ids.device,
        )


class TokenIdStreamer(BaseStreamer):
    _END = object()

    def __init__(self) -> None:
        self.queue: Queue[int | BaseException | object] = Queue()
        self.is_prompt = True

    def put(self, value: torch.Tensor) -> None:
        if self.is_prompt:
            self.is_prompt = False
            return
        token_ids = value.reshape(-1).tolist()
        if len(token_ids) != 1:
            raise ValueError(f"TTS streaming requires batch size 1, got {len(token_ids)} tokens")
        self.queue.put(int(token_ids[0]))

    def end(self) -> None:
        self.queue.put(self._END)

    def fail(self, error: BaseException) -> None:
        self.queue.put(error)
        self.end()

    def __iter__(self) -> Iterator[int]:
        while True:
            item = self.queue.get()
            if item is self._END:
                return
            if isinstance(item, BaseException):
                raise item
            yield int(item)


def build_tts_prompt(text: str, tokenizer: Any) -> str:
    messages = [
        {"role": "system", "content": DEFAULT_SYSTEM_PROMPT},
        {"role": "user", "content": f"{TTS_PREFIX}{text}"},
    ]
    return (
        tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
            enable_thinking=False,
        )
        + "<speechgen_start>"
    )


def build_tts_null_prompt(cond_prompt: str, tokenizer: Any, max_iters: int = 64) -> str:
    target_len = len(tokenizer.encode(cond_prompt))

    def template(null_text: str) -> str:
        return build_tts_prompt(null_text, tokenizer)

    base_len = len(tokenizer.encode(template("")))
    n_unk = max(1, target_len - base_len)
    for _ in range(max_iters):
        prompt = template("<unk>" * n_unk)
        current_len = len(tokenizer.encode(prompt))
        if current_len == target_len:
            return prompt
        n_unk += 1 if current_len < target_len else -1
        n_unk = max(1, n_unk)
    raise ValueError("Unable to construct an equal-length TTS CFG null prompt")


def build_speech_token_map(tokenizer: Any) -> SpeechTokenMap:
    token_map = SpeechTokenMap(
        start=tokenizer.convert_tokens_to_ids("<speechgen_start>"),
        end=tokenizer.convert_tokens_to_ids("<speechgen_end>"),
        codec_start=tokenizer.convert_tokens_to_ids("<speechcodec_0>"),
        codec_end=tokenizer.convert_tokens_to_ids("<speechcodec_65535>"),
        eos=int(tokenizer.eos_token_id),
    )
    expected = (131075, 131076, 131077, 196612)
    actual = (token_map.start, token_map.end, token_map.codec_start, token_map.codec_end)
    if actual != expected:
        raise ValueError(f"Unexpected Audex speech token layout: expected={expected}, actual={actual}")
    return token_map


def load_speech_decoder(model_path: str, device: str = "cuda") -> Any:
    path = Path(model_path)
    for filename in ("config.json", "model.safetensors"):
        if not (path / filename).is_file():
            raise FileNotFoundError(f"Speech decoder file not found: {path / filename}")

    decoder, loading_info = AutoModel.from_pretrained(
        str(path),
        trust_remote_code=True,
        output_loading_info=True,
    )
    invalid_keys = {
        key: loading_info[key]
        for key in ("missing_keys", "unexpected_keys", "mismatched_keys")
        if loading_info[key]
    }
    if invalid_keys:
        raise RuntimeError(f"Speech decoder checkpoint loading was incomplete: {invalid_keys}")

    for module in decoder.modules():
        if {"theta", "cache"} <= getattr(module, "_non_persistent_buffers_set", set()):
            module.rope_init(device=module.theta.device)
            module._rope_ready = False

    sample_rate = int(getattr(decoder.config, "sample_rate", SAMPLE_RATE))
    if sample_rate != SAMPLE_RATE:
        raise ValueError(f"Expected a {SAMPLE_RATE} Hz speech decoder, got {sample_rate} Hz")
    decoder = decoder.to(device).eval()
    for parameter in decoder.parameters():
        parameter.requires_grad = False
    return decoder


def stream_tts(
    model: Any,
    tokenizer: Any,
    decoder: Any,
    text: str,
    max_new_tokens: int,
    *,
    temperature: float = 0.8,
    top_p: float = 1.0,
    top_k: int = 0,
    guidance_scale: float = 2.0,
    seed: int = 0,
    chunk_frames: int = 5,
) -> Iterator[TTSStreamEvent]:
    text = text.strip()
    if not text:
        raise ValueError("Text to synthesize must not be empty")

    token_map = build_speech_token_map(tokenizer)
    cond_prompt = build_tts_prompt(text, tokenizer)
    uncond_prompt = build_tts_null_prompt(cond_prompt, tokenizer)
    cond_ids = tokenizer.encode(cond_prompt, return_tensors="pt").to(model.device)
    uncond_ids = tokenizer.encode(uncond_prompt, return_tensors="pt").to(model.device)
    if cond_ids.shape != uncond_ids.shape:
        raise ValueError(
            f"TTS CFG prompt lengths differ: conditional={cond_ids.shape[-1]}, "
            f"unconditional={uncond_ids.shape[-1]}"
        )

    streamer = TokenIdStreamer()
    cancel_event = Event()
    session = decoder.create_session(
        chunk_frames=chunk_frames,
        sample_rate=SAMPLE_RATE,
        return_numpy=True,
    )
    generation_error: list[BaseException] = []
    started_at = time.perf_counter()

    def generate() -> None:
        try:
            torch.manual_seed(seed)
            torch.cuda.manual_seed_all(seed)
            model.generate(
                input_ids=cond_ids,
                attention_mask=torch.ones_like(cond_ids),
                negative_prompt_ids=uncond_ids,
                negative_prompt_attention_mask=torch.ones_like(uncond_ids),
                guidance_scale=guidance_scale,
                do_sample=True,
                temperature=temperature,
                top_p=top_p,
                **({"top_k": top_k} if top_k > 0 else {}),
                max_new_tokens=max_new_tokens,
                eos_token_id=[token_map.end, token_map.eos],
                pad_token_id=tokenizer.pad_token_id or token_map.eos,
                logits_processor=[SpeechTokenLogitsProcessor(token_map)],
                stopping_criteria=[EventStoppingCriteria(cancel_event)],
                streamer=streamer,
                use_cache=True,
            )
        except BaseException as error:
            generation_error.append(error)
            streamer.fail(error)

    thread = Thread(target=generate, daemon=True)
    thread.start()
    token_count = 0
    terminal_token_seen = False
    try:
        for token_id in streamer:
            if token_id == token_map.end or token_id == token_map.eos:
                terminal_token_seen = True
                break
            if not token_map.codec_start <= token_id <= token_map.codec_end:
                raise RuntimeError(f"Unexpected token in TTS output: {token_id}")

            token_count += 1
            codec_index = token_id - token_map.codec_start
            if token_count == 1:
                yield TTSStreamEvent(
                    pcm=None,
                    token_count=token_count,
                    elapsed_seconds=time.perf_counter() - started_at,
                )
            for _, pcm in session.push([[codec_index]]):
                yield TTSStreamEvent(
                    pcm=np.asarray(pcm, dtype=np.float32),
                    token_count=token_count,
                    elapsed_seconds=time.perf_counter() - started_at,
                )

        if generation_error:
            raise generation_error[0]
        for _, pcm in session.flush():
            yield TTSStreamEvent(
                pcm=np.asarray(pcm, dtype=np.float32),
                token_count=token_count,
                elapsed_seconds=time.perf_counter() - started_at,
            )
        yield TTSStreamEvent(
            pcm=None,
            token_count=token_count,
            elapsed_seconds=time.perf_counter() - started_at,
            done=True,
            truncated=not terminal_token_seen,
        )
    finally:
        cancel_event.set()
        thread.join()


def encode_pcm_chunk(pcm: np.ndarray) -> str:
    samples = np.asarray(pcm, dtype="<f4")
    return base64.b64encode(samples.tobytes()).decode("ascii")


def write_wav(pcm: np.ndarray, sample_rate: int = SAMPLE_RATE) -> str:
    output = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
    output.close()
    sf.write(output.name, np.asarray(pcm, dtype=np.float32), sample_rate, subtype="PCM_16")
    return output.name