whisper_ui.py

#!/usr/bin/env python3
"""
Whisper Web UI with Gradio
Provides a simple interface for audio transcription and translation
"""

import os
import gradio as gr
import whisper
import torch
from pathlib import Path
from fastapi import FastAPI
from fastapi.responses import JSONResponse

# Check if GPU is available
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

if device == "cuda":
    print(f"GPU: {torch.cuda.get_device_name(0)}")
    print(f"VRAM: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.2f} GB")

# Available models with VRAM requirements (in GB)
MODEL_VRAM = {
    "tiny": 1,
    "base": 1,
    "small": 2,
    "medium": 5,
    "large": 10,
    "turbo": 6,
}
MODELS = list(MODEL_VRAM.keys())

# Get available VRAM
available_vram = 0
if torch.cuda.is_available():
    available_vram = torch.cuda.get_device_properties(0).total_memory / 1024**3
    print(f"Available VRAM: {available_vram:.2f} GB")

# Global model cache
current_model = None
current_model_name = None


def check_vram(model_name):
    """Check if model fits in available VRAM"""
    required = MODEL_VRAM.get(model_name, 0)
    # Add 1GB safety margin
    if available_vram > 0 and required + 1 > available_vram:
        return False, f"Model '{model_name}' requires ~{required}GB VRAM but only {available_vram:.1f}GB available"
    return True, ""


def get_available_models():
    """Get list of models that fit in available VRAM"""
    if available_vram == 0:
        return MODELS  # CPU mode, all models available (but slow)
    return [m for m in MODELS if MODEL_VRAM[m] + 1 <= available_vram]


AVAILABLE_MODELS = get_available_models()
print(f"Available models for your VRAM: {AVAILABLE_MODELS}")


def get_models_api(dummy=None):
    """API endpoint to get available models and GPU info"""
    return {
        "available_models": AVAILABLE_MODELS,
        "all_models": MODELS,
        "model_vram": MODEL_VRAM,
        "gpu_vram": round(available_vram, 1),
        "device": device,
    }


def load_model(model_name):
    """Load or return cached model"""
    global current_model, current_model_name

    # Check VRAM before loading
    ok, error_msg = check_vram(model_name)
    if not ok:
        raise RuntimeError(error_msg)

    if current_model_name != model_name:
        print(f"Loading model: {model_name}")
        current_model = whisper.load_model(model_name, device=device)
        current_model_name = model_name
        print(f"Model {model_name} loaded successfully")

    return current_model


def transcribe_audio(audio_file, model_name, task, language, output_format):
    """Transcribe or translate audio file"""
    try:
        if audio_file is None:
            return "Please upload an audio file", ""

        # Load model
        model = load_model(model_name)

        # Transcribe
        print(f"Processing: {audio_file}")
        options = {
            "task": task,
            "fp16": device == "cuda",  # Use FP16 on GPU
        }

        if language and language != "auto":
            options["language"] = language

        result = model.transcribe(audio_file, **options)

        # Format output
        text = result["text"]

        # Save to output directory
        output_dir = Path("/data/output")
        output_dir.mkdir(parents=True, exist_ok=True)

        input_filename = Path(audio_file).stem

        if output_format == "txt":
            output_file = output_dir / f"{input_filename}.txt"
            with open(output_file, "w", encoding="utf-8") as f:
                f.write(text)
        elif output_format == "srt":
            output_file = output_dir / f"{input_filename}.srt"
            with open(output_file, "w", encoding="utf-8") as f:
                f.write(format_srt(result["segments"]))
        elif output_format == "vtt":
            output_file = output_dir / f"{input_filename}.vtt"
            with open(output_file, "w", encoding="utf-8") as f:
                f.write(format_vtt(result["segments"]))

        metadata = f"""
**Model:** {model_name}
**Task:** {task}
**Language:** {result.get('language', 'auto')}
**Output saved to:** {output_file}
"""

        return text, metadata

    except Exception as e:
        return f"Error: {str(e)}", ""


def format_srt(segments):
    """Format segments as SRT subtitles"""
    output = []
    for i, segment in enumerate(segments, 1):
        start = format_timestamp(segment["start"])
        end = format_timestamp(segment["end"])
        text = segment["text"].strip()
        output.append(f"{i}\n{start} --> {end}\n{text}\n")
    return "\n".join(output)


def format_vtt(segments):
    """Format segments as WebVTT subtitles"""
    output = ["WEBVTT\n"]
    for segment in segments:
        start = format_timestamp(segment["start"])
        end = format_timestamp(segment["end"])
        text = segment["text"].strip()
        output.append(f"{start} --> {end}\n{text}\n")
    return "\n".join(output)


def format_timestamp(seconds):
    """Format seconds as timestamp (HH:MM:SS.mmm)"""
    hours = int(seconds // 3600)
    minutes = int((seconds % 3600) // 60)
    secs = int(seconds % 60)
    millis = int((seconds % 1) * 1000)
    return f"{hours:02d}:{minutes:02d}:{secs:02d}.{millis:03d}"


# Create Gradio interface
with gr.Blocks(title="Whisper Transcription") as demo:
    gr.Markdown("# 🎤 Whisper Transcription & Translation")
    gr.Markdown("Upload an audio file to transcribe or translate using OpenAI Whisper with AMD ROCm GPU acceleration")

    with gr.Row():
        with gr.Column():
            audio_input = gr.Audio(
                label="Upload Audio File",
                type="filepath"
            )

            model_dropdown = gr.Dropdown(
                choices=AVAILABLE_MODELS,
                value=AVAILABLE_MODELS[0] if AVAILABLE_MODELS else "tiny",
                label="Model",
                info=f"Models filtered by VRAM ({available_vram:.1f}GB available)"
            )

            task_radio = gr.Radio(
                choices=["transcribe", "translate"],
                value="transcribe",
                label="Task",
                info="Translate converts audio to English"
            )

            language_dropdown = gr.Dropdown(
                choices=["auto", "en", "es", "fr", "de", "it", "pt", "nl", "pl", "ru", "ja", "ko", "zh"],
                value="auto",
                label="Language (optional)",
                info="Auto-detect if not specified"
            )

            format_radio = gr.Radio(
                choices=["txt", "srt", "vtt"],
                value="txt",
                label="Output Format",
                info="Format for saved file"
            )

            submit_btn = gr.Button("Transcribe", variant="primary")

        with gr.Column():
            output_text = gr.Textbox(
                label="Transcription",
                lines=20,
                max_lines=30
            )

            metadata_text = gr.Markdown()

    submit_btn.click(
        fn=transcribe_audio,
        inputs=[audio_input, model_dropdown, task_radio, language_dropdown, format_radio],
        outputs=[output_text, metadata_text]
    )

    # Build model info with availability indicators
    model_info_lines = ["### Model Information"]
    for model, vram in MODEL_VRAM.items():
        available = "✓" if model in AVAILABLE_MODELS else "✗"
        model_info_lines.append(f"- {available} **{model}**: ~{vram}GB VRAM")
    model_info_lines.append(f"\n*Your GPU: {available_vram:.1f}GB VRAM*")
    model_info_lines.append("""
### CLI Usage
You can also use Whisper from the command line:
```bash
whisper /data/input/audio.mp3 --model base --output_dir /data/output
```
""")
    gr.Markdown("\n".join(model_info_lines))


# Create FastAPI app with custom endpoint
app = FastAPI()


@app.get("/api/models")
async def api_get_models():
    """Direct REST endpoint for getting available models"""
    return JSONResponse(content=get_models_api())


# Mount Gradio app
app = gr.mount_gradio_app(app, demo, path="/")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=80)