DramaboxTTS / src /asr_backend.py
Daankular's picture
Auto-transcribe selected voice with CPU-only Whisper to pre-fill the Higgs reference transcript
e13ec41 verified
Raw
History Blame Contribute Delete
2.14 kB
#!/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 ""