piper_train/speaker_encoder/encoder.py

"""High-level Speaker Encoder API.

Provides a unified interface for extracting speaker embeddings from audio
files using either a PyTorch checkpoint or an ONNX model.

Usage (ONNX -- recommended for inference):
    encoder = SpeakerEncoder.from_onnx("speaker_encoder.onnx")
    emb = encoder.encode("audio.wav")

Usage (PyTorch):
    encoder = SpeakerEncoder.from_pytorch("speaker_encoder.ckpt")
    emb = encoder.encode("audio.wav")
"""

from __future__ import annotations

import logging
from pathlib import Path
from typing import TYPE_CHECKING

import numpy as np

from .audio_utils import (
    DEFAULT_FMAX,
    DEFAULT_FMIN,
    DEFAULT_HOP_LENGTH,
    DEFAULT_N_FFT,
    DEFAULT_N_MELS,
    DEFAULT_SR,
    compute_mel_spectrogram,
    load_audio,
    normalize_audio,
)


if TYPE_CHECKING:
    import onnxruntime
    import torch

_LOGGER = logging.getLogger(__name__)


def _infer_hparams(state_dict: dict) -> dict:
    """Infer ECAPA-TDNN hyperparameters from a state_dict.

    Examines key tensor shapes to determine input_dim, channels, emb_dim,
    and scale so that the model can be reconstructed without explicit
    configuration.

    Args:
        state_dict: Model state dictionary.

    Returns:
        Dict of keyword arguments for :class:`ECAPATDNN`.
    """
    # layer1.0.weight has shape (channels, input_dim, kernel_size)
    layer1_weight = state_dict.get("layer1.0.weight")
    if layer1_weight is None:
        raise ValueError(
            "Cannot infer hparams: 'layer1.0.weight' not found in state_dict"
        )

    channels = layer1_weight.shape[0]
    input_dim = layer1_weight.shape[1]

    # fc.weight has shape (emb_dim, channels * 2)
    fc_weight = state_dict.get("fc.weight")
    if fc_weight is None:
        raise ValueError("Cannot infer hparams: 'fc.weight' not found in state_dict")
    emb_dim = fc_weight.shape[0]

    # Infer Res2Net scale from the number of conv modules in layer2.res2net.convs
    # convs has (scale - 1) entries: convs.0, convs.1, ..., convs.(scale-2)
    scale_minus_1 = 0
    for key in state_dict:
        if key.startswith("layer2.res2net.convs.") and key.endswith(".weight"):
            scale_minus_1 += 1
    scale = scale_minus_1 + 1 if scale_minus_1 > 0 else 8

    # Infer SE bottleneck from layer2.se.se.1.weight shape (bottleneck, channels, 1)
    se_weight = state_dict.get("layer2.se.se.1.weight")
    se_bottleneck = se_weight.shape[0] if se_weight is not None else 128

    return {
        "input_dim": input_dim,
        "channels": channels,
        "emb_dim": emb_dim,
        "scale": scale,
        "se_bottleneck": se_bottleneck,
    }


class SpeakerEncoder:
    """Speaker Encoder high-level API.

    Loads a PyTorch or ONNX speaker encoder model and provides methods
    to extract 256-dimensional speaker embeddings from audio files.

    Do not instantiate directly; use :meth:`from_pytorch` or :meth:`from_onnx`.
    """

    def __init__(self) -> None:
        self._mode: str = "none"  # "pytorch" or "onnx"
        self._pytorch_model: torch.nn.Module | None = None
        self._pytorch_device: str = "cpu"
        self._onnx_session: onnxruntime.InferenceSession | None = None

    @classmethod
    def from_pytorch(
        cls,
        checkpoint_path: str | Path,
        device: str = "cpu",
    ) -> SpeakerEncoder:
        """Load a speaker encoder from a PyTorch checkpoint.

        The checkpoint should contain either:
          - A raw state_dict (keys like ``layer1.0.weight``), or
          - A dict with a ``"model_state_dict"`` key.

        Args:
            checkpoint_path: Path to the ``.ckpt`` or ``.pt`` file.
            device: Torch device string (default: ``"cpu"``).

        Returns:
            Configured :class:`SpeakerEncoder` instance.
        """
        import torch  # noqa: PLC0415

        from .ecapa_tdnn import ECAPATDNN  # noqa: PLC0415

        checkpoint_path = Path(checkpoint_path)
        if not checkpoint_path.exists():
            raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")

        ckpt = torch.load(str(checkpoint_path), map_location=device, weights_only=True)

        if isinstance(ckpt, dict) and "model_state_dict" in ckpt:
            state_dict = ckpt["model_state_dict"]
        elif isinstance(ckpt, dict) and all(isinstance(k, str) for k in ckpt.keys()):
            state_dict = ckpt
        else:
            raise ValueError(
                "Checkpoint format not recognised. Expected a state_dict or a "
                "dict with 'model_state_dict' key."
            )

        # Infer model hyperparameters from the state_dict shapes
        hparams = _infer_hparams(state_dict)
        model = ECAPATDNN(**hparams)
        model.load_state_dict(state_dict)
        model.eval()
        model.to(device)

        encoder = cls()
        encoder._mode = "pytorch"
        encoder._pytorch_model = model
        encoder._pytorch_device = device
        _LOGGER.info(
            "Loaded PyTorch speaker encoder from %s (device=%s)",
            checkpoint_path,
            device,
        )
        return encoder

    @classmethod
    def from_onnx(cls, onnx_path: str | Path) -> SpeakerEncoder:
        """Load a speaker encoder from an ONNX model.

        Uses the project's shared ORT session utilities for optimised
        session creation with caching support.

        Args:
            onnx_path: Path to the ``.onnx`` file.

        Returns:
            Configured :class:`SpeakerEncoder` instance.
        """
        onnx_path = Path(onnx_path)
        if not onnx_path.exists():
            raise FileNotFoundError(f"ONNX model not found: {onnx_path}")

        from ..ort_utils import create_session_with_cache  # noqa: PLC0415

        session = create_session_with_cache(onnx_path, device="cpu")

        encoder = cls()
        encoder._mode = "onnx"
        encoder._onnx_session = session
        _LOGGER.info("Loaded ONNX speaker encoder from %s", onnx_path)
        return encoder

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def encode(self, audio_path: str | Path) -> np.ndarray:
        """Extract a 256-dimensional speaker embedding from an audio file.

        The audio is loaded, peak-normalized, converted to a log-mel
        spectrogram, and passed through the encoder model.

        Args:
            audio_path: Path to an audio file (WAV, FLAC, OGG, etc.).

        Returns:
            1-D float32 array of shape ``(256,)``, L2-normalized.
        """
        mel = self._audio_to_mel(audio_path)
        return self._infer(mel)

    def encode_batch(self, audio_paths: list[str | Path]) -> np.ndarray:
        """Extract speaker embeddings for multiple audio files.

        All mel spectrograms are zero-padded to the longest in the batch
        so they can be processed in a single forward pass.

        Args:
            audio_paths: List of audio file paths.

        Returns:
            2-D float32 array of shape ``(len(audio_paths), 256)``.
        """
        if not audio_paths:
            return np.empty((0, 256), dtype=np.float32)

        mels = [self._audio_to_mel(p) for p in audio_paths]

        # Pad to uniform time length
        max_time = max(m.shape[1] for m in mels)
        padded = np.zeros((len(mels), mels[0].shape[0], max_time), dtype=np.float32)
        for i, m in enumerate(mels):
            padded[i, :, : m.shape[1]] = m

        return self._infer_batch(padded)

    @staticmethod
    def similarity(emb1: np.ndarray, emb2: np.ndarray) -> float:
        """Compute cosine similarity between two embeddings.

        Both embeddings should already be L2-normalized (as returned by
        :meth:`encode`), but this method re-normalizes for safety.

        Args:
            emb1: 1-D float32 array of shape ``(emb_dim,)``.
            emb2: 1-D float32 array of shape ``(emb_dim,)``.

        Returns:
            Cosine similarity in [-1, 1].
        """
        emb1 = emb1.flatten().astype(np.float64)
        emb2 = emb2.flatten().astype(np.float64)

        norm1 = np.linalg.norm(emb1)
        norm2 = np.linalg.norm(emb2)
        if norm1 < 1e-12 or norm2 < 1e-12:
            return 0.0

        return float(np.dot(emb1, emb2) / (norm1 * norm2))

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _audio_to_mel(self, audio_path: str | Path) -> np.ndarray:
        """Load audio and compute mel spectrogram.

        Returns:
            (n_mels, time) float32 array.
        """
        audio = load_audio(audio_path, sr=DEFAULT_SR)
        audio = normalize_audio(audio)
        mel = compute_mel_spectrogram(
            audio,
            sr=DEFAULT_SR,
            n_fft=DEFAULT_N_FFT,
            hop_length=DEFAULT_HOP_LENGTH,
            n_mels=DEFAULT_N_MELS,
            fmin=DEFAULT_FMIN,
            fmax=DEFAULT_FMAX,
        )
        return mel

    def _infer(self, mel: np.ndarray) -> np.ndarray:
        """Run inference on a single mel spectrogram.

        Args:
            mel: (n_mels, time) float32 array.

        Returns:
            1-D float32 array of shape ``(emb_dim,)``.
        """
        # Add batch dimension: (1, n_mels, time)
        mel_batch = mel[np.newaxis, :, :]
        return self._infer_batch(mel_batch)[0]

    def _infer_batch(self, mel_batch: np.ndarray) -> np.ndarray:
        """Run inference on a batch of mel spectrograms.

        Args:
            mel_batch: (batch, n_mels, time) float32 array.

        Returns:
            (batch, emb_dim) float32 array.
        """
        if self._mode == "pytorch":
            return self._infer_pytorch(mel_batch)
        elif self._mode == "onnx":
            return self._infer_onnx(mel_batch)
        else:
            raise RuntimeError(
                "SpeakerEncoder not initialised. Use from_pytorch() or from_onnx()."
            )

    def _infer_pytorch(self, mel_batch: np.ndarray) -> np.ndarray:
        """Run inference with the PyTorch model."""
        import torch  # noqa: PLC0415

        assert self._pytorch_model is not None

        tensor = torch.from_numpy(mel_batch).to(self._pytorch_device)
        with torch.no_grad():
            embedding = self._pytorch_model(tensor)

        return embedding.cpu().numpy()

    def _infer_onnx(self, mel_batch: np.ndarray) -> np.ndarray:
        """Run inference with the ONNX model."""
        assert self._onnx_session is not None

        input_name = self._onnx_session.get_inputs()[0].name
        output_name = self._onnx_session.get_outputs()[0].name

        result = self._onnx_session.run(
            [output_name],
            {input_name: mel_batch.astype(np.float32)},
        )
        return result[0]