app.py

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
import spaces

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):
    # Load the model into RAM first
    torch_dtype = torch.float32  # Load onto CPU with float32 precision
    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",
    )
    
    processor = AutoProcessor.from_pretrained(model_id)

    pipe = pipeline(
        "automatic-speech-recognition",
        model=model,
        tokenizer=processor.tokenizer,
        feature_extractor=processor.feature_extractor,
        torch_dtype=torch_dtype,  # Keep in CPU until GPU is requested
        device="cpu",  # Initially stay on CPU
    )
    return pipe, model  # Return both pipe and model for later GPU switch

def move_to_gpu(model):
    if torch.cuda.is_available():
        device = "cuda:0"
        torch_dtype = torch.float16  # Use float16 precision on GPU
        model.to(device, dtype=torch_dtype)
    elif platform == "darwin":
        device = "mps"
        model.to(device)
    else:
        device = "cpu"
    
    return device

@spaces.GPU
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 is None:
        logging.info("first model")
        progress(0.1, desc="Loading Model..")
        pipe, model = create_pipe(modelName, flash)
    elif modelName != last_model:
        logging.info("new model")
        torch.cuda.empty_cache()
        progress(0.1, desc="Loading Model..")
        pipe, model = create_pipe(modelName, flash)
    else:
        logging.info("Model not changed")

    last_model = modelName

    # Now move the model to GPU after the pipe is created, within the function's context
    with torch.inference_mode():
        device = move_to_gpu(pipe.model)

        # Update pipe's device
        pipe.device = torch.device(device)
        pipe.model.to(pipe.device)

        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.en",
        "openai/whisper-base.en",
        "openai/whisper-small.en", "distil-whisper/distil-small.en",
        "openai/whisper-medium.en", "distil-whisper/distil-medium.en",
        "openai/whisper-large-v3", "distil-whisper/distil-large-v3",
    ]
    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-v3",
                                                     label="Model", info="Select whisper model", interactive=True),
                                         gr.Dropdown(choices=["English"], value="English", interactive=False, visible=False,
                                                     label="Language",
                                                     info="Select audio voice language", ),
                                         gr.Text(label="URL", info="(YouTube, etc.)", interactive=False, visible=False),
                                         gr.File(label="Upload Files", file_count="multiple", interactive=False, visible=False),
                                         gr.Audio(sources=["upload", "microphone", ], type="filepath", label="Input",
                                                  waveform_options=waveform_options),
                                         gr.Dropdown(choices=["transcribe", "translate"], label="Task",
                                                     value="transcribe", interactive=False, visible=False),
                                         gr.Checkbox(label='Flash', info='Use Flash Attention 2', interactive=False, visible=False),
                                         gr.Number(label='chunk_length_s', value=30, interactive=False, visible=False),
                                         gr.Number(label='batch_size', value=24, interactive=False, visible=False)
                                     ], outputs=[
                                         gr.File(label="Download"),
                                         gr.Text(label="Transcription"),
                                         gr.Text(label="Segments")
                                     ]
                                     )

if __name__ == "__main__":
    demo.launch()