Update app.py

app.py

@@ -1,65 +1,68 @@
+import os
 import gradio as gr
-from gradio import components
 import whisper
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from gtts import gTTS
+import sentencepiece
+import sounddevice as sd
 import soundfile as sf
-import os
-import numpy as np
-
-def translate_speech_to_speech(input_tuple):
-    input_audio, sample_rate = input_tuple
-    
-    # Save the input audio to a temporary file
-    input_file = "input_audio.wav"
-    sf.write(input_file, input_audio, sample_rate)  # use the sample rate from Gradio
-
-    
-    # Language detection and translation code from the first code snippet
-    model = whisper.load_model("base")
-    audio = whisper.load_audio(input_file)
-    audio = whisper.pad_or_trim(audio)
-    mel = whisper.log_mel_spectrogram(audio).to(model.device)
-    _, probs = model.detect_language(mel)
-    
-    options = whisper.DecodingOptions()
-    result = whisper.decode(model, mel, options)
-    
-    text = result.text
-    lang = max(probs, key=probs.get)
-    
-    # Translation code from the first code snippet
-    to_lang = 'ru'
-    tokenizer = AutoTokenizer.from_pretrained("alirezamsh/small100")
-    model = AutoModelForSeq2SeqLM.from_pretrained("alirezamsh/small100")
-    
-    tokenizer.src_lang = lang
-    encoded_bg = tokenizer(text, return_tensors="pt")
-    generated_tokens = model.generate(**encoded_bg)
-    translated_text = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
-    
-    # Text-to-speech (TTS) code from the first code snippet
-    tts = gTTS(text=translated_text, lang=to_lang)
-    output_file = "translated_speech.wav"
-    tts.save(output_file)
-    
-    # Load the translated audio and return as an output
-    translated_audio, _ = sf.read(output_file, dtype="int16")
-    
-    return translated_audio
-
-title = "Speech-to-Speech Translator"
-
-input_audio = gr.inputs.Audio(source="microphone")
-output_audio = gr.outputs.Audio(type="numpy")
-
-stt_demo = gr.Interface(
-    fn=translate_speech_to_speech,
-    inputs=input_audio,
-    outputs=output_audio,
-    title=title,
-    description="Speak in any language, and the translator will convert it to speech in the target language.",
-)
+import tempfile
+
+
+def translate_voice(audio, target_lang):
+    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_audio:
+        temp_filename = temp_audio.name
+        sf.write(temp_filename, audio, 16000)
+
+        model = whisper.load_model("base").float()
+
+        audio = whisper.load_audio(temp_filename)
+
+        audio = whisper.pad_or_trim(audio)
+
+        mel = whisper.log_mel_spectrogram(audio).to(model.device).float()
+
+        _, probs = model.detect_language(mel)
+        options = whisper.DecodingOptions(fp16=False)
+        result = whisper.decode(model, mel, options)
 
-if __name__ == "__main__":
-    stt_demo.launch()
+        text = result.text
+        lang = max(probs, key=probs.get)
+
+        tokenizer = AutoTokenizer.from_pretrained("alirezamsh/small100")
+        model = AutoModelForSeq2SeqLM.from_pretrained("alirezamsh/small100")
+
+        tokenizer.src_lang = target_lang
+        encoded_bg = tokenizer(text, return_tensors="pt")
+        generated_tokens = model.generate(**encoded_bg)
+        translated_text = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
+
+        tts = gTTS(text=translated_text, lang=target_lang)
+        filename = "to_speech.mp3"
+        tts.save(filename)
+
+        return filename, text, translated_text, target_lang
+
+
+def record_audio():
+    fs = 16000
+    duration = 5  # Record audio for 5 seconds, you can adjust the duration as needed
+    audio = sd.rec(int(duration * fs), samplerate=fs, channels=1)
+    sd.wait()
+    return audio.flatten()
+
+
+iface = gr.Interface(
+    fn=translate_voice,
+    inputs=[
+        gr.inputs.Audio(type="microphone", label="Speak"),
+        gr.inputs.Dropdown(choices=['en', 'ru', 'de', 'fr'], label="Target Language")
+    ],
+    outputs=[
+        gr.outputs.Audio(type="filepath", label="Translated Audio"),
+        gr.outputs.Textbox(label="Original Text"),
+        gr.outputs.Textbox(label="Translated Text"),
+        gr.outputs.Textbox(label="Target Language"),
+    ]
+)
+iface.launch()