""" Varuna STT — Gradio demo Space. """ from __future__ import annotations import os import tempfile from pathlib import Path import gc import gradio as gr import numpy as np import soundfile as sf import librosa import torch import spaces from omegaconf import OmegaConf, open_dict from nemo.collections.asr.models import EncDecRNNTBPEModel # ── Resolve model assets ────────────────────────────────────────────────────── def _resolve_assets(): base = os.environ.get("VARUNA_BASE") tok = os.environ.get("VARUNA_TOK") ckpt = os.environ.get("VARUNA_CKPT") if base and tok and ckpt: return base, tok, ckpt from huggingface_hub import hf_hub_download print("[boot] downloading varuna.ckpt from SkunkWorkLabs/varuna-stt ...") ckpt_path = hf_hub_download("SkunkWorkLabs/varuna-stt", "varuna.ckpt") print("[boot] downloading base nemotron .nemo ...") base_path = hf_hub_download("nvidia/nemotron-speech-streaming-en-0.6b", "nemotron-speech-streaming-en-0.6b.nemo") tok_dir = str(Path(__file__).parent / "tokenizer") return base_path, tok_dir, ckpt_path print("[boot] resolving model assets") BASE, TOKENIZER_DIR, CKPT = _resolve_assets() print(f" base = {BASE}") print(f" tokenizer = {TOKENIZER_DIR}") print(f" ckpt = {CKPT}") print("[boot] loading model ...") _model = EncDecRNNTBPEModel.restore_from(BASE, map_location="cpu") _model.change_vocabulary(new_tokenizer_dir=TOKENIZER_DIR, new_tokenizer_type="bpe") _dec = OmegaConf.to_container(_model.cfg.decoding, resolve=True) _dec = OmegaConf.create(_dec) with open_dict(_dec): _dec.strategy = "greedy_batch" if "greedy" not in _dec: _dec.greedy = {} _dec.greedy.use_cuda_graph_decoder = False _dec.greedy.loop_labels = False # use pure PyTorch decoder, not Numba CUDA kernel for section in ("greedy", "beam"): if section in _dec and "boosting_tree" in _dec[section]: _dec[section].boosting_tree.use_triton = False _model.change_decoding_strategy(_dec) _state = torch.load(CKPT, map_location="cpu", weights_only=False) _sd = _state["state_dict"] if "state_dict" in _state else _state _model.load_state_dict(_sd, strict=False) del _state, _sd _model = _model.eval() print("[boot] model ready") TARGET_SR = 16000 def _numpy_to_wav_paths(sr: int, wav: np.ndarray) -> list[str]: """Convert numpy audio to list of 16kHz mono WAV temp files (one per channel).""" if wav.dtype != np.float32: wav = wav.astype(np.float32) / np.iinfo(wav.dtype).max if wav.ndim == 1: wav = wav[:, None] paths = [] for ch in range(min(wav.shape[1], 2)): mono = wav[:, ch] if sr != TARGET_SR: mono = librosa.resample(mono, orig_sr=sr, target_sr=TARGET_SR) tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False) sf.write(tmp.name, mono, TARGET_SR) paths.append(tmp.name) return paths def _run_nemo(wav_path: str) -> str: with torch.inference_mode(): out = _model.transcribe(audio=[wav_path], batch_size=1, return_hypotheses=False, verbose=False, num_workers=0) if isinstance(out, tuple): out = out[0] h = out[0] return h.text if hasattr(h, "text") else str(h) @spaces.GPU(duration=60) def transcribe(audio_input): if audio_input is None: return "Please upload an audio file or record from the mic.", "", "", gr.update(visible=False) try: # Move model to GPU for this call (NVIDIA ZeroGPU pattern) _model.to("cuda") _model.to(torch.bfloat16) sr, wav = audio_input paths = _numpy_to_wav_paths(sr, wav) n_ch = len(paths) dur = wav.shape[0] / sr if wav.ndim > 1 else len(wav) / sr info_md = ( f"**Channels:** {n_ch} ({'mono' if n_ch == 1 else 'stereo'}) \n" f"**Sample rate:** {sr} Hz \n" f"**Duration:** {dur:.2f}s" ) if n_ch == 1: text = _run_nemo(paths[0]) return info_md, text, "", gr.update(visible=False) left_text = _run_nemo(paths[0]) right_text = _run_nemo(paths[1]) return info_md, left_text, right_text, gr.update(visible=True) except Exception as e: import traceback err = traceback.format_exc() print(err) return f"❌ {type(e).__name__}: {e}", err, "", gr.update(visible=False) finally: _model.cpu() gc.collect() torch.cuda.empty_cache() # ── UI ──────────────────────────────────────────────────────────────────────── CSS = "#title { text-align: center; }" with gr.Blocks(theme=gr.themes.Soft(), css=CSS, title="Varuna STT") as demo: gr.Markdown("# Varuna STT 🌊", elem_id="title") gr.Markdown( "Hindi automatic speech recognition by **SkunkWorks Labs** — a 0.6 B " "Conformer-RNNT fine-tuned from NVIDIA's " "[`nemotron-speech-streaming-en-0.6b`](https://huggingface.co/nvidia/nemotron-speech-streaming-en-0.6b) " "(streaming, cache-aware). Upload or record audio (mono or stereo). " "Stereo audio is transcribed per-channel — useful for call-center " "recordings where the agent and customer are on separate channels.\n\n" "Model: [SkunkWorkLabs/varuna-stt](https://huggingface.co/SkunkWorkLabs/varuna-stt) · " "Benchmark: [SkunkWorkLabs/hindi-asr-benchmark](https://huggingface.co/datasets/SkunkWorkLabs/hindi-asr-benchmark)" ) gr.Markdown( "> ⚡ **Running on ZeroGPU (A10G).** " "**Why streaming + cache-aware?** The encoder consumes audio in " "small chunks (tens of ms) and carries attention KV cache across " "chunks, so it can emit tokens *while* the speaker is still " "talking — low first-token latency, bounded per-chunk compute, " "many concurrent calls per GPU." ) with gr.Row(): with gr.Column(): audio_in = gr.Audio( sources=["upload", "microphone"], type="numpy", label="Audio input (mono or stereo)", ) run_btn = gr.Button("Transcribe", variant="primary") with gr.Column(): info_md = gr.Markdown(label="Audio info") ch1_text = gr.Textbox(label="Transcript (mono / channel 1)", lines=4, max_lines=10) ch2_text = gr.Textbox(label="Channel 2 transcript", lines=4, max_lines=10, visible=False) run_btn.click( fn=transcribe, inputs=audio_in, outputs=[info_md, ch1_text, ch2_text, ch2_text], api_name=False, show_api=False, ) gr.Markdown( "**Output style:** ITN — digits, ordinals (`1st`/`3rd`), " "Indian-numbering commas (`2,50,000`), Devanagari punctuation (`।`).\n\n" "🥇 Best WER on **indictts (9.75 %)** vs ElevenLabs Scribe v1, " "Deepgram Nova-2, and Sarvam Saarika v2.5 on the " "[SkunkWorkLabs Hindi ASR benchmark](https://huggingface.co/datasets/SkunkWorkLabs/hindi-asr-benchmark).\n\n" "📬 **harshris2314@gmail.com**" ) if __name__ == "__main__": demo.queue(api_open=False).launch( server_name="0.0.0.0", server_port=7860, share=os.environ.get("GRADIO_SHARE", "false").lower() == "true", show_api=False, )