app.py

import io
import whisper
import torch
import ffmpeg
import torchaudio
import streamlit as st
LANGUAGES = {
 "en":"english",
 "zh":"chinese",
 "de":"german",
 "es":"spanish",
 "ru":"russian",
 "ko":"korean",
 "fr":"french",
 "ja":"japanese",
 "pt":"portuguese",
 "tr":"turkish",
 "pl":"polish",
 "ca":"catalan",
 "nl":"dutch",
 "ar":"arabic",
 "sv":"swedish",
 "it":"italian",
 "id":"indonesian",
 "hi":"hindi",
 "fi":"finnish",
 "vi":"vietnamese",
 "iw":"hebrew",
 "uk":"ukrainian",
 "el":"greek",
 "ms":"malay",
 "cs":"czech",
 "ro":"romanian",
 "da":"danish",
 "hu":"hungarian",
 "ta":"tamil",
 "no":"norwegian",
 "th":"thai",
 "ur":"urdu",
 "hr":"croatian",
 "bg":"bulgarian",
 "lt":"lithuanian",
 "la":"latin",
 "mi":"maori",
 "ml":"malayalam",
 "cy":"welsh",
 "sk":"slovak",
 "te":"telugu",
 "fa":"persian",
 "lv":"latvian",
 "bn":"bengali",
 "sr":"serbian",
 "az":"azerbaijani",
 "sl":"slovenian",
 "kn":"kannada",
 "et":"estonian",
 "mk":"macedonian",
 "br":"breton",
 "eu":"basque",
 "is":"icelandic",
 "hy":"armenian",
 "ne":"nepali",
 "mn":"mongolian",
 "bs":"bosnian",
 "kk":"kazakh",
 "sq":"albanian",
 "sw":"swahili",
 "gl":"galician",
 "mr":"marathi",
 "pa":"punjabi",
 "si":"sinhala",
 "km":"khmer",
 "sn":"shona",
 "yo":"yoruba",
 "so":"somali",
 "af":"afrikaans",
 "oc":"occitan",
 "ka":"georgian",
 "be":"belarusian",
 "tg":"tajik",
 "sd":"sindhi",
 "gu":"gujarati",
 "am":"amharic",
 "yi":"yiddish",
 "lo":"lao",
 "uz":"uzbek",
 "fo":"faroese",
 "ht":"haitian creole",
 "ps":"pashto",
 "tk":"turkmen",
 "nn":"nynorsk",
 "mt":"maltese",
 "sa":"sanskrit",
 "lb":"luxembourgish",
 "my":"myanmar",
 "bo":"tibetan",
 "tl":"tagalog",
 "mg":"malagasy",
 "as":"assamese",
 "tt":"tatar",
 "haw":"hawaiian",
 "ln":"lingala",
 "ha":"hausa",
 "ba":"bashkir",
 "jw":"javanese",
 "su":"sundanese",
}

def decode(model, mel, options):
 result = whisper.decode(model, mel, options)
 return result.text

def load_audio(audio):
 print(audio.type)
 if audio.type == "audio/wav" or audio.type == "audio/flac":
 wave, sr = torchaudio.load(audio)
 if sr != 16000:
 wave = torchaudio.transforms.Resample(sr, 16000)(wave)
 return wave.squeeze(0)

 elif audio.type == "audio/mpeg":
 audio = audio.read()
 audio, _ = (ffmpeg
 .input('pipe:0')
 .output('pipe:1', format='wav', acodec='pcm_s16le', ac=1, ar='16k')
 .run(capture_stdout=True, input=audio)
 )
 audio = io.BytesIO(audio)
 wave, sr = torchaudio.load(audio)
 if sr != 16000:
 wave = torchaudio.transforms.Resample(sr, 16000)(wave)
 return wave.squeeze(0)

 else:
 st.error("Unsupported audio format")

def detect_language(model, mel):
 _, probs = model.detect_language(mel)
 return max(probs, key=probs.get)

def main():

 st.title("Whisper ASR Demo")
 st.markdown(
 """
 This is a demo of OpenAI's Whisper ASR model. The model is trained on 680,000 hours of dataset. 
 """
 )

 model_selection = st.sidebar.selectbox("Select model", ["tiny", "base", "small", "medium", "large"])
 en_model_selection = st.sidebar.checkbox("English only model", value=False)

 if en_model_selection:
 model_selection += ".en"
 st.sidebar.write(f"Model: {model_selection+' (Multilingual)' if not en_model_selection else model_selection + ' (English only)'}")

 if st.sidebar.checkbox("Show supported languages", value=False):
 st.sidebar.info(list(LANGUAGES.values()))
 st.sidebar.title("Options")
 
 beam_size = st.sidebar.slider("Beam Size", min_value=1, max_value=10, value=5)
 fp16 = st.sidebar.checkbox("Enable FP16 for faster transcription (It may affect performance)", value=False)

 if not en_model_selection:
 task = st.sidebar.selectbox("Select task", ["transcribe", "translate (To English)"], index=0)
 else:
 task = st.sidebar.selectbox("Select task", ["transcribe"], index=0)

 st.title("Audio")
 audio_file = st.file_uploader("Upload Audio", type=["wav", "mp3", "flac"])

 if audio_file is not None:
 st.audio(audio_file, format=audio_file.type)
 with st.spinner("Loading model..."):
 model = whisper.load_model(model_selection)
 model = model.to("cpu") if not torch.cuda.is_available() else model.to("cuda")
 

 audio = load_audio(audio_file)
 with st.spinner("Extracting features..."):
 audio = whisper.pad_or_trim(audio)
 mel = whisper.log_mel_spectrogram(audio).to(model.device)
 if not en_model_selection:
 with st.spinner("Detecting language..."):
 language = detect_language(model, mel)
 st.markdown(f"Detected Language: {LANGUAGES[language]} ({language})")
 else:
 language = "en"
 configuration = {"beam_size": beam_size, "fp16": fp16, "task": task, "language": language}
 with st.spinner("Transcribing..."):
 options = whisper.DecodingOptions(**configuration)
 text = decode(model, mel, options)
 st.markdown(f"**Recognized Text:** {text}")

if __name__ == "__main__":
 main()