Initial commit of WhisperFast project

README.md

@@ -1,13 +1,10 @@
----
 title: WhisperFast
-emoji: 🌍
-colorFrom: green
+emoji: 🗣️
+colorFrom: blue
 colorTo: blue
 sdk: gradio
-sdk_version: 4.42.0
+sdk_version: 4.36.0
 app_file: app.py
 pinned: false
 license: mit
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+short_description: Transcribe audio to subtitles instantly

app.py

@@ -0,0 +1,173 @@
+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
+pipe = 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
+    global pipe
+
+    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"
+   
+    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", "distil-whisper/distil-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=["upload", "microphone",], type="filepath", label="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

@@ -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

@@ -0,0 +1,5 @@
+--extra-index-url https://download.pytorch.org/whl/cu121
+torch>=2.1.1
+gradio==4.16.0
+transformers
+accelerate

subtitle_manager.py

@@ -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)