app.py

import torch
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
import gradio as gr
from pydub import AudioSegment
import os

# Set device and precision for CPU
device = "cpu"
torch_dtype = torch.float32

# Load KB-Whisper model (Large variant)
model_id = "KBLab/kb-whisper-large"

model = AutoModelForSpeechSeq2Seq.from_pretrained(
    model_id, torch_dtype=torch_dtype
).to(device)

processor = AutoProcessor.from_pretrained(model_id)

pipe = pipeline(
    "automatic-speech-recognition",
    model=model,
    tokenizer=processor.tokenizer,
    feature_extractor=processor.feature_extractor,
    device=device,
    torch_dtype=torch_dtype,
)

def transcribe(audio_path):
    # Handle m4a or other formats by converting to wav
    base, ext = os.path.splitext(audio_path)
    if ext.lower() != ".wav":
        try:
            sound = AudioSegment.from_file(audio_path)
            audio_converted_path = base + ".converted.wav"
            sound.export(audio_converted_path, format="wav")
            audio_path = audio_converted_path
        except Exception as e:
            return f"Error converting audio: {str(e)}"

    # Transcribe
    try:
        result = pipe(audio_path, chunk_length_s=30, generate_kwargs={"task": "transcribe", "language": "sv"})
        return result["text"]
    except Exception as e:
        return f"Transcription failed: {str(e)}"

# Build Gradio interface
gr.Interface(
    fn=transcribe,
    inputs=gr.Audio(type="filepath", label="Upload Swedish Audio"),
    outputs=gr.Textbox(label="Transcribed Text"),
    title="KB-Whisper Transcriber (Swedish, Free CPU)",
    description="Upload .m4a, .mp3, or .wav files. Transcribes Swedish speech using KBLab's Whisper Large model.",
).launch(share=True)