From 4cb90efd249bbceca2d6a7cbc6716be1e09c2322 Mon Sep 17 00:00:00 2001 From: Edresson Casanova Date: Sun, 2 Aug 2026 09:36:56 -0700 Subject: [PATCH 01/10] Support TTS speaker id format on SpeakerFilter class Signed-off-by: Edresson Casanova --- nemo/collections/common/data/lhotse/sampling.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/nemo/collections/common/data/lhotse/sampling.py b/nemo/collections/common/data/lhotse/sampling.py index e73039d252f9..3c5293fbac09 100644 --- a/nemo/collections/common/data/lhotse/sampling.py +++ b/nemo/collections/common/data/lhotse/sampling.py @@ -342,6 +342,15 @@ 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 From 1b39526282ce7c1fd9f6c6bf0e7aa8c94c877b79 Mon Sep 17 00:00:00 2001 From: Edresson Casanova Date: Tue, 4 Aug 2026 17:20:20 -0300 Subject: [PATCH 02/10] Add prosody evaluation metrics Signed-off-by: Edresson Casanova --- examples/tts/magpietts_inference.py | 12 +- .../tts/metrics/emotion_encoder.py | 1293 +++++++++++++++++ nemo/collections/tts/metrics/prosody.py | 499 +++++++ .../evaluate_generated_audio.py | 156 +- .../modules/magpietts_inference/evaluation.py | 13 + .../tts/modules/magpietts_inference/utils.py | 43 + tests/collections/tts/metrics/test_prosody.py | 76 + 7 files changed, 2090 insertions(+), 2 deletions(-) create mode 100644 nemo/collections/tts/metrics/emotion_encoder.py create mode 100644 nemo/collections/tts/metrics/prosody.py create mode 100644 tests/collections/tts/metrics/test_prosody.py diff --git a/examples/tts/magpietts_inference.py b/examples/tts/magpietts_inference.py index d93a6c9ba2fb..0d8f556c4756 100644 --- a/examples/tts/magpietts_inference.py +++ b/examples/tts/magpietts_inference.py @@ -169,7 +169,9 @@ def run_inference_and_evaluation( "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," "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 +305,10 @@ 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, + prosody_embedding_type=eval_config.prosody_embedding_type, + prosody_cache_dir=eval_config.prosody_cache_dir, 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,10 @@ 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, + prosody_embedding_type=args.prosody_embedding_type, + prosody_cache_dir=args.prosody_cache_dir, 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/tts/metrics/emotion_encoder.py b/nemo/collections/tts/metrics/emotion_encoder.py new file mode 100644 index 000000000000..d2263420832d --- /dev/null +++ b/nemo/collections/tts/metrics/emotion_encoder.py @@ -0,0 +1,1293 @@ +# 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 + +import argparse +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() + + +# ============================================================================= +# CLI utilities +# ============================================================================= + + +def _tensor_info(tensor: torch.Tensor) -> dict[str, Any]: + return { + "shape": list(tensor.shape), + "dtype": str(tensor.dtype), + "device": str(tensor.device), + } + + +def _parse_labels(labels: Optional[str]) -> Optional[list[str]]: + if labels is None or labels.strip() == "": + return None + + return [item.strip() for item in labels.split(",") if item.strip()] + + +def main() -> None: + parser = argparse.ArgumentParser(description="LAION Empathic Insight Voice embeddings and similarity.") + parser.add_argument( + "--audio", + type=str, + required=True, + help="Input audio path.", + ) + parser.add_argument( + "--audio-b", + type=str, + default=None, + help="Optional second audio path for similarity.", + ) + parser.add_argument( + "--size", + type=str, + default="small", + choices=["small", "large"], + help="Model size.", + ) + parser.add_argument( + "--device", + type=str, + default="cuda", + help="Device for Whisper encoder.", + ) + parser.add_argument( + "--mlp-device", + type=str, + default=None, + help="Device for MLP classifier heads. Defaults to --device.", + ) + parser.add_argument( + "--cache-dir", + type=str, + default=None, + help="Optional Hugging Face cache directory.", + ) + parser.add_argument( + "--embedding-type", + type=str, + default="head_concat", + choices=["head_concat", "head_mean", "score_vector"], + help="Emotion embedding type.", + ) + parser.add_argument( + "--labels", + type=str, + default=None, + help=( + "Comma-separated labels to use. " + "Example: anger,sadness,arousal. " + "If omitted, primary emotion labels are used for similarity." + ), + ) + parser.add_argument( + "--include-auxiliary", + action="store_true", + help="Include auxiliary similarity labels when --labels is omitted.", + ) + parser.add_argument( + "--load-all-classifiers", + action="store_true", + help="Eagerly load all known classifiers at startup.", + ) + parser.add_argument( + "--top-k", + type=int, + default=5, + help="Number of ranked emotions to return.", + ) + + args = parser.parse_args() + + labels = _parse_labels(args.labels) + + model = EmpathicInsightVoice.from_pretrained( + size=args.size, + device=args.device, + mlp_device=args.mlp_device, + cache_dir=args.cache_dir, + cache_classifiers=True, + load_all_classifiers=args.load_all_classifiers, + top_k_emotions=args.top_k, + ) + + result = model.compute( + audio_path=args.audio, + labels=labels, + return_embedding=True, + embedding_type=args.embedding_type, + return_raw_scores=True, + include_auxiliary_for_embedding=args.include_auxiliary, + ) + + printable: dict[str, Any] = { + "audio_path": result["audio_path"], + "model_size": result["model_size"], + "top_emotion": result["top_emotion"], + "embedding_type": result["embedding_type"], + "embedding": _tensor_info(result["embedding"]), + "emotions": result["emotions"], + "raw_scores": result["raw_scores"], + } + + if args.audio_b is not None: + printable["audio_b"] = args.audio_b + printable["similarity"] = model.emotion_similarity( + audio_path_a=args.audio, + audio_path_b=args.audio_b, + labels=labels, + embedding_type=args.embedding_type, + include_auxiliary=args.include_auxiliary, + ) + + result = model.compare_emotion_pair( + audio_path_a=args.audio, + audio_path_b=args.audio_b, + embedding_type="head_concat", + ) + + result_score_vector = model.compare_emotion_pair( + audio_path_a=args.audio, + audio_path_b=args.audio_b, + embedding_type="score_vector", + ) + + result_score_mean = model.compare_emotion_pair( + audio_path_a=args.audio, + audio_path_b=args.audio_b, + embedding_type="head_mean", + ) + + print("embedding_type=head_concat") + print(result["audio_a_top_emotion"]) + print(result["audio_b_top_emotion"]) + print(result["top_emotion_match"]) + print(result["emotion_similarity"]) + + print("embedding_type=score_vector") + print(result_score_vector["audio_a_top_emotion"]) + print(result_score_vector["audio_b_top_emotion"]) + print(result_score_vector["top_emotion_match"]) + print(result_score_vector["emotion_similarity"]) + + print("embedding_type=head_mean") + print(result_score_mean["audio_a_top_emotion"]) + print(result_score_mean["audio_b_top_emotion"]) + print(result_score_mean["top_emotion_match"]) + print(result_score_mean["emotion_similarity"]) + + +if __name__ == "__main__": + main() diff --git a/nemo/collections/tts/metrics/prosody.py b/nemo/collections/tts/metrics/prosody.py new file mode 100644 index 000000000000..92e2aeb47d6b --- /dev/null +++ b/nemo/collections/tts/metrics/prosody.py @@ -0,0 +1,499 @@ +# 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 +import re +from dataclasses import dataclass +from typing import Any, Literal, Optional + +import librosa +import numpy as np + +try: + from numba import njit +except Exception: # pragma: no cover - numba is an optional speedup. + njit = None + + +SpeechRateCharMode = Literal["nonspace", "all", "alnum"] +F0Method = Literal["pyin", "yin", "none"] +F0Normalization = Literal["gt_median", "utterance_median", "none"] +EnergyNormalization = Literal["zscore", "none"] + +_ALNUM_CHAR_RE = re.compile(r"[A-Za-z0-9]") + + +@dataclass(frozen=True) +class ProsodyDistanceConfig: + """Configuration for reference-based acoustic prosody distance metrics. + + The default values are tuned for corpus-level TTS evaluation: pYIN is used + for more stable F0 contours, F0 is converted to semitones relative to the + reference median, intensity uses log-RMS z-scores, and contours are reduced + before DTW to keep evaluation bounded. + """ + + sample_rate: int = 16000 + res_type: str = "soxr_hq" + frame_shift_ms: float = 20.0 + frame_length_ms: float = 64.0 + fmin: float = 55.0 + fmax: float = 450.0 + f0_method: F0Method = "pyin" + yin_silence_db_below_peak: float = 35.0 + pyin_n_thresholds: int = 24 + pyin_beta_a: float = 2.0 + pyin_beta_b: float = 18.0 + pyin_boltzmann_parameter: float = 2.0 + pyin_resolution: float = 0.25 + pyin_max_transition_rate: float = 12.0 + pyin_switch_prob: float = 0.01 + pyin_no_trough_prob: float = 0.01 + pyin_center: bool = True + pyin_pad_mode: str = "constant" + max_dtw_frames: int = 1000 + dtw_band_ratio: float = 0.05 + f0_nan_penalty: float = 6.0 + f0_normalization: F0Normalization = "gt_median" + intensity_normalization: EnergyNormalization = "zscore" + speech_rate_char_mode: SpeechRateCharMode = "nonspace" + min_voiced_frames: int = 5 + + +@dataclass(frozen=True) +class ProsodyDistanceResult: + """Per-pair acoustic prosody distance metrics.""" + + pitch_distance: float + intensity_distance: float + speech_rate_distance: float + gt_duration_sec: float + pred_duration_sec: float + gt_speech_rate_cps: float + pred_speech_rate_cps: float + gt_char_count: int + + def to_dict(self) -> dict[str, float | int]: + """Return a JSON-serializable dictionary.""" + return { + "pitch_distance": self.pitch_distance, + "intensity_distance": self.intensity_distance, + "speech_rate_distance": self.speech_rate_distance, + "gt_duration_sec": self.gt_duration_sec, + "pred_duration_sec": self.pred_duration_sec, + "gt_speech_rate_cps": self.gt_speech_rate_cps, + "pred_speech_rate_cps": self.pred_speech_rate_cps, + "gt_char_count": self.gt_char_count, + } + + +if njit is not None: + + @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 + + if band_radius < 0: + j_start = 1 + j_end = m + 1 + else: + 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) + +else: + _dtw_distance_1d_numba = None + + +def compute_prosody_distances( + gt_audio_path: str, + pred_audio_path: str, + text: Any, + config: Optional[ProsodyDistanceConfig] = None, +) -> 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. + config: Optional prosody distance configuration. + + Returns: + ProsodyDistanceResult with pitch, intensity, and speech-rate distances. + """ + cfg = config or ProsodyDistanceConfig() + gt_audio, sr, gt_duration = _load_audio(gt_audio_path, cfg) + pred_audio, _, pred_duration = _load_audio(pred_audio_path, cfg) + + hop_length, frame_length = _frame_params(sr, cfg) + 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) + + pitch_distance = float("nan") + if cfg.f0_method != "none": + gt_f0 = _compute_f0( + gt_audio, sr=sr, frame_length=frame_length, hop_length=hop_length, log_energy=gt_log_energy, cfg=cfg + ) + pred_f0 = _compute_f0( + pred_audio, + sr=sr, + frame_length=frame_length, + hop_length=hop_length, + log_energy=pred_log_energy, + cfg=cfg, + ) + if np.isfinite(gt_f0).sum() >= cfg.min_voiced_frames and np.isfinite(pred_f0).sum() >= cfg.min_voiced_frames: + gt_pitch, pred_pitch = _prepare_f0_for_metric(gt_f0, pred_f0, cfg) + pitch_distance = _dtw_distance_1d( + _maybe_reduce_for_dtw(gt_pitch, cfg.max_dtw_frames), + _maybe_reduce_for_dtw(pred_pitch, cfg.max_dtw_frames), + nan_penalty=cfg.f0_nan_penalty, + band_ratio=cfg.dtw_band_ratio, + ) + + gt_intensity, pred_intensity = _prepare_intensity_for_metric(gt_log_energy, pred_log_energy, cfg) + intensity_distance = _dtw_distance_1d( + _maybe_reduce_for_dtw(gt_intensity, cfg.max_dtw_frames), + _maybe_reduce_for_dtw(pred_intensity, cfg.max_dtw_frames), + nan_penalty=0.0, + band_ratio=cfg.dtw_band_ratio, + ) + + gt_char_count = _char_count(text, cfg.speech_rate_char_mode) + 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), + gt_duration_sec=_safe_float(gt_duration), + pred_duration_sec=_safe_float(pred_duration), + gt_speech_rate_cps=_safe_float(gt_speech_rate), + pred_speech_rate_cps=_safe_float(pred_speech_rate), + gt_char_count=gt_char_count, + ) + + +def _load_audio(path: str, cfg: ProsodyDistanceConfig) -> 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=cfg.sample_rate, mono=True, res_type=cfg.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, cfg: ProsodyDistanceConfig) -> tuple[int, int]: + hop_length = max(1, int(round(sr * cfg.frame_shift_ms / 1000.0))) + frame_length = max(hop_length * 2, int(round(sr * cfg.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, + log_energy: np.ndarray, + cfg: ProsodyDistanceConfig, +) -> np.ndarray: + if cfg.f0_method == "pyin": + f0, voiced_flag, _ = librosa.pyin( + y=np.asarray(audio, dtype=np.float64), + sr=sr, + fmin=cfg.fmin, + fmax=cfg.fmax, + frame_length=frame_length, + hop_length=hop_length, + center=cfg.pyin_center, + pad_mode=cfg.pyin_pad_mode, + n_thresholds=cfg.pyin_n_thresholds, + beta_parameters=(cfg.pyin_beta_a, cfg.pyin_beta_b), + boltzmann_parameter=cfg.pyin_boltzmann_parameter, + resolution=cfg.pyin_resolution, + max_transition_rate=cfg.pyin_max_transition_rate, + switch_prob=cfg.pyin_switch_prob, + no_trough_prob=cfg.pyin_no_trough_prob, + fill_na=np.nan, + ) + f0 = np.asarray(f0, dtype=np.float64) + if voiced_flag is not None: + 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 + + if cfg.f0_method == "yin": + f0 = librosa.yin( + audio, + sr=sr, + fmin=cfg.fmin, + fmax=cfg.fmax, + frame_length=frame_length, + hop_length=hop_length, + center=True, + ) + f0 = np.asarray(f0, dtype=np.float64) + f0[(f0 < cfg.fmin) | (f0 > cfg.fmax)] = np.nan + return _mask_yin_silence(f0, log_energy, cfg) + + if cfg.f0_method == "none": + return np.asarray([], dtype=np.float64) + + raise ValueError(f"Unsupported f0_method={cfg.f0_method!r}") + + +def _mask_yin_silence(f0: np.ndarray, log_energy: np.ndarray, cfg: ProsodyDistanceConfig) -> np.ndarray: + if len(log_energy) == 0: + return f0 + + min_len = min(len(f0), len(log_energy)) + f0 = f0[:min_len].copy() + energy = log_energy[:min_len] + rms_db = 20.0 * energy / math.log(10.0) + finite = np.isfinite(rms_db) + if finite.any(): + peak_db = float(np.max(rms_db[finite])) + f0[rms_db < peak_db - cfg.yin_silence_db_below_peak] = np.nan + return f0 + + +def _prepare_f0_for_metric( + gt_f0_hz: np.ndarray, + pred_f0_hz: np.ndarray, + cfg: ProsodyDistanceConfig, +) -> tuple[np.ndarray, np.ndarray]: + if cfg.f0_normalization == "none": + return gt_f0_hz, pred_f0_hz + + gt_median = float(np.nanmedian(gt_f0_hz)) if np.isfinite(gt_f0_hz).any() else float("nan") + pred_median = float(np.nanmedian(pred_f0_hz)) if np.isfinite(pred_f0_hz).any() else float("nan") + + if cfg.f0_normalization == "gt_median": + return _hz_to_semitones(gt_f0_hz, gt_median), _hz_to_semitones(pred_f0_hz, gt_median) + if cfg.f0_normalization == "utterance_median": + return _hz_to_semitones(gt_f0_hz, gt_median), _hz_to_semitones(pred_f0_hz, pred_median) + raise ValueError(f"Unsupported f0_normalization={cfg.f0_normalization!r}") + + +def _prepare_intensity_for_metric( + gt_log_energy: np.ndarray, + pred_log_energy: np.ndarray, + cfg: ProsodyDistanceConfig, +) -> tuple[np.ndarray, np.ndarray]: + if cfg.intensity_normalization == "none": + return gt_log_energy, pred_log_energy + if cfg.intensity_normalization == "zscore": + return _zscore(gt_log_energy), _zscore(pred_log_energy) + raise ValueError(f"Unsupported intensity_normalization={cfg.intensity_normalization!r}") + + +def _dtw_distance_1d(x: np.ndarray, y: np.ndarray, nan_penalty: float, band_ratio: 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") + + if band_ratio is None or band_ratio < 0: + band_radius = -1 + else: + band_radius = max(abs(len(x) - len(y)), int(math.ceil(float(band_ratio) * max(len(x), len(y))))) + + if _dtw_distance_1d_numba is not None: + try: + return _safe_float(_dtw_distance_1d_numba(x, y, float(nan_penalty), int(band_radius))) + except Exception: + pass + return _dtw_distance_1d_python(x, y, float(nan_penalty), int(band_radius)) + + +def _dtw_distance_1d_python(x: np.ndarray, y: np.ndarray, nan_penalty: float, band_radius: int) -> float: + n = len(x) + m = len(y) + if n == 0 or m == 0: + return float("nan") + + prev = np.full(m + 1, np.inf, dtype=np.float64) + curr = np.full(m + 1, np.inf, dtype=np.float64) + prev[0] = 0.0 + + for i in range(1, n + 1): + curr.fill(np.inf) + if band_radius < 0: + j_start = 1 + j_end = m + 1 + else: + j_start = max(1, i - band_radius) + j_end = min(m, i + band_radius) + 1 + + for j in range(j_start, j_end): + cost = _frame_distance(x[i - 1], y[j - 1], nan_penalty) + curr[j] = cost + min(prev[j], curr[j - 1], prev[j - 1]) + prev, curr = curr, prev + + total = prev[m] + if not np.isfinite(total): + return float("nan") + return float(total / (n + m)) + + +def _frame_distance(x: float, y: float, nan_penalty: float) -> float: + x_nan = np.isnan(x) + y_nan = np.isnan(y) + if x_nan and y_nan: + return 0.0 + if x_nan or y_nan: + return nan_penalty + return abs(float(x) - float(y)) + + +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, max_frames: int) -> np.ndarray: + values = np.asarray(values, dtype=np.float64) + if max_frames is None or max_frames <= 0 or len(values) <= max_frames: + return values + return _resample_1d_preserve_nans(values, int(max_frames)) + + +def _resample_1d_preserve_nans(values: np.ndarray, target_len: int) -> np.ndarray: + values = np.asarray(values, dtype=np.float64) + if target_len <= 0: + raise ValueError("target_len must be positive") + 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, mode: SpeechRateCharMode) -> int: + if text is None: + return 0 + text = str(text) + if mode == "nonspace": + return sum(1 for ch in text if not ch.isspace()) + if mode == "alnum": + return len(_ALNUM_CHAR_RE.findall(text)) + if mode == "all": + return len(text) + raise ValueError(f"Unsupported speech_rate_char_mode={mode!r}") + + +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..fd02aeca66d0 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 ProsodyDistanceConfig, 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. @@ -100,6 +107,9 @@ def strip_text_annotations_from_text(text: str) -> str: 'cer', 'wer', 'pred_context_ssim', + 'pred_gt_esim', + 'pred_gt_ems', + *PROSODY_DISTANCE_KEYS, 'pred_text', 'gt_audio_text', 'gt_text', @@ -268,6 +278,9 @@ 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", + prosody_cache_dir=None, ): """Load the ASR and speaker-verification models used for evaluation. @@ -278,6 +291,11 @@ 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"``. + prosody_cache_dir: Optional directory used to cache the emotion encoder, + classifiers, and related model files. Returns: Dictionary containing: @@ -291,6 +309,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 +320,7 @@ def load_evaluation_models( 'whisper_model': None, 'whisper_processor': None, 'feature_extractor': None, + 'emotion_model': None, } if asr_model_type == "nemo": @@ -326,9 +347,70 @@ 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_dir=prosody_cache_dir, + 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, embedding_type="score_vector"): + """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, + embedding_type=embedding_type, + ) + 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, + config: Optional[ProsodyDistanceConfig] = None, +): + """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, + config=config, + ).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 +441,10 @@ def evaluate_dir( asr_model_type="nemo", with_utmosv2=True, strip_text_annotations_for_metrics=False, + with_prosody_metrics=False, + prosody_model_size="small", + prosody_embedding_type="score_vector", + prosody_cache_dir=None, asr_batch_size=32, eou_batch_size=32, device="cuda", @@ -396,12 +482,17 @@ 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, + prosody_cache_dir=prosody_cache_dir, ) 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'] + prosody_distance_config = ProsodyDistanceConfig() if with_prosody_metrics else None # 3. EoU classifier (support for English only) if language == "en": @@ -493,6 +584,22 @@ 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, + embedding_type=prosody_embedding_type, + ) + prosody_distance_metrics = compute_acoustic_prosody_metrics( + gt_audio_path=gt_audio_filepath, + pred_audio_path=pred_audio_filepath, + text=gt_text, + config=prosody_distance_config, + ) logging.info(f"{ridx} GT Text: {gt_text}") logging.info(f"{ridx} Pr Text: {pred_text}") # Format cer and wer to 2 decimal places @@ -608,7 +715,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 +736,10 @@ def evaluate( strip_text_annotations_for_metrics=False, with_fcd=True, codec_model_path=None, + with_prosody_metrics=False, + prosody_model_size="small", + prosody_embedding_type="head_concat", + prosody_cache_dir=None, asr_batch_size=32, eou_batch_size=32, device="cuda", @@ -663,6 +777,10 @@ 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, + prosody_embedding_type=prosody_embedding_type, + prosody_cache_dir=prosody_cache_dir, asr_batch_size=asr_batch_size, eou_batch_size=eou_batch_size, device=device, @@ -717,6 +835,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,6 +895,13 @@ 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] @@ -818,6 +955,19 @@ 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( + '--prosody_embedding_type', + type=str, + default="score_vector", + choices=["head_concat", "head_mean", "score_vector"], + ) + parser.add_argument('--prosody_cache_dir', type=str, default=None) parser.add_argument( '--strip_text_annotations_for_metrics', action='store_true', @@ -838,7 +988,11 @@ 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, + prosody_embedding_type=args.prosody_embedding_type, + prosody_cache_dir=args.prosody_cache_dir, ) diff --git a/nemo/collections/tts/modules/magpietts_inference/evaluation.py b/nemo/collections/tts/modules/magpietts_inference/evaluation.py index a4a7da8f7013..cfb55826c526 100644 --- a/nemo/collections/tts/modules/magpietts_inference/evaluation.py +++ b/nemo/collections/tts/modules/magpietts_inference/evaluation.py @@ -43,6 +43,11 @@ 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"). + prosody_embedding_type: Emotion embedding type used for ESIM. + prosody_cache_dir: Optional Hugging Face cache directory for the emotion encoder. 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 +60,10 @@ 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" + prosody_embedding_type: str = "score_vector" + prosody_cache_dir: str = None strip_text_annotations_for_metrics: bool = False device: str = "cuda" asr_batch_size: int = 32 @@ -99,6 +108,10 @@ 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, + prosody_embedding_type=config.prosody_embedding_type, + prosody_cache_dir=config.prosody_cache_dir, 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..6befa831112c 100644 --- a/nemo/collections/tts/modules/magpietts_inference/utils.py +++ b/nemo/collections/tts/modules/magpietts_inference/utils.py @@ -640,6 +640,11 @@ def append_metrics_to_csv(csv_path: str, checkpoint_name: str, dataset: str, met 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', ''), @@ -757,6 +762,11 @@ def turn_sort_key(r): 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] @@ -777,6 +787,11 @@ def turn_sort_key(r): "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), @@ -787,6 +802,11 @@ def turn_sort_key(r): "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, @@ -824,6 +844,11 @@ def _write_grouped_multiturn_filewise_metrics_csv(csv_path: str, grouped_rows: l "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", @@ -833,6 +858,11 @@ def _write_grouped_multiturn_filewise_metrics_csv(csv_path: str, grouped_rows: l "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 +1253,19 @@ 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( + '--prosody_embedding_type', + type=str, + default="score_vector", + choices=["head_concat", "head_mean", "score_vector"], + ) + eval_group.add_argument('--prosody_cache_dir', type=str, default=None) eval_group.add_argument( '--strip_text_annotations_for_metrics', action='store_true', diff --git a/tests/collections/tts/metrics/test_prosody.py b/tests/collections/tts/metrics/test_prosody.py new file mode 100644 index 000000000000..c863e23480b6 --- /dev/null +++ b/tests/collections/tts/metrics/test_prosody.py @@ -0,0 +1,76 @@ +# 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 ProsodyDistanceConfig, compute_prosody_distances + +_SAMPLE_RATE = 16000 +_TEXT = "hello world" + + +def _prosody_config() -> ProsodyDistanceConfig: + return ProsodyDistanceConfig( + f0_method="yin", + dtw_band_ratio=0.15, + max_dtw_frames=128, + min_voiced_frames=2, + ) + + +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, + config=_prosody_config(), + ) + + 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, + config=_prosody_config(), + ) + + 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) From 40bef8e04cb164eea04bb22dc4f457c87545db0e Mon Sep 17 00:00:00 2001 From: Edresson Casanova Date: Wed, 5 Aug 2026 09:59:02 -0300 Subject: [PATCH 03/10] Clean up emotion encoder CLI Signed-off-by: Edresson Casanova --- .../tts/metrics/emotion_encoder.py | 103 +++--------------- 1 file changed, 14 insertions(+), 89 deletions(-) diff --git a/nemo/collections/tts/metrics/emotion_encoder.py b/nemo/collections/tts/metrics/emotion_encoder.py index d2263420832d..25a0f4204a61 100644 --- a/nemo/collections/tts/metrics/emotion_encoder.py +++ b/nemo/collections/tts/metrics/emotion_encoder.py @@ -55,6 +55,7 @@ from __future__ import annotations import argparse +import json from pathlib import Path from typing import Any, Optional, Sequence, Union @@ -1138,79 +1139,40 @@ def _parse_labels(labels: Optional[str]) -> Optional[list[str]]: def main() -> None: parser = argparse.ArgumentParser(description="LAION Empathic Insight Voice embeddings and similarity.") + parser.add_argument("--audio", type=str, required=True, help="Input audio path.") + parser.add_argument("--audio-b", type=str, default=None, help="Optional second audio path for pair comparison.") + parser.add_argument("--size", type=str, default="small", choices=["small", "large"], help="Model size.") + parser.add_argument("--device", type=str, default="cuda", help="Device for Whisper encoder.") parser.add_argument( - "--audio", - type=str, - required=True, - help="Input audio path.", - ) - parser.add_argument( - "--audio-b", - type=str, - default=None, - help="Optional second audio path for similarity.", - ) - parser.add_argument( - "--size", - type=str, - default="small", - choices=["small", "large"], - help="Model size.", - ) - parser.add_argument( - "--device", - type=str, - default="cuda", - help="Device for Whisper encoder.", - ) - parser.add_argument( - "--mlp-device", - type=str, - default=None, - help="Device for MLP classifier heads. Defaults to --device.", - ) - parser.add_argument( - "--cache-dir", - type=str, - default=None, - help="Optional Hugging Face cache directory.", + "--mlp-device", type=str, default=None, help="Device for MLP classifier heads. Defaults to --device." ) + parser.add_argument("--cache-dir", type=str, default=None, help="Optional Hugging Face cache directory.") parser.add_argument( "--embedding-type", type=str, default="head_concat", choices=["head_concat", "head_mean", "score_vector"], - help="Emotion embedding type.", + help="Emotion embedding type used for optional pair comparison.", ) parser.add_argument( "--labels", type=str, default=None, - help=( - "Comma-separated labels to use. " - "Example: anger,sadness,arousal. " - "If omitted, primary emotion labels are used for similarity." - ), + help="Comma-separated labels to use. If omitted, primary emotion labels are used.", ) parser.add_argument( "--include-auxiliary", action="store_true", - help="Include auxiliary similarity labels when --labels is omitted.", + help="Include auxiliary labels for the single-audio embedding when --labels is omitted.", ) parser.add_argument( "--load-all-classifiers", action="store_true", help="Eagerly load all known classifiers at startup.", ) - parser.add_argument( - "--top-k", - type=int, - default=5, - help="Number of ranked emotions to return.", - ) + parser.add_argument("--top-k", type=int, default=5, help="Number of ranked emotions to return.") args = parser.parse_args() - labels = _parse_labels(args.labels) model = EmpathicInsightVoice.from_pretrained( @@ -1231,8 +1193,7 @@ def main() -> None: return_raw_scores=True, include_auxiliary_for_embedding=args.include_auxiliary, ) - - printable: dict[str, Any] = { + output: dict[str, Any] = { "audio_path": result["audio_path"], "model_size": result["model_size"], "top_emotion": result["top_emotion"], @@ -1243,50 +1204,14 @@ def main() -> None: } if args.audio_b is not None: - printable["audio_b"] = args.audio_b - printable["similarity"] = model.emotion_similarity( + output["comparison"] = model.compare_emotion_pair( audio_path_a=args.audio, audio_path_b=args.audio_b, labels=labels, embedding_type=args.embedding_type, - include_auxiliary=args.include_auxiliary, ) - result = model.compare_emotion_pair( - audio_path_a=args.audio, - audio_path_b=args.audio_b, - embedding_type="head_concat", - ) - - result_score_vector = model.compare_emotion_pair( - audio_path_a=args.audio, - audio_path_b=args.audio_b, - embedding_type="score_vector", - ) - - result_score_mean = model.compare_emotion_pair( - audio_path_a=args.audio, - audio_path_b=args.audio_b, - embedding_type="head_mean", - ) - - print("embedding_type=head_concat") - print(result["audio_a_top_emotion"]) - print(result["audio_b_top_emotion"]) - print(result["top_emotion_match"]) - print(result["emotion_similarity"]) - - print("embedding_type=score_vector") - print(result_score_vector["audio_a_top_emotion"]) - print(result_score_vector["audio_b_top_emotion"]) - print(result_score_vector["top_emotion_match"]) - print(result_score_vector["emotion_similarity"]) - - print("embedding_type=head_mean") - print(result_score_mean["audio_a_top_emotion"]) - print(result_score_mean["audio_b_top_emotion"]) - print(result_score_mean["top_emotion_match"]) - print(result_score_mean["emotion_similarity"]) + print(json.dumps(output, indent=2)) if __name__ == "__main__": From 2a8b8f15a733651ac86784e52c38709e51149894 Mon Sep 17 00:00:00 2001 From: Edresson Casanova Date: Wed, 5 Aug 2026 10:09:28 -0300 Subject: [PATCH 04/10] Simplify prosody metric defaults Signed-off-by: Edresson Casanova --- examples/tts/magpietts_inference.py | 4 - nemo/collections/tts/metrics/prosody.py | 437 +++++------------- .../evaluate_generated_audio.py | 31 +- .../modules/magpietts_inference/evaluation.py | 6 - .../tts/modules/magpietts_inference/utils.py | 7 - tests/collections/tts/metrics/test_prosody.py | 13 +- 6 files changed, 125 insertions(+), 373 deletions(-) diff --git a/examples/tts/magpietts_inference.py b/examples/tts/magpietts_inference.py index 0d8f556c4756..b54b1c6c003e 100644 --- a/examples/tts/magpietts_inference.py +++ b/examples/tts/magpietts_inference.py @@ -307,8 +307,6 @@ def run_inference_and_evaluation( codec_model_path=eval_config.codec_model_path, with_prosody_metrics=eval_config.with_prosody_metrics, prosody_model_size=eval_config.prosody_model_size, - prosody_embedding_type=eval_config.prosody_embedding_type, - prosody_cache_dir=eval_config.prosody_cache_dir, strip_text_annotations_for_metrics=eval_config.strip_text_annotations_for_metrics, device=eval_config.device, asr_batch_size=eval_config.asr_batch_size, @@ -464,8 +462,6 @@ def main(argv=None): 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, - prosody_embedding_type=args.prosody_embedding_type, - prosody_cache_dir=args.prosody_cache_dir, 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/tts/metrics/prosody.py b/nemo/collections/tts/metrics/prosody.py index 92e2aeb47d6b..4db446d59041 100644 --- a/nemo/collections/tts/metrics/prosody.py +++ b/nemo/collections/tts/metrics/prosody.py @@ -18,62 +18,30 @@ import math import os -import re from dataclasses import dataclass -from typing import Any, Literal, Optional +from typing import Any import librosa import numpy as np - -try: - from numba import njit -except Exception: # pragma: no cover - numba is an optional speedup. - njit = None - - -SpeechRateCharMode = Literal["nonspace", "all", "alnum"] -F0Method = Literal["pyin", "yin", "none"] -F0Normalization = Literal["gt_median", "utterance_median", "none"] -EnergyNormalization = Literal["zscore", "none"] - -_ALNUM_CHAR_RE = re.compile(r"[A-Za-z0-9]") - - -@dataclass(frozen=True) -class ProsodyDistanceConfig: - """Configuration for reference-based acoustic prosody distance metrics. - - The default values are tuned for corpus-level TTS evaluation: pYIN is used - for more stable F0 contours, F0 is converted to semitones relative to the - reference median, intensity uses log-RMS z-scores, and contours are reduced - before DTW to keep evaluation bounded. - """ - - sample_rate: int = 16000 - res_type: str = "soxr_hq" - frame_shift_ms: float = 20.0 - frame_length_ms: float = 64.0 - fmin: float = 55.0 - fmax: float = 450.0 - f0_method: F0Method = "pyin" - yin_silence_db_below_peak: float = 35.0 - pyin_n_thresholds: int = 24 - pyin_beta_a: float = 2.0 - pyin_beta_b: float = 18.0 - pyin_boltzmann_parameter: float = 2.0 - pyin_resolution: float = 0.25 - pyin_max_transition_rate: float = 12.0 - pyin_switch_prob: float = 0.01 - pyin_no_trough_prob: float = 0.01 - pyin_center: bool = True - pyin_pad_mode: str = "constant" - max_dtw_frames: int = 1000 - dtw_band_ratio: float = 0.05 - f0_nan_penalty: float = 6.0 - f0_normalization: F0Normalization = "gt_median" - intensity_normalization: EnergyNormalization = "zscore" - speech_rate_char_mode: SpeechRateCharMode = "nonspace" - min_voiced_frames: int = 5 +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) @@ -83,92 +51,72 @@ class ProsodyDistanceResult: pitch_distance: float intensity_distance: float speech_rate_distance: float - gt_duration_sec: float - pred_duration_sec: float - gt_speech_rate_cps: float - pred_speech_rate_cps: float - gt_char_count: int - def to_dict(self) -> dict[str, float | int]: + 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, - "gt_duration_sec": self.gt_duration_sec, - "pred_duration_sec": self.pred_duration_sec, - "gt_speech_rate_cps": self.gt_speech_rate_cps, - "pred_speech_rate_cps": self.pred_speech_rate_cps, - "gt_char_count": self.gt_char_count, } -if njit is not None: - - @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 +@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 - inf = 1.0e30 - prev = np.empty(m + 1, dtype=np.float64) - curr = np.empty(m + 1, dtype=np.float64) + for i in range(1, n + 1): 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 - if band_radius < 0: - j_start = 1 - j_end = m + 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: - 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) - -else: - _dtw_distance_1d_numba = None + 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, - config: Optional[ProsodyDistanceConfig] = None, ) -> ProsodyDistanceResult: """Compute acoustic prosody distances between reference and generated audio. @@ -176,50 +124,36 @@ def compute_prosody_distances( 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. - config: Optional prosody distance configuration. Returns: ProsodyDistanceResult with pitch, intensity, and speech-rate distances. """ - cfg = config or ProsodyDistanceConfig() - gt_audio, sr, gt_duration = _load_audio(gt_audio_path, cfg) - pred_audio, _, pred_duration = _load_audio(pred_audio_path, cfg) + 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, cfg) + 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 cfg.f0_method != "none": - gt_f0 = _compute_f0( - gt_audio, sr=sr, frame_length=frame_length, hop_length=hop_length, log_energy=gt_log_energy, cfg=cfg - ) - pred_f0 = _compute_f0( - pred_audio, - sr=sr, - frame_length=frame_length, - hop_length=hop_length, - log_energy=pred_log_energy, - cfg=cfg, + 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, ) - if np.isfinite(gt_f0).sum() >= cfg.min_voiced_frames and np.isfinite(pred_f0).sum() >= cfg.min_voiced_frames: - gt_pitch, pred_pitch = _prepare_f0_for_metric(gt_f0, pred_f0, cfg) - pitch_distance = _dtw_distance_1d( - _maybe_reduce_for_dtw(gt_pitch, cfg.max_dtw_frames), - _maybe_reduce_for_dtw(pred_pitch, cfg.max_dtw_frames), - nan_penalty=cfg.f0_nan_penalty, - band_ratio=cfg.dtw_band_ratio, - ) - - gt_intensity, pred_intensity = _prepare_intensity_for_metric(gt_log_energy, pred_log_energy, cfg) + intensity_distance = _dtw_distance_1d( - _maybe_reduce_for_dtw(gt_intensity, cfg.max_dtw_frames), - _maybe_reduce_for_dtw(pred_intensity, cfg.max_dtw_frames), + _maybe_reduce_for_dtw(_zscore(gt_log_energy)), + _maybe_reduce_for_dtw(_zscore(pred_log_energy)), nan_penalty=0.0, - band_ratio=cfg.dtw_band_ratio, ) - gt_char_count = _char_count(text, cfg.speech_rate_char_mode) + 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) @@ -228,30 +162,25 @@ def compute_prosody_distances( pitch_distance=_safe_float(pitch_distance), intensity_distance=_safe_float(intensity_distance), speech_rate_distance=_safe_float(speech_rate_distance), - gt_duration_sec=_safe_float(gt_duration), - pred_duration_sec=_safe_float(pred_duration), - gt_speech_rate_cps=_safe_float(gt_speech_rate), - pred_speech_rate_cps=_safe_float(pred_speech_rate), - gt_char_count=gt_char_count, ) -def _load_audio(path: str, cfg: ProsodyDistanceConfig) -> tuple[np.ndarray, int, float]: +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=cfg.sample_rate, mono=True, res_type=cfg.res_type) + 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, cfg: ProsodyDistanceConfig) -> tuple[int, int]: - hop_length = max(1, int(round(sr * cfg.frame_shift_ms / 1000.0))) - frame_length = max(hop_length * 2, int(round(sr * cfg.frame_length_ms / 1000.0))) +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 @@ -260,107 +189,42 @@ def _compute_log_energy(audio: np.ndarray, frame_length: int, hop_length: int) - 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, - log_energy: np.ndarray, - cfg: ProsodyDistanceConfig, -) -> np.ndarray: - if cfg.f0_method == "pyin": - f0, voiced_flag, _ = librosa.pyin( - y=np.asarray(audio, dtype=np.float64), - sr=sr, - fmin=cfg.fmin, - fmax=cfg.fmax, - frame_length=frame_length, - hop_length=hop_length, - center=cfg.pyin_center, - pad_mode=cfg.pyin_pad_mode, - n_thresholds=cfg.pyin_n_thresholds, - beta_parameters=(cfg.pyin_beta_a, cfg.pyin_beta_b), - boltzmann_parameter=cfg.pyin_boltzmann_parameter, - resolution=cfg.pyin_resolution, - max_transition_rate=cfg.pyin_max_transition_rate, - switch_prob=cfg.pyin_switch_prob, - no_trough_prob=cfg.pyin_no_trough_prob, - fill_na=np.nan, - ) - f0 = np.asarray(f0, dtype=np.float64) - if voiced_flag is not None: - 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 - - if cfg.f0_method == "yin": - f0 = librosa.yin( - audio, - sr=sr, - fmin=cfg.fmin, - fmax=cfg.fmax, - frame_length=frame_length, - hop_length=hop_length, - center=True, - ) - f0 = np.asarray(f0, dtype=np.float64) - f0[(f0 < cfg.fmin) | (f0 > cfg.fmax)] = np.nan - return _mask_yin_silence(f0, log_energy, cfg) - - if cfg.f0_method == "none": - return np.asarray([], dtype=np.float64) - - raise ValueError(f"Unsupported f0_method={cfg.f0_method!r}") - - -def _mask_yin_silence(f0: np.ndarray, log_energy: np.ndarray, cfg: ProsodyDistanceConfig) -> np.ndarray: - if len(log_energy) == 0: +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 - min_len = min(len(f0), len(log_energy)) - f0 = f0[:min_len].copy() - energy = log_energy[:min_len] - rms_db = 20.0 * energy / math.log(10.0) - finite = np.isfinite(rms_db) - if finite.any(): - peak_db = float(np.max(rms_db[finite])) - f0[rms_db < peak_db - cfg.yin_silence_db_below_peak] = np.nan + 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, - cfg: ProsodyDistanceConfig, -) -> tuple[np.ndarray, np.ndarray]: - if cfg.f0_normalization == "none": - return gt_f0_hz, pred_f0_hz - +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") - pred_median = float(np.nanmedian(pred_f0_hz)) if np.isfinite(pred_f0_hz).any() else float("nan") - - if cfg.f0_normalization == "gt_median": - return _hz_to_semitones(gt_f0_hz, gt_median), _hz_to_semitones(pred_f0_hz, gt_median) - if cfg.f0_normalization == "utterance_median": - return _hz_to_semitones(gt_f0_hz, gt_median), _hz_to_semitones(pred_f0_hz, pred_median) - raise ValueError(f"Unsupported f0_normalization={cfg.f0_normalization!r}") - + return _hz_to_semitones(gt_f0_hz, gt_median), _hz_to_semitones(pred_f0_hz, gt_median) -def _prepare_intensity_for_metric( - gt_log_energy: np.ndarray, - pred_log_energy: np.ndarray, - cfg: ProsodyDistanceConfig, -) -> tuple[np.ndarray, np.ndarray]: - if cfg.intensity_normalization == "none": - return gt_log_energy, pred_log_energy - if cfg.intensity_normalization == "zscore": - return _zscore(gt_log_energy), _zscore(pred_log_energy) - raise ValueError(f"Unsupported intensity_normalization={cfg.intensity_normalization!r}") - -def _dtw_distance_1d(x: np.ndarray, y: np.ndarray, nan_penalty: float, band_ratio: float) -> float: +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: @@ -368,57 +232,8 @@ def _dtw_distance_1d(x: np.ndarray, y: np.ndarray, nan_penalty: float, band_rati if len(x) == 0 or len(y) == 0: return float("nan") - if band_ratio is None or band_ratio < 0: - band_radius = -1 - else: - band_radius = max(abs(len(x) - len(y)), int(math.ceil(float(band_ratio) * max(len(x), len(y))))) - - if _dtw_distance_1d_numba is not None: - try: - return _safe_float(_dtw_distance_1d_numba(x, y, float(nan_penalty), int(band_radius))) - except Exception: - pass - return _dtw_distance_1d_python(x, y, float(nan_penalty), int(band_radius)) - - -def _dtw_distance_1d_python(x: np.ndarray, y: np.ndarray, nan_penalty: float, band_radius: int) -> float: - n = len(x) - m = len(y) - if n == 0 or m == 0: - return float("nan") - - prev = np.full(m + 1, np.inf, dtype=np.float64) - curr = np.full(m + 1, np.inf, dtype=np.float64) - prev[0] = 0.0 - - for i in range(1, n + 1): - curr.fill(np.inf) - if band_radius < 0: - j_start = 1 - j_end = m + 1 - else: - j_start = max(1, i - band_radius) - j_end = min(m, i + band_radius) + 1 - - for j in range(j_start, j_end): - cost = _frame_distance(x[i - 1], y[j - 1], nan_penalty) - curr[j] = cost + min(prev[j], curr[j - 1], prev[j - 1]) - prev, curr = curr, prev - - total = prev[m] - if not np.isfinite(total): - return float("nan") - return float(total / (n + m)) - - -def _frame_distance(x: float, y: float, nan_penalty: float) -> float: - x_nan = np.isnan(x) - y_nan = np.isnan(y) - if x_nan and y_nan: - return 0.0 - if x_nan or y_nan: - return nan_penalty - return abs(float(x) - float(y)) + 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: @@ -446,17 +261,14 @@ def _zscore(values: np.ndarray, eps: float = 1.0e-8) -> np.ndarray: return out -def _maybe_reduce_for_dtw(values: np.ndarray, max_frames: int) -> np.ndarray: +def _maybe_reduce_for_dtw(values: np.ndarray) -> np.ndarray: values = np.asarray(values, dtype=np.float64) - if max_frames is None or max_frames <= 0 or len(values) <= max_frames: + if len(values) <= _MAX_DTW_FRAMES: return values - return _resample_1d_preserve_nans(values, int(max_frames)) + return _resample_1d_preserve_nans(values, _MAX_DTW_FRAMES) def _resample_1d_preserve_nans(values: np.ndarray, target_len: int) -> np.ndarray: - values = np.asarray(values, dtype=np.float64) - if target_len <= 0: - raise ValueError("target_len must be positive") if len(values) == target_len: return values if len(values) == 0: @@ -476,17 +288,10 @@ def _resample_1d_preserve_nans(values: np.ndarray, target_len: int) -> np.ndarra return out.astype(np.float64) -def _char_count(text: Any, mode: SpeechRateCharMode) -> int: +def _char_count(text: Any) -> int: if text is None: return 0 - text = str(text) - if mode == "nonspace": - return sum(1 for ch in text if not ch.isspace()) - if mode == "alnum": - return len(_ALNUM_CHAR_RE.findall(text)) - if mode == "all": - return len(text) - raise ValueError(f"Unsupported speech_rate_char_mode={mode!r}") + return sum(1 for ch in str(text) if not ch.isspace()) def _safe_float(value: Any) -> float: 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 fd02aeca66d0..790f1a48f43a 100644 --- a/nemo/collections/tts/modules/magpietts_inference/evaluate_generated_audio.py +++ b/nemo/collections/tts/modules/magpietts_inference/evaluate_generated_audio.py @@ -36,7 +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 ProsodyDistanceConfig, compute_prosody_distances +from nemo.collections.tts.metrics.prosody import compute_prosody_distances from nemo.collections.tts.parts.utils.tts_dataset_utils import ( JapaneseTextProcessor, NemoTranscriber, @@ -280,7 +280,6 @@ def load_evaluation_models( device="cuda", with_prosody_metrics=False, prosody_model_size="small", - prosody_cache_dir=None, ): """Load the ASR and speaker-verification models used for evaluation. @@ -294,8 +293,6 @@ def load_evaluation_models( 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"``. - prosody_cache_dir: Optional directory used to cache the emotion encoder, - classifiers, and related model files. Returns: Dictionary containing: @@ -356,7 +353,6 @@ def load_evaluation_models( size=prosody_model_size, device=device, mlp_device=device, - cache_dir=prosody_cache_dir, cache_classifiers=True, load_all_classifiers=False, top_k_emotions=1, @@ -367,7 +363,7 @@ def load_evaluation_models( return models -def compute_emotion_pair_metrics(emotion_model, gt_audio_path, pred_audio_path, embedding_type="score_vector"): +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') @@ -376,7 +372,7 @@ def compute_emotion_pair_metrics(emotion_model, gt_audio_path, pred_audio_path, result = emotion_model.compare_emotion_pair( audio_path_a=gt_audio_path, audio_path_b=pred_audio_path, - embedding_type=embedding_type, + embedding_type="score_vector", ) return float(result["emotion_similarity"]), float(result["top_emotion_match"]) except Exception as e: @@ -392,7 +388,6 @@ def compute_acoustic_prosody_metrics( gt_audio_path, pred_audio_path, text, - config: Optional[ProsodyDistanceConfig] = None, ): """Compute reference-based pitch, intensity, and speech-rate distances.""" if gt_audio_path is None or pred_audio_path is None: @@ -403,7 +398,6 @@ def compute_acoustic_prosody_metrics( gt_audio_path=gt_audio_path, pred_audio_path=pred_audio_path, text=text, - config=config, ).to_dict() return {key: metrics[key] for key in PROSODY_DISTANCE_KEYS} except Exception as e: @@ -443,8 +437,6 @@ def evaluate_dir( strip_text_annotations_for_metrics=False, with_prosody_metrics=False, prosody_model_size="small", - prosody_embedding_type="score_vector", - prosody_cache_dir=None, asr_batch_size=32, eou_batch_size=32, device="cuda", @@ -484,7 +476,6 @@ def evaluate_dir( device=device, with_prosody_metrics=with_prosody_metrics, prosody_model_size=prosody_model_size, - prosody_cache_dir=prosody_cache_dir, ) asr_model = models['asr_model'] @@ -492,7 +483,6 @@ def evaluate_dir( speaker_verification_model = models['sv_model'] speaker_verification_model_alternate = models['sv_model_alternate'] emotion_model = models['emotion_model'] - prosody_distance_config = ProsodyDistanceConfig() if with_prosody_metrics else None # 3. EoU classifier (support for English only) if language == "en": @@ -592,13 +582,11 @@ def evaluate_dir( emotion_model, gt_audio_filepath, pred_audio_filepath, - embedding_type=prosody_embedding_type, ) prosody_distance_metrics = compute_acoustic_prosody_metrics( gt_audio_path=gt_audio_filepath, pred_audio_path=pred_audio_filepath, text=gt_text, - config=prosody_distance_config, ) logging.info(f"{ridx} GT Text: {gt_text}") logging.info(f"{ridx} Pr Text: {pred_text}") @@ -738,8 +726,6 @@ def evaluate( codec_model_path=None, with_prosody_metrics=False, prosody_model_size="small", - prosody_embedding_type="head_concat", - prosody_cache_dir=None, asr_batch_size=32, eou_batch_size=32, device="cuda", @@ -779,8 +765,6 @@ def evaluate( strip_text_annotations_for_metrics=strip_text_annotations_for_metrics, with_prosody_metrics=with_prosody_metrics, prosody_model_size=prosody_model_size, - prosody_embedding_type=prosody_embedding_type, - prosody_cache_dir=prosody_cache_dir, asr_batch_size=asr_batch_size, eou_batch_size=eou_batch_size, device=device, @@ -961,13 +945,6 @@ def main(): 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( - '--prosody_embedding_type', - type=str, - default="score_vector", - choices=["head_concat", "head_mean", "score_vector"], - ) - parser.add_argument('--prosody_cache_dir', type=str, default=None) parser.add_argument( '--strip_text_annotations_for_metrics', action='store_true', @@ -991,8 +968,6 @@ def main(): 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, - prosody_embedding_type=args.prosody_embedding_type, - prosody_cache_dir=args.prosody_cache_dir, ) diff --git a/nemo/collections/tts/modules/magpietts_inference/evaluation.py b/nemo/collections/tts/modules/magpietts_inference/evaluation.py index cfb55826c526..8ca3d8c3ccc0 100644 --- a/nemo/collections/tts/modules/magpietts_inference/evaluation.py +++ b/nemo/collections/tts/modules/magpietts_inference/evaluation.py @@ -46,8 +46,6 @@ class EvaluationConfig: 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"). - prosody_embedding_type: Emotion embedding type used for ESIM. - prosody_cache_dir: Optional Hugging Face cache directory for the emotion encoder. 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. """ @@ -62,8 +60,6 @@ class EvaluationConfig: codec_model_path: str = None with_prosody_metrics: bool = False prosody_model_size: str = "small" - prosody_embedding_type: str = "score_vector" - prosody_cache_dir: str = None strip_text_annotations_for_metrics: bool = False device: str = "cuda" asr_batch_size: int = 32 @@ -110,8 +106,6 @@ def evaluate_generated_audio_dir( codec_model_path=config.codec_model_path, with_prosody_metrics=config.with_prosody_metrics, prosody_model_size=config.prosody_model_size, - prosody_embedding_type=config.prosody_embedding_type, - prosody_cache_dir=config.prosody_cache_dir, 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 6befa831112c..9fbf87d995e5 100644 --- a/nemo/collections/tts/modules/magpietts_inference/utils.py +++ b/nemo/collections/tts/modules/magpietts_inference/utils.py @@ -1259,13 +1259,6 @@ def _add_common_args(parser: argparse.ArgumentParser) -> None: 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( - '--prosody_embedding_type', - type=str, - default="score_vector", - choices=["head_concat", "head_mean", "score_vector"], - ) - eval_group.add_argument('--prosody_cache_dir', type=str, default=None) eval_group.add_argument( '--strip_text_annotations_for_metrics', action='store_true', diff --git a/tests/collections/tts/metrics/test_prosody.py b/tests/collections/tts/metrics/test_prosody.py index c863e23480b6..66096ee9b68e 100644 --- a/tests/collections/tts/metrics/test_prosody.py +++ b/tests/collections/tts/metrics/test_prosody.py @@ -16,21 +16,12 @@ import pytest import soundfile as sf -from nemo.collections.tts.metrics.prosody import ProsodyDistanceConfig, compute_prosody_distances +from nemo.collections.tts.metrics.prosody import compute_prosody_distances _SAMPLE_RATE = 16000 _TEXT = "hello world" -def _prosody_config() -> ProsodyDistanceConfig: - return ProsodyDistanceConfig( - f0_method="yin", - dtw_band_ratio=0.15, - max_dtw_frames=128, - min_voiced_frames=2, - ) - - 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 @@ -49,7 +40,6 @@ def test_prosody_distances_are_zero_for_identical_audio(tmp_path): gt_audio_path=str(gt_path), pred_audio_path=str(pred_path), text=_TEXT, - config=_prosody_config(), ) assert metrics.pitch_distance == pytest.approx(0.0, abs=1.0e-6) @@ -68,7 +58,6 @@ def test_speech_rate_distance_tracks_duration_difference(tmp_path): gt_audio_path=str(gt_path), pred_audio_path=str(pred_path), text=_TEXT, - config=_prosody_config(), ) assert np.isfinite(metrics.pitch_distance) From 07c8740385333455c70f8c8d4c15d2834995b733 Mon Sep 17 00:00:00 2001 From: Edresson Casanova Date: Wed, 5 Aug 2026 10:11:54 -0300 Subject: [PATCH 05/10] Use default prosody emotion embedding Signed-off-by: Edresson Casanova --- .../tts/modules/magpietts_inference/evaluate_generated_audio.py | 1 - 1 file changed, 1 deletion(-) 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 790f1a48f43a..447be739b76f 100644 --- a/nemo/collections/tts/modules/magpietts_inference/evaluate_generated_audio.py +++ b/nemo/collections/tts/modules/magpietts_inference/evaluate_generated_audio.py @@ -372,7 +372,6 @@ def compute_emotion_pair_metrics(emotion_model, gt_audio_path, pred_audio_path): result = emotion_model.compare_emotion_pair( audio_path_a=gt_audio_path, audio_path_b=pred_audio_path, - embedding_type="score_vector", ) return float(result["emotion_similarity"]), float(result["top_emotion_match"]) except Exception as e: From ad26f40f4e13207264ce95d1f93c197562d11f0f Mon Sep 17 00:00:00 2001 From: Edresson Casanova Date: Wed, 5 Aug 2026 10:15:28 -0300 Subject: [PATCH 06/10] Remove emotion encoder CLI Signed-off-by: Edresson Casanova --- .../tts/metrics/emotion_encoder.py | 103 ------------------ 1 file changed, 103 deletions(-) diff --git a/nemo/collections/tts/metrics/emotion_encoder.py b/nemo/collections/tts/metrics/emotion_encoder.py index 25a0f4204a61..54dee78083a2 100644 --- a/nemo/collections/tts/metrics/emotion_encoder.py +++ b/nemo/collections/tts/metrics/emotion_encoder.py @@ -54,8 +54,6 @@ from __future__ import annotations -import argparse -import json from pathlib import Path from typing import Any, Optional, Sequence, Union @@ -1115,104 +1113,3 @@ def cleanup(self) -> None: if torch.cuda.is_available(): torch.cuda.empty_cache() - - -# ============================================================================= -# CLI utilities -# ============================================================================= - - -def _tensor_info(tensor: torch.Tensor) -> dict[str, Any]: - return { - "shape": list(tensor.shape), - "dtype": str(tensor.dtype), - "device": str(tensor.device), - } - - -def _parse_labels(labels: Optional[str]) -> Optional[list[str]]: - if labels is None or labels.strip() == "": - return None - - return [item.strip() for item in labels.split(",") if item.strip()] - - -def main() -> None: - parser = argparse.ArgumentParser(description="LAION Empathic Insight Voice embeddings and similarity.") - parser.add_argument("--audio", type=str, required=True, help="Input audio path.") - parser.add_argument("--audio-b", type=str, default=None, help="Optional second audio path for pair comparison.") - parser.add_argument("--size", type=str, default="small", choices=["small", "large"], help="Model size.") - parser.add_argument("--device", type=str, default="cuda", help="Device for Whisper encoder.") - parser.add_argument( - "--mlp-device", type=str, default=None, help="Device for MLP classifier heads. Defaults to --device." - ) - parser.add_argument("--cache-dir", type=str, default=None, help="Optional Hugging Face cache directory.") - parser.add_argument( - "--embedding-type", - type=str, - default="head_concat", - choices=["head_concat", "head_mean", "score_vector"], - help="Emotion embedding type used for optional pair comparison.", - ) - parser.add_argument( - "--labels", - type=str, - default=None, - help="Comma-separated labels to use. If omitted, primary emotion labels are used.", - ) - parser.add_argument( - "--include-auxiliary", - action="store_true", - help="Include auxiliary labels for the single-audio embedding when --labels is omitted.", - ) - parser.add_argument( - "--load-all-classifiers", - action="store_true", - help="Eagerly load all known classifiers at startup.", - ) - parser.add_argument("--top-k", type=int, default=5, help="Number of ranked emotions to return.") - - args = parser.parse_args() - labels = _parse_labels(args.labels) - - model = EmpathicInsightVoice.from_pretrained( - size=args.size, - device=args.device, - mlp_device=args.mlp_device, - cache_dir=args.cache_dir, - cache_classifiers=True, - load_all_classifiers=args.load_all_classifiers, - top_k_emotions=args.top_k, - ) - - result = model.compute( - audio_path=args.audio, - labels=labels, - return_embedding=True, - embedding_type=args.embedding_type, - return_raw_scores=True, - include_auxiliary_for_embedding=args.include_auxiliary, - ) - output: dict[str, Any] = { - "audio_path": result["audio_path"], - "model_size": result["model_size"], - "top_emotion": result["top_emotion"], - "embedding_type": result["embedding_type"], - "embedding": _tensor_info(result["embedding"]), - "emotions": result["emotions"], - "raw_scores": result["raw_scores"], - } - - if args.audio_b is not None: - output["comparison"] = model.compare_emotion_pair( - audio_path_a=args.audio, - audio_path_b=args.audio_b, - labels=labels, - embedding_type=args.embedding_type, - ) - - print(json.dumps(output, indent=2)) - - -if __name__ == "__main__": - main() From c8b9edc2a1d7a023420d63ff553abb0d006e4c8b Mon Sep 17 00:00:00 2001 From: Edresson Casanova Date: Wed, 5 Aug 2026 11:17:02 -0300 Subject: [PATCH 07/10] Fix black formatting Signed-off-by: Edresson Casanova --- nemo/collections/common/data/lhotse/sampling.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/nemo/collections/common/data/lhotse/sampling.py b/nemo/collections/common/data/lhotse/sampling.py index 3c5293fbac09..6e5fcd3bb1f5 100644 --- a/nemo/collections/common/data/lhotse/sampling.py +++ b/nemo/collections/common/data/lhotse/sampling.py @@ -345,11 +345,7 @@ def __call__(self, example) -> bool: # 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() - ) + speaker_id = speaker_id.rsplit("Speaker:", maxsplit=1)[-1].split("|", maxsplit=1)[0].strip() if speaker_id in excluded_speaker_ids: return False From c912db20ed25c223609783467fa37974b219ed8f Mon Sep 17 00:00:00 2001 From: Edresson Casanova Date: Fri, 7 Aug 2026 08:17:28 -0300 Subject: [PATCH 08/10] Update speaker filter test expectation Signed-off-by: Edresson Casanova --- tests/collections/common/test_lhotse_tts_filters.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 From f7db3e52f66bde7ad8cf37378eb2d2a4268f6700 Mon Sep 17 00:00:00 2001 From: Edresson Casanova Date: Mon, 10 Aug 2026 17:09:45 -0300 Subject: [PATCH 09/10] Add generated-to-ground-truth audio CER metric Signed-off-by: Edresson Casanova --- examples/tts/magpietts_inference.py | 3 ++- .../magpietts_inference/evaluate_generated_audio.py | 12 ++++++++++++ .../tts/modules/magpietts_inference/utils.py | 7 +++++++ 3 files changed, 21 insertions(+), 1 deletion(-) diff --git a/examples/tts/magpietts_inference.py b/examples/tts/magpietts_inference.py index b54b1c6c003e..176bea8f5df7 100644 --- a/examples/tts/magpietts_inference.py +++ b/examples/tts/magpietts_inference.py @@ -167,7 +167,8 @@ 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," + "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,esim_pred_gt_avg,ems_pred_gt_avg," "pitch_distance_avg,intensity_distance_avg,speech_rate_distance_avg," 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 447be739b76f..91a6ea83a0a5 100644 --- a/nemo/collections/tts/modules/magpietts_inference/evaluate_generated_audio.py +++ b/nemo/collections/tts/modules/magpietts_inference/evaluate_generated_audio.py @@ -105,6 +105,7 @@ def strip_text_annotations_from_text(text: str) -> str: FILEWISE_METRICS_TO_SAVE = [ 'cer', + 'cer_pred_gt_audio', 'wer', 'pred_context_ssim', 'pred_gt_esim', @@ -561,6 +562,11 @@ 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') + ) # Japanese: additional reading-based CER on Katakana (pyopenjtalk g2p), robust to # kanji/kana spelling differences between reference and ASR hypothesis. @@ -682,6 +688,7 @@ 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], 'katakana_cer': katakana_cer, 'gt_katakana': gt_katakana, @@ -888,7 +895,11 @@ def compute_global_metrics( # 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') 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['cer_gt_audio_cumulative'] = word_error_rate_detail( hypotheses=gt_audio_texts, references=gt_texts, use_cer=True )[0] @@ -896,6 +907,7 @@ 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['cer_gt_audio_cumulative'] = float('NaN') avg_metrics['wer_gt_audio_cumulative'] = float('NaN') logging.warning( diff --git a/nemo/collections/tts/modules/magpietts_inference/utils.py b/nemo/collections/tts/modules/magpietts_inference/utils.py index 9fbf87d995e5..982662449675 100644 --- a/nemo/collections/tts/modules/magpietts_inference/utils.py +++ b/nemo/collections/tts/modules/magpietts_inference/utils.py @@ -634,6 +634,8 @@ 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('ssim_pred_gt_avg', ''), metrics.get('ssim_pred_context_avg', ''), metrics.get('ssim_gt_context_avg', ''), @@ -758,6 +760,7 @@ 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] 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] @@ -783,6 +786,7 @@ 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), "pred_context_ssim": _mean_finite(pred_context_ssim_turns), "pred_gt_ssim": _mean_finite(pred_gt_ssim_turns), @@ -798,6 +802,7 @@ def turn_sort_key(r): # 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, "pred_context_ssim_turns": pred_context_ssim_turns, "pred_gt_ssim_turns": pred_gt_ssim_turns, @@ -840,6 +845,7 @@ def _write_grouped_multiturn_filewise_metrics_csv(csv_path: str, grouped_rows: l "rank", "num_turns", "cer", + "cer_pred_gt_audio", "wer", "pred_context_ssim", "pred_gt_ssim", @@ -854,6 +860,7 @@ def _write_grouped_multiturn_filewise_metrics_csv(csv_path: str, grouped_rows: l "eou_trail_rms_ratio", "turn_ids", "cer_turns", + "cer_pred_gt_audio_turns", "wer_turns", "pred_context_ssim_turns", "pred_gt_ssim_turns", From 75c5827d34b3bc1af3f0613b06b54bdc6d959c38 Mon Sep 17 00:00:00 2001 From: Edresson Casanova Date: Mon, 10 Aug 2026 17:18:49 -0300 Subject: [PATCH 10/10] Add generated-to-ground-truth audio WER metric Signed-off-by: Edresson Casanova --- examples/tts/magpietts_inference.py | 1 + .../magpietts_inference/evaluate_generated_audio.py | 12 ++++++++++++ .../tts/modules/magpietts_inference/utils.py | 7 +++++++ 3 files changed, 20 insertions(+) diff --git a/examples/tts/magpietts_inference.py b/examples/tts/magpietts_inference.py index 176bea8f5df7..cc50159c4752 100644 --- a/examples/tts/magpietts_inference.py +++ b/examples/tts/magpietts_inference.py @@ -168,6 +168,7 @@ def run_inference_and_evaluation( csv_header = ( "checkpoint_name,dataset,cer_filewise_avg,wer_filewise_avg,cer_cumulative," "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,esim_pred_gt_avg,ems_pred_gt_avg," 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 91a6ea83a0a5..70bc8c7379ba 100644 --- a/nemo/collections/tts/modules/magpietts_inference/evaluate_generated_audio.py +++ b/nemo/collections/tts/modules/magpietts_inference/evaluate_generated_audio.py @@ -107,6 +107,7 @@ def strip_text_annotations_from_text(text: str) -> str: 'cer', 'cer_pred_gt_audio', 'wer', + 'wer_pred_gt_audio', 'pred_context_ssim', 'pred_gt_esim', 'pred_gt_ems', @@ -567,6 +568,11 @@ def evaluate_dir( 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. @@ -690,6 +696,7 @@ def evaluate_dir( '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, @@ -896,10 +903,14 @@ def compute_global_metrics( # 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] @@ -908,6 +919,7 @@ def compute_global_metrics( )[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( diff --git a/nemo/collections/tts/modules/magpietts_inference/utils.py b/nemo/collections/tts/modules/magpietts_inference/utils.py index 982662449675..8a9a5ff05e6f 100644 --- a/nemo/collections/tts/modules/magpietts_inference/utils.py +++ b/nemo/collections/tts/modules/magpietts_inference/utils.py @@ -636,6 +636,8 @@ def append_metrics_to_csv(csv_path: str, checkpoint_name: str, dataset: str, met 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', ''), @@ -762,6 +764,7 @@ def turn_sort_key(r): 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] @@ -788,6 +791,7 @@ def turn_sort_key(r): "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), @@ -804,6 +808,7 @@ def turn_sort_key(r): "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, @@ -847,6 +852,7 @@ def _write_grouped_multiturn_filewise_metrics_csv(csv_path: str, grouped_rows: l "cer", "cer_pred_gt_audio", "wer", + "wer_pred_gt_audio", "pred_context_ssim", "pred_gt_ssim", "gt_context_ssim", @@ -862,6 +868,7 @@ def _write_grouped_multiturn_filewise_metrics_csv(csv_path: str, grouped_rows: l "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",