RustyMark/dots.tts

1import math2import random3 4import torch5import torch.nn as nn6import torchaudio7from torch.nn.utils.rnn import pad_sequence8 9from dots_tts.modules.speaker.campplus import CAMPPlus10from dots_tts.modules.speaker.fbank import (11    _SPEAKER_FBANK_N_MELS,12    _SPEAKER_FBANK_SAMPLE_RATE,13    extract_speaker_fbank,14)15 16 17class SpeakerXVectorFeatures(nn.Module):18    """19    Speaker embedding extractor based on 3D-Speaker CAM++.20    """21 22    def __init__(23        self,24        sample_rate=_SPEAKER_FBANK_SAMPLE_RATE,25        campplus_embedding_size=512,26        max_audio_seconds=10.0,27    ):28        super().__init__()29 30        self.sample_rate = sample_rate31        self.max_audio_seconds = float(max_audio_seconds)32        self.model = CAMPPlus(33            feat_dim=_SPEAKER_FBANK_N_MELS,34            embedding_size=campplus_embedding_size,35        )36        self.resample = None37        if self.sample_rate != _SPEAKER_FBANK_SAMPLE_RATE:38            self.resample = torchaudio.transforms.Resample(39                orig_freq=sample_rate,40                new_freq=_SPEAKER_FBANK_SAMPLE_RATE,41            )42 43        for param in self.model.parameters():44            param.requires_grad = False45 46    @staticmethod47    def _normalize_lengths(lengths, batch_size, max_length, device, *, min_length):48        if lengths is None:49            return torch.full(50                (batch_size,),51                max_length,52                device=device,53                dtype=torch.long,54            )55        return lengths.to(device=device, dtype=torch.long).clamp(56            min=min_length,57            max=max_length,58        )59 60    def _crop_audio(self, audio, audio_lengths=None):61        original_lengths = self._normalize_lengths(62            audio_lengths,63            audio.size(0),64            audio.size(-1),65            audio.device,66            min_length=0,67        )68        if self.max_audio_seconds <= 0:69            return audio, original_lengths, original_lengths, torch.zeros_like(70                original_lengths71            )72 73        max_input_length = round(self.sample_rate * self.max_audio_seconds)74        cropped_audio = []75        cropped_lengths = []76        starts = []77 78        for index, total_length_tensor in enumerate(original_lengths):79            total_length = int(total_length_tensor.item())80            cropped_length = min(total_length, max_input_length)81            start = (82                random.randint(0, total_length - cropped_length)83                if total_length > cropped_length84                else 085            )86            cropped_audio.append(audio[index, start : start + cropped_length])87            cropped_lengths.append(cropped_length)88            starts.append(start)89 90        return pad_sequence(91            cropped_audio,92            batch_first=True,93            padding_value=0.0,94        ), original_lengths, torch.tensor(95            cropped_lengths,96            device=audio.device,97            dtype=torch.long,98        ), torch.tensor(starts, device=audio.device, dtype=torch.long)99 100    def _crop_fbank(101        self,102        fbank,103        fbank_lengths,104        original_audio_lengths,105        cropped_audio_lengths,106        starts,107    ):108        original_fbank_lengths = self._normalize_lengths(109            fbank_lengths,110            fbank.size(0),111            fbank.size(1),112            fbank.device,113            min_length=1,114        )115        cropped_fbank = []116        cropped_fbank_lengths = []117 118        for index, total_feat_length_tensor in enumerate(original_fbank_lengths):119            total_audio_length = int(original_audio_lengths[index].item())120            total_feat_length = int(total_feat_length_tensor.item())121            start_audio = int(starts[index].item())122            end_audio = start_audio + int(cropped_audio_lengths[index].item())123 124            if total_audio_length > 0:125                start_feat = math.floor(126                    start_audio * total_feat_length / total_audio_length127                )128                end_feat = math.ceil(end_audio * total_feat_length / total_audio_length)129            else:130                start_feat = 0131                end_feat = 1132 133            start_feat = min(start_feat, total_feat_length - 1)134            end_feat = min(max(end_feat, start_feat + 1), total_feat_length)135            cropped_fbank.append(fbank[index, start_feat:end_feat])136            cropped_fbank_lengths.append(end_feat - start_feat)137 138        return pad_sequence(139            cropped_fbank,140            batch_first=True,141            padding_value=0.0,142        ), torch.tensor(143            cropped_fbank_lengths,144            device=fbank.device,145            dtype=torch.long,146        )147 148    def _extract_fbank_batch(self, audio, audio_lengths):149        if self.resample is not None:150            audio = self.resample(audio)151            audio_lengths = torch.ceil(152                audio_lengths.float()153                * (_SPEAKER_FBANK_SAMPLE_RATE / self.sample_rate)154            ).long()155 156        audio_cpu = audio.detach().cpu()157        features = []158 159        for index, valid_length_tensor in enumerate(audio_lengths):160            valid_length = int(valid_length_tensor.item())161            waveform = audio_cpu[index, :valid_length]162            if waveform.numel() == 0:163                waveform = audio_cpu.new_zeros(1)164            features.append(165                extract_speaker_fbank(166                    waveform,167                    sample_rate=_SPEAKER_FBANK_SAMPLE_RATE,168                )169            )170 171        fbank_lengths = torch.tensor(172            [feature.size(0) for feature in features],173            device=audio.device,174            dtype=torch.long,175        )176        fbank = pad_sequence(177            features,178            batch_first=True,179            padding_value=0.0,180        ).to(device=audio.device, dtype=audio.dtype)181        return fbank, fbank_lengths182 183    @torch.no_grad()184    @torch.autocast(enabled=False, device_type="cuda")185    def forward(186        self, audio, audio_lengths=None, fbank=None, fbank_lengths=None, **_kwargs187    ):188        self.model.eval()189        audio = audio.float()190        if audio.dim() == 3:191            if audio.size(1) != 1:192                raise ValueError(193                    f"Speaker encoder expects mono audio, got shape {tuple(audio.shape)}."194                )195            audio = audio[:, 0]196        elif audio.dim() != 2:197            raise ValueError(198                f"Speaker encoder expects a 2D or 3D audio tensor, got shape {tuple(audio.shape)}."199            )200 201        audio, original_audio_lengths, cropped_audio_lengths, starts = self._crop_audio(202            audio,203            audio_lengths=audio_lengths,204        )205 206        if fbank is None:207            fbank, fbank_lengths = self._extract_fbank_batch(208                audio,209                cropped_audio_lengths,210            )211        else:212            if not isinstance(fbank, torch.Tensor):213                raise TypeError("Speaker encoder expects `fbank` to be a torch.Tensor.")214            if fbank.dim() != 3 or fbank.size(0) != audio.size(0):215                raise ValueError(216                    f"Speaker encoder expects `fbank` with shape (B, T, F) and matching batch size, got {tuple(fbank.shape)}."217                )218            fbank, fbank_lengths = self._crop_fbank(219                fbank.to(device=audio.device, dtype=torch.float32),220                fbank_lengths,221                original_audio_lengths,222                cropped_audio_lengths,223                starts,224            )225 226        return self.model(fbank, lengths=fbank_lengths)227