Spaces:
Running on Zero
Running on Zero
| import os | |
| # ZeroGPU allocconf: chunked transcriptions allocate variable-size feature | |
| # batches; keep transient spikes from fragmenting the caching allocator. | |
| os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") | |
| import shutil | |
| import tempfile | |
| import time | |
| from pathlib import Path | |
| import spaces # MUST come before torch / transformers | |
| import torch | |
| import gradio as gr | |
| from transformers import AutoModel, AutoProcessor | |
| MODEL_ID = "oruk/oruk-loop-153M" | |
| print("Loading processor ...", flush=True) | |
| processor = AutoProcessor.from_pretrained(MODEL_ID) | |
| print("Loading oruk-loop-153M ...", flush=True) | |
| t0 = time.perf_counter() | |
| # Authors' release path: trust_remote_code runs modeling_oruk_loop.py, which | |
| # defaults to bf16 + merge=True on GPU (per-depth loop weights materialised | |
| # for fast decoding, ~4.2 GB of VRAM). | |
| model = AutoModel.from_pretrained(MODEL_ID, trust_remote_code=True, device="cuda") | |
| print(f"Model ready in {time.perf_counter() - t0:.1f}s", flush=True) | |
| def _normalize_audio(audio_path: str, scratch_dir: str) -> str: | |
| """Convert browser-supplied audio (webm/mp4/mp3/stereo) to 16 kHz mono WAV. | |
| Microphone recordings arrive as .webm, which soundfile cannot read; | |
| pydub (ffmpeg) handles every container a browser can produce. Files that | |
| are already 16 kHz mono wav are passed through untouched, so bundled | |
| examples keep their exact bytes. | |
| """ | |
| from pydub import AudioSegment | |
| audio = AudioSegment.from_file(audio_path) | |
| changed = False | |
| if audio.frame_rate != 16000: | |
| audio = audio.set_frame_rate(16000) | |
| changed = True | |
| if audio.channels > 1: | |
| audio = audio.set_channels(1) | |
| changed = True | |
| if not changed and Path(audio_path).suffix.lower() == ".wav": | |
| return audio_path | |
| out = Path(scratch_dir) / (Path(audio_path).stem + "_16k_mono.wav") | |
| audio.export(out, format="wav") | |
| return out.as_posix() | |
| def _audio_seconds(audio_path) -> float: | |
| try: | |
| from pydub import AudioSegment | |
| return AudioSegment.from_file(audio_path).duration_seconds | |
| except Exception: | |
| return 0.0 | |
| def _estimate_duration(audio_path, *args, **kwargs) -> int: | |
| """Callable ZeroGPU duration, sized from measured runs on this Space. | |
| Measured: 0.84 s for 11 s of audio (beam 4) and 4.08 s for 38 s | |
| (beam 4, two chunks) — about 0.12 s per second of audio, plus ~2 s of | |
| weight streaming on the first request after boot. The budget below is | |
| that with margin (0.35 s/s + 15 s floor), capped for very long uploads. | |
| """ | |
| if not audio_path: | |
| return 30 | |
| seconds = _audio_seconds(audio_path) | |
| return min(240, max(30, 15 + int(seconds * 0.35))) | |
| def transcribe(audio_path: str, num_beams: int = 1) -> str: | |
| """Transcribe English speech to text with oruk-loop-153M. | |
| A 153M-parameter looped-transformer ASR model distilled from | |
| Qwen3-ASR-1.7B: one encoder block applied 24 times and one decoder | |
| block applied 28 times. Recordings longer than 30 s are automatically | |
| cut at quiet points into chunks and the chunk transcripts are joined. | |
| Args: | |
| audio_path: Path to an audio file (upload or microphone recording). | |
| num_beams: Beam size for decoding. 1 = greedy, 4 = the setting of | |
| the model card's headline WER numbers. | |
| Returns: | |
| The transcribed text, with casing and punctuation. | |
| """ | |
| if not audio_path: | |
| raise gr.Error("Provide an audio file or record from your microphone first.") | |
| scratch = tempfile.mkdtemp(prefix="oruk_loop_") | |
| try: | |
| wav_path = _normalize_audio(audio_path, scratch) | |
| beams = max(1, int(num_beams)) | |
| start = time.perf_counter() | |
| text = model.transcribe(wav_path, processor, num_beams=beams) | |
| elapsed = time.perf_counter() - start | |
| seconds = _audio_seconds(wav_path) | |
| print( | |
| f"TRANSCRIBE {elapsed:.2f}s for {seconds:.1f}s audio (beam {beams}) " | |
| f"-> {text[:120]!r}", | |
| flush=True, | |
| ) | |
| return text | |
| finally: | |
| shutil.rmtree(scratch, ignore_errors=True) | |
| CSS = """ | |
| #col-container { max-width: 1000px; margin: 0 auto; } | |
| .dark .gradio-container { color: var(--body-text-color); } | |
| """ | |
| EXAMPLE_FILES = [ | |
| "jfk.wav", | |
| "librispeech_sample.wav", | |
| "eu_toxicity_25s.wav", | |
| "eu_ireland_38s.wav", | |
| ] | |
| with gr.Blocks(title="oruk-loop-153M") as demo: | |
| with gr.Column(elem_id="col-container"): | |
| gr.Markdown( | |
| "# 🎙️ oruk-loop-153M — English ASR with a looped transformer\n\n" | |
| "**oruk-loop-153M** ([oruk/oruk-loop-153M](https://huggingface.co/oruk/oruk-loop-153M)) " | |
| "is an English speech recognizer with **153M stored parameters**, distilled from " | |
| "[Qwen3-ASR-1.7B](https://huggingface.co/Qwen/Qwen3-ASR-1.7B-hf) (13.3× larger). " | |
| "Each transformer stack is a single shared block run in a loop — one encoder block " | |
| "applied 24× and one decoder block applied 28× — with low-rank per-depth adapters. " | |
| "Mean WER on the Open ASR Leaderboard's seven English test sets: **7.12%** with beam 4.\n\n" | |
| "Upload or record audio (recordings longer than 30 s are chunked automatically) " | |
| "and get a transcript with casing and punctuation." | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| audio_in = gr.Audio( | |
| label="Audio", | |
| type="filepath", | |
| sources=["upload", "microphone"], | |
| ) | |
| transcribe_btn = gr.Button("Transcribe", variant="primary") | |
| with gr.Column(): | |
| text_out = gr.Textbox( | |
| label="Transcript", | |
| lines=10, | |
| placeholder="The recognized text will appear here ...", | |
| ) | |
| with gr.Accordion("Decoding settings", open=False): | |
| beams = gr.Radio( | |
| choices=[("Greedy (fast)", 1), ("Beam 4 (best accuracy)", 4)], | |
| value=1, | |
| type="value", | |
| label="Decoding", | |
| info="The model card's headline numbers use beam 4; greedy is faster.", | |
| ) | |
| gr.Examples( | |
| examples=[[f] for f in EXAMPLE_FILES], | |
| inputs=audio_in, | |
| fn=transcribe, | |
| outputs=text_out, | |
| label="Or try one of these", | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| ) | |
| transcribe_btn.click( | |
| transcribe, | |
| inputs=[audio_in, beams], | |
| outputs=text_out, | |
| api_name="transcribe", | |
| ) | |
| demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True) | |