File size: 3,345 Bytes
7ef8c9b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | """Standalone VocBulwark speaker encoder.
Loadable with `AutoModel.from_pretrained(path, trust_remote_code=True)` without
the training repository. Maps a raw waveform (mono, at `config.raw_sample_rate`)
to a fixed-size speaker embedding used to condition the companion VocBulwark
vocoder.
"""
from dataclasses import dataclass
import torch
import torch.nn.utils.parametrize as parametrize
from torch.nn.utils import remove_weight_norm
from transformers import PreTrainedModel
from transformers.modeling_outputs import ModelOutput
def _strip_weight_norm(root):
"""Remove weight-norm reparametrizations (both conventions), leaving plain
``.weight`` parameters that match the folded export weights."""
for m in root.modules():
if parametrize.is_parametrized(m, "weight"):
parametrize.remove_parametrizations(m, "weight", leave_parametrized=True)
elif hasattr(m, "weight_g"):
try:
remove_weight_norm(m)
except (ValueError, RuntimeError):
pass
from .configuration_speaker import SpeakerEncoderConfig
from .speaker_embedding_config import SpeakerEmbeddingConfig
from .speaker_embedding_architecture import SpeakerEmbeddingArchitecture
# Force the loader to copy the (transitively used) leaf module into the cache.
from .ge2e_loss import GE2ELoss as _GE2ELoss # noqa: F401
@dataclass
class SpeakerEmbeddingOutput(ModelOutput):
embeddings: torch.Tensor = None
class SpeakerEncoder(PreTrainedModel):
config_class = SpeakerEncoderConfig
def __init__(self, config):
super().__init__(config)
spk_cfg = config.speaker_embed_config
if isinstance(spk_cfg, dict):
spk_cfg = SpeakerEmbeddingConfig.from_dict(spk_cfg)
self.speaker_embedding_model = SpeakerEmbeddingArchitecture(spk_cfg)
self.speaker_embedding_model.eval()
# Export folds weight-norm into plain .weight; strip the reparametrization
# so module keys match and inference is convention-independent.
_strip_weight_norm(self)
@torch.no_grad()
def forward(self, raw_audio=None, attention_mask=None, input_values=None, **kwargs):
"""Encode a raw waveform into a speaker embedding.
Args:
raw_audio: [B, samples] (or [samples]) mono audio @ config.raw_sample_rate.
attention_mask: optional [B, samples] mask.
Returns:
SpeakerEmbeddingOutput with `.embeddings` of shape [B, embedding_size].
"""
audio = raw_audio if raw_audio is not None else input_values
if audio is None:
raise ValueError("Pass raw_audio [B, samples] @ config.raw_sample_rate.")
p = next(self.parameters())
audio = audio.to(device=p.device, dtype=p.dtype)
if audio.dim() == 1:
audio = audio.unsqueeze(0)
if attention_mask is not None:
attention_mask = attention_mask.to(device=p.device)
emb = self.speaker_embedding_model(audio, attention_mask=attention_mask).embeddings
return SpeakerEmbeddingOutput(embeddings=emb)
@torch.no_grad()
def embed(self, raw_audio, attention_mask=None):
"""Convenience: return just the [B, embedding_size] embedding tensor."""
return self.forward(raw_audio=raw_audio, attention_mask=attention_mask).embeddings
|