first commit
This commit is contained in:
@@ -0,0 +1,226 @@
|
||||
import math
|
||||
import random
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torchaudio
|
||||
from torch.nn.utils.rnn import pad_sequence
|
||||
|
||||
from dots_tts.modules.speaker.campplus import CAMPPlus
|
||||
from dots_tts.modules.speaker.fbank import (
|
||||
_SPEAKER_FBANK_N_MELS,
|
||||
_SPEAKER_FBANK_SAMPLE_RATE,
|
||||
extract_speaker_fbank,
|
||||
)
|
||||
|
||||
|
||||
class SpeakerXVectorFeatures(nn.Module):
|
||||
"""
|
||||
Speaker embedding extractor based on 3D-Speaker CAM++.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
sample_rate=_SPEAKER_FBANK_SAMPLE_RATE,
|
||||
campplus_embedding_size=512,
|
||||
max_audio_seconds=10.0,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.sample_rate = sample_rate
|
||||
self.max_audio_seconds = float(max_audio_seconds)
|
||||
self.model = CAMPPlus(
|
||||
feat_dim=_SPEAKER_FBANK_N_MELS,
|
||||
embedding_size=campplus_embedding_size,
|
||||
)
|
||||
self.resample = None
|
||||
if self.sample_rate != _SPEAKER_FBANK_SAMPLE_RATE:
|
||||
self.resample = torchaudio.transforms.Resample(
|
||||
orig_freq=sample_rate,
|
||||
new_freq=_SPEAKER_FBANK_SAMPLE_RATE,
|
||||
)
|
||||
|
||||
for param in self.model.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
@staticmethod
|
||||
def _normalize_lengths(lengths, batch_size, max_length, device, *, min_length):
|
||||
if lengths is None:
|
||||
return torch.full(
|
||||
(batch_size,),
|
||||
max_length,
|
||||
device=device,
|
||||
dtype=torch.long,
|
||||
)
|
||||
return lengths.to(device=device, dtype=torch.long).clamp(
|
||||
min=min_length,
|
||||
max=max_length,
|
||||
)
|
||||
|
||||
def _crop_audio(self, audio, audio_lengths=None):
|
||||
original_lengths = self._normalize_lengths(
|
||||
audio_lengths,
|
||||
audio.size(0),
|
||||
audio.size(-1),
|
||||
audio.device,
|
||||
min_length=0,
|
||||
)
|
||||
if self.max_audio_seconds <= 0:
|
||||
return audio, original_lengths, original_lengths, torch.zeros_like(
|
||||
original_lengths
|
||||
)
|
||||
|
||||
max_input_length = round(self.sample_rate * self.max_audio_seconds)
|
||||
cropped_audio = []
|
||||
cropped_lengths = []
|
||||
starts = []
|
||||
|
||||
for index, total_length_tensor in enumerate(original_lengths):
|
||||
total_length = int(total_length_tensor.item())
|
||||
cropped_length = min(total_length, max_input_length)
|
||||
start = (
|
||||
random.randint(0, total_length - cropped_length)
|
||||
if total_length > cropped_length
|
||||
else 0
|
||||
)
|
||||
cropped_audio.append(audio[index, start : start + cropped_length])
|
||||
cropped_lengths.append(cropped_length)
|
||||
starts.append(start)
|
||||
|
||||
return pad_sequence(
|
||||
cropped_audio,
|
||||
batch_first=True,
|
||||
padding_value=0.0,
|
||||
), original_lengths, torch.tensor(
|
||||
cropped_lengths,
|
||||
device=audio.device,
|
||||
dtype=torch.long,
|
||||
), torch.tensor(starts, device=audio.device, dtype=torch.long)
|
||||
|
||||
def _crop_fbank(
|
||||
self,
|
||||
fbank,
|
||||
fbank_lengths,
|
||||
original_audio_lengths,
|
||||
cropped_audio_lengths,
|
||||
starts,
|
||||
):
|
||||
original_fbank_lengths = self._normalize_lengths(
|
||||
fbank_lengths,
|
||||
fbank.size(0),
|
||||
fbank.size(1),
|
||||
fbank.device,
|
||||
min_length=1,
|
||||
)
|
||||
cropped_fbank = []
|
||||
cropped_fbank_lengths = []
|
||||
|
||||
for index, total_feat_length_tensor in enumerate(original_fbank_lengths):
|
||||
total_audio_length = int(original_audio_lengths[index].item())
|
||||
total_feat_length = int(total_feat_length_tensor.item())
|
||||
start_audio = int(starts[index].item())
|
||||
end_audio = start_audio + int(cropped_audio_lengths[index].item())
|
||||
|
||||
if total_audio_length > 0:
|
||||
start_feat = math.floor(
|
||||
start_audio * total_feat_length / total_audio_length
|
||||
)
|
||||
end_feat = math.ceil(end_audio * total_feat_length / total_audio_length)
|
||||
else:
|
||||
start_feat = 0
|
||||
end_feat = 1
|
||||
|
||||
start_feat = min(start_feat, total_feat_length - 1)
|
||||
end_feat = min(max(end_feat, start_feat + 1), total_feat_length)
|
||||
cropped_fbank.append(fbank[index, start_feat:end_feat])
|
||||
cropped_fbank_lengths.append(end_feat - start_feat)
|
||||
|
||||
return pad_sequence(
|
||||
cropped_fbank,
|
||||
batch_first=True,
|
||||
padding_value=0.0,
|
||||
), torch.tensor(
|
||||
cropped_fbank_lengths,
|
||||
device=fbank.device,
|
||||
dtype=torch.long,
|
||||
)
|
||||
|
||||
def _extract_fbank_batch(self, audio, audio_lengths):
|
||||
if self.resample is not None:
|
||||
audio = self.resample(audio)
|
||||
audio_lengths = torch.ceil(
|
||||
audio_lengths.float()
|
||||
* (_SPEAKER_FBANK_SAMPLE_RATE / self.sample_rate)
|
||||
).long()
|
||||
|
||||
audio_cpu = audio.detach().cpu()
|
||||
features = []
|
||||
|
||||
for index, valid_length_tensor in enumerate(audio_lengths):
|
||||
valid_length = int(valid_length_tensor.item())
|
||||
waveform = audio_cpu[index, :valid_length]
|
||||
if waveform.numel() == 0:
|
||||
waveform = audio_cpu.new_zeros(1)
|
||||
features.append(
|
||||
extract_speaker_fbank(
|
||||
waveform,
|
||||
sample_rate=_SPEAKER_FBANK_SAMPLE_RATE,
|
||||
)
|
||||
)
|
||||
|
||||
fbank_lengths = torch.tensor(
|
||||
[feature.size(0) for feature in features],
|
||||
device=audio.device,
|
||||
dtype=torch.long,
|
||||
)
|
||||
fbank = pad_sequence(
|
||||
features,
|
||||
batch_first=True,
|
||||
padding_value=0.0,
|
||||
).to(device=audio.device, dtype=audio.dtype)
|
||||
return fbank, fbank_lengths
|
||||
|
||||
@torch.no_grad()
|
||||
@torch.autocast(enabled=False, device_type="cuda")
|
||||
def forward(
|
||||
self, audio, audio_lengths=None, fbank=None, fbank_lengths=None, **_kwargs
|
||||
):
|
||||
self.model.eval()
|
||||
audio = audio.float()
|
||||
if audio.dim() == 3:
|
||||
if audio.size(1) != 1:
|
||||
raise ValueError(
|
||||
f"Speaker encoder expects mono audio, got shape {tuple(audio.shape)}."
|
||||
)
|
||||
audio = audio[:, 0]
|
||||
elif audio.dim() != 2:
|
||||
raise ValueError(
|
||||
f"Speaker encoder expects a 2D or 3D audio tensor, got shape {tuple(audio.shape)}."
|
||||
)
|
||||
|
||||
audio, original_audio_lengths, cropped_audio_lengths, starts = self._crop_audio(
|
||||
audio,
|
||||
audio_lengths=audio_lengths,
|
||||
)
|
||||
|
||||
if fbank is None:
|
||||
fbank, fbank_lengths = self._extract_fbank_batch(
|
||||
audio,
|
||||
cropped_audio_lengths,
|
||||
)
|
||||
else:
|
||||
if not isinstance(fbank, torch.Tensor):
|
||||
raise TypeError("Speaker encoder expects `fbank` to be a torch.Tensor.")
|
||||
if fbank.dim() != 3 or fbank.size(0) != audio.size(0):
|
||||
raise ValueError(
|
||||
f"Speaker encoder expects `fbank` with shape (B, T, F) and matching batch size, got {tuple(fbank.shape)}."
|
||||
)
|
||||
fbank, fbank_lengths = self._crop_fbank(
|
||||
fbank.to(device=audio.device, dtype=torch.float32),
|
||||
fbank_lengths,
|
||||
original_audio_lengths,
|
||||
cropped_audio_lengths,
|
||||
starts,
|
||||
)
|
||||
|
||||
return self.model(fbank, lengths=fbank_lengths)
|
||||
Reference in New Issue
Block a user