Deaconhead's picture
Fix audio input parsing for browser/client variants
447647d verified
Raw
History Blame Contribute Delete
3.11 kB
import numpy as np
import torch
import gradio as gr
from transformers import WhisperForConditionalGeneration, WhisperProcessor
from transformers.pipelines.audio_utils import ffmpeg_read
MODEL_ID = "openai/whisper-small"
TARGET_SAMPLE_RATE = 16000
processor = WhisperProcessor.from_pretrained(MODEL_ID)
model = WhisperForConditionalGeneration.from_pretrained(MODEL_ID)
model.eval()
def _normalize_waveform(data: np.ndarray) -> np.ndarray:
if data.ndim > 1:
data = data.mean(axis=1)
if np.issubdtype(data.dtype, np.integer):
max_int = np.iinfo(data.dtype).max
data = data.astype(np.float32) / float(max_int)
else:
data = data.astype(np.float32)
peak = np.max(np.abs(data)) if data.size else 0.0
if peak > 1.0:
data = data / peak
return data
def _extract_audio(audio_input):
# Gradio may send audio as (sample_rate, np.ndarray), a filepath string, or a FileData-like dict.
if isinstance(audio_input, tuple) and len(audio_input) == 2:
sample_rate, data = audio_input
return int(sample_rate), _normalize_waveform(data)
path = None
if isinstance(audio_input, str):
path = audio_input
elif isinstance(audio_input, dict):
path = audio_input.get("path")
if path:
with open(path, "rb") as f:
audio_bytes = f.read()
data = ffmpeg_read(audio_bytes, TARGET_SAMPLE_RATE)
return TARGET_SAMPLE_RATE, _normalize_waveform(data)
raise ValueError("Unsupported audio input format. Please upload a valid audio file.")
def transcribe_audio(audio, task):
if audio is None:
return "Please upload or record audio first."
sample_rate, data = _extract_audio(audio)
inputs = processor(
data,
sampling_rate=sample_rate,
return_tensors="pt",
)
forced_decoder_ids = processor.get_decoder_prompt_ids(task=task)
with torch.inference_mode():
predicted_ids = model.generate(
inputs.input_features,
forced_decoder_ids=forced_decoder_ids,
)
text = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0].strip()
return text or "(No speech detected)"
with gr.Blocks(title="Whisper Small STT CPU") as demo:
gr.Markdown(
"""
# Whisper Small STT (Free CPU)
Public speech-to-text using `openai/whisper-small`.
- `transcribe`: keep original language
- `translate`: translate speech to English
"""
)
audio_input = gr.Audio(
label="Audio Input",
type="numpy",
sources=["upload", "microphone"],
)
task_input = gr.Dropdown(
choices=["transcribe", "translate"],
value="transcribe",
label="Task",
)
run_btn = gr.Button("Convert Speech to Text", variant="primary")
text_output = gr.Textbox(label="Transcript", lines=12)
run_btn.click(
fn=transcribe_audio,
inputs=[audio_input, task_input],
outputs=[text_output],
api_name="transcribe",
)
if __name__ == "__main__":
demo.queue(default_concurrency_limit=1).launch()