app.py

import os
import tempfile
import sys
import subprocess
import gradio as gr
import numpy as np
import soundfile as sf
import librosa
import torch
import torch.cuda
import gc

# Check if required packages are installed, if not install them
try:
    from espnet2.bin.s2t_inference import Speech2Text
    import torchaudio
    # Try importing espnet_model_zoo specifically
    try:
        import espnet_model_zoo
        print("All packages already installed.")
    except ModuleNotFoundError:
        print("Installing espnet_model_zoo. This may take a few minutes...")
        subprocess.check_call([sys.executable, "-m", "pip", "install", "-U", "espnet_model_zoo"])
        import espnet_model_zoo
        print("espnet_model_zoo installed successfully.")
except ModuleNotFoundError as e:
    missing_module = str(e).split("'")[1]
    print(f"Installing missing module: {missing_module}")
    
    if missing_module == "espnet2":
        print("Installing ESPnet. This may take a few minutes...")
        subprocess.check_call([sys.executable, "-m", "pip", "install", "espnet"])
    elif missing_module == "torchaudio":
        print("Installing torchaudio. This may take a few minutes...")
        subprocess.check_call([sys.executable, "-m", "pip", "install", "torchaudio"])
    
    # Try importing again
    try:
        from espnet2.bin.s2t_inference import Speech2Text
        import torchaudio
        # Also check for espnet_model_zoo
        try:
            import espnet_model_zoo
        except ModuleNotFoundError:
            print("Installing espnet_model_zoo. This may take a few minutes...")
            subprocess.check_call([sys.executable, "-m", "pip", "install", "-U", "espnet_model_zoo"])
            import espnet_model_zoo
        print("All required packages installed successfully.")
    except ModuleNotFoundError as e:
        print(f"Failed to install {str(e).split('No module named ')[1]}. Please install manually.")
        raise

# Initialize the model with language option
def load_model():
    # Force garbage collection
    gc.collect()
    torch.cuda.empty_cache()
    
    # Set memory-efficient options
    torch.cuda.set_per_process_memory_fraction(0.95)  # Use 95% of available memory
    
    # Check if CUDA is available
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Using device: {device}")
    
    # For memory efficiency, you could try loading with 8-bit quantization
    # This requires the bitsandbytes library
    # pip install bitsandbytes
    
    model = Speech2Text.from_pretrained(
        "espnet/owls_4B_180K",
        task_sym="<asr>",
        beam_size=1,
        device=device
    )
    return model

# Load the model at startup with English as default
print("Loading multilingual model...")
model = load_model()
print("Model loaded successfully!")

def transcribe_audio(audio_file, language):
    """Process the audio file and return the transcription"""
    if audio_file is None:
        return "Please upload an audio file or record audio."
    
    # If audio is a tuple (from microphone recording)
    if isinstance(audio_file, tuple):
        sr, audio_data = audio_file
        # Create a temporary file to save the audio
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_audio:
            temp_path = temp_audio.name
            sf.write(temp_path, audio_data, sr)
            audio_file = temp_path
    
    # Load and resample the audio file to 16kHz
    speech, _ = librosa.load(audio_file, sr=16000)
    
    # Update the language symbol if needed
    model.beam_search.hyps = None
    model.beam_search.pre_beam_score_key = None

    if language != None:
        model.lang_sym = language
    
    # Perform ASR
    text, *_ = model(speech)[0]
    
    # Clean up temporary file if created
    if isinstance(audio_file, str) and audio_file.startswith(tempfile.gettempdir()):
        os.unlink(audio_file)
    
    return text

# Function to handle English transcription
def transcribe_english(audio_file):
    return transcribe_audio(audio_file, "<eng>")

# Function to handle Chinese transcription
def transcribe_chinese(audio_file):
    return transcribe_audio(audio_file, "<zho>")

# Function to handle Japanese transcription
def transcribe_japanese(audio_file):
    return transcribe_audio(audio_file, "<jpn>")

# Function to handle Korean transcription
def transcribe_korean(audio_file):
    return transcribe_audio(audio_file, "<kor>")

# Function to handle Thai transcription
def transcribe_thai(audio_file):
    return transcribe_audio(audio_file, "<tha>")

# Function to handle Italian transcription
def transcribe_italian(audio_file):
    return transcribe_audio(audio_file, "<ita>")

# Function to handle German transcription
def transcribe_german(audio_file):
    return transcribe_audio(audio_file, "<deu>")

# Create the Gradio interface with tabs
demo = gr.Blocks(title="NVIDIA Research Multilingual Demo")

with demo:
    gr.Markdown("# NVIDIA Research Multilingual Demo")
    gr.Markdown("Upload or record audio to transcribe up to 150 human languages using the NVIDIA Research (NVR) 9B model. Audio will be automatically resampled to 16kHz.")
    
    with gr.Tabs():
        with gr.TabItem("Microphone Recording"):
            language_mic = gr.Radio(
                ["English", "Mandarin", "Japanese", "Korean", "Thai", "Italian", "German"], 
                label="Select Language", 
                value="English"
            )
            
            with gr.Row():
                with gr.Column():
                    mic_input = gr.Audio(sources=["microphone"], type="filepath", label="Record Audio")
                    mic_button = gr.Button("Transcribe Recording")
                with gr.Column():
                    mic_output = gr.Textbox(label="Transcription")
            
            def transcribe_mic(audio, lang):
                lang_map = {
                    "English": "<eng>", 
                    "Chinese": "<zho>", 
                    "Japanese": "<jpn>", 
                    "Korean": "<kor>",
                    "Thai": "<tha>",
                    "Italian": "<ita>",
                    "German": "<deu>"
                }
                return transcribe_audio(audio, lang_map.get(lang, "<eng>"))
            
            mic_button.click(fn=transcribe_mic, inputs=[mic_input, language_mic], outputs=mic_output)
        
        with gr.TabItem("English"):
            with gr.Row():
                with gr.Column():
                    en_input = gr.Audio(sources=["upload"], type="filepath", label="Upload Audio")
                    en_button = gr.Button("Transcribe Speech")
                with gr.Column():
                    en_output = gr.Textbox(label="Speech Transcription")
            
            # Add example if the file exists
            if os.path.exists("wav_en_sample_48k.wav"):
                gr.Examples(
                    examples=[["wav_en_sample_48k.wav"]],
                    inputs=en_input
                )
            
            en_button.click(fn=transcribe_english, inputs=en_input, outputs=en_output)
            
        with gr.TabItem("Mandarin"):
            with gr.Row():
                with gr.Column():
                    zh_input = gr.Audio(sources=["upload"], type="filepath", label="Upload Audio")
                    zh_button = gr.Button("Transcribe Speech")
                with gr.Column():
                    zh_output = gr.Textbox(label="Speech Transcription")
            
            # Add example if the file exists
            if os.path.exists("wav_zh_tw_sample_16k.wav"):
                gr.Examples(
                    examples=[["wav_zh_tw_sample_16k.wav"]],
                    inputs=zh_input
                )
            
            zh_button.click(fn=transcribe_chinese, inputs=zh_input, outputs=zh_output)
            
        with gr.TabItem("Japanese"):
            with gr.Row():
                with gr.Column():
                    jp_input = gr.Audio(sources=["upload"], type="filepath", label="Upload Audio")
                    jp_button = gr.Button("Transcribe Speech")
                with gr.Column():
                    jp_output = gr.Textbox(label="Speech Transcription")
            
            # Add example if the file exists
            if os.path.exists("wav_jp_sample_48k.wav"):
                gr.Examples(
                    examples=[["wav_jp_sample_48k.wav"]],
                    inputs=jp_input
                )
            
            jp_button.click(fn=transcribe_japanese, inputs=jp_input, outputs=jp_output)
            
        with gr.TabItem("Korean"):
            with gr.Row():
                with gr.Column():
                    kr_input = gr.Audio(sources=["upload"], type="filepath", label="Upload Audio")
                    kr_button = gr.Button("Transcribe Speech")
                with gr.Column():
                    kr_output = gr.Textbox(label="Speech Transcription")
            
            # Add example if the file exists
            if os.path.exists("wav_kr_sample_48k.wav"):
                gr.Examples(
                    examples=[["wav_kr_sample_48k.wav"]],
                    inputs=kr_input
                )
            
            kr_button.click(fn=transcribe_korean, inputs=kr_input, outputs=kr_output)
            
        with gr.TabItem("Thai"):
            with gr.Row():
                with gr.Column():
                    th_input = gr.Audio(sources=["upload"], type="filepath", label="Upload Audio")
                    th_button = gr.Button("Transcribe Speech")
                with gr.Column():
                    th_output = gr.Textbox(label="Speech Transcription")
            
            # Add example if the file exists
            if os.path.exists("wav_thai_sample.wav"):
                gr.Examples(
                    examples=[["wav_thai_sample.wav"]],
                    inputs=th_input
                )
            
            th_button.click(fn=transcribe_thai, inputs=th_input, outputs=th_output)
            
        with gr.TabItem("Italian"):
            with gr.Row():
                with gr.Column():
                    it_input = gr.Audio(sources=["upload"], type="filepath", label="Upload Audio")
                    it_button = gr.Button("Transcribe Speech")
                with gr.Column():
                    it_output = gr.Textbox(label="Speech Transcription")
            
            # Add example if the file exists
            if os.path.exists("wav_it_sample.wav"):
                gr.Examples(
                    examples=[["wav_it_sample.wav"]],
                    inputs=it_input
                )
            
            it_button.click(fn=transcribe_italian, inputs=it_input, outputs=it_output)

        with gr.TabItem("German"):
            with gr.Row():
                with gr.Column():
                    de_input = gr.Audio(sources=["upload"], type="filepath", label="Upload Audio")
                    de_button = gr.Button("Transcribe Speech")
                with gr.Column():
                    de_output = gr.Textbox(label="Speech Transcription")
            
            # Add example if the file exists
            if os.path.exists("wav_de_sample.wav"):
                gr.Examples(
                    examples=[["wav_de_sample.wav"]],
                    inputs=de_input
                )
            
            de_button.click(fn=transcribe_german, inputs=de_input, outputs=de_output)

# Launch the app with Hugging Face Spaces compatible settings
if __name__ == "__main__":
    demo.launch(share=False)