diff --git a/examples/tts/magpietts_inference.py b/examples/tts/magpietts_inference.py index d93a6c9ba2fb..cc50159c4752 100644 --- a/examples/tts/magpietts_inference.py +++ b/examples/tts/magpietts_inference.py @@ -167,9 +167,13 @@ def run_inference_and_evaluation( # CSV headers csv_header = ( "checkpoint_name,dataset,cer_filewise_avg,wer_filewise_avg,cer_cumulative," - "wer_cumulative,ssim_pred_gt_avg,ssim_pred_context_avg,ssim_gt_context_avg," + "wer_cumulative,cer_pred_gt_audio_filewise_avg,cer_pred_gt_audio_cumulative," + "wer_pred_gt_audio_filewise_avg,wer_pred_gt_audio_cumulative," + "ssim_pred_gt_avg,ssim_pred_context_avg,ssim_gt_context_avg," "ssim_pred_gt_avg_alternate,ssim_pred_context_avg_alternate," - "ssim_gt_context_avg_alternate,cer_gt_audio_cumulative,wer_gt_audio_cumulative," + "ssim_gt_context_avg_alternate,esim_pred_gt_avg,ems_pred_gt_avg," + "pitch_distance_avg,intensity_distance_avg,speech_rate_distance_avg," + "cer_gt_audio_cumulative,wer_gt_audio_cumulative," "utmosv2_avg,total_gen_audio_seconds,frechet_codec_distance," "eou_cutoff_rate,eou_silence_rate,eou_noise_rate,eou_error_rate," "katakana_cer_filewise_avg,katakana_cer_cumulative" @@ -303,6 +307,8 @@ def run_inference_and_evaluation( with_utmosv2=eval_config.with_utmosv2, with_fcd=eval_config.with_fcd, codec_model_path=eval_config.codec_model_path, + with_prosody_metrics=eval_config.with_prosody_metrics, + prosody_model_size=eval_config.prosody_model_size, strip_text_annotations_for_metrics=eval_config.strip_text_annotations_for_metrics, device=eval_config.device, asr_batch_size=eval_config.asr_batch_size, @@ -456,6 +462,8 @@ def main(argv=None): with_utmosv2=not args.disable_utmosv2, with_fcd=not args.disable_fcd, codec_model_path=args.codecmodel_path if not args.disable_fcd else None, + with_prosody_metrics=args.with_prosody_metrics, + prosody_model_size=args.prosody_model_size, strip_text_annotations_for_metrics=args.strip_text_annotations_for_metrics, asr_batch_size=args.asr_batch_size, eou_batch_size=args.eou_batch_size, diff --git a/nemo/collections/common/data/lhotse/sampling.py b/nemo/collections/common/data/lhotse/sampling.py index e73039d252f9..6e5fcd3bb1f5 100644 --- a/nemo/collections/common/data/lhotse/sampling.py +++ b/nemo/collections/common/data/lhotse/sampling.py @@ -342,6 +342,11 @@ def __call__(self, example) -> bool: else: speaker_id = getattr(supervision, field, None) + # Support the TTS speaker ID format: + # | Language:en Dataset: Speaker: | + if isinstance(speaker_id, str) and "Speaker:" in speaker_id: + speaker_id = speaker_id.rsplit("Speaker:", maxsplit=1)[-1].split("|", maxsplit=1)[0].strip() + if speaker_id in excluded_speaker_ids: return False return True diff --git a/nemo/collections/tts/metrics/emotion_encoder.py b/nemo/collections/tts/metrics/emotion_encoder.py new file mode 100644 index 000000000000..54dee78083a2 --- /dev/null +++ b/nemo/collections/tts/metrics/emotion_encoder.py @@ -0,0 +1,1115 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Lightweight LAION Empathic Insight Voice interface. + +This script provides a Hugging Face-style Python class for LAION's +Empathic Insight Voice models without using ModelScope. + +It supports: + + 1. Restoring the Whisper encoder from Hugging Face. + 2. Restoring LAION classifier MLP heads from Hugging Face .pth files. + 3. Extracting the full Whisper encoder embedding: + [B, 1500, 768] + 4. Extracting one classifier projection embedding: + Small: [B, 64] + Large: [B, 128] + 5. Extracting an SV-style emotion similarity embedding by concatenating + multiple classifier projection embeddings: + Small, 40 labels: [B, 40 * 64] = [B, 2560] + Large, 40 labels: [B, 40 * 128] = [B, 5120] + 6. Extracting an official-style emotion score vector: + [B, num_labels] + 7. Computing ranked emotion predictions and cosine similarity. + +Recommended for emotion similarity: + + model = EmpathicInsightVoice.from_pretrained(size="small", device="cuda") + emb = model.extract_emotion_embedding("audio.wav", embedding_type="head_concat") + sim = model.emotion_similarity("a.wav", "b.wav", embedding_type="head_concat") + +Notes: + + - The official model is a collection of independent expert heads. Each head + predicts one emotion or attribute score. + - The "head_concat" embedding is an engineering adaptation for + speaker-verification-style similarity. It concatenates the learned + projection outputs from multiple classifier heads. + - The "score_vector" embedding is closer to the documented inference output: + a vector of raw emotion intensity scores. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, Optional, Sequence, Union + +import librosa +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from huggingface_hub import hf_hub_download +from transformers import WhisperForConditionalGeneration, WhisperProcessor + + +# ============================================================================= +# Label mapping +# ============================================================================= + +# MIRRORING-style 12 emotions (https://bench.theliva.ai/legacy/mirroring.html): +# Amusement, Anger, Elation, Impatience, Surprise, +# Emotional Numbness, Contemplation, Disappointment, +# Confusion, Pride, Affection, Sadness. +# +# The public-facing labels below are Python-friendly. +# Some labels map to LAION's original longer checkpoint names: +# impatience -> model_Impatience_and_Irritability_best.pth +# surprise -> model_Astonishment_Surprise_best.pth + +LAION_LABEL_TO_FILENAME: dict[str, str] = { + "amusement": "model_Amusement_best.pth", + "anger": "model_Anger_best.pth", + "elation": "model_Elation_best.pth", + "impatience": "model_Impatience_and_Irritability_best.pth", + "surprise": "model_Astonishment_Surprise_best.pth", + "emotional_numbness": "model_Emotional_Numbness_best.pth", + "contemplation": "model_Contemplation_best.pth", + "disappointment": "model_Disappointment_best.pth", + "confusion": "model_Confusion_best.pth", + "pride": "model_Pride_best.pth", + "affection": "model_Affection_best.pth", + "sadness": "model_Sadness_best.pth", +} + + +PRIMARY_EMOTION_LABELS: list[str] = [ + "amusement", + "anger", + "elation", + "impatience", + "surprise", + "emotional_numbness", + "contemplation", + "disappointment", + "confusion", + "pride", + "affection", + "sadness", +] + + +AUXILIARY_SIMILARITY_LABELS: list[str] = [] + +# ============================================================================= +# Model architecture specs +# ============================================================================= + +MODEL_SPECS: dict[str, dict[str, Any]] = { + "small": { + "repo_id": "laion/Empathic-Insight-Voice-Small", + "whisper_model_id": "laion/BUD-E-Whisper", + "sample_rate": 16000, + "max_audio_seconds": 30.0, + "seq_len": 1500, + "embed_dim": 768, + "projection_dim": 64, + "mlp_hidden_dims": [64, 32, 16], + "mlp_dropouts": [0.0, 0.1, 0.1, 0.1], + }, + "large": { + "repo_id": "laion/Empathic-Insight-Voice-Large", + "whisper_model_id": "laion/BUD-E-Whisper", + "sample_rate": 16000, + "max_audio_seconds": 30.0, + "seq_len": 1500, + "embed_dim": 768, + "projection_dim": 128, + "mlp_hidden_dims": [128, 64, 32], + "mlp_dropouts": [0.0, 0.1, 0.1, 0.1], + }, +} + + +# ============================================================================= +# MLP head +# ============================================================================= + + +class FullEmbeddingMLP(nn.Module): + """Classifier head used by Empathic Insight Voice. + + The model receives a full Whisper encoder sequence embedding: + + [batch, seq_len, embed_dim] + + For Empathic Insight Voice this is normally: + + [batch, 1500, 768] + + It then performs: + + flatten -> projection -> MLP -> scalar score + + The projection output is useful as an SV-style emotion embedding: + + Small: [batch, 64] + Large: [batch, 128] + + Each restored classifier head has its own projection layer. Therefore, + "anger" projection, "sadness" projection, and "arousal" projection are + all different learned spaces. + """ + + def __init__( + self, + seq_len: int, + embed_dim: int, + projection_dim: int, + mlp_hidden_dims: Sequence[int], + mlp_dropout_rates: Sequence[float], + ) -> None: + super().__init__() + + if len(mlp_dropout_rates) != len(mlp_hidden_dims) + 1: + raise ValueError( + "Dropout rates length error. " + f"Expected {len(mlp_hidden_dims) + 1}, " + f"got {len(mlp_dropout_rates)}." + ) + + self.seq_len = seq_len + self.embed_dim = embed_dim + self.projection_dim = projection_dim + + self.flatten = nn.Flatten() + self.proj = nn.Linear(seq_len * embed_dim, projection_dim) + + layers: list[nn.Module] = [ + nn.ReLU(), + nn.Dropout(mlp_dropout_rates[0]), + ] + + current_dim = projection_dim + for i, hidden_dim in enumerate(mlp_hidden_dims): + layers.extend( + [ + nn.Linear(current_dim, hidden_dim), + nn.ReLU(), + nn.Dropout(mlp_dropout_rates[i + 1]), + ] + ) + current_dim = hidden_dim + + layers.append(nn.Linear(current_dim, 1)) + self.mlp = nn.Sequential(*layers) + + def extract_projected_embedding(self, x: torch.Tensor) -> torch.Tensor: + """Return the classifier projection embedding before the MLP. + + Args: + x: + Whisper embedding with shape [B, seq_len, embed_dim], or + [B, 1, seq_len, embed_dim]. + + Returns: + Projected embedding with shape [B, projection_dim]. + """ + if x.ndim == 4 and x.shape[1] == 1: + x = x.squeeze(1) + + if x.ndim != 3: + raise ValueError(f"Expected x with shape [B, T, C], got shape {tuple(x.shape)}.") + + return self.proj(self.flatten(x)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Return one scalar score per input example.""" + projected = self.extract_projected_embedding(x) + return self.mlp(projected) + + +# ============================================================================= +# Main class +# ============================================================================= + + +class EmpathicInsightVoice(nn.Module): + """Lightweight Hugging Face-style Empathic Insight Voice class. + + This class intentionally does not inherit from NeMo ModelPT. It behaves + like a normal PyTorch/Hugging Face utility model. + + Main methods: + + - extract_whisper_embedding(audio_path) + Returns [1, 1500, 768]. + + - extract_classifier_projection(audio_path, label) + Returns one head-specific projection: + Small: [1, 64] + Large: [1, 128] + + - extract_emotion_embedding(audio_path, embedding_type="head_concat") + Recommended SV-style emotion similarity embedding. + + - predict_emotions_from_embedding(embedding) + Returns raw scores and ranked softmax-like top emotions. + + - compute(audio_path) + Returns emotion predictions and optionally an embedding. + + - emotion_similarity(audio_path_a, audio_path_b) + Returns cosine similarity between extracted emotion embeddings. + """ + + def __init__( + self, + size: str = "small", + device: Union[str, torch.device] = "cuda", + mlp_device: Optional[Union[str, torch.device]] = None, + cache_dir: Optional[Union[str, Path]] = None, + cache_classifiers: bool = True, + load_all_classifiers: bool = False, + top_k_emotions: int = 5, + torch_dtype: Optional[torch.dtype] = None, + trust_remote_code: bool = False, + ) -> None: + super().__init__() + + if size not in MODEL_SPECS: + raise ValueError(f"Unsupported size={size!r}. Expected one of {sorted(MODEL_SPECS)}.") + + self.size = size + self.spec = MODEL_SPECS[size] + self.cache_classifiers = cache_classifiers + self.top_k_emotions = top_k_emotions + self.cache_dir = Path(cache_dir) if cache_dir is not None else None + + requested_device = torch.device(device) + if requested_device.type == "cuda" and not torch.cuda.is_available(): + requested_device = torch.device("cpu") + + self.device = requested_device + self.mlp_device = torch.device(mlp_device) if mlp_device is not None else self.device + + self.sample_rate = int(self.spec["sample_rate"]) + self.max_audio_seconds = float(self.spec["max_audio_seconds"]) + + # Load Whisper processor and encoder model. + self.processor = WhisperProcessor.from_pretrained( + self.spec["whisper_model_id"], + cache_dir=str(self.cache_dir) if self.cache_dir is not None else None, + trust_remote_code=trust_remote_code, + ) + + whisper_kwargs: dict[str, Any] = { + "cache_dir": str(self.cache_dir) if self.cache_dir is not None else None, + "trust_remote_code": trust_remote_code, + } + if torch_dtype is not None: + whisper_kwargs["torch_dtype"] = torch_dtype + + self.whisper_model = WhisperForConditionalGeneration.from_pretrained( + self.spec["whisper_model_id"], + **whisper_kwargs, + ).to(self.device) + + self.whisper_model.eval() + + # Restored MLP heads are stored here. + # + # ModuleDict keys must be sanitized because labels can contain punctuation. + self.classifiers = nn.ModuleDict() + + if load_all_classifiers: + self.load_classifiers() + + @classmethod + def from_pretrained( + cls, + size: str = "small", + **kwargs: Any, + ) -> "EmpathicInsightVoice": + """Construct the model using Hugging Face checkpoints. + + Example: + model = EmpathicInsightVoice.from_pretrained( + size="small", + device="cuda", + mlp_device="cuda", + ) + """ + return cls(size=size, **kwargs) + + @property + def repo_id(self) -> str: + """Hugging Face repo ID for the selected model size.""" + return str(self.spec["repo_id"]) + + @property + def available_labels(self) -> list[str]: + """All labels known to this script.""" + return list(LAION_LABEL_TO_FILENAME.keys()) + + @property + def projection_dim(self) -> int: + """Classifier projection dimension for the selected model size.""" + return int(self.spec["projection_dim"]) + + # ------------------------------------------------------------------------- + # Audio and Whisper embedding extraction + # ------------------------------------------------------------------------- + + @torch.no_grad() + def extract_whisper_embedding(self, audio_path: Union[str, Path]) -> torch.Tensor: + """Extract the full Whisper encoder embedding from an audio file. + + Args: + audio_path: + Path to an audio file readable by librosa. + + Returns: + Tensor with shape [1, 1500, 768]. + """ + waveform, _ = librosa.load(str(audio_path), sr=self.sample_rate, mono=True) + waveform = self._prepare_waveform(waveform) + return self.extract_whisper_embedding_from_waveform(waveform) + + @torch.no_grad() + def extract_whisper_embedding_from_waveform( + self, + waveform: np.ndarray, + ) -> torch.Tensor: + """Extract the full Whisper encoder embedding from a waveform. + + Args: + waveform: + Mono waveform at self.sample_rate. + + Returns: + Tensor with shape [1, 1500, 768]. + """ + waveform = self._prepare_waveform(waveform) + + input_features = self.processor( + waveform, + sampling_rate=self.sample_rate, + return_tensors="pt", + ).input_features + + input_features = input_features.to(self.device) + input_features = input_features.to(self.whisper_model.dtype) + + encoder_outputs = self.whisper_model.get_encoder()(input_features=input_features) + + embedding = encoder_outputs.last_hidden_state + embedding = self._pad_or_trim_embedding(embedding) + + return embedding + + # ------------------------------------------------------------------------- + # Classifier projection extraction + # ------------------------------------------------------------------------- + + @torch.no_grad() + def extract_classifier_projection( + self, + audio_path: Union[str, Path], + label: str, + normalize: bool = True, + ) -> torch.Tensor: + """Extract one head-specific classifier projection embedding. + + This is the closest equivalent to extracting an x-vector-like embedding + from one specific emotion classifier head. + + Flow: + audio -> Whisper encoder -> label-specific classifier projection + + Args: + audio_path: + Input audio path. + label: + Label whose classifier projection should be used, for example: + "anger", "sadness", "arousal". + normalize: + If True, apply L2 normalization. + + Returns: + Small: [1, 64] + Large: [1, 128] + """ + whisper_embedding = self.extract_whisper_embedding(audio_path) + return self.extract_classifier_projection_from_whisper_embedding( + whisper_embedding=whisper_embedding, + label=label, + normalize=normalize, + ) + + @torch.no_grad() + def extract_classifier_projection_from_whisper_embedding( + self, + whisper_embedding: torch.Tensor, + label: str, + normalize: bool = True, + ) -> torch.Tensor: + """Extract one classifier projection from an existing Whisper embedding.""" + classifier = self._get_classifier(label) + param = next(classifier.parameters()) + + working_embedding = whisper_embedding.to(device=param.device, dtype=param.dtype) + projected = classifier.extract_projected_embedding(working_embedding) + + projected = projected.float() + if normalize: + projected = F.normalize(projected, p=2, dim=-1) + + return projected + + # ------------------------------------------------------------------------- + # Emotion similarity embeddings + # ------------------------------------------------------------------------- + + @torch.no_grad() + def extract_emotion_embedding( + self, + audio_path: Union[str, Path], + labels: Optional[Sequence[str]] = None, + embedding_type: str = "head_concat", + normalize: bool = True, + include_auxiliary: bool = False, + ) -> torch.Tensor: + """Extract a fixed-dimensional emotion embedding. + + Recommended for SV-style emotion similarity: + embedding_type="head_concat" + + Supported embedding types: + + 1. "head_concat" + Concatenate the projection output of each selected classifier head. + + Small, 40 primary labels: + [1, 40 * 64] = [1, 2560] + + Large, 40 primary labels: + [1, 40 * 128] = [1, 5120] + + 2. "head_mean" + Average the projection outputs across selected heads. + + Small: + [1, 64] + + Large: + [1, 128] + + 3. "score_vector" + Use raw scalar outputs from the selected classifier heads. + + Shape: + [1, num_labels] + + This is closest to the official annotation output. + + Args: + audio_path: + Input audio path. + labels: + Labels to use. If None, PRIMARY_EMOTION_LABELS are used. + embedding_type: + "head_concat", "head_mean", or "score_vector". + normalize: + If True, apply L2 normalization to the final embedding. + include_auxiliary: + If labels is None, append AUXILIARY_SIMILARITY_LABELS. + + Returns: + torch.Tensor fixed-dimensional emotion embedding. + """ + whisper_embedding = self.extract_whisper_embedding(audio_path) + labels_to_run = self._default_similarity_labels( + labels=labels, + include_auxiliary=include_auxiliary, + ) + + return self.extract_emotion_embedding_from_whisper_embedding( + whisper_embedding=whisper_embedding, + labels=labels_to_run, + embedding_type=embedding_type, + normalize=normalize, + ) + + @torch.no_grad() + def extract_emotion_embedding_from_whisper_embedding( + self, + whisper_embedding: torch.Tensor, + labels: Optional[Sequence[str]] = None, + embedding_type: str = "head_concat", + normalize: bool = True, + include_auxiliary: bool = False, + ) -> torch.Tensor: + """Extract a fixed-dimensional emotion embedding from Whisper features.""" + labels_to_run = self._default_similarity_labels( + labels=labels, + include_auxiliary=include_auxiliary, + ) + + if embedding_type == "score_vector": + prediction = self.predict_emotions_from_embedding( + embedding=whisper_embedding, + labels=labels_to_run, + return_raw_scores=True, + rank_scores=False, + ) + raw_scores = prediction["raw_scores"] + + output = torch.tensor( + [raw_scores[label] for label in labels_to_run], + dtype=torch.float32, + ).unsqueeze(0) + + elif embedding_type in {"head_concat", "head_mean"}: + projected_embeddings: list[torch.Tensor] = [] + + for label in labels_to_run: + classifier = self._get_classifier(label) + param = next(classifier.parameters()) + + working_embedding = whisper_embedding.to( + device=param.device, + dtype=param.dtype, + ) + + projected = classifier.extract_projected_embedding(working_embedding) + projected_embeddings.append(projected.float().cpu()) + + if embedding_type == "head_concat": + output = torch.cat(projected_embeddings, dim=-1) + else: + output = torch.stack(projected_embeddings, dim=0).mean(dim=0) + + else: + raise ValueError( + f"Unsupported embedding_type={embedding_type!r}. " + "Expected 'head_concat', 'head_mean', or 'score_vector'." + ) + + if normalize: + output = F.normalize(output, p=2, dim=-1) + + return output + + @torch.no_grad() + def emotion_similarity( + self, + audio_path_a: Union[str, Path], + audio_path_b: Union[str, Path], + labels: Optional[Sequence[str]] = None, + embedding_type: str = "head_concat", + include_auxiliary: bool = False, + ) -> float: + """Compute cosine similarity between two audios in emotion space. + + Args: + audio_path_a: + First audio path. + audio_path_b: + Second audio path. + labels: + Optional label subset. + embedding_type: + "head_concat", "head_mean", or "score_vector". + include_auxiliary: + If labels is None, append auxiliary labels. + + Returns: + Cosine similarity as a Python float. + """ + emb_a = self.extract_emotion_embedding( + audio_path=audio_path_a, + labels=labels, + embedding_type=embedding_type, + normalize=True, + include_auxiliary=include_auxiliary, + ) + emb_b = self.extract_emotion_embedding( + audio_path=audio_path_b, + labels=labels, + embedding_type=embedding_type, + normalize=True, + include_auxiliary=include_auxiliary, + ) + + return float(F.cosine_similarity(emb_a, emb_b, dim=-1).item()) + + # ------------------------------------------------------------------------- + # Prediction + # ------------------------------------------------------------------------- + @torch.no_grad() + def compare_emotion_pair( + self, + audio_path_a: Union[str, Path], + audio_path_b: Union[str, Path], + labels: Optional[Sequence[str]] = None, + embedding_type: str = "score_vector", + ) -> dict[str, Any]: + """Compare two audio files using the 12-emotion set. + + This method does not perform corpus-level ranking. It only returns: + + - top emotion for audio A + - top emotion for audio B + - matched emotion label if both top emotions match + - emotion similarity + + Args: + audio_path_a: + First audio file. + audio_path_b: + Second audio file. + labels: + Optional subset of labels. Defaults to PRIMARY_EMOTION_LABELS, + which is the 12-emotion set. + embedding_type: + Similarity representation: + - "score_vector": cosine over raw 12-emotion score vector. + - "head_concat": cosine over concatenated classifier projections. + - "head_mean": cosine over averaged classifier projections. + + For MIRRORING-style emotion-vector similarity, use "score_vector". + + Returns: + { + "audio_path_a": str, + "audio_path_b": str, + "audio_a_top_emotion": str | None, + "audio_b_top_emotion": str | None, + "top_emotion_match": bool , + "emotion_similarity": float, + "audio_a_raw_scores": dict[str, float], + "audio_b_raw_scores": dict[str, float], + } + """ + labels_to_run = self._validate_labels(labels or PRIMARY_EMOTION_LABELS) + + result_a = self.compute( + audio_path=audio_path_a, + labels=labels_to_run, + return_embedding=False, + return_raw_scores=True, + ) + + result_b = self.compute( + audio_path=audio_path_b, + labels=labels_to_run, + return_embedding=False, + return_raw_scores=True, + ) + + top_a = result_a["top_emotion"] + top_b = result_b["top_emotion"] + + similarity = self.emotion_similarity( + audio_path_a=audio_path_a, + audio_path_b=audio_path_b, + labels=labels_to_run, + embedding_type=embedding_type, + include_auxiliary=False, + ) + + return { + "audio_path_a": str(audio_path_a), + "audio_path_b": str(audio_path_b), + "audio_a_top_emotion": top_a, + "audio_b_top_emotion": top_b, + "top_emotion_match": top_a is not None and top_a == top_b, + "emotion_similarity": similarity, + "audio_a_raw_scores": result_a["raw_scores"], + "audio_b_raw_scores": result_b["raw_scores"], + } + + @torch.no_grad() + def compute( + self, + audio_path: Union[str, Path], + labels: Optional[Sequence[str]] = None, + return_embedding: bool = True, + embedding_type: str = "head_concat", + return_raw_scores: bool = True, + include_auxiliary_for_embedding: bool = False, + ) -> dict[str, Any]: + """Compute emotion predictions and optionally an embedding. + + Args: + audio_path: + Input audio file. + labels: + Prediction labels. If None, all known labels are attempted. + For the official 40-emotion profile, pass PRIMARY_EMOTION_LABELS. + return_embedding: + If True, return an emotion embedding. + embedding_type: + Embedding type to return: + "head_concat", "head_mean", or "score_vector". + return_raw_scores: + If True, return raw classifier outputs. + include_auxiliary_for_embedding: + If True and return_embedding=True, include auxiliary labels in the + returned embedding. + + Returns: + { + "audio_path": str, + "model_size": "small" | "large", + "top_emotion": str | None, + "emotions": { + label: {"score": float, "rank": int} + }, + "raw_scores": { + label: float + }, + "embedding": torch.Tensor, + "embedding_type": str + } + """ + whisper_embedding = self.extract_whisper_embedding(audio_path) + + prediction = self.predict_emotions_from_embedding( + embedding=whisper_embedding, + labels=labels, + return_raw_scores=return_raw_scores, + rank_scores=True, + ) + + top_emotion = None + if prediction["emotions"]: + top_emotion = next(iter(prediction["emotions"])) + + output: dict[str, Any] = { + "audio_path": str(audio_path), + "model_size": self.size, + "top_emotion": top_emotion, + "emotions": prediction["emotions"], + } + + if return_raw_scores: + output["raw_scores"] = prediction["raw_scores"] + + if return_embedding: + output["embedding"] = self.extract_emotion_embedding_from_whisper_embedding( + whisper_embedding=whisper_embedding, + labels=None, + embedding_type=embedding_type, + normalize=True, + include_auxiliary=include_auxiliary_for_embedding, + ) + output["embedding_type"] = embedding_type + + return output + + @torch.no_grad() + def predict_emotions_from_embedding( + self, + embedding: torch.Tensor, + labels: Optional[Sequence[str]] = None, + return_raw_scores: bool = True, + rank_scores: bool = True, + ) -> dict[str, Any]: + """Predict emotion or attribute scores from a Whisper embedding. + + Args: + embedding: + Whisper encoder embedding, normally [1, 1500, 768]. + labels: + Labels to evaluate. If None, all known labels are attempted. + return_raw_scores: + If True, include raw classifier scores. + rank_scores: + If True, softmax and rank scores into top-k emotions. + + Returns: + Dict containing: + "emotions": ranked top-k scores if rank_scores=True + "raw_scores": raw scalar outputs if return_raw_scores=True + """ + labels_to_run = self._validate_labels(labels) + + raw_scores: dict[str, float] = {} + + for label in labels_to_run: + classifier = self._get_classifier(label) + param = next(classifier.parameters()) + + working_embedding = embedding.to(device=param.device, dtype=param.dtype) + score = classifier(working_embedding).detach().cpu().item() + raw_scores[label] = float(score) + + if not self.cache_classifiers: + cache_key = self._cache_key(label) + if cache_key in self.classifiers: + del self.classifiers[cache_key] + + output: dict[str, Any] = {} + + if rank_scores: + output["emotions"] = self._softmax_and_rank(raw_scores) + else: + output["emotions"] = {} + + if return_raw_scores: + output["raw_scores"] = raw_scores + + return output + + # ------------------------------------------------------------------------- + # Classifier loading + # ------------------------------------------------------------------------- + + def load_classifiers( + self, + labels: Optional[Sequence[str]] = None, + ) -> None: + """Eagerly download and restore classifier MLPs. + + By default the class lazy-loads heads when needed. This method is useful + when you want to pre-load selected heads before repeated inference. + """ + labels_to_load = self._validate_labels(labels) + for label in labels_to_load: + self._get_classifier(label) + + def _get_classifier(self, label: str) -> FullEmbeddingMLP: + """Download, reconstruct, and restore one classifier MLP head. + + This is where the classifier MLP is restored: + + classifier = FullEmbeddingMLP(...) + state_dict = torch.load(...) + state_dict = strip "_orig_mod." prefix if needed + classifier.load_state_dict(state_dict) + + Args: + label: + Python-friendly label key, such as "anger" or "arousal". + + Returns: + Restored FullEmbeddingMLP. + """ + if label not in LAION_LABEL_TO_FILENAME: + raise ValueError(f"Unknown label {label!r}. Available labels: " f"{sorted(LAION_LABEL_TO_FILENAME)}") + + cache_key = self._cache_key(label) + + if cache_key in self.classifiers: + classifier = self.classifiers[cache_key] + if not isinstance(classifier, FullEmbeddingMLP): + raise TypeError(f"Cached classifier for {label!r} has unexpected type " f"{type(classifier)}.") + return classifier + + filename = LAION_LABEL_TO_FILENAME[label] + + local_path = hf_hub_download( + repo_id=self.repo_id, + filename=filename, + cache_dir=str(self.cache_dir) if self.cache_dir is not None else None, + repo_type="model", + ) + + classifier = FullEmbeddingMLP( + seq_len=int(self.spec["seq_len"]), + embed_dim=int(self.spec["embed_dim"]), + projection_dim=int(self.spec["projection_dim"]), + mlp_hidden_dims=list(self.spec["mlp_hidden_dims"]), + mlp_dropout_rates=list(self.spec["mlp_dropouts"]), + ) + + state_dict = torch.load(local_path, map_location="cpu") + + if not isinstance(state_dict, dict): + raise RuntimeError(f"Expected {local_path} to contain a state_dict, " f"but got {type(state_dict)}.") + + state_dict = self._strip_orig_mod_prefix_if_needed(state_dict) + + try: + classifier.load_state_dict(state_dict) + except RuntimeError as exc: + raise RuntimeError( + f"Failed to load classifier for label={label!r} from {local_path}. " + f"This often means the selected size={self.size!r} has different " + "MLP dimensions than MODEL_SPECS declares." + ) from exc + + classifier.eval() + classifier = classifier.to(self.mlp_device) + + # Keep classifier dtype compatible with Whisper dtype if user loaded + # Whisper in fp16/bf16. + if self.whisper_model.dtype in (torch.float16, torch.bfloat16): + classifier = classifier.to(dtype=self.whisper_model.dtype) + + if self.cache_classifiers: + self.classifiers[cache_key] = classifier + + return classifier + + # ------------------------------------------------------------------------- + # Internal helpers + # ------------------------------------------------------------------------- + + def _validate_labels( + self, + labels: Optional[Sequence[str]], + ) -> list[str]: + """Validate label names and return a concrete list.""" + if labels is None: + return list(LAION_LABEL_TO_FILENAME.keys()) + + labels_list = list(labels) + unknown = sorted(set(labels_list) - set(LAION_LABEL_TO_FILENAME.keys())) + + if unknown: + raise ValueError( + f"Unknown labels: {unknown}. " f"Available labels: {sorted(LAION_LABEL_TO_FILENAME.keys())}" + ) + + return labels_list + + def _default_similarity_labels( + self, + labels: Optional[Sequence[str]], + include_auxiliary: bool, + ) -> list[str]: + """Choose the default labels for emotion similarity embeddings.""" + if labels is not None: + return self._validate_labels(labels) + + labels_to_run = list(PRIMARY_EMOTION_LABELS) + + if include_auxiliary: + labels_to_run.extend(AUXILIARY_SIMILARITY_LABELS) + + return self._validate_labels(labels_to_run) + + def _prepare_waveform(self, waveform: np.ndarray) -> np.ndarray: + """Convert waveform to mono float32 and trim to max_audio_seconds.""" + waveform = np.asarray(waveform, dtype=np.float32) + + if waveform.ndim > 1: + waveform = np.mean(waveform, axis=0).astype(np.float32) + + max_samples = int(self.sample_rate * self.max_audio_seconds) + + if waveform.shape[0] > max_samples: + waveform = waveform[:max_samples] + + return waveform + + def _pad_or_trim_embedding(self, embedding: torch.Tensor) -> torch.Tensor: + """Pad or trim Whisper encoder output to the expected sequence length.""" + seq_len = int(self.spec["seq_len"]) + embed_dim = int(self.spec["embed_dim"]) + + if embedding.ndim != 3: + raise RuntimeError(f"Expected Whisper embedding with shape [B, T, C], " f"got {tuple(embedding.shape)}.") + + if embedding.shape[-1] != embed_dim: + raise RuntimeError(f"Unexpected embedding dim. Expected {embed_dim}, " f"got {embedding.shape[-1]}.") + + current_seq_len = embedding.shape[1] + + if current_seq_len < seq_len: + padding = torch.zeros( + ( + embedding.shape[0], + seq_len - current_seq_len, + embed_dim, + ), + device=embedding.device, + dtype=embedding.dtype, + ) + embedding = torch.cat([embedding, padding], dim=1) + + elif current_seq_len > seq_len: + embedding = embedding[:, :seq_len, :] + + return embedding + + def _softmax_and_rank( + self, + raw_scores: dict[str, float], + ) -> dict[str, dict[str, Union[float, int]]]: + """Convert raw scores to a sorted top-k softmax dictionary. + + The raw MLP outputs are independent regression scores. This method is + mainly for producing an easy top emotion label. For similarity, prefer + raw score vectors or projection embeddings. + """ + if not raw_scores: + return {} + + labels = list(raw_scores.keys()) + values = np.array([raw_scores[label] for label in labels], dtype=np.float32) + + values = values - np.max(values) + exp_values = np.exp(values) + probs = exp_values / np.sum(exp_values) + + ranked = sorted( + zip(labels, probs.tolist()), + key=lambda item: item[1], + reverse=True, + ) + + ranked = ranked[: self.top_k_emotions] + + return { + label: { + "score": float(prob), + "rank": rank, + } + for rank, (label, prob) in enumerate(ranked, start=1) + } + + @staticmethod + def _strip_orig_mod_prefix_if_needed( + state_dict: dict[str, torch.Tensor], + ) -> dict[str, torch.Tensor]: + """Strip torch.compile '_orig_mod.' prefixes if present.""" + if not any(key.startswith("_orig_mod.") for key in state_dict.keys()): + return state_dict + + return { + key[len("_orig_mod.") :] if key.startswith("_orig_mod.") else key: value + for key, value in state_dict.items() + } + + @staticmethod + def _cache_key(label: str) -> str: + """Convert an arbitrary label into a safe ModuleDict key.""" + return label.replace(".", "_").replace("/", "_").replace("-", "_").replace(" ", "_").replace("&", "and") + + def cleanup(self) -> None: + """Move modules to CPU and clear classifier cache.""" + for key in list(self.classifiers.keys()): + self.classifiers[key].cpu() + + self.classifiers.clear() + self.whisper_model.cpu() + + if torch.cuda.is_available(): + torch.cuda.empty_cache() diff --git a/nemo/collections/tts/metrics/prosody.py b/nemo/collections/tts/metrics/prosody.py new file mode 100644 index 000000000000..4db446d59041 --- /dev/null +++ b/nemo/collections/tts/metrics/prosody.py @@ -0,0 +1,304 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Reference-based acoustic prosody distances for TTS evaluation.""" + +from __future__ import annotations + +import math +import os +from dataclasses import dataclass +from typing import Any + +import librosa +import numpy as np +from numba import njit + +_SAMPLE_RATE = 16000 +_RES_TYPE = "soxr_hq" +_FRAME_SHIFT_MS = 20.0 +_FRAME_LENGTH_MS = 64.0 +_FMIN = 55.0 +_FMAX = 450.0 +_PYIN_N_THRESHOLDS = 24 +_PYIN_BETA_PARAMETERS = (2.0, 18.0) +_PYIN_BOLTZMANN_PARAMETER = 2.0 +_PYIN_RESOLUTION = 0.25 +_PYIN_MAX_TRANSITION_RATE = 12.0 +_PYIN_SWITCH_PROB = 0.01 +_PYIN_NO_TROUGH_PROB = 0.01 +_MAX_DTW_FRAMES = 1000 +_DTW_BAND_RATIO = 0.05 +_F0_NAN_PENALTY = 6.0 +_MIN_VOICED_FRAMES = 5 + + +@dataclass(frozen=True) +class ProsodyDistanceResult: + """Per-pair acoustic prosody distance metrics.""" + + pitch_distance: float + intensity_distance: float + speech_rate_distance: float + + def to_dict(self) -> dict[str, float]: + """Return a JSON-serializable dictionary.""" + return { + "pitch_distance": self.pitch_distance, + "intensity_distance": self.intensity_distance, + "speech_rate_distance": self.speech_rate_distance, + } + + +@njit +def _dtw_distance_1d_numba(x: np.ndarray, y: np.ndarray, nan_penalty: float, band_radius: int) -> float: + n = x.shape[0] + m = y.shape[0] + if n == 0 or m == 0: + return np.nan + + inf = 1.0e30 + prev = np.empty(m + 1, dtype=np.float64) + curr = np.empty(m + 1, dtype=np.float64) + for j in range(m + 1): + prev[j] = inf + curr[j] = inf + prev[0] = 0.0 + + for i in range(1, n + 1): + for j in range(m + 1): + curr[j] = inf + + j_start = max(1, i - band_radius) + j_end = min(m, i + band_radius) + 1 + + for j in range(j_start, j_end): + xv = x[i - 1] + yv = y[j - 1] + x_nan = np.isnan(xv) + y_nan = np.isnan(yv) + if x_nan and y_nan: + cost = 0.0 + elif x_nan or y_nan: + cost = nan_penalty + else: + diff = xv - yv + cost = diff if diff >= 0.0 else -diff + + best_prev = prev[j - 1] + if prev[j] < best_prev: + best_prev = prev[j] + if curr[j - 1] < best_prev: + best_prev = curr[j - 1] + curr[j] = cost + best_prev + + tmp = prev + prev = curr + curr = tmp + + total = prev[m] + if total >= inf / 2.0: + return np.nan + return total / float(n + m) + + +def compute_prosody_distances( + gt_audio_path: str, + pred_audio_path: str, + text: Any, +) -> ProsodyDistanceResult: + """Compute acoustic prosody distances between reference and generated audio. + + Args: + gt_audio_path: Ground-truth/reference audio path. + pred_audio_path: Generated/predicted audio path. + text: Reference text used for character-per-second speech rate. + + Returns: + ProsodyDistanceResult with pitch, intensity, and speech-rate distances. + """ + gt_audio, sr, gt_duration = _load_audio(gt_audio_path) + pred_audio, _, pred_duration = _load_audio(pred_audio_path) + + hop_length, frame_length = _frame_params(sr) + gt_log_energy = _compute_log_energy(gt_audio, frame_length=frame_length, hop_length=hop_length) + pred_log_energy = _compute_log_energy(pred_audio, frame_length=frame_length, hop_length=hop_length) + + gt_f0 = _compute_f0(gt_audio, sr=sr, frame_length=frame_length, hop_length=hop_length) + pred_f0 = _compute_f0(pred_audio, sr=sr, frame_length=frame_length, hop_length=hop_length) + + pitch_distance = float("nan") + if np.isfinite(gt_f0).sum() >= _MIN_VOICED_FRAMES and np.isfinite(pred_f0).sum() >= _MIN_VOICED_FRAMES: + gt_pitch, pred_pitch = _prepare_f0_for_metric(gt_f0, pred_f0) + pitch_distance = _dtw_distance_1d( + _maybe_reduce_for_dtw(gt_pitch), + _maybe_reduce_for_dtw(pred_pitch), + nan_penalty=_F0_NAN_PENALTY, + ) + + intensity_distance = _dtw_distance_1d( + _maybe_reduce_for_dtw(_zscore(gt_log_energy)), + _maybe_reduce_for_dtw(_zscore(pred_log_energy)), + nan_penalty=0.0, + ) + + gt_char_count = _char_count(text) + gt_speech_rate = gt_char_count / gt_duration if gt_duration > 0.0 else float("nan") + pred_speech_rate = gt_char_count / pred_duration if pred_duration > 0.0 else float("nan") + speech_rate_distance = abs(gt_speech_rate - pred_speech_rate) + + return ProsodyDistanceResult( + pitch_distance=_safe_float(pitch_distance), + intensity_distance=_safe_float(intensity_distance), + speech_rate_distance=_safe_float(speech_rate_distance), + ) + + +def _load_audio(path: str) -> tuple[np.ndarray, int, float]: + if not path: + raise FileNotFoundError("empty audio filepath") + if not os.path.exists(path): + raise FileNotFoundError(path) + + audio, sr = librosa.load(path, sr=_SAMPLE_RATE, mono=True, res_type=_RES_TYPE) + audio = np.asarray(audio, dtype=np.float32) + if audio.size == 0: + raise ValueError(f"empty audio after loading: {path}") + return audio, int(sr), float(audio.shape[0] / sr) + + +def _frame_params(sr: int) -> tuple[int, int]: + hop_length = max(1, int(round(sr * _FRAME_SHIFT_MS / 1000.0))) + frame_length = max(hop_length * 2, int(round(sr * _FRAME_LENGTH_MS / 1000.0))) + return hop_length, frame_length + + +def _compute_log_energy(audio: np.ndarray, frame_length: int, hop_length: int) -> np.ndarray: + rms = librosa.feature.rms(y=audio, frame_length=frame_length, hop_length=hop_length, center=True)[0] + return np.log(np.maximum(np.asarray(rms, dtype=np.float64), 1.0e-10)) + + +def _compute_f0(audio: np.ndarray, sr: int, frame_length: int, hop_length: int) -> np.ndarray: + f0, voiced_flag, _ = librosa.pyin( + y=np.asarray(audio, dtype=np.float64), + sr=sr, + fmin=_FMIN, + fmax=_FMAX, + frame_length=frame_length, + hop_length=hop_length, + center=True, + pad_mode="constant", + n_thresholds=_PYIN_N_THRESHOLDS, + beta_parameters=_PYIN_BETA_PARAMETERS, + boltzmann_parameter=_PYIN_BOLTZMANN_PARAMETER, + resolution=_PYIN_RESOLUTION, + max_transition_rate=_PYIN_MAX_TRANSITION_RATE, + switch_prob=_PYIN_SWITCH_PROB, + no_trough_prob=_PYIN_NO_TROUGH_PROB, + fill_na=np.nan, + ) + f0 = np.asarray(f0, dtype=np.float64) + if voiced_flag is None: + return f0 + + voiced_flag = np.asarray(voiced_flag, dtype=bool) + min_len = min(len(f0), len(voiced_flag)) + f0 = f0[:min_len] + f0[~voiced_flag[:min_len]] = np.nan + return f0 + + +def _prepare_f0_for_metric(gt_f0_hz: np.ndarray, pred_f0_hz: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + gt_median = float(np.nanmedian(gt_f0_hz)) if np.isfinite(gt_f0_hz).any() else float("nan") + return _hz_to_semitones(gt_f0_hz, gt_median), _hz_to_semitones(pred_f0_hz, gt_median) + + +def _dtw_distance_1d(x: np.ndarray, y: np.ndarray, nan_penalty: float) -> float: + x = np.asarray(x, dtype=np.float64) + y = np.asarray(y, dtype=np.float64) + if x.ndim != 1 or y.ndim != 1: + raise ValueError("DTW expects 1-D arrays") + if len(x) == 0 or len(y) == 0: + return float("nan") + + band_radius = max(abs(len(x) - len(y)), int(math.ceil(_DTW_BAND_RATIO * max(len(x), len(y))))) + return _safe_float(_dtw_distance_1d_numba(x, y, float(nan_penalty), int(band_radius))) + + +def _hz_to_semitones(f0_hz: np.ndarray, ref_hz: float) -> np.ndarray: + f0_hz = np.asarray(f0_hz, dtype=np.float64) + out = np.full_like(f0_hz, np.nan, dtype=np.float64) + if not np.isfinite(ref_hz) or ref_hz <= 0.0: + return out + valid = np.isfinite(f0_hz) & (f0_hz > 0.0) + out[valid] = 12.0 * np.log2(f0_hz[valid] / ref_hz) + return out + + +def _zscore(values: np.ndarray, eps: float = 1.0e-8) -> np.ndarray: + values = np.asarray(values, dtype=np.float64) + finite = np.isfinite(values) + if finite.sum() == 0: + return values + mean = float(np.mean(values[finite])) + std = float(np.std(values[finite])) + out = values.copy() + if std < eps: + out[finite] = out[finite] - mean + else: + out[finite] = (out[finite] - mean) / std + return out + + +def _maybe_reduce_for_dtw(values: np.ndarray) -> np.ndarray: + values = np.asarray(values, dtype=np.float64) + if len(values) <= _MAX_DTW_FRAMES: + return values + return _resample_1d_preserve_nans(values, _MAX_DTW_FRAMES) + + +def _resample_1d_preserve_nans(values: np.ndarray, target_len: int) -> np.ndarray: + if len(values) == target_len: + return values + if len(values) == 0: + return values + if len(values) == 1: + return np.full(target_len, values[0], dtype=np.float64) + + old_t = np.linspace(0.0, 1.0, len(values)) + new_t = np.linspace(0.0, 1.0, target_len) + valid = np.isfinite(values) + if valid.sum() == 0: + return np.full(target_len, np.nan, dtype=np.float64) + + out = np.interp(new_t, old_t[valid], values[valid]) + valid_interp = np.interp(new_t, old_t, valid.astype(np.float64)) + out[valid_interp < 0.5] = np.nan + return out.astype(np.float64) + + +def _char_count(text: Any) -> int: + if text is None: + return 0 + return sum(1 for ch in str(text) if not ch.isspace()) + + +def _safe_float(value: Any) -> float: + try: + value = float(value) + except (TypeError, ValueError): + return float("nan") + if not math.isfinite(value): + return float("nan") + return value diff --git a/nemo/collections/tts/modules/magpietts_inference/evaluate_generated_audio.py b/nemo/collections/tts/modules/magpietts_inference/evaluate_generated_audio.py index f406b52b2bd0..70bc8c7379ba 100644 --- a/nemo/collections/tts/modules/magpietts_inference/evaluate_generated_audio.py +++ b/nemo/collections/tts/modules/magpietts_inference/evaluate_generated_audio.py @@ -36,6 +36,7 @@ from nemo.collections.asr.metrics.wer import word_error_rate_detail from nemo.collections.tts.metrics.eou_classifier import EoUClassification, EoUClassifier, EoUType from nemo.collections.tts.metrics.frechet_codec_distance import FrechetCodecDistance +from nemo.collections.tts.metrics.prosody import compute_prosody_distances from nemo.collections.tts.parts.utils.tts_dataset_utils import ( JapaneseTextProcessor, NemoTranscriber, @@ -65,6 +66,12 @@ 'pred_katakana', ] +PROSODY_DISTANCE_KEYS = [ + 'pitch_distance', + 'intensity_distance', + 'speech_rate_distance', +] + # Regexes mirrored from the IPA preprocessing script that creates # custom["text_without_annotation"]. This is used only for text inputs # during metric computation when requested. @@ -98,8 +105,13 @@ def strip_text_annotations_from_text(text: str) -> str: FILEWISE_METRICS_TO_SAVE = [ 'cer', + 'cer_pred_gt_audio', 'wer', + 'wer_pred_gt_audio', 'pred_context_ssim', + 'pred_gt_esim', + 'pred_gt_ems', + *PROSODY_DISTANCE_KEYS, 'pred_text', 'gt_audio_text', 'gt_text', @@ -268,6 +280,8 @@ def load_evaluation_models( asr_model_name="stt_en_conformer_transducer_large", asr_model_type="nemo", device="cuda", + with_prosody_metrics=False, + prosody_model_size="small", ): """Load the ASR and speaker-verification models used for evaluation. @@ -278,6 +292,9 @@ def load_evaluation_models( asr_model_type: ASR model implementation. Supported values are ``"nemo"``, ``"nemo_with_prompt"``, and ``"whisper"``. device: Device on which the evaluation models are loaded. + with_prosody_metrics: Whether to compute ESIM/EMS plus pitch, + intensity, and speech-rate distance metrics. + prosody_model_size: Size of the emotion encoder. Supported values are ``"small"`` or ``"large"``. Returns: Dictionary containing: @@ -291,6 +308,8 @@ def load_evaluation_models( - ``sv_model``: Primary speaker-verification model. - ``sv_model_alternate``: Alternate ``titanet_small`` speaker-verification model. + - ``emotion_model``: Emotion encoder when ``with_prosody_metrics=True``; + otherwise ``None``. It is also ``None`` if loading fails. Raises: ValueError: If ``asr_model_type`` is unsupported. @@ -300,6 +319,7 @@ def load_evaluation_models( 'whisper_model': None, 'whisper_processor': None, 'feature_extractor': None, + 'emotion_model': None, } if asr_model_type == "nemo": @@ -326,9 +346,66 @@ def load_evaluation_models( ) models['sv_model_alternate'] = models['sv_model_alternate'].to(device).eval() + if with_prosody_metrics: + logging.info("Loading emotion encoder for ESIM/EMS prosody metrics...") + try: + from nemo.collections.tts.metrics.emotion_encoder import EmpathicInsightVoice + + models['emotion_model'] = EmpathicInsightVoice.from_pretrained( + size=prosody_model_size, + device=device, + mlp_device=device, + cache_classifiers=True, + load_all_classifiers=False, + top_k_emotions=1, + ).eval() + except Exception as e: + logging.warning(f"Emotion encoder could not be loaded: {e}. ESIM/EMS metrics will be set to NaN.") + return models +def compute_emotion_pair_metrics(emotion_model, gt_audio_path, pred_audio_path): + """Compute ground-truth to predicted emotion similarity and top-emotion match.""" + if emotion_model is None or gt_audio_path is None or pred_audio_path is None: + return float('NaN'), float('NaN') + + try: + result = emotion_model.compare_emotion_pair( + audio_path_a=gt_audio_path, + audio_path_b=pred_audio_path, + ) + return float(result["emotion_similarity"]), float(result["top_emotion_match"]) + except Exception as e: + logging.warning(f"Could not compute ESIM/EMS for {gt_audio_path} and {pred_audio_path}: {e}") + return float('NaN'), float('NaN') + + +def _empty_prosody_distance_metrics(): + return {key: float('NaN') for key in PROSODY_DISTANCE_KEYS} + + +def compute_acoustic_prosody_metrics( + gt_audio_path, + pred_audio_path, + text, +): + """Compute reference-based pitch, intensity, and speech-rate distances.""" + if gt_audio_path is None or pred_audio_path is None: + return _empty_prosody_distance_metrics() + + try: + metrics = compute_prosody_distances( + gt_audio_path=gt_audio_path, + pred_audio_path=pred_audio_path, + text=text, + ).to_dict() + return {key: metrics[key] for key in PROSODY_DISTANCE_KEYS} + except Exception as e: + logging.warning(f"Could not compute acoustic prosody distances for {gt_audio_path} and {pred_audio_path}: {e}") + return _empty_prosody_distance_metrics() + + def classify_eou_batched( eou_classifier: EoUClassifier, items: list[tuple[Union[str, np.ndarray], str]], batch_size: int = 32 ) -> list[EoUClassification]: @@ -359,6 +436,8 @@ def evaluate_dir( asr_model_type="nemo", with_utmosv2=True, strip_text_annotations_for_metrics=False, + with_prosody_metrics=False, + prosody_model_size="small", asr_batch_size=32, eou_batch_size=32, device="cuda", @@ -396,12 +475,15 @@ def evaluate_dir( asr_model_name=asr_model_name, asr_model_type=asr_model_type, device=device, + with_prosody_metrics=with_prosody_metrics, + prosody_model_size=prosody_model_size, ) asr_model = models['asr_model'] feature_extractor = models['feature_extractor'] speaker_verification_model = models['sv_model'] speaker_verification_model_alternate = models['sv_model_alternate'] + emotion_model = models['emotion_model'] # 3. EoU classifier (support for English only) if language == "en": @@ -481,6 +563,16 @@ def evaluate_dir( detailed_cer = word_error_rate_detail(hypotheses=[pred_text], references=[gt_text], use_cer=True) detailed_wer = word_error_rate_detail(hypotheses=[pred_text], references=[gt_text], use_cer=False) + cer_pred_gt_audio = ( + word_error_rate_detail(hypotheses=[pred_text], references=[gt_audio_text], use_cer=True)[0] + if gt_audio_text is not None + else float('NaN') + ) + wer_pred_gt_audio = ( + word_error_rate_detail(hypotheses=[pred_text], references=[gt_audio_text], use_cer=False)[0] + if gt_audio_text is not None + else float('NaN') + ) # Japanese: additional reading-based CER on Katakana (pyopenjtalk g2p), robust to # kanji/kana spelling differences between reference and ASR hypothesis. @@ -493,6 +585,20 @@ def evaluate_dir( 0 ] + pred_gt_esim = float('NaN') + pred_gt_ems = float('NaN') + prosody_distance_metrics = _empty_prosody_distance_metrics() + if with_prosody_metrics: + pred_gt_esim, pred_gt_ems = compute_emotion_pair_metrics( + emotion_model, + gt_audio_filepath, + pred_audio_filepath, + ) + prosody_distance_metrics = compute_acoustic_prosody_metrics( + gt_audio_path=gt_audio_filepath, + pred_audio_path=pred_audio_filepath, + text=gt_text, + ) logging.info(f"{ridx} GT Text: {gt_text}") logging.info(f"{ridx} Pr Text: {pred_text}") # Format cer and wer to 2 decimal places @@ -588,7 +694,9 @@ def evaluate_dir( 'detailed_cer': detailed_cer, 'detailed_wer': detailed_wer, 'cer': detailed_cer[0], + 'cer_pred_gt_audio': cer_pred_gt_audio, 'wer': detailed_wer[0], + 'wer_pred_gt_audio': wer_pred_gt_audio, 'katakana_cer': katakana_cer, 'gt_katakana': gt_katakana, 'pred_katakana': pred_katakana, @@ -608,7 +716,10 @@ def evaluate_dir( 'total_gen_audio_seconds': file_duration, 'predicted_codes_path': codes_file_lists[ridx] if has_codes else None, } - + if with_prosody_metrics: + metric_row['pred_gt_esim'] = pred_gt_esim + metric_row['pred_gt_ems'] = pred_gt_ems + metric_row.update(prosody_distance_metrics) filewise_metrics.append(metric_row) return filewise_metrics @@ -626,6 +737,8 @@ def evaluate( strip_text_annotations_for_metrics=False, with_fcd=True, codec_model_path=None, + with_prosody_metrics=False, + prosody_model_size="small", asr_batch_size=32, eou_batch_size=32, device="cuda", @@ -663,6 +776,8 @@ def evaluate( asr_model_type=asr_model_type, with_utmosv2=with_utmosv2, strip_text_annotations_for_metrics=strip_text_annotations_for_metrics, + with_prosody_metrics=with_prosody_metrics, + prosody_model_size=prosody_model_size, asr_batch_size=asr_batch_size, eou_batch_size=eou_batch_size, device=device, @@ -717,6 +832,18 @@ def compute_fcd(gt_audio_paths, predicted_codes_paths, codec_model_path, device= return fcd +def _mean_finite_metric(filewise_metrics, key: str) -> float: + values = [] + for metrics in filewise_metrics: + try: + value = float(metrics[key]) + except (KeyError, TypeError, ValueError): + continue + if np.isfinite(value): + values.append(value) + return float('nan') if not values else float(np.mean(values)) + + def compute_global_metrics( filewise_metrics, gt_audio_paths=None, @@ -765,10 +892,25 @@ def compute_global_metrics( sum(m['pred_context_ssim_alternate'] for m in filewise_metrics) / n ) avg_metrics['ssim_gt_context_avg_alternate'] = sum(m['gt_context_ssim_alternate'] for m in filewise_metrics) / n + if 'pred_gt_esim' in filewise_metrics[0]: + avg_metrics['esim_pred_gt_avg'] = sum(m['pred_gt_esim'] for m in filewise_metrics) / n + avg_metrics['ems_pred_gt_avg'] = sum(m['pred_gt_ems'] for m in filewise_metrics) / n + if 'pitch_distance' in filewise_metrics[0]: + avg_metrics['pitch_distance_avg'] = _mean_finite_metric(filewise_metrics, 'pitch_distance') + avg_metrics['intensity_distance_avg'] = _mean_finite_metric(filewise_metrics, 'intensity_distance') + avg_metrics['speech_rate_distance_avg'] = _mean_finite_metric(filewise_metrics, 'speech_rate_distance') # Cumulative WER/CER on ground-truth audio transcriptions (if available) gt_audio_texts = [m['gt_audio_text'] for m in filewise_metrics] + avg_metrics['cer_pred_gt_audio_filewise_avg'] = _mean_finite_metric(filewise_metrics, 'cer_pred_gt_audio') + avg_metrics['wer_pred_gt_audio_filewise_avg'] = _mean_finite_metric(filewise_metrics, 'wer_pred_gt_audio') if None not in gt_audio_texts: + avg_metrics['cer_pred_gt_audio_cumulative'] = word_error_rate_detail( + hypotheses=pred_texts, references=gt_audio_texts, use_cer=True + )[0] + avg_metrics['wer_pred_gt_audio_cumulative'] = word_error_rate_detail( + hypotheses=pred_texts, references=gt_audio_texts, use_cer=False + )[0] avg_metrics['cer_gt_audio_cumulative'] = word_error_rate_detail( hypotheses=gt_audio_texts, references=gt_texts, use_cer=True )[0] @@ -776,6 +918,8 @@ def compute_global_metrics( hypotheses=gt_audio_texts, references=gt_texts, use_cer=False )[0] else: + avg_metrics['cer_pred_gt_audio_cumulative'] = float('NaN') + avg_metrics['wer_pred_gt_audio_cumulative'] = float('NaN') avg_metrics['cer_gt_audio_cumulative'] = float('NaN') avg_metrics['wer_gt_audio_cumulative'] = float('NaN') logging.warning( @@ -818,6 +962,12 @@ def main(): parser.add_argument('--generated_audio_dir', type=str, default=None) parser.add_argument('--language', type=str, default="en") parser.add_argument('--evalset', type=str, default=None) + parser.add_argument( + '--with_prosody_metrics', + action='store_true', + help='Compute ESIM/EMS and pitch, intensity, and speech-rate distance metrics.', + ) + parser.add_argument('--prosody_model_size', type=str, default="small", choices=["small", "large"]) parser.add_argument( '--strip_text_annotations_for_metrics', action='store_true', @@ -838,7 +988,9 @@ def main(): args.language, sv_model_type="wavlm", asr_model_name="nvidia/parakeet-ctc-0.6b", + with_prosody_metrics=args.with_prosody_metrics, strip_text_annotations_for_metrics=args.strip_text_annotations_for_metrics, + prosody_model_size=args.prosody_model_size, ) diff --git a/nemo/collections/tts/modules/magpietts_inference/evaluation.py b/nemo/collections/tts/modules/magpietts_inference/evaluation.py index a4a7da8f7013..8ca3d8c3ccc0 100644 --- a/nemo/collections/tts/modules/magpietts_inference/evaluation.py +++ b/nemo/collections/tts/modules/magpietts_inference/evaluation.py @@ -43,6 +43,9 @@ class EvaluationConfig: with_utmosv2: Whether to compute UTMOSv2 (Mean Opinion Score) metrics. with_fcd: Whether to compute Frechet Codec Distance metric. codec_model_path: Path to the audio codec model. If None, will skip computing Frechet Codec Distance metric. + with_prosody_metrics: Whether to compute ESIM/EMS plus pitch, + intensity, and speech-rate distance metrics. + prosody_model_size: Emotion encoder size ("small" or "large"). strip_text_annotations_for_metrics: Whether to strip annotation/control markers from reference and ASR hypothesis text before text metrics. device: Device to use for running models used during evaluation. """ @@ -55,6 +58,8 @@ class EvaluationConfig: with_utmosv2: bool = True with_fcd: bool = True codec_model_path: str = None + with_prosody_metrics: bool = False + prosody_model_size: str = "small" strip_text_annotations_for_metrics: bool = False device: str = "cuda" asr_batch_size: int = 32 @@ -99,6 +104,8 @@ def evaluate_generated_audio_dir( with_utmosv2=config.with_utmosv2, with_fcd=config.with_fcd, codec_model_path=config.codec_model_path, + with_prosody_metrics=config.with_prosody_metrics, + prosody_model_size=config.prosody_model_size, strip_text_annotations_for_metrics=config.strip_text_annotations_for_metrics, device=config.device, eou_model_name=config.eou_model_name, diff --git a/nemo/collections/tts/modules/magpietts_inference/utils.py b/nemo/collections/tts/modules/magpietts_inference/utils.py index 078c9c8d84c0..8a9a5ff05e6f 100644 --- a/nemo/collections/tts/modules/magpietts_inference/utils.py +++ b/nemo/collections/tts/modules/magpietts_inference/utils.py @@ -634,12 +634,21 @@ def append_metrics_to_csv(csv_path: str, checkpoint_name: str, dataset: str, met metrics.get('wer_filewise_avg', ''), metrics.get('cer_cumulative', ''), metrics.get('wer_cumulative', ''), + metrics.get('cer_pred_gt_audio_filewise_avg', ''), + metrics.get('cer_pred_gt_audio_cumulative', ''), + metrics.get('wer_pred_gt_audio_filewise_avg', ''), + metrics.get('wer_pred_gt_audio_cumulative', ''), metrics.get('ssim_pred_gt_avg', ''), metrics.get('ssim_pred_context_avg', ''), metrics.get('ssim_gt_context_avg', ''), metrics.get('ssim_pred_gt_avg_alternate', ''), metrics.get('ssim_pred_context_avg_alternate', ''), metrics.get('ssim_gt_context_avg_alternate', ''), + metrics.get('esim_pred_gt_avg', ''), + metrics.get('ems_pred_gt_avg', ''), + metrics.get('pitch_distance_avg', ''), + metrics.get('intensity_distance_avg', ''), + metrics.get('speech_rate_distance_avg', ''), metrics.get('cer_gt_audio_cumulative', ''), metrics.get('wer_gt_audio_cumulative', ''), metrics.get('utmosv2_avg', ''), @@ -753,10 +762,17 @@ def turn_sort_key(r): turns = sorted(turns, key=turn_sort_key) cer_turns = [r.get("cer") for r in turns] + cer_pred_gt_audio_turns = [r.get("cer_pred_gt_audio") for r in turns] wer_turns = [r.get("wer") for r in turns] + wer_pred_gt_audio_turns = [r.get("wer_pred_gt_audio") for r in turns] pred_context_ssim_turns = [r.get("pred_context_ssim") for r in turns] pred_gt_ssim_turns = [r.get("pred_gt_ssim") for r in turns] gt_context_ssim_turns = [r.get("gt_context_ssim") for r in turns] + pred_gt_esim_turns = [r.get("pred_gt_esim") for r in turns] + pred_gt_ems_turns = [r.get("pred_gt_ems") for r in turns] + pitch_distance_turns = [r.get("pitch_distance") for r in turns] + intensity_distance_turns = [r.get("intensity_distance") for r in turns] + speech_rate_distance_turns = [r.get("speech_rate_distance") for r in turns] utmosv2_turns = [r.get("utmosv2") for r in turns] eou_type_turns = [r.get("eou_type") for r in turns] eou_trailing_duration_turns = [r.get("eou_trailing_duration") for r in turns] @@ -773,20 +789,34 @@ def turn_sort_key(r): "num_turns": len(turns), # Sample-level averages over all turns. "cer": _mean_finite(cer_turns), + "cer_pred_gt_audio": _mean_finite(cer_pred_gt_audio_turns), "wer": _mean_finite(wer_turns), + "wer_pred_gt_audio": _mean_finite(wer_pred_gt_audio_turns), "pred_context_ssim": _mean_finite(pred_context_ssim_turns), "pred_gt_ssim": _mean_finite(pred_gt_ssim_turns), "gt_context_ssim": _mean_finite(gt_context_ssim_turns), + "pred_gt_esim": _mean_finite(pred_gt_esim_turns), + "pred_gt_ems": _mean_finite(pred_gt_ems_turns), + "pitch_distance": _mean_finite(pitch_distance_turns), + "intensity_distance": _mean_finite(intensity_distance_turns), + "speech_rate_distance": _mean_finite(speech_rate_distance_turns), "utmosv2": _mean_finite(utmosv2_turns), "eou_trailing_duration": _mean_finite(eou_trailing_duration_turns), "eou_trail_rms_ratio": _mean_finite(eou_trail_rms_ratio_turns), # Turn-by-turn values, old-script style. "turn_ids": [r.get("turn_id", i) for i, r in enumerate(turns)], "cer_turns": cer_turns, + "cer_pred_gt_audio_turns": cer_pred_gt_audio_turns, "wer_turns": wer_turns, + "wer_pred_gt_audio_turns": wer_pred_gt_audio_turns, "pred_context_ssim_turns": pred_context_ssim_turns, "pred_gt_ssim_turns": pred_gt_ssim_turns, "gt_context_ssim_turns": gt_context_ssim_turns, + "pred_gt_esim_turns": pred_gt_esim_turns, + "pred_gt_ems_turns": pred_gt_ems_turns, + "pitch_distance_turns": pitch_distance_turns, + "intensity_distance_turns": intensity_distance_turns, + "speech_rate_distance_turns": speech_rate_distance_turns, "utmosv2_turns": utmosv2_turns, "eou_type_turns": eou_type_turns, "eou_trailing_duration_turns": eou_trailing_duration_turns, @@ -820,19 +850,33 @@ def _write_grouped_multiturn_filewise_metrics_csv(csv_path: str, grouped_rows: l "rank", "num_turns", "cer", + "cer_pred_gt_audio", "wer", + "wer_pred_gt_audio", "pred_context_ssim", "pred_gt_ssim", "gt_context_ssim", + "pred_gt_esim", + "pred_gt_ems", + "pitch_distance", + "intensity_distance", + "speech_rate_distance", "utmosv2", "eou_trailing_duration", "eou_trail_rms_ratio", "turn_ids", "cer_turns", + "cer_pred_gt_audio_turns", "wer_turns", + "wer_pred_gt_audio_turns", "pred_context_ssim_turns", "pred_gt_ssim_turns", "gt_context_ssim_turns", + "pred_gt_esim_turns", + "pred_gt_ems_turns", + "pitch_distance_turns", + "intensity_distance_turns", + "speech_rate_distance_turns", "utmosv2_turns", "eou_type_turns", "eou_trailing_duration_turns", @@ -1223,6 +1267,12 @@ def _add_common_args(parser: argparse.ArgumentParser) -> None: eval_group.add_argument('--num_repeats', type=int, default=1) eval_group.add_argument('--confidence_level', type=float, default=0.95) eval_group.add_argument('--disable_utmosv2', action='store_true') + eval_group.add_argument( + '--with_prosody_metrics', + action='store_true', + help='Compute ESIM/EMS and pitch, intensity, and speech-rate distance metrics.', + ) + eval_group.add_argument('--prosody_model_size', type=str, default="small", choices=["small", "large"]) eval_group.add_argument( '--strip_text_annotations_for_metrics', action='store_true', diff --git a/tests/collections/common/test_lhotse_tts_filters.py b/tests/collections/common/test_lhotse_tts_filters.py index 90a08f32d55e..cabba231fa6b 100644 --- a/tests/collections/common/test_lhotse_tts_filters.py +++ b/tests/collections/common/test_lhotse_tts_filters.py @@ -146,7 +146,7 @@ def test_cut_validation_status_filter(cut_example): def test_cut_speaker_filter_by_speaker(cut_example): f = SpeakerFilter( - excluded_speaker_ids=["| Language:en Dataset:nvyt2505 Speaker:Zdud2gXLTXY_SPEAKER_02 |"], + excluded_speaker_ids=["Zdud2gXLTXY_SPEAKER_02"], speaker_fields=["speaker"], ) assert f(cut_example) == False diff --git a/tests/collections/tts/metrics/test_prosody.py b/tests/collections/tts/metrics/test_prosody.py new file mode 100644 index 000000000000..66096ee9b68e --- /dev/null +++ b/tests/collections/tts/metrics/test_prosody.py @@ -0,0 +1,65 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import numpy as np +import pytest +import soundfile as sf + +from nemo.collections.tts.metrics.prosody import compute_prosody_distances + +_SAMPLE_RATE = 16000 +_TEXT = "hello world" + + +def _write_sine(path, duration_sec: float, frequency_hz: float = 220.0, amplitude: float = 0.2) -> None: + sample_count = int(round(_SAMPLE_RATE * duration_sec)) + time = np.arange(sample_count, dtype=np.float32) / _SAMPLE_RATE + audio = amplitude * np.sin(2.0 * np.pi * frequency_hz * time) + sf.write(path, audio.astype(np.float32), _SAMPLE_RATE) + + +@pytest.mark.unit +def test_prosody_distances_are_zero_for_identical_audio(tmp_path): + gt_path = tmp_path / "gt.wav" + pred_path = tmp_path / "pred.wav" + _write_sine(gt_path, duration_sec=1.0) + _write_sine(pred_path, duration_sec=1.0) + + metrics = compute_prosody_distances( + gt_audio_path=str(gt_path), + pred_audio_path=str(pred_path), + text=_TEXT, + ) + + assert metrics.pitch_distance == pytest.approx(0.0, abs=1.0e-6) + assert metrics.intensity_distance == pytest.approx(0.0, abs=1.0e-6) + assert metrics.speech_rate_distance == pytest.approx(0.0, abs=1.0e-6) + + +@pytest.mark.unit +def test_speech_rate_distance_tracks_duration_difference(tmp_path): + gt_path = tmp_path / "gt.wav" + pred_path = tmp_path / "pred.wav" + _write_sine(gt_path, duration_sec=1.0) + _write_sine(pred_path, duration_sec=2.0) + + metrics = compute_prosody_distances( + gt_audio_path=str(gt_path), + pred_audio_path=str(pred_path), + text=_TEXT, + ) + + assert np.isfinite(metrics.pitch_distance) + assert np.isfinite(metrics.intensity_distance) + assert metrics.speech_rate_distance == pytest.approx(5.0, abs=1.0e-6)