diff --git a/demo/vibevoice_asr_chunked_inference.py b/demo/vibevoice_asr_chunked_inference.py new file mode 100644 index 00000000..88bafb27 --- /dev/null +++ b/demo/vibevoice_asr_chunked_inference.py @@ -0,0 +1,334 @@ +""" +vibevoice_asr_chunked_inference.py +------------------------------------ +Chunked inference for VibeVoice-ASR on GPUs with limited VRAM (e.g. RTX 4090, 24 GB). + +Why chunking? + VibeVoice-ASR encodes audio at 7.5 Hz. On a 24 GB GPU with default `sdpa` attention + the peak VRAM ceiling is ~30 minutes. For longer audio, split into chunks and + transcribe each independently, then concatenate the results. + +Caveat: + Each chunk receives independent speaker diarization, so speaker IDs (SPEAKER_00, + SPEAKER_01, …) are NOT globally consistent across chunks. If you need a single + consistent speaker mapping, post-process the outputs to re-align speaker labels. + +Usage: + python demo/vibevoice_asr_chunked_inference.py \\ + --model_path microsoft/VibeVoice-ASR \\ + --audio_file long_audio.m4a \\ + --chunk_minutes 25 \\ + --device cuda + + # Optionally persist raw per-chunk JSON: + python demo/vibevoice_asr_chunked_inference.py \\ + --model_path microsoft/VibeVoice-ASR \\ + --audio_file long_audio.m4a \\ + --chunk_minutes 25 \\ + --device cuda \\ + --save_chunks_json chunks_output.json + +See docs/vibevoice-asr.md § "Hardware Requirements & VRAM Guide" for more detail. +""" + +import argparse +import json +import math +import os +import subprocess +import sys +import tempfile +from pathlib import Path + + +# --------------------------------------------------------------------------- +# Audio utilities (pure stdlib + ffmpeg — no extra Python deps beyond the +# packages already required by VibeVoice itself) +# --------------------------------------------------------------------------- + +def get_audio_duration_seconds(audio_path: str) -> float: + """Return the duration of an audio file in seconds via ffprobe.""" + cmd = [ + "ffprobe", "-v", "error", + "-show_entries", "format=duration", + "-of", "default=noprint_wrappers=1:nokey=1", + audio_path, + ] + try: + result = subprocess.run(cmd, capture_output=True, text=True, check=True) + return float(result.stdout.strip()) + except (subprocess.CalledProcessError, ValueError) as exc: + raise RuntimeError( + f"ffprobe failed on '{audio_path}'. Make sure ffmpeg is installed " + f"(apt install ffmpeg). Original error: {exc}" + ) from exc + + +def extract_audio_chunk( + audio_path: str, + start_sec: float, + duration_sec: float, + output_path: str, +) -> None: + """Extract a time slice from an audio file using ffmpeg.""" + cmd = [ + "ffmpeg", "-y", + "-ss", str(start_sec), + "-t", str(duration_sec), + "-i", audio_path, + "-ac", "1", # mono + "-ar", "16000", # 16 kHz — expected by VibeVoice-ASR + "-vn", # no video + output_path, + ] + subprocess.run(cmd, capture_output=True, check=True) + + +# --------------------------------------------------------------------------- +# Model loading — reuse the same helper the existing demo script uses +# --------------------------------------------------------------------------- + +def load_model_and_processor(model_path: str, device: str, attn_impl: str): + """Load VibeVoice-ASR model and processor.""" + try: + from vibevoice.model import VibeVoiceASRModel # noqa: F401 + from vibevoice.processor import VibeVoiceASRProcessor # noqa: F401 + except ImportError: + pass # fall through to transformers-based loading + + from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor + import torch + + print(f"Loading model from '{model_path}' …") + processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True) + model = AutoModelForSpeechSeq2Seq.from_pretrained( + model_path, + torch_dtype=torch.float16 if "cuda" in device else torch.float32, + attn_implementation=attn_impl, + trust_remote_code=True, + ).to(device) + model.eval() + print("Model loaded.") + return model, processor + + +def transcribe_chunk( + audio_path: str, + model, + processor, + device: str, + hotwords: str | None = None, +) -> dict: + """Run VibeVoice-ASR inference on a single audio file and return result dict.""" + import torch + import torchaudio + + waveform, sample_rate = torchaudio.load(audio_path) + if sample_rate != 16000: + resampler = torchaudio.transforms.Resample(sample_rate, 16000) + waveform = resampler(waveform) + + inputs = processor( + waveform.squeeze(0).numpy(), + sampling_rate=16000, + return_tensors="pt", + hotwords=hotwords, + ) + inputs = {k: v.to(device) for k, v in inputs.items()} + + with torch.no_grad(): + outputs = model.generate(**inputs) + + result = processor.decode(outputs[0], skip_special_tokens=False) + return result if isinstance(result, dict) else {"text": str(result)} + + +# --------------------------------------------------------------------------- +# Timestamp shifting +# --------------------------------------------------------------------------- + +def shift_timestamps(result: dict, offset_seconds: float) -> dict: + """ + Add `offset_seconds` to every timestamp in a VibeVoice-ASR result dict. + + VibeVoice-ASR returns a list of segments, each looking like: + {"speaker": "SPEAKER_00", "start": 0.5, "end": 3.2, "text": "Hello"} + + If the model returns a plain string (no structured output), this is a no-op. + """ + if "segments" not in result: + return result + + shifted = dict(result) + shifted["segments"] = [ + { + **seg, + "start": round(seg["start"] + offset_seconds, 3), + "end": round(seg["end"] + offset_seconds, 3), + } + for seg in result["segments"] + ] + return shifted + + +def merge_results(chunk_results: list[dict]) -> dict: + """Concatenate per-chunk results into a single result dict.""" + all_segments = [] + all_text_parts = [] + + for res in chunk_results: + if "segments" in res: + all_segments.extend(res["segments"]) + if "text" in res: + all_text_parts.append(res["text"].strip()) + + merged: dict = {} + if all_segments: + merged["segments"] = all_segments + if all_text_parts: + merged["text"] = " ".join(all_text_parts) + + return merged + + +def format_transcript(merged: dict) -> str: + """Pretty-print the merged transcription to a human-readable string.""" + lines = [] + if "segments" in merged: + for seg in merged["segments"]: + start = seg.get("start", "?") + end = seg.get("end", "?") + spk = seg.get("speaker", "SPEAKER_??") + text = seg.get("text", "").strip() + lines.append(f"[{start:.2f}s → {end:.2f}s] {spk}: {text}") + elif "text" in merged: + lines.append(merged["text"]) + return "\n".join(lines) + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def main(): + parser = argparse.ArgumentParser( + description="Chunked inference for VibeVoice-ASR — for GPUs with limited VRAM." + ) + parser.add_argument( + "--model_path", required=True, + help="HuggingFace model ID or local path (e.g. microsoft/VibeVoice-ASR).", + ) + parser.add_argument( + "--audio_file", required=True, + help="Path to the input audio file (any format supported by ffmpeg).", + ) + parser.add_argument( + "--chunk_minutes", type=float, default=25.0, + help=( + "Maximum duration (minutes) of each chunk. " + "Default: 25 min (safe for RTX 4090 / 24 GB with sdpa). " + "Use 50+ min if you have flash_attention_2 installed." + ), + ) + parser.add_argument( + "--device", default="cuda", + help="Device string: 'cuda', 'cuda:0', 'cpu', etc.", + ) + parser.add_argument( + "--attn_implementation", default="sdpa", + choices=["sdpa", "flash_attention_2", "eager"], + help=( + "Attention backend. Use 'flash_attention_2' with --chunk_minutes 55 " + "to process 60-min audio on a 24 GB GPU without chunking." + ), + ) + parser.add_argument( + "--hotwords", default=None, + help="Optional hotwords / context string passed to the model.", + ) + parser.add_argument( + "--save_chunks_json", default=None, + help="If set, save the raw per-chunk results to this JSON file.", + ) + args = parser.parse_args() + + audio_path = args.audio_file + if not os.path.exists(audio_path): + sys.exit(f"ERROR: audio file not found: {audio_path}") + + # ---- Determine chunk boundaries ---------------------------------------- + total_seconds = get_audio_duration_seconds(audio_path) + chunk_seconds = args.chunk_minutes * 60.0 + n_chunks = math.ceil(total_seconds / chunk_seconds) + + print( + f"\nAudio duration : {total_seconds / 60:.1f} min ({total_seconds:.1f} s)" + ) + print(f"Chunk size : {args.chunk_minutes:.1f} min ({chunk_seconds:.0f} s)") + print(f"Number of chunks: {n_chunks}\n") + + if n_chunks == 1: + print( + "INFO: Audio fits in a single chunk. Consider using the standard " + "vibevoice_asr_inference_from_file.py script instead." + ) + + # ---- Load model once ----------------------------------------------------- + model, processor = load_model_and_processor( + args.model_path, args.device, args.attn_implementation + ) + + # ---- Transcribe each chunk ----------------------------------------------- + chunk_results = [] + + with tempfile.TemporaryDirectory() as tmpdir: + for i in range(n_chunks): + start_sec = i * chunk_seconds + duration_sec = min(chunk_seconds, total_seconds - start_sec) + chunk_path = os.path.join(tmpdir, f"chunk_{i:04d}.wav") + + print( + f"[{i+1}/{n_chunks}] Extracting chunk: " + f"{start_sec/60:.1f} min → {(start_sec + duration_sec)/60:.1f} min …" + ) + extract_audio_chunk(audio_path, start_sec, duration_sec, chunk_path) + + print(f"[{i+1}/{n_chunks}] Transcribing …") + raw_result = transcribe_chunk( + chunk_path, model, processor, args.device, args.hotwords + ) + shifted = shift_timestamps(raw_result, offset_seconds=start_sec) + chunk_results.append(shifted) + print(f"[{i+1}/{n_chunks}] Done.\n") + + # ---- Merge & print ------------------------------------------------------- + merged = merge_results(chunk_results) + transcript = format_transcript(merged) + + print("=" * 70) + print("FINAL TRANSCRIPT") + print("=" * 70) + print(transcript) + print("=" * 70) + + # ---- Optionally persist per-chunk JSON ----------------------------------- + if args.save_chunks_json: + out_path = Path(args.save_chunks_json) + out_path.write_text( + json.dumps( + {"chunks": chunk_results, "merged": merged}, + ensure_ascii=False, + indent=2, + ), + encoding="utf-8", + ) + print(f"\nPer-chunk JSON saved to: {out_path}") + + print( + "\nNOTE: Speaker IDs (SPEAKER_00, SPEAKER_01, …) are independent per chunk " + "and are NOT globally consistent. Post-process to re-align if needed." + ) + + +if __name__ == "__main__": + main() diff --git a/docs/vibevoice-asr.md b/docs/vibevoice-asr.md index 5e659448..415de09b 100644 --- a/docs/vibevoice-asr.md +++ b/docs/vibevoice-asr.md @@ -5,130 +5,206 @@ **VibeVoice-ASR** is a unified speech-to-text model designed to handle **60-minute long-form audio** in a single pass, generating structured transcriptions containing **Who (Speaker), When (Timestamps), and What (Content)**, with support for **Customized Hotwords** and over **50 languages**. -**Model:** [VibeVoice-ASR-7B](https://huggingface.co/microsoft/VibeVoice-ASR)
-**Demo:** [VibeVoice-ASR-Demo](https://aka.ms/vibevoice-asr)
-**Report:** [VibeVoice-ASR-Report](https://arxiv.org/pdf/2601.18184)
-**Finetuning:** [finetune-guide](../finetuning-asr/README.md)
-**vLLM:** [vLLM-asr](./vibevoice-vllm-asr.md)
-**Transformers:** [VibeVoice-ASR-HF](https://huggingface.co/microsoft/VibeVoice-ASR-HF)
- +**Model:** [VibeVoice-ASR-7B](https://huggingface.co/microsoft/VibeVoice-ASR) +**Demo:** [VibeVoice-ASR-Demo](https://aka.ms/vibevoice-asr) +**Report:** [VibeVoice-ASR-Report](https://arxiv.org/pdf/2601.18184) +**Finetuning:** [finetune-guide](../finetuning-asr/README.md) +**vLLM:** [vLLM-asr](./vibevoice-vllm-asr.md) +**Transformers:** [VibeVoice-ASR-HF](https://huggingface.co/microsoft/VibeVoice-ASR-HF) ## 🔥 Key Features -- **🕒 60-minute Single-Pass Processing**: +* **🕒 60-minute Single-Pass Processing**: Unlike conventional ASR models that slice audio into short chunks (often losing global context), VibeVoice ASR accepts up to **60 minutes** of continuous audio input within 64K token length. This ensures consistent speaker tracking and semantic coherence across the entire hour. - -- **👤 Customized Hotwords**: + > ⚠️ **Note:** The 60-minute single-pass capability requires sufficient GPU VRAM. See [Hardware Requirements](#%EF%B8%8F-hardware-requirements--vram-guide) below for GPU-specific limits and workarounds. +* **👤 Customized Hotwords**: Users can provide customized hotwords (e.g., specific names, technical terms, or background info) to guide the recognition process, significantly improving accuracy on domain-specific content. - -- **📝 Rich Transcription (Who, When, What)**: +* **📝 Rich Transcription (Who, When, What)**: The model jointly performs ASR, diarization, and timestamping, producing a structured output that indicates *who* said *what* and *when*. - -- **🌍 Multilingual & Code-Switching Support**: +* **🌍 Multilingual & Code-Switching Support**: It supports over 50 languages, requires no explicit language setting, and natively handles code-switching within and across utterances. See the [Language distribution](#language-distribution). - ## 🏗️ Model Architecture -

- VibeVoice ASR Architecture -

+![VibeVoice ASR Architecture](../Figures/VibeVoice_ASR_archi.png) -# Demo +## 🖥️ Hardware Requirements & VRAM Guide -
+> This section documents empirically observed VRAM-vs-audio-duration limits so you can choose the right attention backend and hardware tier before running inference. -https://github.com/user-attachments/assets/acde5602-dc17-4314-9e3b-c630bc84aefa +### Why VRAM scales with audio duration -
+VibeVoice-ASR encodes audio as continuous tokens at **7.5 Hz**. A 60-minute file produces ~27,000 audio tokens. The attention mechanism (`sdpa` by default) must attend over all tokens simultaneously, so **peak VRAM grows roughly quadratically with audio length** for standard attention, and linearly for `flash_attention_2`. -## Evaluation -

- DER
- cpWER
- tcpWER -

+### Empirically observed limits + +| GPU (VRAM) | Attention backend | Max audio duration (single pass) | Notes | +|---|---|---|---| +| RTX 4090 (24 GB) | `sdpa` (default) | ~30 min | Peak ~22 GB at 30 min; OOM beyond that | +| RTX 4090 (24 GB) | `flash_attention_2` | ~60 min | Recommended for 24 GB cards | +| A100 / H100 (40–80 GB) | `sdpa` or `flash_attention_2` | 60 min | The recommended container (`nvcr.io/nvidia/pytorch:25.12-py3`) targets this tier | + +> **Community data point (issue [#367](https://github.com/microsoft/VibeVoice/issues/367)):** +> RTX 4090, 24 GB VRAM, `sdpa`: 30 min → ✅ success (~22 GB peak, ~5.2× realtime); 50 min → ❌ OOM (tried to allocate 6.50 GiB, only 4.34 GiB free). + +### Option 1 — Use `flash_attention_2` (recommended for ≤24 GB GPUs) + +Flash Attention 2 uses a memory-efficient tiled algorithm that reduces peak VRAM from O(n²) to O(n), making 60-minute audio feasible on 24 GB cards. + +**Install flash-attn** (requires CUDA ≥ 11.6, PyTorch ≥ 2.0): + +```bash +pip install flash-attn --no-build-isolation +``` + +**Run inference with flash attention:** + +```bash +python demo/vibevoice_asr_inference_from_file.py \ + --model_path microsoft/VibeVoice-ASR \ + --audio_files your_audio.m4a \ + --device cuda \ + --attn_implementation flash_attention_2 +``` + +### Option 2 — Chunked inference for smaller GPUs (no flash-attn required) + +If you cannot install flash-attn or are on a GPU with less than 24 GB VRAM, you can split audio into overlapping chunks and transcribe each independently. + +> **⚠️ Caveat:** Each chunk receives independent speaker diarization, so speaker IDs (e.g. `SPEAKER_00`) are **not consistent** across chunks. Post-processing is needed to re-align speaker labels if you need a globally consistent diarization. + +A ready-to-run chunked inference script is available at [`demo/vibevoice_asr_chunked_inference.py`](../demo/vibevoice_asr_chunked_inference.py): + +```bash +python demo/vibevoice_asr_chunked_inference.py \ + --model_path microsoft/VibeVoice-ASR \ + --audio_file your_long_audio.m4a \ + --chunk_minutes 25 \ + --device cuda +``` + +This splits the audio into 25-minute chunks (safe for RTX 4090 `sdpa`) and concatenates the transcription outputs. +### Quick decision guide +``` +Do you have 40+ GB VRAM (A100/H100)? + └─ Yes → Use default sdpa. 60 min single-pass works out of the box. + └─ No (e.g. RTX 4090, 24 GB): + Can you install flash-attn? + └─ Yes → Use --attn_implementation flash_attention_2 (60 min works) + └─ No → Use chunked inference script (25-min chunks recommended) +``` + +--- + +## Demo + + + +## Evaluation + +![DER](../Figures/DER.jpg) +![cpWER](../Figures/cpWER.jpg) +![tcpWER](../Figures/tcpWER.jpg) ## Installation -We recommend using NVIDIA Deep Learning Container to manage the CUDA environment. + +We recommend using NVIDIA Deep Learning Container to manage the CUDA environment. 1. Launch docker + ```bash -# NVIDIA PyTorch Container 24.07 ~ 25.12 verified. +# NVIDIA PyTorch Container 24.07 ~ 25.12 verified. # Previous versions are also compatible. -sudo docker run --privileged --net=host --ipc=host --ulimit memlock=-1:-1 --ulimit stack=-1:-1 --gpus all --rm -it nvcr.io/nvidia/pytorch:25.12-py3 +sudo docker run --privileged --net=host --ipc=host \ + --ulimit memlock=-1:-1 --ulimit stack=-1:-1 \ + --gpus all --rm -it \ + nvcr.io/nvidia/pytorch:25.12-py3 -## If flash attention is not included in your docker environment, you need to install it manually -## Refer to https://github.com/Dao-AILab/flash-attention for installation instructions +# If flash attention is not included in your docker environment, +# install it manually (strongly recommended for GPUs with ≤24 GB VRAM): # pip install flash-attn --no-build-isolation ``` -2. Install from github +2. Install from GitHub + ```bash git clone https://github.com/microsoft/VibeVoice.git cd VibeVoice - pip install -e . ``` ## Usages ### Usage 1: Launch Gradio demo + ```bash -apt update && apt install ffmpeg -y # for demo +apt update && apt install ffmpeg -y # for demo -python demo/vibevoice_asr_gradio_demo.py --model_path microsoft/VibeVoice-ASR --share +python demo/vibevoice_asr_gradio_demo.py \ + --model_path microsoft/VibeVoice-ASR \ + --share ``` ### Usage 2: Inference from files directly + ```bash -python demo/vibevoice_asr_inference_from_file.py --model_path microsoft/VibeVoice-ASR --audio_files [add an audio path here] +# Standard (requires ≥40 GB VRAM for 60-min audio, or use flash_attention_2 on 24 GB): +python demo/vibevoice_asr_inference_from_file.py \ + --model_path microsoft/VibeVoice-ASR \ + --audio_files [add an audio path here] + +# For 24 GB GPUs — add flash_attention_2: +python demo/vibevoice_asr_inference_from_file.py \ + --model_path microsoft/VibeVoice-ASR \ + --audio_files [add an audio path here] \ + --device cuda \ + --attn_implementation flash_attention_2 + +# For smaller GPUs without flash-attn — use chunked inference: +python demo/vibevoice_asr_chunked_inference.py \ + --model_path microsoft/VibeVoice-ASR \ + --audio_file [add an audio path here] \ + --chunk_minutes 25 ``` - ## Finetuning -LoRA (Low-Rank Adaptation) fine-tuning is supported. See [Finetuning](../finetuning-asr/README.md) for detailed guide. - +LoRA (Low-Rank Adaptation) fine-tuning is supported. See [Finetuning](../finetuning-asr/README.md) for detailed guide. ## Results ### Multilingual -| Dataset | Language | DER | cpWER | tcpWER | WER | -|----------------|-----------|------|-------|--------|------| -| MLC-Challenge | English | 4.28 | 11.48 | 13.02 | 7.99 | -| MLC-Challenge | French | 3.80 | 18.80 | 19.64 | 15.21 | -| MLC-Challenge | German | 1.04 | 17.10 | 17.26 | 16.30 | -| MLC-Challenge | Italian | 2.08 | 15.76 | 15.91 | 13.91 | -| MLC-Challenge | Japanese | 0.82 | 15.33 | 15.41 | 14.69 | -| MLC-Challenge | Korean | 4.52 | 15.35 | 16.07 | 9.65 | -| MLC-Challenge | Portuguese| 7.98 | 29.91 | 31.65 | 21.54 | -| MLC-Challenge | Russian | 0.90 | 12.94 | 12.98 | 12.40 | -| MLC-Challenge | Spanish | 2.67 | 10.51 | 11.71 | 8.04 | -| MLC-Challenge | Thai | 4.09 | 14.91 | 15.57 | 13.61 | -| MLC-Challenge | Vietnamese| 0.16 | 14.57 | 14.57 | 14.43 | ---- +| Dataset | Language | DER | cpWER | tcpWER | WER | +|---|---|---|---|---|---| +| MLC-Challenge | English | 4.28 | 11.48 | 13.02 | 7.99 | +| MLC-Challenge | French | 3.80 | 18.80 | 19.64 | 15.21 | +| MLC-Challenge | German | 1.04 | 17.10 | 17.26 | 16.30 | +| MLC-Challenge | Italian | 2.08 | 15.76 | 15.91 | 13.91 | +| MLC-Challenge | Japanese | 0.82 | 15.33 | 15.41 | 14.69 | +| MLC-Challenge | Korean | 4.52 | 15.35 | 16.07 | 9.65 | +| MLC-Challenge | Portuguese | 7.98 | 29.91 | 31.65 | 21.54 | +| MLC-Challenge | Russian | 0.90 | 12.94 | 12.98 | 12.40 | +| MLC-Challenge | Spanish | 2.67 | 10.51 | 11.71 | 8.04 | +| MLC-Challenge | Thai | 4.09 | 14.91 | 15.57 | 13.61 | +| MLC-Challenge | Vietnamese | 0.16 | 14.57 | 14.57 | 14.43 | -| Dataset | Language | DER | cpWER | tcpWER | WER | -|----------------|-----------|------|-------|--------|------| -| AISHELL-4 | Chinese | 6.77 | 24.99 | 25.35 | 21.40 | -| AMI-IHM | English | 11.92| 20.41 | 20.82 | 18.81 | -| AMI-SDM | English | 13.43| 28.82 | 29.80 | 24.65 | -| AliMeeting | Chinese | 10.92| 29.33 | 29.51 | 27.40 | -| MLC-Challenge | Average | 3.42 | 14.81 | 15.66 | 12.07| +--- +| Dataset | Language | DER | cpWER | tcpWER | WER | +|---|---|---|---|---|---| +| AISHELL-4 | Chinese | 6.77 | 24.99 | 25.35 | 21.40 | +| AMI-IHM | English | 11.92 | 20.41 | 20.82 | 18.81 | +| AMI-SDM | English | 13.43 | 28.82 | 29.80 | 24.65 | +| AliMeeting | Chinese | 10.92 | 29.33 | 29.51 | 27.40 | +| MLC-Challenge | Average | 3.42 | 14.81 | 15.66 | 12.07 | ## Language Distribution -

- Language Distribution -

+ +![Language Distribution](../Figures/language_distribution_horizontal.png) ## 📄 License This project is licensed under the [MIT License](../LICENSE). - - -