gepard_inference/engine.py

"""High-level TTS engine: model + codec + speakers behind one facade.

Designed for a Hugging Face ZeroGPU Space: everything heavy (checkpoint
download, weight loading, codec init) happens ONCE at process startup in
:meth:`GepardEngine.load`; per-request work is only the autoregressive
generation itself. On ZeroGPU the ``.to("cuda")`` calls issued at startup are
intercepted by the ``spaces`` package and replayed when a GPU is attached, so
the model is never reloaded between requests.
"""

from __future__ import annotations

import os
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, List, NamedTuple, Optional, Tuple

import numpy as np
import torch
import yaml

from .runner import GepardRunner
from .speakers import SpeakerLibrary


class Example(NamedTuple):
    """A demo row for the UI examples table.

    Attributes:
        speaker: Preset speaker key (must exist under ``speakers:``); used for
            the codec lookup and to sync the dropdown.
        label: Human-readable text shown in the examples table (e.g.
            "Den from California"). Defaults to ``speaker`` when not given.
        text: Phrase to synthesize.
    """

    speaker: str
    label: str
    text: str


@dataclass
class GenerationParams:
    """User-tunable generation arguments (mirrors ``GepardRunner.generate``).

    ``cfg_frames == 0`` means "guide the whole utterance" and is translated to
    ``None`` for the runner.
    """

    temperature: float = 0.3
    top_k: int = 0
    cfg_scale: float = 3.0
    cfg_frames: int = 0
    stop_threshold: float = 0.5
    max_frames: int = 2000
    repetition_penalty: float = 1.0
    repetition_window: int = 32

    def to_generate_kwargs(self) -> dict:
        """Translate UI-facing values into ``GepardRunner.generate`` kwargs."""
        return {
            "temperature": float(self.temperature),
            "top_k": int(self.top_k),
            "cfg_scale": float(self.cfg_scale),
            "cfg_frames": int(self.cfg_frames) if int(self.cfg_frames) > 0 else None,
            "stop_threshold": float(self.stop_threshold),
            "max_frames": int(self.max_frames),
            "repetition_penalty": float(self.repetition_penalty),
            "repetition_window": int(self.repetition_window),
        }


@dataclass
class AppConfig:
    """Application configuration loaded from ``config.yaml``.

    Attributes:
        checkpoint: HF repo id (or local dir) of the self-describing checkpoint.
        attn_implementation: Backbone attention implementation for inference.
        codec_id: HF repo id of the NeMo audio codec.
        sample_rate: Codec sample rate (Hz).
        speakers: ``{speaker_key: path_to_pt}`` preset voices.
        speaker_labels: ``{speaker_key: display_label}`` human-readable names
            for the dropdown / examples (defaults to the key).
        examples: ``[Example(speaker, label, text), ...]`` demo rows for the UI
            (each speaker must be a key in ``speakers``).
        defaults: Default generation parameters (slider initial values).
        max_ref_seconds: Recorded/uploaded reference audio is truncated to
            this many seconds before encoding.
        gpu_duration: ZeroGPU time budget per request (seconds).
        cfg_max_text_tokens: Disable text-CFG when the input has more than this
            many text tokens (CFG helps short phrases but rushes/flattens long
            ones). None = never gate.
        root: Directory of the config file; relative paths resolve against it.
    """

    checkpoint: str
    attn_implementation: str = "eager"
    codec_id: str = "nvidia/nemo-nano-codec-22khz-1.89kbps-21.5fps"
    sample_rate: int = 22050
    speakers: Dict[str, str] = field(default_factory=dict)
    speaker_labels: Dict[str, str] = field(default_factory=dict)
    examples: List[Example] = field(default_factory=list)
    defaults: GenerationParams = field(default_factory=GenerationParams)
    max_ref_seconds: float = 60.0
    gpu_duration: int = 120
    cfg_max_text_tokens: Optional[int] = None
    root: Path = field(default_factory=Path.cwd)

    @classmethod
    def from_yaml(cls, path: str | Path) -> "AppConfig":
        """Load and validate the application config from a YAML file."""
        path = Path(path)
        with open(path) as f:
            raw = yaml.safe_load(f) or {}
        model = raw.get("model") or {}
        codec = raw.get("codec") or {}
        app = raw.get("app") or {}
        generation = raw.get("generation") or {}
        if not model.get("checkpoint"):
            raise ValueError(f"{path}: model.checkpoint is required")
        cfg_max = generation.get("cfg_max_text_tokens")
        speakers, speaker_labels = cls._parse_speakers(raw.get("speakers"), path)
        examples = cls._parse_examples(raw.get("examples"), speakers, speaker_labels, path)
        return cls(
            checkpoint=str(model["checkpoint"]),
            attn_implementation=str(model.get("attn_implementation", "eager")),
            codec_id=str(codec.get("id", cls.codec_id)),
            sample_rate=int(codec.get("sample_rate", 22050)),
            speakers=speakers,
            speaker_labels=speaker_labels,
            examples=examples,
            defaults=GenerationParams(**(raw.get("defaults") or {})),
            max_ref_seconds=float(app.get("max_ref_seconds", 60.0)),
            gpu_duration=int(app.get("gpu_duration", 120)),
            cfg_max_text_tokens=int(cfg_max) if cfg_max is not None else None,
            root=path.parent.resolve(),
        )

    @staticmethod
    def _parse_speakers(raw_speakers, path) -> Tuple[Dict[str, str], Dict[str, str]]:
        """Normalize the ``speakers:`` section into ``(paths, labels)``.

        Each value may be either a path string (label defaults to the key) or a
        mapping with ``path`` and an optional human-readable ``label``:

            speakers:
              Den:                       # mapping form, custom label
                path: speakers/den.pt
                label: Den from California
              Nurisa: speakers/nurisa.pt  # string form, label = "Nurisa"
        """
        paths: Dict[str, str] = {}
        labels: Dict[str, str] = {}
        for key, val in (raw_speakers or {}).items():
            if isinstance(val, dict):
                p = val.get("path")
                label = val.get("label") or key
            else:
                p = val
                label = key
            if not p:
                raise ValueError(f"{path}: speaker {key!r} has no path")
            paths[str(key)] = str(p)
            labels[str(key)] = str(label)
        return paths, labels

    @staticmethod
    def _parse_examples(
        raw_examples, speakers: Dict[str, str], speaker_labels: Dict[str, str], path
    ) -> List[Example]:
        """Validate the ``examples:`` section into ``[Example, ...]``.

        Each entry must be a mapping with ``speaker`` (a known preset) and
        ``text``; an optional ``label`` overrides the string shown in the UI
        (defaults to the speaker's label). Malformed or unknown-speaker rows are
        skipped with a warning rather than failing the Space — examples are
        cosmetic.
        """
        examples: List[Example] = []
        for i, item in enumerate(raw_examples or []):
            if not isinstance(item, dict):
                print(f"[AppConfig] {path}: examples[{i}] is not a mapping; skipped")
                continue
            speaker = item.get("speaker")
            text = item.get("text")
            if not speaker or not text:
                print(f"[AppConfig] {path}: examples[{i}] missing speaker/text; skipped")
                continue
            if speaker not in speakers:
                print(f"[AppConfig] {path}: examples[{i}] unknown speaker {speaker!r}; skipped")
                continue
            label = item.get("label") or speaker_labels.get(speaker, speaker)
            examples.append(Example(str(speaker), str(label), str(text)))
        return examples


def runtime_device() -> torch.device:
    """The device inference will run on.

    On a ZeroGPU Space CUDA is not initialized in the main process, but the
    ``spaces`` package intercepts ``.to("cuda")`` calls — so "cuda" is the
    right answer whenever the Space is ZeroGPU-backed OR a real GPU is
    visible. Falls back to CPU for local GPU-less runs.
    """
    if os.environ.get("SPACES_ZERO_GPU"):
        return torch.device("cuda")
    return torch.device("cuda" if torch.cuda.is_available() else "cpu")


class GepardEngine:
    """Facade over the Gepard runner, the NeMo codec and the speaker library.

    Lifecycle:
        engine = GepardEngine(AppConfig.from_yaml("config.yaml"))
        engine.load()                       # once, at process startup
        ...
        sr, wave = engine.synthesize(text, ref_codes, params)   # per request
    """

    def __init__(self, config: AppConfig):
        self.config = config
        self.device = runtime_device()
        self.runner: Optional[GepardRunner] = None
        self.codec = None
        self.speakers = SpeakerLibrary(config.root, config.speakers)

    # ------------------------------------------------------------------
    # Startup
    # ------------------------------------------------------------------

    def load(self) -> "GepardEngine":
        """Load the model + codec once and move them to the runtime device."""
        from .codec_wrapper import UnfoldedCodecModel  # NeMo import is heavy

        self.runner = GepardRunner.from_checkpoint(
            self.config.checkpoint,
            device="cpu",  # weights land on CPU; moved below (ZeroGPU-friendly)
            attn_implementation=self.config.attn_implementation,
        )
        self.codec = UnfoldedCodecModel.from_pretrained(self.config.codec_id).eval()

        self.runner.model.to(self.device)
        self.codec.to(self.device)
        self.runner.device = self.device

        self.speakers.preload()
        print(f"[GepardEngine] ready on {self.device} | speakers: {self.speakers.names}")
        return self

    # ------------------------------------------------------------------
    # Per-request work (runs inside the GPU context on ZeroGPU)
    # ------------------------------------------------------------------

    @property
    def fsq_levels(self) -> list:
        """FSQ levels per codebook dimension, read from the loaded model."""
        return list(self.runner.model.ref_compressor.fsq_levels)

    def encode_reference(self, audio_path: str) -> torch.Tensor:
        """Encode a reference clip into unfolded codec codes ``[1, T, C_total]``.

        The clip is loaded mono at the codec sample rate and truncated to
        ``config.max_ref_seconds``.
        """
        import librosa

        wave_np, _ = librosa.load(audio_path, sr=self.config.sample_rate, mono=True)
        max_samples = int(self.config.max_ref_seconds * self.config.sample_rate)
        wave_np = wave_np[:max_samples]
        if wave_np.size == 0:
            raise ValueError("reference audio is empty")

        from .codec_ops import unfold_tokens

        wave = torch.from_numpy(wave_np).unsqueeze(0).to(self.device)
        wave_len = torch.tensor([wave.shape[-1]], device=self.device)
        with torch.inference_mode():
            tokens, _ = self.codec.encode(audio=wave, audio_len=wave_len)  # [1, C, T]
        ref_codes = (
            unfold_tokens(tokens.cpu(), self.fsq_levels).permute(0, 2, 1).contiguous()
        )  # [1, T, C_total]
        return ref_codes

    def synthesize(
        self,
        text: str,
        ref_codes: Optional[torch.Tensor],
        params: GenerationParams,
    ) -> Tuple[int, np.ndarray]:
        """Generate speech for ``text`` conditioned on ``ref_codes``.

        Returns ``(sample_rate, waveform)`` with the waveform as float32
        numpy — the format ``gr.Audio`` consumes directly.
        """
        if self.runner is None:
            raise RuntimeError("GepardEngine.load() must be called before synthesize()")
        text = (text or "").strip()
        if not text:
            raise ValueError("text is empty")

        tokens = self.runner.generate(
            text,
            ref_codes=ref_codes.to(self.device) if ref_codes is not None else None,
            cfg_max_text_tokens=self.config.cfg_max_text_tokens,
            **params.to_generate_kwargs(),
        )  # (num_heads, T)

        codes = tokens.unsqueeze(0).to(self.device)
        codes_len = torch.tensor([codes.shape[-1]], device=self.device)
        with torch.inference_mode():
            audio, _ = self.codec.decode_from_codes(codes, codes_len)
        wave = audio.float().cpu().detach().flatten().numpy()
        return self.config.sample_rate, wave