Update app.py

app.py

@@ -1,59 +1,108 @@
+# import gradio as gr
+# import numpy as np
+# import torch
+# from datasets import load_dataset
+# from transformers import SpeechT5ForTextToSpeech, SpeechT5HifiGan, SpeechT5Processor, pipeline
+
+
+# device = "cuda:0" if torch.cuda.is_available() else "cpu"
+
+# # load speech translation checkpoint
+# asr_pipe = pipeline("automatic-speech-recognition", model="openai/whisper-tiny", device=device)
+
+# # load text-to-speech checkpoint and speaker embeddings
+# model_id = "microsoft/speecht5_tts" #"Ellight/speecht5_finetuned_voxpopuli_nl"  # update with your model id
+# # pipe = pipeline("automatic-speech-recognition", model=model_id)
+# model = SpeechT5ForTextToSpeech.from_pretrained(model_id)
+# vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan")
+# embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation",trust_remote_code=True))
+# speaker_embeddings = torch.tensor(embeddings_dataset[7306]["xvector"]).unsqueeze(0)
+# # speaker_embeddings = torch.tensor(embeddings_dataset[7440]["xvector"]).unsqueeze(0)
+
+# processor = SpeechT5Processor.from_pretrained(model_id)
+
+# replacements = [
+#     ("à", "a"),
+#     ("ç", "c"),
+#     ("è", "e"),
+#     ("ë", "e"),
+#     ("í", "i"),
+#     ("ï", "i"),
+#     ("ö", "o"),
+#     ("ü", "u"),
+# ]
+
+# def cleanup_text(text):
+#     for src, dst in replacements:
+#         text = text.replace(src, dst)
+#     return text
+
+# def synthesize_speech(text):
+#     text = cleanup_text(text)
+#     inputs = processor(text=text, return_tensors="pt")
+#     speech = model.generate_speech(inputs["input_ids"].to(device), speaker_embeddings.to(device), vocoder=vocoder)
+
+#     return gr.Audio.update(value=(16000, speech.cpu().numpy()))
+
+# def translate(audio):
+#     outputs = asr_pipe(audio, max_new_tokens=256, generate_kwargs={"task": "transcribe", "language": "Dutch"})
+#     return outputs["text"]
+
+
+# def synthesise(text):
+#     text = cleanup_text(text)
+#     inputs = processor(text=text, return_tensors="pt")
+#     speech = model.generate_speech(inputs["input_ids"].to(device), speaker_embeddings.to(device), vocoder=vocoder)
+#     return speech.cpu()
+
+
+# def speech_to_speech_translation(audio):
+#     translated_text = translate(audio)
+#     synthesised_speech = synthesise(translated_text)
+#     synthesised_speech = (synthesised_speech.numpy() * 32767).astype(np.int16)
+#     return 16000, synthesised_speech
+
 import gradio as gr
 import numpy as np
 import torch
 from datasets import load_dataset
-from transformers import SpeechT5ForTextToSpeech, SpeechT5HifiGan, SpeechT5Processor, pipeline
 
+from transformers import SpeechT5ForTextToSpeech, SpeechT5HifiGan, SpeechT5Processor, pipeline
+from transformers import VitsModel, VitsTokenizer
 
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
 # load speech translation checkpoint
-asr_pipe = pipeline("automatic-speech-recognition", model="openai/whisper-tiny", device=device)
+asr_pipe = pipeline("automatic-speech-recognition", model="openai/whisper-base", device=device)
 
 # load text-to-speech checkpoint and speaker embeddings
-model_id = "microsoft/speecht5_tts" #"Ellight/speecht5_finetuned_voxpopuli_nl"  # update with your model id
-# pipe = pipeline("automatic-speech-recognition", model=model_id)
-model = SpeechT5ForTextToSpeech.from_pretrained(model_id)
-vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan")
-embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation",trust_remote_code=True))
-speaker_embeddings = torch.tensor(embeddings_dataset[7306]["xvector"]).unsqueeze(0)
-# speaker_embeddings = torch.tensor(embeddings_dataset[7440]["xvector"]).unsqueeze(0)
+# processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts")
+# model = SpeechT5ForTextToSpeech.from_pretrained("sanchit-gandhi/speecht5_tts_vox_nl").to(device)
 
-processor = SpeechT5Processor.from_pretrained(model_id)
+vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan").to(device)
+model = VitsModel.from_pretrained("Matthijs/mms-tts-nld")
+tokenizer = VitsTokenizer.from_pretrained("Matthijs/mms-tts-nld")
 
-replacements = [
-    ("à", "a"),
-    ("ç", "c"),
-    ("è", "e"),
-    ("ë", "e"),
-    ("í", "i"),
-    ("ï", "i"),
-    ("ö", "o"),
-    ("ü", "u"),
-]
-
-def cleanup_text(text):
-    for src, dst in replacements:
-        text = text.replace(src, dst)
-    return text
-
-def synthesize_speech(text):
-    text = cleanup_text(text)
-    inputs = processor(text=text, return_tensors="pt")
-    speech = model.generate_speech(inputs["input_ids"].to(device), speaker_embeddings.to(device), vocoder=vocoder)
+embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
+speaker_embeddings = torch.tensor(embeddings_dataset[7306]["xvector"]).unsqueeze(0)
 
-    return gr.Audio.update(value=(16000, speech.cpu().numpy()))
 
 def translate(audio):
-    outputs = asr_pipe(audio, max_new_tokens=256, generate_kwargs={"task": "transcribe", "language": "Dutch"})
+    outputs = asr_pipe(audio, max_new_tokens=256, generate_kwargs={"task": "transcribe", "language": "nl"})
     return outputs["text"]
 
-
 def synthesise(text):
-    text = cleanup_text(text)
-    inputs = processor(text=text, return_tensors="pt")
-    speech = model.generate_speech(inputs["input_ids"].to(device), speaker_embeddings.to(device), vocoder=vocoder)
+    inputs = tokenizer(text, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model(inputs["input_ids"])
+    speech = outputs.audio[0]
+    
     return speech.cpu()
+    
+# def synthesise(text):
+#     inputs = processor(text=text, return_tensors="pt", padding='max_length', truncation=True)
+#     speech = model.generate_speech(inputs["input_ids"].to(device), speaker_embeddings.to(device), vocoder=vocoder)
+#     return speech.cpu()
 
 
 def speech_to_speech_translation(audio):