Auxiliarytrinket commited on
Commit
773e132
1 Parent(s): d16278f

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +11 -10
app.py CHANGED
@@ -9,13 +9,12 @@ from transformers import SpeechT5ForTextToSpeech, SpeechT5HifiGan, SpeechT5Proce
9
  device = "cuda:0" if torch.cuda.is_available() else "cpu"
10
 
11
  # load speech translation checkpoint
12
- asr_pipe = pipeline("automatic-speech-recognition", model="openai/whisper-tiny", device=device) #тут изменил на tiny
13
 
14
- # load text-to-speech checkpoint and speaker embeddings
15
- processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts")
16
 
17
- model = SpeechT5ForTextToSpeech.from_pretrained("microsoft/speecht5_tts").to(device)
18
- vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan").to(device)
19
 
20
  embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
21
  speaker_embeddings = torch.tensor(embeddings_dataset[7306]["xvector"]).unsqueeze(0)
@@ -23,13 +22,15 @@ speaker_embeddings = torch.tensor(embeddings_dataset[7306]["xvector"]).unsqueeze
23
 
24
  def translate(audio):
25
  outputs = asr_pipe(audio, max_new_tokens=256, generate_kwargs={"task": "translate"})
26
- return outputs["text"]
27
 
28
 
29
- def synthesise(text):
30
- inputs = processor(text=text, return_tensors="pt")
31
- speech = model.generate_speech(inputs["input_ids"].to(device), speaker_embeddings.to(device), vocoder=vocoder)
32
- return speech.cpu()
 
 
33
 
34
 
35
  def speech_to_speech_translation(audio):
 
9
  device = "cuda:0" if torch.cuda.is_available() else "cpu"
10
 
11
  # load speech translation checkpoint
12
+ asr_pipe = pipeline("automatic-speech-recognition", model="openai/whisper-tiny", device=device) #тут изменил на tiny, потому что это более компактная модель
13
 
14
+ transl = pipeline("translation", model="Helsinki-NLP/opus-mt-en-ru")
 
15
 
16
+ model = VitsModel.from_pretrained("facebook/mms-tts-rus")
17
+ token = AutoTokenizer.from_pretrained("facebook/mms-tts-rus")
18
 
19
  embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
20
  speaker_embeddings = torch.tensor(embeddings_dataset[7306]["xvector"]).unsqueeze(0)
 
22
 
23
  def translate(audio):
24
  outputs = asr_pipe(audio, max_new_tokens=256, generate_kwargs={"task": "translate"})
25
+ return transl(outputs['text'])[0]['translation_text']
26
 
27
 
28
+ def synthesise(text: str, tokenizer: AutoTokenizer = token, model: VitsModel = model):
29
+ inputs = token(text, return_tensors="pt")
30
+ # print(inputs)
31
+ with torch.no_grad():
32
+ output = model(**inputs).waveform
33
+ return output.cpu()
34
 
35
 
36
  def speech_to_speech_translation(audio):