app.py

import torch
import gradio as gr
from transformers import pipeline
import numpy as np
import librosa

MODEL_NAME = "openai/whisper-large-v3"
BATCH_SIZE = 8

device = 0 if torch.cuda.is_available() else "cpu"

pipe = pipeline(
    task="automatic-speech-recognition",
    model=MODEL_NAME,
    chunk_length_s=30,
    device=device,
)

def transcribe(audio_file, task):
    if audio_file is None:
        raise gr.Error("No audio file submitted! Please upload an audio file before submitting your request.")

    # Load audio file
    try:
        # Use librosa to load the audio file
        audio, sr = librosa.load(audio_file, sr=16000)  # Whisper expects 16kHz sampling rate
    except Exception as e:
        raise gr.Error(f"Error loading audio file: {str(e)}")

    # Convert to format expected by Whisper
    inputs = {"array": audio, "sampling_rate": sr}

    result = pipe(inputs, batch_size=BATCH_SIZE, generate_kwargs={"task": task}, return_timestamps=True)
    
    output = ""
    for chunk in result["chunks"]:
        start_time = chunk["timestamp"][0]
        end_time = chunk["timestamp"][1]
        text = chunk["text"]
        output += f"[{format_timestamp(start_time)} -> {format_timestamp(end_time)}] {text}\n"

    return output

def format_timestamp(seconds):
    minutes, seconds = divmod(seconds, 60)
    hours, minutes = divmod(minutes, 60)
    return f"{int(hours):02d}:{int(minutes):02d}:{seconds:.2f}"

# Use specific Gradio components
audio_input = gr.components.File(label="Audio file", file_types=["audio"])
task_input = gr.components.Radio(["transcribe", "translate"], label="Task", default="transcribe")
output = gr.components.Textbox(label="Transcription with Timestamps")

demo = gr.Interface(
    fn=transcribe,
    inputs=[audio_input, task_input],
    outputs=output,
    title=f"Whisper Large V3: Transcribe Audio with Timestamps",
    description=(
        f"Transcribe audio files with Whisper Large V3 [{MODEL_NAME}](https://huggingface.co/{MODEL_NAME}). "
        "Upload an audio file and choose whether to transcribe or translate. "
        "The output includes timestamps for each transcribed segment."
    ),
)

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