Spaces:

gobeldan
/

insanely-fast-whisper-webui

Running

App Files Files Community

gobeldan commited on Jan 30, 2024

Commit

3d58577

verified ·

1 Parent(s): 82c30a2

Upload 4 files

Browse files

Files changed (4) hide show

app.py ADDED Viewed

	@@ -0,0 +1,172 @@

+import gradio as gr
+import time
+import logging
+import torch
+from sys import platform
+from transformers import pipeline, AutoModelForSpeechSeq2Seq, AutoProcessor
+from transformers.utils import is_flash_attn_2_available
+from languages import get_language_names
+from subtitle_manager import Subtitle
+logging.basicConfig(level=logging.INFO)
+last_model = None
+def write_file(output_file,subtitle):
+ with open(output_file, 'w', encoding='utf-8') as f:
+ f.write(subtitle)
+def create_pipe(model, flash):
+ if torch.cuda.is_available():
+ device = "cuda:0"
+ elif platform == "darwin":
+ device = "mps"
+ else:
+ device = "cpu"
+ torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+ model_id = model
+ model = AutoModelForSpeechSeq2Seq.from_pretrained(
+ model_id,
+ torch_dtype=torch_dtype,
+ low_cpu_mem_usage=True,
+ use_safetensors=True,
+ attn_implementation="flash_attention_2" if flash and is_flash_attn_2_available() else "sdpa",
+ # eager (manual attention implementation)
+ # flash_attention_2 (implementation using flash attention 2)
+ # sdpa (implementation using torch.nn.functional.scaled_dot_product_attention)
+ # PyTorch SDPA requirements in Transformers are not met. Please install torch>=2.1.1.
+ )
+ model.to(device)
+ processor = AutoProcessor.from_pretrained(model_id)
+ pipe = pipeline(
+ "automatic-speech-recognition",
+ model=model,
+ tokenizer=processor.tokenizer,
+ feature_extractor=processor.feature_extractor,
+ # max_new_tokens=128,
+ # chunk_length_s=15,
+ # batch_size=16,
+ torch_dtype=torch_dtype,
+ device=device,
+ )
+ return pipe
+def transcribe_webui_simple_progress(modelName, languageName, urlData, multipleFiles, microphoneData, task, flash,
+ chunk_length_s, batch_size, progress=gr.Progress()):
+ global last_model
+ progress(0, desc="Loading Audio..")
+ logging.info(f"urlData:{urlData}")
+ logging.info(f"multipleFiles:{multipleFiles}")
+ logging.info(f"microphoneData:{microphoneData}")
+ logging.info(f"task: {task}")
+ logging.info(f"is_flash_attn_2_available: {is_flash_attn_2_available()}")
+ logging.info(f"chunk_length_s: {chunk_length_s}")
+ logging.info(f"batch_size: {batch_size}")
+ if last_model == None:
+ logging.info("first model")
+ progress(0.1, desc="Loading Model..")
+ pipe = create_pipe(modelName, flash)
+ elif modelName != last_model:
+ logging.info("new model")
+ torch.cuda.empty_cache()
+ progress(0.1, desc="Loading Model..")
+ pipe = create_pipe(modelName, flash)
+ else:
+ logging.info("Model not changed")
+ last_model = modelName
+ srt_sub = Subtitle("srt")
+ vtt_sub = Subtitle("vtt")
+ txt_sub = Subtitle("txt")
+ files = []
+ if multipleFiles:
+ files+=multipleFiles
+ if urlData:
+ files.append(urlData)
+ if microphoneData:
+ files.append(microphoneData)
+ logging.info(files)
+ generate_kwargs = {}
+ if languageName != "Automatic Detection" and modelName.endswith(".en") == False:
+ generate_kwargs["language"] = languageName
+ if modelName.endswith(".en") == False:
+ generate_kwargs["task"] = task
+ files_out = []
+ for file in progress.tqdm(files, desc="Working..."):
+ start_time = time.time()
+ logging.info(file)
+ outputs = pipe(
+ file,
+ chunk_length_s=chunk_length_s,#30
+ batch_size=batch_size,#24
+ generate_kwargs=generate_kwargs,
+ return_timestamps=True,
+ )
+ logging.debug(outputs)
+ logging.info(print(f"transcribe: {time.time() - start_time} sec."))
+ file_out = file.split('/')[-1]
+ srt = srt_sub.get_subtitle(outputs["chunks"])
+ vtt = vtt_sub.get_subtitle(outputs["chunks"])
+ txt = txt_sub.get_subtitle(outputs["chunks"])
+ write_file(file_out+".srt",srt)
+ write_file(file_out+".vtt",vtt)
+ write_file(file_out+".txt",txt)
+ files_out += [file_out+".srt", file_out+".vtt", file_out+".txt"]
+ progress(1, desc="Completed!")
+ return files_out, vtt, txt
+with gr.Blocks(title="Insanely Fast Whisper") as demo:
+ description = "An opinionated CLI to transcribe Audio files w/ Whisper on-device! Powered by 🤗 Transformers, Optimum & flash-attn"
+ article = "Read the [documentation here](https://github.com/Vaibhavs10/insanely-fast-whisper#cli-options)."
+ whisper_models = [
+ "openai/whisper-tiny", "openai/whisper-tiny.en",
+ "openai/whisper-base", "openai/whisper-base.en",
+ "openai/whisper-small", "openai/whisper-small.en", "distil-whisper/distil-small.en",
+ "openai/whisper-medium", "openai/whisper-medium.en", "distil-whisper/distil-medium.en",
+ "openai/whisper-large",
+ "openai/whisper-large-v1",
+ "openai/whisper-large-v2", "distil-whisper/distil-large-v2",
+ "openai/whisper-large-v3", "xaviviro/whisper-large-v3-catalan-finetuned-v2",
+ ]
+ waveform_options=gr.WaveformOptions(
+ waveform_color="#01C6FF",
+ waveform_progress_color="#0066B4",
+ skip_length=2,
+ show_controls=False,
+ )
+ simple_transcribe = gr.Interface(fn=transcribe_webui_simple_progress,
+ description=description,
+ article=article,
+ inputs=[
+ gr.Dropdown(choices=whisper_models, value="distil-whisper/distil-large-v2", label="Model", info="Select whisper model", interactive = True,),
+ gr.Dropdown(choices=["Automatic Detection"] + sorted(get_language_names()), value="Automatic Detection", label="Language", info="Select audio voice language", interactive = True,),
+ gr.Text(label="URL", info="(YouTube, etc.)", interactive = True),
+ gr.File(label="Upload Files", file_count="multiple"),
+ gr.Audio(sources=["microphone"], type="filepath", label="Microphone Input", waveform_options = waveform_options),
+ gr.Dropdown(choices=["transcribe", "translate"], label="Task", value="transcribe", interactive = True),
+ gr.Checkbox(label='Flash',info='Use Flash Attention 2'),
+ gr.Number(label='chunk_length_s',value=30, interactive = True),
+ gr.Number(label='batch_size',value=24, interactive = True)
+ ], outputs=[
+ gr.File(label="Download"),
+ gr.Text(label="Transcription"),
+ gr.Text(label="Segments")
+ ]
+ )
+if __name__ == "__main__":
+ demo.launch()

languages.py ADDED Viewed

	@@ -0,0 +1,147 @@

+class Language():
+ def __init__(self, code, name):
+ self.code = code
+ self.name = name
+ def __str__(self):
+ return "Language(code={}, name={})".format(self.code, self.name)
+LANGUAGES = [
+ Language('en', 'English'),
+ Language('zh', 'Chinese'),
+ Language('de', 'German'),
+ Language('es', 'Spanish'),
+ Language('ru', 'Russian'),
+ Language('ko', 'Korean'),
+ Language('fr', 'French'),
+ Language('ja', 'Japanese'),
+ Language('pt', 'Portuguese'),
+ Language('tr', 'Turkish'),
+ Language('pl', 'Polish'),
+ Language('ca', 'Catalan'),
+ Language('nl', 'Dutch'),
+ Language('ar', 'Arabic'),
+ Language('sv', 'Swedish'),
+ Language('it', 'Italian'),
+ Language('id', 'Indonesian'),
+ Language('hi', 'Hindi'),
+ Language('fi', 'Finnish'),
+ Language('vi', 'Vietnamese'),
+ Language('he', 'Hebrew'),
+ Language('uk', 'Ukrainian'),
+ Language('el', 'Greek'),
+ Language('ms', 'Malay'),
+ Language('cs', 'Czech'),
+ Language('ro', 'Romanian'),
+ Language('da', 'Danish'),
+ Language('hu', 'Hungarian'),
+ Language('ta', 'Tamil'),
+ Language('no', 'Norwegian'),
+ Language('th', 'Thai'),
+ Language('ur', 'Urdu'),
+ Language('hr', 'Croatian'),
+ Language('bg', 'Bulgarian'),
+ Language('lt', 'Lithuanian'),
+ Language('la', 'Latin'),
+ Language('mi', 'Maori'),
+ Language('ml', 'Malayalam'),
+ Language('cy', 'Welsh'),
+ Language('sk', 'Slovak'),
+ Language('te', 'Telugu'),
+ Language('fa', 'Persian'),
+ Language('lv', 'Latvian'),
+ Language('bn', 'Bengali'),
+ Language('sr', 'Serbian'),
+ Language('az', 'Azerbaijani'),
+ Language('sl', 'Slovenian'),
+ Language('kn', 'Kannada'),
+ Language('et', 'Estonian'),
+ Language('mk', 'Macedonian'),
+ Language('br', 'Breton'),
+ Language('eu', 'Basque'),
+ Language('is', 'Icelandic'),
+ Language('hy', 'Armenian'),
+ Language('ne', 'Nepali'),
+ Language('mn', 'Mongolian'),
+ Language('bs', 'Bosnian'),
+ Language('kk', 'Kazakh'),
+ Language('sq', 'Albanian'),
+ Language('sw', 'Swahili'),
+ Language('gl', 'Galician'),
+ Language('mr', 'Marathi'),
+ Language('pa', 'Punjabi'),
+ Language('si', 'Sinhala'),
+ Language('km', 'Khmer'),
+ Language('sn', 'Shona'),
+ Language('yo', 'Yoruba'),
+ Language('so', 'Somali'),
+ Language('af', 'Afrikaans'),
+ Language('oc', 'Occitan'),
+ Language('ka', 'Georgian'),
+ Language('be', 'Belarusian'),
+ Language('tg', 'Tajik'),
+ Language('sd', 'Sindhi'),
+ Language('gu', 'Gujarati'),
+ Language('am', 'Amharic'),
+ Language('yi', 'Yiddish'),
+ Language('lo', 'Lao'),
+ Language('uz', 'Uzbek'),
+ Language('fo', 'Faroese'),
+ Language('ht', 'Haitian creole'),
+ Language('ps', 'Pashto'),
+ Language('tk', 'Turkmen'),
+ Language('nn', 'Nynorsk'),
+ Language('mt', 'Maltese'),
+ Language('sa', 'Sanskrit'),
+ Language('lb', 'Luxembourgish'),
+ Language('my', 'Myanmar'),
+ Language('bo', 'Tibetan'),
+ Language('tl', 'Tagalog'),
+ Language('mg', 'Malagasy'),
+ Language('as', 'Assamese'),
+ Language('tt', 'Tatar'),
+ Language('haw', 'Hawaiian'),
+ Language('ln', 'Lingala'),
+ Language('ha', 'Hausa'),
+ Language('ba', 'Bashkir'),
+ Language('jw', 'Javanese'),
+ Language('su', 'Sundanese')
+]
+_TO_LANGUAGE_CODE = {
+ **{language.code: language for language in LANGUAGES},
+ "burmese": "my",
+ "valencian": "ca",
+ "flemish": "nl",
+ "haitian": "ht",
+ "letzeburgesch": "lb",
+ "pushto": "ps",
+ "panjabi": "pa",
+ "moldavian": "ro",
+ "moldovan": "ro",
+ "sinhalese": "si",
+ "castilian": "es",
+}
+_FROM_LANGUAGE_NAME = {
+ **{language.name.lower(): language for language in LANGUAGES}
+}
+def get_language_from_code(language_code, default=None) -> Language:
+ """Return the language name from the language code."""
+ return _TO_LANGUAGE_CODE.get(language_code, default)
+def get_language_from_name(language, default=None) -> Language:
+ """Return the language code from the language name."""
+ return _FROM_LANGUAGE_NAME.get(language.lower() if language else None, default)
+def get_language_names():
+ """Return a list of language names."""
+ return [language.name for language in LANGUAGES]
+if __name__ == "__main__":
+ # Test lookup
+ print(get_language_from_code('en'))
+ print(get_language_from_name('English'))
+ print(get_language_names())

requirements.txt ADDED Viewed

	@@ -0,0 +1,5 @@

+gradio
+--index-url https://download.pytorch.org/whl/cu121
+torch>=2.1.1
+torchvision
+torchaudio

subtitle_manager.py ADDED Viewed

	@@ -0,0 +1,52 @@

+import re
+class Subtitle():
+ def __init__(self,ext="srt"):
+ sub_dict = {
+ "srt":{
+ "coma": ",",
+ "header": "",
+ "format": lambda i,segment : f"{i + 1}\n{self.timeformat(segment['timestamp'][0])} --> {self.timeformat(segment['timestamp'][1] if segment['timestamp'][1] != None else segment['timestamp'][0])}\n{segment['text']}\n\n",
+ },
+ "vtt":{
+ "coma": ".",
+ "header": "WebVTT\n\n",
+ "format": lambda i,segment : f"{self.timeformat(segment['timestamp'][0])} --> {self.timeformat(segment['timestamp'][1] if segment['timestamp'][1] != None else segment['timestamp'][0])}\n{segment['text']}\n\n",
+ },
+ "txt":{
+ "coma": "",
+ "header": "",
+ "format": lambda i,segment : f"{segment['text']}\n",
+ },
+ }
+ self.ext = ext
+ self.coma = sub_dict[ext]["coma"]
+ self.header = sub_dict[ext]["header"]
+ self.format = sub_dict[ext]["format"]
+ def timeformat(self,time):
+ hours = time // 3600
+ minutes = (time - hours * 3600) // 60
+ seconds = time - hours * 3600 - minutes * 60
+ milliseconds = (time - int(time)) * 1000
+ return f"{int(hours):02d}:{int(minutes):02d}:{int(seconds):02d}{self.coma}{int(milliseconds):03d}"
+ def get_subtitle(self,segments):
+ output = self.header
+ for i, segment in enumerate(segments):
+ if segment['text'].startswith(' '):
+ segment['text'] = segment['text'][1:]
+ try:
+ output += self.format(i,segment)
+ except Exception as e:
+ print(e,segment)
+ return output
+ def write_subtitle(self, segments, output_file):
+ output_file += "."+self.ext
+ subtitle = self.get_subtitle(segments)
+ with open(output_file, 'w', encoding='utf-8') as f:
+ f.write(subtitle)