Update app.py

app.py

@@ -2,65 +2,138 @@ import os
 import torch
 import gradio as gr
 from pydub import AudioSegment
-from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor
-import tempfile
-import math
+from transformers import pipeline, AutoModelForSpeechSeq2Seq, AutoProcessor
+from pathlib import Path
+from tempfile import NamedTemporaryFile
+from datetime import timedelta
+import time
 
-from datasets import load_dataset, Audio
-import numpy as np
-import torchaudio
+# Configuration
+MODEL_ID = "KBLab/kb-whisper-large"
+CHUNK_DURATION_MS = 10000
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+TORCH_DTYPE = torch.float16 if torch.cuda.is_available() else torch.float32
 
-# Set up model
-device = "cpu"
-torch_dtype = torch.float32
+# Initialize model and pipeline
+def initialize_pipeline():
+    model = AutoModelForSpeechSeq2Seq.from_pretrained(
+        MODEL_ID,
+        torch_dtype=TORCH_DTYPE,
+        low_cpu_mem_usage=True
+    ).to(DEVICE)
+    
+    processor = AutoProcessor.from_pretrained(MODEL_ID)
+    
+    return pipeline(
+        "automatic-speech-recognition",
+        model=model,
+        tokenizer=processor.tokenizer,
+        feature_extractor=processor.feature_extractor,
+        device=DEVICE,
+        torch_dtype=TORCH_DTYPE,
+        model_kwargs={"use_flash_attention_2": torch.cuda.is_available()}
+    )
 
-model_id = "KBLab/kb-whisper-large"
-
-processor = AutoProcessor.from_pretrained(model_id)
-model = AutoModelForSpeechSeq2Seq.from_pretrained(model_id, torch_dtype=torch_dtype).to(device)
+# Convert audio if needed
+def convert_to_wav(audio_path: str) -> str:
+    ext = str(Path(audio_path).suffix).lower()
+    if ext != ".wav":
+        audio = AudioSegment.from_file(audio_path)
+        wav_path = str(Path(audio_path).with_suffix(".converted.wav"))
+        audio.export(wav_path, format="wav")
+        return wav_path
+    return audio_path
 
-# Helper: Split audio into chunks
-def split_audio(audio_path, chunk_duration_ms=10000):
-    audio = AudioSegment.from_file(audio_path)
-    chunks = [audio[i:i + chunk_duration_ms] for i in range(0, len(audio), chunk_duration_ms)]
-    return chunks
+# Split audio into chunks
+def split_audio(audio_path: str) -> list:
+    try:
+        audio = AudioSegment.from_file(audio_path)
+        if len(audio) == 0:
+            raise ValueError("Audio file is empty or invalid.")
+        return [audio[i:i + CHUNK_DURATION_MS] around(i, len(audio), CHUNK_DURATION_MS) for i in range(0, len(audio), CHUNK_DURATION_MS)]
+    except Exception as e:
+        raise ValueError(f"Failed to process audio: {str(e)}")
 
-# Helper: Transcribe a single chunk
-def transcribe_chunk(chunk):
-    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmpfile:
-        chunk.export(tmpfile.name, format="wav")
-        input_audio, _ = torchaudio.load(tmpfile.name)
-        input_features = processor(input_audio.squeeze(), sampling_rate=16000, return_tensors="pt").input_features
-        input_features = input_features.to(device)
-        predicted_ids = model.generate(input_features, task="transcribe", language="sv")
-        transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
-        os.remove(tmpfile.name)
-        return transcription
+# Helper to compute chunk start time
+def get_chunk_time(index: int, chunk_duration_ms: int) -> str:
+    start_ms = index * chunk_duration_ms
+    return str(timedelta(milliseconds=start_ms))
 
-# Full transcription function with progress
-def transcribe_with_progress(audio_path, progress=gr.Progress()):
-    ext = os.path.splitext(audio_path)[1].lower()
-    if ext != ".wav":
-        sound = AudioSegment.from_file(audio_path)
-        audio_path = audio_path.replace(ext, ".converted.wav")
-        sound.export(audio_path, format="wav")
+# Transcribe audio with progress and timestamps
+def transcribe(audio_path: str, include_timestamps: bool = False, progress=gr.Progress()):
+    try:
+        if not audio_path:
+            return "No audio file provided.", None
+            
+        # Convert to WAV if needed
+        wav_path = convert_to_wav(audio_path)
+        
+        # Split and process
+        chunks = split_audio(wav_path)
+        total_chunks = len(chunks)
+        transcript = []
+        timestamped_transcript = []
+        
+        for i, chunk in enumerate(chunks):
+            try:
+                with NamedTemporaryFile(suffix=".wav", delete=False) as temp_file:
+                    chunk.export(temp_file.name, format="wav")
+                    result = PIPELINE(temp_file.name, 
+                                   generate_kwargs={"task": "transcribe", "language": "sv"})
+                    text = result["text"].strip()
+                    transcript.append(text)
+                    if include_timestamps:
+                        timestamp = get_chunk_time(i, CHUNK_DURATION_MS)
+                        timestamped_transcript.append(f"[{timestamp}] {text}")
+            finally:
+                if os.path.exists(temp_file.name):
+                    os.remove(temp_file.name)
+                    
+            progress((i + 1) / total_chunks)
+            yield " ".join(transcript), None
+            
+        # Clean up converted file if created
+        if wav_path != audio_path and os.path.exists(wav_path):
+            os.remove(wav_path)
+            
+        # Prepare final transcript and downloadable file
+        final_transcript = " ".join(transcript)
+        download_content = "\n".join(timestamped_transcript) if include_timestamps else final_transcript
+        with NamedTemporaryFile(suffix=".txt", delete=False, mode='w', encoding='utf-8') as temp_file:
+            temp_file.write(download_content)
+            download_path = temp_file.name
+            
+        return final_transcript, download_path
+        
+    except Exception as e:
+        return f"Error during transcription: {str(e)}", None
 
-    chunks = split_audio(audio_path, chunk_duration_ms=8000)
-    full_transcript = ""
-    total_chunks = len(chunks)
+# Initialize pipeline globally
+PIPELINE = initialize_pipeline()
 
-    for i, chunk in enumerate(chunks):
-        partial_text = transcribe_chunk(chunk)
-        full_transcript += partial_text + " "
-        progress(i + 1, total_chunks)  # Update progress bar
-        yield full_transcript.strip()  # Stream updated text to UI
+# Gradio Interface with Blocks
+def create_interface():
+    with gr.Blocks(theme=gr.themes.Soft()) as demo:
+        gr.Markdown("# Swedish Whisper Transcriber")
+        gr.Markdown("Upload audio (.wav, .mp3, .m4a) for real-time Swedish speech transcription.")
+        
+        with gr.Row():
+            with gr.Column():
+                audio_input = gr.Audio(type="filepath", label="Upload Audio")
+                timestamp_toggle = gr.Checkbox(label="Include Timestamps in Download", value=False)
+                transcribe_btn = gr.Button("Transcribe")
+            
+            with gr.Column():
+                transcript_output = gr.Textbox(label="Live Transcription", lines=10)
+                download_output = gr.File(label="Download Transcript")
+        
+        transcribe_btn.click(
+            fn=transcribe,
+            inputs=[audio_input, timestamp_toggle],
+            outputs=[transcript_output, download_output]
+        )
+    
+    return demo
 
-# UI
-gr.Interface(
-    fn=transcribe_with_progress,
-    inputs=gr.Audio(type="filepath", label="Upload Swedish Audio"),
-    outputs=gr.Textbox(label="Live Transcript (Swedish)"),
-    title="Live Swedish Transcriber (KB-Whisper)",
-    description="Streams transcription word-by-word with visual progress. Supports .m4a, .mp3, .wav. May be slow on CPU.",
-    live=True
-).launch()
+if __name__ == "__main__":
+    create_interface().launch()