Update app.py

app.py

@@ -8,21 +8,27 @@ from transformers import BarkModel, BarkProcessor
 
 from transformers import Speech2TextProcessor, Speech2TextForConditionalGeneration
 
+SAMPLE_RATE = 16000
+
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
-asr_model = Speech2TextForConditionalGeneration.from_pretrained("facebook/s2t-medium-mustc-multilingual-st")
-asr_processor = Speech2TextProcessor.from_pretrained("facebook/s2t-medium-mustc-multilingual-st")
+# asr_model = Speech2TextForConditionalGeneration.from_pretrained("facebook/s2t-medium-mustc-multilingual-st")
+# asr_processor = Speech2TextProcessor.from_pretrained("facebook/s2t-medium-mustc-multilingual-st")
+
+asr_pipe = pipeline("automatic-speech-recognition", model="openai/whisper-base", device=device)
+
 
 bark_model = BarkModel.from_pretrained("suno/bark")
 bark_processor = BarkProcessor.from_pretrained("suno/bark")
 
 
 def translate(audio):
-    inputs = asr_processor(audio, sampling_rate=16000, return_tensors="pt")
-    generated_ids = asr_model.generate(inputs["input_features"],attention_mask=inputs["attention_mask"],
-    forced_bos_token_id=asr_processor.tokenizer.lang_code_to_id["it"],)
-    translation = asr_processor.batch_decode(generated_ids, skip_special_tokens=True)
-    return translation
+    # inputs = asr_processor(audio, sampling_rate=16000, return_tensors="pt")
+    # generated_ids = asr_model.generate(inputs["input_features"],attention_mask=inputs["attention_mask"],
+    # forced_bos_token_id=asr_processor.tokenizer.lang_code_to_id["it"],)
+    # translation = asr_processor.batch_decode(generated_ids, skip_special_tokens=True)
+    translation = asr_pipe(audio, max_new_tokens=256, generate_kwargs={"task": "transcribe",  "language": "it"})
+    return translation["text"]
 
 
 def synthesise(text):
@@ -35,7 +41,7 @@ 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
+    return SAMPLE_RATE, synthesised_speech
 
 
 title = "Cascaded STST"