Update app.py

app.py

@@ -21,6 +21,109 @@ speaker_embeddings = {
 }
 
 
+from datasets import load_dataset, Audio
+
+dataset = load_dataset(
+    "divakaivan/glaswegian_audio"
+)
+
+dataset = dataset.cast_column("audio", Audio(sampling_rate=16000))['train']
+
+from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech
+
+processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts")
+model = SpeechT5ForTextToSpeech.from_pretrained("microsoft/speecht5_tts")
+
+tokenizer = processor.tokenizer
+
+def extract_all_chars(batch):
+    all_text = " ".join(batch["transcription"])
+    vocab = list(set(all_text))
+    return {"vocab": [vocab], "all_text": [all_text]}
+
+vocabs = dataset.map(
+    extract_all_chars, 
+    batched=True, 
+    batch_size=-1, 
+    keep_in_memory=True, 
+    remove_columns=dataset.column_names,
+)
+
+dataset_vocab = set(vocabs["vocab"][0])
+tokenizer_vocab = {k for k,_ in tokenizer.get_vocab().items()}
+
+replacements = [
+    ('à', 'a'),
+    ('ç', 'c'),
+    ('è', 'e'),
+    ('ë', 'e'),
+    ('í', 'i'),
+    ('ï', 'i'),
+    ('ö', 'o'),
+    ('ü', 'u'),
+]
+
+def cleanup_text(inputs):
+    for src, dst in replacements:
+        inputs["transcription"] = inputs["transcription"].replace(src, dst)
+    return inputs
+
+dataset = dataset.map(cleanup_text)
+
+import os
+import torch
+from speechbrain.inference.speaker import EncoderClassifier
+
+spk_model_name = "speechbrain/spkrec-xvect-voxceleb"
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+speaker_model = EncoderClassifier.from_hparams(
+    source=spk_model_name, 
+    run_opts={"device": device}, 
+    savedir=os.path.join("/tmp", spk_model_name),
+)
+
+def create_speaker_embedding(waveform):
+    with torch.no_grad():
+        speaker_embeddings = speaker_model.encode_batch(torch.tensor(waveform))
+        speaker_embeddings = torch.nn.functional.normalize(speaker_embeddings, dim=2)
+        speaker_embeddings = speaker_embeddings.squeeze().cpu().numpy()
+    return speaker_embeddings
+
+
+def prepare_dataset(example):
+    # load the audio data; if necessary, this resamples the audio to 16kHz
+    audio = example["audio"]
+
+    # feature extraction and tokenization
+    example = processor(
+        text=example["transcription"],
+        audio_target=audio["array"], 
+        sampling_rate=audio["sampling_rate"],
+        return_attention_mask=False,
+    )
+
+    # strip off the batch dimension
+    example["labels"] = example["labels"][0]
+
+    # use SpeechBrain to obtain x-vector
+    example["speaker_embeddings"] = create_speaker_embedding(audio["array"])
+
+    return example
+
+processed_example = prepare_dataset(dataset[0])
+from transformers import SpeechT5HifiGan
+vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan")
+
+spectrogram = torch.tensor(processed_example["labels"])
+with torch.no_grad():
+    speech = vocoder(spectrogram)
+
+dataset = dataset.map(
+    prepare_dataset, remove_columns=dataset.column_names,
+)
+
+
 def predict(text, speaker):
     if len(text.strip()) == 0:
         return (16000, np.zeros(0).astype(np.int16))
@@ -31,23 +134,9 @@ def predict(text, speaker):
     input_ids = inputs["input_ids"]
     input_ids = input_ids[..., :model.config.max_text_positions]
 
-    if speaker == "Surprise Me!":
-        # load one of the provided speaker embeddings at random
-        idx = np.random.randint(len(speaker_embeddings))
-        key = list(speaker_embeddings.keys())[idx]
-        speaker_embedding = np.load(speaker_embeddings[key])
-
-        # randomly shuffle the elements
-        np.random.shuffle(speaker_embedding)
-
-        # randomly flip half the values
-        x = (np.random.rand(512) >= 0.5) * 1.0
-        x[x == 0] = -1.0
-        speaker_embedding *= x
-
-        #speaker_embedding = np.random.rand(512).astype(np.float32) * 0.3 - 0.15
-    else:
-        speaker_embedding = np.load(speaker_embeddings[speaker[:3]])
+    ### ### ### 
+    example = dataset["test"][11]
+    speaker_embeddings = torch.tensor(example["speaker_embeddings"]).unsqueeze(0)
 
     speaker_embedding = torch.tensor(speaker_embedding).unsqueeze(0)