app.py

"""
Voice -> Clean Text (Wispr Flow-style) — Gradio app.

Loads the two fine-tuned models pushed to the Hub by the training
notebook:
  - WHISPER_REPO_ID: verbatim, disfluency-robust ASR
  - LLM_REPO_ID:     cleanup + tone adaptation

ZeroGPU note: no CUDA context exists in the main process — a GPU is attached
only for the duration of a call to a function decorated with @spaces.GPU, then
released. So both models are loaded onto CPU at startup; the only functions
that ever touch CUDA are _transcribe_on_gpu and _clean_up_on_gpu below, which
move the (already-loaded) models onto the GPU the first time they run and
reuse that placement on later calls within the same worker. On a plain CPU
Space (no @spaces available), everything falls back to running on CPU as-is.
"""

import os
import time

import gradio as gr
import torch
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    WhisperForConditionalGeneration,
    WhisperProcessor,
    pipeline,
)

try:
    import spaces
    ZEROGPU_AVAILABLE = True
except ImportError:
    ZEROGPU_AVAILABLE = False

    class _NoOpSpaces:
        @staticmethod
        def GPU(fn=None, **kwargs):
            # No-op decorator so @spaces.GPU still works on a plain CPU Space
            # (no `spaces` package installed / not running on HF infra).
            if fn is not None:
                return fn
            return lambda f: f

    spaces = _NoOpSpaces()

# ---------------------------------------------------------------------------
# Config — point these at the repos the notebook pushed to
# ---------------------------------------------------------------------------

WHISPER_REPO_ID = os.environ.get("WHISPER_REPO_ID", "aijadugar/wisprflow-clone-whisper")
LLM_REPO_ID = os.environ.get("LLM_REPO_ID", "aijadugar/wisprflow-clone-llm")

# CPU-only at import time — deliberately never assume CUDA exists here.
# Under ZeroGPU, torch.cuda.is_available() is unreliable/False in the main
# process anyway; actual device placement happens lazily inside the
# @spaces.GPU-decorated functions below.
#
# bfloat16 (not float32) on CPU: halves RAM for both models. Modern PyTorch
# has native CPU kernels for bf16 matmul/linear (unlike float16, which is
# poorly supported on CPU), so this is safe and meaningfully cuts memory
# pressure — important since the LLM checkpoint here is ~15GB.
LOAD_DEVICE = "cpu"
LOAD_DTYPE = torch.bfloat16

MODE_PROMPTS = {
    "Email": (
        "You are a transcription cleanup assistant. Rewrite the raw speech "
        "transcript into clear, well-punctuated, professional email prose. "
        "Remove filler words and false starts, fix grammar, use complete "
        "sentences. Do not add information that was not said. Output ONLY "
        "the cleaned text."
    ),
    "Chat / Slack": (
        "You are a transcription cleanup assistant. Rewrite the raw speech "
        "transcript into clear, casual chat-message text. Remove filler "
        "words and false starts, fix grammar, but keep it brief and "
        "conversational -- do not over-formalize. Do not add information "
        "that was not said. Output ONLY the cleaned text."
    ),
    "Notes": (
        "You are a transcription cleanup assistant. Rewrite the raw speech "
        "transcript into clear, concise note form. Remove filler words and "
        "false starts, fix grammar, tighten wordy phrasing. Do not add "
        "information that was not said. Output ONLY the cleaned text."
    ),
    "Plain cleanup": (
        "You are a transcription cleanup assistant. Rewrite the raw speech "
        "transcript into clear, well-punctuated text. Remove filler words "
        "and false starts, fix grammar. Do not add information that was "
        "not said. Output ONLY the cleaned text."
    ),
}

# ---------------------------------------------------------------------------
# Model loading (once, at startup) — CPU only. GPU placement is deferred to
# the @spaces.GPU-decorated inference functions below.
# ---------------------------------------------------------------------------

print(f"Loading ASR model from {WHISPER_REPO_ID} (CPU) ...")
whisper_processor = WhisperProcessor.from_pretrained(WHISPER_REPO_ID)
whisper_model = WhisperForConditionalGeneration.from_pretrained(
    WHISPER_REPO_ID,
    torch_dtype=LOAD_DTYPE,
    low_cpu_mem_usage=True,
)
whisper_model.eval()

print(f"Loading cleanup LLM from {LLM_REPO_ID} (CPU) ...")
llm_tokenizer = AutoTokenizer.from_pretrained(LLM_REPO_ID)
llm_model = AutoModelForCausalLM.from_pretrained(
    LLM_REPO_ID,
    torch_dtype=LOAD_DTYPE,
    low_cpu_mem_usage=True,
)
llm_model.eval()

print("Models loaded on CPU. GPU (if any) is attached lazily per-call.")

# Tracks whether each model currently lives on cuda, so repeat calls don't
# re-transfer weights every time within the same ZeroGPU worker.
_whisper_gpu_resident = False
_llm_gpu_resident = False
_asr_pipe = None  # built lazily once we know the device/dtype for this call


# ---------------------------------------------------------------------------
# Pipeline
# ---------------------------------------------------------------------------

@spaces.GPU
def _transcribe_on_gpu(audio_path):
    """The only function that touches CUDA for ASR. ZeroGPU attaches a GPU for
    the duration of this call; whisper_model is moved here once (a no-op on
    later calls once resident) and released back to the caller as plain text,
    since nothing outside this function is guaranteed a live CUDA context."""
    global _whisper_gpu_resident, _asr_pipe

    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = torch.float16 if device == "cuda" else torch.float32

    if not _whisper_gpu_resident and device == "cuda":
        whisper_model.to(device=device, dtype=dtype)
        _whisper_gpu_resident = True
        _asr_pipe = None  # rebuild pipe against the now-GPU model

    if _asr_pipe is None:
        _asr_pipe = pipeline(
            "automatic-speech-recognition",
            model=whisper_model,
            tokenizer=whisper_processor.tokenizer,
            feature_extractor=whisper_processor.feature_extractor,
            torch_dtype=whisper_model.dtype,
            device=device,
        )

    result = _asr_pipe(audio_path)
    return result["text"].strip()


@spaces.GPU
def _clean_up_on_gpu(raw_text, mode):
    """The only function that touches CUDA for the cleanup LLM. Same lazy
    move-once pattern as _transcribe_on_gpu."""
    global _llm_gpu_resident

    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = torch.float16 if device == "cuda" else torch.float32

    if not _llm_gpu_resident and device == "cuda":
        llm_model.to(device=device, dtype=dtype)
        _llm_gpu_resident = True

    system_prompt = MODE_PROMPTS.get(mode, MODE_PROMPTS["Plain cleanup"])
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": raw_text},
    ]
    prompt = llm_tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = llm_tokenizer(prompt, return_tensors="pt").to(llm_model.device)

    with torch.no_grad():
        output_ids = llm_model.generate(
            **inputs,
            max_new_tokens=256,
            do_sample=False,
            repetition_penalty=1.1,
            pad_token_id=llm_tokenizer.pad_token_id or llm_tokenizer.eos_token_id,
        )

    new_tokens = output_ids[0][inputs["input_ids"].shape[1]:]
    cleaned = llm_tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
    return cleaned


def transcribe(audio_path):
    if audio_path is None:
        return "", 0.0, None
    start = time.time()
    try:
        text = _transcribe_on_gpu(audio_path)
    except Exception as e:
        return "", time.time() - start, f"Transcription failed: {e}"
    elapsed = time.time() - start
    return text, elapsed, None


def clean_up(raw_text, mode):
    if not raw_text.strip():
        return "", 0.0, None
    start = time.time()
    try:
        cleaned = _clean_up_on_gpu(raw_text, mode)
    except Exception as e:
        return "", time.time() - start, f"Cleanup failed: {e}"
    elapsed = time.time() - start
    return cleaned, elapsed, None


def run_pipeline(audio_path, mode):
    if audio_path is None:
        return "", "", "Record or upload audio first."

    raw_text, asr_seconds, asr_error = transcribe(audio_path)
    if asr_error:
        return "", "", asr_error
    if not raw_text:
        return "", "", "Couldn't transcribe that clip — try again."

    cleaned_text, llm_seconds, llm_error = clean_up(raw_text, mode)
    if llm_error:
        return raw_text, "", llm_error

    total = asr_seconds + llm_seconds
    stats = (
        f"Transcription: {asr_seconds:.1f}s  ·  "
        f"Cleanup: {llm_seconds:.1f}s  ·  "
        f"Total: {total:.1f}s"
    )
    return raw_text, cleaned_text, stats


# ---------------------------------------------------------------------------
# UI — a plain, single-screen layout: record/upload, pick a style, run.
# ---------------------------------------------------------------------------

with gr.Blocks(title="Voice → Clean Text") as demo:
    gr.Markdown(
        "# 🎙️ Voice → Clean Text\n"
        "Record or upload audio, pick a style, and get a cleaned-up transcript."
    )

    with gr.Row():
        with gr.Column(scale=1):
            audio_input = gr.Audio(
                sources=["microphone", "upload"],
                type="filepath",
                label="Speak or upload audio",
            )
            mode = gr.Radio(
                choices=list(MODE_PROMPTS.keys()),
                value="Plain cleanup",
                label="Cleanup style",
            )
            run_btn = gr.Button("Transcribe & Clean Up", variant="primary")

        with gr.Column(scale=1):
            raw_output = gr.Textbox(label="Raw transcript (verbatim)", lines=5)
            clean_output = gr.Textbox(label="Cleaned text", lines=5)
            stats_output = gr.Markdown()

    run_btn.click(
        fn=run_pipeline,
        inputs=[audio_input, mode],
        outputs=[raw_output, clean_output, stats_output],
    )

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