#!/usr/bin/env python3
"""CPU-only Whisper ASR for auto-transcribing voice reference clips.

Runs on CPU and is never wrapped in @spaces.GPU — keeps it off the ZeroGPU
quota entirely and lets it run any time, independent of whichever TTS
backend currently holds the GPU. Powers the auto-fill of "Reference
transcript" when a voice is picked from the gallery, which improves Higgs
Audio v3's zero-shot cloning quality.

Multilingual `whisper-base` (not the `.en` variant) since the voice gallery
spans many languages.
"""
import logging

import torch

ASR_REPO = "openai/whisper-base"
ASR_SAMPLE_RATE = 16000

_processor = None
_model = None


def load():
    """Load the Whisper processor + model onto CPU. Idempotent."""
    global _processor, _model
    if _model is not None:
        return

    from transformers import AutoProcessor, WhisperForConditionalGeneration

    logging.info(f"Loading Whisper ASR ({ASR_REPO}) on CPU…")
    _processor = AutoProcessor.from_pretrained(ASR_REPO)
    _model = WhisperForConditionalGeneration.from_pretrained(ASR_REPO).eval()
    logging.info("Whisper ASR ready.")


def transcribe(audio_path):
    """Best-effort CPU transcription of a reference clip.

    Returns the stripped transcript, or "" if there's no clip or
    transcription fails — callers treat "" as "leave the field as-is /
    let the user fill it in manually".
    """
    if not audio_path or _model is None:
        return ""

    import soundfile as sf
    import torchaudio

    try:
        data, sr = sf.read(audio_path, dtype="float32", always_2d=True)  # [L, C]
        wav = torch.from_numpy(data).mean(dim=1)  # mono [L]
        if sr != ASR_SAMPLE_RATE:
            wav = torchaudio.functional.resample(wav, orig_freq=sr, new_freq=ASR_SAMPLE_RATE)
        inputs = _processor(wav.numpy(), sampling_rate=ASR_SAMPLE_RATE, return_tensors="pt")
        with torch.no_grad():
            tokens = _model.generate(**inputs)
        return _processor.batch_decode(tokens, skip_special_tokens=True)[0].strip()
    except Exception as e:
        logging.warning(f"Reference transcription failed: {e}")
        return ""