build added

build/lib/deep_voice_cloning/cloning/config.json

@@ -0,0 +1,7 @@
+{
+    "en": {
+      "model_path": "microsoft/speecht5_tts",
+      "vocoder_name": "microsoft/speecht5_hifigan",
+      "speaker_model_name": "speechbrain/spkrec-xvect-voxceleb"
+    }
+}

build/lib/deep_voice_cloning/cloning/model.py

@@ -0,0 +1,57 @@
+import os
+import json
+from typing import Dict
+from pathlib import Path
+
+import numpy as np
+import torch
+from speechbrain.pretrained import EncoderClassifier
+from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech, SpeechT5HifiGan
+
+
+class CloningModel:
+    def __init__(self, config: Dict[str, Dict[str, str]] = None, lang: str = 'en'):
+        super(CloningModel, self).__init__()
+        if config is None:
+            self.speaker_embedding = None
+            with open(os.path.join(os.path.dirname(__file__), 'config.json')) as f:
+                self.config = json.load(f)[lang]
+        else:
+            self.config = config
+            self.speaker_embedding = torch.load(Path(self.config['model_path']) / "speaker_embedding.pt")[0]
+        self.processor = SpeechT5Processor.from_pretrained(self.config['model_path'])
+        self.model = SpeechT5ForTextToSpeech.from_pretrained(self.config['model_path'])
+        self.vocoder = SpeechT5HifiGan.from_pretrained(self.config['vocoder_name'])
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.speaker_model = EncoderClassifier.from_hparams(source=self.config['speaker_model_name'])
+        self.to(self.device)
+
+
+
+    def to(self, device: torch.device):
+        self.model = self.model.to(device)
+        self.vocoder = self.vocoder.to(device)
+
+    def save_pretrained(self, save_directory: str):
+        self.model.save_pretrained(save_directory)
+        self.processor.save_pretrained(save_directory)
+        torch.save(self.speaker_embedding, Path(save_directory) / "speaker_embedding.pt")
+
+    def forward(self, text: str) -> np.array:
+        # tokenize text
+        inputs = self.processor(text=text, return_tensors="pt")
+        # generate spectrogram using backbone model
+        spectrogram = self.model.generate_speech(inputs["input_ids"].to(self.device),
+                                                 self.speaker_embedding.to(self.device))
+        # decode spectrogram into waveform using vocoder
+        with torch.no_grad():
+            waveform_array = self.vocoder(spectrogram).detach().cpu().numpy()
+        return waveform_array
+
+    def create_speaker_embedding(self, waveform: torch.tensor) -> torch.tensor:
+        with torch.no_grad():
+            speaker_embeddings = self.speaker_model.encode_batch(waveform)
+            speaker_embeddings = torch.nn.functional.normalize(speaker_embeddings, dim=2)
+            self.speaker_embedding = speaker_embeddings
+            speaker_embeddings = speaker_embeddings.squeeze()
+        return speaker_embeddings

build/lib/deep_voice_cloning/data/collator.py

@@ -0,0 +1,45 @@
+import torch
+from typing import Any, Dict, List, Union
+
+
+class TTSDataCollatorWithPadding:
+
+    def __init__(self, model, processor):
+        self.model = model
+        self.processor = processor
+
+    def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:
+        input_ids = [{"input_ids": feature["input_ids"]} for feature in features]
+        label_features = [{"input_values": feature["labels"]} for feature in features]
+        speaker_features = [feature["speaker_embeddings"] for feature in features]
+
+        # collate the inputs and targets into a batch
+        batch = self.processor.pad(
+            input_ids=input_ids,
+            labels=label_features,
+            return_tensors="pt",
+        )
+
+        # replace padding with -100 to ignore loss correctly
+        batch["labels"] = batch["labels"].masked_fill(
+            batch.decoder_attention_mask.unsqueeze(-1).ne(1), -100
+        )
+
+        # not used during fine-tuning
+        del batch["decoder_attention_mask"]
+
+        # round down target lengths to multiple of reduction factor
+        if self.model.config.reduction_factor > 1:
+            target_lengths = torch.tensor([
+                len(feature["input_values"]) for feature in label_features
+            ])
+            target_lengths = target_lengths.new([
+                length - length % self.model.config.reduction_factor for length in target_lengths
+            ])
+            max_length = max(target_lengths)
+            batch["labels"] = batch["labels"][:, :max_length]
+
+        # add the speaker embeddings
+        batch["speaker_embeddings"] = torch.tensor(speaker_features)
+
+        return batch

build/lib/deep_voice_cloning/data/dataset.py

@@ -0,0 +1,63 @@
+from typing import Dict, Any
+
+import torch
+import librosa
+import numpy as np
+from datasets import Dataset
+
+from ..cloning.model import CloningModel
+from ..transcriber.model import TranscriberModel
+
+
+def prepare_dataset(example: Dict[str, Any], model: CloningModel) -> Dict[str, Any]:
+    """
+    Prepare a single example for training
+    """
+    # feature extraction and tokenization
+    processed_example = model.processor(
+        text=example["normalized_text"],
+        audio_target=example["audio"]["array"],
+        sampling_rate=16000,
+        return_attention_mask=False,
+    )
+
+    # strip off the batch dimension
+    if len(torch.tensor(processed_example['input_ids']).shape) > 1:
+        processed_example['input_ids'] = processed_example['input_ids'][0]
+
+    processed_example["labels"] = processed_example["labels"][0]
+
+    # use SpeechBrain to obtain x-vector
+    processed_example["speaker_embeddings"] = model.create_speaker_embedding(
+        torch.tensor(example["audio"]["array"])
+    ).numpy()
+
+    return processed_example
+
+
+def get_cloning_dataset(input_audio_path: str,
+                        transcriber_model: TranscriberModel,
+                        cloning_model: CloningModel,
+                        sampling_rate: int = 16000,
+                        window_size_secs: int = 5) -> Dataset:
+    """
+    Create dataset by transcribing an audio file using a pretrained Wav2Vec2 model.
+    """
+    speech_array, _ = librosa.load(input_audio_path, sr=sampling_rate)
+
+    # split a waveform into splits of 5 secs each
+    speech_arrays = np.split(speech_array, range(0, len(speech_array), window_size_secs * sampling_rate))[1:]
+    texts = [transcriber_model.forward(speech_array, sampling_rate=sampling_rate)
+             for speech_array in speech_arrays]
+
+    dataset = Dataset.from_list([
+        {'audio': {'array': speech_arrays[i]}, 'normalized_text': texts[i]}
+        for i in range(len(speech_arrays))]
+    )
+
+    dataset = dataset.map(
+        prepare_dataset, fn_kwargs={'model': cloning_model},
+        remove_columns=dataset.column_names,
+    )
+
+    return dataset

build/lib/deep_voice_cloning/transcriber/config.json

@@ -0,0 +1,7 @@
+{
+    "language_model_names": {
+        "en": "jonatasgrosman/wav2vec2-large-xlsr-53-english",
+        "fr": "jonatasgrosman/wav2vec2-large-xlsr-53-french",
+        "de": "jonatasgrosman/wav2vec2-large-xlsr-53-german"
+    }
+}

build/lib/deep_voice_cloning/transcriber/model.py

@@ -0,0 +1,22 @@
+import os
+import json
+
+import numpy as np
+import torch
+from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC
+
+
+class TranscriberModel:
+    def __init__(self, lang: str = 'en'):
+        with open(os.path.join(os.path.dirname(__file__), 'config.json')) as f:
+            config = json.load(f)
+        self.processor = Wav2Vec2Processor.from_pretrained(config['language_model_names'][lang])
+        self.model = Wav2Vec2ForCTC.from_pretrained(config['language_model_names'][lang])
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    def forward(self, speech_array: np.array, sampling_rate: int = 16000) -> str:
+        model_input = self.processor(speech_array, sampling_rate=sampling_rate, return_tensors="pt", padding=True)
+        with torch.no_grad():
+            logits = self.model(model_input.input_values, attention_mask=model_input.attention_mask).logits
+            predicted_ids = torch.argmax(logits, dim=-1)
+        return self.processor.batch_decode(predicted_ids)