Update app.py

app.py

@@ -1,50 +1,46 @@
 import torch
 import gradio as gr
-from transformers import pipeline
+import ffmpeg
 import numpy as np
-from pydub import AudioSegment
-import io
+import whisper
 
-MODEL_NAME = "openai/whisper-large-v3"
-BATCH_SIZE = 8
+MODEL_NAME = "large-v3"
+SAMPLE_RATE = 16000
 
-device = 0 if torch.cuda.is_available() else "cpu"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = whisper.load_model(MODEL_NAME).to(device)
 
-pipe = pipeline(
-    task="automatic-speech-recognition",
-    model=MODEL_NAME,
-    chunk_length_s=30,
-    device=device,
-)
+def load_audio(file):
+    try:
+        # This launches a subprocess to decode audio while down-mixing and resampling as necessary.
+        # Requires the ffmpeg CLI and `ffmpeg-python` package to be installed.
+        out, _ = (
+            ffmpeg.input(file, threads=0)
+            .output("-", format="s16le", acodec="pcm_s16le", ac=1, ar=SAMPLE_RATE)
+            .run(cmd=["ffmpeg", "-nostdin"], capture_stdout=True, capture_stderr=True)
+        )
+    except ffmpeg.Error as e:
+        raise RuntimeError(f"Failed to load audio: {e.stderr.decode()}") from e
+
+    return np.frombuffer(out, np.int16).flatten().astype(np.float32) / 32768.0
 
 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.")
 
     try:
-        # Read the audio file
-        audio = AudioSegment.from_file(audio_file)
-        
-        # Convert to mono if stereo
-        if audio.channels > 1:
-            audio = audio.set_channels(1)
-
-        # Convert to 16kHz sample rate
-        audio = audio.set_frame_rate(16000)
-
-        # Convert to numpy array
-        samples = np.array(audio.get_array_of_samples()).astype(np.float32) / 32768.0
+        # Load audio
+        audio = load_audio(audio_file.name)
 
-        # Convert to the format expected by Whisper
-        inputs = {"array": samples, "sampling_rate": 16000}
+        # Transcribe
+        result = model.transcribe(audio, task=task, language="en")
 
-        result = pipe(inputs, batch_size=BATCH_SIZE, generate_kwargs={"task": task}, return_timestamps=True)
-        
+        # Format output
         output = ""
-        for chunk in result["chunks"]:
-            start_time = chunk["timestamp"][0]
-            end_time = chunk["timestamp"][1]
-            text = chunk["text"]
+        for segment in result["segments"]:
+            start_time = segment["start"]
+            end_time = segment["end"]
+            text = segment["text"]
             output += f"[{format_timestamp(start_time)} -> {format_timestamp(end_time)}] {text}\n"
 
         return output
@@ -66,9 +62,9 @@ demo = gr.Interface(
     fn=transcribe,
     inputs=[audio_input, task_input],
     outputs=output,
-    title=f"Whisper Large V3: Transcribe Audio with Timestamps",
+    title=f"Whisper {MODEL_NAME}: Transcribe Audio with Timestamps",
     description=(
-        f"Transcribe audio files with Whisper Large V3 [{MODEL_NAME}](https://huggingface.co/{MODEL_NAME}). "
+        f"Transcribe audio files with Whisper {MODEL_NAME}. "
         "Upload an audio file and choose whether to transcribe or translate. "
         "The output includes timestamps for each transcribed segment."
     ),