commit b0897668a95c3a8bc13145c7e4cf9a9b90ec11d1 Author: wehub-resource-sync Date: Mon Jul 13 12:05:25 2026 +0800 chore: import upstream snapshot with attribution diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..10b4c39 --- /dev/null +++ b/.gitignore @@ -0,0 +1,178 @@ +# Initially taken from Github's Python gitignore file + +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# tests and logs +tests/fixtures/cached_*_text.txt +logs/ +lightning_logs/ +lang_code_data/ + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +.hypothesis/ +.pytest_cache/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +.python-version + +# celery beat schedule file +celerybeat-schedule + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# vscode +.vs +.vscode + +# Pycharm +.idea + +# TF code +tensorflow_code + +# Models +proc_data + +# examples +runs +/runs_old +/wandb +/examples/runs +/examples/**/*.args +/examples/rag/sweep + +# data +/data +serialization_dir + +# emacs +*.*~ +debug.env + +# vim +.*.swp + +#ctags +tags + +# pre-commit +.pre-commit* + +# .lock +*.lock + +# DS_Store (MacOS) +.DS_Store + +# ruff +.ruff_cache + +# our proj +/output/ +/outputs/ +/checkpoint/ +/checkpoints/ +exp +.gradio/ +experimental_voices \ No newline at end of file diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..372d08c --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,43 @@ +# Contributing to VibeVoice + +Thank you for your interest in **VibeVoice**. VibeVoice is an academic-oriented research project. +Our core principles are **code minimalism**, **high readability**, and **functional purity**. + +## 1. Preferred Contributions + +- **Bug Fixes & New Features** + These are our primary focus. + +- **Minimalism First** + Code must be concise, clear, and minimal. + +## 2. Rejected Patterns + +- **Over-Engineering** + We will reject unnecessary encapsulation, excessive abstraction, or complex architectural refactoring. + Please remember this is research-oriented code, not a commercial enterprise project. + +- **Style Tweaks** + PRs solely for formatting, beautification, or non-functional style adjustments will be rejected. + +- **Non-English Contributions** + PRs containing non-English code comments, documentation, commit messages, or descriptions will be rejected. + +## 3. Review Policy + +- **Line-by-Line Review** + Maintainers strictly audit every single line of change manually. Before submitting, ensure you have personally scrutinized and fully mastered all code logic. + +- **Caution Regarding AI-Generated Code** + LLM-generated code often contains redundant logic and subtle defects. Large chunks of AI-generated code will be rejected unless they have undergone rigorous human cleaning and verification. + Ensure that every line of code has an absolute necessity to exist. + +## 4. Documentation Standards + +- **Precise & Concise** + VibeVoice's documentation serves a global research community. We hold extremely high standards: changes must be succinct and accurate. + Please eliminate all verbosity and redundancy to ensure maximum information density. + +--- + +We look forward to receiving **precise, streamlined, and substantively valuable** research contributions. diff --git a/Figures/DER.jpg b/Figures/DER.jpg new file mode 100644 index 0000000..6d87843 Binary files /dev/null and b/Figures/DER.jpg differ diff --git a/Figures/MOS-preference.png b/Figures/MOS-preference.png new file mode 100644 index 0000000..3e1d21b Binary files /dev/null and b/Figures/MOS-preference.png differ diff --git a/Figures/VibeVoice-TTS-results.jpg b/Figures/VibeVoice-TTS-results.jpg new file mode 100644 index 0000000..e1e4e78 Binary files /dev/null and b/Figures/VibeVoice-TTS-results.jpg differ diff --git a/Figures/VibeVoice.jpg b/Figures/VibeVoice.jpg new file mode 100644 index 0000000..4a99d78 Binary files /dev/null and b/Figures/VibeVoice.jpg differ diff --git a/Figures/VibeVoice_ASR_archi.png b/Figures/VibeVoice_ASR_archi.png new file mode 100644 index 0000000..024eca5 Binary files /dev/null and b/Figures/VibeVoice_ASR_archi.png differ diff --git a/Figures/VibeVoice_Realtime.png b/Figures/VibeVoice_Realtime.png new file mode 100644 index 0000000..d2233ec Binary files /dev/null and b/Figures/VibeVoice_Realtime.png differ diff --git a/Figures/VibeVoice_logo.png b/Figures/VibeVoice_logo.png new file mode 100644 index 0000000..2619848 Binary files /dev/null and b/Figures/VibeVoice_logo.png differ diff --git a/Figures/VibeVoice_logo_white.png b/Figures/VibeVoice_logo_white.png new file mode 100644 index 0000000..ba60c92 Binary files /dev/null and b/Figures/VibeVoice_logo_white.png differ diff --git a/Figures/cpWER.jpg b/Figures/cpWER.jpg new file mode 100644 index 0000000..d5ef00b Binary files /dev/null and b/Figures/cpWER.jpg differ diff --git a/Figures/language_distribution_horizontal.png b/Figures/language_distribution_horizontal.png new file mode 100644 index 0000000..47f7095 Binary files /dev/null and b/Figures/language_distribution_horizontal.png differ diff --git a/Figures/tcpWER.jpg b/Figures/tcpWER.jpg new file mode 100644 index 0000000..b0dcb01 Binary files /dev/null and b/Figures/tcpWER.jpg differ diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..269a897 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2025 Microsoft + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 index 0000000..3e54638 --- /dev/null +++ b/README.md @@ -0,0 +1,207 @@ +
+ +## 🎙️ VibeVoice: Open-Source Frontier Voice AI +[![Project Page](https://img.shields.io/badge/Project-Page-blue?logo=githubpages)](https://microsoft.github.io/VibeVoice) +[![Hugging Face](https://img.shields.io/badge/HuggingFace-Collection-orange?logo=huggingface)](https://huggingface.co/collections/microsoft/vibevoice-68a2ef24a875c44be47b034f) +[![TTS Report](https://img.shields.io/badge/TTS-Report-red?logo=arxiv)](https://openreview.net/pdf?id=FihSkzyxdv) +[![ASR Report](https://img.shields.io/badge/ASR-Report-yellow?logo=arxiv)](https://arxiv.org/pdf/2601.18184) +[![Colab](https://img.shields.io/badge/StreamingTTS-Colab-green?logo=googlecolab)](https://colab.research.google.com/github/microsoft/VibeVoice/blob/main/demo/VibeVoice_colab.ipynb) +[![ASR Playground](https://img.shields.io/badge/ASR-Playground-6F42C1?logo=gradio)](https://aka.ms/vibevoice-asr) + +[![microsoft%2FVibeVoice | Trendshift](https://trendshift.io/api/badge/repositories/15465)](https://trendshift.io/repositories/15465) + +
+ + +
+ + + VibeVoice Logo + +
+ +
+ +

📰 News

+ + + + +2026-03-06: 🚀 VibeVoice ASR is now part of a Transformers release! You can now use our speech recognition model directly through the Hugging Face Transformers library for seamless integration into your projects. + +2026-01-21: 📣 We open-sourced VibeVoice-ASR, 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 User-Customized Context. Try it in [Playground](https://aka.ms/vibevoice-asr). +- ⭐️ VibeVoice-ASR is natively multilingual, supporting over 50 languages — check the [supported languages](docs/vibevoice-asr.md#language-distribution) for details. +- 🔥 The VibeVoice-ASR [finetuning code](finetuning-asr/README.md) is now available! +- ⚡️ **vLLM inference** is now supported for faster inference; see [vllm-asr](docs/vibevoice-vllm-asr.md) for more details. +- 📑 [VibeVoice-ASR Technique Report](https://arxiv.org/pdf/2601.18184) is available. + +2025-12-16: 📣 We added experimental speakers to VibeVoice‑Realtime‑0.5B for exploration, including multilingual voices in nine languages (DE, FR, IT, JP, KR, NL, PL, PT, ES) and 11 distinct English style voices. [Try it](docs/vibevoice-realtime-0.5b.md#optional-more-experimental-voices). More speaker types will be added over time. + +2025-12-03: 📣 We open-sourced VibeVoice‑Realtime‑0.5B, a real‑time text‑to‑speech model that supports streaming text input and robust long-form speech generation. Try it on [Colab](https://colab.research.google.com/github/microsoft/VibeVoice/blob/main/demo/vibevoice_realtime_colab.ipynb). + + +2025-09-05: VibeVoice is an open-source research framework intended to advance collaboration in the speech synthesis community. After release, we discovered instances where the tool was used in ways inconsistent with the stated intent. Since responsible use of AI is one of Microsoft’s guiding principles, we have removed the VibeVoice-TTS code from this repository. + + +2025-08-25: 📣 We open-sourced VibeVoice-TTS, a long-form multi-speaker text-to-speech model that can synthesize speech up to 90 minutes long with up to 4 distinct speakers. — accepted as an [Oral](https://openreview.net/forum?id=FihSkzyxdv) at ICLR 2026! 🔥 + +
+ +## Overview + +VibeVoice is a **family of open-source frontier voice AI models** that includes both Text-to-Speech (TTS) and Automatic Speech Recognition (ASR) models. + +A core innovation of VibeVoice is its use of continuous speech tokenizers (Acoustic and Semantic) operating at an ultra-low frame rate of **7.5 Hz**. These tokenizers efficiently preserve audio fidelity while significantly boosting computational efficiency for processing long sequences. VibeVoice employs a [next-token diffusion](https://arxiv.org/abs/2412.08635) framework, leveraging a Large Language Model (LLM) to understand textual context and dialogue flow, and a diffusion head to generate high-fidelity acoustic details. + +For more information, demos, and examples, please visit our [Project Page](https://microsoft.github.io/VibeVoice). + + +
+ +| Model | Weight | Quick Try | +|-------|--------------|---------| +| VibeVoice-ASR-7B | [HF Link](https://huggingface.co/microsoft/VibeVoice-ASR) | [Playground](https://aka.ms/vibevoice-asr) | +| VibeVoice-TTS-1.5B | [HF Link](https://huggingface.co/microsoft/VibeVoice-1.5B) | Disabled | +| VibeVoice-Realtime-0.5B | [HF Link](https://huggingface.co/microsoft/VibeVoice-Realtime-0.5B) | [Colab](https://colab.research.google.com/github/microsoft/VibeVoice/blob/main/demo/vibevoice_realtime_colab.ipynb) | + +
+ +## Models + + +### 1. 📖 [VibeVoice-ASR](docs/vibevoice-asr.md) - Long-form Speech Recognition + +**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**. + +- **🕒 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**: + 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)**: + The model jointly performs ASR, diarization, and timestamping, producing a structured output that indicates *who* said *what* and *when*. + +[📖 Documentation](docs/vibevoice-asr.md) | [🤗 Hugging Face](https://huggingface.co/microsoft/VibeVoice-ASR) | [🎮 Playground](https://aka.ms/vibevoice-asr) | [🛠️ Finetuning](finetuning-asr/README.md) | [📊 Paper](docs/VibeVoice-ASR-Report.pdf) + + +

+ DER
+ cpWER
+ tcpWER +

+ + +
+ +https://github.com/user-attachments/assets/acde5602-dc17-4314-9e3b-c630bc84aefa + +
+
+ +### 2. 🎙️ [VibeVoice-TTS](docs/vibevoice-tts.md) - Long-form Multi-speaker TTS + +**Best for**: Long-form conversational audio, podcasts, multi-speaker dialogues + +- **⏱️ 90-minute Long-form Generation**: + Synthesizes conversational/single-speaker speech up to **90 minutes** in a single pass, maintaining speaker consistency and semantic coherence throughout. + +- **👥 Multi-speaker Support**: + Supports up to **4 distinct speakers** in a single conversation, with natural turn-taking and speaker consistency across long dialogues. + +- **🎭 Expressive Speech**: + Generates expressive, natural-sounding speech that captures conversational dynamics and emotional nuances. + +- **🌐 Multi-lingual Support**: + Supports English, Chinese and other languages. + + +[📖 Documentation](docs/vibevoice-tts.md) | [🤗 Hugging Face](https://huggingface.co/microsoft/VibeVoice-1.5B) | [📊 Paper](https://arxiv.org/pdf/2508.19205) + + +
+ VibeVoice Results +
+ + +**English** +
+ +https://github.com/user-attachments/assets/0967027c-141e-4909-bec8-091558b1b784 + +
+ + +**Chinese** +
+ +https://github.com/user-attachments/assets/322280b7-3093-4c67-86e3-10be4746c88f + +
+ +**Cross-Lingual** +
+ +https://github.com/user-attachments/assets/838d8ad9-a201-4dde-bb45-8cd3f59ce722 + +
+ +**Spontaneous Singing** +
+ +https://github.com/user-attachments/assets/6f27a8a5-0c60-4f57-87f3-7dea2e11c730 + +
+ + +**Long Conversation with 4 people** +
+ +https://github.com/user-attachments/assets/a357c4b6-9768-495c-a576-1618f6275727 + +
+ + + + + +
+ +### 3. ⚡ [VibeVoice-Streaming](docs/vibevoice-realtime-0.5b.md) - Real-time Streaming TTS + +VibeVoice-Realtime is a **lightweight real‑time** text-to-speech model supporting **streaming text input** and **robust long-form speech generation**. + +- Parameter size: 0.5B (deployment-friendly) +- Real-time TTS (~300 milliseconds first audible latency) +- Streaming text input +- Robust long-form speech generation (~10 minutes) + +[📖 Documentation](docs/vibevoice-realtime-0.5b.md) | [🤗 Hugging Face](https://huggingface.co/microsoft/VibeVoice-Realtime-0.5B) | [🚀 Colab](https://colab.research.google.com/github/microsoft/VibeVoice/blob/main/demo/vibevoice_realtime_colab.ipynb) + + +
+ +https://github.com/user-attachments/assets/0901d274-f6ae-46ef-a0fd-3c4fba4f76dc + +
+ +
+ +## Contributing + +Please see [CONTRIBUTING.md](CONTRIBUTING.md) for detailed contribution guidelines. + + + +## ⚠️ Risks and Limitations + + +While efforts have been made to optimize it through various techniques, it may still produce outputs that are unexpected, biased, or inaccurate. VibeVoice inherits any biases, errors, or omissions produced by its base model (specifically, Qwen2.5 1.5b in this release). +Potential for Deepfakes and Disinformation: High-quality synthetic speech can be misused to create convincing fake audio content for impersonation, fraud, or spreading disinformation. Users must ensure transcripts are reliable, check content accuracy, and avoid using generated content in misleading ways. Users are expected to use the generated content and to deploy the models in a lawful manner, in full compliance with all applicable laws and regulations in the relevant jurisdictions. It is best practice to disclose the use of AI when sharing AI-generated content. + + +We do not recommend using VibeVoice in commercial or real-world applications without further testing and development. This model is intended for research and development purposes only. Please use responsibly. + +## Star History + +![Star History Chart](https://api.star-history.com/svg?repos=Microsoft/vibevoice&type=date&legend=top-left) diff --git a/README.wehub.md b/README.wehub.md new file mode 100644 index 0000000..2889a73 --- /dev/null +++ b/README.wehub.md @@ -0,0 +1,7 @@ +# WeHub 来源说明 + +- 原始项目：`microsoft/VibeVoice` +- 原始仓库：https://github.com/microsoft/VibeVoice +- 导入方式：上游默认分支的最新快照 +- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准 +- 本文件仅用于记录来源，不代表 WeHub 是原项目作者 diff --git a/SECURITY.md b/SECURITY.md new file mode 100644 index 0000000..e751608 --- /dev/null +++ b/SECURITY.md @@ -0,0 +1,14 @@ + + +## Security + +Microsoft takes the security of our software products and services seriously, which +includes all source code repositories in our GitHub organizations. + +**Please do not report security vulnerabilities through public GitHub issues.** + +For security reporting information, locations, contact information, and policies, +please review the latest guidance for Microsoft repositories at +[https://aka.ms/SECURITY.md](https://aka.ms/SECURITY.md). + + \ No newline at end of file diff --git a/demo/asr_demo/demo1-chat.mp3 b/demo/asr_demo/demo1-chat.mp3 new file mode 100644 index 0000000..0b5b35f Binary files /dev/null and b/demo/asr_demo/demo1-chat.mp3 differ diff --git a/demo/asr_demo/demo1-chat.mp4 b/demo/asr_demo/demo1-chat.mp4 new file mode 100644 index 0000000..2ebefd7 Binary files /dev/null and b/demo/asr_demo/demo1-chat.mp4 differ diff --git a/demo/asr_demo/demo2-song.mp3 b/demo/asr_demo/demo2-song.mp3 new file mode 100644 index 0000000..ba49a8b Binary files /dev/null and b/demo/asr_demo/demo2-song.mp3 differ diff --git a/demo/asr_demo/demo2-song.mp4 b/demo/asr_demo/demo2-song.mp4 new file mode 100644 index 0000000..d4bc559 Binary files /dev/null and b/demo/asr_demo/demo2-song.mp4 differ diff --git a/demo/asr_demo/demo3-hotwords.wav b/demo/asr_demo/demo3-hotwords.wav new file mode 100644 index 0000000..0cefdf3 Binary files /dev/null and b/demo/asr_demo/demo3-hotwords.wav differ diff --git a/demo/download_experimental_voices.sh b/demo/download_experimental_voices.sh new file mode 100644 index 0000000..ce8c3ca --- /dev/null +++ b/demo/download_experimental_voices.sh @@ -0,0 +1,48 @@ +#!/usr/bin/env bash +set -e + +echo "[INFO] Starting download of experimental voices..." + +# Absolute path of the current script directory +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +# Target directory relative to this script location +TARGET_DIR="$SCRIPT_DIR/voices/streaming_model/experimental_voices" + +echo "[INFO] Script directory: $SCRIPT_DIR" +echo "[INFO] Target directory: $TARGET_DIR" + +# Ensure the target directory exists +echo "[INFO] Creating target directory if needed..." +mkdir -p "$TARGET_DIR" + +# List of archives and their URLs +FILES=( + "experimental_voices_de.tar.gz|https://github.com/user-attachments/files/24035887/experimental_voices_de.tar.gz" + "experimental_voices_fr.tar.gz|https://github.com/user-attachments/files/24035880/experimental_voices_fr.tar.gz" + "experimental_voices_jp.tar.gz|https://github.com/user-attachments/files/24035882/experimental_voices_jp.tar.gz" + "experimental_voices_kr.tar.gz|https://github.com/user-attachments/files/24035883/experimental_voices_kr.tar.gz" + "experimental_voices_pl.tar.gz|https://github.com/user-attachments/files/24035885/experimental_voices_pl.tar.gz" + "experimental_voices_pt.tar.gz|https://github.com/user-attachments/files/24035886/experimental_voices_pt.tar.gz" + "experimental_voices_sp.tar.gz|https://github.com/user-attachments/files/24035884/experimental_voices_sp.tar.gz" + "experimental_voices_en1.tar.gz|https://github.com/user-attachments/files/24189272/experimental_voices_en1.tar.gz" + "experimental_voices_en2.tar.gz|https://github.com/user-attachments/files/24189273/experimental_voices_en2.tar.gz" +) + +# Download, extract, and clean up each archive +for entry in "${FILES[@]}"; do + IFS="|" read -r FNAME URL <<< "$entry" + + echo "[INFO] Downloading $FNAME ..." + wget -O "$FNAME" "$URL" + + echo "[INFO] Extracting $FNAME ..." + tar -xzvf "$FNAME" -C "$TARGET_DIR" + + echo "[INFO] Cleaning up $FNAME ..." + rm -f "$FNAME" +done + +echo "[SUCCESS] All experimental speakers installed successfully!" +echo "[SUCCESS] Speakers are located at:" +echo " $TARGET_DIR" diff --git a/demo/realtime_model_inference_from_file.py b/demo/realtime_model_inference_from_file.py new file mode 100644 index 0000000..2a2e711 --- /dev/null +++ b/demo/realtime_model_inference_from_file.py @@ -0,0 +1,311 @@ +import argparse +import os +import re +import traceback +from typing import List, Tuple, Union, Dict, Any +import time +import torch +import copy +import glob + +from vibevoice.modular.modeling_vibevoice_streaming_inference import VibeVoiceStreamingForConditionalGenerationInference +from vibevoice.processor.vibevoice_streaming_processor import VibeVoiceStreamingProcessor +from transformers.cache_utils import DynamicCache +from transformers.modeling_outputs import BaseModelOutputWithPast +from transformers.utils import logging + +logging.set_verbosity_info() +logger = logging.get_logger(__name__) + + +class VoiceMapper: + """Maps speaker names to voice file paths""" + + def __init__(self): + self.setup_voice_presets() + # for k, v in self.voice_presets.items(): + # print(f"{k}: {v}") + + def setup_voice_presets(self): + """Setup voice presets by scanning the voices directory.""" + voices_dir = os.path.join(os.path.dirname(__file__), "voices/streaming_model") + + # Check if voices directory exists + if not os.path.exists(voices_dir): + print(f"Warning: Voices directory not found at {voices_dir}") + self.voice_presets = {} + self.available_voices = {} + return + + # Scan for all VOICE files in the voices directory + self.voice_presets = {} + + # Get all .pt files in the voices directory + pt_files = glob.glob(os.path.join(voices_dir, "**", "*.pt"), recursive=True) + + # Create dictionary with filename (without extension) as key + for pt_file in pt_files: + # key: filename without extension + name = os.path.splitext(os.path.basename(pt_file))[0].lower() + full_path = os.path.abspath(pt_file) + self.voice_presets[name] = full_path + + # Sort the voice presets alphabetically by name for better UI + self.voice_presets = dict(sorted(self.voice_presets.items())) + + # Filter out voices that don't exist (this is now redundant but kept for safety) + self.available_voices = { + name: path for name, path in self.voice_presets.items() + if os.path.exists(path) + } + + print(f"Found {len(self.available_voices)} voice files in {voices_dir}") + print(f"Available voices: {', '.join(self.available_voices.keys())}") + + def get_voice_path(self, speaker_name: str) -> str: + """Get voice file path for a given speaker name""" + # First try exact match + speaker_name = speaker_name.lower() + if speaker_name in self.voice_presets: + return self.voice_presets[speaker_name] + + # Try partial matching (case insensitive) + matched_path = None + for preset_name, path in self.voice_presets.items(): + if preset_name.lower() in speaker_name or speaker_name in preset_name.lower(): + if matched_path is not None: + raise ValueError(f"Multiple voice presets match the speaker name '{speaker_name}', please make the speaker_name more specific.") + matched_path = path + if matched_path is not None: + return matched_path + + # Default to first voice if no match found + default_voice = list(self.voice_presets.values())[0] + print(f"Warning: No voice preset found for '{speaker_name}', using default voice: {default_voice}") + return default_voice + + +def parse_args(): + parser = argparse.ArgumentParser(description="VibeVoiceStreaming Processor TXT Input Test") + parser.add_argument( + "--model_path", + type=str, + default="microsoft/VibeVoice-Realtime-0.5B", + help="Path to the HuggingFace model directory", + ) + parser.add_argument( + "--txt_path", + type=str, + default="demo/text_examples/1p_vibevoice.txt", + help="Path to the txt file containing the script", + ) + parser.add_argument( + "--speaker_name", + type=str, + default="Wayne", + help="Single speaker name (e.g., --speaker_name Wayne)", + ) + parser.add_argument( + "--output_dir", + type=str, + default="./outputs", + help="Directory to save output audio files", + ) + parser.add_argument( + "--device", + type=str, + default=("cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")), + help="Device for inference: cuda | mps | cpu", + ) + parser.add_argument( + "--cfg_scale", + type=float, + default=1.5, + help="CFG (Classifier-Free Guidance) scale for generation (default: 1.5)", + ) + + return parser.parse_args() + +def main(): + args = parse_args() + + # Normalize potential 'mpx' typo to 'mps' + if args.device.lower() == "mpx": + print("Note: device 'mpx' detected, treating it as 'mps'.") + args.device = "mps" + + # Validate mps availability if requested + if args.device == "mps" and not torch.backends.mps.is_available(): + print("Warning: MPS not available. Falling back to CPU.") + args.device = "cpu" + + print(f"Using device: {args.device}") + + # Initialize voice mapper + voice_mapper = VoiceMapper() + + # Check if txt file exists + if not os.path.exists(args.txt_path): + print(f"Error: txt file not found: {args.txt_path}") + return + + # Read and parse txt file + print(f"Reading script from: {args.txt_path}") + with open(args.txt_path, 'r', encoding='utf-8') as f: + scripts = f.read().strip() + + if not scripts: + print("Error: No valid scripts found in the txt file") + return + + full_script = scripts.replace("’", "'").replace('“', '"').replace('”', '"') + + print(f"Loading processor & model from {args.model_path}") + processor = VibeVoiceStreamingProcessor.from_pretrained(args.model_path) + + # Decide dtype & attention implementation + if args.device == "mps": + load_dtype = torch.float32 # MPS requires float32 + attn_impl_primary = "sdpa" # flash_attention_2 not supported on MPS + elif args.device == "cuda": + load_dtype = torch.bfloat16 + attn_impl_primary = "flash_attention_2" + else: # cpu + load_dtype = torch.float32 + attn_impl_primary = "sdpa" + print(f"Using device: {args.device}, torch_dtype: {load_dtype}, attn_implementation: {attn_impl_primary}") + # Load model with device-specific logic + try: + if args.device == "mps": + model = VibeVoiceStreamingForConditionalGenerationInference.from_pretrained( + args.model_path, + torch_dtype=load_dtype, + attn_implementation=attn_impl_primary, + device_map=None, # load then move + ) + model.to("mps") + elif args.device == "cuda": + model = VibeVoiceStreamingForConditionalGenerationInference.from_pretrained( + args.model_path, + torch_dtype=load_dtype, + device_map="cuda", + attn_implementation=attn_impl_primary, + ) + else: # cpu + model = VibeVoiceStreamingForConditionalGenerationInference.from_pretrained( + args.model_path, + torch_dtype=load_dtype, + device_map="cpu", + attn_implementation=attn_impl_primary, + ) + except Exception as e: + if attn_impl_primary == 'flash_attention_2': + print(f"[ERROR] : {type(e).__name__}: {e}") + print(traceback.format_exc()) + print("Error loading the model. Trying to use SDPA. However, note that only flash_attention_2 has been fully tested, and using SDPA may result in lower audio quality.") + model = VibeVoiceStreamingForConditionalGenerationInference.from_pretrained( + args.model_path, + torch_dtype=load_dtype, + device_map=(args.device if args.device in ("cuda", "cpu") else None), + attn_implementation='sdpa' + ) + if args.device == "mps": + model.to("mps") + else: + raise e + + + model.eval() + model.set_ddpm_inference_steps(num_steps=5) + + if hasattr(model.model, 'language_model'): + print(f"Language model attention: {model.model.language_model.config._attn_implementation}") + + target_device = args.device if args.device != "cpu" else "cpu" + voice_sample = voice_mapper.get_voice_path(args.speaker_name) + print(f"Using voice preset for {args.speaker_name}: {voice_sample}") + with torch.serialization.safe_globals([BaseModelOutputWithPast, DynamicCache]): + all_prefilled_outputs = torch.load(voice_sample, map_location=target_device, weights_only=True) + + # Prepare inputs for the model + inputs = processor.process_input_with_cached_prompt( + text=full_script, + cached_prompt=all_prefilled_outputs, + padding=True, + return_tensors="pt", + return_attention_mask=True, + ) + + # Move tensors to target device + for k, v in inputs.items(): + if torch.is_tensor(v): + inputs[k] = v.to(target_device) + + print(f"Starting generation with cfg_scale: {args.cfg_scale}") + + # Generate audio + start_time = time.time() + outputs = model.generate( + **inputs, + max_new_tokens=None, + cfg_scale=args.cfg_scale, + tokenizer=processor.tokenizer, + generation_config={'do_sample': False}, + verbose=True, + all_prefilled_outputs=copy.deepcopy(all_prefilled_outputs) if all_prefilled_outputs is not None else None, + ) + generation_time = time.time() - start_time + print(f"Generation time: {generation_time:.2f} seconds") + + # Calculate audio duration and additional metrics + if outputs.speech_outputs and outputs.speech_outputs[0] is not None: + # Assuming 24kHz sample rate (common for speech synthesis) + sample_rate = 24000 + audio_samples = outputs.speech_outputs[0].shape[-1] if len(outputs.speech_outputs[0].shape) > 0 else len(outputs.speech_outputs[0]) + audio_duration = audio_samples / sample_rate + rtf = generation_time / audio_duration if audio_duration > 0 else float('inf') + + print(f"Generated audio duration: {audio_duration:.2f} seconds") + print(f"RTF (Real Time Factor): {rtf:.2f}x") + else: + print("No audio output generated") + return + + # Calculate token metrics + input_tokens = inputs['tts_text_ids'].shape[1] # Number of input tokens + output_tokens = outputs.sequences.shape[1] # Total tokens (input + generated) + generated_tokens = output_tokens - input_tokens - all_prefilled_outputs['tts_lm']['last_hidden_state'].size(1) + + print(f"Prefilling text tokens: {input_tokens}") + print(f"Generated speech tokens: {generated_tokens}") + print(f"Total tokens: {output_tokens}") + + # Save output (processor handles device internally) + txt_filename = os.path.splitext(os.path.basename(args.txt_path))[0] + output_path = os.path.join(args.output_dir, f"{txt_filename}_generated.wav") + os.makedirs(args.output_dir, exist_ok=True) + + processor.save_audio( + outputs.speech_outputs[0], # First (and only) batch item + output_path=output_path, + ) + print(f"Saved output to {output_path}") + + # Print summary + print("\n" + "="*50) + print("GENERATION SUMMARY") + print("="*50) + print(f"Input file: {args.txt_path}") + print(f"Output file: {output_path}") + print(f"Speaker names: {args.speaker_name}") + print(f"Prefilling text tokens: {input_tokens}") + print(f"Generated speech tokens: {generated_tokens}") + print(f"Total tokens: {output_tokens}") + print(f"Generation time: {generation_time:.2f} seconds") + print(f"Audio duration: {audio_duration:.2f} seconds") + print(f"RTF (Real Time Factor): {rtf:.2f}x") + + print("="*50) + +if __name__ == "__main__": + main() diff --git a/demo/text_examples/1p_abs.txt b/demo/text_examples/1p_abs.txt new file mode 100644 index 0000000..32e5c64 --- /dev/null +++ b/demo/text_examples/1p_abs.txt @@ -0,0 +1,2 @@ +Generating long-form, multi-speaker conversational audio like podcasts poses significant challenges for traditional Text-to-Speech (TTS) systems, particularly in scalability, speaker consistency, and natural turn-taking. This report presents VibeVoice, a novel model designed to synthesize long-form speech with multiple speakers by employing the next-token diffusion framework, a unified method for modeling continuous data by autoregressively generating latent vectors via diffusion. +A core component of our approach is the continuous speech tokenizers operating at an ultra-low frame rate of 7.5. This tokenizer effectively preserves audio fidelity while significantly boosting computational efficiency for processing long sequences. This enables VibeVoice to synthesize long-form speech for up to 90 minutes (in a 64K context window length) with up to 4 speakers, capturing the authentic conversational "vibe" and surpassing all known open-source and closed-source dialogue models (for example, Gemini 2.5 Pro Preview TTS). Code and checkpoint are available now. \ No newline at end of file diff --git a/demo/text_examples/1p_vibevoice.txt b/demo/text_examples/1p_vibevoice.txt new file mode 100644 index 0000000..22d74dd --- /dev/null +++ b/demo/text_examples/1p_vibevoice.txt @@ -0,0 +1 @@ +VibeVoice is a novel framework designed for generating expressive, long-form, multi-speaker conversational audio, such as podcasts, from text. It addresses significant challenges in traditional Text-to-Speech (TTS) systems, particularly in scalability, speaker consistency, and natural turn-taking. A core innovation of VibeVoice is its use of continuous speech tokenizers operating at an ultra-low frame rate of 7.5 Hz. These tokenizers efficiently preserve audio fidelity while significantly boosting computational efficiency for processing long sequences. VibeVoice employs a next-token diffusion framework, leveraging a Large Language Model to understand textual context and dialogue flow, and a diffusion head to generate high-fidelity acoustic details. The model can synthesize speech up to 90 minutes long with up to 4 distinct speakers, surpassing the typical 1-2 speaker limits of many prior models. \ No newline at end of file diff --git a/demo/vibevoice_asr_gradio_demo.py b/demo/vibevoice_asr_gradio_demo.py new file mode 100644 index 0000000..3e54d12 --- /dev/null +++ b/demo/vibevoice_asr_gradio_demo.py @@ -0,0 +1,1268 @@ +#!/usr/bin/env python +""" +VibeVoice ASR Gradio Demo +""" + +import os +import sys +import torch +import numpy as np +import soundfile as sf +from pathlib import Path +import argparse +import time +import json +import gradio as gr +from typing import List, Dict, Tuple, Optional, Generator +import tempfile +import base64 +import io +import traceback +import threading +from concurrent.futures import ThreadPoolExecutor, as_completed + + # Error correction 'pwd' for Windows +if sys.platform == 'win32': + os.environ['USER'] = os.environ.get('USERNAME', 'user') + os.environ['TORCHINDUCTOR_CACHE_DIR'] = os.path.join(os.environ.get('TEMP', '.'), 'torch_cache') + + +# Import TextIteratorStreamer for streaming generation +from transformers import TextIteratorStreamer, StoppingCriteria, StoppingCriteriaList + +try: + from liger_kernel.transformers import apply_liger_kernel_to_qwen2 + # Only apply RoPE, RMSNorm, SwiGLU patches (these affect the underlying Qwen2 layers) + apply_liger_kernel_to_qwen2( + rope=True, + rms_norm=True, + swiglu=True, + cross_entropy=False, + ) + print("✅ Liger Kernel applied to Qwen2 components (RoPE, RMSNorm, SwiGLU)") +except Exception as e: + print(f"⚠️ Failed to apply Liger Kernel: {e}, you can install it with: pip install liger-kernel") + +# Try to import pydub for MP3 conversion +try: + from pydub import AudioSegment + HAS_PYDUB = True +except ImportError: + HAS_PYDUB = False + print("⚠️ Warning: pydub not available, falling back to WAV format") + +from vibevoice.modular.modeling_vibevoice_asr import VibeVoiceASRForConditionalGeneration +from vibevoice.processor.vibevoice_asr_processor import VibeVoiceASRProcessor +from vibevoice.processor.audio_utils import load_audio_use_ffmpeg, COMMON_AUDIO_EXTS + + +class VibeVoiceASRInference: + """Simple inference wrapper for VibeVoice ASR model.""" + + def __init__(self, model_path: str, device: str = "cuda", dtype: torch.dtype = torch.bfloat16, attn_implementation: str = "flash_attention_2"): + """ + Initialize the ASR inference pipeline. + + Args: + model_path: Path to the pretrained model (HuggingFace format directory or model name) + device: Device to run inference on (cuda, mps, xpu, cpu, auto) + dtype: Data type for model weights + attn_implementation: Attention implementation to use ('flash_attention_2', 'sdpa', 'eager') + """ + print(f"Loading VibeVoice ASR model from {model_path}") + + # Load processor + self.processor = VibeVoiceASRProcessor.from_pretrained(model_path) + + # Load model with device-specific handling + print(f"Using attention implementation: {attn_implementation}") + if device == "mps": + # MPS: load onto CPU first, then move (device_map="mps" is not supported) + self.model = VibeVoiceASRForConditionalGeneration.from_pretrained( + model_path, + dtype=dtype, + device_map=None, + attn_implementation=attn_implementation, + trust_remote_code=True + ) + self.model = self.model.to("mps") + elif device == "auto": + self.model = VibeVoiceASRForConditionalGeneration.from_pretrained( + model_path, + dtype=dtype, + device_map="auto", + attn_implementation=attn_implementation, + trust_remote_code=True + ) + else: + self.model = VibeVoiceASRForConditionalGeneration.from_pretrained( + model_path, + dtype=dtype, + device_map=device if device != "auto" else None, + attn_implementation=attn_implementation, + trust_remote_code=True + ) + self.model = self.model.to(device) + + self.device = device if device != "auto" else next(self.model.parameters()).device + self.model.eval() + + # Print model info + total_params = sum(p.numel() for p in self.model.parameters()) + print(f"✅ Model loaded successfully on {self.device}") + print(f"📊 Total parameters: {total_params:,} ({total_params/1e9:.2f}B)") + + def transcribe( + self, + audio_path: str = None, + audio_array: np.ndarray = None, + sample_rate: int = None, + max_new_tokens: int = 512, + temperature: float = 0.0, + top_p: float = 1.0, + do_sample: bool = False, + num_beams: int = 1, + repetition_penalty: float = 1.0, + context_info: str = None, + streamer: Optional[TextIteratorStreamer] = None, + ) -> dict: + """ + Transcribe audio to text. + + Args: + audio_path: Path to audio file + audio_array: Audio array (if not loading from file) + sample_rate: Sample rate of audio array + max_new_tokens: Maximum tokens to generate + temperature: Temperature for sampling (0 for greedy) + top_p: Top-p for nucleus sampling (1.0 for no filtering) + do_sample: Whether to use sampling + num_beams: Number of beams for beam search (1 for greedy) + repetition_penalty: Repetition penalty (1.0 for no penalty) + context_info: Optional context information (e.g., hotwords, speaker names, topics) to help transcription + streamer: Optional TextIteratorStreamer for streaming output + + Returns: + Dictionary with transcription results + """ + # Process audio + inputs = self.processor( + audio=audio_path, + sampling_rate=sample_rate, + return_tensors="pt", + add_generation_prompt=True, + context_info=context_info + ) + + # Move to device + inputs = {k: v.to(self.device) if isinstance(v, torch.Tensor) else v + for k, v in inputs.items()} + + # Generate + generation_config = { + "max_new_tokens": max_new_tokens, + "temperature": temperature if temperature > 0 else None, + "top_p": top_p if do_sample else None, + "do_sample": do_sample, + "num_beams": num_beams, + "repetition_penalty": repetition_penalty, + "pad_token_id": self.processor.pad_id, + "eos_token_id": self.processor.tokenizer.eos_token_id, + } + + # Add streamer if provided + if streamer is not None: + generation_config["streamer"] = streamer + + # Add stopping criteria for stop button support + generation_config["stopping_criteria"] = StoppingCriteriaList([StopOnFlag()]) + + # Remove None values + generation_config = {k: v for k, v in generation_config.items() if v is not None} + + start_time = time.time() + + # Calculate input token statistics before generation + input_ids = inputs['input_ids'][0] # Shape: [seq_len] + total_input_tokens = input_ids.shape[0] + + # Count padding tokens (tokens equal to pad_id) + pad_id = self.processor.pad_id + padding_mask = (input_ids == pad_id) + num_padding_tokens = padding_mask.sum().item() + + # Count speech tokens (tokens between speech_start_id and speech_end_id) + speech_start_id = self.processor.speech_start_id + speech_end_id = self.processor.speech_end_id + + # Find speech regions + input_ids_list = input_ids.tolist() + num_speech_tokens = 0 + in_speech = False + for token_id in input_ids_list: + if token_id == speech_start_id: + in_speech = True + num_speech_tokens += 1 # Count speech_start token + elif token_id == speech_end_id: + in_speech = False + num_speech_tokens += 1 # Count speech_end token + elif in_speech: + num_speech_tokens += 1 + + # Text tokens = total - speech - padding + num_text_tokens = total_input_tokens - num_speech_tokens - num_padding_tokens + + with torch.no_grad(): + output_ids = self.model.generate( + **inputs, + **generation_config + ) + + generation_time = time.time() - start_time + + # Decode output + generated_ids = output_ids[0, inputs['input_ids'].shape[1]:] + generated_text = self.processor.decode(generated_ids, skip_special_tokens=True) + + # Parse structured output + try: + transcription_segments = self.processor.post_process_transcription(generated_text) + except Exception as e: + print(f"Warning: Failed to parse structured output: {e}") + transcription_segments = [] + + return { + "raw_text": generated_text, + "segments": transcription_segments, + "generation_time": generation_time, + "input_tokens": { + "total": total_input_tokens, + "speech": num_speech_tokens, + "text": num_text_tokens, + "padding": num_padding_tokens, + }, + } + + +def clip_and_encode_audio( + audio_data: np.ndarray, + sr: int, + start_time: float, + end_time: float, + segment_idx: int, + use_mp3: bool = True, + target_sr: int = 16000, # Downsample to 16kHz for smaller size + mp3_bitrate: str = "32k" # Use low bitrate for minimal transfer +) -> Tuple[int, Optional[str], Optional[str]]: + """ + Clip audio segment and encode to base64. + + Args: + audio_data: Full audio array + sr: Sample rate + start_time: Start time in seconds + end_time: End time in seconds + segment_idx: Segment index for identification + use_mp3: Whether to use MP3 format (smaller size) + target_sr: Target sample rate for downsampling (lower = smaller) + mp3_bitrate: MP3 bitrate (lower = smaller, e.g., "24k", "32k", "48k") + + Returns: + Tuple of (segment_idx, base64_string, error_message) + """ + try: + # Convert time to sample indices + start_sample = int(start_time * sr) + end_sample = int(end_time * sr) + + # Ensure indices are within bounds + start_sample = max(0, start_sample) + end_sample = min(len(audio_data), end_sample) + + if start_sample >= end_sample: + return segment_idx, None, f"Invalid time range: [{start_time:.2f}s - {end_time:.2f}s]" + + # Extract segment + segment_data = audio_data[start_sample:end_sample] + + # Downsample if needed (reduces data size significantly) + if sr != target_sr and target_sr < sr: + # Simple downsampling using linear interpolation + duration = len(segment_data) / sr + new_length = int(duration * target_sr) + indices = np.linspace(0, len(segment_data) - 1, new_length) + segment_data = np.interp(indices, np.arange(len(segment_data)), segment_data) + sr = target_sr + + # Convert float32 audio to int16 for encoding + segment_data_int16 = (segment_data * 32768.0).astype(np.int16) + + # Convert to MP3 if pydub is available and use_mp3 is True + if use_mp3 and HAS_PYDUB: + try: + # Write to WAV in memory + wav_buffer = io.BytesIO() + sf.write(wav_buffer, segment_data_int16, sr, format='WAV', subtype='PCM_16') + wav_buffer.seek(0) + + # Convert to MP3 with low bitrate + audio_segment = AudioSegment.from_wav(wav_buffer) + # Convert to mono if stereo (halves the size) + if audio_segment.channels > 1: + audio_segment = audio_segment.set_channels(1) + mp3_buffer = io.BytesIO() + audio_segment.export(mp3_buffer, format='mp3', bitrate=mp3_bitrate) + mp3_buffer.seek(0) + + # Encode to base64 + audio_bytes = mp3_buffer.read() + audio_base64 = base64.b64encode(audio_bytes).decode('utf-8') + audio_src = f"data:audio/mp3;base64,{audio_base64}" + + return segment_idx, audio_src, None + except Exception as e: + # Fall back to WAV on error + print(f"MP3 conversion failed for segment {segment_idx}, using WAV: {e}") + + # Fall back to WAV format (no temp file, use in-memory buffer) + wav_buffer = io.BytesIO() + sf.write(wav_buffer, segment_data_int16, sr, format='WAV', subtype='PCM_16') + wav_buffer.seek(0) + + audio_bytes = wav_buffer.read() + audio_base64 = base64.b64encode(audio_bytes).decode('utf-8') + audio_src = f"data:audio/wav;base64,{audio_base64}" + + return segment_idx, audio_src, None + + except Exception as e: + error_msg = f"Error clipping segment {segment_idx}: {str(e)}" + print(error_msg) + return segment_idx, None, error_msg + + +def extract_audio_segments(audio_path: str, segments: List[Dict]) -> List[Tuple[str, str, Optional[str]]]: + """ + Extract multiple segments from audio file efficiently with parallel processing. + + Args: + audio_path: Path to original audio file + segments: List of segment dictionaries with start_time, end_time, etc. + + Returns: + List of tuples (segment_label, audio_base64_src, error_msg) + """ + try: + # Read audio file once using ffmpeg for better format support + print(f"📂 Loading audio file: {audio_path}") + audio_data, sr = load_audio_use_ffmpeg(audio_path, resample=False) + print(f"✅ Audio loaded: {len(audio_data)} samples, {sr} Hz") + + # Prepare tasks + tasks = [] + use_mp3 = HAS_PYDUB # Use MP3 if available + + for i, seg in enumerate(segments): + start_time = seg.get('start_time') + end_time = seg.get('end_time') + + # Skip if times are not available or invalid + if (not isinstance(start_time, (int, float)) or + not isinstance(end_time, (int, float)) or + start_time >= end_time): + tasks.append((i, None, None, None, None, None)) # Will be filtered later + continue + + tasks.append((audio_data, sr, start_time, end_time, i, use_mp3)) + + # Process in parallel using ThreadPoolExecutor + results = [] + total_segments = len(tasks) + completed_count = 0 + + # Use CPU count for max workers + max_workers = os.cpu_count() or 4 + print(f"🚀 Starting parallel processing with {max_workers} threads...") + + with ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = {} + for task in tasks: + if task[0] is None: # Skip invalid tasks + continue + future = executor.submit(clip_and_encode_audio, *task) + futures[future] = task[4] # segment_idx + + for future in as_completed(futures): + try: + result = future.result() + results.append(result) + completed_count += 1 + # Log progress every 100 segments or at completion + if completed_count % 100 == 0 or completed_count == len(futures): + print(f"Progress: {completed_count}/{len(futures)} segments processed ({completed_count*100//len(futures)}%)") + except Exception as e: + idx = futures[future] + results.append((idx, None, f"Processing error: {str(e)}")) + completed_count += 1 + print(f"Error on segment {idx}: {e}") + + print(f"✅ Completed processing all {len(futures)} valid segments") + + # Sort by segment index to maintain order + results.sort(key=lambda x: x[0]) + + # Build output list with labels + audio_segments = [] + for i, (idx, audio_src, error_msg) in enumerate(results): + seg = segments[idx] if idx < len(segments) else {} + start_time = seg.get('start_time', 'N/A') + end_time = seg.get('end_time', 'N/A') + speaker_id = seg.get('speaker_id', 'N/A') + + segment_label = f"Segment {idx+1}: [{start_time:.2f}s - {end_time:.2f}s] Speaker {speaker_id}" + audio_segments.append((segment_label, audio_src, error_msg)) + + return audio_segments + + except Exception as e: + print(f"Error loading audio file: {e}") + return [] + + +# Global variable to store the ASR model +asr_model = None + +# Global stop flag for generation +stop_generation_flag = False + + +class StopOnFlag(StoppingCriteria): + """Custom stopping criteria that checks a global flag.""" + def __call__(self, input_ids, scores, **kwargs): + global stop_generation_flag + return stop_generation_flag + + +def parse_time_to_seconds(val: Optional[str]) -> Optional[float]: + """Parse seconds or hh:mm:ss to float seconds.""" + if val is None: + return None + val = val.strip() + if not val: + return None + try: + return float(val) + except ValueError: + pass + if ":" in val: + parts = val.split(":") + if not all(p.strip().replace(".", "", 1).isdigit() for p in parts): + return None + parts = [float(p) for p in parts] + if len(parts) == 3: + h, m, s = parts + elif len(parts) == 2: + h = 0 + m, s = parts + else: + return None + return h * 3600 + m * 60 + s + return None + + +def slice_audio_to_temp( + audio_data: np.ndarray, + sample_rate: int, + start_sec: Optional[float], + end_sec: Optional[float] +) -> Tuple[Optional[str], Optional[str]]: + """Slice audio_data to [start_sec, end_sec) and write to a temp WAV file.""" + n_samples = len(audio_data) + full_duration = n_samples / float(sample_rate) + start = 0.0 if start_sec is None else max(0.0, start_sec) + end = full_duration if end_sec is None else min(full_duration, end_sec) + if end <= start: + return None, f"Invalid time range: start={start:.2f}s, end={end:.2f}s" + start_idx = int(start * sample_rate) + end_idx = int(end * sample_rate) + segment = audio_data[start_idx:end_idx] + temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".wav") + temp_file.close() + segment_int16 = (segment * 32768.0).astype(np.int16) + sf.write(temp_file.name, segment_int16, sample_rate, subtype='PCM_16') + return temp_file.name, None + + +def initialize_model(model_path: str, device: str = "cuda", attn_implementation: str = "flash_attention_2"): + """Initialize the ASR model.""" + global asr_model + try: + # MPS and CPU require float32; CUDA/XPU can use bfloat16 + if device in ("mps", "cpu"): + dtype = torch.float32 + else: + dtype = torch.bfloat16 + asr_model = VibeVoiceASRInference( + model_path=model_path, + device=device, + dtype=dtype, + attn_implementation=attn_implementation + ) + return f"✅ Model loaded successfully from {model_path}" + except Exception as e: + import traceback + traceback.print_exc() + return f"❌ Error loading model: {str(e)}" + + +def transcribe_audio( + audio_input, + audio_path_input: str, + start_time_input: str, + end_time_input: str, + max_new_tokens: int, + temperature: float, + top_p: float, + do_sample: bool, + repetition_penalty: float = 1.0, + context_info: str = "" +) -> Generator[Tuple[str, str], None, None]: + """ + Transcribe audio and return results with audio segments (streaming version). + + Args: + audio_input: Audio file path or tuple (sample_rate, audio_data) + max_new_tokens: Maximum tokens to generate + temperature: Temperature for sampling (0 for greedy) + top_p: Top-p for nucleus sampling + do_sample: Whether to use sampling + context_info: Optional context information (e.g., hotwords, speaker names, topics) + + Yields: + Tuple of (raw_text, audio_segments_html) + """ + if asr_model is None: + yield "❌ Please load a model first!", "" + return + + if not audio_path_input and audio_input is None: + yield "❌ Please provide audio input!", "" + return + + try: + print("[INFO] Transcription requested") + start_sec = parse_time_to_seconds(start_time_input) + end_sec = parse_time_to_seconds(end_time_input) + print(f"[INFO] Parsed time range: start={start_sec}, end={end_sec}") + if (start_time_input and start_sec is None) or (end_time_input and end_sec is None): + yield "❌ Invalid time format. Use seconds or hh:mm:ss.", "" + return + + audio_path = None + audio_array = None + sample_rate = None + + if audio_path_input: + candidate = Path(audio_path_input.strip()) + # Security: validate file extension to prevent arbitrary file probing + if candidate.suffix.lower() not in {e.lower() for e in COMMON_AUDIO_EXTS}: + yield "❌ Unsupported audio format.", "" + return + if not candidate.exists(): + yield f"❌ Provided path does not exist: {candidate}", "" + return + audio_path = str(candidate) + print(f"[INFO] Using provided audio path: {audio_path}") + # Get audio file path (Gradio Audio component returns tuple (sample_rate, audio_data) or file path) + elif isinstance(audio_input, str): + audio_path = audio_input + print(f"[INFO] Using uploaded audio path: {audio_path}") + elif isinstance(audio_input, tuple): + # Audio from microphone: (sample_rate, audio_data) + sample_rate, audio_array = audio_input + print(f"[INFO] Received microphone audio with sample_rate={sample_rate}") + elif audio_path is None: + yield "❌ Invalid audio input format!", "" + return + + # If slicing is requested, load and slice audio + if start_sec is not None or end_sec is not None: + print("[INFO] Slicing audio per requested time range") + if audio_array is None or sample_rate is None: + try: + audio_array, sample_rate = load_audio_use_ffmpeg(audio_path, resample=False) + print("[INFO] Loaded audio for slicing via ffmpeg") + except Exception as exc: + yield f"❌ Failed to load audio for slicing: {exc}", "" + return + sliced_path, err = slice_audio_to_temp(audio_array, sample_rate, start_sec, end_sec) + if err: + yield f"❌ {err}", "" + return + audio_path = sliced_path + print(f"[INFO] Sliced audio written to temp file: {audio_path}") + elif audio_array is not None and sample_rate is not None: + # no slicing but microphone input: write to temp file + temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".wav") + audio_path = temp_file.name + temp_file.close() + audio_data_int16 = (audio_array * 32768.0).astype(np.int16) + sf.write(audio_path, audio_data_int16, sample_rate, subtype='PCM_16') + print(f"[INFO] Microphone audio saved to temp file: {audio_path}") + + # Create streamer for real-time output + streamer = TextIteratorStreamer( + asr_model.processor.tokenizer, + skip_prompt=True, + skip_special_tokens=True + ) + + # Store result in a mutable container for the thread + result_container = {"result": None, "error": None} + + def run_transcription(): + try: + result_container["result"] = asr_model.transcribe( + audio_path=audio_path, + max_new_tokens=max_new_tokens, + temperature=temperature, + top_p=top_p, + do_sample=do_sample, + repetition_penalty=repetition_penalty, + context_info=context_info if context_info and context_info.strip() else None, + streamer=streamer + ) + except Exception as e: + result_container["error"] = str(e) + traceback.print_exc() + + # Start transcription in background thread + print("[INFO] Starting model transcription (streaming mode)") + start_time = time.time() + transcription_thread = threading.Thread(target=run_transcription) + transcription_thread.start() + + # Yield streaming output + generated_text = "" + token_count = 0 + for new_text in streamer: + generated_text += new_text + token_count += 1 + elapsed = time.time() - start_time + # Show streaming output with live stats, format for readability + formatted_text = generated_text.replace('},', '},\n') + streaming_output = f"--- 🔴 LIVE Streaming Output (tokens: {token_count}, time: {elapsed:.1f}s) ---\n{formatted_text}" + yield streaming_output, "
⏳ Generating transcription... Audio segments will appear after completion.
" + + # Wait for thread to complete + transcription_thread.join() + + if result_container["error"]: + yield f"❌ Error during transcription: {result_container['error']}", "" + return + + result = result_container["result"] + generation_time = time.time() - start_time + + # Get input token statistics + input_tokens = result.get('input_tokens', {}) + speech_tokens = input_tokens.get('speech', 0) + text_tokens = input_tokens.get('text', 0) + padding_tokens = input_tokens.get('padding', 0) + total_input = input_tokens.get('total', 0) + + # Format final raw output with input/output token stats + raw_output = f"--- ✅ Raw Output ---\n" + raw_output += f"📥 Input: {total_input} tokens (🎤 speech: {speech_tokens}, 📝 text: {text_tokens}, ⬜ pad: {padding_tokens})\n" + raw_output += f"📤 Output: {token_count} tokens | ⏱️ Time: {generation_time:.2f}s\n" + raw_output += f"---\n" + # Format raw text for better readability: add newline after each dict (},) + formatted_raw_text = result['raw_text'].replace('},', '},\n') + raw_output += formatted_raw_text + + # Debug: print raw output to console + print(f"[DEBUG] Raw model output:") + print(f"[DEBUG] {result['raw_text']}") + print(f"[DEBUG] Found {len(result['segments'])} segments") + + # Create audio segments with server-side encoding (low quality for minimal transfer) + # Using: 16kHz mono MP3 @ 32kbps = ~4KB per second of audio + audio_segments_html = "" + segments = result['segments'] + + if segments: + num_segments = len(segments) + print(f"[INFO] Creating per-segment audio clips ({num_segments} segments, 16kHz mono MP3 @ 32kbps)") + + # Extract all audio segments efficiently (load audio only once) + audio_segments = extract_audio_segments(audio_path, segments) + print("[INFO] Completed creating audio clips") + + # Calculate approximate total size + total_duration = sum( + (seg.get('end_time', 0) - seg.get('start_time', 0)) + for seg in segments + if isinstance(seg.get('start_time'), (int, float)) and isinstance(seg.get('end_time'), (int, float)) + ) + approx_size_kb = total_duration * 4 # ~4KB per second at 32kbps + + # Add CSS for theme-aware styling + theme_css = """ + + """ + + audio_segments_html = theme_css + audio_segments_html += f"
" + + # Add format info + format_info = "MP3 32kbps 16kHz mono" if HAS_PYDUB else "WAV 16kHz" + audio_segments_html += f"

🔊 Audio Segments ({num_segments} segments)" + audio_segments_html += f"📦 ~{approx_size_kb:.0f}KB ({format_info})

" + audio_segments_html += "

🎵 Click the play button to listen to each segment directly!

" + + for i, (label, audio_src, error_msg) in enumerate(audio_segments): + seg = segments[i] if i < len(segments) else {} + start_time = seg.get('start_time', 'N/A') + end_time = seg.get('end_time', 'N/A') + speaker_id = seg.get('speaker_id', 'N/A') + content = seg.get('text', '') + + # Format times nicely + start_str = f"{start_time:.2f}" if isinstance(start_time, (int, float)) else str(start_time) + end_str = f"{end_time:.2f}" if isinstance(end_time, (int, float)) else str(end_time) + + audio_segments_html += f""" +
+
+

Segment {i+1}

+
+ Time: [{start_str}s - {end_str}s] | + Speaker: {speaker_id} +
+
+ +
+ {content} +
+ """ + + if audio_src: + # Detect format from data URI + audio_type = 'audio/mp3' if 'audio/mp3' in audio_src else 'audio/wav' + audio_segments_html += f""" + + """ + elif error_msg: + audio_segments_html += f""" +
+ ❌ {error_msg} +
+ """ + else: + audio_segments_html += """ +
+ Audio playback unavailable for this segment +
+ """ + + audio_segments_html += "
" + + audio_segments_html += "
" + else: + audio_segments_html = """ + +
+

❌ No audio segments available.

+

This could happen if the model output doesn't contain valid time stamps.

+
+ """ + + # Final yield with complete results + yield raw_output, audio_segments_html + + except Exception as e: + print(f"Error during transcription: {e}") + print(traceback.format_exc()) + yield f"❌ Error during transcription: {str(e)}", "" + + +def _detect_device_and_attn( + device: str = "auto", + attn_implementation: str = "auto", +): + """ + Auto-detect the best device and attention implementation. + + Args: + device: Explicit device or ``"auto"`` for best available. + attn_implementation: Explicit implementation or ``"auto"``. + + Returns: + (device, attn_implementation) tuple. + """ + # --- resolve device --- + if device == "auto": + if torch.cuda.is_available(): + device = "cuda" + elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): + device = "mps" + else: + device = "cpu" + elif device == "mps" and not ( + hasattr(torch.backends, "mps") and torch.backends.mps.is_available() + ): + print("Warning: MPS requested but not available. Falling back to CPU.") + device = "cpu" + + # --- resolve attention --- + if attn_implementation == "auto": + if device == "cuda": + try: + import flash_attn # noqa: F401 + attn_implementation = "flash_attention_2" + except ImportError: + print("flash_attn not installed, falling back to sdpa") + attn_implementation = "sdpa" + else: + # MPS / XPU / CPU don't support flash_attention_2 + attn_implementation = "sdpa" + + print(f"Using device: {device}, attn_implementation: {attn_implementation}") + return device, attn_implementation + + +def create_gradio_interface( + model_path: str, + default_max_tokens: int = 8192, + device: str = "auto", + attn_implementation: str = "auto", +): + """Create and launch Gradio interface. + + Args: + model_path: Path to the model (HuggingFace format directory or model name) + default_max_tokens: Default value for max_new_tokens slider + device: Device to run inference on ('auto', 'cuda', 'mps', 'xpu', 'cpu') + attn_implementation: Attention implementation to use ('auto', 'flash_attention_2', 'sdpa', 'eager') + """ + + # Initialize model at startup + device, attn_implementation = _detect_device_and_attn(device, attn_implementation) + model_status = initialize_model(model_path, device, attn_implementation) + print(model_status) + + # Exit if model loading failed + if model_status.startswith("❌"): + print("\n" + "="*80) + print("💥 FATAL ERROR: Model loading failed!") + print("="*80) + print("Cannot start demo without a valid model. Please check:") + print(" 1. Model path is correct") + print(" 2. Model files are not corrupted") + print(" 3. You have enough GPU memory") + print(" 4. CUDA is properly installed (if using GPU)") + print("="*80) + sys.exit(1) + + # Custom CSS for Stop button styling + custom_css = """ + #stop-btn { + background: linear-gradient(135deg, #ef4444 0%, #dc2626 100%) !important; + border: none !important; + color: white !important; + } + #stop-btn:hover { + background: linear-gradient(135deg, #dc2626 0%, #b91c1c 100%) !important; + } + """ + + # Gradio 6.0+ moved theme/css to launch() + with gr.Blocks(title="VibeVoice ASR Demo") as demo: + gr.Markdown("# 🎙️ VibeVoice ASR Demo") + gr.Markdown("Upload audio files or record from microphone to get speech-to-text transcription with speaker diarization.") + gr.Markdown(f"**Model loaded from:** `{model_path}`") + + with gr.Row(): + with gr.Column(scale=1): + # Generation parameters + gr.Markdown("## ⚙️ Generation Parameters") + max_tokens_slider = gr.Slider( + minimum=4096, + maximum=65536, + value=default_max_tokens, + step=4096, + label="Max New Tokens" + ) + + # Sampling parameters + gr.Markdown("### 🎲 Sampling") + do_sample_checkbox = gr.Checkbox( + value=False, + label="Enable Sampling", + info="Enable random sampling instead of deterministic decoding" + ) + + with gr.Column(visible=False) as sampling_params: + temperature_slider = gr.Slider( + minimum=0.0, + maximum=2.0, + value=0.0, + step=0.1, + label="Temperature", + info="0 = greedy, higher = more random" + ) + top_p_slider = gr.Slider( + minimum=0.0, + maximum=1.0, + value=1.0, + step=0.05, + label="Top-p (Nucleus Sampling)", + info="1.0 = no filtering" + ) + + # Repetition penalty (works with both greedy and sampling) + repetition_penalty_slider = gr.Slider( + minimum=1.0, + maximum=1.2, + value=1.0, + step=0.01, + label="Repetition Penalty", + info="1.0 = no penalty, higher = less repetition (works with greedy & sampling)" + ) + + # Context information section + gr.Markdown("## 📋 Context Info (Optional)") + context_info_input = gr.Textbox( + label="Context Information", + placeholder="Enter hotwords, speaker names, topics, or other context to help transcription...\nExample:\nJohn Smith\nMachine Learning\nOpenAI", + lines=4, + max_lines=8, + interactive=True, + info="Provide context like proper nouns, technical terms, or speaker names to improve accuracy" + ) + + with gr.Column(scale=2): + # Audio input section + gr.Markdown("## 🎵 Audio Input") + audio_input = gr.Audio( + label="Upload Audio File or Record from Microphone", + sources=["upload", "microphone"], + type="filepath", + interactive=True, + buttons=["download"] + ) + + with gr.Accordion("📂 Advanced: Remote Path & Time Slicing", open=False): + audio_path_input = gr.Textbox( + label="Audio path (optional)", + placeholder="Enter remote audio file path", + lines=1 + ) + with gr.Row(): + start_time_input = gr.Textbox( + label="Start time", + placeholder="e.g., 0 or 00:00:00", + lines=1, + info="Leave empty to start from the beginning" + ) + end_time_input = gr.Textbox( + label="End time", + placeholder="e.g., 30.5 or 00:00:30.5", + lines=1, + info="Leave empty to use full length" + ) + + with gr.Row(): + transcribe_button = gr.Button("🎯 Transcribe", variant="primary", size="lg", scale=3) + stop_button = gr.Button("⏹️ Stop", variant="secondary", size="lg", scale=1, elem_id="stop-btn") + + # Results section + gr.Markdown("## 📝 Results") + + with gr.Tabs(): + with gr.TabItem("Raw Output"): + raw_output = gr.Textbox( + label="Raw Transcription Output", + lines=8, + max_lines=20, + interactive=False + ) + + with gr.TabItem("Audio Segments"): + audio_segments_output = gr.HTML( + label="Play individual segments to verify accuracy" + ) + + # Event handlers + do_sample_checkbox.change( + fn=lambda x: gr.update(visible=x), + inputs=[do_sample_checkbox], + outputs=[sampling_params] + ) + + def reset_stop_flag(): + """Reset stop flag before starting transcription.""" + global stop_generation_flag + stop_generation_flag = False + + def set_stop_flag(): + """Set stop flag to interrupt generation.""" + global stop_generation_flag + stop_generation_flag = True + return "⏹️ Stop requested..." + + transcribe_button.click( + fn=reset_stop_flag, + inputs=[], + outputs=[], + queue=False + ).then( + fn=transcribe_audio, + inputs=[ + audio_input, + audio_path_input, + start_time_input, + end_time_input, + max_tokens_slider, + temperature_slider, + top_p_slider, + do_sample_checkbox, + repetition_penalty_slider, + context_info_input + ], + outputs=[raw_output, audio_segments_output] + ) + + stop_button.click( + fn=set_stop_flag, + inputs=[], + outputs=[raw_output], + queue=False + ) + + # Add examples + gr.Markdown("## 📋 Instructions") + gr.Markdown(f""" + 1. **Upload Audio**: Use the audio component to upload a file or record from microphone + - **Supported formats**: {', '.join(sorted(set([ext.lower() for ext in COMMON_AUDIO_EXTS])))} + - Optionally set **Start/End time** (seconds or hh:mm:ss) to clip before transcription + 2. **Context Info (Optional)**: Provide context to improve transcription accuracy + - Add hotwords, proper nouns, speaker names, or technical terms + - One item per line or comma-separated + - Examples: "John Smith", "OpenAI", "machine learning" + 3. **Adjust Parameters**: Configure generation parameters as needed + 4. **Transcribe**: Click "Transcribe" to get results + 5. **Review Results**: + - **Raw Output**: View the model's original output + - **Audio Segments**: Play individual segments directly to verify accuracy + + **Audio Segments**: Each segment shows the time range, speaker ID, transcribed content, and an embedded audio player for immediate verification. + """) + + return demo, custom_css + + +def main(): + parser = argparse.ArgumentParser(description="VibeVoice ASR Gradio Demo") + parser.add_argument( + "--model_path", + type=str, + default="microsoft/VibeVoice-ASR", + help="Path to the model (HuggingFace format directory or model name)" + ) + parser.add_argument( + "--device", + type=str, + default="auto", + choices=["auto", "cuda", "mps", "xpu", "cpu"], + help="Device to run inference on. 'auto' detects the best available (default: auto)" + ) + parser.add_argument( + "--attn_implementation", + type=str, + default="auto", + choices=["auto", "flash_attention_2", "sdpa", "eager"], + help="Attention implementation to use. 'auto' selects the best for your device (default: auto)" + ) + parser.add_argument( + "--max_new_tokens", + type=int, + default=32768, + help="Default max new tokens for generation (default: 32768)" + ) + parser.add_argument( + "--host", + type=str, + default="0.0.0.0", + help="Host to bind the server to" + ) + parser.add_argument( + "--port", + type=int, + default=7860, + help="Port to bind the server to" + ) + parser.add_argument( + "--share", + action="store_true", + help="Create a public link" + ) + + args = parser.parse_args() + + # Create and launch interface + demo, custom_css = create_gradio_interface( + model_path=args.model_path, + default_max_tokens=args.max_new_tokens, + device=args.device, + attn_implementation=args.attn_implementation, + ) + + print(f"🚀 Starting VibeVoice ASR Demo...") + print(f"📍 Server will be available at: http://{args.host}:{args.port}") + + # Gradio 6.0+ moved theme/css to launch() + launch_kwargs = { + "server_name": args.host, + "server_port": args.port, + "share": args.share, + "show_error": True, + "theme": gr.themes.Soft(), + "css": custom_css, + } + + # Enable queue for concurrent request handling + demo.queue(default_concurrency_limit=3) + demo.launch(**launch_kwargs) + + +if __name__ == "__main__": + main() diff --git a/demo/vibevoice_asr_inference_from_file.py b/demo/vibevoice_asr_inference_from_file.py new file mode 100644 index 0000000..bf4f75d --- /dev/null +++ b/demo/vibevoice_asr_inference_from_file.py @@ -0,0 +1,581 @@ +#!/usr/bin/env python +""" +VibeVoice ASR Batch Inference Demo Script + +This script supports batch inference for ASR model and compares results +between batch processing and single-sample processing. +""" + +import os +import sys +import torch +import numpy as np +from pathlib import Path +import argparse +import time +import json +import re +from typing import List, Dict, Any, Optional +from functools import wraps + +from vibevoice.modular.modeling_vibevoice_asr import VibeVoiceASRForConditionalGeneration +from vibevoice.processor.vibevoice_asr_processor import VibeVoiceASRProcessor +from vibevoice.processor.audio_utils import COMMON_AUDIO_EXTS + + +class VibeVoiceASRBatchInference: + """Batch inference wrapper for VibeVoice ASR model.""" + + def __init__( + self, + model_path: str, + device: str = "cuda", + dtype: torch.dtype = torch.bfloat16, + attn_implementation: str = "sdpa" + ): + """ + Initialize the ASR batch inference pipeline. + + Args: + model_path: Path to the pretrained model + device: Device to run inference on (cuda, mps, xpu, cpu, auto) + dtype: Data type for model weights + attn_implementation: Attention implementation to use ('flash_attention_2', 'sdpa', 'eager') + """ + print(f"Loading VibeVoice ASR model from {model_path}") + + # Load processor + self.processor = VibeVoiceASRProcessor.from_pretrained( + model_path, + language_model_pretrained_name="Qwen/Qwen2.5-7B" + ) + + # Load model with specified attention implementation + print(f"Using attention implementation: {attn_implementation}") + self.model = VibeVoiceASRForConditionalGeneration.from_pretrained( + model_path, + dtype=dtype, + device_map=device if device == "auto" else None, + attn_implementation=attn_implementation, + trust_remote_code=True + ) + + if device != "auto": + self.model = self.model.to(device) + + self.device = device if device != "auto" else next(self.model.parameters()).device + self.dtype = dtype + self.model.eval() + + print(f"Model loaded successfully on {self.device}") + + def _prepare_generation_config( + self, + max_new_tokens: int = 512, + temperature: float = 0.0, + top_p: float = 0.9, + do_sample: bool = True, + num_beams: int = 1, + ) -> dict: + """Prepare generation configuration.""" + config = { + "max_new_tokens": max_new_tokens, + "pad_token_id": self.processor.pad_id, + "eos_token_id": self.processor.tokenizer.eos_token_id, + } + + # Beam search vs sampling + if num_beams > 1: + config["num_beams"] = num_beams + config["do_sample"] = False # Beam search doesn't use sampling + else: + config["do_sample"] = do_sample + # Only set temperature and top_p when sampling is enabled + if do_sample: + config["temperature"] = temperature + config["top_p"] = top_p + + return config + + def transcribe_batch( + self, + audio_inputs: List, + max_new_tokens: int = 512, + temperature: float = 0.0, + top_p: float = 1.0, + do_sample: bool = True, + num_beams: int = 1, + ) -> List[Dict[str, Any]]: + """ + Transcribe multiple audio files/arrays in a single batch. + + Args: + audio_inputs: List of audio file paths or (array, sampling_rate) tuples + max_new_tokens: Maximum tokens to generate + temperature: Temperature for sampling + top_p: Top-p for nucleus sampling + do_sample: Whether to use sampling + + Returns: + List of transcription results + """ + if len(audio_inputs) == 0: + return [] + + batch_size = len(audio_inputs) + print(f"\nProcessing batch of {batch_size} audio(s)...") + + # Process all audio together + inputs = self.processor( + audio=audio_inputs, + sampling_rate=None, + return_tensors="pt", + padding=True, + add_generation_prompt=True + ) + + # Move to device + inputs = {k: v.to(self.device) if isinstance(v, torch.Tensor) else v + for k, v in inputs.items()} + + # Print batch info + print(f" Input IDs shape: {inputs['input_ids'].shape}") + print(f" Speech tensors shape: {inputs['speech_tensors'].shape}") + print(f" Attention mask shape: {inputs['attention_mask'].shape}") + + # Generate + generation_config = self._prepare_generation_config( + max_new_tokens=max_new_tokens, + temperature=temperature, + top_p=top_p, + do_sample=do_sample, + num_beams=num_beams, + ) + + start_time = time.time() + + with torch.no_grad(): + output_ids = self.model.generate( + **inputs, + **generation_config + ) + + generation_time = time.time() - start_time + + # Decode outputs for each sample in the batch + results = [] + input_length = inputs['input_ids'].shape[1] + + for i, audio_input in enumerate(audio_inputs): + # Get generated tokens for this sample (excluding input tokens) + generated_ids = output_ids[i, input_length:] + + # Remove padding tokens from the end + # Find the first eos_token or pad_token + eos_positions = (generated_ids == self.processor.tokenizer.eos_token_id).nonzero(as_tuple=True)[0] + if len(eos_positions) > 0: + generated_ids = generated_ids[:eos_positions[0] + 1] + + generated_text = self.processor.decode(generated_ids, skip_special_tokens=True) + + # Parse structured output + try: + transcription_segments = self.processor.post_process_transcription(generated_text) + except Exception as e: + print(f"Warning: Failed to parse structured output: {e}") + transcription_segments = [] + + # Get file name based on input type + if isinstance(audio_input, str): + file_name = audio_input + elif isinstance(audio_input, dict) and 'id' in audio_input: + file_name = audio_input['id'] + else: + file_name = f"audio_{i}" + + results.append({ + "file": file_name, + "raw_text": generated_text, + "segments": transcription_segments, + "generation_time": generation_time / batch_size, + }) + + print(f" Total generation time: {generation_time:.2f}s") + print(f" Average time per sample: {generation_time/batch_size:.2f}s") + + return results + + def transcribe_with_batching( + self, + audio_inputs: List, + batch_size: int = 4, + max_new_tokens: int = 512, + temperature: float = 0.0, + top_p: float = 1.0, + do_sample: bool = True, + num_beams: int = 1, + ) -> List[Dict[str, Any]]: + """ + Transcribe multiple audio files/arrays with automatic batching. + + Args: + audio_inputs: List of audio file paths or (array, sampling_rate) tuples + batch_size: Number of samples per batch + max_new_tokens: Maximum tokens to generate + temperature: Temperature for sampling + top_p: Top-p for nucleus sampling + do_sample: Whether to use sampling + + Returns: + List of transcription results + """ + all_results = [] + + # Process in batches + for i in range(0, len(audio_inputs), batch_size): + batch_inputs = audio_inputs[i:i + batch_size] + print(f"\n{'='*60}") + print(f"Processing batch {i//batch_size + 1}/{(len(audio_inputs) + batch_size - 1)//batch_size}") + + batch_results = self.transcribe_batch( + batch_inputs, + max_new_tokens=max_new_tokens, + temperature=temperature, + top_p=top_p, + do_sample=do_sample, + num_beams=num_beams, + ) + all_results.extend(batch_results) + + return all_results + + +def print_result(result: Dict[str, Any]): + """Pretty print a single transcription result.""" + print(f"\nFile: {result['file']}") + print(f"Generation Time: {result['generation_time']:.2f}s") + print(f"\n--- Raw Output ---") + print(result['raw_text'][:500] + "..." if len(result['raw_text']) > 500 else result['raw_text']) + + if result['segments']: + print(f"\n--- Structured Output ({len(result['segments'])} segments) ---") + for seg in result['segments'][:50]: # Show first 50 segments + print(f"[{seg.get('start_time', 'N/A')} - {seg.get('end_time', 'N/A')}] " + f"Speaker {seg.get('speaker_id', 'N/A')}: {seg.get('text', '')}...") + if len(result['segments']) > 50: + print(f" ... and {len(result['segments']) - 50} more segments") + + +def load_dataset_and_concatenate( + dataset_name: str, + split: str, + max_duration: float, + num_audios: int, + target_sr: int = 24000 +) -> Optional[List[np.ndarray]]: + """ + Load a HuggingFace dataset and concatenate audio samples into long audio chunks. + (Note, just for demo purpose, not for benchmark evaluation) + + Args: + dataset_name: HuggingFace dataset name (e.g., 'openslr/librispeech_asr') + split: Dataset split to use (e.g., 'test', 'test.other') + max_duration: Maximum duration in seconds for each concatenated audio + num_audios: Number of concatenated audios to create + target_sr: Target sample rate (default: 24000) + + Returns: + List of concatenated audio arrays, or None if loading fails + """ + try: + from datasets import load_dataset + import torchcodec # just for decode audio in datasets + except ImportError: + print("Please install it with: pip install datasets torchcodec") + return None + + print(f"\nLoading dataset: {dataset_name} (split: {split})") + print(f"Will create {num_audios} concatenated audio(s), each up to {max_duration:.1f}s ({max_duration/3600:.2f} hours)") + + try: + # Use streaming to avoid downloading the entire dataset + dataset = load_dataset(dataset_name, split=split, streaming=True) + print(f"Dataset loaded in streaming mode") + + concatenated_audios = [] # List of concatenated audio metadata + + # Create multiple concatenated audios based on num_audios + current_chunks = [] + current_duration = 0.0 + current_samples_used = 0 + sample_idx = 0 + + for sample in dataset: + if len(concatenated_audios) >= num_audios: + break + + if 'audio' not in sample: + continue + + audio_data = sample['audio'] + audio_array = audio_data['array'] + sr = audio_data['sampling_rate'] + + # Resample if needed + if sr != target_sr: + duration = len(audio_array) / sr + new_length = int(duration * target_sr) + audio_array = np.interp( + np.linspace(0, len(audio_array) - 1, new_length), + np.arange(len(audio_array)), + audio_array + ) + + chunk_duration = len(audio_array) / target_sr + + # Check if adding this chunk exceeds max_duration + if current_duration + chunk_duration > max_duration: + remaining_duration = max_duration - current_duration + if remaining_duration > 0.5: # Only add if > 0.5s remaining + samples_to_take = int(remaining_duration * target_sr) + current_chunks.append(audio_array[:samples_to_take]) + current_duration += remaining_duration + current_samples_used += 1 + + # Save current concatenated audio and start a new one + if current_chunks: + concatenated_audios.append({ + 'array': np.concatenate(current_chunks), + 'duration': current_duration, + 'samples_used': current_samples_used, + }) + print(f" Created audio {len(concatenated_audios)}: {current_duration:.1f}s from {current_samples_used} samples") + + # Reset for next concatenated audio + current_chunks = [] + current_duration = 0.0 + current_samples_used = 0 + + if len(concatenated_audios) >= num_audios: + break + + current_chunks.append(audio_array) + current_duration += chunk_duration + current_samples_used += 1 + + sample_idx += 1 + if sample_idx % 100 == 0: + print(f" Processed {sample_idx} samples...") + + # Don't forget the last batch if it has content + if current_chunks and len(concatenated_audios) < num_audios: + concatenated_audios.append({ + 'array': np.concatenate(current_chunks), + 'duration': current_duration, + 'samples_used': current_samples_used, + }) + print(f" Created audio {len(concatenated_audios)}: {current_duration:.1f}s from {current_samples_used} samples") + + if not concatenated_audios: + print("Warning: No audio samples found in dataset") + return None + + # Extract arrays and print summary + result = [a['array'] for a in concatenated_audios] + total_duration = sum(a['duration'] for a in concatenated_audios) + total_samples = sum(a['samples_used'] for a in concatenated_audios) + print(f"\nCreated {len(result)} concatenated audio(s), total {total_duration:.1f}s ({total_duration/60:.1f} min) from {total_samples} samples") + + return result + + except Exception as e: + print(f"Error loading dataset: {e}") + import traceback + traceback.print_exc() + return None + + +def main(): + parser = argparse.ArgumentParser(description="VibeVoice ASR Batch Inference Demo") + parser.add_argument( + "--model_path", + type=str, + default="", + help="Path to the model checkpoint" + ) + parser.add_argument( + "--audio_files", + type=str, + nargs='+', + required=False, + help="Paths to audio files for transcription" + ) + parser.add_argument( + "--audio_dir", + type=str, + required=False, + help="Directory containing audio files for batch transcription" + ) + parser.add_argument( + "--dataset", + type=str, + required=False, + help="HuggingFace dataset name (e.g., 'openslr/librispeech_asr')" + ) + parser.add_argument( + "--split", + type=str, + default="test", + help="Dataset split to use (e.g., 'test', 'test.other', 'test.clean')" + ) + parser.add_argument( + "--max_duration", + type=float, + default=3600.0, + help="Maximum duration in seconds for concatenated dataset audio (default: 3600 = 1 hour)" + ) + parser.add_argument( + "--batch_size", + type=int, + default=2, + help="Batch size for processing multiple files" + ) + parser.add_argument( + "--device", + type=str, + default="cuda" if torch.cuda.is_available() else ("xpu" if torch.backends.xpu.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu") ), + choices=["cuda", "cpu", "mps","xpu", "auto"], + help="Device to run inference on" + ) + parser.add_argument( + "--max_new_tokens", + type=int, + default=32768, + help="Maximum number of tokens to generate" + ) + parser.add_argument( + "--temperature", + type=float, + default=0.0, + help="Temperature for sampling (0 = greedy decoding)" + ) + parser.add_argument( + "--top_p", + type=float, + default=1.0, + help="Top-p for nucleus sampling" + ) + parser.add_argument( + "--num_beams", + type=int, + default=1, + help="Number of beams for beam search. Use 1 for greedy/sampling" + ) + parser.add_argument( + "--attn_implementation", + type=str, + default="auto", + choices=["flash_attention_2", "sdpa", "eager", "auto"], + help="Attention implementation to use. 'auto' will select the best available for your device (flash_attention_2 for CUDA, sdpa for MPS/CPU/XPU)" + ) + + args = parser.parse_args() + + # Auto-detect best attention implementation based on device + if args.attn_implementation == "auto": + if args.device == "cuda" and torch.cuda.is_available(): + try: + import flash_attn + args.attn_implementation = "flash_attention_2" + except ImportError: + print("flash_attn not installed, falling back to sdpa") + args.attn_implementation = "sdpa" + else: + # MPS/XPU/CPU don't support flash_attention_2 + args.attn_implementation = "sdpa" + print(f"Auto-detected attention implementation: {args.attn_implementation}") + + # Collect audio files + audio_files = [] + concatenated_audio = None # For storing concatenated dataset audio + + if args.audio_files: + audio_files.extend(args.audio_files) + + if args.audio_dir: + supported = set(e.lower() for e in COMMON_AUDIO_EXTS) + for f in os.listdir(args.audio_dir): + if os.path.splitext(f)[1].lower() in supported: + audio_files.append(os.path.join(args.audio_dir, f)) + + if args.dataset: + concatenated_audio = load_dataset_and_concatenate( + dataset_name=args.dataset, + split=args.split, + max_duration=args.max_duration, + num_audios=args.batch_size, + ) + if concatenated_audio is None: + return + + if len(audio_files) == 0 and concatenated_audio is None: + print("No audio files provided. Please specify --audio_files, --audio_dir, or --dataset.") + return + + if audio_files: + print(f"\nAudio files to process ({len(audio_files)}):") + for f in audio_files: + print(f" - {f}") + + if concatenated_audio: + print(f"\nConcatenated dataset audios: {len(concatenated_audio)} audio(s)") + + # Initialize model + # Handle MPS device and dtype + if args.device == "mps": + model_dtype = torch.float32 # MPS works better with float32 + elif args.device == "xpu": + model_dtype = torch.float32 + elif args.device == "cpu": + model_dtype = torch.float32 + else: + model_dtype = torch.bfloat16 + + asr = VibeVoiceASRBatchInference( + model_path=args.model_path, + device=args.device, + dtype=model_dtype, + attn_implementation=args.attn_implementation + ) + + # If temperature is 0, use greedy decoding (no sampling) + do_sample = args.temperature > 0 + + # Combine all audio inputs + all_audio_inputs = audio_files + (concatenated_audio or []) + + print("\n" + "="*80) + print(f"Processing {len(all_audio_inputs)} audio(s)") + print("="*80) + + all_results = asr.transcribe_with_batching( + all_audio_inputs, + batch_size=args.batch_size, + max_new_tokens=args.max_new_tokens, + temperature=args.temperature, + top_p=args.top_p, + do_sample=do_sample, + num_beams=args.num_beams, + ) + + # Print results + print("\n" + "="*80) + print("Results") + print("="*80) + for result in all_results: + print("\n" + "-"*60) + print_result(result) + + +if __name__ == "__main__": + main() diff --git a/demo/vibevoice_realtime_colab.ipynb b/demo/vibevoice_realtime_colab.ipynb new file mode 100644 index 0000000..dc6d762 --- /dev/null +++ b/demo/vibevoice_realtime_colab.ipynb @@ -0,0 +1,216 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "d1785adb", + "metadata": { + "colab_type": "text", + "id": "view-in-github" + }, + "source": [ + "\"Open" + ] + }, + { + "cell_type": "markdown", + "id": "WvIaUJD2y0yU", + "metadata": { + "id": "WvIaUJD2y0yU" + }, + "source": [ + "# VibeVoice-Realtime Colab — T4 Quickstart\n", + "\n" + ] + }, + { + "cell_type": "markdown", + "id": "e8fTKYGx7DZk", + "metadata": { + "id": "e8fTKYGx7DZk" + }, + "source": [ + "## Step 1: Setup Environment" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4wxJ6QHM-ZOb", + "metadata": { + "id": "4wxJ6QHM-ZOb" + }, + "outputs": [], + "source": [ + "# Check for T4 GPU\n", + "import torch\n", + "if torch.cuda.is_available() and \"T4\" in torch.cuda.get_device_name(0):\n", + " print(\"✅ T4 GPU detected\")\n", + "else:\n", + " print(\"\"\"\n", + " ⚠️ WARNING: T4 GPU not detected\n", + "\n", + " The recommended runtime for this Colab notebook is \"T4 GPU\".\n", + "\n", + " To change the runtime type:\n", + "\n", + " 1. Click on \"Runtime\" in the top navigation menu\n", + " 2. Click on \"Change runtime type\"\n", + " 3. Select \"T4 GPU\"\n", + " 4. Click \"OK\" if a \"Disconnect and delete runtime\" window appears\n", + " 5. Click on \"Save\"\n", + "\n", + " \"\"\")\n", + "\n", + "# Clone the VibeVoice repository\n", + "![ -d /content/VibeVoice ] || git clone --quiet --branch main --depth 1 https://github.com/microsoft/VibeVoice.git /content/VibeVoice\n", + "print(\"✅ Cloned VibeVoice repository\")\n", + "\n", + "# Install project dependencies\n", + "!uv pip --quiet install --system -e /content/VibeVoice[streamingtts]\n", + "!wget -q https://github.com/cloudflare/cloudflared/releases/latest/download/cloudflared-linux-amd64 -O cloudflared && chmod +x cloudflared\n", + "print(\"✅ Installed dependencies\")\n", + "\n", + "# Download model\n", + "from huggingface_hub import snapshot_download\n", + "snapshot_download(\"microsoft/VibeVoice-Realtime-0.5B\", local_dir=\"/content/models/VibeVoice-Realtime-0.5B\")\n", + "print(\"✅ Downloaded model: microsoft/VibeVoice-Realtime-0.5B\")\n" + ] + }, + { + "cell_type": "markdown", + "id": "88c727ab", + "metadata": {}, + "source": [ + "[Optional] If the download exceeds 1 minute, it is probably stuck. You can: (1) interrupt the execution, (2) log in to Hugging Face, and (3) try download again." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "dec6b870", + "metadata": {}, + "outputs": [], + "source": [ + "from huggingface_hub import login, snapshot_download\n", + "login()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c579654b", + "metadata": {}, + "outputs": [], + "source": [ + "snapshot_download(\"microsoft/VibeVoice-Realtime-0.5B\", local_dir=\"/content/models/VibeVoice-Realtime-0.5B\")\n", + "print(\"✅ Downloaded model: microsoft/VibeVoice-Realtime-0.5B\")" + ] + }, + { + "cell_type": "markdown", + "id": "dfe30d6f", + "metadata": {}, + "source": [ + "[Optional] More experimental voices" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bb33c9ce", + "metadata": {}, + "outputs": [], + "source": [ + "!bash /content/VibeVoice/demo/download_experimental_voices.sh" + ] + }, + { + "cell_type": "markdown", + "id": "pgKlV7153Ifi", + "metadata": { + "id": "pgKlV7153Ifi" + }, + "source": [ + "## Step 2: Launch VibeVoice-Realtime Demo" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "Yc1N9EHswFxA", + "metadata": { + "id": "Yc1N9EHswFxA" + }, + "outputs": [], + "source": [ + "import subprocess, re, time, threading\n", + "\n", + "srv = subprocess.Popen(\n", + " \"python /content/VibeVoice/demo/vibevoice_realtime_demo.py --model_path /content/models/VibeVoice-Realtime-0.5B --port 8000\",\n", + " shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1, universal_newlines=True,\n", + ")\n", + "cf = subprocess.Popen(\n", + " \"./cloudflared tunnel --url http://localhost:8000 --no-autoupdate\",\n", + " shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1, universal_newlines=True,\n", + ")\n", + "\n", + "public_url = None\n", + "server_ready = False\n", + "url_pattern = re.compile(r\"(https://[a-z0-9-]+\\.trycloudflare\\.com)\")\n", + "\n", + "def read_srv():\n", + " global server_ready\n", + " for ln in srv.stdout:\n", + " print(ln.strip())\n", + " if \"Uvicorn running on\" in ln:\n", + " server_ready = True\n", + "\n", + "def read_cf():\n", + " global public_url\n", + " for ln in cf.stdout:\n", + " m = url_pattern.search(ln)\n", + " if m:\n", + " public_url = m.group(1)\n", + " break\n", + "\n", + "threading.Thread(target=read_srv, daemon=True).start()\n", + "threading.Thread(target=read_cf, daemon=True).start()\n", + "\n", + "\n", + "while True:\n", + " if server_ready and public_url:\n", + " print(f\"✅ Public URL: {public_url}\\n\");\n", + " public_url = None\n", + " time.sleep(0.25)" + ] + } + ], + "metadata": { + "accelerator": "GPU", + "colab": { + "gpuType": "T4", + "include_colab_link": true, + "machine_shape": "hm", + "name": "VibeVoice_Colab.ipynb", + "provenance": [] + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/demo/vibevoice_realtime_demo.py b/demo/vibevoice_realtime_demo.py new file mode 100644 index 0000000..a2beb24 --- /dev/null +++ b/demo/vibevoice_realtime_demo.py @@ -0,0 +1,17 @@ +import argparse, os, uvicorn + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--port", type=int, default=3000) + p.add_argument("--model_path", type=str, default="microsoft/VibeVoice-Realtime-0.5B") + p.add_argument("--device", type=str, default="cuda", choices=["cpu", "cuda", "mpx", "mps"]) + p.add_argument("--reload", action="store_true", help="Reload the model or not") + args = p.parse_args() + + os.environ["MODEL_PATH"] = args.model_path + os.environ["MODEL_DEVICE"] = args.device + + uvicorn.run("web.app:app", host="0.0.0.0", port=args.port, reload=args.reload) + +if __name__ == "__main__": + main() diff --git a/demo/voices/streaming_model/de-Spk0_man.pt b/demo/voices/streaming_model/de-Spk0_man.pt new file mode 100755 index 0000000..7407314 Binary files /dev/null and b/demo/voices/streaming_model/de-Spk0_man.pt differ diff --git a/demo/voices/streaming_model/de-Spk1_woman.pt b/demo/voices/streaming_model/de-Spk1_woman.pt new file mode 100755 index 0000000..c67e669 Binary files /dev/null and b/demo/voices/streaming_model/de-Spk1_woman.pt differ diff --git a/demo/voices/streaming_model/en-Carter_man.pt b/demo/voices/streaming_model/en-Carter_man.pt new file mode 100644 index 0000000..1d795ef Binary files /dev/null and b/demo/voices/streaming_model/en-Carter_man.pt differ diff --git a/demo/voices/streaming_model/en-Davis_man.pt b/demo/voices/streaming_model/en-Davis_man.pt new file mode 100644 index 0000000..a259516 Binary files /dev/null and b/demo/voices/streaming_model/en-Davis_man.pt differ diff --git a/demo/voices/streaming_model/en-Emma_woman.pt b/demo/voices/streaming_model/en-Emma_woman.pt new file mode 100644 index 0000000..9876ca3 Binary files /dev/null and b/demo/voices/streaming_model/en-Emma_woman.pt differ diff --git a/demo/voices/streaming_model/en-Frank_man.pt b/demo/voices/streaming_model/en-Frank_man.pt new file mode 100644 index 0000000..2c37e82 Binary files /dev/null and b/demo/voices/streaming_model/en-Frank_man.pt differ diff --git a/demo/voices/streaming_model/en-Grace_woman.pt b/demo/voices/streaming_model/en-Grace_woman.pt new file mode 100644 index 0000000..366bda8 Binary files /dev/null and b/demo/voices/streaming_model/en-Grace_woman.pt differ diff --git a/demo/voices/streaming_model/en-Mike_man.pt b/demo/voices/streaming_model/en-Mike_man.pt new file mode 100644 index 0000000..f02ef03 Binary files /dev/null and b/demo/voices/streaming_model/en-Mike_man.pt differ diff --git a/demo/voices/streaming_model/fr-Spk0_man.pt b/demo/voices/streaming_model/fr-Spk0_man.pt new file mode 100755 index 0000000..dcc72ff Binary files /dev/null and b/demo/voices/streaming_model/fr-Spk0_man.pt differ diff --git a/demo/voices/streaming_model/fr-Spk1_woman.pt b/demo/voices/streaming_model/fr-Spk1_woman.pt new file mode 100755 index 0000000..48756ff Binary files /dev/null and b/demo/voices/streaming_model/fr-Spk1_woman.pt differ diff --git a/demo/voices/streaming_model/in-Samuel_man.pt b/demo/voices/streaming_model/in-Samuel_man.pt new file mode 100644 index 0000000..9338cb6 Binary files /dev/null and b/demo/voices/streaming_model/in-Samuel_man.pt differ diff --git a/demo/voices/streaming_model/it-Spk0_woman.pt b/demo/voices/streaming_model/it-Spk0_woman.pt new file mode 100755 index 0000000..c57fc61 Binary files /dev/null and b/demo/voices/streaming_model/it-Spk0_woman.pt differ diff --git a/demo/voices/streaming_model/it-Spk1_man.pt b/demo/voices/streaming_model/it-Spk1_man.pt new file mode 100755 index 0000000..68b658e Binary files /dev/null and b/demo/voices/streaming_model/it-Spk1_man.pt differ diff --git a/demo/voices/streaming_model/jp-Spk0_man.pt b/demo/voices/streaming_model/jp-Spk0_man.pt new file mode 100755 index 0000000..bf710af Binary files /dev/null and b/demo/voices/streaming_model/jp-Spk0_man.pt differ diff --git a/demo/voices/streaming_model/jp-Spk1_woman.pt b/demo/voices/streaming_model/jp-Spk1_woman.pt new file mode 100755 index 0000000..4ea2def Binary files /dev/null and b/demo/voices/streaming_model/jp-Spk1_woman.pt differ diff --git a/demo/voices/streaming_model/kr-Spk0_woman.pt b/demo/voices/streaming_model/kr-Spk0_woman.pt new file mode 100755 index 0000000..4264666 Binary files /dev/null and b/demo/voices/streaming_model/kr-Spk0_woman.pt differ diff --git a/demo/voices/streaming_model/kr-Spk1_man.pt b/demo/voices/streaming_model/kr-Spk1_man.pt new file mode 100755 index 0000000..c0eaafa Binary files /dev/null and b/demo/voices/streaming_model/kr-Spk1_man.pt differ diff --git a/demo/voices/streaming_model/nl-Spk0_man.pt b/demo/voices/streaming_model/nl-Spk0_man.pt new file mode 100755 index 0000000..edeb893 Binary files /dev/null and b/demo/voices/streaming_model/nl-Spk0_man.pt differ diff --git a/demo/voices/streaming_model/nl-Spk1_woman.pt b/demo/voices/streaming_model/nl-Spk1_woman.pt new file mode 100755 index 0000000..5014702 Binary files /dev/null and b/demo/voices/streaming_model/nl-Spk1_woman.pt differ diff --git a/demo/voices/streaming_model/pl-Spk0_man.pt b/demo/voices/streaming_model/pl-Spk0_man.pt new file mode 100755 index 0000000..f81b24d Binary files /dev/null and b/demo/voices/streaming_model/pl-Spk0_man.pt differ diff --git a/demo/voices/streaming_model/pl-Spk1_woman.pt b/demo/voices/streaming_model/pl-Spk1_woman.pt new file mode 100755 index 0000000..a55ed5a Binary files /dev/null and b/demo/voices/streaming_model/pl-Spk1_woman.pt differ diff --git a/demo/voices/streaming_model/pt-Spk0_woman.pt b/demo/voices/streaming_model/pt-Spk0_woman.pt new file mode 100755 index 0000000..0316c77 Binary files /dev/null and b/demo/voices/streaming_model/pt-Spk0_woman.pt differ diff --git a/demo/voices/streaming_model/pt-Spk1_man.pt b/demo/voices/streaming_model/pt-Spk1_man.pt new file mode 100755 index 0000000..db7ca1b Binary files /dev/null and b/demo/voices/streaming_model/pt-Spk1_man.pt differ diff --git a/demo/voices/streaming_model/sp-Spk0_woman.pt b/demo/voices/streaming_model/sp-Spk0_woman.pt new file mode 100755 index 0000000..4b84fa2 Binary files /dev/null and b/demo/voices/streaming_model/sp-Spk0_woman.pt differ diff --git a/demo/voices/streaming_model/sp-Spk1_man.pt b/demo/voices/streaming_model/sp-Spk1_man.pt new file mode 100755 index 0000000..ec5698a Binary files /dev/null and b/demo/voices/streaming_model/sp-Spk1_man.pt differ diff --git a/demo/web/app.py b/demo/web/app.py new file mode 100644 index 0000000..4985c7d --- /dev/null +++ b/demo/web/app.py @@ -0,0 +1,518 @@ +import datetime +import builtins +import asyncio +import json +import os +import threading +import traceback +from pathlib import Path +from queue import Empty, Queue +from typing import Any, Callable, Dict, Iterator, Optional, Tuple, cast + +import numpy as np +import torch +from transformers.cache_utils import DynamicCache +from transformers.modeling_outputs import BaseModelOutputWithPast +from fastapi import FastAPI, WebSocket +from fastapi.responses import FileResponse +from fastapi.staticfiles import StaticFiles +from starlette.websockets import WebSocketDisconnect, WebSocketState + +from vibevoice.modular.modeling_vibevoice_streaming_inference import ( + VibeVoiceStreamingForConditionalGenerationInference, +) +from vibevoice.processor.vibevoice_streaming_processor import ( + VibeVoiceStreamingProcessor, +) +from vibevoice.modular.streamer import AudioStreamer + +import copy + +BASE = Path(__file__).parent +SAMPLE_RATE = 24_000 + + +def get_timestamp(): + timestamp = datetime.datetime.utcnow().replace( + tzinfo=datetime.timezone.utc + ).astimezone( + datetime.timezone(datetime.timedelta(hours=8)) + ).strftime("%Y-%m-%d %H:%M:%S.%f")[:-3] + return timestamp + +class StreamingTTSService: + def __init__( + self, + model_path: str, + device: str = "cuda", + inference_steps: int = 5, + ) -> None: + # Keep model_path as string for HuggingFace repo IDs (Path() converts / to \ on Windows) + self.model_path = model_path + self.inference_steps = inference_steps + self.sample_rate = SAMPLE_RATE + + self.processor: Optional[VibeVoiceStreamingProcessor] = None + self.model: Optional[VibeVoiceStreamingForConditionalGenerationInference] = None + self.voice_presets: Dict[str, Path] = {} + self.default_voice_key: Optional[str] = None + self._voice_cache: Dict[str, Tuple[object, Path, str]] = {} + + if device == "mpx": + print("Note: device 'mpx' detected, treating it as 'mps'.") + device = "mps" + if device == "mps" and not torch.backends.mps.is_available(): + print("Warning: MPS not available. Falling back to CPU.") + device = "cpu" + self.device = device + self._torch_device = torch.device(device) + + def load(self) -> None: + print(f"[startup] Loading processor from {self.model_path}") + self.processor = VibeVoiceStreamingProcessor.from_pretrained(self.model_path) + + + # Decide dtype & attention + if self.device == "mps": + load_dtype = torch.float32 + device_map = None + attn_impl_primary = "sdpa" + elif self.device == "cuda": + load_dtype = torch.bfloat16 + device_map = 'cuda' + attn_impl_primary = "flash_attention_2" + else: + load_dtype = torch.float32 + device_map = 'cpu' + attn_impl_primary = "sdpa" + print(f"Using device: {device_map}, torch_dtype: {load_dtype}, attn_implementation: {attn_impl_primary}") + # Load model + try: + self.model = VibeVoiceStreamingForConditionalGenerationInference.from_pretrained( + self.model_path, + torch_dtype=load_dtype, + device_map=device_map, + attn_implementation=attn_impl_primary, + ) + + if self.device == "mps": + self.model.to("mps") + except Exception as e: + if attn_impl_primary == 'flash_attention_2': + print("Error loading the model. Trying to use SDPA. However, note that only flash_attention_2 has been fully tested, and using SDPA may result in lower audio quality.") + + self.model = VibeVoiceStreamingForConditionalGenerationInference.from_pretrained( + self.model_path, + torch_dtype=load_dtype, + device_map=self.device, + attn_implementation='sdpa', + ) + print("Load model with SDPA successfully ") + else: + raise e + + self.model.eval() + + self.model.model.noise_scheduler = self.model.model.noise_scheduler.from_config( + self.model.model.noise_scheduler.config, + algorithm_type="sde-dpmsolver++", + beta_schedule="squaredcos_cap_v2", + ) + self.model.set_ddpm_inference_steps(num_steps=self.inference_steps) + + self.voice_presets = self._load_voice_presets() + preset_name = os.environ.get("VOICE_PRESET") + self.default_voice_key = self._determine_voice_key(preset_name) + self._ensure_voice_cached(self.default_voice_key) + + def _load_voice_presets(self) -> Dict[str, Path]: + voices_dir = BASE.parent / "voices" / "streaming_model" + if not voices_dir.exists(): + raise RuntimeError(f"Voices directory not found: {voices_dir}") + + presets: Dict[str, Path] = {} + for pt_path in voices_dir.rglob("*.pt"): + presets[pt_path.stem] = pt_path + + if not presets: + raise RuntimeError(f"No voice preset (.pt) files found in {voices_dir}") + + print(f"[startup] Found {len(presets)} voice presets") + return dict(sorted(presets.items())) + + def _determine_voice_key(self, name: Optional[str]) -> str: + if name and name in self.voice_presets: + return name + + default_key = "en-Carter_man" + if default_key in self.voice_presets: + return default_key + + first_key = next(iter(self.voice_presets)) + print(f"[startup] Using fallback voice preset: {first_key}") + return first_key + + def _ensure_voice_cached(self, key: str) -> Tuple[object, Path, str]: + if key not in self.voice_presets: + raise RuntimeError(f"Voice preset {key!r} not found") + + if key not in self._voice_cache: + preset_path = self.voice_presets[key] + print(f"[startup] Loading voice preset {key} from {preset_path}") + print(f"[startup] Loading prefilled prompt from {preset_path}") + with torch.serialization.safe_globals([BaseModelOutputWithPast, DynamicCache]): + prefilled_outputs = torch.load( + preset_path, + map_location=self._torch_device, + weights_only=True, + ) + self._voice_cache[key] = prefilled_outputs + + return self._voice_cache[key] + + def _get_voice_resources(self, requested_key: Optional[str]) -> Tuple[str, object, Path, str]: + key = requested_key if requested_key and requested_key in self.voice_presets else self.default_voice_key + if key is None: + key = next(iter(self.voice_presets)) + self.default_voice_key = key + + prefilled_outputs = self._ensure_voice_cached(key) + return key, prefilled_outputs + + def _prepare_inputs(self, text: str, prefilled_outputs: object): + if not self.processor or not self.model: + raise RuntimeError("StreamingTTSService not initialized") + + processor_kwargs = { + "text": text.strip(), + "cached_prompt": prefilled_outputs, + "padding": True, + "return_tensors": "pt", + "return_attention_mask": True, + } + + processed = self.processor.process_input_with_cached_prompt(**processor_kwargs) + + prepared = { + key: value.to(self._torch_device) if hasattr(value, "to") else value + for key, value in processed.items() + } + return prepared + + def _run_generation( + self, + inputs, + audio_streamer: AudioStreamer, + errors, + cfg_scale: float, + do_sample: bool, + temperature: float, + top_p: float, + refresh_negative: bool, + prefilled_outputs, + stop_event: threading.Event, + ) -> None: + try: + self.model.generate( + **inputs, + max_new_tokens=None, + cfg_scale=cfg_scale, + tokenizer=self.processor.tokenizer, + generation_config={ + "do_sample": do_sample, + "temperature": temperature if do_sample else 1.0, + "top_p": top_p if do_sample else 1.0, + }, + audio_streamer=audio_streamer, + stop_check_fn=stop_event.is_set, + verbose=False, + refresh_negative=refresh_negative, + all_prefilled_outputs=copy.deepcopy(prefilled_outputs), + ) + except Exception as exc: # pragma: no cover - diagnostic logging + errors.append(exc) + traceback.print_exc() + audio_streamer.end() + + def stream( + self, + text: str, + cfg_scale: float = 1.5, + do_sample: bool = False, + temperature: float = 0.9, + top_p: float = 0.9, + refresh_negative: bool = True, + inference_steps: Optional[int] = None, + voice_key: Optional[str] = None, + log_callback: Optional[Callable[[str, Dict[str, Any]], None]] = None, + stop_event: Optional[threading.Event] = None, + ) -> Iterator[np.ndarray]: + if not text.strip(): + return + text = text.replace("’", "'") + selected_voice, prefilled_outputs = self._get_voice_resources(voice_key) + + def emit(event: str, **payload: Any) -> None: + if log_callback: + try: + log_callback(event, **payload) + except Exception as exc: + print(f"[log_callback] Error while emitting {event}: {exc}") + + steps_to_use = self.inference_steps + if inference_steps is not None: + try: + parsed_steps = int(inference_steps) + if parsed_steps > 0: + steps_to_use = parsed_steps + except (TypeError, ValueError): + pass + if self.model: + self.model.set_ddpm_inference_steps(num_steps=steps_to_use) + self.inference_steps = steps_to_use + + inputs = self._prepare_inputs(text, prefilled_outputs) + audio_streamer = AudioStreamer(batch_size=1, stop_signal=None, timeout=None) + errors: list = [] + stop_signal = stop_event or threading.Event() + + thread = threading.Thread( + target=self._run_generation, + kwargs={ + "inputs": inputs, + "audio_streamer": audio_streamer, + "errors": errors, + "cfg_scale": cfg_scale, + "do_sample": do_sample, + "temperature": temperature, + "top_p": top_p, + "refresh_negative": refresh_negative, + "prefilled_outputs": prefilled_outputs, + "stop_event": stop_signal, + }, + daemon=True, + ) + thread.start() + + generated_samples = 0 + + try: + stream = audio_streamer.get_stream(0) + for audio_chunk in stream: + if torch.is_tensor(audio_chunk): + audio_chunk = audio_chunk.detach().cpu().to(torch.float32).numpy() + else: + audio_chunk = np.asarray(audio_chunk, dtype=np.float32) + + if audio_chunk.ndim > 1: + audio_chunk = audio_chunk.reshape(-1) + + peak = np.max(np.abs(audio_chunk)) if audio_chunk.size else 0.0 + if peak > 1.0: + audio_chunk = audio_chunk / peak + + generated_samples += int(audio_chunk.size) + emit( + "model_progress", + generated_sec=generated_samples / self.sample_rate, + chunk_sec=audio_chunk.size / self.sample_rate, + ) + + chunk_to_yield = audio_chunk.astype(np.float32, copy=False) + + yield chunk_to_yield + finally: + stop_signal.set() + audio_streamer.end() + thread.join() + if errors: + emit("generation_error", message=str(errors[0])) + raise errors[0] + + def chunk_to_pcm16(self, chunk: np.ndarray) -> bytes: + chunk = np.clip(chunk, -1.0, 1.0) + pcm = (chunk * 32767.0).astype(np.int16) + return pcm.tobytes() + + +app = FastAPI() + + +@app.on_event("startup") +async def _startup() -> None: + model_path = os.environ.get("MODEL_PATH") + if not model_path: + raise RuntimeError("MODEL_PATH not set in environment") + + device = os.environ.get("MODEL_DEVICE", "cuda") + + service = StreamingTTSService( + model_path=model_path, + device=device + ) + service.load() + + app.state.tts_service = service + app.state.model_path = model_path + app.state.device = device + app.state.websocket_lock = asyncio.Lock() + print("[startup] Model ready.") + + +def streaming_tts(text: str, **kwargs) -> Iterator[np.ndarray]: + service: StreamingTTSService = app.state.tts_service + yield from service.stream(text, **kwargs) + +@app.websocket("/stream") +async def websocket_stream(ws: WebSocket) -> None: + await ws.accept() + text = ws.query_params.get("text", "") + print(f"Client connected, text={text!r}") + cfg_param = ws.query_params.get("cfg") + steps_param = ws.query_params.get("steps") + voice_param = ws.query_params.get("voice") + + try: + cfg_scale = float(cfg_param) if cfg_param is not None else 1.5 + except ValueError: + cfg_scale = 1.5 + if cfg_scale <= 0: + cfg_scale = 1.5 + try: + inference_steps = int(steps_param) if steps_param is not None else None + if inference_steps is not None and inference_steps <= 0: + inference_steps = None + except ValueError: + inference_steps = None + + service: StreamingTTSService = app.state.tts_service + lock: asyncio.Lock = app.state.websocket_lock + + if lock.locked(): + busy_message = { + "type": "log", + "event": "backend_busy", + "data": {"message": "Please wait for the other requests to complete."}, + "timestamp": get_timestamp(), + } + print("Please wait for the other requests to complete.") + try: + await ws.send_text(json.dumps(busy_message)) + except Exception: + pass + await ws.close(code=1013, reason="Service busy") + return + + acquired = False + try: + await lock.acquire() + acquired = True + + log_queue: "Queue[Dict[str, Any]]" = Queue() + + def enqueue_log(event: str, **data: Any) -> None: + log_queue.put({"event": event, "data": data}) + + async def flush_logs() -> None: + while True: + try: + entry = log_queue.get_nowait() + except Empty: + break + message = { + "type": "log", + "event": entry.get("event"), + "data": entry.get("data", {}), + "timestamp": get_timestamp(), + } + try: + await ws.send_text(json.dumps(message)) + except Exception: + break + + enqueue_log( + "backend_request_received", + text_length=len(text or ""), + cfg_scale=cfg_scale, + inference_steps=inference_steps, + voice=voice_param, + ) + + stop_signal = threading.Event() + + iterator = streaming_tts( + text, + cfg_scale=cfg_scale, + inference_steps=inference_steps, + voice_key=voice_param, + log_callback=enqueue_log, + stop_event=stop_signal, + ) + sentinel = object() + first_ws_send_logged = False + + await flush_logs() + + try: + while ws.client_state == WebSocketState.CONNECTED: + await flush_logs() + chunk = await asyncio.to_thread(next, iterator, sentinel) + if chunk is sentinel: + break + chunk = cast(np.ndarray, chunk) + payload = service.chunk_to_pcm16(chunk) + await ws.send_bytes(payload) + if not first_ws_send_logged: + first_ws_send_logged = True + enqueue_log("backend_first_chunk_sent") + await flush_logs() + except WebSocketDisconnect: + print("Client disconnected (WebSocketDisconnect)") + enqueue_log("client_disconnected") + stop_signal.set() + except Exception as e: + print(f"Error in websocket stream: {e}") + traceback.print_exc() + enqueue_log("backend_error", message=str(e)) + stop_signal.set() + finally: + stop_signal.set() + enqueue_log("backend_stream_complete") + await flush_logs() + try: + iterator_close = getattr(iterator, "close", None) + if callable(iterator_close): + iterator_close() + except Exception: + pass + # clear the log queue + while not log_queue.empty(): + try: + log_queue.get_nowait() + except Empty: + break + try: + if ws.client_state == WebSocketState.CONNECTED: + await ws.close() + except Exception as e: + print(f"Error closing websocket: {e}") + print("WS handler exit") + finally: + if acquired: + lock.release() + + +@app.get("/") +def index(): + return FileResponse(BASE / "index.html") + + +@app.get("/config") +def get_config(): + service: StreamingTTSService = app.state.tts_service + voices = sorted(service.voice_presets.keys()) + return { + "voices": voices, + "default_voice": service.default_voice_key, + } + diff --git a/demo/web/index.html b/demo/web/index.html new file mode 100644 index 0000000..daf9df8 --- /dev/null +++ b/demo/web/index.html @@ -0,0 +1,1017 @@ + + + +VibeVoice-Realtime TTS Demo + + +
+

VibeVoice-Realtime TTS Demo

+ +
+ + +
+
+ Streaming Input Text +
+
This area will display the streaming input text in real time.
+
+
+ This demo requires the full text to be provided upfront. The model then receives the text via streaming input during synthesis.
+ For non-punctuation special characters, applying text normalization before processing often yields better results. + +
+
+ Speaker + +
+ +
+ + + +
+ +
+ + +
+
+ +
+
+ Model Generated Audio0.00s + Audio Played0.00s +
+
+ +
+ Runtime Logs + 

+    

+  

+
+
+
+
+
diff --git a/docs/setup_gradio_demo.md b/docs/setup_gradio_demo.md
new file mode 100644
index 0000000..efe8728
--- /dev/null
+++ b/docs/setup_gradio_demo.md
@@ -0,0 +1,196 @@
+# VibeVoice Gradio Demo Setup Guide
+
+End-to-end instructions to deploy the VibeVoice ASR server and launch the Gradio web demo.
+
+## Prerequisites
+
+- CUDA-capable GPU(s)
+- Docker with GPU support (`nvidia-docker`)
+- VibeVoice repository cloned locally
+
+```bash
+git clone https://github.com/microsoft/VibeVoice.git
+cd VibeVoice
+```
+
+---
+
+## Step 1 — Start the ASR Server
+
+Launch a Docker container running the vLLM ASR server. The launcher script handles everything automatically (system deps, pip install, model download, tokenizer generation, server start).
+
+### Single GPU (default)
+
+```bash
+docker run -d --gpus '"device=0"' --name vibevoice-asr-demo \
+  --ipc=host \
+  -p 6001:6001 \
+  -e VIBEVOICE_FFMPEG_MAX_CONCURRENCY=64 \
+  -e PYTORCH_ALLOC_CONF=expandable_segments:True \
+  -v $(pwd):/app \
+  -w /app \
+  --entrypoint bash \
+  vllm/vllm-openai:v0.14.1 \
+  -c "python3 /app/vllm_plugin/scripts/start_server.py --port 6001"
+```
+
+### Multi-GPU with Data Parallel (load balancing)
+
+Run 4 independent replicas, one per GPU. vLLM distributes requests automatically:
+
+```bash
+docker run -d --gpus '"device=0,1,2,3"' --name vibevoice-asr-demo \
+  --ipc=host \
+  -p 6001:6001 \
+  -e VIBEVOICE_FFMPEG_MAX_CONCURRENCY=64 \
+  -e PYTORCH_ALLOC_CONF=expandable_segments:True \
+  -v $(pwd):/app \
+  -w /app \
+  --entrypoint bash \
+  vllm/vllm-openai:v0.14.1 \
+  -c "python3 /app/vllm_plugin/scripts/start_server.py --port 6001 --dp 4"
+```
+
+> **Tip**: Use `--dp N` for N-way data parallel (throughput scaling). Use `--tp N` for tensor parallel (large models). See `docs/vibevoice-vllm-asr.md` for details.
+
+### Check Logs
+
+```bash
+docker logs -f vibevoice-asr-demo
+```
+
+Wait until you see `Application startup complete.` — this means the server is ready.
+
+---
+
+## Step 2 — Verify the Server
+
+```bash
+# Check the model is loaded
+curl http://localhost:6001/v1/models
+```
+
+Expected output:
+```json
+{
+  "data": [{ "id": "vibevoice", ... }]
+}
+```
+
+### Quick Test with Audio File
+
+```bash
+docker exec -it vibevoice-asr-demo \
+  python3 /app/vllm_plugin/tests/test_api.py /app/en-Alice_woman.wav \
+  --url http://localhost:6001
+```
+
+---
+
+## Step 3 — Launch the Gradio Demo
+
+### Install tmux inside the container (to keep it running)
+
+```bash
+docker exec vibevoice-asr-demo apt-get install -y tmux
+```
+
+### Start Gradio in tmux
+
+```bash
+docker exec vibevoice-asr-demo bash -c \
+  "PYTHONUNBUFFERED=1 tmux new-session -d -s gradio \
+  'PYTHONUNBUFFERED=1 python3 /app/vllm_plugin/scripts/gradio_asr_demo_api_video.py \
+  --api_url http://localhost:6001 --share \
+  2>&1 | tee /tmp/gradio.log'"
+```
+
+### Get the Share Link
+
+Wait ~20 seconds, then:
+
+```bash
+docker exec vibevoice-asr-demo cat /tmp/gradio.log
+```
+
+You should see:
+```
+✅ Connected to API: http://localhost:6001 | Model: vibevoice
+🚀 Starting VibeVoice ASR Demo
+* Running on local URL:  http://0.0.0.0:7860
+* Running on public URL: https://xxxxxx.gradio.live
+```
+
+The `gradio.live` link is publicly accessible (valid for 1 week).
+
+### Gradio Options
+
+| Flag | Description | Default |
+|------|-------------|---------|
+| `--api_url URL` | vLLM server URL | `http://localhost:8000` |
+| `--share` | Create a public Gradio link | off |
+| `--port PORT` | Local Gradio port | `7860` |
+| `--cloudflared` | Use Cloudflare tunnel instead of Gradio share | off |
+| `--max_video_size MB` | Max upload video size | `50` |
+
+---
+
+## Managing the Service
+
+### Stop Gradio (keep ASR server running)
+
+```bash
+docker exec vibevoice-asr-demo tmux kill-session -t gradio
+```
+
+### Restart Gradio
+
+Re-run the tmux command from Step 3.
+
+### Stop Everything
+
+```bash
+docker stop vibevoice-asr-demo
+docker rm vibevoice-asr-demo
+```
+
+---
+
+## Example: Full Setup on GPU 0 with Port 6001
+
+```bash
+# 1. Start server
+docker run -d --gpus '"device=0"' --name vibevoice-asr-demo \
+  --ipc=host -p 6001:6001 \
+  -e VIBEVOICE_FFMPEG_MAX_CONCURRENCY=64 \
+  -e PYTORCH_ALLOC_CONF=expandable_segments:True \
+  -v $(pwd):/app -w /app \
+  --entrypoint bash \
+  vllm/vllm-openai:v0.14.1 \
+  -c "python3 /app/vllm_plugin/scripts/start_server.py --port 6001"
+
+# 2. Wait for startup (~2 min), then verify
+docker logs -f vibevoice-asr-demo  # wait for "Application startup complete."
+curl http://localhost:6001/v1/models
+
+# 3. Install tmux and launch Gradio
+docker exec vibevoice-asr-demo apt-get install -y tmux
+docker exec vibevoice-asr-demo bash -c \
+  "PYTHONUNBUFFERED=1 tmux new-session -d -s gradio \
+  'PYTHONUNBUFFERED=1 python3 /app/vllm_plugin/scripts/gradio_asr_demo_api_video.py \
+  --api_url http://localhost:6001 --share \
+  2>&1 | tee /tmp/gradio.log'"
+
+# 4. Get the public link
+sleep 20 && docker exec vibevoice-asr-demo cat /tmp/gradio.log
+```
+
+## Troubleshooting
+
+| Issue | Fix |
+|-------|-----|
+| `CUDA out of memory` | Use a different GPU (`device=X`) or reduce `--gpu-memory-utilization 0.7` in `start_server.py` |
+| Gradio log is empty | Wait longer (~30s); Gradio buffers output. Use `PYTHONUNBUFFERED=1` as shown above |
+| `Port already in use` | Pick a different port or stop the existing container: `docker stop  && docker rm ` |
+| Share link shows "No interface" | Gradio is still loading. Wait for `Application startup complete` in the log |
+| `tmux: command not found` | Run `docker exec  apt-get install -y tmux` first |
diff --git a/docs/vibevoice-asr.md b/docs/vibevoice-asr.md
new file mode 100644
index 0000000..5e65944
--- /dev/null
+++ b/docs/vibevoice-asr.md
@@ -0,0 +1,134 @@
+# VibeVoice-ASR
+
+[![Hugging Face](https://img.shields.io/badge/HuggingFace-Collection-orange?logo=huggingface)](https://huggingface.co/microsoft/VibeVoice-ASR)
+[![Live Playground](https://img.shields.io/badge/Live-Playground-green?logo=gradio)](https://aka.ms/vibevoice-asr)
+
+**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)

+
+
+## 🔥 Key Features
+
+- **🕒 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**:
+  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)**:
+  The model jointly performs ASR, diarization, and timestamping, producing a structured output that indicates *who* said *what* and *when*.
+  
+- **🌍 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
+

+
+# Demo
+
+

+
+https://github.com/user-attachments/assets/acde5602-dc17-4314-9e3b-c630bc84aefa
+
+

+
+## Evaluation
+

+  DER

+  cpWER

+  tcpWER
+

+
+
+
+## Installation
+We recommend using NVIDIA Deep Learning Container to manage the CUDA environment. 
+
+1. Launch docker
+```bash
+# 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
+
+## 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
+# pip install flash-attn --no-build-isolation
+```
+
+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
+
+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] 
+```
+
+
+## Finetuning
+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  |
+|----------------|-----------|------|-------|--------|------|
+| 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
+

+
+## 📄 License
+
+This project is licensed under the [MIT License](../LICENSE).
+
+
+
diff --git a/docs/vibevoice-realtime-0.5b.md b/docs/vibevoice-realtime-0.5b.md
new file mode 100644
index 0000000..6ca2deb
--- /dev/null
+++ b/docs/vibevoice-realtime-0.5b.md
@@ -0,0 +1,142 @@
+

+
+## 🎙️ VibeVoice-Realtime: Real-time Long‑Form Text‑to‑Speech with Streaming Input
+[![Hugging Face](https://img.shields.io/badge/HuggingFace-Collection-orange?logo=huggingface)](https://huggingface.co/microsoft/VibeVoice-Realtime-0.5B)
+[![Colab](https://img.shields.io/badge/Run-Colab-orange?logo=googlecolab)](https://colab.research.google.com/github/microsoft/VibeVoice/blob/main/demo/vibevoice_realtime_colab.ipynb)
+

+
+VibeVoice-Realtime is a **lightweight real‑time** text-to-speech model supporting **streaming text input** and **robust long-form speech generation**. It can be used to build real-time TTS services, narrate live data streams, and let different LLMs start speaking from their very first tokens (plug in your preferred model) long before a full answer is generated. It produces initial audible speech in **~200 milliseconds** (hardware dependent).
+
+
+**Model:** [VibeVoice Realtime 0.5B (on Hugging Face)](https://huggingface.co/microsoft/VibeVoice-Realtime-0.5B)

+**Colab:** [VibeVoice Realtime Colab (Jupyter Notebook)](https://colab.research.google.com/github/microsoft/VibeVoice/blob/main/demo/vibevoice_realtime_colab.ipynb)

+
+
+> Note (multilingual exploration): Although the model is primarily built for English, we found that it still exhibits a certain level of multilingual capability—and even performs reasonably well in some languages. We provide nine additional languages (German, French, Italian, Japanese, Korean, Dutch, Polish, Portuguese, and Spanish) for users to explore. These multilingual behaviors have not been extensively tested; use with caution and share observations.
+
+The model uses an interleaved, windowed design: it incrementally encodes incoming text chunks while, in parallel, continuing diffusion-based acoustic latent generation from prior context. Unlike the full multi-speaker long-form variants, this streaming model removes the semantic tokenizer and relies solely on an efficient acoustic tokenizer operating at an ultra-low frame rate (7.5 Hz).
+
+

+	
+		
+		VibeVoice Realtime Overview
+	
+	

+	Overview of VibeVoice Realtime Model.
+

+
+Key features:
+- Parameter size: 0.5B (deployment-friendly)
+- Real-time TTS (~200 milliseconds first audible latency)
+- Streaming text input
+- Robust long-form speech generation
+- 8k context window( ~10 minutes audio generation)
+
+This real-time variant supports only a single speaker. For multi‑speaker conversational speech generation, please use other VibeVoice models (long‑form multi‑speaker variants). The model is currently intended for English speech only; other languages may produce unpredictable results.
+
+To mitigate deepfake risks and ensure low latency for the first speech chunk, voice prompts are provided in an embedded format. For users requiring voice customization, please reach out to our team. We will also be expanding the range of available speakers.
+
+
+### 📋 TODO
+
+- [√] Add more voices (expand available speakers/voice timbres)
+- [ ] Implement streaming text input function to feed new tokens while audio is still being generated
+- [ ] Merge models into official HuggingFace's `transformers` repository 
+
+
+### 🎵 Demo Examples
+
+

+
+https://github.com/user-attachments/assets/9aa8ab3c-681d-4a02-b9ea-3f54ffd180b2
+
+

+
+
+## Results
+
+The model achieves satisfactory performance on short-sentence benchmarks, while the model is more focused on long‑form speech generation.
+
+### Zero-shot TTS performance on LibriSpeech test-clean set
+
+| Model | WER (%) ↓ | Speaker Similarity ↑ |
+|:--------------------|:---------:|:----------------:|
+| VALL-E 2            | 2.40      | 0.643            |
+| Voicebox            | 1.90      | 0.662            |
+| MELLE               | 2.10      | 0.625            |
+| **VibeVoice-Realtime-0.5B** | 2.00 | 0.695            |
+
+### Zero-shot TTS performance on SEED test-en set
+
+| Model | WER (%) ↓ | Speaker Similarity ↑ |
+|:--------------------|:---------:|:----------------:|
+| MaskGCT             | 2.62      | 0.714            |
+| Seed-TTS            | 2.25      | 0.762            |
+| FireRedTTS          | 3.82      | 0.460            |
+| SparkTTS            | 1.98      | 0.584            |
+| CosyVoice2          | 2.57      | 0.652            |
+| **VibeVoice-Realtime-0.5B** | 2.05 | 0.633            | 
+
+
+## Installation
+We recommend using NVIDIA Deep Learning Container to manage the CUDA environment. 
+
+1. Launch docker
+```bash
+# NVIDIA PyTorch Container 24.07 / 24.10 / 24.12 verified. 
+# Later 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:24.07-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
+# pip install flash-attn --no-build-isolation
+```
+
+2. Install from github
+```bash
+git clone https://github.com/microsoft/VibeVoice.git
+cd VibeVoice/
+
+pip install -e .[streamingtts]
+```
+
+
+## Usages
+
+
+### Usage 1: Launch real-time websocket demo
+Note: NVIDIA T4 / Mac M4 Pro achieve real-time performance in our tests; other devices with weaker inference capability may require further testing and speed optimizations. 
+
+Due to network latency, the time when audio playback is heard may exceed the ~300 ms first speech chunk generation latency.
+```bash
+python demo/vibevoice_realtime_demo.py --model_path microsoft/VibeVoice-Realtime-0.5B
+```
+
+Tip: Just try it on [Colab](https://colab.research.google.com/github/microsoft/VibeVoice/blob/main/demo/vibevoice_realtime_colab.ipynb).
+
+### Usage 2: Inference from files directly
+```bash
+# We provide some example scripts under demo/text_examples/ for demo
+python demo/realtime_model_inference_from_file.py --model_path microsoft/VibeVoice-Realtime-0.5B --txt_path demo/text_examples/1p_vibevoice.txt --speaker_name Carter
+```
+
+### [Optional] More experimental voices 
+Download additional experimental multi-lingual speakers before launching demo or inference from files.
+```bash
+bash demo/download_experimental_voices.sh
+```
+## Risks and limitations
+
+While efforts have been made to optimize it through various techniques, it may still produce outputs that are unexpected, biased, or inaccurate. VibeVoice inherits any biases, errors, or omissions produced by its base model (specifically, Qwen2.5 0.5b in this release).
+
+Potential for Deepfakes and Disinformation: High-quality synthetic speech can be misused to create convincing fake audio content for impersonation, fraud, or spreading disinformation. Users must ensure transcripts are reliable, check content accuracy, and avoid using generated content in misleading ways. Users are expected to use the generated content and to deploy the models in a lawful manner, in full compliance with all applicable laws and regulations in the relevant jurisdictions. It is best practice to disclose the use of AI when sharing AI-generated content.
+
+English only: Transcripts in languages other than English may result in unexpected audio outputs.
+
+Non-Speech Audio: The model focuses solely on speech synthesis and does not handle background noise, music, or other sound effects.
+
+Code, formulas, and special symbols: The model does not currently support reading code, mathematical formulas, or uncommon symbols. Please pre‑process input text to remove or normalize such content to avoid unpredictable results. 
+
+Very short inputs: When the input text is extremely short (three words or fewer), the model’s stability may degrade.
+
+We do not recommend using VibeVoice in commercial or real-world applications without further testing and development. This model is intended for research and development purposes only. Please use responsibly.
diff --git a/docs/vibevoice-tts.md b/docs/vibevoice-tts.md
new file mode 100644
index 0000000..a3a5f77
--- /dev/null
+++ b/docs/vibevoice-tts.md
@@ -0,0 +1,147 @@
+# VibeVoice-TTS
+
+[![Hugging Face](https://img.shields.io/badge/HuggingFace-Collection-orange?logo=huggingface)](https://huggingface.co/microsoft/VibeVoice-1.5B)
+[![Technical Report](https://img.shields.io/badge/Technical-Report-red?logo=arxiv)](https://arxiv.org/pdf/2508.19205)
+
+**VibeVoice-TTS** is a **long-form**, **multi-speaker** text-to-speech model designed to generate **expressive conversational audio** such as podcasts from text. It can synthesize speech up to **90 minutes** long with up to **4 distinct speakers**, surpassing the typical 1–2 speaker limits of many prior models.
+
+
+**Model:** [VibeVoice-1.5B](https://huggingface.co/microsoft/VibeVoice-1.5B)

+**Report:** [Technical Report](https://arxiv.org/pdf/2508.19205)

+
+
+

+
+| Model | Context Length | Generation Length | Weight |
+|-------|----------------|-------------------|--------|
+| VibeVoice-1.5B | 64K | ~90 min | [HF link](https://huggingface.co/microsoft/VibeVoice-1.5B) |
+| VibeVoice-Large | 32K | ~45 min | Disabled |
+
+

+
+## 🔥 Key Features
+
+- **⏱️ 90-minute Long-form Generation**:
+  Synthesizes conversational/single-speaker speech up to **90 minutes** in a single pass, maintaining speaker consistency and semantic coherence throughout.
+
+- **👥 Multi-speaker Support**:
+  Supports up to **4 distinct speakers** in a single conversation, with natural turn-taking and speaker consistency across long dialogues.
+
+- **🎭 Expressive Speech**:
+  Generates expressive, natural-sounding speech that captures conversational dynamics and emotional nuances.
+
+- **🌐 Multi-lingual Support**:
+  Supports English, Chinese and other languages.
+
+## 🏗️ Model Architecture
+
+VibeVoice-TTS employs a [next-token diffusion](https://arxiv.org/pdf/2508.19205) framework that combines:
+
+- **Large Language Model (LLM)**: Based on Qwen2.5, understands textual context and dialogue flow
+- **Continuous Speech Tokenizers**: Acoustic and Semantic tokenizers operating at an ultra-low frame rate of **7.5 Hz**, efficiently preserving audio fidelity while boosting computational efficiency
+- **Diffusion Head**: Generates high-fidelity acoustic details through diffusion-based generation
+
+

+  VibeVoice Overview
+

+
+
+## 🎵 Demo Examples
+
+**English**
+

+
+https://github.com/user-attachments/assets/0967027c-141e-4909-bec8-091558b1b784
+
+

+
+**Chinese**
+

+
+https://github.com/user-attachments/assets/322280b7-3093-4c67-86e3-10be4746c88f
+
+

+
+**Cross-Lingual**
+

+
+https://github.com/user-attachments/assets/838d8ad9-a201-4dde-bb45-8cd3f59ce722
+
+

+
+**Spontaneous Singing**
+

+
+https://github.com/user-attachments/assets/6f27a8a5-0c60-4f57-87f3-7dea2e11c730
+
+

+
+**Long Conversation with 4 people**
+

+
+https://github.com/user-attachments/assets/a357c4b6-9768-495c-a576-1618f6275727
+
+

+
+For more examples, see the [Project Page](https://microsoft.github.io/VibeVoice).
+
+## Installation and Usage
+Disabled due to widespread misuse.
+
+## Results
+
+The model achieves state-of-the-art performance on long-form multi-speaker speech generation tasks. For detailed evaluation results, please refer to the [paper](https://arxiv.org/pdf/2508.19205).
+

+  VibeVoice Results
+

+
+
+
+## 🚨 Tips
+We observed users may encounter occasional instability when synthesizing Chinese speech. We recommend:
+
+- Using English punctuation even for Chinese text, preferably only commas and periods.
+- Using the Large model variant, which is considerably more stable.
+- If you find that the generated voice speaks too fast, try chunking your text into multiple turns with the same speaker label.
+
+We'd like to thank [PsiPi](https://huggingface.co/PsiPi) for sharing an interesting way for emotion control. Details can be found via [discussion12](https://huggingface.co/microsoft/VibeVoice-1.5B/discussions/12).
+
+
+## FAQ
+#### Q1: Is this a pretrained model?
+**A:** Yes, it's a pretrained model without any post-training or benchmark-specific optimizations. In a way, this makes VibeVoice very versatile and fun to use.
+
+#### Q2: Randomly trigger Sounds / Music / BGM.
+**A:** As you can see from our demo page, the background music or sounds are spontaneous. This means we can't directly control whether they are generated or not. The model is content-aware, and these sounds are triggered based on the input text and the chosen voice prompt.
+
+Here are a few things we've noticed:
+*   If the voice prompt you use contains background music, the generated speech is more likely to have it as well. (The Large model is quite stable and effective at this—give it a try on the demo!)
+*   If the voice prompt is clean (no BGM), but the input text includes introductory words or phrases like "Welcome to," "Hello," or "However," background music might still appear.
+*   Speaker voice related, using "Alice" results in random BGM than others (fixed).
+*   In other scenarios, the Large model is more stable and has a lower probability of generating unexpected background music.
+
+In fact, we intentionally decided not to denoise our training data because we think it's an interesting feature for BGM to show up at just the right moment. You can think of it as a little easter egg we left for you.
+
+#### Q3: Text normalization?
+**A:** We don't perform any text normalization during training or inference. Our philosophy is that a large language model should be able to handle complex user inputs on its own. However, due to the nature of the training data, you might still run into some corner cases.
+
+#### Q4: Singing Capability.
+**A:** Our training data **doesn't contain any music data**. The ability to sing is an emergent capability of the model (which is why it might sound off-key, even on a famous song like 'See You Again'). (The Large model is more likely to exhibit this than the 1.5B).
+
+#### Q5: Some Chinese pronunciation errors.
+**A:** The volume of Chinese data in our training set is significantly smaller than the English data. Additionally, certain special characters (e.g., Chinese quotation marks) may occasionally cause pronunciation issues.
+
+#### Q6: Instability of cross-lingual transfer.
+**A:** The model does exhibit strong cross-lingual transfer capabilities, including the preservation of accents, but its performance can be unstable. This is an emergent ability of the model that we have not specifically optimized. It's possible that a satisfactory result can be achieved through repeated sampling.
+
+## Risks and Limitations
+
+While efforts have been made to optimize it through various techniques, it may still produce outputs that are unexpected, biased, or inaccurate. VibeVoice inherits any biases, errors, or omissions produced by its base model (specifically, Qwen2.5 1.5b in this release). Potential for Deepfakes and Disinformation: High-quality synthetic speech can be misused to create convincing fake audio content for impersonation, fraud, or spreading disinformation. Users must ensure transcripts are reliable, check content accuracy, and avoid using generated content in misleading ways. Users are expected to use the generated content and to deploy the models in a lawful manner, in full compliance with all applicable laws and regulations in the relevant jurisdictions. It is best practice to disclose the use of AI when sharing AI-generated content.
+
+English and Chinese only: Transcripts in languages other than English or Chinese may result in unexpected audio outputs.
+
+Non-Speech Audio: The model focuses solely on speech synthesis and does not handle background noise, music, or other sound effects.
+
+Overlapping Speech: The current model does not explicitly model or generate overlapping speech segments in conversations.
+
+We do not recommend using VibeVoice in commercial or real-world applications without further testing and development. This model is intended for research and development purposes only. Please use responsibly.
diff --git a/docs/vibevoice-vllm-asr.md b/docs/vibevoice-vllm-asr.md
new file mode 100644
index 0000000..9ed738b
--- /dev/null
+++ b/docs/vibevoice-vllm-asr.md
@@ -0,0 +1,189 @@
+# VibeVoice vLLM ASR Deployment
+
+Huggingface
+
+Deploy VibeVoice ASR model as a high-performance API service using [vLLM](https://github.com/vllm-project/vllm). This plugin provides OpenAI-compatible API endpoints for speech-to-text transcription with streaming support.
+
+## 🔥 Key Features
+
+- **🚀 High-Performance Serving**: Optimized for high-throughput ASR inference with vLLM's continuous batching
+- **📡 OpenAI-Compatible API**: Standard `/v1/chat/completions` endpoint with streaming support
+- **🎵 Long Audio Support**: Process up to 60+ minutes of audio in a single request
+- **🔌 Plugin Architecture**: No vLLM source code modification required - just install and run
+- **⚡ Data Parallel (DP)**: Run independent model replicas across multiple GPUs with automatic load balancing behind a single port
+
+## 🛠️ Installation
+
+Using Official vLLM Docker Image (Recommended)
+
+1. Clone the repository
+```bash
+git clone https://github.com/microsoft/VibeVoice.git
+cd VibeVoice
+```
+
+2. Launch the server (background mode)
+```bash
+docker run -d --gpus all --name vibevoice-vllm \
+  --ipc=host \
+  -p 8000:8000 \
+  -e VIBEVOICE_FFMPEG_MAX_CONCURRENCY=64 \
+  -e PYTORCH_ALLOC_CONF=expandable_segments:True \
+  -v $(pwd):/app \
+  -w /app \
+  --entrypoint bash \
+  vllm/vllm-openai:v0.14.1 \
+  -c "python3 /app/vllm_plugin/scripts/start_server.py"
+```
+
+## ⚡ Multi-GPU Deployment
+
+The launcher supports two types of GPU parallelism via `--tp` and `--dp` flags:
+
+| Flag | Name | What it does |
+|------|------|-------------|
+| `--tp N` | Tensor Parallel | Splits **one model** across N GPUs (for models too large for a single GPU) |
+| `--dp N` | Data Parallel | Runs **N independent replicas**, one per GPU, with automatic load balancing behind a single port |
+
+### Data Parallel (Recommended for scaling throughput)
+
+Run N independent replicas on N GPUs with automatic load balancing behind a single port.
+When `--dp N` is specified (N > 1), the launcher automatically starts N independent vLLM
+processes behind an nginx reverse proxy (2×N workers) for optimal throughput:
+
+```bash
+docker run -d --gpus '"device=0,1,2,3"' --name vibevoice-vllm \
+  --ipc=host \
+  -p 8000:8000 \
+  -e VIBEVOICE_FFMPEG_MAX_CONCURRENCY=64 \
+  -e PYTORCH_ALLOC_CONF=expandable_segments:True \
+  -v $(pwd):/app \
+  -w /app \
+  --entrypoint bash \
+  vllm/vllm-openai:v0.14.1 \
+  -c "python3 /app/vllm_plugin/scripts/start_server.py --dp 4"
+```
+
+Run on all 8 GPUs:
+
+```bash
+docker run -d --gpus all --name vibevoice-vllm \
+  --ipc=host \
+  -p 8000:8000 \
+  -e VIBEVOICE_FFMPEG_MAX_CONCURRENCY=64 \
+  -e PYTORCH_ALLOC_CONF=expandable_segments:True \
+  -v $(pwd):/app \
+  -w /app \
+  --entrypoint bash \
+  vllm/vllm-openai:v0.14.1 \
+  -c "python3 /app/vllm_plugin/scripts/start_server.py --dp 8"
+```
+
+### Tensor Parallel
+
+Split a single model across 2 GPUs (useful if GPU memory is limited):
+
+```bash
+docker run -d --gpus '"device=0,1"' --name vibevoice-vllm \
+  --ipc=host \
+  -p 8000:8000 \
+  -e VIBEVOICE_FFMPEG_MAX_CONCURRENCY=64 \
+  -e PYTORCH_ALLOC_CONF=expandable_segments:True \
+  -v $(pwd):/app \
+  -w /app \
+  --entrypoint bash \
+  vllm/vllm-openai:v0.14.1 \
+  -c "python3 /app/vllm_plugin/scripts/start_server.py --tp 2"
+```
+
+### Hybrid (DP × TP)
+
+Combine both — e.g., 2 replicas, each split across 2 GPUs (4 GPUs total):
+
+```bash
+docker run -d --gpus '"device=0,1,2,3"' --name vibevoice-vllm \
+  --ipc=host \
+  -p 8000:8000 \
+  -v $(pwd):/app \
+  -w /app \
+  --entrypoint bash \
+  vllm/vllm-openai:v0.14.1 \
+  -c "python3 /app/vllm_plugin/scripts/start_server.py --dp 2 --tp 2"
+```
+
+> **Note**: Total GPUs required = `dp × tp`. Make sure to expose enough GPU devices in the Docker `--gpus` flag.
+
+3. View logs
+```bash
+docker logs -f vibevoice-vllm
+```
+
+> **Note**: 
+> - The `-d` flag runs the container in background (detached mode)
+> - Use `docker stop vibevoice-vllm` to stop the service
+> - The model will be downloaded to HuggingFace cache (`~/.cache/huggingface`) inside the container
+
+## 🚀 Usages
+
+### Test the API
+
+Once the vLLM server is running, test it with the provided script:
+
+```bash
+# Basic transcription
+docker exec -it vibevoice-vllm python3 vllm_plugin/tests/test_api.py /app/audio.wav
+
+# With hotwords for better recognition of specific terms
+docker exec -it vibevoice-vllm python3 vllm_plugin/tests/test_api.py /app/audio.wav --hotwords "Microsoft,VibeVoice"
+
+```
+
+```bash
+# With auto-recovery from repetition loops (for long audio)
+docker exec -it vibevoice-vllm python3 vllm_plugin/tests/test_api_auto_recover.py /app/audio.wav
+
+# Auto-recover with hotwords
+docker exec -it vibevoice-vllm python3 vllm_plugin/tests/test_api_auto_recover.py /app/audio.wav --hotwords "Microsoft,VibeVoice"
+```
+
+> **Note**: 
+> - The audio/video file must be inside the mounted directory (`/app` in the container). Copy your files to the VibeVoice folder before testing.
+> - Hotwords help improve recognition of domain-specific terms like proper nouns, technical terms, and speaker names.
+
+### Environment Variables
+
+| Variable | Description | Default |
+|----------|-------------|---------|
+| `VIBEVOICE_FFMPEG_MAX_CONCURRENCY` | Maximum FFmpeg processes for audio decoding | `64` |
+| `PYTORCH_ALLOC_CONF` | PyTorch memory allocator config | `expandable_segments:True` |
+
+
+
+## 📊 Performance Tips
+
+1. **GPU Memory**: Use `--gpu-memory-utilization 0.9` for maximum throughput if you have dedicated GPU
+2. **Batch Size**: Increase `--max-num-seqs` for higher concurrency (requires more GPU memory)
+3. **FFmpeg Concurrency**: Tune `VIBEVOICE_FFMPEG_MAX_CONCURRENCY` based on CPU cores
+
+## 🚨 Troubleshooting
+
+### Common Issues
+
+1. **"CUDA out of memory"**
+   - Reduce `--gpu-memory-utilization`
+   - Reduce `--max-num-seqs`
+   - Use smaller `--max-model-len`
+
+2. **"Audio decoding failed"**
+   - Ensure FFmpeg is installed: `ffmpeg -version`
+   - Check audio file format is supported 
+
+3. **"Model not found"**
+   - Ensure model path contains `config.json` and model weights
+   - Generate tokenizer files if missing
+
+4. **"Plugin not loaded"**
+   - Verify installation: `pip show vibevoice`
+   - Check entry point: `pip show -f vibevoice | grep entry`
+
+
diff --git a/finetuning-asr/README.md b/finetuning-asr/README.md
new file mode 100644
index 0000000..32768b4
--- /dev/null
+++ b/finetuning-asr/README.md
@@ -0,0 +1,154 @@
+# VibeVoice ASR LoRA Fine-tuning
+
+This directory contains scripts for LoRA (Low-Rank Adaptation) fine-tuning of the VibeVoice ASR model.
+
+## Requirements
+
+```bash
+# Install vibevoice first
+pip install -e .
+
+pip install peft
+```
+
+## Toy Dataset
+
+> **Note**: The `toy_dataset/` included in this directory contains **synthetic audio generated by VibeVoice TTS** for demonstration purposes only. It is NOT a full finetuning dataset.
+>
+> When using your own data, you should:
+> - Prepare real audio recordings with accurate transcriptions
+> - Adjust hyperparameters (learning rate, epochs, LoRA rank) based on your dataset size and domain
+> - Consider the audio quality and speaker diversity in your data
+
+## Data Format
+
+Training data should be organized as pairs of audio files and JSON labels in the same directory:
+
+```
+toy_dataset/
+├── 0.mp3
+├── 0.json
+├── 1.mp3
+├── 1.json
+└── ...
+```
+
+### JSON Label Format
+
+Each JSON file should have the following structure:
+
+```json
+{
+  "audio_duration": 351.73,
+  "audio_path": "0.mp3",
+  "segments": [
+    {
+      "speaker": 0,
+      "text": "Hey everyone, welcome back...",
+      "start": 0.0,
+      "end": 38.68
+    },
+    {
+      "speaker": 1,
+      "text": "Thanks for having me...",
+      "start": 38.75,
+      "end": 77.88
+    }
+  ],
+  "customized_context": ["Tea Brew", "Aiden Host", "The property is near Meter Street."]  // optional, domain-specific terms or context sentences
+}
+```
+
+## Training
+
+### Basic
+
+```bash
+# 1 GPU
+torchrun --nproc_per_node=1 lora_finetune.py \
+    --model_path microsoft/VibeVoice-ASR \
+    --data_dir ./toy_dataset \
+    --output_dir ./output \
+    --num_train_epochs 3 \
+    --per_device_train_batch_size 1 \
+    --learning_rate 1e-4 \
+    --bf16 \
+    --report_to none
+
+# Specific GPUs (e.g., GPU 0,1,2,3)
+CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --nproc_per_node=4 lora_finetune.py \
+    --model_path microsoft/VibeVoice-ASR \
+    --data_dir ./toy_dataset \
+    --output_dir ./output \
+    --num_train_epochs 3 \
+    --per_device_train_batch_size 1 \
+    --learning_rate 1e-4 \
+    --bf16 \
+    --report_to none
+```
+
+### Full Options
+
+The script uses HuggingFace's `TrainingArguments`, so all standard options are available:
+
+```bash
+torchrun --nproc_per_node=4 lora_finetune.py \
+    --model_path microsoft/VibeVoice-ASR \
+    --data_dir ./toy_dataset \
+    --output_dir ./output \
+    --lora_r 16 \
+    --lora_alpha 32 \
+    --lora_dropout 0.05 \
+    --num_train_epochs 3 \
+    --per_device_train_batch_size 1 \
+    --gradient_accumulation_steps 4 \
+    --learning_rate 1e-4 \
+    --warmup_ratio 0.1 \
+    --weight_decay 0.01 \
+    --max_grad_norm 1.0 \
+    --logging_steps 10 \
+    --save_steps 100 \
+    --gradient_checkpointing \
+    --bf16 \
+    --report_to none
+```
+
+### Key Parameters
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `--lora_r` | 16 | LoRA rank (lower = fewer params, higher = more expressive) |
+| `--lora_alpha` | 32 | LoRA scaling factor (typically 2x rank) |
+| `--lora_dropout` | 0.05 | Dropout for LoRA layers |
+| `--per_device_train_batch_size` | 8 | Batch size per device |
+| `--gradient_accumulation_steps` | 1 | Effective batch size = batch_size × grad_accum |
+| `--learning_rate` | 5e-5 | Learning rate (1e-4 to 2e-4 typical for LoRA) |
+| `--gradient_checkpointing` | False | Enable to reduce memory usage |
+| `--use_customized_context` | True | Include customized_context from JSON as additional context |
+| `--max_audio_length` | None | Skip audio longer than this (seconds) |
+
+## Inference with Fine-tuned Model
+
+```bash
+python inference_lora.py \
+    --base_model microsoft/VibeVoice-ASR \
+    --lora_path ./output \
+    --audio_file ./toy_dataset/0.mp3 \
+    --context_info "Tea Brew, Aiden Host"
+```
+
+## Merging LoRA Weights (Optional)
+
+To merge LoRA weights into the base model for faster inference:
+
+```python
+from peft import PeftModel
+
+# Load base model + LoRA
+model = VibeVoiceASRForConditionalGeneration.from_pretrained("microsoft/VibeVoice-ASR", ...)
+model = PeftModel.from_pretrained(model, "./output")
+
+# Merge and save
+model = model.merge_and_unload()
+model.save_pretrained("./merged_model")
+```
diff --git a/finetuning-asr/inference_lora.py b/finetuning-asr/inference_lora.py
new file mode 100644
index 0000000..d6d8163
--- /dev/null
+++ b/finetuning-asr/inference_lora.py
@@ -0,0 +1,234 @@
+#!/usr/bin/env python
+"""
+Inference with LoRA Fine-tuned VibeVoice ASR Model
+
+This script loads a LoRA fine-tuned model and runs inference.
+
+Usage:
+    python inference_lora.py \
+        --base_model microsoft/VibeVoice-ASR \
+        --lora_path ./output \
+        --audio_file ./toy_dataset/0.mp3
+"""
+
+import argparse
+import torch
+
+from peft import PeftModel
+
+from vibevoice.modular.modeling_vibevoice_asr import VibeVoiceASRForConditionalGeneration
+from vibevoice.processor.vibevoice_asr_processor import VibeVoiceASRProcessor
+
+
+def load_lora_model(
+    base_model_path: str,
+    lora_path: str,
+    device: str = "cuda",
+    dtype: torch.dtype = torch.bfloat16,
+):
+    """
+    Load base model and merge with LoRA weights.
+    
+    Args:
+        base_model_path: Path to base pretrained model
+        lora_path: Path to LoRA adapter weights
+        device: Device to load model on
+        dtype: Data type for model
+        
+    Returns:
+        Tuple of (model, processor)
+    """
+    print(f"Loading base model from {base_model_path}")
+    
+    # Load processor
+    processor = VibeVoiceASRProcessor.from_pretrained(
+        base_model_path,
+        language_model_pretrained_name="Qwen/Qwen2.5-7B"
+    )
+    
+    # Load base model
+    model = VibeVoiceASRForConditionalGeneration.from_pretrained(
+        base_model_path,
+        dtype=dtype,
+        device_map=device if device == "auto" else None,
+        attn_implementation="flash_attention_2",
+        trust_remote_code=True,
+    )
+    
+    if device != "auto":
+        model = model.to(device)
+    
+    # Load LoRA adapter
+    print(f"Loading LoRA adapter from {lora_path}")
+    model = PeftModel.from_pretrained(model, lora_path)
+    
+    # Optionally merge LoRA weights into base model for faster inference
+    # model = model.merge_and_unload()
+    
+    model.eval()
+    print("Model loaded successfully")
+    
+    return model, processor
+
+
+def transcribe(
+    model,
+    processor,
+    audio_path: str,
+    max_new_tokens: int = 4096,
+    temperature: float = 0.0,
+    context_info: str = None,
+    device: str = "cuda",
+):
+    """
+    Transcribe an audio file using the LoRA fine-tuned model.
+    
+    Args:
+        model: The LoRA fine-tuned model
+        processor: The processor
+        audio_path: Path to audio file
+        max_new_tokens: Maximum tokens to generate
+        temperature: Sampling temperature (0 = greedy)
+        context_info: Optional context info (e.g., hotwords)
+        device: Device
+        
+    Returns:
+        Transcription result
+    """
+    print(f"\nTranscribing: {audio_path}")
+    
+    # Process audio
+    inputs = processor(
+        audio=audio_path,
+        sampling_rate=None,
+        return_tensors="pt",
+        padding=True,
+        add_generation_prompt=True,
+        context_info=context_info,
+    )
+    
+    # Move to device
+    inputs = {k: v.to(device) if isinstance(v, torch.Tensor) else v 
+              for k, v in inputs.items()}
+    
+    # Generation config
+    gen_config = {
+        "max_new_tokens": max_new_tokens,
+        "pad_token_id": processor.pad_id,
+        "eos_token_id": processor.tokenizer.eos_token_id,
+        "do_sample": temperature > 0,
+    }
+    if temperature > 0:
+        gen_config["temperature"] = temperature
+        gen_config["top_p"] = 0.9
+    
+    # Generate
+    with torch.no_grad():
+        output_ids = model.generate(**inputs, **gen_config)
+    
+    # Decode
+    input_length = inputs['input_ids'].shape[1]
+    generated_ids = output_ids[0, input_length:]
+    generated_text = processor.decode(generated_ids, skip_special_tokens=True)
+    
+    # Parse structured output
+    try:
+        segments = processor.post_process_transcription(generated_text)
+    except Exception as e:
+        print(f"Warning: Failed to parse structured output: {e}")
+        segments = []
+    
+    return {
+        "raw_text": generated_text,
+        "segments": segments,
+    }
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Inference with LoRA Fine-tuned VibeVoice ASR")
+    
+    parser.add_argument(
+        "--base_model",
+        type=str,
+        default="microsoft/VibeVoice-ASR",
+        help="Path to base pretrained model"
+    )
+    parser.add_argument(
+        "--lora_path",
+        type=str,
+        required=True,
+        help="Path to LoRA adapter weights"
+    )
+    parser.add_argument(
+        "--audio_file",
+        type=str,
+        required=True,
+        help="Path to audio file to transcribe"
+    )
+    parser.add_argument(
+        "--context_info",
+        type=str,
+        default=None,
+        help="Optional context info (e.g., 'Hotwords: Tea Brew, Aiden Host')"
+    )
+    parser.add_argument(
+        "--max_new_tokens",
+        type=int,
+        default=4096,
+        help="Maximum tokens to generate"
+    )
+    parser.add_argument(
+        "--temperature",
+        type=float,
+        default=0.0,
+        help="Sampling temperature (0 = greedy)"
+    )
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="cuda" if torch.cuda.is_available() else "cpu",
+        help="Device to use"
+    )
+    
+    args = parser.parse_args()
+    
+    # Load model
+    dtype = torch.bfloat16 if args.device != "cpu" else torch.float32
+    model, processor = load_lora_model(
+        base_model_path=args.base_model,
+        lora_path=args.lora_path,
+        device=args.device,
+        dtype=dtype,
+    )
+    
+    # Transcribe
+    result = transcribe(
+        model=model,
+        processor=processor,
+        audio_path=args.audio_file,
+        max_new_tokens=args.max_new_tokens,
+        temperature=args.temperature,
+        context_info=args.context_info,
+        device=args.device,
+    )
+    
+    # Print results
+    print("\n" + "="*60)
+    print("Transcription Result")
+    print("="*60)
+    
+    print("\n--- Raw Output ---")
+    raw_text = result['raw_text']
+    print(raw_text[:2000] + "..." if len(raw_text) > 2000 else raw_text)
+    
+    if result['segments']:
+        print(f"\n--- Structured Output ({len(result['segments'])} segments) ---")
+        for seg in result['segments'][:20]:
+            print(f"[{seg.get('start_time', 'N/A')} - {seg.get('end_time', 'N/A')}] "
+                  f"Speaker {seg.get('speaker_id', 'N/A')}: {seg.get('text', '')[:80]}...")
+        if len(result['segments']) > 20:
+            print(f"  ... and {len(result['segments']) - 20} more segments")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/finetuning-asr/lora_finetune.py b/finetuning-asr/lora_finetune.py
new file mode 100644
index 0000000..0c4f9c9
--- /dev/null
+++ b/finetuning-asr/lora_finetune.py
@@ -0,0 +1,556 @@
+#!/usr/bin/env python
+"""
+VibeVoice ASR LoRA Fine-tuning Script
+
+This script implements LoRA (Low-Rank Adaptation) fine-tuning for the VibeVoice ASR model.
+It uses PEFT (Parameter-Efficient Fine-Tuning) library for efficient training.
+"""
+
+import json
+import logging
+from pathlib import Path
+from typing import List, Dict, Any, Optional, Tuple
+from dataclasses import dataclass, field
+
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader
+import numpy as np
+
+from transformers import (
+    TrainingArguments,
+    Trainer,
+    HfArgumentParser,
+)
+from peft import (
+    LoraConfig,
+    get_peft_model,
+    prepare_model_for_kbit_training,
+    TaskType,
+)
+
+from vibevoice.modular.modeling_vibevoice_asr import VibeVoiceASRForConditionalGeneration
+from vibevoice.processor.vibevoice_asr_processor import VibeVoiceASRProcessor
+
+# Setup logging
+logging.basicConfig(
+    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+    datefmt="%Y-%m-%d %H:%M:%S",
+    level=logging.INFO,
+)
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ModelArguments:
+    """Arguments for model configuration."""
+    model_path: str = field(
+        default="microsoft/VibeVoice-ASR",
+        metadata={"help": "Path to pretrained model (HuggingFace model ID or local path)"}
+    )
+
+
+@dataclass
+class DataArguments:
+    """Arguments for data configuration."""
+    data_dir: str = field(
+        default="./toy_dataset",
+        metadata={"help": "Directory containing training data"}
+    )
+    max_audio_length: Optional[float] = field(
+        default=None,
+        metadata={"help": "Maximum audio length in seconds (default: no limit)"}
+    )
+    use_customized_context: bool = field(
+        default=True,
+        metadata={"help": "Whether to use customized_context from JSON as additional context"}
+    )
+
+
+@dataclass
+class LoraArguments:
+    """Arguments for LoRA configuration."""
+    lora_r: int = field(
+        default=16,
+        metadata={"help": "LoRA rank"}
+    )
+    lora_alpha: int = field(
+        default=32,
+        metadata={"help": "LoRA alpha (scaling factor)"}
+    )
+    lora_dropout: float = field(
+        default=0.05,
+        metadata={"help": "LoRA dropout"}
+    )
+
+
+@dataclass
+class VibeVoiceASRDataCollator:
+    """
+    Data collator for VibeVoice ASR fine-tuning.
+    Handles batching of variable-length audio and text sequences.
+    """
+    processor: VibeVoiceASRProcessor
+    pad_token_id: int
+    label_pad_token_id: int = -100
+    
+    def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, torch.Tensor]:
+        """
+        Collate a batch of features into model inputs.
+        """
+        # Separate inputs and labels
+        input_ids_list = [f["input_ids"] for f in features]
+        labels_list = [f["labels"] for f in features]
+        acoustic_mask_list = [f["acoustic_input_mask"] for f in features]
+        speech_list = [f["speech"] for f in features]
+        vae_tok_lens = [f["vae_tok_len"] for f in features]
+        
+        # Determine max lengths
+        max_seq_len = max(len(ids) for ids in input_ids_list)
+        max_speech_len = max(len(s) for s in speech_list)
+        max_vae_len = max(vae_tok_lens)
+        
+        batch_size = len(features)
+        
+        # Initialize padded tensors
+        input_ids = torch.full((batch_size, max_seq_len), self.pad_token_id, dtype=torch.long)
+        attention_mask = torch.zeros((batch_size, max_seq_len), dtype=torch.long)
+        labels = torch.full((batch_size, max_seq_len), self.label_pad_token_id, dtype=torch.long)
+        acoustic_input_mask = torch.zeros((batch_size, max_seq_len), dtype=torch.bool)
+        speech_tensors = torch.zeros((batch_size, max_speech_len), dtype=torch.float32)
+        speech_masks = torch.zeros((batch_size, max_vae_len), dtype=torch.bool)
+        
+        # Fill in the tensors (right padding for training)
+        # Note: processor uses left padding for inference/generation, but training uses right padding
+        for i, (ids, lbls, amask, speech, vae_len) in enumerate(
+            zip(input_ids_list, labels_list, acoustic_mask_list, speech_list, vae_tok_lens)
+        ):
+            seq_len = len(ids)
+            
+            # Right padding for input_ids and labels
+            input_ids[i, :seq_len] = torch.tensor(ids, dtype=torch.long)
+            attention_mask[i, :seq_len] = 1
+            labels[i, :seq_len] = torch.tensor(lbls, dtype=torch.long)
+            acoustic_input_mask[i, :seq_len] = torch.tensor(amask, dtype=torch.bool)
+            
+            # Speech tensors (right padding, zeros work as padding)
+            speech_len = len(speech)
+            speech_tensors[i, :speech_len] = torch.tensor(speech, dtype=torch.float32)
+            speech_masks[i, :vae_len] = True
+        
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "labels": labels,
+            "acoustic_input_mask": acoustic_input_mask,
+            "speech_tensors": speech_tensors,
+            "speech_masks": speech_masks,
+        }
+
+
+class VibeVoiceASRDataset(Dataset):
+    """
+    Dataset for VibeVoice ASR fine-tuning.
+    
+    Expected data format:
+        - Audio files: .mp3, .wav, .flac, etc.
+        - Label files: .json with matching name
+        
+    JSON format:
+        {
+            "audio_path": "0.mp3",
+            "audio_duration": 351.73,
+            "segments": [
+                {
+                    "speaker": 0,
+                    "text": "Hey everyone, welcome back...",
+                    "start": 0.0,
+                    "end": 38.68
+                },
+                ...
+            ],
+            "customized_context": ["Tea Brew", "The property is near Meter Street."]  # optional
+        }
+    """
+    
+    def __init__(
+        self,
+        data_dir: str,
+        processor: VibeVoiceASRProcessor,
+        max_audio_length: Optional[float] = None,  # in seconds
+        use_customized_context: bool = True,
+    ):
+        """
+        Initialize the dataset.
+        
+        Args:
+            data_dir: Directory containing audio files and JSON labels
+            processor: VibeVoice ASR processor
+            max_audio_length: Maximum audio length in seconds (None = no limit)
+            use_customized_context: Whether to include customized_context in prompt
+        """
+        self.data_dir = Path(data_dir)
+        self.processor = processor
+        self.max_audio_length = max_audio_length
+        self.use_customized_context = use_customized_context
+        
+        # Find all JSON files
+        self.samples = self._load_samples()
+        logger.info(f"Loaded {len(self.samples)} samples from {data_dir}")
+    
+    def _load_samples(self) -> List[Dict[str, Any]]:
+        """Load and validate all samples from data directory."""
+        samples = []
+        
+        for json_path in sorted(self.data_dir.glob("*.json")):
+            try:
+                with open(json_path, "r", encoding="utf-8") as f:
+                    data = json.load(f)
+                
+                # Get audio path from JSON
+                audio_filename = data.get("audio_path")
+                if not audio_filename:
+                    logger.warning(f"No audio_path specified in {json_path}")
+                    continue
+                    
+                audio_path = self.data_dir / audio_filename
+                if not audio_path.exists():
+                    logger.warning(f"Audio file not found: {audio_path}")
+                    continue
+                
+                # Optional: filter by duration
+                if self.max_audio_length is not None:
+                    duration = data.get("audio_duration", float("inf"))
+                    if duration > self.max_audio_length:
+                        logger.info(f"Skipping {json_path.stem}: duration {duration:.1f}s > max {self.max_audio_length}s")
+                        continue
+                
+                samples.append({
+                    "audio_path": str(audio_path),
+                    "json_path": str(json_path),
+                    "data": data,
+                })
+                
+            except Exception as e:
+                logger.warning(f"Error loading {json_path}: {e}")
+                continue
+        
+        return samples
+    
+    def _format_transcription(self, segments: List[Dict], audio_duration: float) -> str:
+        """
+        Format transcription segments into JSON output format.
+        
+        This matches the expected model output format used in training.
+        """
+        formatted_segments = []
+        
+        for seg in segments:
+            formatted_seg = {}
+            # Add timestamp
+            formatted_seg["Start"] = round(seg['start'], 2)
+            formatted_seg["End"] = round(seg['end'], 2)
+            # Add speaker if available
+            if "speaker" in seg:
+                formatted_seg["Speaker"] = seg["speaker"]
+            # Add content
+            formatted_seg["Content"] = seg.get("text", "")
+            formatted_segments.append(formatted_seg)
+        
+        # Return as compact JSON string (no spaces after separators)
+        return json.dumps(formatted_segments, ensure_ascii=False, separators=(',', ':'))
+    
+    def __len__(self) -> int:
+        return len(self.samples)
+    
+    def __getitem__(self, idx: int) -> Dict[str, Any]:
+        """
+        Get a single sample for training.
+        
+        Returns:
+            Dict with:
+                - input_ids: Token IDs for input (system + user + assistant prompt)
+                - labels: Token IDs for labels (-100 for non-predicted tokens)
+                - acoustic_input_mask: Mask for speech token positions
+                - speech: Raw audio array
+                - vae_tok_len: Number of speech tokens
+        """
+        sample = self.samples[idx]
+        data = sample["data"]
+        audio_path = sample["audio_path"]
+        
+        # Prepare context info (customized_context)
+        context_info = None
+        if self.use_customized_context and "customized_context" in data:
+            customized_context = data["customized_context"]
+            if customized_context:
+                context_info = "\n".join(customized_context)
+        
+        # Process audio using the processor's internal method
+        encoding = self.processor._process_single_audio(
+            audio_path,
+            sampling_rate=None,
+            add_generation_prompt=True,
+            use_streaming=True,
+            context_info=context_info,
+        )
+        
+        # Get the input tokens (system + user + generation prompt)
+        input_ids = encoding["input_ids"]
+        acoustic_input_mask = encoding["acoustic_input_mask"]
+        speech = encoding["speech"]
+        vae_tok_len = encoding["vae_tok_len"]
+        
+        # Format the target transcription
+        target_text = self._format_transcription(
+            data["segments"],
+            data.get("audio_duration", len(speech) / 24000)
+        )
+        
+        # Encode target using apply_chat_template to match training format
+        # This adds the assistant role tokens (e.g., <|im_start|>assistant\n...<|im_end|>)
+        target_tokens = self.processor.tokenizer.apply_chat_template(
+            [{"role": "assistant", "content": target_text}],
+            tokenize=True,
+            add_generation_prompt=False,
+        )
+        
+        # Combine input and target
+        full_input_ids = input_ids + target_tokens
+        full_acoustic_mask = acoustic_input_mask + [0] * len(target_tokens)
+        
+        # Create labels: -100 for input tokens, actual tokens for target
+        # We mask the input portion so loss is only computed on the response
+        labels = [-100] * len(input_ids) + target_tokens
+        
+        return {
+            "input_ids": full_input_ids,
+            "labels": labels,
+            "acoustic_input_mask": full_acoustic_mask,
+            "speech": speech,
+            "vae_tok_len": vae_tok_len,
+        }
+
+
+def get_lora_config(
+    r: int = 16,
+    lora_alpha: int = 32,
+    lora_dropout: float = 0.05,
+    target_modules: Optional[List[str]] = None,
+) -> LoraConfig:
+    """
+    Create LoRA configuration for VibeVoice ASR model.
+    
+    We apply LoRA to the language model's attention layers and MLP,
+    following common practices for LLM fine-tuning.
+    
+    Args:
+        r: LoRA rank
+        lora_alpha: LoRA scaling factor
+        lora_dropout: Dropout for LoRA layers
+        target_modules: List of module names to apply LoRA to
+        
+    Returns:
+        LoraConfig object
+    """
+    if target_modules is None:
+        # Target Qwen2 attention and MLP layers
+        # These are the common targets for language model fine-tuning
+        target_modules = [
+            "q_proj",
+            "k_proj", 
+            "v_proj",
+            "o_proj",
+            "gate_proj",
+            "up_proj",
+            "down_proj",
+        ]
+    
+    return LoraConfig(
+        r=r,
+        lora_alpha=lora_alpha,
+        target_modules=target_modules,
+        lora_dropout=lora_dropout,
+        bias="none",
+        task_type=TaskType.CAUSAL_LM,
+    )
+
+
+def setup_model_for_training(
+    model_path: str,
+    lora_config: LoraConfig,
+    device: str = "cuda",
+    dtype: torch.dtype = torch.bfloat16,
+    gradient_checkpointing: bool = True,
+) -> Tuple[nn.Module, VibeVoiceASRProcessor]:
+    """
+    Load and prepare model for LoRA training.
+    
+    Args:
+        model_path: Path to pretrained model
+        lora_config: LoRA configuration
+        device: Device to use
+        dtype: Data type for model
+        gradient_checkpointing: Whether to use gradient checkpointing
+        
+    Returns:
+        Tuple of (model, processor)
+    """
+    logger.info(f"Loading model from {model_path}")
+    
+    # Load processor
+    processor = VibeVoiceASRProcessor.from_pretrained(
+        model_path,
+        language_model_pretrained_name="Qwen/Qwen2.5-7B"
+    )
+    
+    # Load model
+    model = VibeVoiceASRForConditionalGeneration.from_pretrained(
+        model_path,
+        dtype=dtype,
+        device_map=device if device == "auto" else None,
+        attn_implementation="flash_attention_2",
+        trust_remote_code=True,
+    )
+    
+    if device != "auto":
+        model = model.to(device)
+    
+    # Freeze speech tokenizers (we only want to fine-tune the language model)
+    for name, param in model.named_parameters():
+        if "acoustic_tokenizer" in name or "semantic_tokenizer" in name:
+            param.requires_grad = False
+            logger.debug(f"Frozen: {name}")
+    
+    # Apply LoRA
+    logger.info(f"Applying LoRA with config: r={lora_config.r}, alpha={lora_config.lora_alpha}")
+    model = get_peft_model(model, lora_config)
+    
+    # Print trainable parameters
+    model.print_trainable_parameters()
+    
+    # Enable gradient checkpointing if requested
+    if gradient_checkpointing:
+        model.enable_input_require_grads()
+        model.gradient_checkpointing_enable()
+    
+    return model, processor
+
+
+def train(
+    model_args: ModelArguments,
+    data_args: DataArguments,
+    lora_args: LoraArguments,
+    training_args: TrainingArguments,
+    gradient_checkpointing: bool = True,
+):
+    """
+    Main training function for LoRA fine-tuning.
+    
+    Args:
+        model_args: Model configuration arguments
+        data_args: Data configuration arguments
+        lora_args: LoRA configuration arguments
+        training_args: HuggingFace TrainingArguments
+        gradient_checkpointing: Whether to use gradient checkpointing
+    """
+    # Set seed
+    torch.manual_seed(training_args.seed)
+    np.random.seed(training_args.seed)
+    
+    # Setup LoRA config
+    lora_config = get_lora_config(
+        r=lora_args.lora_r,
+        lora_alpha=lora_args.lora_alpha,
+        lora_dropout=lora_args.lora_dropout,
+    )
+    
+    # Determine device
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    # Load model and processor
+    dtype = torch.bfloat16 if device != "cpu" else torch.float32
+    model, processor = setup_model_for_training(
+        model_path=model_args.model_path,
+        lora_config=lora_config,
+        device=device,
+        dtype=dtype,
+        gradient_checkpointing=gradient_checkpointing,
+    )
+    
+    # Create dataset
+    train_dataset = VibeVoiceASRDataset(
+        data_dir=data_args.data_dir,
+        processor=processor,
+        max_audio_length=data_args.max_audio_length,
+        use_customized_context=data_args.use_customized_context,
+    )
+    
+    if len(train_dataset) == 0:
+        logger.error("No training samples found!")
+        return
+    
+    # Create data collator
+    data_collator = VibeVoiceASRDataCollator(
+        processor=processor,
+        pad_token_id=processor.pad_id,
+    )
+    
+    # Set some sensible defaults for audio training
+    training_args.dataloader_num_workers = 0  # Audio loading can be tricky with multiprocessing
+    training_args.remove_unused_columns = False  # Keep all columns
+    
+    # Create trainer
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        data_collator=data_collator,
+    )
+    
+    # Train
+    logger.info("Starting training...")
+    logger.info(f"  Num samples = {len(train_dataset)}")
+    logger.info(f"  Num epochs = {training_args.num_train_epochs}")
+    logger.info(f"  Batch size = {training_args.per_device_train_batch_size}")
+    logger.info(f"  Gradient accumulation steps = {training_args.gradient_accumulation_steps}")
+    total_steps = len(train_dataset) * int(training_args.num_train_epochs) // (
+        training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps
+    )
+    logger.info(f"  Total optimization steps = {total_steps}")
+    
+    train_result = trainer.train()
+    
+    # Save final model
+    logger.info(f"Saving model to {training_args.output_dir}")
+    trainer.save_model(training_args.output_dir)
+    
+    # Save training metrics
+    metrics = train_result.metrics
+    trainer.log_metrics("train", metrics)
+    trainer.save_metrics("train", metrics)
+    
+    # Save processor config
+    processor.save_pretrained(training_args.output_dir)
+    
+    logger.info("Training complete!")
+    
+    return model, processor
+
+
+def main():
+    # Use HfArgumentParser to parse all argument dataclasses
+    parser = HfArgumentParser((ModelArguments, DataArguments, LoraArguments, TrainingArguments))
+    model_args, data_args, lora_args, training_args = parser.parse_args_into_dataclasses()
+    
+    # Run training
+    train(
+        model_args=model_args,
+        data_args=data_args,
+        lora_args=lora_args,
+        training_args=training_args,
+    )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/finetuning-asr/toy_dataset/0.json b/finetuning-asr/toy_dataset/0.json
new file mode 100755
index 0000000..86a3498
--- /dev/null
+++ b/finetuning-asr/toy_dataset/0.json
@@ -0,0 +1,73 @@
+{
+  "audio_duration": 351.73333333333335,
+  "audio_path": "0.mp3",
+  "segments": [
+    {
+      "speaker": 0,
+      "text": "Hey everyone, welcome back to Tea Brew. I’m Aidan Host, and, uh, sorry we’re a tad late today. Travel was a bit wild on the roads and our guest came in from near Meeter Street, so, you know, life happens. But we’re here now and buzzing to talk property tech. We’ve had loads of messages from landlords about spiraling energy costs in HMOs and, at the same time, a bunch of questions on streamlining tenant recruitment. So today we’re kinda merging both worlds: how a smart heating setup can cut bills, and how a software assistant like Rent Byte can take the grind out of advertising, onboarding, and maintenance reporting. And sitting with me is Sayid Guest, who’s, uh, a landlord turned gadget builder—haha—ready to demystify the thermostatic stuff that tenants can actually live with.",
+      "start": 0.0,
+      "end": 38.68
+    },
+    {
+      "speaker": 1,
+      "text": "Thanks, Aidan Host, and, yeah, hi everyone. I’m Sayid Guest, and this topic is personal for me. I started out managing a handful of HMOs, and—um—the classic scene was radiators roaring away with nobody home, windows cracked wide open. You pay for the heat and watch it pour right out. Not fun. There’s also that line we can’t cross: you can’t lock tenants out of controls. So I began tinkering, looking for a way to respect tenant comfort while, uh, controlling the waste. The device I built evolved from a scrappy prototype to a solid system, and the surprising bit was how immediate the savings were. Like, right from the first month, bills were trending down hard, and I thought, wow, landlords need something like this across the board.",
+      "start": 38.75,
+      "end": 77.88
+    },
+    {
+      "speaker": 0,
+      "text": "Exactly. And, actually, on the software side, we’ve seen a similar DIY-to-pro evolution. Rent Byte started because landlords were fed up with juggling spreadsheets, emails, and random listings. It’s, um, designed by landlords with tenants in mind—trust and transparency baked in. With Rent Byte, you can push your property ads out fast, track leads, run checks, and glide folks through the onboarding without the headache. Then, when people move in, it doesn’t stop; you’ve got in-app maintenance reporting, job tracking, and clear timelines, so tenants don’t feel, you know, ignored. Tea Brew listeners keep telling us the pain isn’t just finding residents, it’s keeping the whole machine humming. And if you pair that workflow with smart heating controls, you’re hitting cost, comfort, and communication all in one go.",
+      "start": 77.92,
+      "end": 115.18
+    },
+    {
+      "speaker": 2,
+      "text": "Hold on, can I jump in? So, um, I manage two HMOs over by Meeter Street, and the energy bills last winter were brutal. Everyone online was shouting tips—weather compensation, schedule tweaks, the whole shebang—but tenants still cranked it up when they felt chilly. We can’t, like, lock the thermostat, right? So how does your approach, Sayid Guest, keep the system fair? Tenants get access, but landlords don’t get burned—haha, terrible pun. Does it do occupancy sensing, or is it just a timer? And how do you stop the classic boost button from becoming a permanent ‘on’?",
+      "start": 115.19,
+      "end": 147.56
+    },
+    {
+      "speaker": 1,
+      "text": "Great questions. So it’s a combo. Think of it as a comfort-first schedule with protections. You set reasonable heating windows—morning and evening, say—and tenants can press a boost for extra time. But the boost is capped and resets, so it won’t, uh, run the boiler all day. For empty rooms and the, you know, window-wide-open scenario, if you add sensors, the system detects rapid drops or no movement and tapers the heat until conditions make sense. It’s not about denying warmth; it’s about stopping waste that literally no one benefits from. With HMOs you need communal logic—landings and kitchens matter—so multi-zone control helps keep spaces balanced. And landlords, like me, get analytics: you can see where energy is leaking and fine-tune settings. When I first tested it, I saw 30–50% reductions. That’s not a promise for every building—each setup is its own puzzle—but the pattern has been, um, consistently strong.",
+      "start": 147.56,
+      "end": 193.57
+    },
+    {
+      "speaker": 0,
+      "text": "Yeah, that resonates. I walked into a shared house once—near Meeter Street—and the thermostat was at 29, like sauna levels, with the window propped open. The tenants weren’t being malicious, they were just coping: drafty room, quick fix, crank the dial. So, as Tea Brew keeps saying, we need systems that, uh, balance tenant agency with sensible guardrails. And the nice thing with Rent Byte is that it keeps the conversation flowing. Tenants can raise an issue through maintenance reporting, and as soon as the ticket is created, the timeline starts. If it’s a cold spot or a radiator fault, you’ve got the history at your fingertips. That way, you don’t blame behavior when it’s actually a hardware problem.",
+      "start": 193.57,
+      "end": 225.73
+    },
+    {
+      "speaker": 2,
+      "text": "Okay, so picture this: I advertise a new room, get flooded with inquiries, then I’m buried in emails. The software side—like Rent Byte—would pull those leads into a pipeline, right? And if someone mentions the room feels chilly during a viewing, I could, uh, flag that straight away. Here’s the kicker: can your device talk to the platform? Like, if there’s a temp anomaly or a sensor alert, could it auto-create a maintenance ticket so I don’t miss it? And, sorry, I’m thinking out loud here, but could Rent Byte show those energy graphs inside the tenant portal without freaking people out—more like, you know, helpful nudges than lecturing?",
+      "start": 225.74,
+      "end": 258.74
+    },
+    {
+      "speaker": 1,
+      "text": "Haha, I love the way you think. Yes, integration is the future. We’ve built an API so platforms like Rent Byte can pull summary data—no one needs to drown in charts, just the, um, useful stuff. For example, you can surface a gentle insight: “Heating is already scheduled; boost is available for 30 minutes.” Tenants see options, not rules. If a sensor flags a stuck valve or a window open for ages, Rent Byte can spin up a maintenance ticket, assign it, and track the fix. And because Tea Brew listeners keep asking about transparency, we’ve found tenants appreciate seeing that there’s a fair schedule in place. By the way, thanks again for the invite, Aidan Host; getting this across without jargon is half the battle.",
+      "start": 258.74,
+      "end": 292.52
+    },
+    {
+      "speaker": 0,
+      "text": "Totally. And, uh, we’ve noticed another benefit: advertising feels cleaner when you can tell prospects, right up front, that the home uses sensible heating with tenant-controlled boosts. It sounds small, but it sets expectations and, you know, avoids future friction. On the admin side, Rent Byte keeps everything documented—from viewing notes to audit trails—so if someone on Meeter Street says their radiator’s been weird for weeks, you can point to the timeline and fix history quickly. Um, we’ve had a bunch of landlords message Tea Brew after implementing this kind of setup, saying the combo of software plus smart heating saved money and, honestly, reduced arguments. That’s the vibe we want.",
+      "start": 292.53,
+      "end": 323.63
+    },
+    {
+      "speaker": 2,
+      "text": "Same here. To wrap, uh, I’m thinking: start with clear, humane policies, pair them with tech that respects tenant comfort, and keep the communication channel open. Aidan Host, if folks want to try Rent Byte, what’s the first step? And, Sayid Guest, for the device, is there like a starter kit guide so, um, non-tech landlords don’t panic? Maybe put links under the episode—sorry—under Tea Brew show notes. I’ve got two more rooms to fill near Meeter Street, and it’d be great to kick this off before the cold snaps hit again.",
+      "start": 323.64,
+      "end": 351.73
+    }
+  ],
+  "customized_context": [
+    "Tea Brew",
+    "Aiden Host",
+    "Saeed Guest",
+    "Rent Byte",
+    "The property is located near Meter Street."
+  ]
+}
\ No newline at end of file
diff --git a/finetuning-asr/toy_dataset/0.mp3 b/finetuning-asr/toy_dataset/0.mp3
new file mode 100755
index 0000000..a994394
Binary files /dev/null and b/finetuning-asr/toy_dataset/0.mp3 differ
diff --git a/finetuning-asr/toy_dataset/1.json b/finetuning-asr/toy_dataset/1.json
new file mode 100755
index 0000000..1cdf118
--- /dev/null
+++ b/finetuning-asr/toy_dataset/1.json
@@ -0,0 +1,79 @@
+{
+  "audio_duration": 328.26666666666665,
+  "audio_path": "1.mp3",
+  "segments": [
+    {
+      "speaker": 0,
+      "text": "Welcome back to our Youth Month special. Um, before we dive in, Tandi, you ready? We’re honoring young folks who, like, you know, shook the ground. The day itself marks those anti–language policy protests in the 70s—students across the country, campuses, townships—standing up. And today, we’re asking you all on the WhatsApp line to share who inspired you: a teacher, a cousin, someone you admire. Also, we’ll talk about campaigns like Crown Wrights, because identity’s not cosmetic; it’s, uh, core.",
+      "start": 0.0,
+      "end": 33.15
+    },
+    {
+      "speaker": 1,
+      "text": "Yeah, yeah, I’m totally in, haha. Thanks, Leila. So, um, I’ve been thinking about how the current youth still carry heavy stuff—joblessness, violence against women, and, heartbreakingly, queer kids being targeted. It’s, like, gutting. Yet they keep going. I mean, this is Tandi speaking from the heart here. The person I want to celebrate is Zahra. She stood up at Coyl High when policies tried to tame her natural hair, and she, uh, didn’t flinch. She even wrote a children’s story as part of Crown Wrights, to help little ones see their coils as power instead of problem.",
+      "start": 33.24,
+      "end": 71.33
+    },
+    {
+      "speaker": 2,
+      "text": "Wow, okay, that hits hard. Zahra at Coyl High—I remember, like, seeing clips where she just, you know, stood there calmly while adults were telling her to “fix” herself. It gave me chills. Identity isn’t some minor detail; it’s a major, uh, principle—wait, I always mix that with principal, haha. Anyway, the bravery at Crown Wrights events has ripple effects. And Leila, you’ve talked about how hair, especially coily textures, can be policed as a way to, um, shrink someone’s confidence.",
+      "start": 71.52,
+      "end": 102.9
+    },
+    {
+      "speaker": 0,
+      "text": "Exactly! And thanks, Tandi, for, like, naming the hard stuff. The way that activist reframed “acceptable” appearance shows young people don’t need permission to be whole. When you own your look, you walk into boardrooms you never imagined. You, uh, sit at tables, you speak up. It’s the same energy that wins pageants and policies—like, a confidence that shifts rooms. And for everyone listening, send a voice note about your own champion and share stories about reclaiming a “crown.”",
+      "start": 102.9,
+      "end": 132.14
+    },
+    {
+      "speaker": 1,
+      "text": "Hold on—just to paint the scene. At Coyl High, Zahra didn’t wait for a senior to intervene; she, um, chose the moment. No hesitation. She knew the risk: future opportunities, social backlash, being labeled “difficult.” And she still stood up. For me, Tandi, that shows the lesson the old protests tried to teach—use your voice now, because silence, like, steals time. And you always say confidence is contagious; when one girl lifts her chin, dozens follow.",
+      "start": 132.17,
+      "end": 161.49
+    },
+    {
+      "speaker": 2,
+      "text": "Yeah, yes—exactly. And, mm, the thing that gets me is the dignity piece. People act like hair is trivial, but for Black girls, policing curls is a way to control presence. That campaign reframes it: from “problem” to “pride.” When that flip happens, you start applying for roles you thought were off-limits, you choose your course, you don’t break—uh, brake—just because someone says you don’t fit. And, you mentioned queer youth; the intersections matter. That stand speaks beyond curls, into, like, the right to be fully yourself.",
+      "start": 161.52,
+      "end": 196.53
+    },
+    {
+      "speaker": 0,
+      "text": "Speaking of which, Leila, could you share a listener shout-out? We got a bunch of WhatsApps about teachers, neighbors, and, uh, nurses who kept kids afloat. Also there’s a note praising the book in the Crown Wrights series—apparently the illustrations made a little one feel seen at school.",
+      "start": 196.91,
+      "end": 212.93
+    },
+    {
+      "speaker": 2,
+      "text": "This is Leila—oh yeah, totally. So, um, there’s a message from a parent who says their child used to hide her coils under a cap, and after reading that Crown Wrights story she walked into class with no fear. And another listener’s like, “my teacher changed my life,” which—haha—yes. The past protests weren’t just about language policy; they set a template: organize, insist, repeat. Honestly, Tandi, I feel like honoring Zahra honors every student who refuses to be edited at Coyl High-type schools, you know?",
+      "start": 213.59,
+      "end": 245.25
+    },
+    {
+      "speaker": 1,
+      "text": "Exactly! And, um, can I just say: elders sometimes tell us to wait for leaders to fix things. But the lesson, from then till now, is don’t outsource your voice. The moment you speak, the course of the moment changes—like, the crowd pivots from idle to action. And for anyone feeling alone, tag us and the campaign; you’ll find community fast.",
+      "start": 245.25,
+      "end": 266.64
+    },
+    {
+      "speaker": 0,
+      "text": "Right, and, wow, we’ve barely scratched the surface. I want to circle back because someone asked if honoring hair is, um, frivolous compared to, say, jobs or safety. My answer: it’s linked. Coyl High-style rules are part of a bigger system that decides who is welcome. Undo that, and you expand opportunity. And if you’re looking for names, Leila suggested a list: local mentors, teacher heroes, youth organizers. That point about resilience—phew—needed.",
+      "start": 266.64,
+      "end": 297.57
+    },
+    {
+      "speaker": 2,
+      "text": "Mhm, and one more thought, then we’ll wrap. Zahra didn’t become a symbol because she wanted fame; she became one because she refused to break. That stubborn joy? It’s, like, fuel. I’m Leila, and I’m grateful we still celebrate that day in June as a reminder to keep going. To every listener—send your piece, um, peace—haha—on the line. And thanks, Tandi, for bringing this story, and thanks to campaigns like Crown Wrights for keeping the flame on.",
+      "start": 297.78,
+      "end": 328.27
+    }
+  ],
+  "customized_context": [
+    "Thandie",
+    "Leila",
+    "Zara",
+    "The story takes place at Coyle High school.",
+    "Crown Rites is a campaign for hair rights."
+  ]
+}
\ No newline at end of file
diff --git a/finetuning-asr/toy_dataset/1.mp3 b/finetuning-asr/toy_dataset/1.mp3
new file mode 100755
index 0000000..fa6a27e
Binary files /dev/null and b/finetuning-asr/toy_dataset/1.mp3 differ
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 0000000..ea0f837
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,56 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "vibevoice"
+version = "1.0.0"
+authors = [
+  { name="vibevoice team", email="VibeVoice@microsoft.com" },
+]
+description = "Open-Source Frontier Voice AI."
+readme = "README.md"
+requires-python = ">=3.10"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    # "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+]
+dependencies = [
+    "torch",
+    "transformers>=4.51.3,<5.0.0",
+    "accelerate",
+    "llvmlite>=0.40.0",
+    "numba>=0.57.0",
+    "diffusers",
+    "tqdm",
+    "numpy",
+    "scipy",
+    "librosa",
+    "ml-collections",
+    "absl-py",
+    "gradio",
+    "av",
+    "aiortc",
+    "uvicorn[standard]",
+    "fastapi",
+    "pydub",
+    "requests",
+]
+
+
+[project.optional-dependencies]
+streamingtts = [
+  "transformers==4.51.3", 
+]
+
+[project.entry-points."vllm.general_plugins"]
+vibevoice = "vllm_plugin:register_vibevoice"
+
+[project.urls]
+"Homepage" = "https://github.com/microsoft/VibeVoice"
+"Bug Tracker" = "https://github.com/microsoft/VibeVoice/issues"
+
+[tool.setuptools.packages.find]
+where = ["."]
+include = ["vibevoice*", "vllm_plugin*"]
diff --git a/vibevoice/__init__.py b/vibevoice/__init__.py
new file mode 100644
index 0000000..b00e91a
--- /dev/null
+++ b/vibevoice/__init__.py
@@ -0,0 +1,16 @@
+# vibevoice/__init__.py
+from vibevoice.modular import (
+    VibeVoiceStreamingForConditionalGenerationInference,
+    VibeVoiceStreamingConfig,
+)
+from vibevoice.processor import (
+    VibeVoiceStreamingProcessor,
+    VibeVoiceTokenizerProcessor,
+)
+
+__all__ = [
+    "VibeVoiceStreamingForConditionalGenerationInference",
+    "VibeVoiceStreamingConfig",
+    "VibeVoiceStreamingProcessor",
+    "VibeVoiceTokenizerProcessor",
+]
\ No newline at end of file
diff --git a/vibevoice/configs/qwen2.5_1.5b_64k.json b/vibevoice/configs/qwen2.5_1.5b_64k.json
new file mode 100644
index 0000000..febd05c
--- /dev/null
+++ b/vibevoice/configs/qwen2.5_1.5b_64k.json
@@ -0,0 +1,112 @@
+{
+  "_attn_implementation_autoset": true,
+  "acoustic_vae_dim": 64,
+  "acoustic_tokenizer_config": {
+    "causal": true,
+    "channels": 1,
+    "conv_bias": true,
+    "conv_norm": "none",
+    "corpus_normalize": 0.0,
+    "decoder_depths": null,
+    "decoder_n_filters": 32,
+    "decoder_ratios": [
+      8,
+      5,
+      5,
+      4,
+      2,
+      2
+    ],
+    "disable_last_norm": true,
+    "encoder_depths": "3-3-3-3-3-3-8",
+    "encoder_n_filters": 32,
+    "encoder_ratios": [
+      8,
+      5,
+      5,
+      4,
+      2,
+      2
+    ],
+    "fix_std": 0.5,
+    "layer_scale_init_value": 1e-06,
+    "layernorm": "RMSNorm",
+    "layernorm_elementwise_affine": true,
+    "layernorm_eps": 1e-05,
+    "mixer_layer": "depthwise_conv",
+    "model_type": "vibepod_acoustic_tokenizer",
+    "pad_mode": "constant",
+    "std_dist_type": "gaussian",
+    "vae_dim": 64,
+    "weight_init_value": 0.01
+  },
+  "decoder_config": {
+    "attention_dropout": 0.0,
+    "hidden_act": "silu",
+    "hidden_size": 1536,
+    "initializer_range": 0.02,
+    "intermediate_size": 8960,
+    "max_position_embeddings": 65536,
+    "max_window_layers": 28,
+    "model_type": "qwen2",
+    "num_attention_heads": 12,
+    "num_hidden_layers": 28,
+    "num_key_value_heads": 2,
+    "rms_norm_eps": 1e-06,
+    "rope_scaling": null,
+    "rope_theta": 1000000.0,
+    "sliding_window": null,
+    "tie_word_embeddings": true,
+    "torch_dtype": "bfloat16",
+    "use_cache": true,
+    "use_sliding_window": false,
+    "vocab_size": 151936
+  },
+  "diffusion_head_config": {
+    "ddpm_batch_mul": 4,
+    "ddpm_beta_schedule": "cosine",
+    "ddpm_num_inference_steps": 20,
+    "ddpm_num_steps": 1000,
+    "diffusion_type": "ddpm",
+    "head_ffn_ratio": 3.0,
+    "head_layers": 4,
+    "hidden_size": 1536,
+    "latent_size": 64,
+    "model_type": "vibepod_diffusion_head",
+    "prediction_type": "v_prediction",
+    "rms_norm_eps": 1e-05,
+    "speech_vae_dim": 64
+  },
+  "model_type": "vibepod",
+  "semantic_tokenizer_config": {
+    "causal": true,
+    "channels": 1,
+    "conv_bias": true,
+    "conv_norm": "none",
+    "corpus_normalize": 0.0,
+    "disable_last_norm": true,
+    "encoder_depths": "3-3-3-3-3-3-8",
+    "encoder_n_filters": 32,
+    "encoder_ratios": [
+      8,
+      5,
+      5,
+      4,
+      2,
+      2
+    ],
+    "fix_std": 0,
+    "layer_scale_init_value": 1e-06,
+    "layernorm": "RMSNorm",
+    "layernorm_elementwise_affine": true,
+    "layernorm_eps": 1e-05,
+    "mixer_layer": "depthwise_conv",
+    "model_type": "vibepod_semantic_tokenizer",
+    "pad_mode": "constant",
+    "std_dist_type": "none",
+    "vae_dim": 128,
+    "weight_init_value": 0.01
+  },
+  "semantic_vae_dim": 128,
+  "torch_dtype": "bfloat16"
+}
diff --git a/vibevoice/configs/qwen2.5_7b_32k.json b/vibevoice/configs/qwen2.5_7b_32k.json
new file mode 100644
index 0000000..d39952c
--- /dev/null
+++ b/vibevoice/configs/qwen2.5_7b_32k.json
@@ -0,0 +1,113 @@
+{
+  "_attn_implementation_autoset": true,
+  "acoustic_vae_dim": 64,
+  "acoustic_tokenizer_config": {
+    "causal": true,
+    "channels": 1,
+    "conv_bias": true,
+    "conv_norm": "none",
+    "corpus_normalize": 0.0,
+    "decoder_depths": null,
+    "decoder_n_filters": 32,
+    "decoder_ratios": [
+      8,
+      5,
+      5,
+      4,
+      2,
+      2
+    ],
+    "disable_last_norm": true,
+    "encoder_depths": "3-3-3-3-3-3-8",
+    "encoder_n_filters": 32,
+    "encoder_ratios": [
+      8,
+      5,
+      5,
+      4,
+      2,
+      2
+    ],
+    "fix_std": 0.5,
+    "layer_scale_init_value": 1e-06,
+    "layernorm": "RMSNorm",
+    "layernorm_elementwise_affine": true,
+    "layernorm_eps": 1e-05,
+    "mixer_layer": "depthwise_conv",
+    "model_type": "vibepod_acoustic_tokenizer",
+    "pad_mode": "constant",
+    "std_dist_type": "gaussian",
+    "vae_dim": 64,
+    "weight_init_value": 0.01
+  },
+  "decoder_config": {
+    "attention_dropout": 0.0,
+    "hidden_act": "silu",
+    "hidden_size": 3584,
+    "initializer_range": 0.02,
+    "intermediate_size": 18944,
+    "max_position_embeddings": 32768,
+    "max_window_layers": 28,
+    "model_type": "qwen2",
+    "num_attention_heads": 28,
+    "num_hidden_layers": 28,
+    "num_key_value_heads": 4,
+    "rms_norm_eps": 1e-06,
+    "rope_theta": 1000000.0,
+    "sliding_window": null,
+    "tie_word_embeddings": false,
+    "torch_dtype": "bfloat16",
+    "transformers_version": "4.40.1",
+    "use_cache": true,
+    "use_mrope": false,
+    "use_sliding_window": false,
+    "vocab_size": 152064
+  },
+  "diffusion_head_config": {
+    "ddpm_batch_mul": 4,
+    "ddpm_beta_schedule": "cosine",
+    "ddpm_num_inference_steps": 20,
+    "ddpm_num_steps": 1000,
+    "diffusion_type": "ddpm",
+    "head_ffn_ratio": 3.0,
+    "head_layers": 4,
+    "hidden_size": 3584,
+    "latent_size": 64,
+    "model_type": "vibepod_diffusion_head",
+    "prediction_type": "v_prediction",
+    "rms_norm_eps": 1e-05,
+    "speech_vae_dim": 64
+  },
+  "model_type": "vibepod",
+  "semantic_tokenizer_config": {
+    "causal": true,
+    "channels": 1,
+    "conv_bias": true,
+    "conv_norm": "none",
+    "corpus_normalize": 0.0,
+    "disable_last_norm": true,
+    "encoder_depths": "3-3-3-3-3-3-8",
+    "encoder_n_filters": 32,
+    "encoder_ratios": [
+      8,
+      5,
+      5,
+      4,
+      2,
+      2
+    ],
+    "fix_std": 0,
+    "layer_scale_init_value": 1e-06,
+    "layernorm": "RMSNorm",
+    "layernorm_elementwise_affine": true,
+    "layernorm_eps": 1e-05,
+    "mixer_layer": "depthwise_conv",
+    "model_type": "vibepod_semantic_tokenizer",
+    "pad_mode": "constant",
+    "std_dist_type": "none",
+    "vae_dim": 128,
+    "weight_init_value": 0.01
+  },
+  "semantic_vae_dim": 128,
+  "torch_dtype": "bfloat16"
+}
diff --git a/vibevoice/modular/__init__.py b/vibevoice/modular/__init__.py
new file mode 100644
index 0000000..c92b1a7
--- /dev/null
+++ b/vibevoice/modular/__init__.py
@@ -0,0 +1,14 @@
+# vibevoice/modular/__init__.py
+from .modeling_vibevoice_streaming_inference import VibeVoiceStreamingForConditionalGenerationInference
+from .configuration_vibevoice_streaming import VibeVoiceStreamingConfig
+from .modeling_vibevoice_streaming import VibeVoiceStreamingModel, VibeVoiceStreamingPreTrainedModel
+from .streamer import AudioStreamer, AsyncAudioStreamer
+
+__all__ = [
+    "VibeVoiceStreamingForConditionalGenerationInference",
+    "VibeVoiceStreamingConfig",
+    "VibeVoiceStreamingModel",
+    "VibeVoiceStreamingPreTrainedModel",
+    "AudioStreamer",
+    "AsyncAudioStreamer",
+]
\ No newline at end of file
diff --git a/vibevoice/modular/configuration_vibevoice.py b/vibevoice/modular/configuration_vibevoice.py
new file mode 100644
index 0000000..1845113
--- /dev/null
+++ b/vibevoice/modular/configuration_vibevoice.py
@@ -0,0 +1,406 @@
+""" VibeVoice_AcousticTokenizer model configuration"""
+
+from typing import Dict, List, Optional, Tuple
+
+import torch
+from transformers.configuration_utils import PretrainedConfig 
+from transformers.utils import logging
+
+from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
+
+logger = logging.get_logger(__name__)
+
+
+def _convert_dtype_to_string(config_dict: dict) -> dict:
+    """
+    Convert torch.dtype objects to their string representation for JSON serialization.
+    
+    This fixes the "Object of type dtype is not JSON serializable" error that occurs
+    when transformers tries to log/serialize the config with torch_dtype as a torch.dtype object.
+    
+    See: https://github.com/microsoft/VibeVoice/issues/199
+    """
+    if "torch_dtype" in config_dict and config_dict["torch_dtype"] is not None:
+        dtype = config_dict["torch_dtype"]
+        if isinstance(dtype, torch.dtype):
+            # Convert torch.dtype to string (e.g., torch.bfloat16 -> "bfloat16")
+            config_dict["torch_dtype"] = str(dtype).replace("torch.", "")
+    return config_dict
+
+
+class VibeVoiceAcousticTokenizerConfig(PretrainedConfig):
+    model_type = "vibevoice_acoustic_tokenizer"
+
+    def __init__(
+        self,
+        channels: int = 1,
+        corpus_normalize: float = 0.0,
+        causal: bool = True,
+        vae_dim: int = 64,
+        fix_std: float = 0.5,
+        std_dist_type: str = 'gaussian',
+        # common 
+        mixer_layer: str = 'depthwise_conv',
+        conv_norm: str = 'none',
+        pad_mode: str = 'constant',
+        disable_last_norm: bool = True,
+        layernorm: str = 'RMSNorm',
+        layernorm_eps: float = 1e-5,
+        layernorm_elementwise_affine: bool = True,
+        conv_bias: bool = True,
+        layer_scale_init_value: float = 1e-6,
+        weight_init_value: float = 1e-2,
+        # encoder specific
+        encoder_n_filters: int = 32,
+        encoder_ratios: Optional[List[int]] = [8,5,5,4,2,2],
+        encoder_depths: str = "3-3-3-3-3-3-8",
+        # decoder specific
+        decoder_n_filters: int = 32,
+        decoder_ratios: Optional[List[int]] = None, # if None, same as encoder
+        decoder_depths: Optional[str] = None,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.channels = channels
+        self.corpus_normalize = corpus_normalize
+        self.causal = causal
+        self.vae_dim = vae_dim
+        self.fix_std = fix_std
+        self.std_dist_type = std_dist_type
+        
+        # common parameters
+        self.conv_norm = conv_norm
+        self.pad_mode = pad_mode
+        self.layernorm_eps = layernorm_eps
+        self.disable_last_norm = disable_last_norm
+        self.layernorm = layernorm
+        self.layernorm_elementwise_affine = layernorm_elementwise_affine
+        self.conv_bias = conv_bias
+        self.layer_scale_init_value = layer_scale_init_value
+        self.weight_init_value = weight_init_value
+        self.mixer_layer = mixer_layer
+
+        # encoder specific parameters
+        self.encoder_n_filters = encoder_n_filters
+        self.encoder_ratios = encoder_ratios
+        self.encoder_depths = encoder_depths
+        
+        # decoder specific parameters
+        self.decoder_ratios = decoder_ratios if decoder_ratios is not None else encoder_ratios
+        self.decoder_n_filters = decoder_n_filters
+        self.decoder_depths = decoder_depths
+
+
+class VibeVoiceSemanticTokenizerConfig(PretrainedConfig):
+    model_type = "vibevoice_semantic_tokenizer"
+    
+    def __init__(
+        self,
+        channels: int = 1,
+        corpus_normalize: float = 0.0,
+        causal: bool = True,
+        vae_dim: int = 64,
+        fix_std: float = 0,
+        std_dist_type: str = 'none',
+        # common 
+        mixer_layer: str = 'depthwise_conv',
+        conv_norm: str = 'none',
+        pad_mode: str = 'constant',
+        disable_last_norm: bool = True,
+        layernorm: str = 'RMSNorm',
+        layernorm_eps: float = 1e-5,
+        layernorm_elementwise_affine: bool = True,
+        conv_bias: bool = True,
+        layer_scale_init_value: float = 1e-6,
+        weight_init_value: float = 1e-2,
+        # encoder specific
+        encoder_n_filters: int = 32,
+        encoder_ratios: Optional[List[int]] = [8,5,5,4,2,2],
+        encoder_depths: str = "3-3-3-3-3-3-8",
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.channels = channels
+        self.corpus_normalize = corpus_normalize
+        self.causal = causal
+        self.vae_dim = vae_dim
+        self.fix_std = fix_std
+        self.std_dist_type = std_dist_type
+        
+        # common parameters
+        self.conv_norm = conv_norm
+        self.pad_mode = pad_mode
+        self.layernorm_eps = layernorm_eps
+        self.disable_last_norm = disable_last_norm
+        self.layernorm = layernorm
+        self.layernorm_elementwise_affine = layernorm_elementwise_affine
+        self.conv_bias = conv_bias
+        self.layer_scale_init_value = layer_scale_init_value
+        self.weight_init_value = weight_init_value
+        self.mixer_layer = mixer_layer
+
+        # encoder specific parameters
+        self.encoder_n_filters = encoder_n_filters
+        self.encoder_ratios = encoder_ratios
+        self.encoder_depths = encoder_depths
+        
+
+class VibeVoiceDiffusionHeadConfig(PretrainedConfig):
+    model_type = "vibevoice_diffusion_head"
+
+    def __init__(
+        self,
+        hidden_size=768,
+        head_layers=4,
+        head_ffn_ratio=3.0,
+        rms_norm_eps=1e-5,
+        latent_size=64,
+        speech_vae_dim=None,
+        prediction_type="v_prediction",
+        diffusion_type="ddpm",
+        ddpm_num_steps=1000,
+        ddpm_num_inference_steps=20,
+        ddpm_beta_schedule="cosine",
+        ddpm_batch_mul=4,
+        **kwargs
+    ):
+        self.hidden_size = hidden_size
+        self.head_layers = head_layers
+        self.head_ffn_ratio = head_ffn_ratio
+        self.rms_norm_eps = rms_norm_eps
+        self.latent_size = latent_size
+        self.speech_vae_dim = speech_vae_dim
+        self.prediction_type = prediction_type
+        self.diffusion_type = diffusion_type
+        self.ddpm_num_steps = ddpm_num_steps
+        self.ddpm_num_inference_steps = ddpm_num_inference_steps
+        self.ddpm_beta_schedule = ddpm_beta_schedule
+        self.ddpm_batch_mul = ddpm_batch_mul
+        
+        super().__init__(**kwargs)
+
+class VibeVoiceConfig(PretrainedConfig):
+    model_type = "vibevoice"
+    is_composition = True
+    sub_configs = {
+        "acoustic_tokenizer_config": VibeVoiceAcousticTokenizerConfig, 
+        "semantic_tokenizer_config": VibeVoiceSemanticTokenizerConfig,
+        "decoder_config": Qwen2Config,
+        "diffusion_head_config": VibeVoiceDiffusionHeadConfig,
+    }
+    # keys_to_ignore_at_inference = ["past_key_values"]
+    # Default tensor parallel plan for base model `Qwen2`
+    base_model_tp_plan = {
+        "layers.*.self_attn.q_proj": "colwise",
+        "layers.*.self_attn.k_proj": "colwise",
+        "layers.*.self_attn.v_proj": "colwise",
+        "layers.*.self_attn.o_proj": "rowwise",
+        "layers.*.mlp.gate_proj": "colwise",
+        "layers.*.mlp.up_proj": "colwise",
+        "layers.*.mlp.down_proj": "rowwise",
+    }
+    
+    def __init__(
+        self,
+        acoustic_tokenizer_config=None,
+        semantic_tokenizer_config=None,
+        decoder_config=None,
+        diffusion_head_config=None,
+        **kwargs
+    ):
+
+        # kwargs["_attn_implementation"] = "flash_attention_2"
+        kwargs["_attn_implementation_autoset"] = False 
+
+        if acoustic_tokenizer_config is None:
+            self.acoustic_tokenizer_config = self.sub_configs["acoustic_tokenizer_config"]()
+        elif isinstance(acoustic_tokenizer_config, dict):
+            acoustic_tokenizer_config["model_type"] = "vibevoice_acoustic_tokenizer"
+            self.acoustic_tokenizer_config = self.sub_configs["acoustic_tokenizer_config"](**acoustic_tokenizer_config)
+        elif isinstance(acoustic_tokenizer_config, VibeVoiceAcousticTokenizerConfig):
+            # If an instance of the config class is provided
+            self.acoustic_tokenizer_config = acoustic_tokenizer_config
+
+        if semantic_tokenizer_config is None:
+            self.semantic_tokenizer_config = self.sub_configs["semantic_tokenizer_config"]()
+        elif isinstance(semantic_tokenizer_config, dict):
+            semantic_tokenizer_config["model_type"] = "vibevoice_semantic_tokenizer"
+            self.semantic_tokenizer_config = self.sub_configs["semantic_tokenizer_config"](**semantic_tokenizer_config)
+        elif isinstance(semantic_tokenizer_config, VibeVoiceSemanticTokenizerConfig):
+            # If an instance of the config class is provided
+            self.semantic_tokenizer_config = semantic_tokenizer_config
+
+        if decoder_config is None:
+            self.decoder_config = self.sub_configs["decoder_config"]()
+        elif isinstance(decoder_config, dict):
+            # If a dictionary is provided, instantiate the config class with it
+            # self.decoder_config = self.sub_configs["decoder_config"](**decoder_config)
+            if decoder_config.get("model_type", '') == "qwen2":
+                self.decoder_config = Qwen2Config(**decoder_config)
+            else:
+                raise ValueError(f"Unsupported decoder model type: {decoder_config.get('model_type', '')}")
+        elif isinstance(decoder_config, (Qwen2Config,)):
+            # If an instance of the config class is provided
+            self.decoder_config = decoder_config
+
+        if diffusion_head_config is None:
+            self.diffusion_head_config = self.sub_configs["diffusion_head_config"]()
+        elif isinstance(diffusion_head_config, dict):
+            diffusion_head_config["model_type"] = "vibevoice_diffusion_head"
+            self.diffusion_head_config = self.sub_configs["diffusion_head_config"](**diffusion_head_config)
+        elif isinstance(diffusion_head_config, VibeVoiceDiffusionHeadConfig):
+            # If an instance of the config class is provided
+            self.diffusion_head_config = diffusion_head_config
+
+        # other parameters
+        self.acoustic_vae_dim = getattr(self.acoustic_tokenizer_config, 'vae_dim', 64)
+        self.semantic_vae_dim = getattr(self.semantic_tokenizer_config, 'vae_dim', 128)
+
+        super().__init__(**kwargs)
+
+    def get_text_config(self, decoder=False):
+        """
+        Returns the text config for this model.
+        
+        vLLM uses this method to get the text configuration from multimodal models.
+        This allows vLLM to correctly determine hidden_size, num_attention_heads,
+        and other properties needed for memory profiling and model execution.
+        
+        For VibeVoice, the "text config" is the decoder_config (Qwen2Config).
+        
+        Args:
+            decoder: If True, return the decoder config (for encoder-decoder models).
+                    For VibeVoice, this is always the decoder_config.
+        
+        Returns:
+            The decoder configuration (Qwen2Config) which contains hidden_size, etc.
+        """
+        return self.decoder_config
+
+    def to_dict(self):
+        """
+        Override to_dict to handle torch.dtype serialization.
+        
+        Fixes: https://github.com/microsoft/VibeVoice/issues/199
+        """
+        output = super().to_dict()
+        return _convert_dtype_to_string(output)
+
+class VibeVoiceASRConfig(PretrainedConfig):
+    model_type = "vibevoice"
+    is_composition = True
+    sub_configs = {
+        "acoustic_tokenizer_config": VibeVoiceAcousticTokenizerConfig, 
+        "semantic_tokenizer_config": VibeVoiceSemanticTokenizerConfig,
+        "decoder_config": Qwen2Config,
+    }
+    # keys_to_ignore_at_inference = ["past_key_values"]
+    # Default tensor parallel plan for base model `Qwen2`
+    base_model_tp_plan = {
+        "layers.*.self_attn.q_proj": "colwise",
+        "layers.*.self_attn.k_proj": "colwise",
+        "layers.*.self_attn.v_proj": "colwise",
+        "layers.*.self_attn.o_proj": "rowwise",
+        "layers.*.mlp.gate_proj": "colwise",
+        "layers.*.mlp.up_proj": "colwise",
+        "layers.*.mlp.down_proj": "rowwise",
+    }
+    
+    def __init__(
+        self,
+        acoustic_tokenizer_config=None,
+        semantic_tokenizer_config=None,
+        decoder_config=None,
+        **kwargs
+    ):
+
+        # kwargs["_attn_implementation"] = "flash_attention_2"
+        kwargs["_attn_implementation_autoset"] = False 
+
+        if acoustic_tokenizer_config is None:
+            self.acoustic_tokenizer_config = self.sub_configs["acoustic_tokenizer_config"]()
+        elif isinstance(acoustic_tokenizer_config, dict):
+            acoustic_tokenizer_config["model_type"] = "vibevoice_acoustic_tokenizer"
+            self.acoustic_tokenizer_config = self.sub_configs["acoustic_tokenizer_config"](**acoustic_tokenizer_config)
+        elif isinstance(acoustic_tokenizer_config, VibeVoiceAcousticTokenizerConfig):
+            # If an instance of the config class is provided
+            self.acoustic_tokenizer_config = acoustic_tokenizer_config
+
+        if semantic_tokenizer_config is None:
+            self.semantic_tokenizer_config = self.sub_configs["semantic_tokenizer_config"]()
+        elif isinstance(semantic_tokenizer_config, dict):
+            semantic_tokenizer_config["model_type"] = "vibevoice_semantic_tokenizer"
+            self.semantic_tokenizer_config = self.sub_configs["semantic_tokenizer_config"](**semantic_tokenizer_config)
+        elif isinstance(semantic_tokenizer_config, VibeVoiceSemanticTokenizerConfig):
+            # If an instance of the config class is provided
+            self.semantic_tokenizer_config = semantic_tokenizer_config
+
+        if decoder_config is None:
+            self.decoder_config = self.sub_configs["decoder_config"]()
+        elif isinstance(decoder_config, dict):
+            # If a dictionary is provided, instantiate the config class with it
+            # self.decoder_config = self.sub_configs["decoder_config"](**decoder_config)
+            if decoder_config.get("model_type", '') == "qwen2":
+                self.decoder_config = Qwen2Config(**decoder_config)
+            else:
+                raise ValueError(f"Unsupported decoder model type: {decoder_config.get('model_type', '')}")
+        elif isinstance(decoder_config, Qwen2Config):
+            # If an instance of the config class is provided
+            self.decoder_config = decoder_config
+
+        # other parameters
+        self.acoustic_vae_dim = getattr(self.acoustic_tokenizer_config, 'vae_dim', 64)
+        self.semantic_vae_dim = getattr(self.semantic_tokenizer_config, 'vae_dim', 128)
+
+        super().__init__(**kwargs)
+
+    def to_dict(self):
+        """
+        Override to_dict to handle torch.dtype serialization.
+        
+        Fixes: https://github.com/microsoft/VibeVoice/issues/199
+        """
+        output = super().to_dict()
+        return _convert_dtype_to_string(output)
+
+    def get_text_config(self, decoder: bool = False):
+        """Return the text (decoder) config for generation."""
+        return self.decoder_config
+
+    @property
+    def vocab_size(self):
+        """Return vocab_size from decoder config for generation compatibility."""
+        return self.decoder_config.vocab_size
+
+    @property
+    def num_attention_heads(self):
+        """Return num_attention_heads from decoder config for Ulysses SP compatibility."""
+        return self.decoder_config.num_attention_heads
+    
+    @property
+    def num_key_value_heads(self):
+        """Return num_key_value_heads from decoder config for Ulysses SP compatibility."""
+        return self.decoder_config.num_key_value_heads
+    
+    @property
+    def hidden_size(self):
+        """Return hidden_size from decoder config for model compatibility."""
+        return self.decoder_config.hidden_size
+    
+    @property
+    def num_hidden_layers(self):
+        """Return num_hidden_layers from decoder config for Ulysses SP compatibility."""
+        return self.decoder_config.num_hidden_layers
+    
+    @property
+    def head_dim(self):
+        """Return head_dim from decoder config for Ulysses SP compatibility."""
+        return getattr(self.decoder_config, 'head_dim', self.hidden_size // self.num_attention_heads)
+    
+__all__ = [
+    "VibeVoiceAcousticTokenizerConfig", 
+    "VibeVoiceSemanticTokenizerConfig", 
+    "VibeVoiceDiffusionHeadConfig", 
+    "VibeVoiceConfig",
+    "VibeVoiceASRConfig"
+]
\ No newline at end of file
diff --git a/vibevoice/modular/configuration_vibevoice_streaming.py b/vibevoice/modular/configuration_vibevoice_streaming.py
new file mode 100644
index 0000000..2bd9d6e
--- /dev/null
+++ b/vibevoice/modular/configuration_vibevoice_streaming.py
@@ -0,0 +1,104 @@
+""" VibeVoice Streaming model configuration"""
+
+import torch
+from transformers.configuration_utils import PretrainedConfig 
+from transformers.utils import logging
+
+from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
+
+from .configuration_vibevoice import VibeVoiceAcousticTokenizerConfig, VibeVoiceDiffusionHeadConfig, _convert_dtype_to_string
+
+logger = logging.get_logger(__name__)
+
+
+class VibeVoiceStreamingConfig(PretrainedConfig):
+    model_type = "vibevoice_streaming"
+    is_composition = True
+    sub_configs = {
+        "acoustic_tokenizer_config": VibeVoiceAcousticTokenizerConfig,
+        "decoder_config": Qwen2Config,
+        "diffusion_head_config": VibeVoiceDiffusionHeadConfig,
+    }
+    # keys_to_ignore_at_inference = ["past_key_values"]
+    # Default tensor parallel plan for base model `Qwen2`
+    base_model_tp_plan = {
+        "layers.*.self_attn.q_proj": "colwise",
+        "layers.*.self_attn.k_proj": "colwise",
+        "layers.*.self_attn.v_proj": "colwise",
+        "layers.*.self_attn.o_proj": "rowwise",
+        "layers.*.mlp.gate_proj": "colwise",
+        "layers.*.mlp.up_proj": "colwise",
+        "layers.*.mlp.down_proj": "rowwise",
+    }
+    
+    def __init__(
+        self,
+        acoustic_tokenizer_config=None,
+        decoder_config=None,
+        diffusion_head_config=None,
+        tts_backbone_num_hidden_layers=20,
+        **kwargs
+    ):
+
+        # kwargs["_attn_implementation"] = "flash_attention_2"
+        kwargs["_attn_implementation_autoset"] = False 
+
+        if acoustic_tokenizer_config is None:
+            self.acoustic_tokenizer_config = self.sub_configs["acoustic_tokenizer_config"]()
+        elif isinstance(acoustic_tokenizer_config, dict):
+            acoustic_tokenizer_config["model_type"] = "vibevoice_acoustic_tokenizer"
+            self.acoustic_tokenizer_config = self.sub_configs["acoustic_tokenizer_config"](**acoustic_tokenizer_config)
+        elif isinstance(acoustic_tokenizer_config, VibeVoiceAcousticTokenizerConfig):
+            # If an instance of the config class is provided
+            self.acoustic_tokenizer_config = acoustic_tokenizer_config
+
+        if decoder_config is None:
+            self.decoder_config = self.sub_configs["decoder_config"]()
+        elif isinstance(decoder_config, dict):
+            # If a dictionary is provided, instantiate the config class with it
+            # self.decoder_config = self.sub_configs["decoder_config"](**decoder_config)
+            if decoder_config.get("model_type", '') == "qwen2":
+                self.decoder_config = Qwen2Config(**decoder_config)
+            else:
+                raise ValueError(f"Unsupported decoder model type: {decoder_config.get('model_type', '')}")
+        elif isinstance(decoder_config, (Qwen2Config,)):
+            # If an instance of the config class is provided
+            self.decoder_config = decoder_config
+
+        if diffusion_head_config is None:
+            self.diffusion_head_config = self.sub_configs["diffusion_head_config"]()
+        elif isinstance(diffusion_head_config, dict):
+            diffusion_head_config["model_type"] = "vibevoice_diffusion_head"
+            self.diffusion_head_config = self.sub_configs["diffusion_head_config"](**diffusion_head_config)
+        elif isinstance(diffusion_head_config, VibeVoiceDiffusionHeadConfig):
+            # If an instance of the config class is provided
+            self.diffusion_head_config = diffusion_head_config
+
+        # other parameters
+        self.acoustic_vae_dim = getattr(self.acoustic_tokenizer_config, 'vae_dim', 64)
+        # The decoder of the model is divided into two components. The lower Transformer layers are only used for encoding text, while the upper Transformer layers are used for encoding text and generating speech. `tts_backbone_num_hidden_layers` indicates the number of upper layers used for TTS.
+        self.tts_backbone_num_hidden_layers = tts_backbone_num_hidden_layers
+
+        super().__init__(**kwargs)
+
+    def get_text_config(self, decoder=False):
+        """Returns the decoder config (required for transformers >= 4.57 cache compatibility)."""
+        return self.decoder_config
+
+    @property
+    def num_hidden_layers(self):
+        """Proxy to decoder_config.num_hidden_layers (required for transformers >= 4.57)."""
+        return self.decoder_config.num_hidden_layers
+
+    def to_dict(self):
+        """
+        Override to_dict to handle torch.dtype serialization.
+        
+        Fixes: https://github.com/microsoft/VibeVoice/issues/199
+        """
+        output = super().to_dict()
+        return _convert_dtype_to_string(output)
+
+__all__ = [
+    "VibeVoiceStreamingConfig"
+]
\ No newline at end of file
diff --git a/vibevoice/modular/modeling_vibevoice.py b/vibevoice/modular/modeling_vibevoice.py
new file mode 100644
index 0000000..a4ecbab
--- /dev/null
+++ b/vibevoice/modular/modeling_vibevoice.py
@@ -0,0 +1,496 @@
+# copied from https://github.com/vibevoice-community/VibeVoice/blob/main/vibevoice/modular/modeling_vibevoice.py
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Tuple, Union, Callable
+from tqdm import tqdm
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.distributed as dist
+
+from transformers.models.auto import AutoModel, AutoModelForCausalLM
+
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import CausalLMOutput, BaseModelOutputWithPast, ModelOutput
+from transformers.models.llama.modeling_llama import LlamaRMSNorm
+from transformers import modeling_utils
+from transformers.modeling_utils import PreTrainedModel
+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+from transformers.utils import logging
+
+
+from .modular_vibevoice_tokenizer import VibeVoiceTokenizerStreamingCache, VibeVoiceAcousticTokenizerModel, VibeVoiceSemanticTokenizerModel
+from .modular_vibevoice_diffusion_head import VibeVoiceDiffusionHead
+from vibevoice.schedule.dpm_solver import DPMSolverMultistepScheduler
+
+from .configuration_vibevoice import VibeVoiceConfig
+
+
+logger = logging.get_logger(__name__)
+
+if not hasattr(modeling_utils, "ALL_PARALLEL_STYLES") or modeling_utils.ALL_PARALLEL_STYLES is None:
+    modeling_utils.ALL_PARALLEL_STYLES = ["tp", "none", "colwise", "rowwise"]
+
+@dataclass
+class VibeVoiceCausalLMOutputWithPast(ModelOutput):
+    loss: Optional[torch.FloatTensor] = None
+    diffusion_loss: Optional[torch.FloatTensor] = None
+    speech_token_num: Optional[int] = None
+    logits: torch.FloatTensor = None
+    past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+
+
+@dataclass
+class VibeVoiceGenerationOutput(ModelOutput):
+    """
+    Output type for VibeVoice generation.
+    
+    Args:
+        sequences (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+            The generated sequences. 
+        speech_outputs (`List[torch.FloatTensor]`, *optional*):
+            List of generated speech waveforms or latents for each speech segment.
+    """
+    sequences: torch.LongTensor = None
+    speech_outputs: Optional[List[torch.FloatTensor]] = None
+
+
+class SpeechConnector(nn.Module):
+    def __init__(self, input_dim, output_dim):
+        super().__init__()
+        self.fc1 = nn.Linear(input_dim, output_dim)
+        self.norm = LlamaRMSNorm(output_dim, eps=1e-6)
+        self.fc2 = nn.Linear(output_dim, output_dim)
+
+    def forward(self, features, **kwargs):    
+        x = self.fc1(features)
+        x = self.norm(x)
+        x = self.fc2(x)
+        return x
+
+
+# @auto_docstring
+class VibeVoicePreTrainedModel(PreTrainedModel):
+    config_class = VibeVoiceConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _skip_keys_device_placement = "past_key_values"
+    _supports_cache_class = True
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _supports_quantized_cache = True
+    _supports_static_cache = True
+    _supports_attention_backend = True
+
+    def _init_weights(self, module):
+        if isinstance(module, VibeVoiceDiffusionHead):
+            module.initialize_weights()
+            return
+
+        # Use the language model's initializer_range if available
+        if hasattr(self.config, 'language_model_config') and hasattr(self.config.language_model_config, 'initializer_range'):
+            std = self.config.language_model_config.initializer_range
+        elif hasattr(self.config, 'decoder_config') and hasattr(self.config.decoder_config, 'initializer_range'):
+            std = self.config.decoder_config.initializer_range
+        else:
+            std = 0.02  # Default value
+            
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.LayerNorm):
+            module.weight.data.fill_(1.0)
+            module.bias.data.zero_()
+
+# @auto_docstring
+class VibeVoiceModel(VibeVoicePreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        
+        if hasattr(config, 'torch_dtype') and config.torch_dtype is not None:
+            if isinstance(config.torch_dtype, str):
+                dtype = getattr(torch, config.torch_dtype)
+            else:
+                dtype = config.torch_dtype
+        else:
+            dtype = torch.float32
+        
+        # Initialize Qwen2 model for language modeling
+        lm_config = config.decoder_config 
+        self.language_model = AutoModel.from_config(lm_config)
+        
+        # Initialize speech components if needed
+        self.acoustic_tokenizer = AutoModel.from_config(config.acoustic_tokenizer_config).to(dtype)
+        self.semantic_tokenizer = AutoModel.from_config(config.semantic_tokenizer_config).to(dtype)
+
+        self.acoustic_connector = SpeechConnector(config.acoustic_vae_dim, lm_config.hidden_size).to(dtype)
+        self.semantic_connector = SpeechConnector(config.semantic_vae_dim, lm_config.hidden_size).to(dtype)
+        
+        # Register scaling factors as buffers - use 1D tensors for FSDP compatibility
+        self.register_buffer('speech_scaling_factor', torch.tensor(float('nan')))  
+        self.register_buffer('speech_bias_factor', torch.tensor(float('nan')))
+
+        # Initialize prediction head for speech generation
+        self.prediction_head = AutoModel.from_config(config.diffusion_head_config).to(dtype)
+
+        # Initialize noise scheduler
+        self.noise_scheduler = DPMSolverMultistepScheduler(
+            num_train_timesteps=config.diffusion_head_config.ddpm_num_steps,
+            beta_schedule=config.diffusion_head_config.ddpm_beta_schedule,
+            prediction_type=config.diffusion_head_config.prediction_type
+        )
+    
+    def get_input_embeddings(self):
+        if hasattr(self.language_model, 'embed_tokens'):
+            # If the language model has an embed_tokens attribute, return it
+            return self.language_model.embed_tokens
+        
+        for name, attr in self.language_model.fullmap.items(): # parallel by nnscaler, the name is changed
+            if attr.orig_name == 'embed_tokens.weight':
+                return getattr(self.language_model, name)
+        assert False, 'should not arrive here'
+
+    def set_input_embeddings(self, value):
+        self.language_model.embed_tokens = value
+    
+    def set_speech_tokenizers(self, acoustic_tokenizer=None, semantic_tokenizer=None):
+        """Set the speech tokenizers used for encoding and decoding speech."""
+        self.acoustic_tokenizer = acoustic_tokenizer
+        self.semantic_tokenizer = semantic_tokenizer
+        
+        # Reset the encoder to evaluation mode
+        if self.acoustic_tokenizer is not None:
+            self.acoustic_tokenizer.eval()
+            
+        if self.semantic_tokenizer is not None:
+            self.semantic_tokenizer.eval()
+    
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        
+        # Forward through language model
+        outputs = self.language_model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            cache_position=cache_position,
+            **kwargs,
+        )
+        
+        if not return_dict:
+            return outputs
+            
+        return BaseModelOutputWithPast(
+            last_hidden_state=outputs.last_hidden_state,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+class VibeVoiceForConditionalGeneration(VibeVoicePreTrainedModel):
+    _tied_weights_keys = ["lm_head.weight"]
+    _tp_plan = {"lm_head": "colwise_rep"}
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = VibeVoiceModel(config)
+        self.vocab_size = config.decoder_config.vocab_size
+        self.lm_head = nn.Linear(config.decoder_config.hidden_size, self.vocab_size, bias=False)
+
+        self.post_init()
+        
+    def get_input_embeddings(self):
+        return self.model.get_input_embeddings()
+
+    def set_input_embeddings(self, value):
+        self.model.set_input_embeddings(value)
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_decoder(self, decoder):
+        self.model.language_model = decoder
+
+    def get_decoder(self):
+        return self.model.language_model
+
+    def tie_weights(self):
+        """
+        Tie the weights between the input embeddings and the output embeddings.
+        """
+        if getattr(self.config.decoder_config, 'tie_word_embeddings', False):
+            # The standard PreTrainedModel method will handle the tying.
+            # It typically does a simple parameter object assignment, which is
+            # CORRECT to do BEFORE FSDP wraps the model.
+            output_embeddings = self.get_output_embeddings()
+            input_embeddings = self.get_input_embeddings()
+            if hasattr(input_embeddings, 'weight'):
+                output_embeddings.weight = input_embeddings.weight
+            else:
+                # maybe returned input_embeddings a tensor directly
+                output_embeddings.weight = input_embeddings
+
+            if getattr(output_embeddings, "bias", None) is not None:
+                output_embeddings.bias.data = nn.functional.pad(
+                    output_embeddings.bias.data,
+                    (0, output_embeddings.weight.shape[0] - output_embeddings.bias.shape[0]),
+                    "constant",
+                    0,
+                )
+            print("Tied input and output embeddings using standard assignment.")
+        else:
+            print("tie_word_embeddings is False, not tying weights.")
+
+    # Also, ensure set_output_embeddings is safe, though your implementation looks okay.
+    # The key is to avoid calling it after accelerator.prepare().
+    def set_output_embeddings(self, new_embeddings):
+        # Your current implementation using data.copy_ is good practice,
+        # but the best way is to not call this after prepare().
+        self.lm_head = new_embeddings
+
+    def forward_speech_features(
+            self, 
+            speech_tensors=None, 
+            speech_masks=None, 
+            speech_type="audio", 
+            return_unmask=False
+        ):
+        if speech_tensors is None:
+            # Use config to get vae_dim instead of non-existent self.args
+            vae_dim = self.config.acoustic_tokenizer_config.vae_dim
+            audio_features = torch.zeros(1, 1, vae_dim).to(self.get_input_embeddings().weight)
+            connect_features = self.model.acoustic_connector(audio_features)
+            return audio_features, connect_features
+        else:
+            with torch.no_grad():
+                if speech_type == "audio":
+                    with torch.no_grad():
+                        frames = self.model.acoustic_tokenizer.encode(speech_tensors.unsqueeze(1))[0][0]
+                    audio_tokens = frames.sample(self.model.acoustic_tokenizer.std_dist_type)[0]
+
+                elif speech_type == "vae":
+                    # Use config to get vae_dim instead of non-existent self.args
+                    vae_dim = self.config.acoustic_tokenizer_config.vae_dim
+                    speech_mode = speech_tensors.reshape(speech_tensors.size(0), -1, vae_dim)
+
+                    # gaussian sample from the speech_mode
+                    batch_size = speech_mode.size(0)
+                    value = self.model.acoustic_tokenizer.fix_std / 0.8
+                    std = torch.randn(batch_size, dtype=speech_mode.dtype, device=speech_mode.device) * value
+                    std = std.view(-1, *[1] * (speech_mode.dim() - 1))
+                    audio_tokens = speech_mode + std * torch.randn(speech_mode.shape).to(speech_mode)
+                else:
+                    raise NotImplementedError(f"Speech type {speech_type} not implemented")
+                
+                if torch.isnan(self.model.speech_scaling_factor) or torch.isnan(self.model.speech_bias_factor):
+                    scaling_factor = 1. / audio_tokens[speech_masks].flatten().std()
+                    bias_factor = -audio_tokens[speech_masks].flatten().mean()
+                    
+                    # Only use distributed operations if the process group is initialized
+                    if dist.is_available() and dist.is_initialized():
+                        dist.all_reduce(scaling_factor, op=dist.ReduceOp.SUM)
+                        dist.all_reduce(bias_factor, op=dist.ReduceOp.SUM)
+                        world_size = dist.get_world_size()
+                        self.model.speech_scaling_factor.copy_(scaling_factor / world_size)  
+                        self.model.speech_bias_factor.copy_(bias_factor / world_size)
+                        print(f"Speech scaling factor (distributed): {self.model.speech_scaling_factor}, bias factor: {self.model.speech_bias_factor}", flush=True)
+                    else:
+                        # Single process case
+                        self.model.speech_scaling_factor.copy_(scaling_factor)  
+                        self.model.speech_bias_factor.copy_(bias_factor)
+                        print(f"Speech scaling factor (single process): {self.model.speech_scaling_factor}, bias factor: {self.model.speech_bias_factor}", flush=True)
+                    
+                audio_features = (audio_tokens + self.model.speech_bias_factor) * self.model.speech_scaling_factor
+            
+            connect_features = self.model.acoustic_connector(audio_features)
+            if return_unmask:
+                return audio_features, connect_features
+            return audio_features[speech_masks], connect_features[speech_masks]
+        
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = False,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        # New arguments for speech processing and loss calculation
+        speech_tensors: Optional[torch.FloatTensor] = None,
+        speech_masks: Optional[torch.BoolTensor] = None,
+        speeches_loss_input: Optional[torch.FloatTensor] = None,
+        speech_semantic_tensors: Optional[torch.FloatTensor] = None, 
+        acoustic_input_mask: Optional[torch.BoolTensor] = None,
+        acoustic_loss_mask: Optional[torch.BoolTensor] = None,
+        ddpm_batch_mul: int = 1,
+        **kwargs: Optional[Dict[str, Union[torch.Tensor, str]]],
+        ) -> Union[Tuple, VibeVoiceCausalLMOutputWithPast]:
+        
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        
+        x = self.get_input_embeddings()(input_ids)
+
+        semantic_speech_all_connect_features = self.model.semantic_connector(speech_semantic_tensors)
+        if speeches_loss_input is not None:
+            # only part audio need diffuse
+            speech_all_features, speech_all_connect_features = self.forward_speech_features(
+                    speech_tensors=speech_tensors.type_as(x) if speech_tensors is not None else None,
+                    speech_masks=speech_masks,
+                    speech_type=kwargs.get("speech_type", "audio"),
+                    return_unmask=True
+                )
+            if speech_tensors is not None:
+                if semantic_speech_all_connect_features is not None:
+                    x[acoustic_input_mask] = (
+                        speech_all_connect_features[speech_masks]
+                        + semantic_speech_all_connect_features[speech_masks]
+                    )
+                else:
+                    x[acoustic_input_mask] = speech_all_connect_features[speech_masks]
+
+                # Select only the target segments' latents for diffusion loss.
+                # Both masks are [num_segments, max_latent_len]; using 2D mask on [B,T,D] selects [N_true, D].
+                target_latent_mask = speeches_loss_input & speech_masks
+                speech_features = speech_all_features[target_latent_mask]
+                speech_connect_features = speech_all_connect_features[target_latent_mask]
+        else:
+            speech_features, speech_connect_features = self.forward_speech_features(
+                    speech_tensors=speech_tensors.type_as(x) if speech_tensors is not None else None,
+                    speech_masks=speech_masks,
+                    speech_type=kwargs.get("speech_type", "audio"),
+                )
+            if speech_tensors is not None:
+                x[acoustic_input_mask] = speech_connect_features
+
+        outputs = self.model(
+            input_ids=None,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=x,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=False,
+            return_dict=return_dict,
+            cache_position=cache_position,
+        )
+
+        hidden_states = outputs.last_hidden_state
+        logits = self.lm_head(hidden_states)
+        # logits = logits.float()
+
+        loss = None
+        if labels is not None:
+            # The custom CE loss with masking is calculated in the training script.
+            # We leave the standard loss calculation here as None.
+            pass
+
+        # --- Diffusion Loss Calculation ---
+        diffusion_loss = None
+        # This block is executed only if we are in a context that involves speech.
+        if speech_tensors is not None and acoustic_loss_mask.sum().item() > 0:
+            condition_features = hidden_states[acoustic_loss_mask]
+            
+            speech_len, latent_size = speech_features.shape
+            
+            noise = torch.randn(
+                (speech_len * ddpm_batch_mul, latent_size),
+                device=hidden_states.device,
+                dtype=hidden_states.dtype
+            )
+            
+            timesteps = torch.multinomial(
+                torch.ones(self.config.diffusion_head_config.ddpm_num_steps),
+                speech_len * ddpm_batch_mul,
+                replacement=True,
+            ).to(hidden_states.device)
+
+            speech_features_repeated = speech_features.repeat_interleave(ddpm_batch_mul, dim=0)
+            condition_features_repeated = condition_features.repeat_interleave(ddpm_batch_mul, dim=0)
+
+            noisy_speech_features = self.model.noise_scheduler.add_noise(
+                speech_features_repeated, noise, timesteps
+            )
+            
+            model_output = self.model.prediction_head(
+                noisy_speech_features, 
+                timesteps.type_as(x), 
+                condition_features_repeated
+            )
+
+            prediction_type = self.config.diffusion_head_config.prediction_type
+            if prediction_type == "epsilon":
+                target_for_loss = noise
+            elif prediction_type == "v_prediction":
+                target_for_loss = self.model.noise_scheduler.get_velocity(
+                    speech_features_repeated, noise, timesteps
+                )
+            else:
+                raise NotImplementedError(f"Prediction type {prediction_type} not implemented")
+
+            diffusion_loss = F.mse_loss(model_output.float(), target_for_loss.float(), reduction='sum')
+            if latent_size > 0 and ddpm_batch_mul > 0:
+                diffusion_loss = diffusion_loss / latent_size / ddpm_batch_mul
+            else:
+                diffusion_loss = torch.tensor(0.0, device=diffusion_loss.device)
+        
+        else:
+            # Dummy loss for DDP to work when there are no speech samples in a batch,
+            # but we are in a speech context.
+            diffusion_loss = sum(p.sum() for p in self.model.prediction_head.parameters()) * 0.0
+            diffusion_loss += sum(p.sum() for p in self.model.acoustic_connector.parameters()) * 0.0
+            diffusion_loss += sum(p.sum() for p in self.model.semantic_connector.parameters()) * 0.0
+        # --- End Diffusion Loss Calculation ---
+
+        if not return_dict:
+            output = (logits, speech_len) + outputs.to_tuple()[1:]
+            return (loss, diffusion_loss) + output
+
+        return VibeVoiceCausalLMOutputWithPast(
+            loss=loss,
+            diffusion_loss=diffusion_loss,
+            speech_token_num=speech_len if speech_tensors is not None else 0,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+AutoModel.register(VibeVoiceConfig, VibeVoiceModel)
+AutoModelForCausalLM.register(VibeVoiceConfig, VibeVoiceForConditionalGeneration)
+
+__all__ = [
+    "VibeVoiceModel",
+    "VibeVoicePreTrainedModel",
+    "VibeVoiceForConditionalGeneration",
+    "VibeVoiceCausalLMOutputWithPast",
+    "VibeVoiceGenerationOutput",
+]
\ No newline at end of file
diff --git a/vibevoice/modular/modeling_vibevoice_asr.py b/vibevoice/modular/modeling_vibevoice_asr.py
new file mode 100644
index 0000000..4663d3f
--- /dev/null
+++ b/vibevoice/modular/modeling_vibevoice_asr.py
@@ -0,0 +1,522 @@
+from typing import List, Optional, Tuple, Union
+import torch
+import torch.nn as nn
+
+from transformers.models.auto import AutoModel, AutoModelForCausalLM
+
+from transformers.modeling_outputs import CausalLMOutput, BaseModelOutputWithPast
+from transformers import modeling_utils
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import logging
+from transformers.generation import GenerationMixin
+
+from .modular_vibevoice_tokenizer import (
+    VibeVoiceTokenizerStreamingCache, 
+    VibeVoiceTokenizerEncoderOutput
+)
+
+from .configuration_vibevoice import VibeVoiceASRConfig
+from .modeling_vibevoice import (
+    VibeVoiceCausalLMOutputWithPast,    
+    SpeechConnector
+)
+
+logger = logging.get_logger(__name__)
+
+if not hasattr(modeling_utils, "ALL_PARALLEL_STYLES") or modeling_utils.ALL_PARALLEL_STYLES is None:
+    modeling_utils.ALL_PARALLEL_STYLES = ["tp", "none", "colwise", "rowwise"]
+
+# @auto_docstring
+class VibeVoiceASRPreTrainedModel(PreTrainedModel):
+    config_class = VibeVoiceASRConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _skip_keys_device_placement = "past_key_values"
+    _supports_cache_class = True
+    _supports_flash_attn = True
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _supports_quantized_cache = True
+    _supports_static_cache = True
+    _supports_attention_backend = True
+
+    def _init_weights(self, module):
+
+        # Use the language model's initializer_range if available
+        if hasattr(self.config, 'language_model_config') and hasattr(self.config.language_model_config, 'initializer_range'):
+            std = self.config.language_model_config.initializer_range
+        elif hasattr(self.config, 'decoder_config') and hasattr(self.config.decoder_config, 'initializer_range'):
+            std = self.config.decoder_config.initializer_range
+        else:
+            std = 0.02  # Default value
+            
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.LayerNorm):
+            module.weight.data.fill_(1.0)
+            module.bias.data.zero_()
+
+# @auto_docstring
+class VibeVoiceASRModel(VibeVoiceASRPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        
+        if hasattr(config, 'torch_dtype') and config.torch_dtype is not None:
+            if isinstance(config.torch_dtype, str):
+                dtype = getattr(torch, config.torch_dtype)
+            else:
+                dtype = config.torch_dtype
+        else:
+            dtype = torch.float32
+        
+        # Initialize Qwen2 model for language modeling
+        lm_config = config.decoder_config 
+        self.language_model = AutoModel.from_config(lm_config)
+        
+        # Initialize speech components if needed
+        self.acoustic_tokenizer = AutoModel.from_config(config.acoustic_tokenizer_config).to(dtype)
+        self.semantic_tokenizer = AutoModel.from_config(config.semantic_tokenizer_config).to(dtype)
+
+        self.acoustic_connector = SpeechConnector(config.acoustic_vae_dim, lm_config.hidden_size).to(dtype)
+        self.semantic_connector = SpeechConnector(config.semantic_vae_dim, lm_config.hidden_size).to(dtype)
+    
+    def get_input_embeddings(self):
+        if hasattr(self.language_model, 'embed_tokens'):
+            # If the language model has an embed_tokens attribute, return it
+            return self.language_model.embed_tokens
+        
+        for name, attr in self.language_model.fullmap.items(): # parallel by nnscaler, the name is changed
+            if attr.orig_name == 'embed_tokens.weight':
+                return getattr(self.language_model, name)
+        assert False, 'should not arrive here'
+
+    def set_input_embeddings(self, value):
+        self.language_model.embed_tokens = value
+    
+    def set_speech_tokenizers(self, acoustic_tokenizer=None, semantic_tokenizer=None):
+        """Set the speech tokenizers used for encoding and decoding speech."""
+        self.acoustic_tokenizer = acoustic_tokenizer
+        self.semantic_tokenizer = semantic_tokenizer
+        
+        # Reset the encoder to evaluation mode
+        if self.acoustic_tokenizer is not None:
+            self.acoustic_tokenizer.eval()
+            
+        if self.semantic_tokenizer is not None:
+            self.semantic_tokenizer.eval()
+    
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        
+        # Forward through language model
+        outputs = self.language_model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            cache_position=cache_position,
+            **kwargs,
+        )
+        
+        if not return_dict:
+            return outputs
+            
+        return BaseModelOutputWithPast(
+            last_hidden_state=outputs.last_hidden_state,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+class VibeVoiceASRForConditionalGeneration(VibeVoiceASRPreTrainedModel, GenerationMixin):
+    """
+    VibeVoice model for Automatic Speech Recognition (ASR) with language modeling head for conditional generation.
+    This class is designed for inference and generation tasks.
+    """
+    _tied_weights_keys = ["lm_head.weight"]
+    _tp_plan = {"lm_head": "colwise_rep"}
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = VibeVoiceASRModel(config)
+        self.vocab_size = config.decoder_config.vocab_size
+        
+        # Determine the dtype to use
+        if hasattr(config, 'torch_dtype') and config.torch_dtype is not None:
+            if isinstance(config.torch_dtype, str):
+                dtype = getattr(torch, config.torch_dtype)
+            else:
+                dtype = config.torch_dtype
+        else:
+            dtype = torch.float32
+        
+        # Initialize lm_head with the correct dtype
+        self.lm_head = nn.Linear(config.decoder_config.hidden_size, self.vocab_size, bias=False).to(dtype)
+        
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.get_input_embeddings()
+
+    def set_input_embeddings(self, value):
+        self.model.set_input_embeddings(value)
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model.language_model = decoder
+
+    def get_decoder(self):
+        return self.model.language_model
+
+    def tie_weights(self):
+        """Tie the weights between the input embeddings and the output embeddings."""
+        if getattr(self.config.decoder_config, 'tie_word_embeddings', False):
+            output_embeddings = self.get_output_embeddings()
+            input_embeddings = self.get_input_embeddings()
+            if hasattr(input_embeddings, 'weight'):
+                output_embeddings.weight = input_embeddings.weight
+            else:
+                output_embeddings.weight = input_embeddings
+
+    def encode_speech(
+        self,
+        speech_tensors: torch.FloatTensor,
+        speech_masks: Optional[torch.BoolTensor] = None,
+        speech_semantic_tensors: Optional[torch.FloatTensor] = None,
+        streaming_segment_duration: float = 60.0,  # seconds
+    ):
+        """
+        Encode speech input into features that can be used by the language model.
+        This method is called once before generation to process the speech input.
+        
+        For long audio (>600s by default), uses streaming processing to avoid conv overflow (>2^32).
+        Segments are processed independently, then concatenated before final sampling.
+        
+        Args:
+            speech_tensors: Input audio tensor [batch_size, samples]
+            speech_masks: Optional mask for speech features
+            speech_semantic_tensors: Optional pre-computed semantic tokens
+            streaming_segment_duration: Segment duration in seconds for streaming processing (default: 60s)
+        """
+        if hasattr(self.config, 'torch_dtype') and self.config.torch_dtype is not None:
+            if isinstance(self.config.torch_dtype, str):
+                dtype = getattr(torch, self.config.torch_dtype)
+            else:
+                dtype = self.config.torch_dtype
+        else:
+            dtype = torch.float32
+            
+        speech_tensors = speech_tensors.to(dtype)
+        
+        # Ensure proper shape: (batch, samples)
+        if speech_tensors.ndim == 1:
+            speech_tensors = speech_tensors.unsqueeze(0)
+        
+        batch_size, total_samples = speech_tensors.shape
+        sample_rate = 24000  # fix 24kHz sample rate
+        
+        # Calculate segment size in samples
+        segment_samples = int(streaming_segment_duration * sample_rate)
+        
+        # Decide whether to use streaming based on audio length
+        use_streaming = total_samples > segment_samples
+        
+        with torch.no_grad():
+            if not use_streaming:
+                # Short audio: direct processing (original behavior)
+                encoder_output = self.model.acoustic_tokenizer.encode(speech_tensors.unsqueeze(1))
+                audio_tokens = encoder_output.sample(dist_type=self.model.acoustic_tokenizer.std_dist_type)[0]
+                acoustic_features = self.model.acoustic_connector(audio_tokens)
+                
+                # Encode semantic features
+                if speech_semantic_tensors is not None:
+                    semantic_features = self.model.semantic_connector(speech_semantic_tensors)
+                else:
+                    semantic_tokens = self.model.semantic_tokenizer.encode(speech_tensors.unsqueeze(1)).mean
+                    semantic_features = self.model.semantic_connector(semantic_tokens)
+            else:
+                # Long audio: streaming processing
+                # print(f"Using streaming processing for long audio: {total_samples/sample_rate:.1f}s "
+                #       f"(segment size: {streaming_segment_duration}s)")
+                
+                # Initialize caches for both tokenizers
+                acoustic_encoder_cache = VibeVoiceTokenizerStreamingCache()
+                semantic_encoder_cache = VibeVoiceTokenizerStreamingCache()
+                acoustic_mean_segments = []
+                semantic_mean_segments = []
+                sample_indices = torch.arange(batch_size, device=speech_tensors.device)
+                
+                # Helper function from batch_asr_sft_cache.py
+                def _iter_segments(total_length: int, segment_length: int):
+                    """Iterate over audio segments with a given segment length."""
+                    if segment_length <= 0:
+                        raise ValueError("segment_length must be positive")
+                    for start in range(0, total_length, segment_length):
+                        end = min(start + segment_length, total_length)
+                        if end > start:
+                            yield start, end
+                
+                # Process each segment for both acoustic and semantic tokenizers
+                segments = list(_iter_segments(total_samples, segment_samples))
+                num_segments = len(segments)
+                for seg_idx, (start, end) in enumerate(segments):
+                    chunk = speech_tensors[:, start:end].contiguous()
+                    if chunk.numel() == 0:
+                        continue
+                    
+                    # Check if this is the final segment
+                    is_final = (seg_idx == num_segments - 1)
+                    
+                    # Encode chunk for acoustic tokenizer (don't sample yet)
+                    acoustic_encoder_output = self.model.acoustic_tokenizer.encode(
+                        chunk.unsqueeze(1),
+                        cache=acoustic_encoder_cache,
+                        sample_indices=sample_indices,
+                        use_cache=True,
+                        is_final_chunk=is_final,
+                    )
+                    acoustic_mean_segments.append(acoustic_encoder_output.mean)
+                    
+                    # Encode chunk for semantic tokenizer (take mean directly)
+                    semantic_encoder_output = self.model.semantic_tokenizer.encode(
+                        chunk.unsqueeze(1),
+                        cache=semantic_encoder_cache,
+                        sample_indices=sample_indices,
+                        use_cache=True,
+                        is_final_chunk=is_final,
+                    )
+                    semantic_mean_segments.append(semantic_encoder_output.mean)
+                
+                # print(f"Processed {len(acoustic_mean_segments)} segments.")
+                # Concatenate all acoustic means and sample once
+                acoustic_mean_full = torch.cat(acoustic_mean_segments, dim=1).contiguous()
+                acoustic_encoder_output = VibeVoiceTokenizerEncoderOutput(
+                    mean=acoustic_mean_full,
+                    std=self.model.acoustic_tokenizer.fix_std
+                )
+                audio_tokens = acoustic_encoder_output.sample(
+                    dist_type=self.model.acoustic_tokenizer.std_dist_type
+                )[0]
+                acoustic_features = self.model.acoustic_connector(audio_tokens)
+                
+                # Concatenate all semantic means
+                semantic_tokens = torch.cat(semantic_mean_segments, dim=1).contiguous()
+                semantic_features = self.model.semantic_connector(semantic_tokens)
+            
+            # Combine acoustic and semantic features
+            if speech_masks is not None:
+                combined_features = acoustic_features[speech_masks] + semantic_features[speech_masks]
+            else:
+                combined_features = acoustic_features + semantic_features
+                
+        return combined_features
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        # Speech-specific arguments
+        speech_tensors: Optional[torch.FloatTensor] = None,
+        speech_masks: Optional[torch.BoolTensor] = None,
+        speech_semantic_tensors: Optional[torch.FloatTensor] = None,
+        acoustic_input_mask: Optional[torch.BoolTensor] = None,
+        **kwargs,
+    ) -> Union[Tuple, CausalLMOutput]:
+        """
+        Forward pass for the model. Handles both training and generation scenarios.
+        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        use_cache = use_cache if use_cache is not None else getattr(self.config, 'use_cache', False)
+
+        # Process inputs
+        if inputs_embeds is None and input_ids is not None:
+            inputs_embeds = self.get_input_embeddings()(input_ids)
+            
+        # If we have speech input and acoustic_input_mask, encode and insert speech features
+        if speech_tensors is not None and acoustic_input_mask is not None:
+            speech_features = self.encode_speech(
+                speech_tensors=speech_tensors,
+                speech_masks=speech_masks,
+                speech_semantic_tensors=speech_semantic_tensors,
+            )
+            # Clone to avoid in-place operation on leaf variable during training
+            inputs_embeds = inputs_embeds.clone()
+            inputs_embeds[acoustic_input_mask] = speech_features
+
+        # Forward through the model
+        outputs = self.model(
+            input_ids=None,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            cache_position=cache_position,
+        )
+
+        hidden_states = outputs[0] if not return_dict else outputs.last_hidden_state
+        logits = self.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
+            shift_logits = shift_logits.view(-1, self.vocab_size)
+            shift_labels = shift_labels.view(-1)
+            # Enable model parallelism
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return VibeVoiceCausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past_key_values=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        cache_position=None,
+        position_ids=None,
+        use_cache=True,
+        speech_tensors=None,
+        speech_masks=None,
+        speech_semantic_tensors=None,
+        acoustic_input_mask=None,
+        **kwargs,
+    ):
+        """
+        Prepare inputs for generation step. This method is called by generate() 
+        for each token generation step.
+        
+        Following Qwen2-VL's approach: speech inputs are only forwarded on the first pass
+        (when cache_position[0] == 0), and are excluded in subsequent generation steps.
+        """
+        # If we have past key values, we only need to process the new tokens
+        if past_key_values is not None:
+            if isinstance(past_key_values, tuple):
+                past_length = past_key_values[0][0].shape[2]
+            else:
+                past_length = past_key_values.get_seq_length()
+            
+            # Keep only the new tokens
+            if input_ids is not None and input_ids.shape[1] > past_length:
+                input_ids = input_ids[:, past_length:]
+
+        # Prepare position ids
+        if position_ids is None and attention_mask is not None:
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past_key_values is not None and input_ids is not None:
+                position_ids = position_ids[:, -input_ids.shape[1]:]
+
+        # Prepare cache position
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens, 
+                past_seen_tokens + (input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1]), 
+                device=input_ids.device if input_ids is not None else inputs_embeds.device
+            )
+
+        # Prepare model inputs
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        model_inputs.update(
+            {
+                "position_ids": position_ids,
+                "cache_position": cache_position,
+                "past_key_values": past_key_values,
+                "use_cache": use_cache,
+                "attention_mask": attention_mask,
+            }
+        )
+        
+        # Following Qwen2-VL pattern: only include speech inputs on the first forward pass
+        # (when cache_position[0] == 0), exclude them in subsequent generation steps
+        if cache_position is not None and len(cache_position) > 0 and cache_position[0] == 0:
+            # First forward pass - include speech inputs if provided
+            model_inputs.update({
+                "speech_tensors": speech_tensors,
+                "speech_masks": speech_masks,
+                "speech_semantic_tensors": speech_semantic_tensors,
+                "acoustic_input_mask": acoustic_input_mask,
+            })
+        else:
+            # Subsequent generation steps - exclude speech inputs
+            model_inputs.update({
+                "speech_tensors": None,
+                "speech_masks": None, 
+                "speech_semantic_tensors": None,
+                "acoustic_input_mask": None,
+            })
+        
+        # Include any remaining kwargs that might be needed
+        model_inputs.update(kwargs)
+        
+        return model_inputs
+
+AutoModel.register(VibeVoiceASRConfig, VibeVoiceASRModel)
+AutoModelForCausalLM.register(VibeVoiceASRConfig, VibeVoiceASRForConditionalGeneration)
+
+__all__ = [
+    "VibeVoiceASRPreTrainedModel",
+    "VibeVoiceASRModel",
+    "VibeVoiceASRForConditionalGeneration",
+]
\ No newline at end of file
diff --git a/vibevoice/modular/modeling_vibevoice_streaming.py b/vibevoice/modular/modeling_vibevoice_streaming.py
new file mode 100644
index 0000000..c4488c8
--- /dev/null
+++ b/vibevoice/modular/modeling_vibevoice_streaming.py
@@ -0,0 +1,190 @@
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Tuple, Union, Callable
+from tqdm import tqdm
+import copy
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.distributed as dist
+
+from transformers.models.auto import AutoModel, AutoModelForCausalLM
+
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import CausalLMOutput, BaseModelOutputWithPast, ModelOutput
+from transformers.models.llama.modeling_llama import LlamaRMSNorm
+from transformers import modeling_utils
+from transformers.modeling_utils import PreTrainedModel
+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+from transformers.utils import logging
+
+from .modular_vibevoice_diffusion_head import VibeVoiceDiffusionHead
+from vibevoice.schedule.dpm_solver import DPMSolverMultistepScheduler
+
+from .configuration_vibevoice_streaming import VibeVoiceStreamingConfig
+
+
+logger = logging.get_logger(__name__)
+
+if not hasattr(modeling_utils, "ALL_PARALLEL_STYLES") or modeling_utils.ALL_PARALLEL_STYLES is None:
+    modeling_utils.ALL_PARALLEL_STYLES = ["tp", "none", "colwise", "rowwise"]
+
+
+class BinaryClassifier(nn.Module):
+    def __init__(self, hidden_size):
+        super(BinaryClassifier, self).__init__()
+        self.fc1 = nn.Linear(hidden_size, hidden_size)
+        self.fc2 = nn.Linear(hidden_size, 1)
+
+    def forward(self, x):
+        x = torch.relu(self.fc1(x))
+        x = self.fc2(x)
+        return x
+
+
+class SpeechConnector(nn.Module):
+    def __init__(self, input_dim, output_dim):
+        super().__init__()
+        self.fc1 = nn.Linear(input_dim, output_dim)
+        self.norm = LlamaRMSNorm(output_dim, eps=1e-6)
+        self.fc2 = nn.Linear(output_dim, output_dim)
+
+    def forward(self, features, **kwargs):    
+        x = self.fc1(features)
+        x = self.norm(x)
+        x = self.fc2(x)
+        return x
+
+
+# @auto_docstring
+class VibeVoiceStreamingPreTrainedModel(PreTrainedModel):
+    config_class = VibeVoiceStreamingConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _skip_keys_device_placement = "past_key_values"
+    _supports_cache_class = True
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _supports_quantized_cache = True
+    _supports_static_cache = True
+    _supports_attention_backend = True
+
+    def _init_weights(self, module):
+        if isinstance(module, VibeVoiceDiffusionHead):
+            module.initialize_weights()
+            return
+
+        # Use the language model's initializer_range if available
+        if hasattr(self.config, 'language_model_config') and hasattr(self.config.language_model_config, 'initializer_range'):
+            std = self.config.language_model_config.initializer_range
+        elif hasattr(self.config, 'decoder_config') and hasattr(self.config.decoder_config, 'initializer_range'):
+            std = self.config.decoder_config.initializer_range
+        else:
+            std = 0.02  # Default value
+            
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.LayerNorm):
+            module.weight.data.fill_(1.0)
+            module.bias.data.zero_()
+
+
+# @auto_docstring
+class VibeVoiceStreamingModel(VibeVoiceStreamingPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        
+        if hasattr(config, 'torch_dtype') and config.torch_dtype is not None:
+            if isinstance(config.torch_dtype, str):
+                dtype = getattr(torch, config.torch_dtype)
+            else:
+                dtype = config.torch_dtype
+        else:
+            dtype = torch.float32
+        
+        # Initialize Qwen2 model for language modeling. 
+        # The lower Transformer layers are only used for encoding text, while the upper Transformer layers are used for encoding text and generating speech. 
+        # To keep the code clean, we constructs two language models.
+        # The final norm layer of the first language_model is set to identity and will not be used in inference.
+        lm_config = copy.deepcopy(config.decoder_config)
+        lm_backbone_num_hidden_layers = getattr(lm_config, 'num_hidden_layers', 24) - config.tts_backbone_num_hidden_layers
+        lm_config.num_hidden_layers = lm_backbone_num_hidden_layers
+        self.language_model = AutoModel.from_config(lm_config)
+        self.language_model.norm = nn.Identity()
+        
+        # We only need the Transformer layers here. Note that embed_tokens in tts_language_model is unused
+        tts_lm_config = copy.deepcopy(lm_config)
+        tts_lm_config.num_hidden_layers = config.tts_backbone_num_hidden_layers
+        self.tts_language_model = AutoModel.from_config(tts_lm_config)
+
+        # Marks the text that needs to be spoken by the TTS model.
+        self.tts_input_types = nn.Embedding(num_embeddings=2, embedding_dim=config.decoder_config.hidden_size)
+        
+        # Initialize speech components if needed
+        self.acoustic_tokenizer = AutoModel.from_config(config.acoustic_tokenizer_config).to(dtype)
+        self.acoustic_connector = SpeechConnector(config.acoustic_vae_dim, lm_config.hidden_size).to(dtype)
+        
+        # Register scaling factors as buffers - use 1D tensors for FSDP compatibility
+        self.register_buffer('speech_scaling_factor', torch.tensor(float('nan')))  
+        self.register_buffer('speech_bias_factor', torch.tensor(float('nan')))
+
+        # Initialize prediction head for speech generation
+        self.prediction_head = AutoModel.from_config(config.diffusion_head_config).to(dtype)
+
+        # Initialize noise scheduler
+        self.noise_scheduler = DPMSolverMultistepScheduler(
+            num_train_timesteps=config.diffusion_head_config.ddpm_num_steps,
+            beta_schedule=config.diffusion_head_config.ddpm_beta_schedule,
+            prediction_type=config.diffusion_head_config.prediction_type
+        )
+    
+    def get_input_embeddings(self):
+        if hasattr(self.language_model, 'embed_tokens'):
+            # If the language model has an embed_tokens attribute, return it
+            return self.language_model.embed_tokens
+        
+        for name, attr in self.language_model.fullmap.items(): # parallel by nnscaler, the name is changed
+            if attr.orig_name == 'embed_tokens.weight':
+                return getattr(self.language_model, name)
+        assert False, 'should not arrive here'
+
+    def set_input_embeddings(self, value):
+        self.language_model.embed_tokens = value
+    
+    def set_speech_tokenizers(self, acoustic_tokenizer=None):
+        """Set the speech tokenizers used for encoding and decoding speech."""
+        self.acoustic_tokenizer = acoustic_tokenizer
+        
+        # Reset the encoder to evaluation mode
+        if self.acoustic_tokenizer is not None:
+            self.acoustic_tokenizer.eval()
+    
+    def forward(self, *args, **kwargs):
+        """
+        Intentionally not implemented.
+
+        This streaming model is split into two explicit submodules:
+          - `language_model`      for plain text processing (lower layers).
+          - `tts_language_model`  for TTS-related upper layers.
+
+        We deliberately avoid a unified `forward` to prevent accidental calls
+        that mix responsibilities.
+
+        To use the model:
+          - Call `self.language_model(...)` for text embeddings / hidden states.
+          - Call `self.tts_language_model(...)` for the TTS portion.
+          - Use the dedicated inference class for combined generation logic.
+        """
+        raise RuntimeError(
+            "VibeVoiceStreamingModel.forward is intentionally disabled. "
+            "Use `model.language_model(...)` or `model.tts_language_model(...)` instead."
+        )
+
+
+AutoModel.register(VibeVoiceStreamingConfig, VibeVoiceStreamingModel)
+
+__all__ = [
+    "VibeVoiceStreamingPreTrainedModel",
+    "VibeVoiceStreamingModel",
+]
\ No newline at end of file
diff --git a/vibevoice/modular/modeling_vibevoice_streaming_inference.py b/vibevoice/modular/modeling_vibevoice_streaming_inference.py
new file mode 100644
index 0000000..70a4895
--- /dev/null
+++ b/vibevoice/modular/modeling_vibevoice_streaming_inference.py
@@ -0,0 +1,906 @@
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple, Union, Callable
+from tqdm import tqdm
+import inspect
+import torch
+import torch.nn as nn
+
+from transformers.models.auto import AutoModel, AutoModelForCausalLM
+from transformers.generation import GenerationMixin, GenerationConfig, LogitsProcessor, LogitsProcessorList, StoppingCriteriaList
+from transformers.modeling_outputs import BaseModelOutputWithPast, ModelOutput
+from transformers import modeling_utils
+from transformers.modeling_utils import PreTrainedModel
+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+from transformers.utils import logging
+
+from .modular_vibevoice_tokenizer import VibeVoiceTokenizerStreamingCache
+from .modular_vibevoice_diffusion_head import VibeVoiceDiffusionHead
+from vibevoice.schedule.dpm_solver import DPMSolverMultistepScheduler
+from .configuration_vibevoice_streaming import VibeVoiceStreamingConfig
+from .modular_vibevoice_text_tokenizer import VibeVoiceTextTokenizer, VibeVoiceTextTokenizerFast
+from .modeling_vibevoice_streaming import VibeVoiceStreamingPreTrainedModel, VibeVoiceStreamingModel, BinaryClassifier
+from .streamer import AudioStreamer, AsyncAudioStreamer
+
+logger = logging.get_logger(__name__)
+
+if not hasattr(modeling_utils, "ALL_PARALLEL_STYLES") or modeling_utils.ALL_PARALLEL_STYLES is None:
+    modeling_utils.ALL_PARALLEL_STYLES = ["tp", "none", "colwise", "rowwise"]
+
+TTS_TEXT_WINDOW_SIZE = 5
+TTS_SPEECH_WINDOW_SIZE = 6
+
+
+# ============================================================================
+# Transformers >= 4.57 Compatibility Layer
+# The cache system was refactored in transformers 4.57, requiring these helpers.
+# ============================================================================
+
+class MockCacheLayer:
+    """
+    Mock cache layer for transformers >= 4.57 compatibility.
+    Provides the `layers` interface expected by DynamicCache in newer versions.
+    """
+    
+    def __init__(self, key_cache, value_cache, parent_cache=None, layer_idx=0):
+        self.key_cache = key_cache
+        self.value_cache = value_cache
+        self._parent_cache = parent_cache
+        self._layer_idx = layer_idx
+    
+    def get_mask_sizes(self, cache_position):
+        """Return KV length and offset for mask creation."""
+        seq_length = self.key_cache.shape[2] if self.key_cache is not None else 0
+        query_length = cache_position.shape[0]
+        return seq_length + query_length, 0
+    
+    def update(self, key_states, value_states, cache_kwargs=None):
+        """Update the cache with new key/value states."""
+        if self._parent_cache is None:
+            return self.key_cache, self.value_cache
+        
+        parent = self._parent_cache
+        idx = self._layer_idx
+        
+        # Extend cache lists if needed
+        while len(parent.key_cache) <= idx:
+            parent.key_cache.append(None)
+            parent.value_cache.append(None)
+        
+        # Concatenate or initialize cache
+        if parent.key_cache[idx] is not None:
+            parent.key_cache[idx] = torch.cat([parent.key_cache[idx], key_states], dim=2)
+            parent.value_cache[idx] = torch.cat([parent.value_cache[idx], value_states], dim=2)
+        else:
+            parent.key_cache[idx] = key_states
+            parent.value_cache[idx] = value_states
+        
+        # Update local references
+        self.key_cache = parent.key_cache[idx]
+        self.value_cache = parent.value_cache[idx]
+        return self.key_cache, self.value_cache
+
+
+def _ensure_cache_has_layers(cache):
+    """
+    Ensure the cache has all required attributes for transformers >= 4.57.
+    Creates MockCacheLayer wrappers to provide the expected `layers` interface.
+    """
+    if cache is None:
+        return cache
+    
+    # Add required attributes (skip if read-only)
+    for attr, default in [('layer_class_to_replicate', None), ('offloading', False), ('is_compileable', False)]:
+        if not hasattr(cache, attr):
+            try:
+                setattr(cache, attr, default)
+            except AttributeError:
+                pass
+    
+    # Build layers list from key_cache/value_cache
+    if hasattr(cache, 'key_cache') and hasattr(cache, 'value_cache'):
+        try:
+            cache.layers = [
+                MockCacheLayer(cache.key_cache[i], cache.value_cache[i], parent_cache=cache, layer_idx=i)
+                for i in range(len(cache.key_cache))
+            ]
+        except AttributeError:
+            pass
+    elif not hasattr(cache, 'layers'):
+        try:
+            cache.layers = []
+        except AttributeError:
+            pass
+    
+    return cache
+
+
+def _update_model_kwargs_for_generation(
+    outputs: ModelOutput,
+    model_kwargs: Dict[str, Any],
+    num_new_tokens: int = 1,
+) -> Dict[str, Any]:
+    """
+    Update model_kwargs after adding new tokens (supports multi-token windows).
+    
+    Updates past_key_values, attention_mask, and cache_position for the next forward pass.
+    """
+    model_kwargs["past_key_values"] = _ensure_cache_has_layers(outputs.past_key_values)
+    
+    attention_mask = model_kwargs["attention_mask"]
+    model_kwargs["attention_mask"] = torch.cat(
+        [attention_mask, attention_mask.new_ones((attention_mask.shape[0], num_new_tokens))], dim=-1
+    )
+    
+    cache_pos = model_kwargs["cache_position"]
+    model_kwargs["cache_position"] = torch.arange(
+        cache_pos[-1] + 1, cache_pos[-1] + num_new_tokens + 1, device=cache_pos.device
+    )
+    
+    return model_kwargs
+
+
+@dataclass
+class VibeVoiceCausalLMOutputWithPast(BaseModelOutputWithPast):
+    logits: Optional[torch.FloatTensor] = None
+
+
+@dataclass
+class VibeVoiceGenerationOutput(ModelOutput):
+    """
+    Output type for VibeVoice generation.
+    
+    Args:
+        sequences (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+            The generated sequences. 
+        speech_outputs (`List[torch.FloatTensor]`, *optional*):
+            List of generated speech waveforms or latents for each speech segment.
+    """
+    sequences: torch.LongTensor = None
+    speech_outputs: Optional[List[torch.FloatTensor]] = None
+    reach_max_step_sample: Optional[torch.BoolTensor] = None
+
+
+class VibeVoiceStreamingForConditionalGenerationInference(VibeVoiceStreamingPreTrainedModel, GenerationMixin):
+
+    def __init__(self, config):
+        super().__init__(config)
+        
+        # Initialize the base model
+        self.model = VibeVoiceStreamingModel(config)
+
+        # TTS generation EOS classifier
+        self.tts_eos_classifier = BinaryClassifier(config.decoder_config.hidden_size)
+        
+        # inference configuration
+        self.ddpm_inference_steps = config.diffusion_head_config.ddpm_num_inference_steps
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @property
+    def noise_scheduler(self):
+        return self.model.noise_scheduler
+
+    @property
+    def prediction_head(self):
+        return self.model.prediction_head
+    
+    @property
+    def speech_scaling_factor(self):
+        return self.model.speech_scaling_factor
+
+    @property
+    def speech_bias_factor(self):
+        return self.model.speech_bias_factor
+
+    @property
+    def acoustic_tokenizer(self):
+        return self.model.acoustic_tokenizer
+    
+    @property
+    def acoustic_connector(self):
+        return self.model.acoustic_connector
+        
+    def tie_weights(self):
+        """
+        Tie the weights between the input embeddings and the output embeddings.
+        """
+        # Tie lm_head.weight to language_model.embed_tokens.weight
+        if not getattr(self.config, 'tie_word_embeddings', False):
+            return
+         
+        if hasattr(self, 'lm_head') and hasattr(self.model.language_model, 'embed_tokens'):
+            self.lm_head.weight = self.model.language_model.embed_tokens.weight
+        
+    def get_input_embeddings(self):
+        return self.model.get_input_embeddings()
+    
+    def set_input_embeddings(self, value):
+        self.model.set_input_embeddings(value)
+    
+    def get_output_embeddings(self):
+        """
+        This model does not define an `lm_head` (vocabulary projection).
+        """
+        return None
+    
+    def set_output_embeddings(self, new_embeddings):
+        """
+        No-op because there is no `lm_head`. Provided only to satisfy optional API calls.
+        To enable, first create `self.lm_head` then allow assignment.
+        """
+        raise RuntimeError("Output embeddings (lm_head) are not defined for this model. "
+                           "Create one before calling set_output_embeddings if needed.")
+    
+    def set_speech_tokenizers(self, acoustic_tokenizer=None):
+        """Set the speech tokenizers used for encoding and decoding speech."""
+        self.model.set_speech_tokenizers(acoustic_tokenizer)
+    
+    def set_ddpm_inference_steps(self, num_steps=None):
+        self.ddpm_inference_steps = num_steps or self.config.diffusion_head_config.ddpm_num_inference_steps
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.LongTensor,
+        past_key_values=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        cache_position=None,
+        **kwargs,
+    ):
+        """Prepare model inputs for generation (transformers >= 4.57 compatible)."""
+        model_inputs = {"cache_position": cache_position}
+
+        # Slice inputs when using cache
+        if past_key_values is not None:
+            model_inputs["past_key_values"] = past_key_values
+            if inputs_embeds is not None and input_ids.shape[1] == 0:
+                inputs_embeds = inputs_embeds[:, -cache_position.shape[0]:]
+            elif inputs_embeds is not None or (cache_position is not None and cache_position[-1] >= input_ids.shape[1]):
+                input_ids = input_ids[:, -cache_position.shape[0]:]
+            elif cache_position is not None and input_ids.shape[1] != cache_position.shape[0]:
+                input_ids = input_ids[:, cache_position]
+
+        # Set input_ids or inputs_embeds
+        use_embeds = inputs_embeds is not None and (
+            past_key_values is None or (cache_position is not None and len(cache_position) == inputs_embeds.shape[1])
+        )
+        if use_embeds:
+            model_inputs["input_ids"] = None
+            model_inputs["inputs_embeds"] = inputs_embeds
+        else:
+            model_inputs["input_ids"] = input_ids.clone(memory_format=torch.contiguous_format) if input_ids is not None else None
+            model_inputs["inputs_embeds"] = None
+
+        if attention_mask is not None:
+            model_inputs["attention_mask"] = attention_mask
+
+        # Create position_ids from attention_mask
+        if attention_mask is not None and kwargs.get("position_ids") is None:
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            kwargs["position_ids"] = position_ids
+
+        # Slice position_ids when using cache
+        if kwargs.get("position_ids") is not None:
+            if past_key_values is not None:
+                seq_len = model_inputs["inputs_embeds"].shape[1] if model_inputs.get("inputs_embeds") is not None else model_inputs["input_ids"].shape[1]
+                model_inputs["position_ids"] = kwargs["position_ids"][:, -seq_len:].clone(memory_format=torch.contiguous_format)
+            else:
+                model_inputs["position_ids"] = kwargs.pop("position_ids").clone(memory_format=torch.contiguous_format)
+
+        # Forward remaining kwargs
+        for key, value in kwargs.items():
+            if key not in model_inputs:
+                model_inputs[key] = value
+
+        model_inputs.pop("labels", None)
+        return model_inputs
+
+    def _update_model_kwargs_for_generation(
+        self,
+        outputs,
+        model_kwargs,
+        is_encoder_decoder=False,
+        num_new_tokens=1,
+    ):
+        """Override to ensure cache compatibility with transformers >= 4.57."""
+        model_kwargs = super()._update_model_kwargs_for_generation(
+            outputs, model_kwargs, is_encoder_decoder=is_encoder_decoder, num_new_tokens=num_new_tokens
+        )
+        if "past_key_values" in model_kwargs:
+            model_kwargs["past_key_values"] = _ensure_cache_has_layers(model_kwargs["past_key_values"])
+        return model_kwargs
+
+    def _init_cache_for_generation(self, generation_config, model_kwargs, batch_size, max_cache_length, device):
+        """
+        Initialize cache for generation, handling different transformers versions.
+        For transformers >= 4.57, returns None to let the model create the cache dynamically.
+        """
+        try:
+            from transformers.cache_utils import DynamicCache
+            sig = inspect.signature(DynamicCache.__init__)
+            if 'config' in sig.parameters:
+                # transformers >= 4.57: let model handle cache creation
+                return None
+            else:
+                # Older versions: use parent method
+                prep_sig = inspect.signature(self._prepare_cache_for_generation)
+                if 'device' in prep_sig.parameters:
+                    self._prepare_cache_for_generation(generation_config, model_kwargs, None, batch_size, max_cache_length, device)
+                else:
+                    self._prepare_cache_for_generation(generation_config, model_kwargs, None, batch_size, max_cache_length)
+                return model_kwargs.get("past_key_values")
+        except Exception:
+            return None
+
+    # @can_return_tuple
+    def forward_lm(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        """
+        Single pass of the base text LM.
+
+        - Builds embeddings if `inputs_embeds` not provided.
+        - Uses (and returns) `past_key_values` when `use_cache=True`.
+        - No loss / no lm_head / no speech logic.
+
+        Args:
+            input_ids: (B, S) token ids.
+            attention_mask: (B, S) mask.
+            past_key_values: cache from previous steps.
+            cache_position: positions for cached tokens.
+            labels: unsupported (will raise).
+
+        Returns:
+            BaseModelOutputWithPast with `last_hidden_state` and `past_key_values`.
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        
+        # Get embeddings
+        if inputs_embeds is None:
+            inputs_embeds = self.model.get_input_embeddings()(input_ids)
+
+        outputs = self.model.language_model(
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        hidden_states = outputs[0] if not return_dict else outputs.last_hidden_state
+                
+        if labels is not None:
+            raise NotImplementedError("Loss computation is not implemented in this version.")
+
+        return BaseModelOutputWithPast(
+            past_key_values=outputs.past_key_values,
+            last_hidden_state=hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    # @can_return_tuple
+    def forward_tts_lm(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        lm_last_hidden_state: Optional[torch.FloatTensor] = None,
+        tts_text_masks: Optional[torch.BoolTensor] = None,
+        **kwargs,
+    ) -> Union[Tuple, VibeVoiceCausalLMOutputWithPast]:
+        """
+        Single pass of the TTS LM.
+
+        - Overwrites tail embeddings with `lm_last_hidden_state`.
+        - Adds type embedding via `tts_text_masks` (1=text, 0=speech).
+        - Predicts EOS from last hidden state (binary classifier).
+        - No loss / no full acoustic decoding here.
+
+        Args:
+            input_ids: (B, S) token ids.
+            attention_mask: (B, S) mask.
+            lm_last_hidden_state: (B, K, H) hidden states to splice into the tail.
+            tts_text_masks: (B, 1) mask marking current position as text(1)/speech(0).
+            past_key_values: cache from previous TTS steps.
+            cache_position: positions for cached tokens.
+            labels: unsupported (will raise).
+
+        Returns:
+            VibeVoiceCausalLMOutputWithPast with `logits` (EOS), `last_hidden_state`, `past_key_values`.
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        
+        # Get embeddings
+        if inputs_embeds is None:
+            # Will be replaced with lm_last_hidden_state
+            inputs_embeds = self.model.get_input_embeddings()(input_ids)
+        
+        # Replace the last part of inputs_embeds with lm_last_hidden_state
+        start_idx = inputs_embeds.shape[1] - lm_last_hidden_state.shape[1]
+        inputs_embeds[:, start_idx:, :] = lm_last_hidden_state
+        
+        # Adds type embedding via `tts_text_masks`.
+        inputs_embeds = inputs_embeds + self.model.tts_input_types(tts_text_masks.long())
+
+        outputs = self.model.tts_language_model(
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        hidden_states = outputs[0] if not return_dict else outputs.last_hidden_state
+        logits = self.tts_eos_classifier(hidden_states[:, -1, :])
+                
+        if labels is not None:
+            raise NotImplementedError("Loss computation is not implemented in this version.")
+
+        return VibeVoiceCausalLMOutputWithPast(
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            last_hidden_state=hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def forward(self, *args, **kwargs):
+        """
+        Unified forward is intentionally disabled.
+
+        Reasons:
+          1. The inference pipeline is staged: base text LM, then TTS LM, plus streaming & diffusion handled in `generate`.
+          2. A monolithic call would hide required sequencing (prefill, window stepping, speech diffusion sampling).
+
+        Use instead:
+          - self.forward_lm(...)       for a base text LM step (prefill or incremental).
+          - self.forward_tts_lm(...)   for a single TTS LM step (needs LM hidden states).
+          - self.generate(...)         for full streaming (text + speech + diffusion + audio assembly).
+
+        Raises:
+            RuntimeError: Always (by design).
+        """
+        raise RuntimeError(
+            "Unified forward is disabled. Use `forward_lm`, `forward_tts_lm`, or `generate` instead."
+        )
+
+    def _build_generate_config_model_kwargs(self, generation_config, inputs, tokenizer, return_processors=False, **kwargs):
+        if generation_config is None:
+            generation_config = GenerationConfig(
+                bos_token_id=tokenizer.bos_token_id,
+                eos_token_id=tokenizer.eos_token_id,
+                pad_token_id = tokenizer.pad_token_id
+            )
+        else:
+            generation_config = GenerationConfig(
+                **generation_config,
+                bos_token_id=tokenizer.bos_token_id,
+                eos_token_id=tokenizer.eos_token_id,
+                pad_token_id = tokenizer.pad_token_id
+            )
+
+        generation_config, model_kwargs = self._prepare_generation_config(
+            generation_config, 
+            True, 
+            speech_start_id=tokenizer.speech_start_id, 
+            speech_end_id=tokenizer.speech_end_id, 
+            speech_diffusion_id=tokenizer.speech_diffusion_id, 
+            **kwargs
+        )
+        generation_config.speech_start_id = tokenizer.speech_start_id
+        generation_config.speech_end_id = tokenizer.speech_end_id
+        generation_config.speech_diffusion_id = tokenizer.speech_diffusion_id
+
+        inputs_tensor, model_input_name, model_kwargs = self._prepare_model_inputs(inputs, generation_config.bos_token_id, model_kwargs)
+        batch_size = inputs_tensor.shape[0]
+        device = self.device
+        
+        self._prepare_special_tokens(generation_config, True, device=device)
+        generation_config.use_cache = True
+        model_kwargs["use_cache"] = generation_config.use_cache
+        input_ids = inputs_tensor.to(self.device)
+
+        input_ids_length = input_ids.shape[1]
+        has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None
+        has_default_min_length = kwargs.get("min_length") is None and generation_config.min_length is not None
+        generation_config = self._prepare_generated_length(
+            generation_config=generation_config,
+            has_default_max_length=has_default_max_length,
+            has_default_min_length=has_default_min_length,
+            model_input_name=model_input_name,
+            inputs_tensor=inputs_tensor,
+            input_ids_length=input_ids_length,
+        )
+
+        max_cache_length = generation_config.max_length - 1
+        # Handle cache initialization for different transformers versions
+        model_kwargs["past_key_values"] = self._init_cache_for_generation(
+            generation_config, model_kwargs, batch_size, max_cache_length, device
+        )
+        model_kwargs['cache_position'] = torch.arange(input_ids_length, device=device, dtype=torch.long)
+        for k, v in model_kwargs.items():
+            if isinstance(v, torch.Tensor):
+                model_kwargs[k] = v.to(device=device)
+        
+        if return_processors:
+            logits_processor = self._get_logits_processor(
+                generation_config=generation_config,
+                input_ids_seq_length=input_ids_length,
+                encoder_input_ids=inputs_tensor,
+                prefix_allowed_tokens_fn=None,
+                logits_processor=LogitsProcessorList(),
+                device=inputs_tensor.device,
+                model_kwargs=model_kwargs,
+            )
+
+            stopping_criteria = self._get_stopping_criteria(generation_config=generation_config, stopping_criteria=StoppingCriteriaList())
+        
+            return generation_config, model_kwargs, input_ids, logits_processor, stopping_criteria
+        else:
+            return generation_config, model_kwargs, input_ids
+
+    @torch.no_grad()
+    def generate(
+        self,
+        inputs: Optional[torch.Tensor] = None,
+        generation_config: Optional[GenerationConfig] = None,
+        logits_processor: Optional[LogitsProcessorList] = None,
+        stopping_criteria: Optional[StoppingCriteriaList] = None,
+        prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None,
+        synced_gpus: Optional[bool] = None,
+        assistant_model: Optional["PreTrainedModel"] = None,
+        audio_streamer: Optional[Union[AudioStreamer, AsyncAudioStreamer]] = None,
+        negative_prompt_ids: Optional[torch.Tensor] = None,
+        negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+        speech_tensors: Optional[torch.FloatTensor] = None,
+        speech_masks: Optional[torch.BoolTensor] = None,
+        speech_input_mask: Optional[torch.BoolTensor] = None,
+        tts_text_ids: Optional[torch.LongTensor] = None,
+        return_speech: bool = True,
+        cfg_scale: float = 1.0,
+        stop_check_fn: Optional[Callable[[], bool]] = None,
+        **kwargs,
+    ) -> Union[torch.LongTensor, VibeVoiceGenerationOutput]:
+        """
+        Text is fed in small windows (dynamic slicing of `tts_text_ids`), which enables streaming text input: you don’t need the full text upfront. After each text window, a loop samples several speech latents (diffusion). The interleaved text encoding + speech generation enables streaming text input and realtime speech output.
+        The function only supports batch size = 1 currently.
+
+        - Windowed text prefill → incremental LM + TTS LM updates.
+        - Interleave speech token diffusion sampling (`sample_speech_tokens`).
+        - Stops on EOS (binary classifier) or max length / external `stop_check_fn`.
+        - Returns final token `sequences` and (optionally) concatenated speech audio.
+
+        Args (selected):
+            tts_text_ids: Full text tokens to stream in windows.
+            audio_streamer: If provided, emits audio chunks during generation.
+            cfg_scale: Classifier-free guidance scale for speech diffusion.
+            return_speech: If False, skips audio decode concatenation.
+            stop_check_fn: External early-stop hook (returns True to halt).
+
+        Returns:
+            VibeVoiceGenerationOutput with:
+              - sequences: final token ids
+              - speech_outputs: list of concatenated audio tensors (or None)
+              - reach_max_step_sample: flags for samples stopped by max length
+        """
+        # 1. Handle `generation_config` and kwargs that might update it, and validate the `.generate()` call
+        tokenizer = kwargs.pop("tokenizer", None)
+        neg_text_input_id = tokenizer.convert_tokens_to_ids("<|image_pad|>")
+        
+        tts_lm_input_ids = kwargs.pop("tts_lm_input_ids", None)
+        tts_lm_attention_mask = kwargs.pop("tts_lm_attention_mask", None)
+        # all_prefilled_outputs: cached prefilled prompt outputs for lm, tts_lm, neg_lm, neg_tts_lm
+        all_prefilled_outputs = kwargs.pop("all_prefilled_outputs", None)
+        tts_text_ids = tts_text_ids.to(self.device)
+
+        if kwargs.get('max_new_tokens', None) is None:
+            kwargs['max_new_tokens'] = self.config.decoder_config.max_position_embeddings - tts_lm_input_ids.shape[-1]
+
+        generation_config, model_kwargs, input_ids, logits_processor, stopping_criteria = self._build_generate_config_model_kwargs(
+            generation_config, inputs, tokenizer, return_processors=True, **kwargs
+        )
+        
+        negative_kwargs = {
+            'input_ids': torch.full((kwargs['input_ids'].shape[0], 1), neg_text_input_id, dtype=torch.long, device=kwargs['input_ids'].device),
+            'attention_mask':  torch.ones((kwargs['input_ids'].shape[0], 1), dtype=torch.long, device=kwargs['input_ids'].device),
+            'max_new_tokens': kwargs.get('max_new_tokens', 100) 
+        }
+        negative_generation_config, negative_model_kwargs, negative_input_ids = self._build_generate_config_model_kwargs(
+            None, None, tokenizer, return_processors=False, **negative_kwargs
+        )
+
+        tts_lm_kwargs = {
+            'input_ids': tts_lm_input_ids,
+            'attention_mask': tts_lm_attention_mask,
+            'max_new_tokens': kwargs.get('max_new_tokens', 100) 
+        }
+        tts_lm_generation_config, tts_lm_model_kwargs, tts_lm_input_ids = self._build_generate_config_model_kwargs(
+            None, None, tokenizer, return_processors=False, **tts_lm_kwargs
+        )
+
+        tts_lm_negative_kwargs = {
+            'input_ids': torch.full((kwargs['input_ids'].shape[0], 1), neg_text_input_id, dtype=torch.long, device=kwargs['input_ids'].device),
+            'attention_mask':  torch.ones((kwargs['input_ids'].shape[0], 1), dtype=torch.long, device=kwargs['input_ids'].device),
+            'max_new_tokens': kwargs.get('max_new_tokens', 100) 
+        }
+        tts_lm_negative_generation_config, tts_lm_negative_model_kwargs, tts_lm_negative_input_ids = self._build_generate_config_model_kwargs(
+            None, None, tokenizer, return_processors=False, **tts_lm_negative_kwargs
+        )
+
+        acoustic_cache = VibeVoiceTokenizerStreamingCache()
+        batch_size = input_ids.shape[0]
+        assert batch_size == 1, "Currently only supports batch size == 1"
+        device = input_ids.device
+        finished_tags = torch.zeros(batch_size, dtype=torch.bool, device=device)
+        verbose = kwargs.get("verbose", False)
+
+        # Initialize audio chunks storage for each sample
+        audio_chunks = [[] for _ in range(batch_size)]
+        tts_text_window_index = 0
+        reach_max_step_sample = torch.zeros(batch_size, dtype=torch.bool, device=device)
+        first_text_window_size = TTS_TEXT_WINDOW_SIZE if tts_text_ids.shape[1] >= TTS_TEXT_WINDOW_SIZE else tts_text_ids.shape[1]
+
+        outputs = all_prefilled_outputs["lm"]
+        tts_lm_outputs = all_prefilled_outputs["tts_lm"]
+        negative_outputs = all_prefilled_outputs["neg_lm"]
+        tts_lm_negative_outputs = all_prefilled_outputs["neg_tts_lm"]
+
+        model_kwargs = _update_model_kwargs_for_generation(
+            outputs, model_kwargs, num_new_tokens=first_text_window_size,
+        )
+        tts_lm_model_kwargs = _update_model_kwargs_for_generation(
+            tts_lm_outputs, tts_lm_model_kwargs, num_new_tokens=first_text_window_size,
+        )
+        negative_model_kwargs = self._update_model_kwargs_for_generation(
+            negative_outputs, negative_model_kwargs, is_encoder_decoder=False,
+        )
+        tts_lm_negative_model_kwargs = self._update_model_kwargs_for_generation(
+            tts_lm_negative_outputs, tts_lm_negative_model_kwargs, is_encoder_decoder=False,
+        )
+
+        step = tts_lm_input_ids.shape[1]
+        total_generated_speech_tokens = 0
+        total_prefilled_text_tokens = 0
+        if kwargs.get("show_progress_bar", True):
+            progress_bar = tqdm(
+                total=tts_lm_generation_config.max_length,
+                desc=f"Prefilled {step} tokens, current step ({step} / {tts_lm_generation_config.max_length})",
+                initial=step,
+                leave=False
+            )
+        else:
+            progress_bar = None
+
+        while True:
+            # Check for external stop signal
+            if stop_check_fn is not None and stop_check_fn():
+                if verbose:
+                    print(f"Generation stopped externally at step {step + 1}")
+                # End the audio streamer if it exists
+                if audio_streamer is not None:
+                    audio_streamer.end()
+                break
+            
+            # # Check if audio_streamer has been ended (stopped externally)
+            # if audio_streamer is not None and hasattr(audio_streamer, 'finished_flags'):
+            #     if any(audio_streamer.finished_flags):
+            #         if verbose:
+            #             print(f"Audio generation stopped externally at step {step + 1}")
+            #         break
+            
+            if finished_tags.all():
+                if hasattr(progress_bar, 'set_description'):
+                    progress_bar.set_description("Generation complete")
+                break
+
+            cur_input_tts_text_ids = tts_text_ids[:, tts_text_window_index*TTS_TEXT_WINDOW_SIZE:(tts_text_window_index+1)*TTS_TEXT_WINDOW_SIZE]
+            next_text_window_size = tts_text_ids[:, (tts_text_window_index+1)*TTS_TEXT_WINDOW_SIZE:(tts_text_window_index+2)*TTS_TEXT_WINDOW_SIZE].shape[1]
+            tts_text_window_index += 1
+
+            if cur_input_tts_text_ids.shape[1] > 0:
+                input_ids = torch.cat([input_ids, cur_input_tts_text_ids], dim=-1)
+                tts_lm_input_ids = torch.cat([tts_lm_input_ids, cur_input_tts_text_ids], dim=-1)
+
+                if tts_lm_input_ids.shape[1] > tts_lm_generation_config.max_length:
+                    if verbose:
+                        print(f"Reached maximum generation length {generation_config.max_length}, stopped it.")
+                    reached_samples = torch.arange(batch_size, device=device)[~finished_tags]
+                    if reached_samples.numel() > 0:
+                        reach_max_step_sample[reached_samples] = True
+                    break
+                
+                step += cur_input_tts_text_ids.shape[1]
+                total_prefilled_text_tokens += cur_input_tts_text_ids.shape[1]
+                if progress_bar is not None:
+                    progress_bar.update(cur_input_tts_text_ids.shape[1])
+                    progress_bar.set_description(f"Prefilled {total_prefilled_text_tokens} text tokens, generated {total_generated_speech_tokens} speech tokens, current step ({step} / {tts_lm_generation_config.max_length})")
+
+                model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
+                # Forward pass through the model
+                outputs = self.forward_lm(
+                    **model_inputs, return_dict=True, output_attentions=False, output_hidden_states=False,
+                )
+                model_kwargs = _update_model_kwargs_for_generation(
+                    outputs, model_kwargs, num_new_tokens=next_text_window_size,
+                )
+
+                tts_lm_model_inputs = self.prepare_inputs_for_generation(tts_lm_input_ids, **tts_lm_model_kwargs)
+                tts_lm_additional_inputs = {
+                    "tts_text_masks": torch.ones_like(tts_lm_input_ids[:, -1:]),
+                    "lm_last_hidden_state": outputs.last_hidden_state,
+                }
+                # Forward pass through the model
+                tts_lm_outputs = self.forward_tts_lm(
+                    **tts_lm_model_inputs, **tts_lm_additional_inputs, return_dict=True, output_attentions=False, output_hidden_states=False,
+                )
+                tts_lm_model_kwargs = self._update_model_kwargs_for_generation(
+                    tts_lm_outputs, tts_lm_model_kwargs, is_encoder_decoder=False,
+                )
+
+            diffusion_indices = torch.LongTensor([0])
+            for cur_speech_index in range(TTS_SPEECH_WINDOW_SIZE):
+                positive_condition = tts_lm_outputs.last_hidden_state[diffusion_indices, -1, :]
+                negative_condition = tts_lm_negative_outputs.last_hidden_state[diffusion_indices, -1, :]
+                
+                speech_latent = self.sample_speech_tokens(
+                    positive_condition,
+                    negative_condition,
+                    cfg_scale=cfg_scale,
+                ).unsqueeze(1)
+                                
+                # Decode acoustic latent to audio using acoustic streaming cache
+                scaled_latent = speech_latent / self.model.speech_scaling_factor.to(speech_latent.device) - self.model.speech_bias_factor.to(speech_latent.device)
+                audio_chunk = self.model.acoustic_tokenizer.decode(
+                    scaled_latent.to(self.model.acoustic_tokenizer.device),
+                    cache=acoustic_cache,  # Use acoustic-specific cache
+                    sample_indices=diffusion_indices.to(self.model.acoustic_tokenizer.device),
+                    use_cache=True,
+                    debug=False
+                )
+                
+                # Store audio chunks for each sample
+                for i, sample_idx in enumerate(diffusion_indices):
+                    idx = sample_idx.item()
+                    # Only append audio chunk if the sample is not finished
+                    if not finished_tags[idx]:
+                        audio_chunks[idx].append(audio_chunk[i])
+
+                 # Add streaming support here
+                if audio_streamer is not None:
+                    # Stream the audio chunks immediately
+                    audio_streamer.put(audio_chunk, diffusion_indices)
+
+                acoustic_embed = self.model.acoustic_connector(speech_latent)
+                tts_lm_input_ids = torch.cat([tts_lm_input_ids, torch.ones_like(tts_lm_input_ids[:, -1:])], dim=-1)
+
+                if tts_lm_input_ids.shape[1] > tts_lm_generation_config.max_length:
+                    break
+                
+                step += 1
+                total_generated_speech_tokens += 1
+                if progress_bar is not None:
+                    progress_bar.update(1)
+                    progress_bar.set_description(f"Prefilled {total_prefilled_text_tokens} text tokens, generated {total_generated_speech_tokens} speech tokens, current step ({step} / {tts_lm_generation_config.max_length})")
+
+                tts_lm_model_inputs = self.prepare_inputs_for_generation(tts_lm_input_ids, **tts_lm_model_kwargs)
+                tts_lm_additional_inputs = {
+                    "tts_text_masks": torch.zeros_like(tts_lm_input_ids[:, -1:]),
+                    "lm_last_hidden_state": acoustic_embed,
+                }
+                # Forward pass through the model
+                tts_lm_outputs = self.forward_tts_lm(
+                    **tts_lm_model_inputs, **tts_lm_additional_inputs, return_dict=True, output_attentions=False, output_hidden_states=False,
+                )
+                if cur_speech_index == TTS_SPEECH_WINDOW_SIZE - 1 and next_text_window_size > 0:
+                    tts_lm_model_kwargs = _update_model_kwargs_for_generation(
+                        tts_lm_outputs, tts_lm_model_kwargs, num_new_tokens=next_text_window_size,
+                    )
+                else:
+                    tts_lm_model_kwargs = self._update_model_kwargs_for_generation(
+                        tts_lm_outputs, tts_lm_model_kwargs, is_encoder_decoder=False,
+                    )
+
+                tts_lm_negative_input_ids = torch.cat([tts_lm_negative_input_ids, torch.ones_like(tts_lm_input_ids[:, -1:])], dim=-1)
+                tts_lm_negative_model_inputs = self.prepare_inputs_for_generation(tts_lm_negative_input_ids, **tts_lm_negative_model_kwargs)
+                # Forward negative pass through the model
+                tts_lm_negative_additional_inputs = {
+                    "tts_text_masks": torch.zeros_like(tts_lm_negative_input_ids[:, -1:]),
+                    "lm_last_hidden_state": acoustic_embed,
+                }
+                tts_lm_negative_outputs = self.forward_tts_lm(
+                    **tts_lm_negative_model_inputs, **tts_lm_negative_additional_inputs, return_dict=True, output_attentions=False, output_hidden_states=False,
+                )
+                tts_lm_negative_model_kwargs = self._update_model_kwargs_for_generation(
+                    tts_lm_negative_outputs, tts_lm_negative_model_kwargs, is_encoder_decoder=False,
+                )
+
+                tts_eos_logits = torch.sigmoid(self.tts_eos_classifier(tts_lm_outputs.last_hidden_state[diffusion_indices, -1, :]))
+                if tts_eos_logits[0].item() > 0.5:
+                    # If EOS token is predicted, we can stop generation for this sample
+                    finished_tags[diffusion_indices] = True
+                    if audio_streamer is not None:
+                        audio_streamer.end(diffusion_indices)
+
+            if tts_lm_input_ids.shape[1] > tts_lm_generation_config.max_length:
+                if verbose:
+                    print(f"Reached maximum generation length {tts_lm_generation_config.max_length}, stopped it.")
+                reached_samples = torch.arange(batch_size, device=device)[~finished_tags]
+                if reached_samples.numel() > 0:
+                    reach_max_step_sample[reached_samples] = True
+                break
+
+        if audio_streamer is not None:
+            audio_streamer.end()
+
+        # Concatenate audio chunks for each sample
+        final_audio_outputs = []
+        for sample_chunks in audio_chunks:
+            if sample_chunks:
+                # Concatenate all chunks along the time dimension (assumed to be the last dimension)
+                concatenated_audio = torch.cat(sample_chunks, dim=-1)
+                final_audio_outputs.append(concatenated_audio)
+            else:
+                # If no audio was generated for this sample, append None
+                final_audio_outputs.append(None)
+        
+        if reach_max_step_sample is not None and reach_max_step_sample.any():
+            print(f"Reached maximum generation length {tts_lm_generation_config.max_length}, stopped it.")
+
+        return VibeVoiceGenerationOutput(
+            sequences=tts_lm_input_ids,
+            speech_outputs=final_audio_outputs if return_speech else None,
+            reach_max_step_sample=reach_max_step_sample,
+        )
+
+    @torch.no_grad()
+    def sample_speech_tokens(self, condition, neg_condition, cfg_scale=3.0):
+        self.model.noise_scheduler.set_timesteps(self.ddpm_inference_steps)
+        condition = torch.cat([condition, neg_condition], dim=0).to(self.model.prediction_head.device)
+        speech = torch.randn(condition.shape[0], self.config.acoustic_vae_dim).to(condition)
+        for t in self.model.noise_scheduler.timesteps:
+            half = speech[: len(speech) // 2]
+            combined = torch.cat([half, half], dim=0)
+            eps = self.model.prediction_head(combined, t.repeat(combined.shape[0]).to(combined), condition=condition)
+            cond_eps, uncond_eps = torch.split(eps, len(eps) // 2, dim=0)
+            half_eps = uncond_eps + cfg_scale * (cond_eps - uncond_eps)
+            eps = torch.cat([half_eps, half_eps], dim=0)
+            speech = self.model.noise_scheduler.step(eps, t, speech).prev_sample
+        return speech[: len(speech) // 2]
+    
+
+AutoModelForCausalLM.register(VibeVoiceStreamingConfig, VibeVoiceStreamingForConditionalGenerationInference)
+
+__all__ = [
+    "VibeVoiceStreamingForConditionalGenerationInference",
+]
diff --git a/vibevoice/modular/modular_vibevoice_diffusion_head.py b/vibevoice/modular/modular_vibevoice_diffusion_head.py
new file mode 100644
index 0000000..59de50f
--- /dev/null
+++ b/vibevoice/modular/modular_vibevoice_diffusion_head.py
@@ -0,0 +1,287 @@
+import math
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers.models.auto import AutoModel
+from transformers.modeling_utils import PreTrainedModel
+# from transformers.modeling_layers import GradientCheckpointingLayer
+from transformers.activations import ACT2FN
+from transformers.utils import logging
+
+from .configuration_vibevoice import VibeVoiceDiffusionHeadConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6, elementwise_affine=True, memory_efficient=False):
+        super().__init__()
+        self.dim = dim
+        self.eps = eps
+        self.elementwise_affine = elementwise_affine
+        if self.elementwise_affine:
+            self.weight = nn.Parameter(torch.ones(dim))
+        else:
+            self.register_parameter('weight', None)
+
+    def _norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        output = self._norm(x.float()).type_as(x)
+        if self.weight is not None:
+            output = output * self.weight
+        return output
+
+    def extra_repr(self) -> str:
+        return f'dim={self.dim}, eps={self.eps}, elementwise_affine={self.elementwise_affine}'
+    
+def modulate(x, shift, scale):
+    """Apply modulation to input tensor."""
+    return x * (1 + scale) + shift
+
+
+class TimestepEmbedder(nn.Module):
+    """
+    Embeds scalar timesteps into vector representations.
+    
+    Args:
+        hidden_size (`int`): Size of the output embedding
+        frequency_embedding_size (`int`, optional): Size of the intermediate frequency embedding
+    """
+    def __init__(self, hidden_size, frequency_embedding_size=256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=False),
+            # nn.SiLU(),
+            ACT2FN['silu'],
+            nn.Linear(hidden_size, hidden_size, bias=False),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    @staticmethod
+    def timestep_embedding(t, dim, max_period=10000):
+        """
+        Create sinusoidal timestep embeddings.
+        
+        Args:
+            t (`torch.Tensor`): A 1-D Tensor of N indices, one per batch element.
+                            These may be fractional.
+            dim (`int`): The dimension of the output.
+            max_period (`int`, optional): Controls the minimum frequency of the embeddings.
+            
+        Returns:
+            `torch.Tensor`: An [N, D] Tensor of positional embeddings.
+        """
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+        ).to(t.device)
+        args = t[:, None].float() * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        return embedding.to(t.dtype)
+
+    def forward(self, t):
+        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+        t_emb = self.mlp(t_freq)
+        return t_emb
+
+
+class FeedForwardNetwork(nn.Module):
+    """
+    Standard feed-forward network with SwiGLU activation.
+    
+    Args:
+        embed_dim (`int`): Input dimension
+        ffn_dim (`int`): Hidden dimension
+    """
+    def __init__(
+        self,
+        embed_dim,
+        ffn_dim,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.gate_proj = nn.Linear(self.embed_dim, ffn_dim, bias=False)
+        self.up_proj = nn.Linear(self.embed_dim, ffn_dim, bias=False)
+        self.down_proj = nn.Linear(ffn_dim, self.embed_dim, bias=False)
+        self.act_fn = ACT2FN['silu']  # Using SiLU as the activation function
+
+    def forward(self, x):
+        gate = self.gate_proj(x)
+        up = self.up_proj(x)
+        
+        # SwiGLU activation
+        # gate = F.silu(gate)
+        gate = self.act_fn(gate)
+        return self.down_proj(gate * up)
+
+    
+class HeadLayer(nn.Module):
+    """
+    A layer in the diffusion head.
+    
+    Args:
+        embed_dim (`int`): Input dimension
+        ffn_dim (`int`): Hidden dimension
+        cond_dim (`int`): Condition embedding dimension
+        norm_eps (`float`, optional): Epsilon for normalization
+    """
+    def __init__(
+        self,
+        embed_dim,
+        ffn_dim,
+        cond_dim,
+        norm_eps=1e-5,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.cond_dim = cond_dim
+        self.ffn_dim = ffn_dim
+        self.ffn = FeedForwardNetwork(
+            self.embed_dim,
+            self.ffn_dim,
+        )
+        self.norm = RMSNorm(self.embed_dim, eps=norm_eps)
+        self.adaLN_modulation = nn.Sequential(
+            # nn.SiLU(),
+            ACT2FN['silu'],
+            nn.Linear(cond_dim, 3 * self.embed_dim, bias=False)
+        )
+
+    def forward(self, x, c):
+        shift_ffn, scale_ffn, gate_ffn = self.adaLN_modulation(c).chunk(3, dim=-1)
+        x = x + gate_ffn * self.ffn(modulate(self.norm(x), shift_ffn, scale_ffn))
+        return x
+
+
+class FinalLayer(nn.Module):
+    """
+    Final layer in the diffusion head.
+    
+    Args:
+        hidden_size (`int`): Input dimension
+        output_size (`int`): Output dimension
+        cond_size (`int`): Condition embedding dimension
+        norm_eps (`float`, optional): Epsilon for normalization
+    """
+    def __init__(self, hidden_size, output_size, cond_size, norm_eps=1e-5):
+        super().__init__()
+        self.norm_final = RMSNorm(hidden_size, eps=norm_eps, elementwise_affine=False)
+        self.linear = nn.Linear(hidden_size, output_size, bias=False)
+        self.adaLN_modulation = nn.Sequential(
+            # nn.SiLU(),
+            ACT2FN['silu'],
+            nn.Linear(cond_size, 2 * hidden_size, bias=False)
+        )
+
+    def forward(self, x, c):
+        shift, scale = self.adaLN_modulation(c).chunk(2, dim=-1)
+        x = modulate(self.norm_final(x), shift, scale)
+        x = self.linear(x)
+        return x
+
+
+class VibeVoiceDiffusionHead(PreTrainedModel):
+    """
+    Diffusion head model for vibevoice.
+    
+    Args:
+        config (`VibeVoiceDiffusionHeadConfig`): Model configuration
+        latent_size (`int`, optional): Size of the latent space. If not provided, uses `config.latent_size`.
+    """
+    config_class = VibeVoiceDiffusionHeadConfig
+    supports_gradient_checkpointing = True
+    _supports_flash_attn_2 = True  
+    _supports_sdpa = True  
+    
+    def __init__(
+        self,
+        config,
+    ):
+        super().__init__(config)
+        self.config = config
+        self.cond_dim = config.hidden_size
+        latent_size = config.latent_size
+        
+        self.noisy_images_proj = nn.Linear(latent_size, config.hidden_size, bias=False)
+        self.cond_proj = nn.Linear(config.hidden_size, self.cond_dim, bias=False)
+        self.t_embedder = TimestepEmbedder(self.cond_dim)
+        
+        ffn_dim = int(config.hidden_size * config.head_ffn_ratio)
+        
+        # Create the intermediate layers
+        self.layers = nn.ModuleList([
+            HeadLayer(
+                embed_dim=config.hidden_size,
+                ffn_dim=ffn_dim,
+                cond_dim=self.cond_dim,
+                norm_eps=config.rms_norm_eps
+            )
+            for _ in range(config.head_layers)
+        ])
+        
+        # Final layer for output
+        self.final_layer = FinalLayer(
+            hidden_size=config.hidden_size, 
+            output_size=latent_size,
+            cond_size=self.cond_dim,
+            norm_eps=config.rms_norm_eps
+        )
+        
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        """Initialize the weights of the model."""
+        # Initialize timestep embedder
+        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+
+        # Zero-out adaLN modulation layers
+        for layer in self.layers:
+            nn.init.constant_(layer.adaLN_modulation[-1].weight, 0)
+
+        # Zero-out output layers
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
+        nn.init.constant_(self.final_layer.linear.weight, 0)
+
+    def forward(
+        self,
+        noisy_images,
+        timesteps,
+        condition,
+    ):
+        """
+        Forward pass of the prediction head.
+        
+        Args:
+            noisy_images (`torch.Tensor`): Noisy images/latents to denoise
+            timesteps (`torch.Tensor`): Timesteps for diffusion
+            condition (`torch.Tensor`): Conditioning information
+            
+        Returns:
+            `torch.Tensor`: The predicted noise/velocity
+        """
+        x = self.noisy_images_proj(noisy_images)
+        t = self.t_embedder(timesteps)
+        condition = self.cond_proj(condition)
+        c = condition + t
+        
+        for layer in self.layers:
+            x = layer(x, c)
+            
+        x = self.final_layer(x, c)
+        return x
+
+
+AutoModel.register(VibeVoiceDiffusionHeadConfig, VibeVoiceDiffusionHead)
+
+__all__ = [
+    "VibeVoiceDiffusionHead",
+]
\ No newline at end of file
diff --git a/vibevoice/modular/modular_vibevoice_text_tokenizer.py b/vibevoice/modular/modular_vibevoice_text_tokenizer.py
new file mode 100644
index 0000000..9532d9f
--- /dev/null
+++ b/vibevoice/modular/modular_vibevoice_text_tokenizer.py
@@ -0,0 +1,313 @@
+"""Tokenization classes for vibevoice."""
+
+from typing import List, Optional, Union
+
+from transformers.utils import logging
+from transformers.models.qwen2.tokenization_qwen2 import Qwen2Tokenizer
+from transformers.models.qwen2.tokenization_qwen2_fast import Qwen2TokenizerFast
+
+logger = logging.get_logger(__name__)
+
+
+class VibeVoiceTextTokenizer(Qwen2Tokenizer):
+    """
+    Construct a VibeVoice tokenizer. Based on the Qwen2 tokenizer with additional special tokens for speech.
+    
+    Args:
+        vocab_file (`str`):
+            Path to the vocabulary file.
+        merges_file (`str`):
+            Path to the merges file.
+        errors (`str`, *optional*, defaults to `"replace"`):
+            Paradigm to follow when decoding bytes to UTF-8.
+        unk_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
+            The unknown token.
+        bos_token (`str`, *optional*):
+            The beginning of sequence token. Not used for vibevoice.
+        eos_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
+            The end of sequence token.
+        pad_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
+            The token used for padding.
+        add_special_tokens (`bool`, *optional*, defaults to `True`):
+            Whether or not to add special tokens when encoding.
+    """
+
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(
+        self,
+        vocab_file,
+        merges_file,
+        errors="replace",
+        unk_token="<|endoftext|>",
+        bos_token=None,
+        eos_token="<|endoftext|>",
+        pad_token="<|endoftext|>",
+        add_prefix_space=False,
+        add_special_tokens=True,
+        **kwargs,
+    ):
+        super().__init__(
+            vocab_file=vocab_file,
+            merges_file=merges_file,
+            errors=errors,
+            unk_token=unk_token,
+            bos_token=bos_token,
+            eos_token=eos_token,
+            pad_token=pad_token,
+            add_prefix_space=add_prefix_space,
+            add_special_tokens=add_special_tokens,
+            **kwargs,
+        )
+        
+        # Add VibeVoice-specific special tokens
+        self._add_vibevoice_special_tokens()
+        
+    def _add_vibevoice_special_tokens(self):
+        """Add VibeVoice-specific special tokens."""
+        special_tokens = {
+            "additional_special_tokens": [
+                "<|vision_start|>",  # Speech start (reusing vision tokens)
+                "<|vision_end|>",  # Speech end
+                "<|vision_pad|>",  # Speech diffusion pad
+            ]
+        }
+        num_added = self.add_special_tokens(special_tokens)
+        
+        # Cache special token IDs
+        self._speech_start_id = self.convert_tokens_to_ids("<|vision_start|>")
+        self._speech_end_id = self.convert_tokens_to_ids("<|vision_end|>")
+        self._speech_diffusion_id = self.convert_tokens_to_ids("<|vision_pad|>")
+        
+        self._eos_id = self.convert_tokens_to_ids('<|endoftext|>')
+
+        return num_added
+    
+    @property
+    def eos_id(self) -> int:
+        """ID of the end of sequence token."""
+        return self._eos_id
+    
+    @property
+    def speech_start_id(self) -> int:
+        """ID of the speech start token."""
+        return self._speech_start_id
+    
+    @property
+    def speech_end_id(self) -> int:
+        """ID of the speech end token."""
+        return self._speech_end_id
+    
+    @property
+    def speech_diffusion_id(self) -> int:
+        """ID of the speech diffusion token."""
+        return self._speech_diffusion_id
+    
+    @property
+    def pad_id(self) -> int:
+        """ID used for padding (returns -100 for loss masking)."""
+        return -100
+
+
+class VibeVoiceTextTokenizerFast(Qwen2TokenizerFast):
+    """
+    Construct a "fast" VibeVoice tokenizer (backed by HuggingFace's *tokenizers* library).
+    Based on the Qwen2 tokenizer with additional special tokens for speech.
+    
+    Args:
+        vocab_file (`str`, *optional*):
+            Path to the vocabulary file.
+        merges_file (`str`, *optional*):
+            Path to the merges file.
+        tokenizer_file (`str`, *optional*):
+            Path to [tokenizers](https://github.com/huggingface/tokenizers) file.
+        unk_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
+            The unknown token.
+        bos_token (`str`, *optional*):
+            The beginning of sequence token. Not used for vibevoice.
+        eos_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
+            The end of sequence token.
+        pad_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
+            The token used for padding.
+    """
+
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(
+        self,
+        vocab_file=None,
+        merges_file=None,
+        tokenizer_file=None,
+        unk_token="<|endoftext|>",
+        bos_token=None,
+        eos_token="<|endoftext|>",
+        pad_token="<|endoftext|>",
+        add_prefix_space=False,
+        **kwargs,
+    ):
+        super().__init__(
+            vocab_file=vocab_file,
+            merges_file=merges_file,
+            tokenizer_file=tokenizer_file,
+            unk_token=unk_token,
+            bos_token=bos_token,
+            eos_token=eos_token,
+            pad_token=pad_token,
+            add_prefix_space=add_prefix_space,
+            **kwargs,
+        )
+        
+        # Add VibeVoice-specific special tokens
+        self._add_vibevoice_special_tokens()
+        
+    def _add_vibevoice_special_tokens(self):
+        """Add VibeVoice-specific special tokens."""
+        special_tokens = {
+            "additional_special_tokens": [
+                "<|vision_start|>",  # Speech start (reusing vision tokens)
+                "<|vision_end|>",  # Speech end
+                "<|vision_pad|>",  # Speech diffusion pad
+            ]
+        }
+        num_added = self.add_special_tokens(special_tokens)
+        
+        # Cache special token IDs
+        self._speech_start_id = self.convert_tokens_to_ids("<|vision_start|>")
+        self._speech_end_id = self.convert_tokens_to_ids("<|vision_end|>")
+        self._speech_diffusion_id = self.convert_tokens_to_ids("<|vision_pad|>")
+
+        # self._eos_id = self.convert_tokens_to_ids('<|endoftext|>')
+        self._eos_id = self.eos_token_id # qwen2 / qwen3
+        self._pad_id = self.convert_tokens_to_ids('<|image_pad|>')
+        
+        return num_added
+    
+    @property
+    def eos_id(self) -> int:
+        """ID of the end of sequence token."""
+        return self._eos_id
+    
+    @property
+    def speech_start_id(self) -> int:
+        """ID of the speech start token."""
+        return self._speech_start_id
+    
+    @property
+    def speech_end_id(self) -> int:
+        """ID of the speech end token."""
+        return self._speech_end_id
+    
+    @property
+    def speech_diffusion_id(self) -> int:
+        """ID of the speech diffusion token."""
+        return self._speech_diffusion_id
+    
+    @property
+    def pad_id(self) -> int:
+        """ID used for padding (returns -100 for loss masking)."""
+        return self._pad_id
+
+class VibeVoiceASRTextTokenizerFast(Qwen2TokenizerFast):
+    """
+    Construct a "fast" VibeVoice tokenizer (backed by HuggingFace's *tokenizers* library).
+    Based on the Qwen2 tokenizer with additional special tokens for speech.
+    
+    Args:
+        vocab_file (`str`, *optional*):
+            Path to the vocabulary file.
+        merges_file (`str`, *optional*):
+            Path to the merges file.
+        tokenizer_file (`str`, *optional*):
+            Path to [tokenizers](https://github.com/huggingface/tokenizers) file.
+        unk_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
+            The unknown token.
+        bos_token (`str`, *optional*):
+            The beginning of sequence token. Not used for vibevoice.
+        eos_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
+            The end of sequence token.
+        pad_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
+            The token used for padding.
+    """
+
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(
+        self,
+        vocab_file=None,
+        merges_file=None,
+        tokenizer_file=None,
+        unk_token="<|endoftext|>",
+        bos_token=None,
+        eos_token="<|endoftext|>",
+        pad_token="<|endoftext|>",
+        add_prefix_space=False,
+        **kwargs,
+    ):
+        super().__init__(
+            vocab_file=vocab_file,
+            merges_file=merges_file,
+            tokenizer_file=tokenizer_file,
+            unk_token=unk_token,
+            bos_token=bos_token,
+            eos_token=eos_token,
+            pad_token=pad_token,
+            add_prefix_space=add_prefix_space,
+            **kwargs,
+        )
+        
+        # Add VibeVoice-specific special tokens
+        self._add_vibevoice_special_tokens()
+        
+        # https://github.com/QwenLM/Qwen2.5-VL/blob/d2240f11656bfe404b9ba56db4e51cd09f522ff1/qwen-vl-finetune/qwenvl/data/data_qwen_packed.py#L57C5-L57C222
+        self.chat_template = "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}"
+        
+    def _add_vibevoice_special_tokens(self):
+        """Add VibeVoice-specific special tokens."""
+        special_tokens = {
+            "additional_special_tokens": [
+                "<|object_ref_start|>",  # Speech start (reusing vision tokens)
+                "<|object_ref_end|>",  # Speech end
+                "<|box_start|>",  # Speech diffusion pad
+            ]
+        }
+        num_added = self.add_special_tokens(special_tokens)
+        
+        # Cache special token IDs
+        self._speech_start_id = self.convert_tokens_to_ids("<|object_ref_start|>")
+        self._speech_end_id = self.convert_tokens_to_ids("<|object_ref_end|>")
+        self._speech_pad_id = self.convert_tokens_to_ids("<|box_start|>")
+
+        self._eos_id = self.eos_token_id # qwen2 / qwen3
+        self._pad_id = self.convert_tokens_to_ids('<|image_pad|>')
+        
+        return num_added    
+    
+    @property
+    def eos_id(self) -> int:
+        """ID of the end of sequence token."""
+        return self._eos_id
+    
+    @property
+    def speech_start_id(self) -> int:
+        """ID of the speech start token."""
+        return self._speech_start_id
+    
+    @property
+    def speech_end_id(self) -> int:
+        """ID of the speech end token."""
+        return self._speech_end_id
+    
+    @property
+    def speech_pad_id(self) -> int:
+        """ID of the speech diffusion token."""
+        return self._speech_pad_id
+    
+    @property
+    def pad_id(self) -> int:
+        return self._pad_id
+    
+__all__ = [
+    "VibeVoiceTextTokenizer", 
+    "VibeVoiceTextTokenizerFast", 
+    "VibeVoiceASRTextTokenizerFast",
+]
\ No newline at end of file
diff --git a/vibevoice/modular/modular_vibevoice_tokenizer.py b/vibevoice/modular/modular_vibevoice_tokenizer.py
new file mode 100644
index 0000000..454f9c1
--- /dev/null
+++ b/vibevoice/modular/modular_vibevoice_tokenizer.py
@@ -0,0 +1,1207 @@
+import math
+import typing as tp
+from functools import partial
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Tuple, Union
+import copy
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers.models.auto import AutoModel
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+from transformers.modeling_utils import PreTrainedModel
+from transformers.activations import ACT2FN
+
+from .configuration_vibevoice import VibeVoiceAcousticTokenizerConfig, VibeVoiceSemanticTokenizerConfig
+
+logger = logging.get_logger(__name__)
+
+import os
+# Try to import APEX FusedRMSNorm
+try:
+    from apex.normalization.fused_layer_norm import fused_rms_norm_affine
+    APEX_AVAILABLE = True
+    # logger.info("APEX FusedRMSNorm is available and will be used for optimization")
+    if int(os.getenv("OPTIMIZE_FOR_SPEED", "0")) == 0:
+        APEX_AVAILABLE = False
+        # logger.warning("APEX FusedRMSNorm is disabled by environment variable OPTIMIZE_FOR_SPEED=0")
+except ImportError:
+    APEX_AVAILABLE = False
+    # logger.warning("APEX FusedRMSNorm not available, using native implementation")
+
+# Normalization modules
+class ConvLayerNorm(nn.LayerNorm):
+    """
+    Convolution-friendly LayerNorm that moves channels to last dimensions
+    before running the normalization and moves them back to original position right after.
+    """
+    def __init__(self, normalized_shape: tp.Union[int, tp.List[int], torch.Size], **kwargs):
+        super().__init__(normalized_shape, **kwargs)
+
+    def forward(self, x):
+        x = x.transpose(1, 2)  # b ... t -> b t ...
+        x = nn.functional.layer_norm(x.float(), self.normalized_shape, self.weight.float(), self.bias.float(), self.eps).type_as(x) 
+        x = x.transpose(1, 2)  # b t ... -> b ... t
+        return x
+    
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-5, elementwise_affine=True, weight_shape=None):
+        super().__init__()
+        self.dim = dim
+        self.eps = eps
+        self.elementwise_affine = elementwise_affine
+        if self.elementwise_affine:
+            weight_shape = (dim,) if weight_shape is None else weight_shape
+            self.weight = nn.Parameter(torch.ones(weight_shape))
+        else:
+            self.register_parameter('weight', None)
+
+    def _norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        output = self._norm(x.float()).type_as(x)
+        if self.weight is not None:
+            output = output * self.weight
+        return output
+
+    def extra_repr(self) -> str:
+        return f'dim={self.dim}, eps={self.eps}, elementwise_affine={self.elementwise_affine}'
+
+class ConvRMSNorm(RMSNorm):
+    def __init__(self, dim: int, eps: float = 1e-5, elementwise_affine=True, weight_shape=None):
+        super().__init__(dim, eps, elementwise_affine, weight_shape)
+
+    def forward(self, x):
+        x = x.transpose(1, 2)  # b ... t -> b t ...
+        if (not APEX_AVAILABLE) or (not self.elementwise_affine):
+            # Fallback to native implementation
+            output = self._norm(x.float()).type_as(x)
+            if self.weight is not None:
+                output = output * self.weight
+        else:
+            output = fused_rms_norm_affine(x, self.weight, self.weight.shape, self.eps)
+        output = output.transpose(1, 2)  # b t ... -> b ... t
+        return output
+
+# Convolutional layers and utilities
+CONV_NORMALIZATIONS = frozenset(['none', 'weight_norm', 'spectral_norm',
+                                'time_layer_norm', 'layer_norm', 'time_group_norm'])
+
+
+def apply_parametrization_norm(module: nn.Module, norm: str = 'none') -> nn.Module:
+    assert norm in CONV_NORMALIZATIONS
+    if norm == 'weight_norm':
+        return nn.utils.weight_norm(module)
+    elif norm == 'spectral_norm':
+        return nn.utils.spectral_norm(module)
+    else:
+        # We already check was in CONV_NORMALIZATION, so any other choice
+        # doesn't need reparametrization.
+        return module
+
+
+def get_norm_module(module: nn.Module, causal: bool = False, norm: str = 'none', **norm_kwargs) -> nn.Module:
+    """Return the proper normalization module. If causal is True, this will ensure the returned
+    module is causal, or return an error if the normalization doesn't support causal evaluation.
+    """
+    assert norm in CONV_NORMALIZATIONS
+    if norm == 'layer_norm':
+        assert isinstance(module, nn.modules.conv._ConvNd)
+        return ConvLayerNorm(module.out_channels, **norm_kwargs)
+    elif norm == 'time_group_norm':
+        if causal:
+            raise ValueError("GroupNorm doesn't support causal evaluation.")
+        assert isinstance(module, nn.modules.conv._ConvNd)
+        return nn.GroupNorm(1, module.out_channels, **norm_kwargs)
+    else:
+        return nn.Identity()
+
+
+def get_extra_padding_for_conv1d(x: torch.Tensor, kernel_size: int, stride: int,
+                                padding_total: int = 0) -> int:
+    """Calculate extra padding needed for convolution to have the same output length"""
+    length = x.shape[-1]
+    n_frames = (length - kernel_size + padding_total) / stride + 1
+    ideal_length = (math.ceil(n_frames) - 1) * stride + (kernel_size - padding_total)
+    return ideal_length - length
+
+
+def pad1d(x: torch.Tensor, paddings: tp.Tuple[int, int], mode: str = 'zero', value: float = 0.):
+    """Pad 1D input with handling for small inputs in reflect mode"""
+    length = x.shape[-1]
+    padding_left, padding_right = paddings
+    assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)
+    if mode == 'reflect':
+        max_pad = max(padding_left, padding_right)
+        extra_pad = 0
+        if length <= max_pad:
+            extra_pad = max_pad - length + 1
+            x = F.pad(x, (0, extra_pad))
+        padded = F.pad(x, paddings, mode, value)
+        end = padded.shape[-1] - extra_pad
+        return padded[..., :end]
+    else:
+        return F.pad(x, paddings, mode, value)
+
+
+def unpad1d(x: torch.Tensor, paddings: tp.Tuple[int, int]):
+    """Remove padding from x, handling properly zero padding. Only for 1d!"""
+    padding_left, padding_right = paddings
+    assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)
+    assert (padding_left + padding_right) <= x.shape[-1]
+    end = x.shape[-1] - padding_right
+    return x[..., padding_left: end]
+
+
+class NormConv1d(nn.Module):
+    """Wrapper around Conv1d and normalization applied to this conv"""
+    def __init__(self, *args, causal: bool = False, norm: str = 'none',
+                norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):
+        super().__init__()
+        self.conv = apply_parametrization_norm(nn.Conv1d(*args, **kwargs), norm)
+        self.norm = get_norm_module(self.conv, causal, norm, **norm_kwargs)
+        self.norm_type = norm
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.norm(x)
+        return x
+
+
+class NormConvTranspose1d(nn.Module):
+    """Wrapper around ConvTranspose1d and normalization applied to this conv"""
+    def __init__(self, *args, causal: bool = False, norm: str = 'none',
+                norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):
+        super().__init__()
+        self.convtr = apply_parametrization_norm(nn.ConvTranspose1d(*args, **kwargs), norm)
+        self.norm = get_norm_module(self.convtr, causal, norm, **norm_kwargs)
+        self.norm_type = norm
+
+    def forward(self, x):
+        x = self.convtr(x)
+        x = self.norm(x)
+        return x
+
+
+class VibeVoiceTokenizerStreamingCache:
+    """Cache for streaming convolution, similar to KV cache in attention"""
+    def __init__(self):
+        self.cache = {}  # Dict mapping (layer_id, sample_idx) to state tensor
+        
+    def get(self, layer_id: str, sample_indices: torch.Tensor) -> Optional[torch.Tensor]:
+        """Get cached states for given layer and sample indices"""
+        states = []
+        max_length = 0
+        
+        # First pass: collect states and find max length
+        for idx in sample_indices.tolist():
+            key = (layer_id, idx)
+            if key not in self.cache:
+                return None  # If any sample is missing, return None
+            state = self.cache[key]
+            states.append(state)
+            max_length = max(max_length, state.shape[-1])
+        
+        # Second pass: pad states to max length if needed
+        if len(states) > 0 and states[0].dim() >= 2:
+            padded_states = []
+            for state in states:
+                if state.shape[-1] < max_length:
+                    # Pad on the time dimension (last dimension)
+                    pad_size = max_length - state.shape[-1]
+                    # Pad with zeros on the LEFT to align the most recent samples
+                    padded_state = F.pad(state, (pad_size, 0), mode='constant', value=0)
+                    padded_states.append(padded_state)
+                else:
+                    padded_states.append(state)
+            return torch.stack(padded_states, dim=0)
+        else:
+            return torch.stack(states, dim=0)
+    
+    def set(self, layer_id: str, sample_indices: torch.Tensor, states: torch.Tensor):
+        """Set cached states for given layer and sample indices"""
+        for i, idx in enumerate(sample_indices.tolist()):
+            key = (layer_id, idx)
+            self.cache[key] = states[i].detach()
+
+    def set_to_zero(self, sample_indices: torch.Tensor):
+        """Set all cached states to zero for given sample indices"""
+        for key in list(self.cache.keys()):
+            layer_id, sample_idx = key
+            if sample_idx in sample_indices.tolist():
+                # Create zero tensor with same shape and dtype as cached tensor
+                cached_tensor = self.cache[key]
+                self.cache[key] = torch.zeros_like(cached_tensor)
+                
+    def clear(self, layer_id: Optional[str] = None, sample_indices: Optional[torch.Tensor] = None):
+        """Clear cache for specific layer/samples or everything"""
+        if layer_id is None and sample_indices is None:
+            self.cache.clear()
+        elif layer_id is not None and sample_indices is None:
+            # Clear all samples for a specific layer
+            keys_to_remove = [k for k in self.cache.keys() if k[0] == layer_id]
+            for k in keys_to_remove:
+                del self.cache[k]
+        elif layer_id is not None and sample_indices is not None:
+            # Clear specific samples for a specific layer
+            for idx in sample_indices.tolist():
+                key = (layer_id, idx)
+                self.cache.pop(key, None)
+
+class SConv1d(nn.Module):
+    """Conv1d with built-in handling of asymmetric or causal padding and normalization."""
+    def __init__(self, in_channels: int, out_channels: int,
+                kernel_size: int, stride: int = 1, dilation: int = 1,
+                groups: int = 1, bias: bool = True, causal: bool = False,
+                norm: str = 'none', norm_kwargs: tp.Dict[str, tp.Any] = {},
+                pad_mode: str = 'reflect'):
+        super().__init__()
+        self.conv = NormConv1d(in_channels, out_channels, kernel_size, stride,
+                            dilation=dilation, groups=groups, bias=bias, causal=causal,
+                            norm=norm, norm_kwargs=norm_kwargs)
+        self.causal = causal
+        self.pad_mode = pad_mode
+        
+        # Store configuration
+        self.kernel_size = kernel_size
+        self.dilation = dilation
+        self.stride = stride
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        
+        # For causal convolution, we need to maintain kernel_size - 1 samples as context
+        # need to check use which context_size is more suitable
+        # self.context_size = (kernel_size - 1) * dilation
+        self.context_size = (kernel_size - 1) * dilation - (stride - 1)
+        
+        # For non-streaming mode, calculate padding
+        self.padding_total = (kernel_size - 1) * dilation - (stride - 1)
+        
+        # Create a unique layer ID for cache management
+        self._layer_id = None
+                  
+    @property
+    def layer_id(self):
+        if self._layer_id is None:
+            self._layer_id = f"sconv1d_{id(self)}"
+        return self._layer_id
+        
+    def forward(self, x: torch.Tensor, 
+                cache: Optional[VibeVoiceTokenizerStreamingCache] = None,
+                sample_indices: Optional[torch.Tensor] = None,
+                use_cache: bool = False,
+                debug: bool = False,
+                is_final_chunk: bool = False) -> torch.Tensor:
+        """
+        Forward pass with optional streaming support via cache.
+        
+        Args:
+            x: Input tensor [batch_size, channels, time]
+            cache: VibeVoiceTokenizerStreamingCache object for maintaining states
+            sample_indices: Indices identifying each sample for cache management
+            use_cache: Whether to use cached states for streaming
+            debug: Whether to print debug information
+            is_final_chunk: Whether this is the final chunk (adds extra padding for alignment)
+            
+        Returns:
+            Output tensor
+        """
+        B, C, T = x.shape
+        
+        # Non-streaming mode
+        if not use_cache or cache is None:
+            return self._forward_non_streaming(x, debug=debug)
+        
+        # Streaming mode
+        assert self.causal, "Streaming mode is only supported for causal convolutions"
+        assert sample_indices is not None, "sample_indices must be provided for streaming mode"
+        assert len(sample_indices) == B, "sample_indices must match batch size"
+        
+        return self._forward_streaming(x, cache, sample_indices, debug, is_final_chunk)
+    
+    def _forward_streaming(self, x: torch.Tensor, 
+                          cache: VibeVoiceTokenizerStreamingCache,
+                          sample_indices: torch.Tensor,
+                          debug: bool = False,
+                          is_final_chunk: bool = False) -> torch.Tensor:
+        """Streaming forward pass with cache operations kept separate from compiled code"""
+        B, C, T = x.shape
+        
+        # Cache operations (not compiled)
+        cached_states = cache.get(self.layer_id, sample_indices)
+        
+        if cached_states is None:
+            # First chunk - initialize with zeros for context
+            if self.context_size > 0:
+                cached_states = torch.zeros(B, C, self.context_size, device=x.device, dtype=x.dtype)
+                if debug:
+                    print(f"[DEBUG] Initialized cache with shape: {cached_states.shape}, context_size={self.context_size}")
+            else:
+                cached_states = torch.zeros(B, C, 0, device=x.device, dtype=x.dtype)
+                if debug:
+                    print(f"[DEBUG] No context needed (kernel_size=stride)")
+        
+        # Concatenate cached states with input
+        if cached_states.shape[2] > 0:
+            input_with_context = torch.cat([cached_states, x], dim=2)
+        else:
+            input_with_context = x
+        
+        # For final chunk, add extra padding to ensure ceil behavior (same as non-streaming)
+        if is_final_chunk:
+            extra_padding = get_extra_padding_for_conv1d(
+                input_with_context, self.kernel_size, self.stride, self.padding_total
+            )
+            if extra_padding > 0:
+                input_with_context = pad1d(input_with_context, (0, extra_padding), mode=self.pad_mode)
+                if debug:
+                    print(f"[DEBUG] Final chunk: added extra_padding={extra_padding}")
+            
+        if debug:
+            print(f"[DEBUG] Input shape: {x.shape}, Cache shape: {cached_states.shape}, Combined: {input_with_context.shape}")
+        
+        # Apply convolution directly - no extra padding in streaming mode
+        # The conv layer will handle its own padding internally
+        output = self.conv(input_with_context)
+
+        if debug:
+            print(f"[DEBUG] Output shape: {output.shape}")
+        
+        # Update cache for next chunk
+        if self.context_size > 0:
+            # Calculate how many samples to keep
+            total_input_length = input_with_context.shape[2]
+            
+            # Keep the last context_size samples
+            if total_input_length >= self.context_size:
+                new_cache_start = total_input_length - self.context_size
+                new_cache = input_with_context[:, :, new_cache_start:]
+            else:
+                # If we have less than context_size samples, keep everything
+                new_cache = input_with_context
+                
+            if debug:
+                print(f"[DEBUG] New cache shape: {new_cache.shape}")
+                
+            cache.set(self.layer_id, sample_indices, new_cache)
+        
+        return output
+    
+    def _forward_non_streaming(self, x: torch.Tensor, debug: bool = False) -> torch.Tensor:
+        """Standard forward pass without streaming"""
+        B, C, T = x.shape
+        kernel_size = self.kernel_size
+        stride = self.stride
+        dilation = self.dilation
+        padding_total = self.padding_total
+        
+        # Compute extra padding for stride alignment
+        extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total)
+        
+        if debug:
+            print(f"[DEBUG NON-STREAMING] Input shape: {x.shape}, padding_total={padding_total}, extra_padding={extra_padding}")
+        
+        if self.causal:
+            # Left padding for causal
+            if self.pad_mode == 'constant':
+                x = pad1d(x, (padding_total, extra_padding), mode=self.pad_mode, value=0)
+            else:
+                x = pad1d(x, (padding_total, extra_padding), mode=self.pad_mode)
+        else:
+            # Symmetric padding for non-causal
+            padding_right = padding_total // 2
+            padding_left = padding_total - padding_right
+            x = pad1d(x, (padding_left, padding_right + extra_padding), mode=self.pad_mode)
+        
+        if debug:
+            print(f"[DEBUG NON-STREAMING] After padding: {x.shape}")
+            
+        output = self.conv(x)
+        
+        if debug:
+            print(f"[DEBUG NON-STREAMING] Output shape: {output.shape}")
+        
+        return output
+
+
+class SConvTranspose1d(nn.Module):
+    """ConvTranspose1d with built-in handling of asymmetric or causal padding and normalization."""
+    def __init__(self, in_channels: int, out_channels: int,
+                kernel_size: int, stride: int = 1, causal: bool = False,
+                norm: str = 'none', trim_right_ratio: float = 1.,
+                norm_kwargs: tp.Dict[str, tp.Any] = {}, bias: bool = True):
+        super().__init__()
+        self.convtr = NormConvTranspose1d(in_channels, out_channels, kernel_size, stride,
+                                        causal=causal, norm=norm, norm_kwargs=norm_kwargs, bias=bias)
+        self.causal = causal
+        self.trim_right_ratio = trim_right_ratio
+        assert self.causal or self.trim_right_ratio == 1., \
+            "`trim_right_ratio` != 1.0 only makes sense for causal convolutions"
+        assert self.trim_right_ratio >= 0. and self.trim_right_ratio <= 1.
+
+        # Store configuration
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        
+        # For transposed convolution, padding calculation is different
+        self.padding_total = kernel_size - stride
+        
+        # For streaming, we need to keep track of input history
+        # Transposed conv needs to see multiple input samples to produce correct output
+        self.context_size = kernel_size - 1
+        
+        # Create a unique layer ID for cache management
+        self._layer_id = None
+
+    @property
+    def layer_id(self):
+        if self._layer_id is None:
+            self._layer_id = f"sconvtr1d_{id(self)}"
+        return self._layer_id
+    
+    def forward(self, x: torch.Tensor,
+                cache: Optional[VibeVoiceTokenizerStreamingCache] = None,
+                sample_indices: Optional[torch.Tensor] = None,
+                use_cache: bool = False,
+                debug: bool = False) -> torch.Tensor:
+        """
+        Forward pass with optional streaming support via cache.
+        """
+        B, C, T = x.shape
+        
+        # Non-streaming mode
+        if not use_cache or cache is None:
+            return self._forward_non_streaming(x, debug=debug)
+        
+        # Streaming mode
+        assert sample_indices is not None, "sample_indices must be provided for streaming mode"
+        assert len(sample_indices) == B, "sample_indices must match batch size"
+        
+        return self._forward_streaming(x, cache, sample_indices, debug)
+    
+    def _forward_streaming(self, x: torch.Tensor,
+                          cache: VibeVoiceTokenizerStreamingCache,
+                          sample_indices: torch.Tensor,
+                          debug: bool = False) -> torch.Tensor:
+        """Streaming forward pass with cache operations kept separate from compiled code"""
+        B, C, T = x.shape
+        
+        # Cache operations (not compiled)
+        cached_input = cache.get(self.layer_id, sample_indices)
+        
+        if cached_input is None:
+            # First chunk - no history yet
+            cached_input = torch.zeros(B, C, 0, device=x.device, dtype=x.dtype)
+            if debug:
+                print(f"[DEBUG] Initialized empty cache for transposed conv")
+        
+        # Concatenate cached input with new input
+        full_input = torch.cat([cached_input, x], dim=2)
+        
+        if debug:
+            print(f"[DEBUG] Input shape: {x.shape}, Cache shape: {cached_input.shape}, Combined: {full_input.shape}")
+        
+        # First chunk or debug mode - use uncompiled version
+        full_output = self.convtr(full_input)
+        
+        if debug:
+            print(f"[DEBUG] Full transposed conv output shape: {full_output.shape}")
+        
+        # Calculate padding to remove
+        if self.causal:
+            padding_right = math.ceil(self.padding_total * self.trim_right_ratio)
+            padding_left = self.padding_total - padding_right
+        else:
+            padding_right = self.padding_total // 2
+            padding_left = self.padding_total - padding_right
+        
+        # Remove padding
+        if padding_left + padding_right > 0:
+            full_output = unpad1d(full_output, (padding_left, padding_right))
+        
+        if debug:
+            print(f"[DEBUG] After unpadding: {full_output.shape}")
+        
+        # Determine which part of the output corresponds to the new input
+        if cached_input.shape[2] == 0:
+            # First chunk - return all output
+            output = full_output
+        else:
+            # Subsequent chunks - return only the new output
+            expected_new_output = T * self.stride
+            
+            # Take the last expected_new_output samples
+            if full_output.shape[2] >= expected_new_output:
+                output = full_output[:, :, -expected_new_output:]
+            else:
+                output = full_output
+        
+        if debug:
+            print(f"[DEBUG] Final streaming output shape: {output.shape}")
+        
+        # Update cache
+        if full_input.shape[2] > self.context_size:
+            new_cache = full_input[:, :, -self.context_size:]
+        else:
+            new_cache = full_input
+        
+        if debug:
+            print(f"[DEBUG] New cache shape: {new_cache.shape}")
+            
+        cache.set(self.layer_id, sample_indices, new_cache)
+        
+        return output
+    
+    def _forward_non_streaming(self, x: torch.Tensor, debug: bool = False) -> torch.Tensor:
+        """Standard forward pass without streaming"""
+        if debug:
+            print(f"[DEBUG NON-STREAMING] Input shape: {x.shape}")
+        
+        # Apply transposed convolution
+        y = self.convtr(x)
+        
+        if debug:
+            print(f"[DEBUG NON-STREAMING] After transposed conv: {y.shape}")
+        
+        # Calculate and remove padding
+        if self.causal:
+            padding_right = math.ceil(self.padding_total * self.trim_right_ratio)
+            padding_left = self.padding_total - padding_right
+        else:
+            padding_right = self.padding_total // 2
+            padding_left = self.padding_total - padding_right
+        
+        if padding_left + padding_right > 0:
+            y = unpad1d(y, (padding_left, padding_right))
+        
+        if debug:
+            print(f"[DEBUG NON-STREAMING] Final output shape: {y.shape}")
+            
+        return y
+    
+# FFN 
+class FFN(nn.Module):
+    def __init__(
+        self,
+        embed_dim,
+        ffn_dim,
+        bias=False,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.linear1 = nn.Linear(self.embed_dim, ffn_dim, bias=bias) 
+        self.gelu = ACT2FN["gelu"]
+        self.linear2 = nn.Linear(ffn_dim, self.embed_dim, bias=bias)
+
+    def forward(self, x):
+        x = self.linear1(x)
+        x = self.gelu(x)
+        x = self.linear2(x)
+        return x
+
+
+class Convlayer(nn.Module):
+    def __init__(
+            self, 
+            in_channels, 
+            out_channels, 
+            kernel_size, 
+            stride=1, 
+            dilation=1, 
+            groups=1, 
+            bias=True, 
+            pad_mode='zeros', 
+            norm='weight_norm', 
+            causal=True, 
+        ):
+        super().__init__()
+        self.conv = SConv1d(in_channels, out_channels, kernel_size, stride=stride, dilation=dilation, 
+                           groups=groups, bias=bias, pad_mode=pad_mode, norm=norm, causal=causal)
+
+    def forward(self, x):
+        return self.conv(x)
+
+class Block1D(nn.Module):
+    def __init__(self, dim, kernel_size=7, drop_path=0., mixer_layer='conv',  
+                layer_scale_init_value=1e-6, **kwargs):
+        super().__init__()
+        
+        if kwargs.get('layernorm', 'LN') == 'LN':
+            self.norm = ConvLayerNorm(dim, eps=kwargs.get('eps', 1e-6))
+            self.ffn_norm = ConvLayerNorm(dim, eps=kwargs.get('eps', 1e-6))               
+        elif kwargs.get('layernorm', 'RMSNorm') == 'RMSNorm':
+            self.norm = ConvRMSNorm(dim, eps=kwargs.get('eps', 1e-6))
+            self.ffn_norm = ConvRMSNorm(dim, eps=kwargs.get('eps', 1e-6))
+
+        if mixer_layer == 'conv':
+            self.mixer = Convlayer(dim, dim, groups=kwargs.get('groups', 1),
+                                kernel_size=kernel_size, 
+                                pad_mode=kwargs.get('pad_mode', 'reflect'), 
+                                norm=kwargs.get('norm', 'none'), 
+                                causal=kwargs.get('causal', True), 
+                                bias=kwargs.get('bias', True),
+                                )
+        elif mixer_layer == 'depthwise_conv':
+            self.mixer = Convlayer(dim, dim, groups=dim,
+                                kernel_size=kernel_size, 
+                                pad_mode=kwargs.get('pad_mode', 'reflect'), 
+                                norm=kwargs.get('norm', 'none'), 
+                                causal=kwargs.get('causal', True), 
+                                bias=kwargs.get('bias', True),
+                                )
+        else:
+            raise ValueError(f"Unsupported mixer layer: {mixer_layer}")
+        
+        self.ffn = FFN(
+            dim, 
+            kwargs.get('ffn_expansion', 4) * dim, 
+            bias=kwargs.get('bias', False),
+        )
+        self.drop_path = nn.Identity() if drop_path <= 0. else nn.modules.DropPath(drop_path)
+
+        if layer_scale_init_value > 0:
+            self.gamma = nn.Parameter(layer_scale_init_value * torch.ones((dim)), requires_grad=True)
+            self.ffn_gamma = nn.Parameter(layer_scale_init_value * torch.ones((dim)), requires_grad=True)
+        else:
+            self.gamma = None
+            self.ffn_gamma = None
+
+    def forward(self, x):
+        # mixer
+        residual = x
+        x = self.norm(x)
+        x = self.mixer(x)
+        if self.gamma is not None:
+            x = x * self.gamma.unsqueeze(-1)
+        x = residual + self.drop_path(x)
+
+        # ffn
+        residual = x
+        x = self.ffn_norm(x)
+        x = x.permute(0, 2, 1)
+        x = self.ffn(x)
+        x = x.permute(0, 2, 1)
+        if self.ffn_gamma is not None:
+            x = x * self.ffn_gamma.unsqueeze(-1)
+        x = residual + self.drop_path(x)
+
+        return x
+
+
+class TokenizerEncoder(nn.Module):
+    """
+    Encoder component for the VibeVoice tokenizer that converts audio to latent representations.
+    
+    Args:
+        config: Configuration object with model parameters
+    """
+    def __init__(self, config):
+        super().__init__()
+        
+        # Extract parameters from config
+        self.channels = config.channels
+        self.dimension = config.dimension
+        self.n_filters = config.n_filters
+        self.ratios = list(reversed(config.ratios))
+        self.depths = config.depths
+        self.n_residual_layers = getattr(config, "n_residual_layers", 1)
+        self.hop_length = np.prod(self.ratios)
+        self.causal = config.causal
+        
+        # Additional config parameters with defaults
+        kernel_size = getattr(config, "kernel_size", 7)
+        last_kernel_size = getattr(config, "last_kernel_size", 7)
+        norm = getattr(config, "norm", "none")
+        norm_params = getattr(config, "norm_params", {})
+        pad_mode = getattr(config, "pad_mode", "reflect")
+        bias = getattr(config, "bias", True)
+        layernorm = getattr(config, "layernorm", "LN")
+        layernorm_eps = getattr(config, "layernorm_eps", 1e-6)
+        layernorm_elementwise_affine = getattr(config, "layernorm_elementwise_affine", True)
+        drop_path_rate = getattr(config, "drop_path_rate", 0.0)
+        mixer_layer = getattr(config, "mixer_layer", "conv")
+        layer_scale_init_value = getattr(config, "layer_scale_init_value", 0)
+        disable_last_norm = getattr(config, "disable_last_norm", False)
+        
+        # determine the norm type based on layernorm
+        if layernorm == 'LN':
+            norm_type = ConvLayerNorm
+        elif layernorm == 'RMSNorm':
+            norm_type = partial(ConvRMSNorm, elementwise_affine=layernorm_elementwise_affine)
+        else:
+            raise ValueError(f"Unsupported norm type: {layernorm}")
+        
+        # stem and intermediate downsampling conv layers
+        stem = nn.Sequential(
+                SConv1d(self.channels, self.n_filters, kernel_size, norm=norm, norm_kwargs=norm_params, causal=self.causal, pad_mode=pad_mode, bias=bias),
+            )
+        
+        self.downsample_layers = nn.ModuleList()
+        self.downsample_layers.append(stem)
+        for i in range(len(self.ratios)):
+            in_ch = self.n_filters * (2 ** i)
+            out_ch = self.n_filters * (2 ** (i + 1))
+            downsample_layer = nn.Sequential(
+                SConv1d(in_ch, out_ch, kernel_size=self.ratios[i] * 2, stride=self.ratios[i], causal=self.causal, pad_mode=pad_mode, norm=norm, bias=bias)
+            )
+            self.downsample_layers.append(downsample_layer)
+
+        # configure the transformer blocks
+        layer_type = partial(
+            Block1D,
+            mixer_layer=mixer_layer,
+            layernorm=layernorm,
+            eps=layernorm_eps,
+            causal=self.causal,
+            pad_mode=pad_mode,
+            norm=norm,
+            bias=bias,
+            layer_scale_init_value=layer_scale_init_value,
+        )
+        
+        self.stages = nn.ModuleList()
+        dp_rates = [x.item() for x in torch.linspace(0, drop_path_rate, sum(self.depths))] 
+        cur = 0
+
+        for i in range(len(self.depths)):
+            in_ch = self.n_filters * (2 ** i)
+            stage = nn.Sequential(
+                *[layer_type(dim=in_ch, drop_path=dp_rates[cur + j]) for j in range(self.depths[i])]
+            )
+            self.stages.append(stage)
+            cur += self.depths[i]
+        
+        if not disable_last_norm:
+            self.norm = norm_type(in_ch, eps=layernorm_eps)
+        else:
+            self.norm = nn.Identity()
+        self.head = SConv1d(in_ch, self.dimension, kernel_size=last_kernel_size, causal=self.causal, pad_mode=pad_mode, norm=norm, bias=bias)
+
+    def forward_features(self, x, cache=None, sample_indices=None, use_cache=False, debug=False, is_final_chunk=False):
+        for i in range(len(self.depths)):
+            # Apply downsampling
+            for layer in self.downsample_layers[i]:
+                if isinstance(layer, SConv1d):
+                    x = layer(x, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug, is_final_chunk=is_final_chunk)
+                else:
+                    x = layer(x)
+            
+            # Apply stage (Block1D contains Convlayer which contains SConv1d)
+            for block in self.stages[i]:
+                if hasattr(block, 'mixer') and hasattr(block.mixer, 'conv') and isinstance(block.mixer.conv, SConv1d):
+                    # Block1D forward with cache support
+                    residual = x
+                    x = block.norm(x)
+                    x = block.mixer.conv(x, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug, is_final_chunk=is_final_chunk)
+                    if block.gamma is not None:
+                        x = x * block.gamma.unsqueeze(-1)
+                    x = residual + x
+                    
+                    # FFN part
+                    residual = x
+                    x = block.ffn_norm(x)
+                    x = x.permute(0, 2, 1)
+                    x = block.ffn(x)
+                    x = x.permute(0, 2, 1)
+                    if block.ffn_gamma is not None:
+                        x = x * block.ffn_gamma.unsqueeze(-1)
+                    x = residual + x
+                else:
+                    x = block(x)
+
+        return self.norm(x)
+
+    def forward(self, x, cache=None, sample_indices=None, use_cache=False, debug=False, is_final_chunk=False):
+        x = self.forward_features(x, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug, is_final_chunk=is_final_chunk)
+        x = self.head(x, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug, is_final_chunk=is_final_chunk)
+        return x
+
+
+class TokenizerDecoder(nn.Module):
+    """
+    Decoder component for the VibeVoice tokenizer that converts latent representations back to audio.
+    
+    Args:
+        config: Configuration object with model parameters
+    """
+    def __init__(self, config):
+        super().__init__()
+        
+        # Extract parameters from config
+        self.dimension = config.dimension
+        self.channels = config.channels
+        self.n_filters = config.n_filters
+        self.ratios = config.ratios
+        
+        # IMPORTANT CHANGE: Don't reverse depths again since they're already reversed in VibeVoiceAcousticTokenizerModel
+        self.depths = config.depths  # Changed from list(reversed(config.depths))
+        
+        self.n_residual_layers = getattr(config, "n_residual_layers", 1)
+        self.hop_length = np.prod(self.ratios)
+        self.causal = config.causal
+        
+        # Additional config parameters with defaults
+        kernel_size = getattr(config, "kernel_size", 7)
+        last_kernel_size = getattr(config, "last_kernel_size", 7)
+        norm = getattr(config, "norm", "none")
+        norm_params = getattr(config, "norm_params", {})
+        pad_mode = getattr(config, "pad_mode", "reflect")
+        bias = getattr(config, "bias", True)
+        layernorm = getattr(config, "layernorm", "LN")
+        layernorm_eps = getattr(config, "layernorm_eps", 1e-6)
+        trim_right_ratio = getattr(config, "trim_right_ratio", 1.0)
+        layernorm_elementwise_affine = getattr(config, "layernorm_elementwise_affine", True)
+        drop_path_rate = getattr(config, "drop_path_rate", 0.0)
+        mixer_layer = getattr(config, "mixer_layer", "conv")
+        layer_scale_init_value = getattr(config, "layer_scale_init_value", 0)
+        disable_last_norm = getattr(config, "disable_last_norm", False)
+
+        # determine the norm type based on layernorm
+        if layernorm == 'LN':
+            norm_type = ConvLayerNorm
+        elif layernorm == 'RMSNorm':
+            norm_type = partial(ConvRMSNorm, elementwise_affine=layernorm_elementwise_affine)
+        else:
+            raise ValueError(f"Unsupported norm type: {layernorm}")
+        
+        # stem and upsampling layers
+        stem = nn.Sequential(
+                SConv1d(self.dimension, self.n_filters * 2 ** (len(self.depths) - 1), kernel_size, norm=norm, 
+                        norm_kwargs=norm_params, causal=self.causal, pad_mode=pad_mode, bias=bias),
+            )
+        
+        self.upsample_layers = nn.ModuleList()
+        self.upsample_layers.append(stem)
+        for i in range(len(self.ratios)):
+            in_ch = self.n_filters * (2 ** (len(self.depths) - 1 - i))
+            out_ch = self.n_filters * (2 ** (len(self.depths) - 1 - i - 1))
+            upsample_layer = nn.Sequential(
+                SConvTranspose1d(in_ch, out_ch,
+                                kernel_size=self.ratios[i] * 2, stride=self.ratios[i],
+                                norm=norm, norm_kwargs=norm_params, bias=bias,
+                                causal=self.causal, trim_right_ratio=trim_right_ratio),
+            )
+            self.upsample_layers.append(upsample_layer)
+
+        # configure transformer blocks
+        layer_type = partial(
+            Block1D,
+            mixer_layer=mixer_layer,
+            layernorm=layernorm,
+            eps=layernorm_eps,
+            causal=self.causal,
+            pad_mode=pad_mode,
+            norm=norm,
+            bias=bias,
+            layer_scale_init_value=layer_scale_init_value,
+        )
+
+        self.stages = nn.ModuleList()
+        dp_rates = [x.item() for x in torch.linspace(0, drop_path_rate, sum(self.depths))] 
+        cur = 0
+        
+        # Create stages in the same order as the original model
+        for i in range(len(self.depths)):
+            in_ch = self.n_filters * (2 ** (len(self.depths) - 1 - i))
+            stage = nn.Sequential(
+                *[layer_type(dim=in_ch, drop_path=dp_rates[cur + j]) for j in range(self.depths[i])]
+            )
+            self.stages.append(stage)
+            cur += self.depths[i]
+
+        if not disable_last_norm:
+            self.norm = norm_type(in_ch, eps=layernorm_eps)
+        else:
+            self.norm = nn.Identity()
+        self.head = SConv1d(in_ch, self.channels, kernel_size=last_kernel_size, causal=self.causal, pad_mode=pad_mode, norm=norm, bias=bias)
+
+    def forward_features(self, x, cache=None, sample_indices=None, use_cache=False, debug=False):
+        for i in range(len(self.depths)):
+            # Apply upsampling
+            for layer in self.upsample_layers[i]:
+                if isinstance(layer, (SConv1d, SConvTranspose1d)):
+                    x = layer(x, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug)
+                else:
+                    x = layer(x)
+            
+            # Apply stage (Block1D contains Convlayer which contains SConv1d)
+            for block in self.stages[i]:
+                if hasattr(block, 'mixer') and hasattr(block.mixer, 'conv') and isinstance(block.mixer.conv, SConv1d):
+                    # Block1D forward with cache support
+                    residual = x
+                    x = block.norm(x)
+                    x = block.mixer.conv(x, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug)
+                    if block.gamma is not None:
+                        x = x * block.gamma.unsqueeze(-1)
+                    x = residual + x
+                    
+                    # FFN part
+                    residual = x
+                    x = block.ffn_norm(x)
+                    x = x.permute(0, 2, 1)
+                    x = block.ffn(x)
+                    x = x.permute(0, 2, 1)
+                    if block.ffn_gamma is not None:
+                        x = x * block.ffn_gamma.unsqueeze(-1)
+                    x = residual + x
+                else:
+                    x = block(x)
+
+        return self.norm(x)
+    
+    def forward(self, x, cache=None, sample_indices=None, use_cache=False, debug=False):
+        x = self.forward_features(x, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug)
+        x = self.head(x, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug)
+        return x
+    
+
+@dataclass
+class VibeVoiceTokenizerEncoderOutput:
+    """
+    Output of VibeVoice tokenizer encoder, representing a Gaussian distribution with fixed variance.
+    
+    Args:
+        mean (`torch.FloatTensor`): The mean parameters of the distribution.
+        std (`float` or `torch.FloatTensor`): Fixed standard deviation value.
+    """
+    mean: torch.Tensor
+    std: Optional[Union[float, torch.Tensor]] = None
+    
+    def sample(self, dist_type='fix'):
+        """
+        Sample from the distribution.
+        
+        Args:
+            dist_type (`str`): Sampling method, either 'fix' or 'gaussian'.
+                
+        Returns:
+            `torch.FloatTensor`: Sampled values.
+            `torch.FloatTensor` (optional): Standard deviation used (only when dist_type='gaussian').
+        """
+        if dist_type == 'fix':
+            x = self.mean + self.std * torch.randn_like(self.mean)
+            return x, self.std
+        elif dist_type == 'gaussian':
+            batch_size = self.mean.size(0)
+            value = self.std / 0.8
+            std = torch.randn(batch_size, device=self.mean.device, dtype=self.mean.dtype) * value
+
+            while std.dim() < self.mean.dim():
+                std = std.unsqueeze(-1)
+
+            x = self.mean + std * torch.randn_like(self.mean)
+            return x, std
+        else:
+            return self.mean, self.std
+
+    def kl(self):
+        """Compute KL divergence between this distribution and a standard normal."""
+        target = torch.zeros_like(self.mean)
+        return F.mse_loss(self.mean, target, reduction='none')
+
+    def mode(self):
+        """Return the distribution mode (which is the mean for Gaussian)."""
+        return self.mean
+    
+class VibeVoiceAcousticTokenizerModel(PreTrainedModel):
+    """VibeVoice speech tokenizer model combining encoder and decoder for acoustic tokens"""
+    
+    config_class = VibeVoiceAcousticTokenizerConfig
+    base_model_prefix = "vibevoice_acoustic_tokenizer"
+    _supports_flash_attn_2 = True  
+    _supports_sdpa = True  
+    _no_split_modules = ["TokenizerEncoder", "TokenizerDecoder"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        
+        self.register_buffer('fix_std', torch.tensor(config.fix_std), persistent=False)
+        self.std_dist_type = getattr(config, "std_dist_type", "fix")
+        
+        # Parse encoder depths
+        if isinstance(config.encoder_depths, str):
+            encoder_depths = [int(d) for d in config.encoder_depths.split('-')]
+        else:
+            encoder_depths = config.encoder_depths
+            
+        # Parse decoder depths if provided
+        if config.decoder_depths is not None and isinstance(config.decoder_depths, str):
+            decoder_depths = [int(d) for d in config.decoder_depths.split('-')]
+        else:
+            # Default: use reversed encoder depths if decoder_depths is None
+            decoder_depths = list(reversed(encoder_depths))
+        
+        # Create encoder config
+        encoder_config = copy.deepcopy(config)
+        encoder_config.dimension = config.vae_dim
+        encoder_config.n_filters = config.encoder_n_filters
+        encoder_config.ratios = config.encoder_ratios
+        encoder_config.depths = encoder_depths
+        encoder_config.norm = config.conv_norm
+        encoder_config.pad_mode = config.pad_mode
+        encoder_config.bias = config.conv_bias
+        encoder_config.layernorm_eps = config.layernorm_eps
+        encoder_config.layernorm_elementwise_affine = config.layernorm_elementwise_affine
+        encoder_config.mixer_layer = config.mixer_layer
+        encoder_config.layer_scale_init_value = config.layer_scale_init_value
+        encoder_config.disable_last_norm = config.disable_last_norm
+        
+        # Create decoder config
+        decoder_config = copy.deepcopy(config)
+        decoder_config.dimension = config.vae_dim
+        decoder_config.n_filters = config.decoder_n_filters
+        decoder_config.ratios = config.decoder_ratios
+        decoder_config.depths = decoder_depths
+        decoder_config.norm = config.conv_norm
+        decoder_config.pad_mode = config.pad_mode
+        decoder_config.bias = config.conv_bias
+        decoder_config.layernorm_eps = config.layernorm_eps
+        decoder_config.layernorm_elementwise_affine = config.layernorm_elementwise_affine
+        decoder_config.mixer_layer = config.mixer_layer
+        decoder_config.layer_scale_init_value = config.layer_scale_init_value
+        decoder_config.disable_last_norm = config.disable_last_norm
+        
+        # Initialize encoder and decoder
+        self.encoder = TokenizerEncoder(encoder_config)
+        self.decoder = TokenizerDecoder(decoder_config)
+        
+        # Initialize weights
+        self.apply(self._init_weights)
+    
+    def _init_weights(self, module):
+        """Initialize weights for the model"""
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, std=self.config.weight_init_value)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.LayerNorm):
+            nn.init.ones_(module.weight)
+            nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Conv1d):
+            nn.init.normal_(module.weight, std=self.config.weight_init_value)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+    
+    @torch.no_grad()
+    def encode(self, audio, cache=None, sample_indices=None, use_cache=False, debug=False, is_final_chunk=False):
+        """Convert audio to latent representations"""
+        latents = self.encoder(audio, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug, is_final_chunk=is_final_chunk)
+        return VibeVoiceTokenizerEncoderOutput(mean=latents.permute(0, 2, 1), std=self.fix_std)
+    
+    @torch.no_grad()
+    def sampling(self, encoder_output, dist_type=None):
+        """Sample from the encoder output distribution"""
+        dist_type = dist_type or self.std_dist_type
+    
+        if dist_type == 'fix':
+            return encoder_output.sample(dist_type='fix')
+        elif dist_type == 'gaussian':
+            return encoder_output.sample(dist_type='gaussian')
+        else:
+            raise ValueError(f"Unsupported dist_type: {dist_type}, expected 'fix' or 'gaussian'")
+    
+    @torch.no_grad()
+    def decode(self, latents, cache=None, sample_indices=None, use_cache=False, debug=False):
+        """Convert latent representations back to audio"""
+        if latents.shape[1] == self.config.vae_dim:
+            pass
+        else:
+            latents = latents.permute(0, 2, 1)
+
+        audio = self.decoder(latents, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug)
+        return audio
+
+    def forward(self, audio, cache=None, sample_indices=None, use_cache=False, debug=False):
+        """Full forward pass: encode audio to latents, then decode back to audio"""
+        encoder_output = self.encode(audio, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug)
+        sampled_latents, _ = self.sampling(encoder_output)
+        reconstructed = self.decode(sampled_latents, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug)
+        return reconstructed, sampled_latents
+
+
+class VibeVoiceSemanticTokenizerModel(PreTrainedModel):
+    """VibeVoice speech tokenizer model with only encoder for semantic tokens"""
+    
+    config_class = VibeVoiceSemanticTokenizerConfig
+    base_model_prefix = "vibevoice_semantic_tokenizer"
+    _supports_flash_attn_2 = True  
+    _supports_sdpa = True  
+    _no_split_modules = ["TokenizerEncoder"]
+    
+    def __init__(self, config):
+        super().__init__(config)
+        
+        # Parse encoder depths
+        if isinstance(config.encoder_depths, str):
+            encoder_depths = [int(d) for d in config.encoder_depths.split('-')]
+        else:
+            encoder_depths = config.encoder_depths
+        
+        # Create encoder config
+        encoder_config = copy.deepcopy(config)
+        encoder_config.dimension = config.vae_dim
+        encoder_config.n_filters = config.encoder_n_filters
+        encoder_config.ratios = config.encoder_ratios
+        encoder_config.depths = encoder_depths
+        encoder_config.norm = config.conv_norm
+        encoder_config.pad_mode = config.pad_mode
+        encoder_config.bias = config.conv_bias
+        encoder_config.layernorm_eps = config.layernorm_eps
+        encoder_config.layernorm_elementwise_affine = config.layernorm_elementwise_affine
+        encoder_config.mixer_layer = config.mixer_layer
+        encoder_config.layer_scale_init_value = config.layer_scale_init_value
+        encoder_config.disable_last_norm = config.disable_last_norm
+        
+        # Initialize encoder and decoder
+        self.encoder = TokenizerEncoder(encoder_config)
+        
+        # Initialize weights
+        self.apply(self._init_weights)
+    
+    def _init_weights(self, module):
+        """Initialize weights for the model"""
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, std=self.config.weight_init_value)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.LayerNorm):
+            nn.init.ones_(module.weight)
+            nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Conv1d):
+            nn.init.normal_(module.weight, std=self.config.weight_init_value)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+    
+    @torch.no_grad()
+    def encode(self, audio, cache=None, sample_indices=None, use_cache=False, debug=False, is_final_chunk=False):
+        """Convert audio to latent representations"""
+        latents = self.encoder(audio, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug, is_final_chunk=is_final_chunk)
+        return VibeVoiceTokenizerEncoderOutput(mean=latents.permute(0, 2, 1))
+    
+    @torch.no_grad()
+    def sampling(self, encoder_output, dist_type=None):
+        """Sample from the encoder output distribution"""
+        return encoder_output.sample(dist_type='none')
+
+    def forward(self, audio, cache=None, sample_indices=None, use_cache=False, debug=False):
+        """Full forward pass: encode audio to latents, then decode back to audio"""
+        encoder_output = self.encode(audio, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug)
+        sampled_latents, _ = self.sampling(encoder_output, dist_type='none')
+        return None, sampled_latents
+
+AutoModel.register(VibeVoiceAcousticTokenizerConfig, VibeVoiceAcousticTokenizerModel)
+AutoModel.register(VibeVoiceSemanticTokenizerConfig, VibeVoiceSemanticTokenizerModel)
+
+__all__ = [
+    "VibeVoiceTokenizerStreamingCache",
+    "VibeVoiceAcousticTokenizerModel",
+    "VibeVoiceSemanticTokenizerModel",
+]
\ No newline at end of file
diff --git a/vibevoice/modular/streamer.py b/vibevoice/modular/streamer.py
new file mode 100644
index 0000000..5dd7892
--- /dev/null
+++ b/vibevoice/modular/streamer.py
@@ -0,0 +1,264 @@
+from __future__ import annotations
+
+import torch
+
+import asyncio
+from queue import Empty, Queue
+from typing import TYPE_CHECKING, Optional
+
+
+from transformers.generation import BaseStreamer
+
+
+class AudioStreamer(BaseStreamer):
+    """
+    Audio streamer that stores audio chunks in queues for each sample in the batch.
+    This allows streaming audio generation for multiple samples simultaneously.
+    
+    Parameters:
+        batch_size (`int`):
+            The batch size for generation
+        stop_signal (`any`, *optional*):
+            The signal to put in the queue when generation ends. Defaults to None.
+        timeout (`float`, *optional*):
+            The timeout for the audio queue. If `None`, the queue will block indefinitely.
+    """
+    
+    def __init__(
+        self, 
+        batch_size: int,
+        stop_signal: Optional[any] = None,
+        timeout: Optional[float] = None,
+    ):
+        self.batch_size = batch_size
+        self.stop_signal = stop_signal
+        self.timeout = timeout
+        
+        # Create a queue for each sample in the batch
+        self.audio_queues = [Queue() for _ in range(batch_size)]
+        self.finished_flags = [False for _ in range(batch_size)]
+        self.sample_indices_map = {}  # Maps from sample index to queue index
+        
+    def put(self, audio_chunks: torch.Tensor, sample_indices: torch.Tensor):
+        """
+        Receives audio chunks and puts them in the appropriate queues.
+        
+        Args:
+            audio_chunks: Tensor of shape (num_samples, ...) containing audio chunks
+            sample_indices: Tensor indicating which samples these chunks belong to
+        """
+        for i, sample_idx in enumerate(sample_indices):
+            idx = sample_idx.item()
+            if idx < self.batch_size and not self.finished_flags[idx]:
+                # Convert to numpy or keep as tensor based on preference
+                audio_chunk = audio_chunks[i].detach().cpu()
+                self.audio_queues[idx].put(audio_chunk, timeout=self.timeout)
+    
+    def end(self, sample_indices: Optional[torch.Tensor] = None):
+        """
+        Signals the end of generation for specified samples or all samples.
+        
+        Args:
+            sample_indices: Optional tensor of sample indices to end. If None, ends all.
+        """
+        if sample_indices is None:
+            # End all samples
+            for idx in range(self.batch_size):
+                if not self.finished_flags[idx]:
+                    self.audio_queues[idx].put(self.stop_signal, timeout=self.timeout)
+                    self.finished_flags[idx] = True
+        else:
+            # End specific samples
+            for sample_idx in sample_indices:
+                idx = sample_idx.item() if torch.is_tensor(sample_idx) else sample_idx
+                if idx < self.batch_size and not self.finished_flags[idx]:
+                    self.audio_queues[idx].put(self.stop_signal, timeout=self.timeout)
+                    self.finished_flags[idx] = True
+    
+    def __iter__(self):
+        """Returns an iterator over the batch of audio streams."""
+        return AudioBatchIterator(self)
+    
+    def get_stream(self, sample_idx: int):
+        """Get the audio stream for a specific sample."""
+        if sample_idx >= self.batch_size:
+            raise ValueError(f"Sample index {sample_idx} exceeds batch size {self.batch_size}")
+        return AudioSampleIterator(self, sample_idx)
+
+
+class AudioSampleIterator:
+    """Iterator for a single audio stream from the batch."""
+    
+    def __init__(self, streamer: AudioStreamer, sample_idx: int):
+        self.streamer = streamer
+        self.sample_idx = sample_idx
+        
+    def __iter__(self):
+        return self
+    
+    def __next__(self):
+        value = self.streamer.audio_queues[self.sample_idx].get(timeout=self.streamer.timeout)
+        if value == self.streamer.stop_signal:
+            raise StopIteration()
+        return value
+
+
+class AudioBatchIterator:
+    """Iterator that yields audio chunks for all samples in the batch."""
+    
+    def __init__(self, streamer: AudioStreamer):
+        self.streamer = streamer
+        self.active_samples = set(range(streamer.batch_size))
+        
+    def __iter__(self):
+        return self
+    
+    def __next__(self):
+        if not self.active_samples:
+            raise StopIteration()
+            
+        batch_chunks = {}
+        samples_to_remove = set()
+        
+        # Try to get chunks from all active samples
+        for idx in self.active_samples:
+            try:
+                value = self.streamer.audio_queues[idx].get(block=False)
+                if value == self.streamer.stop_signal:
+                    samples_to_remove.add(idx)
+                else:
+                    batch_chunks[idx] = value
+            except Empty:
+                # Queue is empty for this sample, skip it this iteration
+                pass
+        
+        # Remove finished samples
+        self.active_samples -= samples_to_remove
+        
+        if batch_chunks:
+            return batch_chunks
+        elif self.active_samples:
+            # If no chunks were ready but we still have active samples, 
+            # wait a bit and try again
+            import time
+            time.sleep(0.01)
+            return self.__next__()
+        else:
+            raise StopIteration()
+
+
+class AsyncAudioStreamer(AudioStreamer):
+    """
+    Async version of AudioStreamer for use in async contexts.
+    """
+    
+    def __init__(
+        self, 
+        batch_size: int,
+        stop_signal: Optional[any] = None,
+        timeout: Optional[float] = None,
+    ):
+        super().__init__(batch_size, stop_signal, timeout)
+        # Replace regular queues with async queues
+        self.audio_queues = [asyncio.Queue() for _ in range(batch_size)]
+        self.loop = asyncio.get_running_loop()
+        
+    def put(self, audio_chunks: torch.Tensor, sample_indices: torch.Tensor):
+        """Put audio chunks in the appropriate async queues."""
+        for i, sample_idx in enumerate(sample_indices):
+            idx = sample_idx.item()
+            if idx < self.batch_size and not self.finished_flags[idx]:
+                audio_chunk = audio_chunks[i].detach().cpu()
+                self.loop.call_soon_threadsafe(
+                    self.audio_queues[idx].put_nowait, audio_chunk
+                )
+    
+    def end(self, sample_indices: Optional[torch.Tensor] = None):
+        """Signal the end of generation for specified samples."""
+        if sample_indices is None:
+            indices_to_end = range(self.batch_size)
+        else:
+            indices_to_end = [s.item() if torch.is_tensor(s) else s for s in sample_indices]
+            
+        for idx in indices_to_end:
+            if idx < self.batch_size and not self.finished_flags[idx]:
+                self.loop.call_soon_threadsafe(
+                    self.audio_queues[idx].put_nowait, self.stop_signal
+                )
+                self.finished_flags[idx] = True
+    
+    async def get_stream(self, sample_idx: int):
+        """Get async iterator for a specific sample's audio stream."""
+        if sample_idx >= self.batch_size:
+            raise ValueError(f"Sample index {sample_idx} exceeds batch size {self.batch_size}")
+            
+        while True:
+            value = await self.audio_queues[sample_idx].get()
+            if value == self.stop_signal:
+                break
+            yield value
+    
+    def __aiter__(self):
+        """Returns an async iterator over all audio streams."""
+        return AsyncAudioBatchIterator(self)
+
+
+class AsyncAudioBatchIterator:
+    """Async iterator for batch audio streaming."""
+    
+    def __init__(self, streamer: AsyncAudioStreamer):
+        self.streamer = streamer
+        self.active_samples = set(range(streamer.batch_size))
+        
+    def __aiter__(self):
+        return self
+        
+    async def __anext__(self):
+        if not self.active_samples:
+            raise StopAsyncIteration()
+            
+        batch_chunks = {}
+        samples_to_remove = set()
+        
+        # Create tasks for all active samples
+        tasks = {
+            idx: asyncio.create_task(self._get_chunk(idx)) 
+            for idx in self.active_samples
+        }
+        
+        # Wait for at least one chunk to be ready
+        done, pending = await asyncio.wait(
+            tasks.values(), 
+            return_when=asyncio.FIRST_COMPLETED,
+            timeout=self.streamer.timeout
+        )
+        
+        # Cancel pending tasks
+        for task in pending:
+            task.cancel()
+            
+        # Process completed tasks
+        for idx, task in tasks.items():
+            if task in done:
+                try:
+                    value = await task
+                    if value == self.streamer.stop_signal:
+                        samples_to_remove.add(idx)
+                    else:
+                        batch_chunks[idx] = value
+                except asyncio.CancelledError:
+                    pass
+                    
+        self.active_samples -= samples_to_remove
+        
+        if batch_chunks:
+            return batch_chunks
+        elif self.active_samples:
+            # Try again if we still have active samples
+            return await self.__anext__()
+        else:
+            raise StopAsyncIteration()
+    
+    async def _get_chunk(self, idx):
+        """Helper to get a chunk from a specific queue."""
+        return await self.streamer.audio_queues[idx].get()
\ No newline at end of file
diff --git a/vibevoice/processor/__init__.py b/vibevoice/processor/__init__.py
new file mode 100644
index 0000000..c945e6d
--- /dev/null
+++ b/vibevoice/processor/__init__.py
@@ -0,0 +1,11 @@
+# vibevoice/processor/__init__.py
+from .vibevoice_processor import VibeVoiceProcessor
+from .vibevoice_streaming_processor import VibeVoiceStreamingProcessor
+from .vibevoice_tokenizer_processor import VibeVoiceTokenizerProcessor, AudioNormalizer
+
+__all__ = [
+    "VibeVoiceProcessor",
+    "VibeVoiceStreamingProcessor",
+    "VibeVoiceTokenizerProcessor",
+    "AudioNormalizer",
+]
\ No newline at end of file
diff --git a/vibevoice/processor/audio_utils.py b/vibevoice/processor/audio_utils.py
new file mode 100644
index 0000000..3f9d112
--- /dev/null
+++ b/vibevoice/processor/audio_utils.py
@@ -0,0 +1,217 @@
+import os
+import threading
+
+import numpy as np
+from subprocess import run
+from typing import List, Optional, Union, Dict, Any
+
+COMMON_AUDIO_EXTS = [
+    '.mp3', '.MP3', '.Mp3',  # All case variations of mp3
+    '.m4a', 
+    '.mp4', '.MP4',
+    '.wav', '.WAV',
+    '.m4v',
+    '.aac',
+    '.ogg',
+    '.mov', '.MOV',
+    '.opus',
+    '.m4b',
+    '.flac',
+    '.wma', '.WMA',
+    '.rm', '.3gp', '.mpeg', '.flv', '.webm', '.mp2', '.aif', '.aiff', '.oga', '.ogv', '.mpga', '.m3u8', '.amr'
+]
+
+def load_audio_use_ffmpeg(file: str, resample: bool = False, target_sr: int = 24000):
+    """
+    Open an audio file and read as mono waveform, optionally resampling.
+    Returns both the audio data and the original sample rate.
+
+    Parameters
+    ----------
+    file: str
+        The audio file to open
+    resample: bool
+        Whether to resample the audio
+    target_sr: int
+        The target sample rate if resampling is requested
+
+    Returns
+    -------
+    A tuple containing:
+    - A NumPy array with the audio waveform in float32 dtype
+    - The original sample rate of the audio file
+    """
+    if not resample:
+        # First, get the original sample rate
+        cmd_probe = [
+            "ffprobe",
+            "-v", "quiet",
+            "-show_entries", "stream=sample_rate",
+            "-of", "default=noprint_wrappers=1:nokey=1",
+            file
+        ]
+        
+        original_sr = int(run(cmd_probe, capture_output=True, check=True).stdout.decode().strip())
+    else:
+        original_sr = None
+
+    # Now load the audio
+    sr_to_use = target_sr if resample else original_sr
+    
+    cmd = [
+        "ffmpeg",
+        "-loglevel", "error",
+        "-nostdin",
+        "-threads", "0",
+        "-i", file,
+        "-f", "s16le",
+        "-ac", "1",
+        "-acodec", "pcm_s16le",
+        "-ar", str(sr_to_use),
+        "-",
+    ]
+    
+    out = _run_ffmpeg(cmd).stdout
+    audio_data = np.frombuffer(out, np.int16).flatten().astype(np.float32) / 32768.0
+    
+    return audio_data, sr_to_use
+
+
+def _get_ffmpeg_max_concurrency() -> int:
+    """Get the maximum FFmpeg concurrency from environment variable."""
+    v = os.getenv("VIBEVOICE_FFMPEG_MAX_CONCURRENCY", "")
+    try:
+        n = int(v) if v.strip() else 0
+    except Exception:
+        n = 0
+    # 0/negative means no explicit limit.
+    return n
+
+
+_FFMPEG_MAX_CONCURRENCY = _get_ffmpeg_max_concurrency()
+_FFMPEG_SEM = threading.Semaphore(_FFMPEG_MAX_CONCURRENCY) if _FFMPEG_MAX_CONCURRENCY > 0 else None
+
+
+def _run_ffmpeg(cmd: list, *, stdin_bytes: bytes = None):
+    """Run ffmpeg with optional global concurrency limiting.
+
+    This is important for vLLM multi-request concurrency: spawning too many
+    ffmpeg processes can saturate CPU/IO and cause request failures/timeouts.
+    """
+    if _FFMPEG_SEM is None:
+        return run(cmd, capture_output=True, check=True, input=stdin_bytes)
+    with _FFMPEG_SEM:
+        return run(cmd, capture_output=True, check=True, input=stdin_bytes)
+
+
+def load_audio_bytes_use_ffmpeg(data: bytes, *, resample: bool = False, target_sr: int = 24000):
+    """Decode audio bytes via ffmpeg stdin pipe.
+
+    Compared to writing bytes to a temp file, this avoids filesystem IO and
+    reduces contention under high request concurrency.
+    
+    Parameters
+    ----------
+    data: bytes
+        The audio data bytes
+    resample: bool
+        Whether to resample the audio (must be True)
+    target_sr: int
+        The target sample rate if resampling is requested
+
+    Returns
+    -------
+    A tuple containing:
+    - A NumPy array with the audio waveform in float32 dtype
+    - The sample rate
+    """
+    if not resample:
+        # For stdin bytes, we don't have a cheap/robust way to probe original sr.
+        # Keep behavior explicit.
+        raise ValueError("load_audio_bytes_use_ffmpeg requires resample=True")
+
+    cmd = [
+        "ffmpeg",
+        "-loglevel", "error",
+        "-threads", "0",
+        "-i", "pipe:0",
+        "-f", "s16le",
+        "-ac", "1",
+        "-acodec", "pcm_s16le",
+        "-ar", str(target_sr),
+        "-",
+    ]
+    out = _run_ffmpeg(cmd, stdin_bytes=data).stdout
+    audio_data = np.frombuffer(out, np.int16).flatten().astype(np.float32) / 32768.0
+    return audio_data, target_sr
+
+
+class AudioNormalizer:
+    """
+    Audio normalization class for VibeVoice tokenizer.
+    
+    This class provides audio normalization to ensure consistent input levels
+    for the VibeVoice tokenizer while maintaining audio quality.
+    """
+    
+    def __init__(self, target_dB_FS: float = -25, eps: float = 1e-6):
+        """
+        Initialize the audio normalizer.
+        
+        Args:
+            target_dB_FS (float): Target dB FS level for the audio. Default: -25
+            eps (float): Small value to avoid division by zero. Default: 1e-6
+        """
+        self.target_dB_FS = target_dB_FS
+        self.eps = eps
+    
+    def tailor_dB_FS(self, audio: np.ndarray) -> tuple:
+        """
+        Adjust the audio to the target dB FS level.
+        
+        Args:
+            audio (np.ndarray): Input audio signal
+            
+        Returns:
+            tuple: (normalized_audio, rms, scalar)
+        """
+        rms = np.sqrt(np.mean(audio**2))
+        scalar = 10 ** (self.target_dB_FS / 20) / (rms + self.eps)
+        normalized_audio = audio * scalar
+        return normalized_audio, rms, scalar
+    
+    def avoid_clipping(self, audio: np.ndarray, scalar: Optional[float] = None) -> tuple:
+        """
+        Avoid clipping by scaling down if necessary.
+        
+        Args:
+            audio (np.ndarray): Input audio signal
+            scalar (float, optional): Explicit scaling factor
+            
+        Returns:
+            tuple: (normalized_audio, scalar)
+        """
+        if scalar is None:
+            max_val = np.max(np.abs(audio))
+            if max_val > 1.0:
+                scalar = max_val + self.eps
+            else:
+                scalar = 1.0
+        
+        return audio / scalar, scalar
+    
+    def __call__(self, audio: np.ndarray) -> np.ndarray:
+        """
+        Normalize the audio by adjusting to target dB FS and avoiding clipping.
+        
+        Args:
+            audio (np.ndarray): Input audio signal
+            
+        Returns:
+            np.ndarray: Normalized audio signal
+        """
+        # First adjust to target dB FS
+        audio, _, _ = self.tailor_dB_FS(audio)
+        # Then avoid clipping
+        audio, _ = self.avoid_clipping(audio)
+        return audio
\ No newline at end of file
diff --git a/vibevoice/processor/vibevoice_asr_processor.py b/vibevoice/processor/vibevoice_asr_processor.py
new file mode 100644
index 0000000..cacb116
--- /dev/null
+++ b/vibevoice/processor/vibevoice_asr_processor.py
@@ -0,0 +1,572 @@
+"""
+Processor class for VibeVoice ASR models.
+"""
+
+import os
+import json
+import math
+import warnings
+from typing import List, Optional, Union, Dict, Any, Tuple
+
+import numpy as np
+import torch
+
+from transformers.tokenization_utils_base import BatchEncoding
+from transformers.utils import TensorType, logging
+from .vibevoice_tokenizer_processor import VibeVoiceTokenizerProcessor, AudioNormalizer
+
+try:
+    from .audio_utils import load_audio_use_ffmpeg
+    HAS_FFMPEG_UTILS = True
+except ImportError:
+    HAS_FFMPEG_UTILS = False
+    warnings.warn("audio_utils not available, will fall back to soundfile for audio loading")
+
+logger = logging.get_logger(__name__)
+
+SYSTEM_PROMPT = "You are a helpful assistant that transcribes audio input into text output in JSON format."
+
+
+class VibeVoiceASRProcessor: 
+    """
+    Processor for VibeVoice ASR (Automatic Speech Recognition) models.
+    
+    This processor handles audio preprocessing and tokenization for ASR tasks,
+    following the exact format used in training with proper chat templates.
+    
+    Args:
+        tokenizer: The text tokenizer for processing text
+        audio_processor: The audio processor for processing speech
+        speech_tok_compress_ratio (int): Compression ratio for speech tokenization. Default: 3200 (product of encoder ratios [8,5,5,4,2,2])
+        target_sample_rate (int): Target sample rate for audio
+        normalize_audio (bool): Whether to normalize audio input
+    """
+    
+    def __init__(
+        self,
+        tokenizer=None,
+        audio_processor=None,
+        speech_tok_compress_ratio=3200,
+        target_sample_rate=24000,
+        normalize_audio=True,
+        **kwargs
+    ):
+        self.tokenizer = tokenizer
+        self.audio_processor = audio_processor or VibeVoiceTokenizerProcessor(
+            sampling_rate=target_sample_rate,
+            normalize_audio=normalize_audio
+        )
+        self.speech_tok_compress_ratio = speech_tok_compress_ratio
+        self.target_sample_rate = target_sample_rate
+        self.normalize_audio = normalize_audio
+        
+        if normalize_audio:
+            self.audio_normalizer = AudioNormalizer()
+        else:
+            self.audio_normalizer = None
+        
+        # Cache special token IDs
+        self._cache_special_tokens()
+    
+    def _cache_special_tokens(self):
+        """Cache special token IDs for efficiency."""
+        # Add safety checks for special tokens
+        if hasattr(self.tokenizer, 'speech_start_id'):
+            self.speech_start_id = self.tokenizer.speech_start_id
+        else:
+            self.speech_start_id = self.tokenizer.convert_tokens_to_ids("<|speech_start|>")
+            
+        if hasattr(self.tokenizer, 'speech_end_id'):
+            self.speech_end_id = self.tokenizer.speech_end_id
+        else:
+            self.speech_end_id = self.tokenizer.convert_tokens_to_ids("<|speech_end|>")
+            
+        if hasattr(self.tokenizer, 'speech_pad_id'):
+            self.speech_pad_id = self.tokenizer.speech_pad_id
+        else:
+            self.speech_pad_id = self.tokenizer.convert_tokens_to_ids("<|speech_pad|>")
+            
+        if hasattr(self.tokenizer, 'pad_id'):
+            self.pad_id = self.tokenizer.pad_id
+        elif hasattr(self.tokenizer, 'pad_token_id'):
+            self.pad_id = self.tokenizer.pad_token_id
+        else:
+            self.pad_id = self.tokenizer.convert_tokens_to_ids("<|endoftext|>")
+        
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        """
+        Load processor from a pretrained model path.
+        
+        Args:
+            pretrained_model_name_or_path: Path to the pretrained model
+            **kwargs: Additional keyword arguments
+            
+        Returns:
+            VibeVoiceASRProcessor: The loaded processor
+        """
+        import json
+        from transformers.utils import cached_file
+        from vibevoice.modular.modular_vibevoice_text_tokenizer import VibeVoiceASRTextTokenizerFast
+        
+        # Try to load configuration
+        config_path = os.path.join(pretrained_model_name_or_path, "preprocessor_config.json")
+        config = {}
+        
+        if os.path.exists(config_path):
+            with open(config_path, 'r') as f:
+                config = json.load(f)
+        else:
+            try:
+                config_file = cached_file(
+                    pretrained_model_name_or_path,
+                    "preprocessor_config.json",
+                    **kwargs
+                )
+                with open(config_file, 'r') as f:
+                    config = json.load(f)
+            except Exception as e:
+                logger.warning(f"Could not load preprocessor_config.json: {e}")
+                logger.warning("Using default configuration")
+        
+        # Extract parameters
+        speech_tok_compress_ratio = config.get("speech_tok_compress_ratio", 3200)
+        target_sample_rate = config.get("target_sample_rate", 24000)
+        normalize_audio = config.get("normalize_audio", True)
+        
+        # Load tokenizer
+        language_model_pretrained_name = config.get("language_model_pretrained_name", None) or kwargs.pop("language_model_pretrained_name", "Qwen/Qwen2.5-1.5B")
+        logger.info(f"Loading tokenizer from {language_model_pretrained_name}")
+        
+        if 'qwen' in language_model_pretrained_name.lower():
+            tokenizer = VibeVoiceASRTextTokenizerFast.from_pretrained(
+                language_model_pretrained_name,
+                **kwargs
+            )
+        else:
+            raise ValueError(f"Unsupported tokenizer type for {language_model_pretrained_name}")
+        
+        # Load audio processor
+        audio_processor = VibeVoiceTokenizerProcessor(
+            sampling_rate=target_sample_rate,
+            normalize_audio=normalize_audio,
+            target_dB_FS=config.get("target_dB_FS", -25),
+            eps=config.get("eps", 1e-6),
+        )
+        
+        return cls(
+            tokenizer=tokenizer,
+            audio_processor=audio_processor,
+            speech_tok_compress_ratio=speech_tok_compress_ratio,
+            target_sample_rate=target_sample_rate,
+            normalize_audio=normalize_audio,
+        )
+    
+    def save_pretrained(self, save_directory: Union[str, os.PathLike], **kwargs):
+        """
+        Save processor configuration to a directory.
+        
+        Args:
+            save_directory: Directory to save the configuration
+            **kwargs: Additional keyword arguments
+        """
+        import json
+        
+        os.makedirs(save_directory, exist_ok=True)
+        
+        # Save processor configuration
+        processor_config = {
+            "processor_class": "VibeVoiceASRProcessor",
+            "speech_tok_compress_ratio": self.speech_tok_compress_ratio,
+            "target_sample_rate": self.target_sample_rate,
+            "normalize_audio": self.normalize_audio,
+            "target_dB_FS": -25,
+            "eps": 1e-6,
+        }
+        
+        config_path = os.path.join(save_directory, "preprocessor_config.json")
+        with open(config_path, 'w') as f:
+            json.dump(processor_config, f, indent=2)
+        
+        logger.info(f"Processor configuration saved in {config_path}")
+    
+    def __call__(
+        self,
+        audio: Optional[Union[str, np.ndarray, torch.Tensor, List[Union[str, np.ndarray, torch.Tensor]]]] = None,
+        sampling_rate: Optional[int] = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        padding: bool = True,
+        max_length: Optional[int] = None,
+        truncation: bool = False,
+        add_generation_prompt: bool = True,
+        use_streaming: bool = True,
+        context_info: Optional[str] = None,
+        **kwargs
+    ) -> BatchEncoding:
+        """
+        Process audio input for ASR model.
+        
+        Args:
+            audio: Audio input(s). Can be:
+                - str: Path to audio file
+                - np.ndarray: Audio array
+                - torch.Tensor: Audio tensor
+                - List of the above for batch processing
+            sampling_rate: Sampling rate of input audio
+            return_tensors: Output format ('pt' for PyTorch, 'np' for NumPy)
+            padding: Whether to pad batch inputs
+            max_length: Maximum sequence length
+            truncation: Whether to truncate long sequences
+            add_generation_prompt: Whether to add generation prompt for inference
+            use_streaming: Whether to use streaming mode (True by default, auto False if <60s)
+            context_info: Optional context information (e.g., hotwords, metadata) to help transcription
+            
+        Returns:
+            BatchEncoding with:
+                - input_ids: Token IDs for the model
+                - attention_mask: Attention mask
+                - acoustic_input_mask: Mask indicating speech token positions
+                - speech_tensors: Processed speech features
+                - speech_masks: Valid speech masks
+                - vae_tok_seqlens: Length of each speech segment in tokens
+        """
+        if audio is None:
+            raise ValueError("Audio input is required for ASR processing")
+        
+        # Handle single vs batch input
+        if isinstance(audio, list):
+            is_batched = True
+            audio_list = audio
+        else:
+            is_batched = False
+            audio_list = [audio]
+        
+        # Process each audio input
+        all_encodings = []
+        for audio_input in audio_list:
+            encoding = self._process_single_audio(
+                audio_input,
+                sampling_rate=sampling_rate,
+                add_generation_prompt=add_generation_prompt,
+                use_streaming=use_streaming,
+                context_info=context_info,
+            )
+            all_encodings.append(encoding)
+        
+        # Combine into batch
+        batch_encoding = self._batch_encode(
+            all_encodings,
+            padding=padding,
+            max_length=max_length,
+            truncation=truncation,
+            return_tensors=return_tensors,
+        )
+        
+        return batch_encoding
+    
+    def _process_single_audio(
+        self,
+        audio: Union[str, np.ndarray, torch.Tensor],
+        sampling_rate: Optional[int] = None,
+        add_generation_prompt: bool = True,
+        use_streaming: bool = True,
+        context_info: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """
+        Process a single audio input.
+        
+        Args:
+            audio: Single audio input
+            sampling_rate: Audio sampling rate
+            add_generation_prompt: Whether to add generation prompt
+            context_info: Optional context information (e.g., hotwords, metadata) to help transcription
+            
+        Returns:
+            Dictionary with processed tokens and audio features
+        """
+        # Process audio through audio processor
+        if isinstance(audio, str):
+            # Load from file using ffmpeg for better format support
+            if HAS_FFMPEG_UTILS:
+                try:
+                    audio_array, file_sr = load_audio_use_ffmpeg(audio, resample=False)
+                except Exception as e:
+                    # Fall back to soundfile if ffmpeg fails
+                    warnings.warn(f"ffmpeg loading failed, falling back to soundfile: {e}")
+                    import soundfile as sf
+                    audio_array, file_sr = sf.read(audio)
+                    if audio_array.ndim > 1:
+                        audio_array = audio_array.mean(axis=1)  # Convert to mono
+            else:
+                import soundfile as sf
+                audio_array, file_sr = sf.read(audio)
+                if audio_array.ndim > 1:
+                    audio_array = audio_array.mean(axis=1)  # Convert to mono
+            
+            # Resample if needed
+            if file_sr != self.target_sample_rate:
+                import librosa
+                audio_array = librosa.resample(
+                    audio_array, 
+                    orig_sr=file_sr, 
+                    target_sr=self.target_sample_rate
+                )
+        elif isinstance(audio, torch.Tensor):
+            audio_array = audio.cpu().numpy()
+            if audio_array.ndim > 1:
+                audio_array = audio_array.squeeze()
+        else:
+            audio_array = np.array(audio, dtype=np.float32)
+            if audio_array.ndim > 1:
+                audio_array = audio_array.squeeze()
+        
+        # Ensure float32
+        audio_array = audio_array.astype(np.float32)
+        
+        # Normalize if needed
+        if self.normalize_audio and self.audio_normalizer:
+            audio_array = self.audio_normalizer(audio_array)
+        
+        # Calculate audio duration
+        audio_duration = len(audio_array) / self.target_sample_rate
+        
+        # Auto-disable streaming for short audio (<60s)
+        if use_streaming and audio_duration < 60.0:
+            use_streaming = False
+        
+        # Calculate token length based on streaming mode
+        # Non-streaming: uses ceil (encoder adds extra_padding for stride alignment)
+        # Streaming: uses floor (segments processed independently, no global alignment)
+        # if use_streaming:
+        #     vae_tok_len = len(audio_array) // self.speech_tok_compress_ratio
+        # else:
+        vae_tok_len = math.ceil(len(audio_array) / self.speech_tok_compress_ratio)
+        
+        # Build token sequence following training format
+        # 1. System prompt - use apply_chat_template then encode like in training
+        system_prompt_text = self.tokenizer.apply_chat_template(
+            [{"role": "system", "content": SYSTEM_PROMPT}],
+            tokenize=False
+        )
+        system_tokens = self.tokenizer.encode(system_prompt_text)
+        
+        # 2. User input with speech tokens
+        # Build speech placeholder string
+        sp_start_token = self.tokenizer.convert_ids_to_tokens(self.speech_start_id)
+        sp_pad_token = self.tokenizer.convert_ids_to_tokens(self.speech_pad_id)
+        sp_end_token = self.tokenizer.convert_ids_to_tokens(self.speech_end_id)
+        
+        # User suffix with audio duration info
+        show_keys = ['Start time', 'End time', 'Speaker ID', 'Content']
+        if context_info and context_info.strip():
+            user_suffix = f"This is a {audio_duration:.2f} seconds audio, with extra info: {context_info.strip()}\n\nPlease transcribe it with these keys: " + ", ".join(show_keys)
+        else:
+            user_suffix = f"This is a {audio_duration:.2f} seconds audio, please transcribe it with these keys: " + ", ".join(show_keys)
+        
+        user_input_string = ''.join(
+            [sp_start_token] + [sp_pad_token] * vae_tok_len + [sp_end_token]
+        ) + '\n' + user_suffix
+        
+        user_tokens = self.tokenizer.apply_chat_template(
+            [{"role": "user", "content": user_input_string}],
+            tokenize=True
+        )
+        
+        # Combine tokens
+        full_tokens = system_tokens + user_tokens
+        
+        # Create acoustic input mask
+        acoustic_input_mask = [1 if token == self.speech_pad_id else 0 for token in full_tokens]
+        
+        return {
+            "input_ids": full_tokens,
+            "acoustic_input_mask": acoustic_input_mask,
+            "speech": audio_array,
+            "vae_tok_len": vae_tok_len,
+        }
+    
+    def _batch_encode(
+        self,
+        encodings: List[Dict[str, Any]],
+        padding: bool = True,
+        max_length: Optional[int] = None,
+        truncation: bool = False,
+        return_tensors: Optional[str] = None,
+    ) -> BatchEncoding:
+        """
+        Combine multiple encodings into a batch.
+        
+        Args:
+            encodings: List of encoded samples
+            padding: Whether to pad sequences
+            max_length: Maximum sequence length
+            truncation: Whether to truncate
+            return_tensors: Output format
+            
+        Returns:
+            BatchEncoding with batched data
+        """
+        # Extract components
+        input_ids_list = [enc["input_ids"] for enc in encodings]
+        acoustic_masks_list = [enc["acoustic_input_mask"] for enc in encodings]
+        speech_list = [enc["speech"] for enc in encodings]
+        vae_tok_lens = [enc["vae_tok_len"] for enc in encodings]
+        
+        # Determine max length for padding
+        if padding:
+            if max_length is not None:
+                target_length = max_length
+            else:
+                target_length = max(len(ids) for ids in input_ids_list)
+            
+            # Pad sequences
+            padded_input_ids = []
+            padded_acoustic_masks = []
+            attention_masks = []
+            
+            for input_ids, acoustic_mask in zip(input_ids_list, acoustic_masks_list):
+                # Truncate if needed
+                if truncation and len(input_ids) > target_length:
+                    input_ids = input_ids[:target_length]
+                    acoustic_mask = acoustic_mask[:target_length]
+                
+                # Pad sequences to left (for autoregressive generation)
+                padding_length = target_length - len(input_ids)
+                padded_ids = [self.pad_id] * padding_length + input_ids
+                padded_acoustic = [0] * padding_length + acoustic_mask
+                attention_mask = [0] * padding_length + [1] * len(input_ids)
+                
+                padded_input_ids.append(padded_ids)
+                padded_acoustic_masks.append(padded_acoustic)
+                attention_masks.append(attention_mask)
+            
+            input_ids_list = padded_input_ids
+            acoustic_masks_list = padded_acoustic_masks
+        else:
+            attention_masks = [[1] * len(ids) for ids in input_ids_list]
+        
+        # Process speech tensors - raw audio is 1D, so we keep it as is
+        max_speech_length = max(len(s) for s in speech_list)
+        padded_speeches = np.zeros((len(speech_list), max_speech_length), dtype=np.float32)
+        speech_masks = np.zeros((len(speech_list), max(vae_tok_lens)), dtype=bool)
+        
+        for i, (speech, vae_len) in enumerate(zip(speech_list, vae_tok_lens)):
+            padded_speeches[i, :len(speech)] = speech
+            speech_masks[i, :vae_len] = True
+        
+        # Create batch encoding
+        batch_encoding = BatchEncoding()
+        
+        if return_tensors == "pt":
+            batch_encoding["input_ids"] = torch.tensor(input_ids_list, dtype=torch.long)
+            batch_encoding["attention_mask"] = torch.tensor(attention_masks, dtype=torch.long)
+            batch_encoding["acoustic_input_mask"] = torch.tensor(acoustic_masks_list, dtype=torch.bool)
+            batch_encoding["speech_tensors"] = torch.tensor(padded_speeches, dtype=torch.float32)
+            batch_encoding["speech_masks"] = torch.tensor(speech_masks, dtype=torch.bool)
+            # Note: vae_tok_seqlens and speech_type are not included as they are not model inputs
+        else:
+            batch_encoding["input_ids"] = input_ids_list if len(input_ids_list) > 1 else input_ids_list[0]
+            batch_encoding["attention_mask"] = attention_masks if len(attention_masks) > 1 else attention_masks[0]
+            batch_encoding["acoustic_input_mask"] = acoustic_masks_list if len(acoustic_masks_list) > 1 else acoustic_masks_list[0]
+            batch_encoding["speech_tensors"] = padded_speeches if len(padded_speeches) > 1 else padded_speeches[0]
+            batch_encoding["speech_masks"] = speech_masks if len(speech_masks) > 1 else speech_masks[0]
+        
+        return batch_encoding
+    
+    def batch_decode(self, *args, **kwargs):
+        """
+        Decode batch of token IDs to text.
+        Forwards to tokenizer's batch_decode method.
+        """
+        return self.tokenizer.batch_decode(*args, **kwargs)
+    
+    def decode(self, *args, **kwargs):
+        """
+        Decode token IDs to text.
+        Forwards to tokenizer's decode method.
+        """
+        return self.tokenizer.decode(*args, **kwargs)
+    
+    def post_process_transcription(self, text: str) -> List[Dict[str, Any]]:
+        """
+        Post-process the generated transcription text to extract structured data.
+        
+        Args:
+            text: Generated text from the model
+            
+        Returns:
+            List of dictionaries with transcription segments
+        """
+        try:
+            # Try to parse as JSON
+            if "```json" in text:
+                # Extract JSON from markdown code block
+                json_start = text.find("```json") + 7
+                json_end = text.find("```", json_start)
+                json_str = text[json_start:json_end].strip()
+            else:
+                # Try to find JSON array or object
+                json_start = text.find("[")
+                if json_start == -1:
+                    json_start = text.find("{")
+                if json_start != -1:
+                    # Find matching closing bracket
+                    bracket_count = 0
+                    json_end = json_start
+                    for i in range(json_start, len(text)):
+                        if text[i] in "[{":
+                            bracket_count += 1
+                        elif text[i] in "]}":
+                            bracket_count -= 1
+                            if bracket_count == 0:
+                                json_end = i + 1
+                                break
+                    json_str = text[json_start:json_end]
+                else:
+                    json_str = text
+            
+            # Parse JSON
+            result = json.loads(json_str)
+            
+            # Ensure it's a list
+            if isinstance(result, dict):
+                result = [result]
+            
+            # Validate and clean up the result
+            cleaned_result = []
+            for item in result:
+                if isinstance(item, dict):
+                    cleaned_item = {}
+                    # Map keys to expected format
+                    key_mapping = {
+                        "Start time": "start_time",
+                        "Start": "start_time",
+                        "End time": "end_time",
+                        "End": "end_time",
+                        "Speaker ID": "speaker_id",
+                        "Speaker": "speaker_id",
+                        "Content": "text",
+                    }
+                    for key, mapped_key in key_mapping.items():
+                        if key in item:
+                            cleaned_item[mapped_key] = item[key]
+                    
+                    if cleaned_item:
+                        cleaned_result.append(cleaned_item)
+            
+            return cleaned_result
+            
+        except json.JSONDecodeError as e:
+            logger.warning(f"Failed to parse JSON from transcription: {e}")
+            logger.debug(f"Raw text: {text}")
+            return []
+        except Exception as e:
+            logger.warning(f"Error post-processing transcription: {e}")
+            return []
+    
+    @property
+    def model_input_names(self):
+        """Return the list of inputs accepted by the model."""
+        return ["input_ids", "attention_mask", "acoustic_input_mask", "speech_tensors", "speech_masks"]
+
+__all__ = ["VibeVoiceASRProcessor"]
diff --git a/vibevoice/processor/vibevoice_processor.py b/vibevoice/processor/vibevoice_processor.py
new file mode 100644
index 0000000..178af53
--- /dev/null
+++ b/vibevoice/processor/vibevoice_processor.py
@@ -0,0 +1,692 @@
+import math
+import warnings
+from typing import List, Optional, Union, Dict, Any, Tuple
+import os
+import re
+
+import numpy as np
+import torch
+
+from transformers.tokenization_utils_base import BatchEncoding, PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
+from transformers.utils import TensorType, logging
+from .vibevoice_tokenizer_processor import AudioNormalizer
+
+logger = logging.get_logger(__name__)
+
+
+class VibeVoiceProcessor:
+    r"""
+    Constructs a VibeVoice processor which wraps a VibeVoice tokenizer and audio processor into a single processor.
+
+    [`VibeVoiceProcessor`] offers all the functionalities of [`VibeVoiceTokenizer`] and [`VibeVoiceTokenizerProcessor`]. 
+    See the [`~VibeVoiceProcessor.__call__`] and [`~VibeVoiceProcessor.decode`] for more information.
+
+    Args:
+        tokenizer (`VibeVoiceTextTokenizer` or `VibeVoiceTextTokenizerFast`):
+            The tokenizer for text processing.
+        audio_processor (`VibeVoiceTokenizerProcessor`):
+            The audio processor for speech processing.
+        speech_tok_compress_ratio (`int`, *optional*, defaults to 3200):
+            The compression ratio for speech tokenization.
+        db_normalize (`bool`, *optional*, defaults to True):
+            Whether to apply decibel normalization to audio inputs.
+    """
+
+    def __init__(self, tokenizer=None, audio_processor=None, speech_tok_compress_ratio=3200, db_normalize=True, **kwargs):
+        self.tokenizer = tokenizer
+        self.audio_processor = audio_processor
+        self.speech_tok_compress_ratio = speech_tok_compress_ratio
+        self.db_normalize = db_normalize
+        self.audio_normalizer = AudioNormalizer() if db_normalize else None
+        self.system_prompt = " Transform the text provided by various speakers into speech output, utilizing the distinct voice of each respective speaker.\n"
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        """
+        Instantiate a VibeVoiceProcessor from a pretrained VibeVoice processor.
+
+        Args:
+            pretrained_model_name_or_path (`str` or `os.PathLike`):
+                This can be either:
+                - a string, the *model id* of a pretrained model
+                - a path to a *directory* containing processor config
+
+        Returns:
+            [`VibeVoiceProcessor`]: The processor object instantiated from pretrained model.
+        """
+        import os
+        import json
+        from transformers.utils import cached_file
+        from .vibevoice_tokenizer_processor import VibeVoiceTokenizerProcessor
+        from vibevoice.modular.modular_vibevoice_text_tokenizer import (
+            VibeVoiceTextTokenizer, 
+            VibeVoiceTextTokenizerFast
+        )
+        
+        # Try to load from local path first, then from HF hub
+        config_path = os.path.join(pretrained_model_name_or_path, "preprocessor_config.json")
+        config = None
+        
+        if os.path.exists(config_path):
+            # Local path exists
+            with open(config_path, 'r') as f:
+                config = json.load(f)
+        else:
+            # Try to load from HF hub
+            try:
+                config_file = cached_file(
+                    pretrained_model_name_or_path,
+                    "preprocessor_config.json",
+                    **kwargs
+                )
+                with open(config_file, 'r') as f:
+                    config = json.load(f)
+            except Exception as e:
+                logger.warning(f"Could not load preprocessor_config.json from {pretrained_model_name_or_path}: {e}")
+                logger.warning("Using default configuration")
+                config = {
+                    "speech_tok_compress_ratio": 3200,
+                    "db_normalize": True,
+                }
+        
+        # Extract main processor parameters
+        speech_tok_compress_ratio = config.get("speech_tok_compress_ratio", 3200)
+        db_normalize = config.get("db_normalize", True)
+        
+        # Load tokenizer - try from model path first, then fall back to Qwen        
+        language_model_pretrained_name = config.get("language_model_pretrained_name", None) or kwargs.pop("language_model_pretrained_name", "Qwen/Qwen2.5-1.5B")
+        logger.info(f"Loading tokenizer from {language_model_pretrained_name}")
+        if 'qwen' in language_model_pretrained_name.lower():
+            tokenizer = VibeVoiceTextTokenizerFast.from_pretrained(
+                language_model_pretrained_name,
+                **kwargs
+            )
+        else:
+            raise ValueError(f"Unsupported tokenizer type for {language_model_pretrained_name}. Supported types: Qwen, Llama, Gemma.")
+        
+        # Load audio processor
+        if "audio_processor" in config:
+            # Create audio processor from config
+            audio_config = config["audio_processor"]
+            audio_processor = VibeVoiceTokenizerProcessor(
+                sampling_rate=audio_config.get("sampling_rate", 24000),
+                normalize_audio=audio_config.get("normalize_audio", True),
+                target_dB_FS=audio_config.get("target_dB_FS", -25),
+                eps=audio_config.get("eps", 1e-6),
+            )
+        else:
+            # Create default audio processor
+            audio_processor = VibeVoiceTokenizerProcessor()
+        
+        # Create and return the processor
+        return cls(
+            tokenizer=tokenizer,
+            audio_processor=audio_processor,
+            speech_tok_compress_ratio=speech_tok_compress_ratio,
+            db_normalize=db_normalize,
+        )
+    
+    def save_pretrained(self, save_directory: Union[str, os.PathLike], **kwargs):
+        """
+        Save a processor to a directory, so that it can be re-loaded using the
+        [`~VibeVoiceProcessor.from_pretrained`] class method.
+
+        Args:
+            save_directory (`str` or `os.PathLike`):
+                Directory where the processor will be saved.
+        """
+        import os
+        import json
+        
+        os.makedirs(save_directory, exist_ok=True)
+        
+        # Save processor configuration
+        processor_config = {
+            "processor_class": "VibeVoiceProcessor",
+            "speech_tok_compress_ratio": self.speech_tok_compress_ratio,
+            "db_normalize": self.db_normalize,
+            "audio_processor": {
+                "feature_extractor_type": "VibeVoiceTokenizerProcessor",
+                "sampling_rate": getattr(self.audio_processor, 'sampling_rate', 24000),
+                "normalize_audio": getattr(self.audio_processor, 'normalize_audio', True),
+                "target_dB_FS": getattr(self.audio_processor, 'target_dB_FS', -25),
+                "eps": getattr(self.audio_processor, 'eps', 1e-6),
+            }
+        }
+        
+        config_path = os.path.join(save_directory, "preprocessor_config.json")
+        with open(config_path, 'w') as f:
+            json.dump(processor_config, f, indent=2)
+        
+        logger.info(f"Processor configuration saved in {config_path}")
+    
+    def __call__(
+        self,
+        text: Optional[Union[str, List[str], TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]]] = None,
+        voice_samples: Optional[Union[List[Union[str, np.ndarray]], List[List[Union[str, np.ndarray]]]]] = None,
+        padding: Union[bool, str, PaddingStrategy] = True,
+        truncation: Union[bool, str, TruncationStrategy] = False,
+        max_length: Optional[int] = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        return_attention_mask: bool = True,
+        **kwargs,
+    ) -> BatchEncoding:
+        """
+        Main method to process one or more podcast scripts with optional voice samples.
+
+        Args:
+            text (`str`, `List[str]`):
+                The input text(s) to process. Can be:
+                - A single script string
+                - A list of script strings for batch processing
+                - A path to a .json or .txt file
+                - A list of paths
+            voice_samples (`List[Union[str, np.ndarray]]`, `List[List[Union[str, np.ndarray]]]`, *optional*):
+                Voice samples for each script. Can be:
+                - A list of samples for a single script
+                - A list of lists for batch processing
+            padding (`bool`, `str` or `PaddingStrategy`, defaults to `True`):
+                Whether to pad sequences to the same length
+            truncation (`bool`, `str` or `TruncationStrategy`, defaults to `False`):
+                Whether to truncate sequences
+            max_length (`int`, *optional*):
+                Maximum length of the returned sequences
+            return_tensors (`str` or `TensorType`, *optional*):
+                If set, will return tensors of a particular framework
+            return_attention_mask (`bool`, defaults to `True`):
+                Whether to return the attention mask
+
+        Returns:
+            `BatchEncoding`: A BatchEncoding with the following fields:
+                - **input_ids** -- List of token id sequences or tensor
+                - **attention_mask** -- List of attention masks or tensor
+                - **speech_tensors** -- Padded speech inputs (if voice_samples provided)
+                - **speech_masks** -- Speech masks (if voice_samples provided)
+                - **speech_input_mask** -- Boolean masks indicating speech token positions
+        """
+        # Handle single vs batch input
+        if isinstance(text, str) or (isinstance(text, list) and len(text) > 0 and not isinstance(text[0], str)):
+            # Single input
+            texts = [text]
+            is_batched = False
+        else:
+            # Batch input
+            texts = text
+            is_batched = True
+            
+        # Handle voice samples
+        if voice_samples is not None:
+            if not is_batched or (isinstance(voice_samples[0], (str, np.ndarray))):
+                # Single set of voice samples
+                voice_samples_list = [voice_samples]
+            else:
+                # Batch of voice samples
+                voice_samples_list = voice_samples
+        else:
+            voice_samples_list = [None] * len(texts)
+        
+        # Process each input
+        all_encodings = []
+        for text_input, voice_input in zip(texts, voice_samples_list):
+            encoding = self._process_single(text_input, voice_input)
+            all_encodings.append(encoding)
+            
+        # Combine batch
+        batch_encoding = self._batch_encode(
+            all_encodings,
+            padding=padding,
+            truncation=truncation,
+            max_length=max_length,
+            return_tensors=return_tensors,
+            return_attention_mask=return_attention_mask,
+        )
+        
+        return batch_encoding
+    
+    def _process_single(
+        self,
+        text: Union[str, TextInput],
+        voice_samples: Optional[List[Union[str, np.ndarray]]] = None,
+    ) -> Dict[str, Any]:
+        """Process a single podcast script."""
+        # Determine if text is a file path or direct script
+        script = None
+        if isinstance(text, str):
+            # Check if it's a file path
+            if text.endswith('.json') and os.path.exists(text):
+                script = self._convert_json_to_script(text)
+            elif text.endswith('.txt') and os.path.exists(text):
+                script = self._convert_text_to_script(text)
+            else:
+                # Assume it's the script content directly
+                script = text
+        
+        if script is None:
+            raise ValueError(f"Could not process input text: {text}")
+        
+        # Parse the script
+        parsed_lines = self._parse_script(script)
+        all_speakers = list(set(speaker_id for speaker_id, _ in parsed_lines))
+        
+        # Create system prompt
+        # system_tokens = self.tokenizer.encode(self.system_prompt, add_special_tokens=False)
+        system_tokens = self.tokenizer.encode(self.system_prompt)
+        
+        # Process voice samples if provided
+        if voice_samples:
+            voice_tokens, voice_speech_inputs, voice_speech_masks = self._create_voice_prompt(voice_samples[:len(all_speakers)])
+        else:
+            voice_tokens, voice_speech_inputs, voice_speech_masks = [], [], []
+        
+        # Build full token sequence
+        full_tokens = system_tokens + voice_tokens
+        speech_input_mask = [False] * len(system_tokens) + voice_speech_masks
+        
+        # Add text input section
+        full_tokens += self.tokenizer.encode(' Text input:\n', add_special_tokens=False)
+        speech_input_mask += [False] * len(self.tokenizer.encode(' Text input:\n', add_special_tokens=False))
+        
+        for speaker_id, speaker_text in parsed_lines:
+            speaker_text_tokens = self.tokenizer.encode(f" Speaker {speaker_id}:{speaker_text}\n", add_special_tokens=False)
+            full_tokens += speaker_text_tokens
+            speech_input_mask += [False] * len(speaker_text_tokens)
+        
+        # Add speech output section
+        full_tokens += self.tokenizer.encode(' Speech output:\n', add_special_tokens=False) + [self.tokenizer.speech_start_id]
+        speech_input_mask += [False] * (len(self.tokenizer.encode(' Speech output:\n', add_special_tokens=False)) + 1)
+        
+        return {
+            "input_ids": full_tokens,
+            "speech_inputs": voice_speech_inputs if voice_speech_inputs else None,
+            "speech_input_mask": speech_input_mask,
+            "parsed_script": parsed_lines,
+            "all_speakers": all_speakers,
+        }
+    
+    def _batch_encode(
+        self,
+        encodings: List[Dict[str, Any]],
+        padding: Union[bool, str, PaddingStrategy] = True,
+        truncation: Union[bool, str, TruncationStrategy] = False,
+        max_length: Optional[int] = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        return_attention_mask: bool = True,
+    ) -> BatchEncoding:
+        """Combine multiple encodings into a batch with padding."""
+        # Extract input_ids and create attention_mask
+        input_ids_list = [enc["input_ids"] for enc in encodings]
+        speech_input_masks_list = [enc["speech_input_mask"] for enc in encodings]
+        
+        # Determine padding strategy
+        if isinstance(padding, bool):
+            padding_strategy = PaddingStrategy.LONGEST if padding else PaddingStrategy.DO_NOT_PAD
+        elif isinstance(padding, str):
+            padding_strategy = PaddingStrategy(padding)
+        else:
+            padding_strategy = padding
+            
+        # Apply padding to input_ids
+        if padding_strategy != PaddingStrategy.DO_NOT_PAD:
+            if padding_strategy == PaddingStrategy.LONGEST:
+                max_len = max(len(ids) for ids in input_ids_list)
+            elif padding_strategy == PaddingStrategy.MAX_LENGTH and max_length is not None:
+                max_len = max_length
+            else:
+                max_len = max(len(ids) for ids in input_ids_list)
+                
+            # Pad sequences
+            padded_input_ids = []
+            attention_masks = []
+            padded_speech_input_masks = []
+            
+            for input_ids, speech_mask in zip(input_ids_list, speech_input_masks_list):
+                # Truncate if needed
+                if truncation and len(input_ids) > max_len:
+                    input_ids = input_ids[:max_len]
+                    speech_mask = speech_mask[:max_len]
+                    
+                # Pad
+                padding_length = max_len - len(input_ids)
+                # padded_ids = [self.tokenizer.pad_token_id] * padding_length + input_ids
+                padded_ids = [self.tokenizer.pad_id] * padding_length + input_ids
+                attention_mask = [0] * padding_length + [1] * len(input_ids)
+                padded_speech_mask = [False] * padding_length + speech_mask
+                
+                padded_input_ids.append(padded_ids)
+                attention_masks.append(attention_mask)
+                padded_speech_input_masks.append(padded_speech_mask)
+                
+            input_ids_list = padded_input_ids
+            speech_input_masks_list = padded_speech_input_masks
+        else:
+            # No padding, just create attention masks
+            attention_masks = [[1] * len(ids) for ids in input_ids_list] if return_attention_mask else None
+            
+        # Process speech inputs
+        all_speech_inputs = []
+        has_speech = False
+        for enc in encodings:
+            if enc["speech_inputs"] is not None:
+                all_speech_inputs.extend(enc["speech_inputs"])
+                has_speech = True
+                
+        # Prepare batch encoding
+        batch_encoding = BatchEncoding()
+        
+        # Handle tensor conversion
+        if return_tensors is not None:
+            batch_encoding["input_ids"] = torch.tensor(input_ids_list, dtype=torch.long)
+            if return_attention_mask and attention_masks is not None:
+                batch_encoding["attention_mask"] = torch.tensor(attention_masks, dtype=torch.long)
+            batch_encoding["speech_input_mask"] = torch.tensor(speech_input_masks_list, dtype=torch.bool)
+        else:
+            batch_encoding["input_ids"] = input_ids_list
+            if return_attention_mask and attention_masks is not None:
+                batch_encoding["attention_mask"] = attention_masks
+            batch_encoding["speech_input_mask"] = speech_input_masks_list
+            
+        # Process speech tensors if present
+        if has_speech:
+            speech_dict = self.prepare_speech_inputs(
+                all_speech_inputs,
+                return_tensors=return_tensors,
+            )
+            batch_encoding["speech_tensors"] = speech_dict["padded_speeches"]
+            batch_encoding["speech_masks"] = speech_dict["speech_masks"]
+        else:
+            batch_encoding["speech_tensors"] = None
+            batch_encoding["speech_masks"] = None
+            
+        # Add metadata
+        batch_encoding["parsed_scripts"] = [enc["parsed_script"] for enc in encodings]
+        batch_encoding["all_speakers_list"] = [enc["all_speakers"] for enc in encodings]
+        
+        return batch_encoding
+
+    def _create_voice_prompt(
+        self, 
+        speaker_samples: List[Union[str, np.ndarray]]
+    ) -> Tuple[List[int], List[np.ndarray], List[bool]]:
+        """
+        Create voice prompt tokens and process audio samples.
+        
+        Returns:
+            tuple: (voice_tokens, voice_speech_inputs, voice_speech_masks)
+        """
+        vae_token_id = self.tokenizer.speech_diffusion_id
+        
+        voice_full_tokens = self.tokenizer.encode(' Voice input:\n', add_special_tokens=False)
+        voice_speech_inputs = []
+        voice_speech_masks = [False] * len(voice_full_tokens)
+        
+        for speaker_id, speaker_audio in enumerate(speaker_samples):
+            prefix_tokens = self.tokenizer.encode(f" Speaker {speaker_id}:", add_special_tokens=False)
+            
+            # Process audio
+            if isinstance(speaker_audio, str):
+                # Load audio from file
+                wav = self.audio_processor._load_audio_from_path(speaker_audio)
+            else:
+                wav = np.array(speaker_audio, dtype=np.float32)
+            
+            # Apply normalization if needed
+            if self.db_normalize and self.audio_normalizer:
+                wav = self.audio_normalizer(wav)
+            
+            # Calculate token length based on compression ratio
+            # if speaker_audio.endswith('.pt') or speaker_audio.endswith('.npy'):
+            #     vae_tok_len = wav.shape[0]
+            # else:
+            vae_tok_len = math.ceil(wav.shape[0] / self.speech_tok_compress_ratio)
+            
+            # Build tokens and masks
+            speaker_tokens = (prefix_tokens + 
+                            [self.tokenizer.speech_start_id] + 
+                            [vae_token_id] * vae_tok_len + 
+                            [self.tokenizer.speech_end_id] + 
+                            self.tokenizer.encode('\n', add_special_tokens=False))
+            
+            vae_input_mask = ([False] * len(prefix_tokens) + 
+                            [False] + 
+                            [True] * vae_tok_len + 
+                            [False] + 
+                            [False])
+            
+            voice_full_tokens.extend(speaker_tokens)
+            voice_speech_masks.extend(vae_input_mask)
+            voice_speech_inputs.append(wav)
+            
+        return voice_full_tokens, voice_speech_inputs, voice_speech_masks
+
+    def prepare_speech_inputs(
+        self,
+        speech_inputs: List[np.ndarray],
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        device: Optional[Union[str, torch.device]] = None,
+        dtype: Optional[torch.dtype] = None,
+    ) -> Dict[str, Any]:
+        """
+        Prepare speech inputs for model consumption.
+        
+        Args:
+            speech_inputs: List of speech arrays
+            return_tensors: Output tensor type
+            device: Device to place tensors on
+            dtype: Data type for tensors
+            
+        Returns:
+            Dictionary with padded_speeches and speech_masks
+        """
+        if not speech_inputs:
+            return {"padded_speeches": None, "speech_masks": None}
+        
+        # Calculate sequence lengths
+        vae_tok_seqlens = [math.ceil(s.shape[0] / self.speech_tok_compress_ratio) for s in speech_inputs]
+        # vae_tok_seqlens = [math.ceil(s.shape[0] / self.speech_tok_compress_ratio) if s.ndim == 1 else s.shape[0] for s in speech_inputs]
+        max_speech_length = max(s.shape[0] for s in speech_inputs)
+        
+        # Pad speeches
+        if speech_inputs[0].ndim == 1:
+            padded_speeches = np.full((len(speech_inputs), max_speech_length), fill_value=0, dtype=np.float32)
+        else:
+            padded_speeches = np.full((len(speech_inputs), max_speech_length, speech_inputs[0].shape[-1]), fill_value=0, dtype=np.float32)
+        speech_masks = np.zeros((len(speech_inputs), max(vae_tok_seqlens)), dtype=np.bool_)
+        
+        for i, (speech, vae_tok_length) in enumerate(zip(speech_inputs, vae_tok_seqlens)):
+            padded_speeches[i, :len(speech)] = speech
+            speech_masks[i, :vae_tok_length] = True
+        
+        result = {
+            "padded_speeches": padded_speeches,
+            "speech_masks": speech_masks,
+        }
+        
+        # Convert to tensors if requested
+        if return_tensors == "pt":
+            result["padded_speeches"] = torch.tensor(padded_speeches, device=device, dtype=dtype or torch.float32)
+            result["speech_masks"] = torch.tensor(speech_masks, device=device, dtype=torch.bool)
+            
+        return result
+        
+    def _convert_json_to_script(self, json_file: str) -> str:
+        """
+        Convert JSON format to script format.
+        Expected JSON format:
+        [
+            {"speaker": "1", "text": "Hello everyone..."},
+            {"speaker": "2", "text": "Great to be here..."}
+        ]
+        """
+        import json
+        
+        with open(json_file, 'r', encoding='utf-8') as f:
+            data = json.load(f)
+        
+        if not isinstance(data, list):
+            raise ValueError("JSON file must contain a list of speaker entries")
+        
+        script_lines = []
+        for item in data:
+            if not isinstance(item, dict):
+                logger.warning(f"Skipping non-dict entry: {item}")
+                continue
+                
+            speaker = item.get('speaker')
+            text = item.get('text')
+            
+            if speaker is None or text is None:
+                logger.warning(f"Skipping entry missing speaker or text: {item}")
+                continue
+            
+            # Ensure speaker ID is valid
+            try:
+                speaker_id = int(speaker)
+            except (ValueError, TypeError):
+                logger.warning(f"Invalid speaker ID: {speaker}, skipping entry")
+                continue
+            
+            # Clean up text
+            text = text.strip()
+            if text:
+                script_lines.append(f"Speaker {speaker_id}: {text}")
+        
+        if not script_lines:
+            raise ValueError("No valid entries found in JSON file")
+            
+        return "\n".join(script_lines)
+
+    def _convert_text_to_script(self, text_file: str) -> str:
+        """
+        Convert text file to script format.
+        Handles multiple formats:
+        1. Already formatted as "Speaker X: text"
+        2. Plain text (assigns to Speaker 1)
+        
+        Handles edge cases like multiple colons in a line.
+        """
+        with open(text_file, 'r', encoding='utf-8') as f:
+            lines = f.readlines()
+        
+        script_lines = []
+        current_speaker = 1
+        
+        for line in lines:
+            line = line.strip()
+            if not line:
+                continue
+            
+            # Try to parse as "Speaker X: text" format
+            # Use regex to be more robust
+            speaker_match = re.match(r'^Speaker\s+(\d+)\s*:\s*(.*)$', line, re.IGNORECASE)
+            
+            if speaker_match:
+                speaker_id = int(speaker_match.group(1))
+                text = speaker_match.group(2).strip()
+                if text:
+                    script_lines.append(f"Speaker {speaker_id}: {text}")
+            else:
+                # Treat as plain text - assign to current speaker
+                script_lines.append(f"Speaker {current_speaker}: {line}")
+        
+        if not script_lines:
+            raise ValueError("No valid content found in text file")
+            
+        return "\n".join(script_lines)
+
+    def _parse_script(self, script: str) -> List[Tuple[int, str]]:
+        """Parse script into list of (speaker_id, text) tuples."""
+        lines = script.strip().split("\n")
+        parsed_lines = []
+        speaker_ids = []
+                
+        # First pass: parse all lines and collect speaker IDs
+        for line in lines:
+            if not line.strip():
+                continue
+                
+            # Use regex to handle edge cases like multiple colons
+            match = re.match(r'^Speaker\s+(\d+)\s*:\s*(.*)$', line.strip(), re.IGNORECASE)
+            
+            if match:
+                speaker_id = int(match.group(1))
+                text = ' ' + match.group(2).strip()
+                parsed_lines.append((speaker_id, text))
+                speaker_ids.append(speaker_id)
+            else:
+                logger.warning(f"Could not parse line: '{line}'")
+        
+        if not parsed_lines:
+            raise ValueError("No valid speaker lines found in script")
+        
+        # Check if we need to normalize speaker IDs (only if all are > 0)
+        min_speaker_id = min(speaker_ids)
+        if min_speaker_id > 0:
+            # Normalize to start from 0
+            normalized_lines = []
+            for speaker_id, text in parsed_lines:
+                normalized_lines.append((speaker_id - 1, text))
+            return normalized_lines
+        else:
+            # Keep original IDs
+            return parsed_lines
+
+    def _merge_inputs(self, text_inputs: BatchEncoding, audio_inputs: Dict) -> BatchEncoding:
+        """Merge text and audio inputs into a single BatchEncoding."""
+        # Start with text inputs
+        merged = BatchEncoding(text_inputs)
+        
+        # Add audio-specific fields
+        if "audio" in audio_inputs:
+            merged["speech_inputs"] = audio_inputs["audio"]
+        if "streaming" in audio_inputs:
+            merged["streaming"] = audio_inputs["streaming"]
+            
+        return merged
+
+    def batch_decode(self, *args, **kwargs):
+        """
+        This method forwards all its arguments to VibeVoiceTextTokenizer's [`~PreTrainedTokenizer.batch_decode`].
+        Please refer to the docstring of this method for more information.
+        """
+        return self.tokenizer.batch_decode(*args, **kwargs)
+
+    def decode(self, *args, **kwargs):
+        """
+        This method forwards all its arguments to VibeVoiceTextTokenizer's [`~PreTrainedTokenizer.decode`].
+        Please refer to the docstring of this method for more information.
+        """
+        return self.tokenizer.decode(*args, **kwargs)
+
+    @property
+    def model_input_names(self):
+        """
+        Return the list of inputs accepted by the model.
+        """
+        tokenizer_input_names = self.tokenizer.model_input_names
+        audio_processor_input_names = self.audio_processor.model_input_names
+        return list(dict.fromkeys(tokenizer_input_names + audio_processor_input_names + ["speech_inputs", "speech_input_mask"]))
+
+    def save_audio(self, 
+        audio: Union[torch.Tensor, np.ndarray, List[Union[torch.Tensor, np.ndarray]]],
+        output_path: str = "output.wav",
+        sampling_rate: Optional[int] = None,
+        normalize: bool = False,
+        batch_prefix: str = "audio_",
+    ) -> str:
+        """
+        Save audio data to a file.
+        Args:
+            audio (Union[torch.Tensor, np.ndarray, List[Union[torch.Tensor, np.ndarray]]]):
+                The audio data to save. Can be a single tensor/array or a list of them.
+            output_path (str, optional): Path to save the audio file. Defaults to "output.wav".
+            sampling_rate (int, optional): Sampling rate for the audio. If None, uses the processor's default.
+            normalize (bool, optional): Whether to normalize the audio before saving. Defaults to False.
+            batch_prefix (str, optional): Prefix for batch audio files. Defaults to "audio_".
+        Returns:
+            str: The path to the saved audio file.
+        """
+        return self.audio_processor.save_audio(audio, output_path=output_path, sampling_rate=sampling_rate, normalize=normalize, batch_prefix=batch_prefix)
+    
+__all__ = [
+    "VibeVoiceProcessor",
+]
\ No newline at end of file
diff --git a/vibevoice/processor/vibevoice_streaming_processor.py b/vibevoice/processor/vibevoice_streaming_processor.py
new file mode 100644
index 0000000..39c262b
--- /dev/null
+++ b/vibevoice/processor/vibevoice_streaming_processor.py
@@ -0,0 +1,409 @@
+import math
+import warnings
+from typing import List, Optional, Union, Dict, Any, Tuple
+import os
+import re
+
+import numpy as np
+import torch
+
+from transformers.tokenization_utils_base import BatchEncoding, PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
+from transformers.utils import TensorType, logging
+from .vibevoice_tokenizer_processor import AudioNormalizer
+
+logger = logging.get_logger(__name__)
+
+
+class VibeVoiceStreamingProcessor:
+    r"""
+    Constructs a VibeVoice Streaming processor which wraps a VibeVoice tokenizer and audio processor into a single processor.
+
+    Args:
+        tokenizer (`VibeVoiceTextTokenizer` or `VibeVoiceTextTokenizerFast`):
+            The tokenizer for text processing.
+        audio_processor (`VibeVoiceTokenizerProcessor`):
+            The audio processor for speech processing.
+        speech_tok_compress_ratio (`int`, *optional*, defaults to 3200):
+            The compression ratio for speech tokenization.
+        db_normalize (`bool`, *optional*, defaults to True):
+            Whether to apply decibel normalization to audio inputs.
+    """
+
+    def __init__(self, tokenizer=None, audio_processor=None, speech_tok_compress_ratio=3200, db_normalize=True, **kwargs):
+        self.tokenizer = tokenizer
+        self.audio_processor = audio_processor
+        self.speech_tok_compress_ratio = speech_tok_compress_ratio
+        self.db_normalize = db_normalize
+        self.audio_normalizer = AudioNormalizer() if db_normalize else None
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        """
+        Instantiate a VibeVoiceStreamingProcessor from a pretrained VibeVoice Streaming processor.
+
+        Args:
+            pretrained_model_name_or_path (`str` or `os.PathLike`):
+                This can be either:
+                - a string, the *model id* of a pretrained model
+                - a path to a *directory* containing processor config
+
+        Returns:
+            [`VibeVoiceStreamingProcessor`]: The processor object instantiated from pretrained model.
+        """
+        import os
+        import json
+        from transformers.utils import cached_file
+        from .vibevoice_tokenizer_processor import VibeVoiceTokenizerProcessor
+        from vibevoice.modular.modular_vibevoice_text_tokenizer import (
+            VibeVoiceTextTokenizer, 
+            VibeVoiceTextTokenizerFast
+        )
+        
+        # Try to load from local path first, then from HF hub
+        config_path = os.path.join(pretrained_model_name_or_path, "preprocessor_config.json")
+        config = None
+        
+        if os.path.exists(config_path):
+            # Local path exists
+            with open(config_path, 'r') as f:
+                config = json.load(f)
+        else:
+            # Try to load from HF hub
+            try:
+                config_file = cached_file(
+                    pretrained_model_name_or_path,
+                    "preprocessor_config.json",
+                    **kwargs
+                )
+                with open(config_file, 'r') as f:
+                    config = json.load(f)
+            except Exception as e:
+                logger.warning(f"Could not load preprocessor_config.json from {pretrained_model_name_or_path}: {e}")
+                logger.warning("Using default configuration")
+                config = {
+                    "speech_tok_compress_ratio": 3200,
+                    "db_normalize": True,
+                }
+        
+        # Extract main processor parameters
+        speech_tok_compress_ratio = config.get("speech_tok_compress_ratio", 3200)
+        db_normalize = config.get("db_normalize", True)
+        
+        # Load tokenizer - try from model path first, then fall back to Qwen        
+        language_model_pretrained_name = config.get("language_model_pretrained_name", None) or kwargs.pop("language_model_pretrained_name", "Qwen/Qwen2.5-1.5B")
+        logger.info(f"Loading tokenizer from {language_model_pretrained_name}")
+        if 'qwen' in language_model_pretrained_name.lower():
+            tokenizer = VibeVoiceTextTokenizerFast.from_pretrained(
+                language_model_pretrained_name,
+                **kwargs
+            )
+        else:
+            raise ValueError(f"Unsupported tokenizer type for {language_model_pretrained_name}. Supported types: Qwen, Llama, Gemma.")
+        
+        # Load audio processor
+        if "audio_processor" in config:
+            # Create audio processor from config
+            audio_config = config["audio_processor"]
+            audio_processor = VibeVoiceTokenizerProcessor(
+                sampling_rate=audio_config.get("sampling_rate", 24000),
+                normalize_audio=audio_config.get("normalize_audio", True),
+                target_dB_FS=audio_config.get("target_dB_FS", -25),
+                eps=audio_config.get("eps", 1e-6),
+            )
+        else:
+            # Create default audio processor
+            audio_processor = VibeVoiceTokenizerProcessor()
+        
+        # Create and return the processor
+        return cls(
+            tokenizer=tokenizer,
+            audio_processor=audio_processor,
+            speech_tok_compress_ratio=speech_tok_compress_ratio,
+            db_normalize=db_normalize,
+        )
+    
+    def save_pretrained(self, save_directory: Union[str, os.PathLike], **kwargs):
+        """
+        Save a processor to a directory, so that it can be re-loaded using the
+        [`~VibeVoiceStreamingProcessor.from_pretrained`] class method.
+
+        Args:
+            save_directory (`str` or `os.PathLike`):
+                Directory where the processor will be saved.
+        """
+        import os
+        import json
+        
+        os.makedirs(save_directory, exist_ok=True)
+        
+        # Save processor configuration
+        processor_config = {
+            "processor_class": "VibeVoiceStreamingProcessor",
+            "speech_tok_compress_ratio": self.speech_tok_compress_ratio,
+            "db_normalize": self.db_normalize,
+            "audio_processor": {
+                "feature_extractor_type": "VibeVoiceTokenizerProcessor",
+                "sampling_rate": getattr(self.audio_processor, 'sampling_rate', 24000),
+                "normalize_audio": getattr(self.audio_processor, 'normalize_audio', True),
+                "target_dB_FS": getattr(self.audio_processor, 'target_dB_FS', -25),
+                "eps": getattr(self.audio_processor, 'eps', 1e-6),
+            }
+        }
+        
+        config_path = os.path.join(save_directory, "preprocessor_config.json")
+        with open(config_path, 'w') as f:
+            json.dump(processor_config, f, indent=2)
+        
+        logger.info(f"Processor configuration saved in {config_path}")
+    
+    def __call__(self) -> BatchEncoding:
+        """
+        Note:
+            This method is intentionally not implemented in the streaming processor.
+            Use `process_input_with_cached_prompt` for streaming use cases.
+        """
+        raise NotImplementedError(
+            "VibeVoiceStreamingProcessor.__call__ is not implemented. "
+            "Use process_input_with_cached_prompt for streaming inputs."
+        )
+
+    def process_input_with_cached_prompt(
+        self,
+        text: Optional[str] = None,
+        cached_prompt: Optional[Dict[str, Any]] = None,
+        padding: Union[bool, str, PaddingStrategy] = True,
+        truncation: Union[bool, str, TruncationStrategy] = False,
+        max_length: Optional[int] = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        return_attention_mask: bool = True,
+        **kwargs,
+    ) -> BatchEncoding:
+        """
+        Main method to process one text script based on cached prompt. The function currently only supports single examples.
+
+        Args:
+            text (`str`):
+                The input text to process.
+            cached_prompt (`Dict[str, Any]`, *optional*):
+                The cached prompt to use for processing. It contains the kv cache of the voice prompt.
+            padding (`bool`, `str` or `PaddingStrategy`, defaults to `True`):
+                Whether to pad sequences to the same length
+            truncation (`bool`, `str` or `TruncationStrategy`, defaults to `False`):
+                Whether to truncate sequences
+            max_length (`int`, *optional*):
+                Maximum length of the returned sequences
+            return_tensors (`str` or `TensorType`, *optional*):
+                If set, will return tensors of a particular framework
+            return_attention_mask (`bool`, defaults to `True`):
+                Whether to return the attention mask
+
+        Returns:
+            `BatchEncoding`: A BatchEncoding with the following fields:
+                - **input_ids** -- List of token id sequences or tensor
+                - **attention_mask** -- List of attention masks or tensor
+                - **tts_lm_input_ids** -- List of token id sequences or tensor used for TTS LM
+                - **tts_lm_attention_mask** -- List of attention masks or tensor used for TTS LM
+                - **tts_text_ids** -- List of token id sequences or tensor for TTS text input
+                - **speech_tensors** -- Padded speech inputs (if voice_samples provided)
+                - **speech_masks** -- Speech masks (if voice_samples provided)
+                - **speech_input_mask** -- Boolean masks indicating speech token positions
+        """
+        # Only support single example
+        texts = [text]
+        cached_prompts = [cached_prompt]
+        is_batched = False
+        
+        # Process each input
+        all_encodings = []
+        for text_input, cached_prompt_input in zip(texts, cached_prompts):
+            script_tokens = self.tokenizer.encode(text_input.strip() + "\n", add_special_tokens=False)
+            input_id_length = cached_prompt_input['lm']['last_hidden_state'].size(1)
+            tts_lm_input_id_length = cached_prompt_input['tts_lm']['last_hidden_state'].size(1)
+
+            # pseudo input ids and masks
+            input_ids = [self.tokenizer.pad_id] * input_id_length
+            tts_lm_input_ids = [self.tokenizer.pad_id] * tts_lm_input_id_length
+            speech_input_mask = [False] * tts_lm_input_id_length
+
+            encoding = {
+                        "input_ids": input_ids,
+                        "tts_lm_input_ids": tts_lm_input_ids,
+                        "tts_text_ids": script_tokens,
+                        "speech_inputs": None,
+                        "speech_input_mask": speech_input_mask,
+                    }
+            all_encodings.append(encoding)
+            
+        # Combine batch
+        batch_encoding = self._batch_encode(
+            all_encodings,
+            padding=padding,
+            truncation=truncation,
+            max_length=max_length,
+            return_tensors=return_tensors,
+            return_attention_mask=return_attention_mask,
+        )
+        
+        return batch_encoding
+    
+    def _batch_encode(
+        self,
+        encodings: List[Dict[str, Any]],
+        padding: Union[bool, str, PaddingStrategy] = True,
+        truncation: Union[bool, str, TruncationStrategy] = False,
+        max_length: Optional[int] = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        return_attention_mask: bool = True,
+    ) -> BatchEncoding:
+        """Combine multiple encodings into a batch with padding."""
+        # Extract input_ids and create attention_mask
+        input_ids_list = [enc["input_ids"] for enc in encodings]
+        tts_lm_input_ids_list = [enc["tts_lm_input_ids"] for enc in encodings]
+        tts_text_ids_list = [enc["tts_text_ids"] for enc in encodings]
+        speech_input_masks_list = [enc["speech_input_mask"] for enc in encodings]
+        
+        attention_masks = [[1] * len(ids) for ids in input_ids_list] if return_attention_mask else None
+        tts_lm_attention_masks = [[1] * len(ids) for ids in tts_lm_input_ids_list] if return_attention_mask else None
+            
+        # Process speech inputs
+        all_speech_inputs = []
+        has_speech = False
+        for enc in encodings:
+            if enc["speech_inputs"] is not None:
+                all_speech_inputs.extend(enc["speech_inputs"])
+                has_speech = True
+                
+        # Prepare batch encoding
+        batch_encoding = BatchEncoding()
+        
+        # Handle tensor conversion
+        if return_tensors is not None:
+            batch_encoding["input_ids"] = torch.tensor(input_ids_list, dtype=torch.long)
+            batch_encoding["tts_lm_input_ids"] = torch.tensor(tts_lm_input_ids_list, dtype=torch.long)
+            batch_encoding["tts_text_ids"] = torch.tensor(tts_text_ids_list, dtype=torch.long)
+
+            if return_attention_mask and attention_masks is not None:
+                batch_encoding["attention_mask"] = torch.tensor(attention_masks, dtype=torch.long)
+                batch_encoding["tts_lm_attention_mask"] = torch.tensor(tts_lm_attention_masks, dtype=torch.long)
+            
+            batch_encoding["speech_input_mask"] = torch.tensor(speech_input_masks_list, dtype=torch.bool)
+        else:
+            batch_encoding["input_ids"] = input_ids_list
+            batch_encoding["tts_lm_input_ids"] = tts_lm_input_ids_list
+            batch_encoding["tts_text_ids"] = tts_text_ids_list
+            if return_attention_mask and attention_masks is not None:
+                batch_encoding["attention_mask"] = attention_masks
+                batch_encoding["tts_lm_attention_mask"] = tts_lm_attention_masks
+            batch_encoding["speech_input_mask"] = speech_input_masks_list
+            
+        # Process speech tensors if present
+        if has_speech:
+            speech_dict = self.prepare_speech_inputs(
+                all_speech_inputs,
+                return_tensors=return_tensors,
+            )
+            batch_encoding["speech_tensors"] = speech_dict["padded_speeches"]
+            batch_encoding["speech_masks"] = speech_dict["speech_masks"]
+        else:
+            batch_encoding["speech_tensors"] = None
+            batch_encoding["speech_masks"] = None
+            
+        return batch_encoding
+
+    def prepare_speech_inputs(
+        self,
+        speech_inputs: List[np.ndarray],
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        device: Optional[Union[str, torch.device]] = None,
+        dtype: Optional[torch.dtype] = None,
+    ) -> Dict[str, Any]:
+        """
+        Prepare speech inputs for model consumption.
+        
+        Args:
+            speech_inputs: List of speech arrays
+            return_tensors: Output tensor type
+            device: Device to place tensors on
+            dtype: Data type for tensors
+            
+        Returns:
+            Dictionary with padded_speeches and speech_masks
+        """
+        if not speech_inputs:
+            return {"padded_speeches": None, "speech_masks": None}
+        
+        # Calculate sequence lengths
+        vae_tok_seqlens = [math.ceil(s.shape[0] / self.speech_tok_compress_ratio) for s in speech_inputs]
+        # vae_tok_seqlens = [math.ceil(s.shape[0] / self.speech_tok_compress_ratio) if s.ndim == 1 else s.shape[0] for s in speech_inputs]
+        max_speech_length = max(s.shape[0] for s in speech_inputs)
+        
+        # Pad speeches
+        if speech_inputs[0].ndim == 1:
+            padded_speeches = np.full((len(speech_inputs), max_speech_length), fill_value=0, dtype=np.float32)
+        else:
+            padded_speeches = np.full((len(speech_inputs), max_speech_length, speech_inputs[0].shape[-1]), fill_value=0, dtype=np.float32)
+        speech_masks = np.zeros((len(speech_inputs), max(vae_tok_seqlens)), dtype=np.bool_)
+        
+        for i, (speech, vae_tok_length) in enumerate(zip(speech_inputs, vae_tok_seqlens)):
+            padded_speeches[i, :len(speech)] = speech
+            speech_masks[i, :vae_tok_length] = True
+        
+        result = {
+            "padded_speeches": padded_speeches,
+            "speech_masks": speech_masks,
+        }
+        
+        # Convert to tensors if requested
+        if return_tensors == "pt":
+            result["padded_speeches"] = torch.tensor(padded_speeches, device=device, dtype=dtype or torch.float32)
+            result["speech_masks"] = torch.tensor(speech_masks, device=device, dtype=torch.bool)
+            
+        return result
+
+    def batch_decode(self, *args, **kwargs):
+        """
+        This method forwards all its arguments to VibeVoiceTextTokenizer's [`~PreTrainedTokenizer.batch_decode`].
+        Please refer to the docstring of this method for more information.
+        """
+        return self.tokenizer.batch_decode(*args, **kwargs)
+
+    def decode(self, *args, **kwargs):
+        """
+        This method forwards all its arguments to VibeVoiceTextTokenizer's [`~PreTrainedTokenizer.decode`].
+        Please refer to the docstring of this method for more information.
+        """
+        return self.tokenizer.decode(*args, **kwargs)
+
+    @property
+    def model_input_names(self):
+        """
+        Return the list of inputs accepted by the model.
+        """
+        tokenizer_input_names = self.tokenizer.model_input_names
+        audio_processor_input_names = self.audio_processor.model_input_names
+        return list(dict.fromkeys(tokenizer_input_names + audio_processor_input_names + ["speech_inputs", "speech_input_mask"]))
+
+    def save_audio(self, 
+        audio: Union[torch.Tensor, np.ndarray, List[Union[torch.Tensor, np.ndarray]]],
+        output_path: str = "output.wav",
+        sampling_rate: Optional[int] = None,
+        normalize: bool = False,
+        batch_prefix: str = "audio_",
+    ) -> str:
+        """
+        Save audio data to a file.
+        Args:
+            audio (Union[torch.Tensor, np.ndarray, List[Union[torch.Tensor, np.ndarray]]]):
+                The audio data to save. Can be a single tensor/array or a list of them.
+            output_path (str, optional): Path to save the audio file. Defaults to "output.wav".
+            sampling_rate (int, optional): Sampling rate for the audio. If None, uses the processor's default.
+            normalize (bool, optional): Whether to normalize the audio before saving. Defaults to False.
+            batch_prefix (str, optional): Prefix for batch audio files. Defaults to "audio_".
+        Returns:
+            str: The path to the saved audio file.
+        """
+        return self.audio_processor.save_audio(audio, output_path=output_path, sampling_rate=sampling_rate, normalize=normalize, batch_prefix=batch_prefix)
+    
+__all__ = [
+    "VibeVoiceStreamingProcessor",
+]
\ No newline at end of file
diff --git a/vibevoice/processor/vibevoice_tokenizer_processor.py b/vibevoice/processor/vibevoice_tokenizer_processor.py
new file mode 100644
index 0000000..67f61a6
--- /dev/null
+++ b/vibevoice/processor/vibevoice_tokenizer_processor.py
@@ -0,0 +1,413 @@
+"""
+Processor class for VibeVoice models.
+"""
+
+import os
+import json
+import warnings
+from typing import List, Optional, Union, Dict, Any
+
+import numpy as np
+import torch
+
+from transformers.feature_extraction_utils import FeatureExtractionMixin
+from transformers.utils import logging
+
+from .audio_utils import AudioNormalizer
+
+logger = logging.get_logger(__name__)
+
+# Change from ProcessorMixin to FeatureExtractionMixin which is designed for single components
+class VibeVoiceTokenizerProcessor(FeatureExtractionMixin):
+    """
+    Processor for VibeVoice acoustic tokenizer models.
+    
+    This processor handles audio preprocessing for VibeVoice models, including:
+    - Audio format conversion (stereo to mono)
+    - Optional audio normalization
+    - Streaming support for infinite-length audio
+    
+    Args:
+        sampling_rate (int, optional): Expected sampling rate. Defaults to 24000.
+        normalize_audio (bool, optional): Whether to normalize audio. Defaults to True.
+        target_dB_FS (float, optional): Target dB FS for normalization. Defaults to -25.
+        eps (float, optional): Small value for numerical stability. Defaults to 1e-6.
+    """
+    model_input_names = ["input_features"]
+    
+    def __init__(
+        self,
+        sampling_rate: int = 24000,
+        normalize_audio: bool = True,
+        target_dB_FS: float = -25,
+        eps: float = 1e-6,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        
+        self.sampling_rate = sampling_rate
+        self.normalize_audio = normalize_audio
+        
+        # Initialize audio normalizer if needed
+        if self.normalize_audio:
+            self.normalizer = AudioNormalizer(target_dB_FS=target_dB_FS, eps=eps)
+        else:
+            self.normalizer = None
+        
+        # Save config
+        self.feature_extractor_dict = {
+            "sampling_rate": sampling_rate,
+            "normalize_audio": normalize_audio,
+            "target_dB_FS": target_dB_FS,
+            "eps": eps,
+        }
+    
+    def _ensure_mono(self, audio: np.ndarray) -> np.ndarray:
+        """
+        Convert stereo audio to mono if needed.
+        
+        Args:
+            audio (np.ndarray): Input audio array
+            
+        Returns:
+            np.ndarray: Mono audio array
+        """
+        if len(audio.shape) == 1:
+            return audio
+        elif len(audio.shape) == 2:
+            if audio.shape[0] == 2:  # (2, time)
+                return np.mean(audio, axis=0)
+            elif audio.shape[1] == 2:  # (time, 2)
+                return np.mean(audio, axis=1)
+            else:
+                # If one dimension is 1, squeeze it
+                if audio.shape[0] == 1:
+                    return audio.squeeze(0)
+                elif audio.shape[1] == 1:
+                    return audio.squeeze(1)
+                else:
+                    raise ValueError(f"Unexpected audio shape: {audio.shape}")
+        else:
+            raise ValueError(f"Audio should be 1D or 2D, got shape: {audio.shape}")
+    
+    def _process_single_audio(self, audio: Union[np.ndarray, List[float]]) -> np.ndarray:
+        """
+        Process a single audio array.
+        
+        Args:
+            audio: Single audio input
+            
+        Returns:
+            np.ndarray: Processed audio
+        """
+        # Convert to numpy array
+        if not isinstance(audio, np.ndarray):
+            audio = np.array(audio, dtype=np.float32)
+        else:
+            audio = audio.astype(np.float32)
+        
+        # Ensure mono
+        audio = self._ensure_mono(audio)
+        
+        # Normalize if requested
+        if self.normalize_audio and self.normalizer is not None:
+            audio = self.normalizer(audio)
+        
+        return audio
+    
+    def __call__(
+        self,
+        audio: Union[str, np.ndarray, List[float], List[np.ndarray], List[List[float]], List[str]] = None,
+        sampling_rate: Optional[int] = None,
+        return_tensors: Optional[str] = None,
+        **kwargs,
+    ):
+        """
+        Process audio for VibeVoice models.
+        
+        Args:
+            audio: Audio input(s) to process. Can be:
+                - str: Path to audio file
+                - np.ndarray: Audio array
+                - List[float]: Audio as list of floats
+                - List[np.ndarray]: Batch of audio arrays
+                - List[str]: Batch of audio file paths
+            sampling_rate (int, optional): Sampling rate of the input audio
+            return_tensors (str, optional): Return format ('pt' for PyTorch, 'np' for NumPy)
+            
+        Returns:
+            dict: Processed audio inputs with keys:
+                - input_features: Audio tensor(s) ready for the model
+        """
+        if audio is None:
+            raise ValueError("Audio input is required")
+        
+        # Validate sampling rate
+        if sampling_rate is not None and sampling_rate != self.sampling_rate:
+            logger.warning(
+                f"Input sampling rate ({sampling_rate}) differs from expected "
+                f"sampling rate ({self.sampling_rate}). Please resample your audio."
+            )
+        
+        # Handle different input types
+        if isinstance(audio, str):
+            # Single audio file path
+            audio = self._load_audio_from_path(audio)
+            is_batched = False
+        elif isinstance(audio, list):
+            if len(audio) == 0:
+                raise ValueError("Empty audio list provided")
+            
+            # Check if it's a list of file paths
+            if all(isinstance(item, str) for item in audio):
+                # Batch of audio file paths
+                audio = [self._load_audio_from_path(path) for path in audio]
+                is_batched = True
+            else:
+                # Check if it's batched audio arrays
+                is_batched = isinstance(audio[0], (np.ndarray, list))
+        else:
+            # Single audio array or list
+            is_batched = False
+        
+        # Process audio
+        if is_batched:
+            processed_audio = [self._process_single_audio(a) for a in audio]
+        else:
+            processed_audio = [self._process_single_audio(audio)]
+        
+        # Convert to tensors if requested
+        if return_tensors == "pt":
+            if len(processed_audio) == 1:
+                # Create a proper batch dimension (B, T)
+                input_features = torch.from_numpy(processed_audio[0]).unsqueeze(0).unsqueeze(1)
+            else:
+                # For batched input with different lengths, create a batch properly
+                input_features = torch.stack([torch.from_numpy(a) for a in processed_audio]).unsqueeze(1)
+        elif return_tensors == "np":
+            if len(processed_audio) == 1:
+                input_features = processed_audio[0][np.newaxis, np.newaxis, :]
+            else:
+                input_features = np.stack(processed_audio)[:, np.newaxis, :]
+        else:
+            input_features = processed_audio[0] if len(processed_audio) == 1 else processed_audio
+        
+        outputs = {
+            "audio": input_features,  # Use "audio" instead of "input_features"
+        }
+        
+        return outputs
+
+    def _load_audio_from_path(self, audio_path: str) -> np.ndarray:
+        """
+        Load audio from file path.
+        
+        Args:
+            audio_path (str): Path to audio file
+            
+        Returns:
+            np.ndarray: Loaded audio array
+        """
+        # Get file extension to determine loading method
+        file_ext = os.path.splitext(audio_path)[1].lower()
+        
+        if file_ext in ['.wav', '.mp3', '.flac', '.m4a', '.ogg']:
+            # Audio file - use librosa
+            import librosa
+            audio_array, sr = librosa.load(
+                audio_path, 
+                sr=self.sampling_rate, 
+                mono=True
+            )
+            return audio_array
+        elif file_ext == '.pt':
+            # PyTorch tensor file
+            audio_tensor = torch.load(audio_path, map_location='cpu', weights_only=True).squeeze()
+            if isinstance(audio_tensor, torch.Tensor):
+                audio_array = audio_tensor.numpy()
+            else:
+                audio_array = np.array(audio_tensor)
+            return audio_array.astype(np.float32)
+        elif file_ext == '.npy':
+            # NumPy file
+            audio_array = np.load(audio_path)
+            return audio_array.astype(np.float32)
+        else:
+            raise ValueError(
+                f"Unsupported file format: {file_ext}. "
+                f"Supported formats: .wav, .mp3, .flac, .m4a, .ogg, .pt, .npy, .npz"
+            )
+    
+    def preprocess_audio(
+        self, 
+        audio_path_or_array: Union[str, np.ndarray],
+        normalize: Optional[bool] = None,
+    ) -> np.ndarray:
+        """
+        Convenience method to preprocess audio from file path or array.
+        This method is kept for backward compatibility but __call__ is recommended.
+        
+        Args:
+            audio_path_or_array: Path to audio file or numpy array
+            normalize: Whether to normalize (overrides default setting)
+            
+        Returns:
+            np.ndarray: Preprocessed audio array
+        """
+        if isinstance(audio_path_or_array, str):
+            audio_array = self._load_audio_from_path(audio_path_or_array)
+        else:
+            audio_array = np.array(audio_path_or_array, dtype=np.float32)
+        
+        # Override normalization setting if specified
+        original_normalize = self.normalize_audio
+        if normalize is not None:
+            self.normalize_audio = normalize
+        
+        try:
+            processed = self._process_single_audio(audio_array)
+        finally:
+            # Restore original setting
+            self.normalize_audio = original_normalize
+        
+        return processed
+    
+    # Override to_dict method for configuration saving
+    def to_dict(self) -> Dict[str, Any]:
+        """
+        Convert the object to a dict containing all attributes needed for serialization.
+        """
+        return self.feature_extractor_dict
+
+    def save_audio(
+        self,
+        audio: Union[torch.Tensor, np.ndarray, List[Union[torch.Tensor, np.ndarray]]],
+        output_path: str = "output.wav",
+        sampling_rate: Optional[int] = None,
+        normalize: bool = False,
+        batch_prefix: str = "audio_",
+    ):
+        """
+        Save audio data to WAV file(s).
+        
+        Args:
+            audio: Audio data to save. Can be:
+                - torch.Tensor: PyTorch tensor with shape (B, C, T) or (B, T) or (T)
+                - np.ndarray: NumPy array with shape (B, C, T) or (B, T) or (T)
+                - List of tensors or arrays
+            output_path: Path where to save the audio. If saving multiple files,
+                this is treated as a directory and individual files will be saved inside.
+            sampling_rate: Sampling rate for the saved audio. Defaults to the processor's rate.
+            normalize: Whether to normalize audio before saving.
+            batch_prefix: Prefix for batch files when saving multiple audios.
+                
+        Returns:
+            List[str]: Paths to the saved audio files.
+        """
+        if sampling_rate is None:
+            sampling_rate = self.sampling_rate
+        
+        try:
+            import soundfile as sf
+        except ImportError:
+            raise ImportError(
+                "soundfile is required to save audio files. "
+                "Install it with: pip install soundfile"
+            )
+        
+        # Ensure audio is in the right format
+        if isinstance(audio, torch.Tensor):
+            # Convert PyTorch tensor to numpy
+            audio_np = audio.float().detach().cpu().numpy()
+        elif isinstance(audio, np.ndarray):
+            audio_np = audio
+        elif isinstance(audio, list):
+            # Handle list of tensors or arrays
+            if all(isinstance(a, torch.Tensor) for a in audio):
+                audio_np = [a.float().detach().cpu().numpy() for a in audio]
+            else:
+                audio_np = audio
+        else:
+            raise ValueError(f"Unsupported audio type: {type(audio)}")
+        
+        saved_paths = []
+        
+        # Handle based on shape or type
+        if isinstance(audio_np, list):
+            # Multiple separate audios to save
+            output_dir = output_path
+            
+            # Ensure output directory exists
+            os.makedirs(output_dir, exist_ok=True)
+            
+            # Save each audio
+            for i, audio_item in enumerate(audio_np):
+                audio_item = self._prepare_audio_for_save(audio_item, normalize)
+                file_path = os.path.join(output_dir, f"{batch_prefix}{i}.wav")
+                sf.write(file_path, audio_item, sampling_rate)
+                saved_paths.append(file_path)
+                
+        else:
+            # Handle different dimensions
+            if len(audio_np.shape) >= 3:  # (B, C, T) or similar
+                # Get batch size
+                batch_size = audio_np.shape[0]
+                
+                if batch_size > 1:
+                    # Multiple audios in a batch
+                    output_dir = output_path
+                    
+                    # Ensure output directory exists
+                    os.makedirs(output_dir, exist_ok=True)
+                    
+                    # Save each audio in the batch
+                    for i in range(batch_size):
+                        # Extract single audio and remove channel dim if present
+                        single_audio = audio_np[i]
+                        if len(single_audio.shape) > 1:
+                            if single_audio.shape[0] == 1:  # (1, T)
+                                single_audio = single_audio.squeeze(0)
+                        
+                        single_audio = self._prepare_audio_for_save(single_audio, normalize)
+                        file_path = os.path.join(output_dir, f"{batch_prefix}{i}.wav")
+                        sf.write(file_path, single_audio, sampling_rate)
+                        saved_paths.append(file_path)
+                else:
+                    # Single audio with batch and channel dims
+                    audio_item = audio_np.squeeze()  # Remove batch and channel dimensions
+                    audio_item = self._prepare_audio_for_save(audio_item, normalize)
+                    sf.write(output_path, audio_item, sampling_rate)
+                    saved_paths.append(output_path)
+            else:
+                # Single audio without batch dimension
+                audio_item = self._prepare_audio_for_save(audio_np, normalize)
+                sf.write(output_path, audio_item, sampling_rate)
+                saved_paths.append(output_path)
+        
+        return saved_paths
+
+    def _prepare_audio_for_save(self, audio: np.ndarray, normalize: bool) -> np.ndarray:
+        """
+        Prepare audio for saving by ensuring it's the right shape and optionally normalizing.
+        
+        Args:
+            audio: Audio data as numpy array
+            normalize: Whether to normalize audio
+            
+        Returns:
+            np.ndarray: Processed audio ready for saving
+        """
+        # Ensure right dimensionality
+        if len(audio.shape) > 1 and audio.shape[0] == 1:  # (1, T)
+            audio = audio.squeeze(0)
+        
+        # Normalize if requested
+        if normalize:
+            max_val = np.abs(audio).max()
+            if max_val > 0:
+                audio = audio / max_val
+        
+        return audio
+
+
+__all__ = ["VibeVoiceTokenizerProcessor", "AudioNormalizer"]
\ No newline at end of file
diff --git a/vibevoice/schedule/__init__.py b/vibevoice/schedule/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/vibevoice/schedule/dpm_solver.py b/vibevoice/schedule/dpm_solver.py
new file mode 100644
index 0000000..b392a48
--- /dev/null
+++ b/vibevoice/schedule/dpm_solver.py
@@ -0,0 +1,1065 @@
+# Copyright 2024 TSAIL Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# DISCLAIMER: This file is strongly influenced by https://github.com/LuChengTHU/dpm-solver
+
+import math
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.utils import deprecate
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.schedulers.scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin, SchedulerOutput
+
+def betas_for_alpha_bar(
+    num_diffusion_timesteps,
+    max_beta=0.999,
+    alpha_transform_type="cosine",
+):
+    """
+    Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of
+    (1-beta) over time from t = [0,1].
+
+    Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up
+    to that part of the diffusion process.
+
+
+    Args:
+        num_diffusion_timesteps (`int`): the number of betas to produce.
+        max_beta (`float`): the maximum beta to use; use values lower than 1 to
+                     prevent singularities.
+        alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar.
+                     Choose from `cosine` or `exp`
+
+    Returns:
+        betas (`np.ndarray`): the betas used by the scheduler to step the model outputs
+    """
+    if alpha_transform_type == "cosine":
+
+        def alpha_bar_fn(t):
+            return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2
+            # return math.cos(t * math.pi / 2 * 0.95) ** 2
+
+    elif alpha_transform_type == "exp":
+
+        def alpha_bar_fn(t):
+            return math.exp(t * -12.0)
+
+    elif alpha_transform_type == "cauchy":
+        # µ + γ tan (π (0.5 - x))  γ = 1, µ = 3
+        # alpha^2 = 1-1/(exp(λ)+1)
+        def alpha_bar_fn(t, gamma=1, mu=3):
+            snr = mu + gamma * math.tan(math.pi * (0.5 - t) * 0.9)
+            return 1 - 1 / (math.exp(snr) + 1.1)
+
+    elif alpha_transform_type == "laplace":
+        # µ − bsgn(0.5 − t) log(1 − 2|t − 0.5|) µ = 0, b = 1
+        def alpha_bar_fn(t, mu=0, b=1):
+            snr = mu - b * math.copysign(1, 0.5 - t) * math.log(1 - 2 * abs(t - 0.5) * 0.98)
+            return 1 - 1 / (math.exp(snr) + 1.02)
+
+    else:
+        raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}")
+
+    betas = []
+    for i in range(num_diffusion_timesteps):
+        t1 = i / num_diffusion_timesteps
+        t2 = (i + 1) / num_diffusion_timesteps
+        betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta))
+    return torch.tensor(betas, dtype=torch.float32)
+
+
+# Copied from diffusers.schedulers.scheduling_ddim.rescale_zero_terminal_snr
+def rescale_zero_terminal_snr(betas):
+    """
+    Rescales betas to have zero terminal SNR Based on https://arxiv.org/pdf/2305.08891.pdf (Algorithm 1)
+
+
+    Args:
+        betas (`torch.Tensor`):
+            the betas that the scheduler is being initialized with.
+
+    Returns:
+        `torch.Tensor`: rescaled betas with zero terminal SNR
+    """
+    # Convert betas to alphas_bar_sqrt
+    alphas = 1.0 - betas
+    alphas_cumprod = torch.cumprod(alphas, dim=0)
+    alphas_bar_sqrt = alphas_cumprod.sqrt()
+
+    # Store old values.
+    alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone()
+    alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone()
+
+    # Shift so the last timestep is zero.
+    alphas_bar_sqrt -= alphas_bar_sqrt_T
+
+    # Scale so the first timestep is back to the old value.
+    alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T)
+
+    # Convert alphas_bar_sqrt to betas
+    alphas_bar = alphas_bar_sqrt**2  # Revert sqrt
+    alphas = alphas_bar[1:] / alphas_bar[:-1]  # Revert cumprod
+    alphas = torch.cat([alphas_bar[0:1], alphas])
+    betas = 1 - alphas
+
+    return betas
+
+class DPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
+    """
+    `DPMSolverMultistepScheduler` is a fast dedicated high-order solver for diffusion ODEs.
+
+    This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
+    methods the library implements for all schedulers such as loading and saving.
+
+    Args:
+        num_train_timesteps (`int`, defaults to 1000):
+            The number of diffusion steps to train the model.
+        beta_start (`float`, defaults to 0.0001):
+            The starting `beta` value of inference.
+        beta_end (`float`, defaults to 0.02):
+            The final `beta` value.
+        beta_schedule (`str`, defaults to `"linear"`):
+            The beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from
+            `linear`, `scaled_linear`, or `squaredcos_cap_v2`.
+        trained_betas (`np.ndarray`, *optional*):
+            Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`.
+        solver_order (`int`, defaults to 2):
+            The DPMSolver order which can be `1` or `2` or `3`. It is recommended to use `solver_order=2` for guided
+            sampling, and `solver_order=3` for unconditional sampling.
+        prediction_type (`str`, defaults to `epsilon`, *optional*):
+            Prediction type of the scheduler function; can be `epsilon` (predicts the noise of the diffusion process),
+            `sample` (directly predicts the noisy sample`) or `v_prediction` (see section 2.4 of [Imagen
+            Video](https://imagen.research.google/video/paper.pdf) paper).
+        thresholding (`bool`, defaults to `False`):
+            Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such
+            as Stable Diffusion.
+        dynamic_thresholding_ratio (`float`, defaults to 0.995):
+            The ratio for the dynamic thresholding method. Valid only when `thresholding=True`.
+        sample_max_value (`float`, defaults to 1.0):
+            The threshold value for dynamic thresholding. Valid only when `thresholding=True` and
+            `algorithm_type="dpmsolver++"`.
+        algorithm_type (`str`, defaults to `dpmsolver++`):
+            Algorithm type for the solver; can be `dpmsolver`, `dpmsolver++`, `sde-dpmsolver` or `sde-dpmsolver++`. The
+            `dpmsolver` type implements the algorithms in the [DPMSolver](https://huggingface.co/papers/2206.00927)
+            paper, and the `dpmsolver++` type implements the algorithms in the
+            [DPMSolver++](https://huggingface.co/papers/2211.01095) paper. It is recommended to use `dpmsolver++` or
+            `sde-dpmsolver++` with `solver_order=2` for guided sampling like in Stable Diffusion.
+        solver_type (`str`, defaults to `midpoint`):
+            Solver type for the second-order solver; can be `midpoint` or `heun`. The solver type slightly affects the
+            sample quality, especially for a small number of steps. It is recommended to use `midpoint` solvers.
+        lower_order_final (`bool`, defaults to `True`):
+            Whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. This can
+            stabilize the sampling of DPMSolver for steps < 15, especially for steps <= 10.
+        euler_at_final (`bool`, defaults to `False`):
+            Whether to use Euler's method in the final step. It is a trade-off between numerical stability and detail
+            richness. This can stabilize the sampling of the SDE variant of DPMSolver for small number of inference
+            steps, but sometimes may result in blurring.
+        use_karras_sigmas (`bool`, *optional*, defaults to `False`):
+            Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`,
+            the sigmas are determined according to a sequence of noise levels {σi}.
+        use_lu_lambdas (`bool`, *optional*, defaults to `False`):
+            Whether to use the uniform-logSNR for step sizes proposed by Lu's DPM-Solver in the noise schedule during
+            the sampling process. If `True`, the sigmas and time steps are determined according to a sequence of
+            `lambda(t)`.
+        final_sigmas_type (`str`, defaults to `"zero"`):
+            The final `sigma` value for the noise schedule during the sampling process. If `"sigma_min"`, the final
+            sigma is the same as the last sigma in the training schedule. If `zero`, the final sigma is set to 0.
+        lambda_min_clipped (`float`, defaults to `-inf`):
+            Clipping threshold for the minimum value of `lambda(t)` for numerical stability. This is critical for the
+            cosine (`squaredcos_cap_v2`) noise schedule.
+        variance_type (`str`, *optional*):
+            Set to "learned" or "learned_range" for diffusion models that predict variance. If set, the model's output
+            contains the predicted Gaussian variance.
+        timestep_spacing (`str`, defaults to `"linspace"`):
+            The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
+            Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
+        steps_offset (`int`, defaults to 0):
+            An offset added to the inference steps, as required by some model families.
+        rescale_betas_zero_snr (`bool`, defaults to `False`):
+            Whether to rescale the betas to have zero terminal SNR. This enables the model to generate very bright and
+            dark samples instead of limiting it to samples with medium brightness. Loosely related to
+            [`--offset_noise`](https://github.com/huggingface/diffusers/blob/74fd735eb073eb1d774b1ab4154a0876eb82f055/examples/dreambooth/train_dreambooth.py#L506).
+    """
+
+    _compatibles = [e.name for e in KarrasDiffusionSchedulers]
+    order = 1
+
+    @register_to_config
+    def __init__(
+        self,
+        num_train_timesteps: int = 1000,
+        beta_start: float = 0.0001,
+        beta_end: float = 0.02,
+        beta_schedule: str = "linear",
+        trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
+        solver_order: int = 2,
+        prediction_type: str = "epsilon",
+        thresholding: bool = False,
+        dynamic_thresholding_ratio: float = 0.995,
+        sample_max_value: float = 1.0,
+        algorithm_type: str = "dpmsolver++",
+        solver_type: str = "midpoint",
+        lower_order_final: bool = True,
+        euler_at_final: bool = False,
+        use_karras_sigmas: Optional[bool] = False,
+        use_lu_lambdas: Optional[bool] = False,
+        final_sigmas_type: Optional[str] = "zero",  # "zero", "sigma_min"
+        lambda_min_clipped: float = -float("inf"),
+        variance_type: Optional[str] = None,
+        timestep_spacing: str = "linspace",
+        steps_offset: int = 0,
+        rescale_betas_zero_snr: bool = False,
+    ):
+        if algorithm_type in ["dpmsolver", "sde-dpmsolver"]:
+            deprecation_message = f"algorithm_type {algorithm_type} is deprecated and will be removed in a future version. Choose from `dpmsolver++` or `sde-dpmsolver++` instead"
+            deprecate("algorithm_types dpmsolver and sde-dpmsolver", "1.0.0", deprecation_message)
+
+        if trained_betas is not None:
+            self.betas = torch.tensor(trained_betas, dtype=torch.float32)
+        elif beta_schedule == "linear":
+            self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32)
+        elif beta_schedule == "scaled_linear":
+            # this schedule is very specific to the latent diffusion model.
+            self.betas = torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2
+        elif beta_schedule == "squaredcos_cap_v2" or beta_schedule == "cosine":
+            # Glide cosine schedule
+            self.betas = betas_for_alpha_bar(num_train_timesteps, alpha_transform_type="cosine")
+        elif beta_schedule == "cauchy":
+            self.betas = betas_for_alpha_bar(num_train_timesteps, alpha_transform_type="cauchy")
+        elif beta_schedule == "laplace":
+            self.betas = betas_for_alpha_bar(num_train_timesteps, alpha_transform_type="laplace")
+        else:
+            raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}")
+
+        if rescale_betas_zero_snr:
+            self.betas = rescale_zero_terminal_snr(self.betas)
+
+        self.alphas = 1.0 - self.betas
+        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
+
+        if rescale_betas_zero_snr:
+            # Close to 0 without being 0 so first sigma is not inf
+            # FP16 smallest positive subnormal works well here
+            self.alphas_cumprod[-1] = 2**-24
+
+        # Currently we only support VP-type noise schedule
+        self.alpha_t = torch.sqrt(self.alphas_cumprod)
+        self.sigma_t = torch.sqrt(1 - self.alphas_cumprod)
+        self.lambda_t = torch.log(self.alpha_t) - torch.log(self.sigma_t)
+        self.sigmas = ((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5
+
+        # standard deviation of the initial noise distribution
+        self.init_noise_sigma = 1.0
+
+        # settings for DPM-Solver
+        if algorithm_type not in ["dpmsolver", "dpmsolver++", "sde-dpmsolver", "sde-dpmsolver++"]:
+            if algorithm_type == "deis":
+                self.register_to_config(algorithm_type="dpmsolver++")
+            else:
+                raise NotImplementedError(f"{algorithm_type} is not implemented for {self.__class__}")
+
+        if solver_type not in ["midpoint", "heun"]:
+            if solver_type in ["logrho", "bh1", "bh2"]:
+                self.register_to_config(solver_type="midpoint")
+            else:
+                raise NotImplementedError(f"{solver_type} is not implemented for {self.__class__}")
+
+        if algorithm_type not in ["dpmsolver++", "sde-dpmsolver++"] and final_sigmas_type == "zero":
+            raise ValueError(
+                f"`final_sigmas_type` {final_sigmas_type} is not supported for `algorithm_type` {algorithm_type}. Please choose `sigma_min` instead."
+            )
+
+        # settable values
+        self.num_inference_steps = None
+        timesteps = np.linspace(0, num_train_timesteps - 1, num_train_timesteps, dtype=np.float32)[::-1].copy()
+        self.timesteps = torch.from_numpy(timesteps)
+        self.model_outputs = [None] * solver_order
+        self.lower_order_nums = 0
+        self._step_index = None
+        self._begin_index = None
+        self.sigmas = self.sigmas.to("cpu")  # to avoid too much CPU/GPU communication
+
+    @property
+    def step_index(self):
+        """
+        The index counter for current timestep. It will increase 1 after each scheduler step.
+        """
+        return self._step_index
+
+    @property
+    def begin_index(self):
+        """
+        The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
+        """
+        return self._begin_index
+
+    def set_begin_index(self, begin_index: int = 0):
+        """
+        Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
+
+        Args:
+            begin_index (`int`):
+                The begin index for the scheduler.
+        """
+        self._begin_index = begin_index
+
+    def set_timesteps(
+        self,
+        num_inference_steps: int = None,
+        device: Union[str, torch.device] = None,
+        timesteps: Optional[List[int]] = None,
+    ):
+        """
+        Sets the discrete timesteps used for the diffusion chain (to be run before inference).
+
+        Args:
+            num_inference_steps (`int`):
+                The number of diffusion steps used when generating samples with a pre-trained model.
+            device (`str` or `torch.device`, *optional*):
+                The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+            timesteps (`List[int]`, *optional*):
+                Custom timesteps used to support arbitrary timesteps schedule. If `None`, timesteps will be generated
+                based on the `timestep_spacing` attribute. If `timesteps` is passed, `num_inference_steps` and `sigmas`
+                must be `None`, and `timestep_spacing` attribute will be ignored.
+        """
+        if num_inference_steps is None and timesteps is None:
+            raise ValueError("Must pass exactly one of `num_inference_steps` or `timesteps`.")
+        if num_inference_steps is not None and timesteps is not None:
+            raise ValueError("Can only pass one of `num_inference_steps` or `custom_timesteps`.")
+        if timesteps is not None and self.config.use_karras_sigmas:
+            raise ValueError("Cannot use `timesteps` with `config.use_karras_sigmas = True`")
+        if timesteps is not None and self.config.use_lu_lambdas:
+            raise ValueError("Cannot use `timesteps` with `config.use_lu_lambdas = True`")
+
+        if timesteps is not None:
+            timesteps = np.array(timesteps).astype(np.int64)
+        else:
+            # Clipping the minimum of all lambda(t) for numerical stability.
+            # This is critical for cosine (squaredcos_cap_v2) noise schedule.
+            clipped_idx = torch.searchsorted(torch.flip(self.lambda_t, [0]), self.config.lambda_min_clipped)
+            last_timestep = ((self.config.num_train_timesteps - clipped_idx).numpy()).item()
+
+            # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://arxiv.org/abs/2305.08891
+            if self.config.timestep_spacing == "linspace":
+                timesteps = (
+                    np.linspace(0, last_timestep - 1, num_inference_steps + 1)
+                    .round()[::-1][:-1]
+                    .copy()
+                    .astype(np.int64)
+                )
+            elif self.config.timestep_spacing == "leading":
+                step_ratio = last_timestep // (num_inference_steps + 1)
+                # creates integer timesteps by multiplying by ratio
+                # casting to int to avoid issues when num_inference_step is power of 3
+                timesteps = (
+                    (np.arange(0, num_inference_steps + 1) * step_ratio).round()[::-1][:-1].copy().astype(np.int64)
+                )
+                timesteps += self.config.steps_offset
+            elif self.config.timestep_spacing == "trailing":
+                step_ratio = self.config.num_train_timesteps / num_inference_steps
+                # creates integer timesteps by multiplying by ratio
+                # casting to int to avoid issues when num_inference_step is power of 3
+                timesteps = np.arange(last_timestep, 0, -step_ratio).round().copy().astype(np.int64)
+                timesteps -= 1
+            else:
+                raise ValueError(
+                    f"{self.config.timestep_spacing} is not supported. Please make sure to choose one of 'linspace', 'leading' or 'trailing'."
+                )
+
+        sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
+        log_sigmas = np.log(sigmas)
+
+        if self.config.use_karras_sigmas:
+            sigmas = np.flip(sigmas).copy()
+            sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=num_inference_steps)
+            timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]).round()
+        elif self.config.use_lu_lambdas:
+            lambdas = np.flip(log_sigmas.copy())
+            lambdas = self._convert_to_lu(in_lambdas=lambdas, num_inference_steps=num_inference_steps)
+            sigmas = np.exp(lambdas)
+            timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]).round()
+        else:
+            sigmas = np.interp(timesteps, np.arange(0, len(sigmas)), sigmas)
+
+        if self.config.final_sigmas_type == "sigma_min":
+            sigma_last = ((1 - self.alphas_cumprod[0]) / self.alphas_cumprod[0]) ** 0.5
+        elif self.config.final_sigmas_type == "zero":
+            sigma_last = 0
+        else:
+            raise ValueError(
+                f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}"
+            )
+
+        sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32)
+
+        self.sigmas = torch.from_numpy(sigmas)
+        self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=torch.int64)
+
+        self.num_inference_steps = len(timesteps)
+
+        self.model_outputs = [
+            None,
+        ] * self.config.solver_order
+        self.lower_order_nums = 0
+
+        # add an index counter for schedulers that allow duplicated timesteps
+        self._step_index = None
+        self._begin_index = None
+        self.sigmas = self.sigmas.to("cpu")  # to avoid too much CPU/GPU communication
+
+    # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample
+    def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor:
+        """
+        "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the
+        prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by
+        s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing
+        pixels from saturation at each step. We find that dynamic thresholding results in significantly better
+        photorealism as well as better image-text alignment, especially when using very large guidance weights."
+
+        https://arxiv.org/abs/2205.11487
+        """
+        dtype = sample.dtype
+        batch_size, channels, *remaining_dims = sample.shape
+
+        if dtype not in (torch.float32, torch.float64):
+            sample = sample.float()  # upcast for quantile calculation, and clamp not implemented for cpu half
+
+        # Flatten sample for doing quantile calculation along each image
+        sample = sample.reshape(batch_size, channels * np.prod(remaining_dims))
+
+        abs_sample = sample.abs()  # "a certain percentile absolute pixel value"
+
+        s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1)
+        s = torch.clamp(
+            s, min=1, max=self.config.sample_max_value
+        )  # When clamped to min=1, equivalent to standard clipping to [-1, 1]
+        s = s.unsqueeze(1)  # (batch_size, 1) because clamp will broadcast along dim=0
+        sample = torch.clamp(sample, -s, s) / s  # "we threshold xt0 to the range [-s, s] and then divide by s"
+
+        sample = sample.reshape(batch_size, channels, *remaining_dims)
+        sample = sample.to(dtype)
+
+        return sample
+
+    # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._sigma_to_t
+    def _sigma_to_t(self, sigma, log_sigmas):
+        # get log sigma
+        log_sigma = np.log(np.maximum(sigma, 1e-10))
+
+        # get distribution
+        dists = log_sigma - log_sigmas[:, np.newaxis]
+
+        # get sigmas range
+        low_idx = np.cumsum((dists >= 0), axis=0).argmax(axis=0).clip(max=log_sigmas.shape[0] - 2)
+        high_idx = low_idx + 1
+
+        low = log_sigmas[low_idx]
+        high = log_sigmas[high_idx]
+
+        # interpolate sigmas
+        w = (low - log_sigma) / (low - high)
+        w = np.clip(w, 0, 1)
+
+        # transform interpolation to time range
+        t = (1 - w) * low_idx + w * high_idx
+        t = t.reshape(sigma.shape)
+        return t
+
+    def _sigma_to_alpha_sigma_t(self, sigma):
+        alpha_t = 1 / ((sigma**2 + 1) ** 0.5)
+        sigma_t = sigma * alpha_t
+
+        return alpha_t, sigma_t
+
+    # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras
+    def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps) -> torch.Tensor:
+        """Constructs the noise schedule of Karras et al. (2022)."""
+
+        # Hack to make sure that other schedulers which copy this function don't break
+        # TODO: Add this logic to the other schedulers
+        if hasattr(self.config, "sigma_min"):
+            sigma_min = self.config.sigma_min
+        else:
+            sigma_min = None
+
+        if hasattr(self.config, "sigma_max"):
+            sigma_max = self.config.sigma_max
+        else:
+            sigma_max = None
+
+        sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item()
+        sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item()
+
+        rho = 7.0  # 7.0 is the value used in the paper
+        ramp = np.linspace(0, 1, num_inference_steps)
+        min_inv_rho = sigma_min ** (1 / rho)
+        max_inv_rho = sigma_max ** (1 / rho)
+        sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
+        return sigmas
+
+    def _convert_to_lu(self, in_lambdas: torch.Tensor, num_inference_steps) -> torch.Tensor:
+        """Constructs the noise schedule of Lu et al. (2022)."""
+
+        lambda_min: float = in_lambdas[-1].item()
+        lambda_max: float = in_lambdas[0].item()
+
+        rho = 1.0  # 1.0 is the value used in the paper
+        ramp = np.linspace(0, 1, num_inference_steps)
+        min_inv_rho = lambda_min ** (1 / rho)
+        max_inv_rho = lambda_max ** (1 / rho)
+        lambdas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
+        return lambdas
+
+    def convert_model_output(
+        self,
+        model_output: torch.Tensor,
+        *args,
+        sample: torch.Tensor = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        Convert the model output to the corresponding type the DPMSolver/DPMSolver++ algorithm needs. DPM-Solver is
+        designed to discretize an integral of the noise prediction model, and DPM-Solver++ is designed to discretize an
+        integral of the data prediction model.
+
+        
+
+        The algorithm and model type are decoupled. You can use either DPMSolver or DPMSolver++ for both noise
+        prediction and data prediction models.
+
+        
+
+        Args:
+            model_output (`torch.Tensor`):
+                The direct output from the learned diffusion model.
+            sample (`torch.Tensor`):
+                A current instance of a sample created by the diffusion process.
+
+        Returns:
+            `torch.Tensor`:
+                The converted model output.
+        """
+        timestep = args[0] if len(args) > 0 else kwargs.pop("timestep", None)
+        if sample is None:
+            if len(args) > 1:
+                sample = args[1]
+            else:
+                raise ValueError("missing `sample` as a required keyword argument")
+        if timestep is not None:
+            deprecate(
+                "timesteps",
+                "1.0.0",
+                "Passing `timesteps` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
+        # DPM-Solver++ needs to solve an integral of the data prediction model.
+        if self.config.algorithm_type in ["dpmsolver++", "sde-dpmsolver++"]:
+            if self.config.prediction_type == "epsilon":
+                # DPM-Solver and DPM-Solver++ only need the "mean" output.
+                if self.config.variance_type in ["learned", "learned_range"]:
+                    model_output = model_output[:, :3]
+                sigma = self.sigmas[self.step_index]
+                alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
+                x0_pred = (sample - sigma_t * model_output) / alpha_t
+            elif self.config.prediction_type == "sample":
+                x0_pred = model_output
+            elif self.config.prediction_type == "v_prediction":
+                sigma = self.sigmas[self.step_index]
+                alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
+                x0_pred = alpha_t * sample - sigma_t * model_output
+            else:
+                raise ValueError(
+                    f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or"
+                    " `v_prediction` for the DPMSolverMultistepScheduler."
+                )
+
+            if self.config.thresholding:
+                x0_pred = self._threshold_sample(x0_pred)
+
+            return x0_pred
+
+        # DPM-Solver needs to solve an integral of the noise prediction model.
+        elif self.config.algorithm_type in ["dpmsolver", "sde-dpmsolver"]:
+            if self.config.prediction_type == "epsilon":
+                # DPM-Solver and DPM-Solver++ only need the "mean" output.
+                if self.config.variance_type in ["learned", "learned_range"]:
+                    epsilon = model_output[:, :3]
+                else:
+                    epsilon = model_output
+            elif self.config.prediction_type == "sample":
+                sigma = self.sigmas[self.step_index]
+                alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
+                epsilon = (sample - alpha_t * model_output) / sigma_t
+            elif self.config.prediction_type == "v_prediction":
+                sigma = self.sigmas[self.step_index]
+                alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
+                epsilon = alpha_t * model_output + sigma_t * sample
+            else:
+                raise ValueError(
+                    f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or"
+                    " `v_prediction` for the DPMSolverMultistepScheduler."
+                )
+
+            if self.config.thresholding:
+                sigma = self.sigmas[self.step_index]
+                alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
+                x0_pred = (sample - sigma_t * epsilon) / alpha_t
+                x0_pred = self._threshold_sample(x0_pred)
+                epsilon = (sample - alpha_t * x0_pred) / sigma_t
+
+            return epsilon
+
+    def dpm_solver_first_order_update(
+        self,
+        model_output: torch.Tensor,
+        *args,
+        sample: torch.Tensor = None,
+        noise: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        One step for the first-order DPMSolver (equivalent to DDIM).
+
+        Args:
+            model_output (`torch.Tensor`):
+                The direct output from the learned diffusion model.
+            sample (`torch.Tensor`):
+                A current instance of a sample created by the diffusion process.
+
+        Returns:
+            `torch.Tensor`:
+                The sample tensor at the previous timestep.
+        """
+        timestep = args[0] if len(args) > 0 else kwargs.pop("timestep", None)
+        prev_timestep = args[1] if len(args) > 1 else kwargs.pop("prev_timestep", None)
+        if sample is None:
+            if len(args) > 2:
+                sample = args[2]
+            else:
+                raise ValueError(" missing `sample` as a required keyword argument")
+        if timestep is not None:
+            deprecate(
+                "timesteps",
+                "1.0.0",
+                "Passing `timesteps` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
+        if prev_timestep is not None:
+            deprecate(
+                "prev_timestep",
+                "1.0.0",
+                "Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
+        sigma_t, sigma_s = self.sigmas[self.step_index + 1], self.sigmas[self.step_index]
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
+        alpha_s, sigma_s = self._sigma_to_alpha_sigma_t(sigma_s)
+        lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
+        lambda_s = torch.log(alpha_s) - torch.log(sigma_s)
+
+        h = lambda_t - lambda_s
+        if self.config.algorithm_type == "dpmsolver++":
+            x_t = (sigma_t / sigma_s) * sample - (alpha_t * (torch.exp(-h) - 1.0)) * model_output
+        elif self.config.algorithm_type == "dpmsolver":
+            x_t = (alpha_t / alpha_s) * sample - (sigma_t * (torch.exp(h) - 1.0)) * model_output
+        elif self.config.algorithm_type == "sde-dpmsolver++":
+            assert noise is not None
+            x_t = (
+                (sigma_t / sigma_s * torch.exp(-h)) * sample
+                + (alpha_t * (1 - torch.exp(-2.0 * h))) * model_output
+                + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise
+            )
+        elif self.config.algorithm_type == "sde-dpmsolver":
+            assert noise is not None
+            x_t = (
+                (alpha_t / alpha_s) * sample
+                - 2.0 * (sigma_t * (torch.exp(h) - 1.0)) * model_output
+                + sigma_t * torch.sqrt(torch.exp(2 * h) - 1.0) * noise
+            )
+        return x_t
+
+    def multistep_dpm_solver_second_order_update(
+        self,
+        model_output_list: List[torch.Tensor],
+        *args,
+        sample: torch.Tensor = None,
+        noise: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        One step for the second-order multistep DPMSolver.
+
+        Args:
+            model_output_list (`List[torch.Tensor]`):
+                The direct outputs from learned diffusion model at current and latter timesteps.
+            sample (`torch.Tensor`):
+                A current instance of a sample created by the diffusion process.
+
+        Returns:
+            `torch.Tensor`:
+                The sample tensor at the previous timestep.
+        """
+        timestep_list = args[0] if len(args) > 0 else kwargs.pop("timestep_list", None)
+        prev_timestep = args[1] if len(args) > 1 else kwargs.pop("prev_timestep", None)
+        if sample is None:
+            if len(args) > 2:
+                sample = args[2]
+            else:
+                raise ValueError(" missing `sample` as a required keyword argument")
+        if timestep_list is not None:
+            deprecate(
+                "timestep_list",
+                "1.0.0",
+                "Passing `timestep_list` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
+        if prev_timestep is not None:
+            deprecate(
+                "prev_timestep",
+                "1.0.0",
+                "Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
+        sigma_t, sigma_s0, sigma_s1 = (
+            self.sigmas[self.step_index + 1],
+            self.sigmas[self.step_index],
+            self.sigmas[self.step_index - 1],
+        )
+
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
+        alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
+        alpha_s1, sigma_s1 = self._sigma_to_alpha_sigma_t(sigma_s1)
+
+        lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
+        lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
+        lambda_s1 = torch.log(alpha_s1) - torch.log(sigma_s1)
+
+        m0, m1 = model_output_list[-1], model_output_list[-2]
+
+        h, h_0 = lambda_t - lambda_s0, lambda_s0 - lambda_s1
+        r0 = h_0 / h
+        D0, D1 = m0, (1.0 / r0) * (m0 - m1)
+        if self.config.algorithm_type == "dpmsolver++":
+            # See https://arxiv.org/abs/2211.01095 for detailed derivations
+            if self.config.solver_type == "midpoint":
+                x_t = (
+                    (sigma_t / sigma_s0) * sample
+                    - (alpha_t * (torch.exp(-h) - 1.0)) * D0
+                    - 0.5 * (alpha_t * (torch.exp(-h) - 1.0)) * D1
+                )
+            elif self.config.solver_type == "heun":
+                x_t = (
+                    (sigma_t / sigma_s0) * sample
+                    - (alpha_t * (torch.exp(-h) - 1.0)) * D0
+                    + (alpha_t * ((torch.exp(-h) - 1.0) / h + 1.0)) * D1
+                )
+        elif self.config.algorithm_type == "dpmsolver":
+            # See https://arxiv.org/abs/2206.00927 for detailed derivations
+            if self.config.solver_type == "midpoint":
+                x_t = (
+                    (alpha_t / alpha_s0) * sample
+                    - (sigma_t * (torch.exp(h) - 1.0)) * D0
+                    - 0.5 * (sigma_t * (torch.exp(h) - 1.0)) * D1
+                )
+            elif self.config.solver_type == "heun":
+                x_t = (
+                    (alpha_t / alpha_s0) * sample
+                    - (sigma_t * (torch.exp(h) - 1.0)) * D0
+                    - (sigma_t * ((torch.exp(h) - 1.0) / h - 1.0)) * D1
+                )
+        elif self.config.algorithm_type == "sde-dpmsolver++":
+            assert noise is not None
+            if self.config.solver_type == "midpoint":
+                x_t = (
+                    (sigma_t / sigma_s0 * torch.exp(-h)) * sample
+                    + (alpha_t * (1 - torch.exp(-2.0 * h))) * D0
+                    + 0.5 * (alpha_t * (1 - torch.exp(-2.0 * h))) * D1
+                    + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise
+                )
+            elif self.config.solver_type == "heun":
+                x_t = (
+                    (sigma_t / sigma_s0 * torch.exp(-h)) * sample
+                    + (alpha_t * (1 - torch.exp(-2.0 * h))) * D0
+                    + (alpha_t * ((1.0 - torch.exp(-2.0 * h)) / (-2.0 * h) + 1.0)) * D1
+                    + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise
+                )
+        elif self.config.algorithm_type == "sde-dpmsolver":
+            assert noise is not None
+            if self.config.solver_type == "midpoint":
+                x_t = (
+                    (alpha_t / alpha_s0) * sample
+                    - 2.0 * (sigma_t * (torch.exp(h) - 1.0)) * D0
+                    - (sigma_t * (torch.exp(h) - 1.0)) * D1
+                    + sigma_t * torch.sqrt(torch.exp(2 * h) - 1.0) * noise
+                )
+            elif self.config.solver_type == "heun":
+                x_t = (
+                    (alpha_t / alpha_s0) * sample
+                    - 2.0 * (sigma_t * (torch.exp(h) - 1.0)) * D0
+                    - 2.0 * (sigma_t * ((torch.exp(h) - 1.0) / h - 1.0)) * D1
+                    + sigma_t * torch.sqrt(torch.exp(2 * h) - 1.0) * noise
+                )
+        return x_t
+
+    def multistep_dpm_solver_third_order_update(
+        self,
+        model_output_list: List[torch.Tensor],
+        *args,
+        sample: torch.Tensor = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        One step for the third-order multistep DPMSolver.
+
+        Args:
+            model_output_list (`List[torch.Tensor]`):
+                The direct outputs from learned diffusion model at current and latter timesteps.
+            sample (`torch.Tensor`):
+                A current instance of a sample created by diffusion process.
+
+        Returns:
+            `torch.Tensor`:
+                The sample tensor at the previous timestep.
+        """
+
+        timestep_list = args[0] if len(args) > 0 else kwargs.pop("timestep_list", None)
+        prev_timestep = args[1] if len(args) > 1 else kwargs.pop("prev_timestep", None)
+        if sample is None:
+            if len(args) > 2:
+                sample = args[2]
+            else:
+                raise ValueError(" missing`sample` as a required keyword argument")
+        if timestep_list is not None:
+            deprecate(
+                "timestep_list",
+                "1.0.0",
+                "Passing `timestep_list` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
+        if prev_timestep is not None:
+            deprecate(
+                "prev_timestep",
+                "1.0.0",
+                "Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
+        sigma_t, sigma_s0, sigma_s1, sigma_s2 = (
+            self.sigmas[self.step_index + 1],
+            self.sigmas[self.step_index],
+            self.sigmas[self.step_index - 1],
+            self.sigmas[self.step_index - 2],
+        )
+
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
+        alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
+        alpha_s1, sigma_s1 = self._sigma_to_alpha_sigma_t(sigma_s1)
+        alpha_s2, sigma_s2 = self._sigma_to_alpha_sigma_t(sigma_s2)
+
+        lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
+        lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
+        lambda_s1 = torch.log(alpha_s1) - torch.log(sigma_s1)
+        lambda_s2 = torch.log(alpha_s2) - torch.log(sigma_s2)
+
+        m0, m1, m2 = model_output_list[-1], model_output_list[-2], model_output_list[-3]
+
+        h, h_0, h_1 = lambda_t - lambda_s0, lambda_s0 - lambda_s1, lambda_s1 - lambda_s2
+        r0, r1 = h_0 / h, h_1 / h
+        D0 = m0
+        D1_0, D1_1 = (1.0 / r0) * (m0 - m1), (1.0 / r1) * (m1 - m2)
+        D1 = D1_0 + (r0 / (r0 + r1)) * (D1_0 - D1_1)
+        D2 = (1.0 / (r0 + r1)) * (D1_0 - D1_1)
+        if self.config.algorithm_type == "dpmsolver++":
+            # See https://arxiv.org/abs/2206.00927 for detailed derivations
+            x_t = (
+                (sigma_t / sigma_s0) * sample
+                - (alpha_t * (torch.exp(-h) - 1.0)) * D0
+                + (alpha_t * ((torch.exp(-h) - 1.0) / h + 1.0)) * D1
+                - (alpha_t * ((torch.exp(-h) - 1.0 + h) / h**2 - 0.5)) * D2
+            )
+        elif self.config.algorithm_type == "dpmsolver":
+            # See https://arxiv.org/abs/2206.00927 for detailed derivations
+            x_t = (
+                (alpha_t / alpha_s0) * sample
+                - (sigma_t * (torch.exp(h) - 1.0)) * D0
+                - (sigma_t * ((torch.exp(h) - 1.0) / h - 1.0)) * D1
+                - (sigma_t * ((torch.exp(h) - 1.0 - h) / h**2 - 0.5)) * D2
+            )
+        return x_t
+
+    def index_for_timestep(self, timestep, schedule_timesteps=None):
+        if schedule_timesteps is None:
+            schedule_timesteps = self.timesteps
+
+        index_candidates = (schedule_timesteps == timestep).nonzero()
+
+        if len(index_candidates) == 0:
+            step_index = len(self.timesteps) - 1
+        # The sigma index that is taken for the **very** first `step`
+        # is always the second index (or the last index if there is only 1)
+        # This way we can ensure we don't accidentally skip a sigma in
+        # case we start in the middle of the denoising schedule (e.g. for image-to-image)
+        elif len(index_candidates) > 1:
+            step_index = index_candidates[1].item()
+        else:
+            step_index = index_candidates[0].item()
+
+        return step_index
+
+    def _init_step_index(self, timestep):
+        """
+        Initialize the step_index counter for the scheduler.
+        """
+
+        if self.begin_index is None:
+            if isinstance(timestep, torch.Tensor):
+                timestep = timestep.to(self.timesteps.device)
+            self._step_index = self.index_for_timestep(timestep)
+        else:
+            self._step_index = self._begin_index
+
+    def step(
+        self,
+        model_output: torch.Tensor,
+        timestep: int,
+        sample: torch.Tensor,
+        generator=None,
+        variance_noise: Optional[torch.Tensor] = None,
+        return_dict: bool = True,
+    ) -> Union[SchedulerOutput, Tuple]:
+        """
+        Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with
+        the multistep DPMSolver.
+
+        Args:
+            model_output (`torch.Tensor`):
+                The direct output from learned diffusion model.
+            timestep (`int`):
+                The current discrete timestep in the diffusion chain.
+            sample (`torch.Tensor`):
+                A current instance of a sample created by the diffusion process.
+            generator (`torch.Generator`, *optional*):
+                A random number generator.
+            variance_noise (`torch.Tensor`):
+                Alternative to generating noise with `generator` by directly providing the noise for the variance
+                itself. Useful for methods such as [`LEdits++`].
+            return_dict (`bool`):
+                Whether or not to return a [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`.
+
+        Returns:
+            [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`:
+                If return_dict is `True`, [`~schedulers.scheduling_utils.SchedulerOutput`] is returned; otherwise, a
+                tuple is returned where the first element is the sample tensor.
+
+        """
+        if self.num_inference_steps is None:
+            raise ValueError(
+                "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"
+            )
+
+        if self.step_index is None:
+            self._init_step_index(timestep)
+
+        # Improve numerical stability for small number of steps
+        lower_order_final = (self.step_index == len(self.timesteps) - 1) and (
+            self.config.euler_at_final
+            or (self.config.lower_order_final and len(self.timesteps) < 15)
+            or self.config.final_sigmas_type == "zero"
+        )
+        lower_order_second = (
+            (self.step_index == len(self.timesteps) - 2) and self.config.lower_order_final and len(self.timesteps) < 15
+        )
+
+        model_output = self.convert_model_output(model_output, sample=sample)
+        for i in range(self.config.solver_order - 1):
+            self.model_outputs[i] = self.model_outputs[i + 1]
+        self.model_outputs[-1] = model_output
+
+        # Upcast to avoid precision issues when computing prev_sample
+        sample = sample.to(torch.float32)
+        if self.config.algorithm_type in ["sde-dpmsolver", "sde-dpmsolver++"] and variance_noise is None:
+            noise = randn_tensor(
+                model_output.shape, generator=generator, device=model_output.device, dtype=torch.float32
+            )
+        elif self.config.algorithm_type in ["sde-dpmsolver", "sde-dpmsolver++"]:
+            noise = variance_noise.to(device=model_output.device, dtype=torch.float32)
+        else:
+            noise = None
+
+        if self.config.solver_order == 1 or self.lower_order_nums < 1 or lower_order_final:
+            prev_sample = self.dpm_solver_first_order_update(model_output, sample=sample, noise=noise)
+        elif self.config.solver_order == 2 or self.lower_order_nums < 2 or lower_order_second:
+            prev_sample = self.multistep_dpm_solver_second_order_update(self.model_outputs, sample=sample, noise=noise)
+        else:
+            prev_sample = self.multistep_dpm_solver_third_order_update(self.model_outputs, sample=sample)
+
+        if self.lower_order_nums < self.config.solver_order:
+            self.lower_order_nums += 1
+
+        # Cast sample back to expected dtype
+        prev_sample = prev_sample.to(model_output.dtype)
+
+        # upon completion increase step index by one
+        self._step_index += 1
+
+        if not return_dict:
+            return (prev_sample,)
+
+        return SchedulerOutput(prev_sample=prev_sample)
+
+    def add_noise(
+        self,
+        original_samples: torch.Tensor,
+        noise: torch.Tensor,
+        timesteps: torch.IntTensor,
+    ) -> torch.Tensor:
+        # Make sure sigmas and timesteps have the same device and dtype as original_samples
+        # alpha_t = self.alpha_t.to(device=original_samples.device, dtype=original_samples.dtype)
+        # sigma_t = self.sigma_t.to(device=original_samples.device, dtype=original_samples.dtype)
+        alpha_t = self.alpha_t.to(original_samples.device).to(original_samples.dtype)
+        sigma_t = self.sigma_t.to(original_samples.device).to(original_samples.dtype)
+        timesteps = timesteps.to(original_samples.device)
+        alpha_t = alpha_t[timesteps].flatten()
+        while len(alpha_t.shape) < len(original_samples.shape):
+            alpha_t = alpha_t.unsqueeze(-1)
+
+        sigma_t = sigma_t[timesteps].flatten()
+        while len(sigma_t.shape) < len(original_samples.shape):
+            sigma_t = sigma_t.unsqueeze(-1)
+        noisy_samples = alpha_t * original_samples + sigma_t * noise
+        return noisy_samples
+    
+    def get_velocity(self, original_samples: torch.Tensor, noise: torch.Tensor, timesteps: torch.IntTensor) -> torch.Tensor:
+        # alpha_t = self.alpha_t.to(device=original_samples.device, dtype=original_samples.dtype)
+        # sigma_t = self.sigma_t.to(device=original_samples.device, dtype=original_samples.dtype)
+        alpha_t = self.alpha_t.to(original_samples.device).to(original_samples.dtype)
+        sigma_t = self.sigma_t.to(original_samples.device).to(original_samples.dtype)
+
+        timesteps = timesteps.to(original_samples.device)
+        alpha_t = alpha_t[timesteps].flatten()
+        while len(alpha_t.shape) < len(original_samples.shape):
+            alpha_t = alpha_t.unsqueeze(-1)
+
+        sigma_t = sigma_t[timesteps].flatten()
+        while len(sigma_t.shape) < len(original_samples.shape):
+            sigma_t = sigma_t.unsqueeze(-1)
+
+        velocity = alpha_t * noise - sigma_t * original_samples
+        return velocity
+
+    def __len__(self):
+        return self.config.num_train_timesteps
\ No newline at end of file
diff --git a/vibevoice/schedule/timestep_sampler.py b/vibevoice/schedule/timestep_sampler.py
new file mode 100644
index 0000000..177b66f
--- /dev/null
+++ b/vibevoice/schedule/timestep_sampler.py
@@ -0,0 +1,19 @@
+import math
+import torch
+
+
+class UniformSampler:
+    def __init__(self, timesteps = 1000):
+        self.timesteps = timesteps
+    def sample(self, batch_size, device):
+        return torch.randint(0, self.timesteps, (batch_size,), device=device)
+    
+class LogitNormalSampler:
+    def __init__(self, timesteps = 1000, m = 0, s = 1):
+        self.timesteps = timesteps
+        timesteps = torch.linspace(0, 1, timesteps)
+        logit = torch.log(timesteps / (1 - timesteps))
+        self.prob = torch.exp(-0.5 * (logit - m) ** 2 / s ** 2) / (s * math.sqrt(2 * math.pi))
+    def sample(self, batch_size, device):
+        return torch.multinomial(self.prob, batch_size, replacement=True).to(device)
+    
\ No newline at end of file
diff --git a/vibevoice/scripts/__init__.py b/vibevoice/scripts/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/vllm_plugin/__init__.py b/vllm_plugin/__init__.py
new file mode 100644
index 0000000..989acb4
--- /dev/null
+++ b/vllm_plugin/__init__.py
@@ -0,0 +1,62 @@
+"""VibeVoice vLLM Plugin - Registers VibeVoice model for vLLM inference.
+
+This plugin enables VibeVoice ASR models to be loaded and served through vLLM.
+It registers the model architecture, configuration, tokenizer, and processor
+with their respective registries.
+
+The plugin is automatically loaded by vLLM via the 'vllm.general_plugins'
+entry point defined in pyproject.toml.
+"""
+
+from vllm.model_executor.models import ModelRegistry
+from transformers import AutoConfig, AutoTokenizer, Qwen2Tokenizer, AutoProcessor, Qwen2AudioProcessor
+
+from vibevoice.modular.configuration_vibevoice import VibeVoiceConfig
+from vibevoice.modular.modular_vibevoice_text_tokenizer import VibeVoiceASRTextTokenizerFast
+
+from .model import VibeVoiceForCausalLM
+
+
+def register_vibevoice():
+    """Register VibeVoice model with vLLM and transformers.
+    
+    This function is called automatically by vLLM through the entry point
+    mechanism. It registers:
+    - VibeVoiceConfig with AutoConfig
+    - VibeVoiceASRTextTokenizerFast with AutoTokenizer (for ASR)
+    - Qwen2AudioProcessor with AutoProcessor
+    - VibeVoiceForCausalLM with vLLM ModelRegistry
+    """
+    # Register the configuration class with transformers
+    AutoConfig.register("vibevoice", VibeVoiceConfig)
+
+    # Register the tokenizer with transformers.
+    # IMPORTANT (ASR): Align with the PyTorch ASR path.
+    # VibeVoiceASRTextTokenizerFast maps:
+    #   speech_start_id -> <|object_ref_start|>
+    #   speech_pad_id   -> <|box_start|>
+    #   speech_end_id   -> <|object_ref_end|>
+    # This significantly affects ASR quality even when requests succeed.
+    try:
+        AutoTokenizer.register(
+            VibeVoiceConfig,
+            slow_tokenizer_class=Qwen2Tokenizer,
+            fast_tokenizer_class=VibeVoiceASRTextTokenizerFast,
+        )
+    except Exception:
+        pass  # May already be registered
+
+    # Register the processor with transformers
+    try:
+        AutoProcessor.register(VibeVoiceConfig, processor_class=Qwen2AudioProcessor)
+    except Exception:
+        pass  # May already be registered
+
+    # Register the model class with the architecture name "VibeVoice"
+    # This name must match the "architectures" list in config.json
+    ModelRegistry.register_model("VibeVoice", VibeVoiceForCausalLM)
+    ModelRegistry.register_model("VibeVoiceForASRTraining", VibeVoiceForCausalLM)
+
+
+# Note: This function is called via vllm.general_plugins entry point
+# defined in pyproject.toml, ensuring it runs in all vLLM processes
diff --git a/vllm_plugin/inputs.py b/vllm_plugin/inputs.py
new file mode 100644
index 0000000..841e28f
--- /dev/null
+++ b/vllm_plugin/inputs.py
@@ -0,0 +1,101 @@
+"""Audio input mapper for vLLM multimodal pipeline.
+
+This module handles audio data loading and preprocessing for VibeVoice ASR inference.
+It converts various audio input formats (path, bytes, numpy array) into tensors
+that can be processed by the VibeVoice model.
+"""
+import os
+import logging
+import torch
+import numpy as np
+from typing import Union, List
+from vllm.multimodal.inputs import MultiModalInputs
+from vibevoice.processor.audio_utils import load_audio_use_ffmpeg, load_audio_bytes_use_ffmpeg, AudioNormalizer
+
+logger = logging.getLogger(__name__)
+
+# Maximum audio duration in seconds. Default 3660s (61 minutes) matches the
+# model's designed capacity. Override via environment variable to guard against
+# OOM on GPUs with less VRAM. See https://github.com/microsoft/VibeVoice/issues/210
+_MAX_AUDIO_DURATION = float(os.environ.get("VIBEVOICE_MAX_AUDIO_DURATION", "3660"))
+
+
+def load_audio(audio_path: str, target_sr: int = 24000) -> np.ndarray:
+    """Load and normalize audio from file path.
+    
+    Args:
+        audio_path: Path to audio file
+        target_sr: Target sample rate (default 24kHz for VibeVoice)
+        
+    Returns:
+        Normalized audio waveform as numpy array
+    """
+    # Load with FFmpeg (handles various formats)
+    audio, sr = load_audio_use_ffmpeg(audio_path, resample=True, target_sr=target_sr)
+    
+    # Normalize audio
+    normalizer = AudioNormalizer()
+    audio = normalizer(audio)
+    
+    return audio
+
+
+def vibevoice_audio_input_mapper(ctx, data: Union[str, bytes, np.ndarray, List[str]]) -> MultiModalInputs:
+    """Map audio input data to vLLM MultiModalInputs format.
+    
+    This function is registered as the input mapper for VibeVoice audio processing.
+    It handles multiple input formats and converts them to normalized tensors.
+    
+    Args:
+        ctx: vLLM context (unused)
+        data: Audio data in one of these formats:
+            - str: Path to audio file
+            - bytes: Raw audio bytes (any format FFmpeg supports)
+            - np.ndarray: Pre-loaded audio waveform
+            - List[str]: List of audio paths (only first is used)
+            
+    Returns:
+        MultiModalInputs containing:
+            - audio: Audio tensor (float32)
+            - audio_length: Length of audio in samples
+    """
+    # Handle list input (take first item)
+    if isinstance(data, list):
+        data = data[0]
+
+    audio_waveform = None
+    
+    if isinstance(data, str):
+        # Load from file path
+        audio_waveform = load_audio(data)
+        
+    elif isinstance(data, bytes):
+        # Decode bytes directly via ffmpeg stdin pipe to avoid temp-file IO
+        audio_waveform, _sr = load_audio_bytes_use_ffmpeg(data, resample=True, target_sr=24000)
+        normalizer = AudioNormalizer()
+        audio_waveform = normalizer(audio_waveform)
+                
+    elif isinstance(data, np.ndarray):
+        # Already loaded numpy array
+        audio_waveform = data
+    else:
+        raise ValueError(f"Unsupported audio data type: {type(data)}")
+
+    # Validate audio duration before tensor conversion to catch OOM early
+    duration_sec = len(audio_waveform) / 24000
+    if duration_sec > _MAX_AUDIO_DURATION:
+        raise ValueError(
+            f"Audio duration ({duration_sec:.1f}s) exceeds the configured "
+            f"limit ({_MAX_AUDIO_DURATION:.0f}s). Set the "
+            f"VIBEVOICE_MAX_AUDIO_DURATION environment variable to adjust "
+            f"this limit, or use shorter audio."
+        )
+
+    # Convert to tensor
+    audio_tensor = torch.from_numpy(audio_waveform).float()
+    audio_length = audio_tensor.shape[0]
+    
+    return MultiModalInputs({
+        "audio": audio_tensor,
+        "audio_length": audio_length
+    })
diff --git a/vllm_plugin/model.py b/vllm_plugin/model.py
new file mode 100644
index 0000000..bd9060f
--- /dev/null
+++ b/vllm_plugin/model.py
@@ -0,0 +1,1251 @@
+"""
+VibeVoice vLLM Plugin Model - Native Multimodal Integration
+
+This module implements the VibeVoice ASR model with full vLLM multimodal registry
+integration for speech-to-text inference.
+"""
+
+from typing import List, Optional, Tuple, Union, Dict, Any, Iterable, Mapping, Sequence
+import os
+import torch
+import torch.nn as nn
+import numpy as np
+import base64
+
+
+# ============================================================================
+# Audio Loading: FFmpeg-based AudioMediaIO
+# ============================================================================
+# VibeVoice uses FFmpeg for audio decoding to ensure consistent behavior
+# across different audio formats (MP3, WAV, FLAC, etc.).
+
+
+from vibevoice.processor.audio_utils import load_audio_use_ffmpeg, load_audio_bytes_use_ffmpeg, AudioNormalizer
+
+
+def _ffmpeg_load_bytes(data: bytes) -> tuple[np.ndarray, int]:
+    """Load audio bytes using FFmpeg via stdin-pipe decoding.
+    
+    Returns:
+        Tuple of (audio_waveform, sample_rate). Sample rate is always 24000.
+    """
+    audio, sr = load_audio_bytes_use_ffmpeg(data, resample=True, target_sr=24000)
+    normalizer = AudioNormalizer()
+    audio = normalizer(audio)
+    return audio, sr
+
+def _ffmpeg_load_file(filepath) -> tuple[np.ndarray, int]:
+    """Load audio file using FFmpeg.
+    
+    Returns:
+        Tuple of (audio_waveform, sample_rate). Sample rate is always 24000.
+    """
+    audio, sr = load_audio_use_ffmpeg(str(filepath), resample=True, target_sr=24000)
+    normalizer = AudioNormalizer()
+    audio = normalizer(audio)
+    return audio, sr
+
+# Register FFmpeg-based audio loader
+try:
+    # Try new location (vLLM >= 0.6.x)
+    from vllm.multimodal.media.audio import AudioMediaIO as _OriginalAudioMediaIO
+except ImportError:
+    # Fall back to old location (vLLM < 0.6.x)
+    import vllm.multimodal.audio as _vllm_audio_module
+    _OriginalAudioMediaIO = _vllm_audio_module.AudioMediaIO
+
+class _PatchedAudioMediaIO(_OriginalAudioMediaIO):
+    """AudioMediaIO implementation using FFmpeg for audio decoding."""
+    
+    def load_bytes(self, data: bytes) -> tuple[np.ndarray, int]:
+        return _ffmpeg_load_bytes(data)
+    
+    def load_base64(self, media_type: str, data: str) -> tuple[np.ndarray, int]:
+        return _ffmpeg_load_bytes(base64.b64decode(data))
+    
+    def load_file(self, filepath) -> tuple[np.ndarray, int]:
+        return _ffmpeg_load_file(filepath)
+
+# Replace globally
+try:
+    # For new vLLM versions
+    import vllm.multimodal.media.audio as _vllm_audio_module
+    _vllm_audio_module.AudioMediaIO = _PatchedAudioMediaIO
+except ImportError:
+    # For old vLLM versions
+    import vllm.multimodal.audio as _vllm_audio_module
+    _vllm_audio_module.AudioMediaIO = _PatchedAudioMediaIO
+
+# Also patch in utils module where it's imported
+try:
+    import vllm.multimodal.utils as _vllm_utils_module
+    _vllm_utils_module.AudioMediaIO = _PatchedAudioMediaIO
+except (ImportError, AttributeError):
+    # AudioMediaIO might not be imported in utils in newer versions
+    pass
+
+# ============================================================================
+
+from transformers import BatchFeature
+from transformers.models.whisper import WhisperFeatureExtractor
+from vllm.config import VllmConfig
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.multimodal.parse import MultiModalDataParser
+from vllm.sequence import IntermediateTensors
+from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP, MultiModalEmbeddings
+from vllm.model_executor.models.utils import (
+    init_vllm_registered_model,
+    maybe_prefix,
+    AutoWeightsLoader,
+    WeightsMapper,
+)
+from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems
+from vllm.multimodal.processing import (
+    BaseMultiModalProcessor,
+    BaseProcessingInfo,
+    PromptReplacement,
+    PromptUpdate,
+    PromptUpdateDetails,
+)
+try:
+    # Try new location (vLLM >= 0.6.x)
+    from vllm.multimodal.processing import BaseDummyInputsBuilder, ProcessorInputs
+except ImportError:
+    # Fall back to old location (vLLM < 0.6.x)
+    try:
+        from vllm.multimodal.profiling import BaseDummyInputsBuilder, ProcessorInputs
+    except ImportError:
+        # If neither location works, try individual imports
+        from vllm.multimodal.processing.dummy_inputs import BaseDummyInputsBuilder
+        from vllm.multimodal.processing.inputs import ProcessorInputs
+
+# Import VibeVoice components
+from vibevoice.modular.modular_vibevoice_tokenizer import (
+    VibeVoiceAcousticTokenizerModel,
+    VibeVoiceSemanticTokenizerModel,
+    VibeVoiceTokenizerStreamingCache,
+    VibeVoiceTokenizerEncoderOutput,
+)
+from vibevoice.modular.configuration_vibevoice import (
+    VibeVoiceAcousticTokenizerConfig,
+    VibeVoiceSemanticTokenizerConfig,
+)
+
+
+class SpeechConnector(nn.Module):
+    """Projects speech features to language model hidden dimension.
+    
+    Architecture: fc1 -> RMSNorm -> fc2 (no activation function)
+    """
+    def __init__(self, input_dim: int, output_dim: int):
+        super().__init__()
+        self.fc1 = nn.Linear(input_dim, output_dim)
+        self.norm = LlamaRMSNorm(output_dim, eps=1e-6)
+        self.fc2 = nn.Linear(output_dim, output_dim)
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.fc1(x)
+        x = self.norm(x)
+        x = self.fc2(x)
+        return x
+
+
+class LlamaRMSNorm(nn.Module):
+    """RMSNorm layer used in SpeechConnector."""
+    def __init__(self, hidden_size: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+
+class VibeVoiceAudioEncoder(nn.Module):
+    """
+    VibeVoice Audio Encoder module.
+    
+    Encapsulates Acoustic/Semantic VAE Tokenizers and projection Connectors.
+    Converts raw audio waveforms into embeddings compatible with the language model.
+    
+    Features:
+        - Streaming support for long audio (>60s by default)
+        - Configurable dtype for numerical precision
+        - Supports both sampling and deterministic (mean) modes
+    """
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        
+        import sys 
+
+        def get_cfg(obj, key, default=None):
+            if isinstance(obj, dict):
+                return obj.get(key, default)
+            return getattr(obj, key, default)
+
+        self.acoustic_vae_dim = get_cfg(config, "acoustic_vae_dim", 64)
+        self.semantic_vae_dim = get_cfg(config, "semantic_vae_dim", 128)
+
+        decoder_config = get_cfg(config, "decoder_config")
+        text_config = get_cfg(config, "text_config")
+        
+        target_hidden_size = None
+        
+        if decoder_config is not None:
+            target_hidden_size = get_cfg(decoder_config, "hidden_size")
+
+        if target_hidden_size is None and text_config is not None:
+            target_hidden_size = get_cfg(text_config, "hidden_size")
+
+        if target_hidden_size is None:
+            target_hidden_size = get_cfg(config, "hidden_size")
+
+        if target_hidden_size is None:
+            print("[VibeVoice] WARN: Could not find hidden_size in config! Defaulting to 3584 (7B).", file=sys.stderr)
+            self.hidden_size = 3584
+        else:
+            self.hidden_size = target_hidden_size
+
+        ac_cfg = get_cfg(config, "acoustic_tokenizer_config")
+        sc_cfg = get_cfg(config, "semantic_tokenizer_config")
+        
+        if ac_cfg is None or sc_cfg is None:
+            raise ValueError("Missing acoustic/semantic tokenizer config in model config")
+
+        # Handle both dict and already-constructed config objects
+        if isinstance(ac_cfg, VibeVoiceAcousticTokenizerConfig):
+            acoustic_config = ac_cfg
+        elif isinstance(ac_cfg, dict):
+            acoustic_config = VibeVoiceAcousticTokenizerConfig(**ac_cfg)
+        else:
+            raise TypeError(f"acoustic_tokenizer_config has unexpected type: {type(ac_cfg)}")
+        
+        if isinstance(sc_cfg, VibeVoiceSemanticTokenizerConfig):
+            semantic_config = sc_cfg
+        elif isinstance(sc_cfg, dict):
+            semantic_config = VibeVoiceSemanticTokenizerConfig(**sc_cfg)
+        else:
+            raise TypeError(f"semantic_tokenizer_config has unexpected type: {type(sc_cfg)}")
+        
+        # Tokenizers use float32 for numerical precision
+        self.acoustic_tokenizer = VibeVoiceAcousticTokenizerModel(acoustic_config)
+        self.semantic_tokenizer = VibeVoiceSemanticTokenizerModel(semantic_config)
+        
+        # Get audio encoder dtype from config (defaults to float32 for precision)
+        root_torch_dtype = get_cfg(config, "torch_dtype", None)
+        if root_torch_dtype is not None:
+            if isinstance(root_torch_dtype, str):
+                self._audio_encoder_dtype = getattr(torch, root_torch_dtype)
+            else:
+                self._audio_encoder_dtype = root_torch_dtype
+        else:
+            self._audio_encoder_dtype = torch.float32
+        
+        self.acoustic_connector = SpeechConnector(self.acoustic_vae_dim, self.hidden_size)
+        self.semantic_connector = SpeechConnector(self.semantic_vae_dim, self.hidden_size)
+        
+        self.compress_ratio = get_cfg(config, "speech_tok_compress_ratio", 3200)
+        
+        # Streaming controls
+        self.sample_rate = get_cfg(config, "target_sample_rate", 24000)
+
+        # Default to True (per requirement): segment + cache inside one forward call.
+        self.enable_streaming = get_cfg(config, "enable_streaming", True)
+        self.streaming_segment_duration = get_cfg(config, "streaming_segment_duration", 60.0)
+
+        # Control whether to use sample() or .mean for acoustic tokens
+        # Default: use sample() for training-consistent behavior
+        # Set VIBEVOICE_USE_MEAN=1 for deterministic output
+        use_mean_env = os.getenv("VIBEVOICE_USE_MEAN", "").strip().lower()
+        self.use_sample = use_mean_env not in ("1", "true", "yes")
+        
+        # Language model dtype (set by VibeVoiceForCausalLM.__init__)
+        # This is the dtype that audio embeddings will be converted to before
+        # being passed to the language model. Defaults to bfloat16.
+        self._lm_dtype: torch.dtype = torch.bfloat16
+
+    def _ensure_audio_encoder_dtype(self):
+        """Ensure all audio encoder components use the correct dtype from config.
+        
+        vLLM may convert weights to a different dtype (e.g., bfloat16) during loading.
+        This method converts audio encoder components back to the config-specified dtype
+        (typically float32) for numerical precision during audio encoding.
+        """
+        import sys
+        target_dtype = self._audio_encoder_dtype
+        
+        # Check and convert acoustic_tokenizer
+        try:
+            acoustic_dtype = next(self.acoustic_tokenizer.parameters()).dtype
+            if acoustic_dtype != target_dtype:
+                self.acoustic_tokenizer = self.acoustic_tokenizer.to(dtype=target_dtype)
+                print(f"[VibeVoice] Converted acoustic_tokenizer to {target_dtype} (was {acoustic_dtype})", file=sys.stderr)
+        except StopIteration:
+            pass
+        
+        # Check and convert semantic_tokenizer
+        try:
+            semantic_dtype = next(self.semantic_tokenizer.parameters()).dtype
+            if semantic_dtype != target_dtype:
+                self.semantic_tokenizer = self.semantic_tokenizer.to(dtype=target_dtype)
+                print(f"[VibeVoice] Converted semantic_tokenizer to {target_dtype} (was {semantic_dtype})", file=sys.stderr)
+        except StopIteration:
+            pass
+        
+        # Check and convert acoustic_connector
+        try:
+            ac_conn_dtype = next(self.acoustic_connector.parameters()).dtype
+            if ac_conn_dtype != target_dtype:
+                self.acoustic_connector = self.acoustic_connector.to(dtype=target_dtype)
+                print(f"[VibeVoice] Converted acoustic_connector to {target_dtype} (was {ac_conn_dtype})", file=sys.stderr)
+        except StopIteration:
+            pass
+        
+        # Check and convert semantic_connector
+        try:
+            sc_conn_dtype = next(self.semantic_connector.parameters()).dtype
+            if sc_conn_dtype != target_dtype:
+                self.semantic_connector = self.semantic_connector.to(dtype=target_dtype)
+                print(f"[VibeVoice] Converted semantic_connector to {target_dtype} (was {sc_conn_dtype})", file=sys.stderr)
+        except StopIteration:
+            pass
+
+    def forward(
+        self,
+        audio: torch.Tensor,
+        *,
+        use_streaming: bool = True,
+        segment_duration_s: Optional[float] = None,
+        use_sample: Optional[bool] = None,
+    ) -> torch.Tensor:
+        """Encode audio with optional streaming for long clips.
+        
+        Args:
+            audio: Input audio tensor [B, T] or [T]
+            use_streaming: Whether to enable segmented encoding for long audio
+            segment_duration_s: Segment length in seconds (defaults to 60s)
+            use_sample: If True, use sampling for acoustic tokens; if False, use mean
+                       Defaults to self.use_sample (controlled by VIBEVOICE_USE_MEAN env var)
+        
+        Returns:
+            Audio embeddings tensor compatible with the language model
+        """
+        # Ensure audio encoder components use correct dtype
+        self._ensure_audio_encoder_dtype()
+        
+        # Audio input should match the audio encoder dtype
+        audio = audio.to(dtype=self._audio_encoder_dtype)
+        
+        if audio.ndim == 1:
+            audio = audio.unsqueeze(0)
+
+        # Resolve streaming options
+        segment_duration = segment_duration_s or self.streaming_segment_duration
+        sample_rate = self.sample_rate
+        total_samples = audio.shape[-1]
+        segment_samples = int(segment_duration * sample_rate)
+        
+        use_streaming = use_streaming and self.enable_streaming and total_samples > segment_samples
+        
+        # Resolve use_sample flag
+        if use_sample is None:
+            use_sample = self.use_sample
+
+        # Keep encoding in inference mode to avoid autograd build-up
+        with torch.no_grad():
+            if not use_streaming:
+                acoustic_input = audio.unsqueeze(1)
+                acoustic_out = self.acoustic_tokenizer.encode(acoustic_input)
+                # Use sample() or .mean based on use_sample flag
+                if use_sample:
+                    acoustic_tokens = acoustic_out.sample(
+                        dist_type=self.acoustic_tokenizer.std_dist_type
+                    )[0]
+                else:
+                    acoustic_tokens = acoustic_out.mean
+                
+                # Connector is now float32, no conversion needed
+                acoustic_embeds = self.acoustic_connector(acoustic_tokens)
+
+                semantic_out = self.semantic_tokenizer.encode(acoustic_input)
+                # Semantic always uses .mean for consistency
+                semantic_tokens = semantic_out.mean
+                # Connector is now float32, no conversion needed
+                semantic_embeds = self.semantic_connector(semantic_tokens)
+            else:
+                # ==========================================
+                # Streaming path (Retained for future use)
+                # ==========================================
+                acoustic_cache = VibeVoiceTokenizerStreamingCache()
+                semantic_cache = VibeVoiceTokenizerStreamingCache()
+                acoustic_mean_segments = []
+                semantic_mean_segments = []
+                batch_size = audio.shape[0]
+                sample_indices = torch.arange(batch_size, device=audio.device)
+
+                def _iter_segments(total_length: int, segment_length: int):
+                    for start in range(0, total_length, segment_length):
+                        end = min(start + segment_length, total_length)
+                        if end > start:
+                            yield start, end
+
+                segments = list(_iter_segments(total_samples, segment_samples))
+                num_segments = len(segments)
+                for seg_idx, (start, end) in enumerate(segments):
+                    chunk = audio[:, start:end].contiguous()
+                    if chunk.numel() == 0:
+                        continue
+
+                    # Check if this is the final segment
+                    is_final = (seg_idx == num_segments - 1)
+
+                    # --- Acoustic Encode ---
+                    acoustic_enc_out = self.acoustic_tokenizer.encode(
+                        chunk.unsqueeze(1),
+                        cache=acoustic_cache,
+                        sample_indices=sample_indices,
+                        use_cache=True,
+                        is_final_chunk=is_final,
+                    )
+                    acoustic_mean_segments.append(acoustic_enc_out.mean)
+
+                    # --- Semantic Encode ---
+                    semantic_enc_out = self.semantic_tokenizer.encode(
+                        chunk.unsqueeze(1),
+                        cache=semantic_cache,
+                        sample_indices=sample_indices,
+                        use_cache=True,
+                        is_final_chunk=is_final,
+                    )
+                    semantic_mean_segments.append(semantic_enc_out.mean)
+
+                # Concatenate sequence outputs (Acoustic)
+                if len(acoustic_mean_segments) == 0:
+                    acoustic_mean_full = torch.zeros(
+                        (batch_size, 0, self.acoustic_vae_dim), 
+                        device=audio.device, 
+                        dtype=self._audio_encoder_dtype  # Use config dtype
+                    )
+                else:
+                    acoustic_mean_full = torch.cat(acoustic_mean_segments, dim=1).contiguous()
+                
+                # Get acoustic tokens based on use_sample flag
+                acoustic_enc_full = VibeVoiceTokenizerEncoderOutput(
+                    mean=acoustic_mean_full,
+                    std=self.acoustic_tokenizer.fix_std,
+                )
+                if use_sample:
+                    acoustic_tokens = acoustic_enc_full.sample(
+                        dist_type=self.acoustic_tokenizer.std_dist_type
+                    )[0]
+                else:
+                    acoustic_tokens = acoustic_enc_full.mean
+                # Connector uses same dtype as tokenizer
+                acoustic_embeds = self.acoustic_connector(acoustic_tokens)
+
+                # Concatenate sequence outputs (Semantic)
+                if len(semantic_mean_segments) == 0:
+                    semantic_tokens = torch.zeros(
+                        (batch_size, 0, self.semantic_vae_dim), 
+                        device=audio.device, 
+                        dtype=self._audio_encoder_dtype  # Use config dtype
+                    )
+                else:
+                    semantic_tokens = torch.cat(semantic_mean_segments, dim=1).contiguous()
+                # Connector uses same dtype as tokenizer
+                semantic_embeds = self.semantic_connector(semantic_tokens)
+
+        # Combine acoustic and semantic embeddings
+        combined_embeds = acoustic_embeds + semantic_embeds
+        
+        # Convert to language model dtype for compatibility
+        # Audio encoder uses config.torch_dtype (typically float32) for numerical precision,
+        # but LM expects the dtype specified by vLLM's --dtype flag (e.g., bfloat16, float16)
+        combined_embeds = combined_embeds.to(dtype=self._lm_dtype)
+        
+        return combined_embeds
+
+# ============================================================================
+# vLLM Multimodal Processing Infrastructure
+# ============================================================================
+
+class VibeVoiceProcessingInfo(BaseProcessingInfo):
+    """Processing info for VibeVoice multimodal model."""
+    
+    def get_hf_config(self):
+        return self.ctx.get_hf_config()
+
+    def get_feature_extractor(self, **kwargs) -> WhisperFeatureExtractor:
+        """
+        Get a WhisperFeatureExtractor for vLLM profiling compatibility.
+        
+        IMPORTANT: This is NOT used in actual inference!
+        VibeVoice uses its own acoustic/semantic VAE tokenizers operating on
+        raw 24kHz waveforms, NOT Whisper mel spectrograms.
+        
+        This feature extractor exists only to satisfy vLLM's multimodal
+        profiling infrastructure which may query audio parameters like
+        sampling_rate and chunk_length for memory estimation.
+        """
+        # Read config from preprocessor_config.json if available
+        import json
+        import os
+        model_path = self.ctx.model_config.model
+        preprocessor_path = os.path.join(model_path, "preprocessor_config.json")
+        
+        # Default values: keep a coherent (sr, hop) pair.
+        # VibeVoice runs at 24kHz in this repo (see demo/asr_transcribe_file.py).
+        config = {
+            "sampling_rate": 24000,
+            "feature_size": 128,
+            # 10ms hop at 24kHz
+            "hop_length": 240,
+            "chunk_length": 30,
+            "n_fft": 400,
+            "padding_value": 0.0,
+        }
+        
+        # Try to load from config file
+        if os.path.exists(preprocessor_path):
+            try:
+                with open(preprocessor_path, "r") as f:
+                    file_config = json.load(f)
+                    config.update({k: file_config[k] for k in config.keys() if k in file_config})
+            except Exception:
+                pass  # Use defaults
+        
+        return WhisperFeatureExtractor(
+            feature_size=config["feature_size"],
+            sampling_rate=config["sampling_rate"],
+            hop_length=config["hop_length"],
+            chunk_length=config["chunk_length"],
+            n_fft=config["n_fft"],
+            padding_value=config["padding_value"],
+        )
+
+    def get_audio_token_info(self) -> dict:
+        """
+        Get audio special tokens and their IDs.
+        
+        Returns dict with:
+            audio_token, audio_bos_token, audio_eos_token,
+            audio_token_id, audio_bos_id, audio_eos_id
+        """
+        tokenizer = self.get_tokenizer()
+        vocab = tokenizer.get_vocab()
+        
+        # VibeVoice special tokens
+        tokens = {
+            "audio_token": "<|AUDIO|>",
+            "audio_bos_token": "<|audio_bos|>",
+            "audio_eos_token": "<|audio_eos|>",
+        }
+        
+        # Get IDs
+        tokens["audio_token_id"] = vocab.get(tokens["audio_token"])
+        tokens["audio_bos_id"] = vocab.get(tokens["audio_bos_token"])
+        tokens["audio_eos_id"] = vocab.get(tokens["audio_eos_token"])
+        
+        return tokens
+
+    def get_supported_mm_limits(self) -> Mapping[str, int | None]:
+        return {"audio": 1}
+
+    def get_mm_max_tokens_per_item(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+    ) -> Mapping[str, int]:
+        """Return the maximum number of audio tokens per item.
+
+        This tells vLLM's scheduler the upper bound so that
+        ``encoder_compute_budget`` is large enough for any audio length
+        the model can handle, preventing the silent scheduling deadlock
+        described in docs/max_num_batched_tokens_issue.md.
+
+        Formula: audio_tokens = ceil(audio_samples / compress_ratio) + 3
+        where +3 accounts for speech_start, speech_end, and newline tokens.
+        The max audio samples is bounded by seq_len (the model's context
+        window cannot hold more tokens than that).
+        """
+        hf_config = self.get_hf_config()
+
+        def _cfg(key: str, default):
+            if isinstance(hf_config, dict):
+                return hf_config.get(key, default)
+            return getattr(hf_config, key, default)
+
+        compress_ratio = int(_cfg("speech_tok_compress_ratio", 3200))
+        sample_rate = int(_cfg("target_sample_rate", 24000))
+
+        # Upper bound: 61-minute audio at 24 kHz
+        max_audio_samples = 61 * 60 * sample_rate  # 88,464,000
+        max_audio_tokens = int(np.ceil(max_audio_samples / compress_ratio)) + 3
+
+        # Cannot exceed the model's context window
+        max_audio_tokens = min(max_audio_tokens, seq_len)
+
+        return {"audio": max_audio_tokens}
+
+
+class VibeVoiceDummyInputsBuilder(BaseDummyInputsBuilder[VibeVoiceProcessingInfo]):
+    """
+    Build dummy inputs for multimodal profiling.
+    
+    vLLM uses dummy inputs to:
+    1. Measure peak GPU activation memory → determine KV cache capacity
+    2. Warm up CUDA graphs
+    
+    The dummy audio length must be consistent with ``get_mm_max_tokens_per_item``
+    so that the memory estimate covers the worst-case (longest audio) scenario.
+    """
+    
+    def _get_max_audio_samples(self, seq_len: int) -> int:
+        """Compute maximum audio samples consistent with ``get_mm_max_tokens_per_item``.
+        
+        Uses the same formula: max_tokens = min(ceil(61min * sr / ratio) + 3, seq_len),
+        then converts back to samples.
+        """
+        hf_config = self.info.get_hf_config()
+
+        def _cfg(key: str, default):
+            if isinstance(hf_config, dict):
+                return hf_config.get(key, default)
+            return getattr(hf_config, key, default)
+
+        compress_ratio = int(_cfg("speech_tok_compress_ratio", 3200))
+        sample_rate = int(_cfg("target_sample_rate", 24000))
+
+        # Upper bound: 61-minute audio at 24 kHz
+        max_hour_samples = 61 * 60 * sample_rate  # 88,464,000
+        max_tokens_from_audio = int(np.ceil(max_hour_samples / compress_ratio)) + 3
+        # Cannot exceed model context window
+        max_tokens = min(max_tokens_from_audio, seq_len)
+        # Convert tokens back to samples
+        return max_tokens * compress_ratio
+
+    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
+        num_audios = mm_counts.get("audio", 0)
+        if num_audios <= 0:
+            return ""
+        
+        token_info = self.info.get_audio_token_info()
+        audio_token = token_info["audio_token"]
+        return audio_token * num_audios
+
+    def get_dummy_mm_data(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+        mm_options: Mapping[str, Any] | None = None,
+    ) -> Dict[str, Any]:
+        """Generate dummy audio data for profiling.
+        
+        The audio length is derived from ``seq_len`` so that profiling
+        accurately measures memory for the longest audio the model can handle.
+        Supports ``AudioDummyOptions.length`` override for faster startup.
+        """
+        num_audios = mm_counts.get("audio", 0)
+        max_audio_len = self._get_max_audio_samples(seq_len)
+
+        audio_overrides = mm_options.get("audio") if mm_options else None
+
+        return {
+            "audio": self._get_dummy_audios(
+                length=max_audio_len,
+                num_audios=num_audios,
+                overrides=audio_overrides,
+            )
+        }
+
+    def get_dummy_processor_inputs(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+        mm_options: Mapping[str, Any] | None = None,
+    ) -> ProcessorInputs:
+        """Build ProcessorInputs for dummy profiling."""
+        return ProcessorInputs(
+            prompt=self.get_dummy_text(mm_counts),
+            mm_data=self.get_dummy_mm_data(seq_len, mm_counts, mm_options),
+        )
+
+
+def _vibevoice_field_config(hf_inputs: Mapping[str, torch.Tensor]):
+    """Map HF processor output keys to audio modality.
+    
+    Returns a config dict that tells vLLM how to batch multimodal data.
+    """
+    # Always define the config for all fields we use
+    # Even if the field isn't in hf_inputs, vLLM needs to know how to batch it
+    config = {
+        # These are our custom fields for VibeVoice
+        "raw_audio": MultiModalFieldConfig.batched("audio"),
+        "raw_audio_lengths": MultiModalFieldConfig.batched("audio"),
+        "salt": MultiModalFieldConfig.batched("audio"),
+    }
+    
+    # Add optional Whisper features if present
+    if "input_features" in hf_inputs:
+        config["input_features"] = MultiModalFieldConfig.batched("audio")
+    if "feature_attention_mask" in hf_inputs:
+        config["feature_attention_mask"] = MultiModalFieldConfig.batched("audio")
+    
+    return config
+
+
+class VibeVoiceMultiModalProcessor(BaseMultiModalProcessor[VibeVoiceProcessingInfo]):
+    """
+    Multimodal processor for VibeVoice.
+    
+    Handles the conversion of raw audio inputs to model-ready features,
+    and manages the prompt token replacement for audio placeholders.
+    """
+    
+    def _get_data_parser(self) -> MultiModalDataParser:
+        """Create a data parser with the correct target sample rate (24kHz)."""
+        # VibeVoice requires 24kHz, not 16kHz (Whisper default)
+        target_sr = 24000
+        return MultiModalDataParser(target_sr=target_sr)
+    
+    def _call_hf_processor(
+        self,
+        prompt: str,
+        mm_data: Mapping[str, object],
+        mm_kwargs: Mapping[str, object],
+        tok_kwargs: Mapping[str, object],
+    ) -> BatchFeature:
+        """
+        Process prompt and audio for vLLM multimodal pipeline.
+        
+        We intentionally do NOT run a HF processor that would pre-expand
+        `<|AUDIO|>` inside this method. Instead we:
+        1) Tokenize the prompt as-is (so `<|AUDIO|>` stays a single token)
+        2) Store raw audio tensors for `embed_multimodal` to encode later
+        3) Let vLLM call `_get_prompt_updates` to expand the single `<|AUDIO|>`
+           into the full ASR format: [speech_start] + N*[speech_pad] + [speech_end] + [\\n]
+        """
+        # Handle the case where 'audios' key is used (transformers deprecation)
+        mm_data = dict(mm_data)  # Make a mutable copy
+        audios = mm_data.pop("audios", None)
+        if audios is not None and "audio" not in mm_data:
+            mm_data["audio"] = audios
+        
+        # Text-only input handling
+        if not mm_data.get("audio"):
+            prompt_ids = self.info.get_tokenizer().encode(prompt)
+            prompt_ids = self._apply_hf_processor_tokens_only(prompt_ids)
+            return BatchFeature(dict(input_ids=[prompt_ids]), tensor_type="pt")
+        
+        # Get raw audio data
+        raw_audio_list = mm_data.get("audio")
+        if isinstance(raw_audio_list, np.ndarray):
+            raw_audio_list = [raw_audio_list]
+        elif not isinstance(raw_audio_list, list):
+            raw_audio_list = list(raw_audio_list)
+        
+        # Tokenize prompt directly to preserve <|AUDIO|> as a single token
+        # vLLM will expand it via _get_prompt_updates
+        tokenizer = self.info.get_tokenizer()
+        prompt_ids = tokenizer.encode(prompt, add_special_tokens=False)
+        prompt_ids = self._apply_hf_processor_tokens_only(prompt_ids)
+        
+        # Create result with input_ids
+        result = BatchFeature(dict(input_ids=[prompt_ids]), tensor_type="pt")
+        
+        # Add raw audio tensors for VibeVoice encoder
+        # Stack into a single tensor for vLLM's batched field config
+        max_len = max(len(a) for a in raw_audio_list)
+        raw_audio_tensors = []
+        audio_lengths = []
+        for audio in raw_audio_list:
+            audio_len = len(audio)
+            audio_lengths.append(audio_len)
+            if audio_len < max_len:
+                audio = np.pad(audio, (0, max_len - audio_len), mode='constant')
+            raw_audio_tensors.append(torch.from_numpy(audio).float())
+        
+        # Stack into [num_audios, max_len] tensor
+        stacked_audio = torch.stack(raw_audio_tensors, dim=0)  # Shape: [num_audios, max_len]
+        result["raw_audio"] = stacked_audio
+        # Convert lengths to tensor as well
+        result["raw_audio_lengths"] = torch.tensor(audio_lengths, dtype=torch.long)
+        
+        # Add a random salt to ensure unique hash and bypass cache
+        import uuid
+        # Use a random integer for salt
+        salt_val = hash(str(uuid.uuid4())) % 100000
+        result["salt"] = torch.tensor([salt_val], dtype=torch.long).expand(len(raw_audio_list))
+        
+        return result
+
+    def _hf_processor_applies_updates(
+        self,
+        prompt_text: str,
+        mm_items,
+        hf_processor_mm_kwargs: Mapping[str, object],
+        tokenization_kwargs: Mapping[str, object],
+    ) -> bool:
+        """Return whether the HF processor applies prompt updates.
+        
+        Returns False because we handle token expansion via _get_prompt_updates.
+        """
+        return False
+
+    def _get_mm_fields_config(
+        self,
+        hf_inputs: BatchFeature,
+        hf_processor_mm_kwargs: Mapping[str, object],
+    ) -> Mapping[str, MultiModalFieldConfig]:
+        """Configure which HF output fields map to which modality."""
+        return _vibevoice_field_config(hf_inputs)
+
+    def _get_prompt_updates(
+        self,
+        mm_items,
+        hf_processor_mm_kwargs: Mapping[str, object],
+        out_mm_kwargs: MultiModalKwargsItems,
+    ) -> Sequence[PromptUpdate]:
+        """
+        Define how to replace the audio placeholder in the prompt.
+
+        vLLM's OpenAI multimodal parsing inserts the model placeholder string
+        from `get_placeholder_str` (here: `<|AUDIO|>`) into the conversation.
+        We expand that single token into N repeated `<|AUDIO|>` tokens, where N
+        is derived from waveform length and `speech_tok_compress_ratio`.
+        """
+        token_info = self.info.get_audio_token_info()
+        audio_token = token_info["audio_token"]
+        audio_token_id = token_info["audio_token_id"]
+        audio_bos_id = token_info.get("audio_bos_id")
+        audio_eos_id = token_info.get("audio_eos_id")
+
+        tokenizer = self.info.get_tokenizer()
+        vocab = tokenizer.get_vocab()
+
+        def _tok_id(name: str) -> int | None:
+            return vocab.get(name)
+
+        # Look up speech token IDs from vocabulary
+        # These tokens mark the start/end of audio embeddings in the prompt
+        speech_start_id = (
+            _tok_id("<|object_ref_start|>")
+            or getattr(tokenizer, "speech_start_id", None)
+            or _tok_id("<|speech_start|>")
+        )
+        speech_end_id = (
+            _tok_id("<|object_ref_end|>")
+            or getattr(tokenizer, "speech_end_id", None)
+            or _tok_id("<|speech_end|>")
+        )
+        speech_pad_id = (
+            _tok_id("<|box_start|>")
+            or getattr(tokenizer, "speech_pad_id", None)
+            or _tok_id("<|speech_pad|>")
+        )
+
+        if audio_token_id is None:
+            return []
+        
+        # Get raw audio lengths (in samples, after resampling to 24kHz) from our stored data
+        out_mm_data = out_mm_kwargs.get_data()
+        raw_audio_lengths = out_mm_data.get("raw_audio_lengths", [])
+
+        # Fetch defaults from model config when available (falls back to 3200)
+        hf_config = self.info.get_hf_config()
+        if isinstance(hf_config, dict):
+            compress_ratio = int(hf_config.get("speech_tok_compress_ratio", 3200))
+        else:
+            compress_ratio = int(getattr(hf_config, "speech_tok_compress_ratio", 3200))
+
+        def _to_int_len(x) -> int:
+            if x is None:
+                return 0
+            if isinstance(x, torch.Tensor):
+                # Accept 0-dim or 1-dim scalar-like tensors
+                if x.numel() == 1:
+                    return int(x.item())
+                # If a full tensor is passed accidentally, fall back to its length
+                return int(x.shape[0])
+            return int(x)
+        
+        def get_replacement(item_idx: int):
+            if raw_audio_lengths and item_idx < len(raw_audio_lengths):
+                audio_len = _to_int_len(raw_audio_lengths[item_idx])
+                num_features = max(1, int(np.ceil(audio_len / compress_ratio)))
+            else:
+                # Fallback: estimate for 30 second audio at 24kHz
+                num_features = int(np.ceil(30 * 24000 / compress_ratio))
+            
+            if num_features == 0:
+                raise ValueError(
+                    f"Audio at index {item_idx} is too short to be represented"
+                )
+            
+            # Build replacement token sequence:
+            #   <|speech_start|> + N * <|speech_pad|> + <|speech_end|> + \n
+            # The newline is important for correct prompt structure.
+            newline_id = 198  # '\n' token
+            if speech_start_id is not None and speech_pad_id is not None and speech_end_id is not None:
+                embed_id = int(speech_pad_id)
+                replacement_ids = [int(speech_start_id)] + [embed_id] * num_features + [int(speech_end_id), newline_id]
+            # Fallback: add audio BOS/EOS boundaries around repeated <|AUDIO|>.
+            elif audio_bos_id is not None and audio_eos_id is not None:
+                embed_id = int(audio_token_id)
+                replacement_ids = [int(audio_bos_id)] + [embed_id] * num_features + [int(audio_eos_id)]
+            else:
+                embed_id = int(audio_token_id)
+                replacement_ids = [embed_id] * num_features
+
+            return PromptUpdateDetails.select_token_id(
+                replacement_ids,
+                embed_token_id=int(embed_id),
+            )
+        
+        return [
+            PromptReplacement(
+                modality="audio",
+                # Keep string placeholder matching for maximum vLLM compatibility.
+                target=audio_token,
+                replacement=get_replacement,
+            )
+        ]
+
+
+# ============================================================================
+# Main Model Class
+# ============================================================================
+
+@MULTIMODAL_REGISTRY.register_processor(
+    VibeVoiceMultiModalProcessor,
+    info=VibeVoiceProcessingInfo,
+    dummy_inputs=VibeVoiceDummyInputsBuilder,
+)
+class VibeVoiceForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
+    """
+    VibeVoice ASR model with native vLLM multimodal integration.
+    
+    This model combines VibeVoice acoustic/semantic tokenizers for audio encoding
+    with a causal language model for text generation.
+    """
+    
+    @classmethod
+    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
+        """Return the placeholder string format for a given modality.
+        
+        Returns "<|AUDIO|>" which vLLM inserts into the conversation prompt.
+        This single placeholder is later expanded by `_get_prompt_updates` into:
+            [speech_start_id] + [speech_pad_id] * N + [speech_end_id] + [newline_id]
+        where N = ceil(audio_samples / compress_ratio).
+        """
+        if modality.startswith("audio"):
+            return "<|AUDIO|>"
+        raise ValueError(f"Unsupported modality: {modality}")
+    
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        config = vllm_config.model_config.hf_config
+        self.config = config
+        
+        self.audio_encoder = VibeVoiceAudioEncoder(config)
+        
+        # Pass decoder_config to the language model initialization
+        decoder_config = getattr(config, "decoder_config", config)
+        self.language_model = init_vllm_registered_model(
+            vllm_config=vllm_config,
+            hf_config=decoder_config,
+            prefix=maybe_prefix(prefix, "language_model"),
+            architectures=["Qwen2ForCausalLM"],
+        )
+        
+        self.make_empty_intermediate_tensors = (
+            self.language_model.make_empty_intermediate_tensors
+        )
+
+        # Set the language model dtype for audio encoder output conversion
+        # This should match vLLM's --dtype flag (e.g., bfloat16, float16, float32)
+        # Audio encoder internal computation stays in fp32 for numerical precision,
+        # but output is converted to LM dtype for compatibility
+        lm_dtype = vllm_config.model_config.dtype
+        if lm_dtype is not None:
+            self.audio_encoder._lm_dtype = lm_dtype
+        
+        # Ensure audio encoder uses correct dtype (typically fp32 for precision)
+        try:
+            self.audio_encoder._ensure_audio_encoder_dtype()
+        except Exception:
+            pass
+
+    def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
+        return self.language_model.compute_logits(hidden_states)
+
+
+    def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
+        """
+        Extract audio embeddings using VibeVoice's acoustic/semantic tokenizers.
+        
+        Called by vLLM to get audio embeddings that replace audio placeholder tokens.
+        
+        Returns:
+            Tuple of embedding tensors, one per audio input.
+        """
+        # Get raw audio data (stored by our processor)
+        raw_audio = kwargs.get("raw_audio")
+        raw_audio_lengths = kwargs.get("raw_audio_lengths")
+        
+        # Handle no audio input - this happens during memory profiling
+        if raw_audio is None:
+            return []
+        
+        # Handle empty audio list
+        if isinstance(raw_audio, (list, tuple)) and len(raw_audio) == 0:
+            return []
+        
+        # Flatten raw_audio_lengths if it's nested
+        def flatten_lengths(lengths):
+            """Flatten nested lists/tensors of lengths to a single list."""
+            if lengths is None:
+                return []
+            
+            result = []
+            if isinstance(lengths, torch.Tensor):
+                lengths = lengths.tolist()
+            
+            if isinstance(lengths, (list, tuple)):
+                for item in lengths:
+                    if isinstance(item, (list, tuple)):
+                        result.extend(flatten_lengths(item))
+                    elif isinstance(item, torch.Tensor):
+                        if item.dim() == 0:
+                            result.append(item.item())
+                        else:
+                            result.extend(item.tolist())
+                    else:
+                        result.append(item)
+            else:
+                result.append(lengths)
+            return result
+        
+        raw_audio_lengths = flatten_lengths(raw_audio_lengths)
+
+        # Streaming controls. Enabled by default; can be overridden per-call.
+        use_streaming_flag = bool(
+            kwargs.get(
+                "use_streaming",
+                getattr(self.audio_encoder, "enable_streaming", True),
+            )
+        )
+        streaming_segment_duration = kwargs.get(
+            "streaming_segment_duration",
+            getattr(self.audio_encoder, "streaming_segment_duration", 60.0),
+        )
+        
+        # Process each audio through the VibeVoice encoder
+        embeddings = []
+        
+        # Get model device for tensor placement.
+        try:
+            device = next(self.audio_encoder.parameters()).device
+        except StopIteration:
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        
+        # Handle both stacked tensor and list of tensors
+        # vLLM batches as: [batch_size, 1, seq_len] or [batch_size, seq_len]
+        if isinstance(raw_audio, torch.Tensor):
+            if raw_audio.dim() == 3:
+                num_audios = raw_audio.shape[0]
+                audio_list = [raw_audio[i].squeeze(0) for i in range(num_audios)]
+            elif raw_audio.dim() == 2:
+                num_audios = raw_audio.shape[0]
+                audio_list = [raw_audio[i] for i in range(num_audios)]
+            else:
+                audio_list = [raw_audio]
+        elif isinstance(raw_audio, (list, tuple)):
+            audio_list = list(raw_audio)
+        else:
+            audio_list = [raw_audio]
+        
+        for i, audio_tensor in enumerate(audio_list):
+            try:
+                if isinstance(audio_tensor, list):
+                    audio_tensor = torch.stack(audio_tensor)
+                if not isinstance(audio_tensor, torch.Tensor):
+                    audio_tensor = torch.tensor(audio_tensor)
+                audio_tensor = audio_tensor.to(device=device)
+                if raw_audio_lengths and i < len(raw_audio_lengths):
+                    actual_len = int(raw_audio_lengths[i])
+                    if actual_len > 0 and actual_len <= audio_tensor.shape[-1]:
+                        audio_tensor = audio_tensor[..., :actual_len]
+                if audio_tensor.numel() < 160:
+                    continue
+                
+                audio_embeds = self.audio_encoder(
+                    audio_tensor,
+                    use_streaming=use_streaming_flag,
+                    segment_duration_s=streaming_segment_duration,
+                )
+                final_embed = audio_embeds.squeeze(0)
+                embeddings.append(final_embed)
+                
+            except Exception as e:
+                print(f"[VibeVoice] Error encoding audio {i}: {e}")
+                import traceback
+                traceback.print_exc()
+                continue
+        
+        return tuple(embeddings)
+
+    def get_input_embeddings(self) -> torch.nn.Module:
+        """Return the text embedding layer (embed_tokens).
+        
+        vLLM uses this to get the embedding module for converting token IDs 
+        to embeddings during decode phase.
+        
+        Returns:
+            The embed_tokens module from the language model
+        """
+        # Get embed_tokens from the language model
+        if hasattr(self.language_model, 'model') and hasattr(self.language_model.model, 'embed_tokens'):
+            return self.language_model.model.embed_tokens
+        elif hasattr(self.language_model, 'embed_tokens'):
+            return self.language_model.embed_tokens
+        else:
+            # Try to get from inner model
+            inner = self.language_model
+            if hasattr(inner, 'language_model'):
+                inner = inner.language_model
+            if hasattr(inner, 'model') and hasattr(inner.model, 'embed_tokens'):
+                return inner.model.embed_tokens
+        
+        raise AttributeError("Cannot find embed_tokens layer")
+    
+    def embed_input_ids(
+        self, 
+        input_ids: torch.Tensor,
+        multimodal_embeddings: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None,
+        is_multimodal: Optional[torch.Tensor] = None,
+        **kwargs,  # Accept any additional kwargs for compatibility
+    ) -> torch.Tensor:
+        """Apply token embeddings to input_ids and merge with multimodal embeddings.
+        
+        This is the preferred method in vLLM V1 for converting token IDs
+        to embeddings and merging multimodal (audio) embeddings.
+        
+        Args:
+            input_ids: Tensor of token IDs to embed
+            multimodal_embeddings: Pre-computed multimodal embeddings (audio).
+                                   Can be a Tensor or a List of Tensors (vLLM standard).
+            is_multimodal: Boolean mask indicating which positions are multimodal
+            **kwargs: Additional arguments for compatibility
+            
+        Returns:
+            Tensor of embeddings with multimodal content merged in
+        """
+        from vllm.model_executor.models.utils import _merge_multimodal_embeddings
+        
+        # Get text embeddings
+        embed_tokens = self.get_input_embeddings()
+        inputs_embeds = embed_tokens(input_ids)
+        
+        # Merge multimodal embeddings if provided
+        if multimodal_embeddings is not None and is_multimodal is not None:
+            # Use vLLM's standard merge function which handles List[Tensor] correctly
+            inputs_embeds = _merge_multimodal_embeddings(
+                inputs_embeds,
+                multimodal_embeddings,
+                is_multimodal,
+            )
+        
+        return inputs_embeds
+
+    def get_language_model(self) -> torch.nn.Module:
+        """Return the language model backbone."""
+        return self.language_model
+
+    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> set[str]:
+        """Load model weights from checkpoint.
+        
+        The checkpoint has weights named like:
+        - lm_head.weight                    -> language_model.lm_head.weight
+        - model.language_model.layers.X...  -> language_model.model.layers.X...
+        - model.acoustic_tokenizer...       -> audio_encoder.acoustic_tokenizer...
+        - model.semantic_tokenizer...       -> audio_encoder.semantic_tokenizer...
+        - model.acoustic_connector...       -> audio_encoder.acoustic_connector...
+        - model.semantic_connector...       -> audio_encoder.semantic_connector...
+        
+        Let vLLM handle all dtype conversions according to --dtype flag.
+        """
+        # Map weight prefixes for VibeVoice
+        # The checkpoint uses "model." prefix, we need to remap it
+        mapper = WeightsMapper(
+            orig_to_new_prefix={
+                # Audio encoder components: model.X -> audio_encoder.X
+                "model.acoustic_tokenizer.": "audio_encoder.acoustic_tokenizer.",
+                "model.semantic_tokenizer.": "audio_encoder.semantic_tokenizer.",
+                "model.acoustic_connector.": "audio_encoder.acoustic_connector.",
+                "model.semantic_connector.": "audio_encoder.semantic_connector.",
+                # Language model: model.language_model.X -> language_model.model.X
+                "model.language_model.": "language_model.model.",
+                # LM head: lm_head.X -> language_model.lm_head.X
+                "lm_head.": "language_model.lm_head.",
+            }
+        )
+        
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights, mapper=mapper)
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor],
+        positions: torch.Tensor,
+        intermediate_tensors: Optional[IntermediateTensors] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        **kwargs: object,
+    ) -> Union[torch.Tensor, IntermediateTensors]:
+        """
+        Forward pass for VibeVoice ASR model.
+        
+        Handles embedding computation and language model forward pass.
+        Uses inputs_embeds if provided (from vLLM multimodal merge), 
+        otherwise computes embeddings from input_ids.
+        
+        Args:
+            input_ids: Token IDs. May be None when inputs_embeds is provided.
+            positions: Position indices for the input tokens.
+            intermediate_tensors: Intermediate tensors for pipeline parallelism.
+            inputs_embeds: Pre-computed embeddings (from multimodal merge or decode).
+        """
+        # PRIORITY: Use inputs_embeds if provided (from vLLM multimodal merge or decode)
+        # Only compute from input_ids if inputs_embeds is not available
+        if inputs_embeds is None and input_ids is not None:
+            # Compute embeddings from input_ids
+            inputs_embeds = self.get_input_embeddings()(input_ids)
+        
+        # If we have intermediate tensors (pipeline parallelism), don't use inputs_embeds
+        if intermediate_tensors is not None:
+            inputs_embeds = None
+        
+        # Get the inner model - handle both wrapped and direct language models
+        language_model = self.language_model
+        if hasattr(language_model, "language_model"):
+            language_model = language_model.language_model
+        
+        # Call the language model's model (Qwen2Model)
+        # vLLM V1 passes kv_caches and attn_metadata via context, not arguments
+        # IMPORTANT: Pass input_ids=None when using inputs_embeds to avoid double embedding
+        hidden_states = language_model.model(
+            input_ids=None,  # Always None when we have inputs_embeds
+            positions=positions,
+            intermediate_tensors=intermediate_tensors,
+            inputs_embeds=inputs_embeds
+        )
+        return hidden_states
+
+
+# Alias for training checkpoint compatibility
+VibeVoiceForASRTraining = VibeVoiceForCausalLM
diff --git a/vllm_plugin/scripts/gradio_asr_demo_api_video.py b/vllm_plugin/scripts/gradio_asr_demo_api_video.py
new file mode 100755
index 0000000..80f72a6
--- /dev/null
+++ b/vllm_plugin/scripts/gradio_asr_demo_api_video.py
@@ -0,0 +1,2209 @@
+#!/usr/bin/env python
+"""
+VibeVoice ASR Gradio Demo
+
+This demo uses the vLLM API server instead of loading the model directly.
+Supports concurrent requests (non-blocking) and streaming output.
+
+Usage:
+    python gradio_asr_demo_api.py --api_url http://localhost:8000
+"""
+import os
+import sys
+import io
+import json
+import time
+import base64
+import asyncio
+import tempfile
+import argparse
+import threading
+import subprocess
+import traceback
+import shutil
+import uuid
+from datetime import datetime
+from pathlib import Path
+from typing import List, Dict, Tuple, Optional, Generator
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+import httpx
+import numpy as np
+import soundfile as sf
+import gradio as gr
+from typing import AsyncGenerator
+
+# Try to import pydub for MP3 conversion
+try:
+    from pydub import AudioSegment
+    HAS_PYDUB = True
+except ImportError:
+    HAS_PYDUB = False
+    print("⚠️ Warning: pydub not available, falling back to WAV format")
+
+# Common audio extensions supported
+COMMON_AUDIO_EXTS = {
+    '.wav', '.mp3', '.flac', '.ogg', '.opus', '.m4a', '.aac',
+    '.wma', '.aiff', '.aif'
+}
+
+# Common video extensions supported
+COMMON_VIDEO_EXTS = {
+    '.mp4', '.webm', '.mov', '.avi', '.mkv', '.flv', '.wmv',
+    '.m4v', '.mpeg', '.mpg', '.3gp', '.ts'
+}
+
+# Default max video size in MB
+DEFAULT_MAX_VIDEO_SIZE_MB = 50
+
+# Default directory to save uploaded files
+DEFAULT_UPLOAD_SAVE_DIR = "local/from_custom"
+
+# Custom temporary directory
+CUSTOM_TEMP_DIR = os.environ.get("VIBEVOICE_TEMP_DIR", "/tmp/vibevoice_demo")
+os.makedirs(CUSTOM_TEMP_DIR, exist_ok=True)
+
+# ============================================================================
+# Cloudflared Tunnel Support
+# ============================================================================
+CLOUDFLARED_PATH = os.path.expanduser("~/.local/bin/cloudflared")
+
+def download_cloudflared():
+    """Download cloudflared binary if not exists"""
+    if os.path.exists(CLOUDFLARED_PATH):
+        return True
+    
+    print("📥 Downloading cloudflared...")
+    os.makedirs(os.path.dirname(CLOUDFLARED_PATH), exist_ok=True)
+    
+    download_url = "https://github.com/cloudflare/cloudflared/releases/latest/download/cloudflared-linux-amd64"
+    try:
+        subprocess.run(
+            ["wget", "-q", download_url, "-O", CLOUDFLARED_PATH],
+            check=True, timeout=120
+        )
+        os.chmod(CLOUDFLARED_PATH, 0o755)
+        print("✅ cloudflared downloaded successfully")
+        return True
+    except Exception as e:
+        print(f"❌ Failed to download cloudflared: {e}")
+        return False
+
+def start_cloudflared_tunnel(port: int):
+    """Start cloudflared tunnel and return the process"""
+    if not download_cloudflared():
+        print("❌ Cannot start cloudflared tunnel")
+        return None
+    
+    print(f"🌐 Starting cloudflared tunnel for port {port}...")
+    
+    process = subprocess.Popen(
+        [CLOUDFLARED_PATH, "tunnel", "--url", f"http://localhost:{port}"],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        text=True
+    )
+    
+    # Read output in background to find the URL
+    def read_output():
+        for line in process.stdout:
+            print(f"[cloudflared] {line.strip()}")
+    
+    thread = threading.Thread(target=read_output, daemon=True)
+    thread.start()
+    
+    # Give it a moment to start
+    time.sleep(3)
+    
+    return process
+
+
+# ============================================================================
+# Audio Utilities
+# ============================================================================
+
+def _guess_mime_type(path: str) -> str:
+    """Guess MIME type from file extension."""
+    ext = os.path.splitext(path)[1].lower()
+    mime_map = {
+        ".wav": "audio/wav",
+        ".mp3": "audio/mpeg",
+        ".m4a": "audio/mp4",
+        ".mp4": "video/mp4",
+        ".m4v": "video/mp4",
+        ".mov": "video/mp4",
+        ".webm": "video/webm",
+        ".flac": "audio/flac",
+        ".ogg": "audio/ogg",
+        ".opus": "audio/ogg",
+        ".aac": "audio/aac",
+    }
+    return mime_map.get(ext, "application/octet-stream")
+
+
+def _get_duration_seconds_ffprobe(path: str) -> float:
+    """Get audio duration using ffprobe."""
+    cmd = [
+        "ffprobe",
+        "-v", "error",
+        "-show_entries", "format=duration",
+        "-of", "default=noprint_wrappers=1:nokey=1",
+        path,
+    ]
+    try:
+        out = subprocess.check_output(cmd, stderr=subprocess.STDOUT).decode("utf-8").strip()
+        return float(out)
+    except Exception:
+        # Fallback: try soundfile
+        try:
+            info = sf.info(path)
+            return info.duration
+        except Exception:
+            return 0.0
+
+
+def load_audio_ffmpeg(path: str, target_sr: int = None) -> Tuple[np.ndarray, int]:
+    """Load audio file using ffmpeg for better format support."""
+    try:
+        # Use soundfile first (faster for supported formats)
+        audio_data, sr = sf.read(path, dtype='float32')
+        
+        # Debug: log audio info
+        print(f"[DEBUG] soundfile loaded: shape={audio_data.shape}, sr={sr}, dtype={audio_data.dtype}")
+        print(f"[DEBUG] audio range: min={audio_data.min():.6f}, max={audio_data.max():.6f}")
+        
+        # Convert to mono if multi-channel
+        if len(audio_data.shape) > 1:
+            audio_data = audio_data.mean(axis=1)  # Convert to mono
+            print(f"[DEBUG] Converted to mono: shape={audio_data.shape}")
+        
+        # Ensure data is in [-1, 1] range (soundfile should do this, but verify)
+        max_val = max(abs(audio_data.max()), abs(audio_data.min()))
+        if max_val > 10.0:
+            # Likely int16 or other integer format, normalize it
+            audio_data = audio_data / max_val
+            print(f"[DEBUG] Normalized audio (int format detected), original max_val={max_val}")
+        elif max_val > 1.0:
+            # Float format with slight overflow, just clip it
+            audio_data = np.clip(audio_data, -1.0, 1.0)
+            print(f"[DEBUG] Clipped audio (slight overflow), original max_val={max_val}")
+        
+        # Check for silent audio
+        if audio_data.max() == 0 and audio_data.min() == 0:
+            print(f"[WARNING] Audio appears to be completely silent!")
+        
+        return audio_data, sr
+    except Exception as e:
+        print(f"[DEBUG] soundfile failed: {e}, trying ffmpeg...")
+    
+    # Fallback to ffmpeg
+    try:
+        target_sr = target_sr or 16000
+        cmd = [
+            "ffmpeg", "-i", path,
+            "-f", "f32le", "-acodec", "pcm_f32le",
+            "-ac", "1", "-ar", str(target_sr),
+            "-"
+        ]
+        process = subprocess.Popen(
+            cmd, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL
+        )
+        audio_bytes, _ = process.communicate()
+        audio_data = np.frombuffer(audio_bytes, dtype=np.float32)
+        
+        print(f"[DEBUG] ffmpeg loaded: shape={audio_data.shape}, sr={target_sr}")
+        print(f"[DEBUG] audio range: min={audio_data.min():.6f}, max={audio_data.max():.6f}")
+        
+        # Check for silent audio
+        if len(audio_data) == 0:
+            raise RuntimeError("ffmpeg returned empty audio data")
+        if audio_data.max() == 0 and audio_data.min() == 0:
+            print(f"[WARNING] Audio appears to be completely silent!")
+        
+        return audio_data, target_sr
+    except Exception as e:
+        raise RuntimeError(f"Failed to load audio: {e}")
+
+
+def get_file_size_mb(file_path: str) -> float:
+    """Get file size in MB."""
+    try:
+        return os.path.getsize(file_path) / (1024 * 1024)
+    except Exception:
+        return 0.0
+
+
+def is_video_file(file_path: str) -> bool:
+    """Check if the file is a video file based on extension."""
+    ext = os.path.splitext(file_path)[1].lower()
+    return ext in COMMON_VIDEO_EXTS
+
+
+def format_srt_time(seconds: float) -> str:
+    """Convert seconds to SRT time format (HH:MM:SS,mmm)."""
+    hours = int(seconds // 3600)
+    minutes = int((seconds % 3600) // 60)
+    secs = int(seconds % 60)
+    millis = int((seconds - int(seconds)) * 1000)
+    return f"{hours:02d}:{minutes:02d}:{secs:02d},{millis:03d}"
+
+
+def segments_to_srt(segments: List[Dict]) -> str:
+    """
+    Convert ASR segments to SRT subtitle format.
+    
+    Args:
+        segments: List of segment dictionaries with Start, End, Content keys
+    
+    Returns:
+        SRT formatted string
+    """
+    srt_lines = []
+    
+    for i, seg in enumerate(segments, 1):
+        start = seg.get('Start', seg.get('start', seg.get('Start time', 0)))
+        end = seg.get('End', seg.get('end', seg.get('End time', 0)))
+        content = seg.get('Content', seg.get('content', seg.get('text', '')))
+        speaker = seg.get('Speaker', seg.get('speaker', seg.get('Speaker ID', None)))
+        
+        if start is None or end is None:
+            continue
+        
+        start_time = format_srt_time(float(start))
+        end_time = format_srt_time(float(end))
+        
+        # Add speaker prefix if available
+        text = f"[Speaker {speaker}] {content}" if speaker is not None else content
+        
+        srt_lines.append(f"{i}")
+        srt_lines.append(f"{start_time} --> {end_time}")
+        srt_lines.append(text)
+        srt_lines.append("")  # Empty line between entries
+    
+    return "\n".join(srt_lines)
+
+
+def segments_to_vtt(segments: List[Dict]) -> str:
+    """
+    Convert ASR segments to WebVTT subtitle format (for HTML5 video).
+    
+    Args:
+        segments: List of segment dictionaries with Start, End, Content keys
+    
+    Returns:
+        WebVTT formatted string
+    """
+    vtt_lines = ["WEBVTT", ""]
+    
+    for i, seg in enumerate(segments, 1):
+        start = seg.get('Start', seg.get('start', seg.get('Start time', 0)))
+        end = seg.get('End', seg.get('end', seg.get('End time', 0)))
+        content = seg.get('Content', seg.get('content', seg.get('text', '')))
+        speaker = seg.get('Speaker', seg.get('speaker', seg.get('Speaker ID', None)))
+        
+        if start is None or end is None:
+            continue
+        
+        # WebVTT uses HH:MM:SS.mmm format (dot instead of comma)
+        start_time = format_srt_time(float(start)).replace(',', '.')
+        end_time = format_srt_time(float(end)).replace(',', '.')
+        
+        # Add speaker prefix if available
+        text = f"[Speaker {speaker}] {content}" if speaker is not None else content
+        
+        vtt_lines.append(f"{i}")
+        vtt_lines.append(f"{start_time} --> {end_time}")
+        vtt_lines.append(text)
+        vtt_lines.append("")  # Empty line between entries
+    
+    return "\n".join(vtt_lines)
+
+
+# Audio formats that need conversion (browsers and some APIs may not support them directly)
+AUDIO_FORMATS_NEED_CONVERSION = {'.opus', '.ogg', '.flac', '.aiff', '.aif', '.wma'}
+
+# Audio formats that can be used directly
+AUDIO_FORMATS_DIRECT = {'.wav', '.mp3', '.m4a', '.aac'}
+
+
+def convert_audio_to_mp3(
+    audio_path: str,
+    output_path: Optional[str] = None,
+    sample_rate: int = 16000,
+    bitrate: str = "128k"
+) -> Tuple[Optional[str], Optional[str]]:
+    """
+    Convert audio file to MP3 format using ffmpeg.
+    
+    This is useful for converting formats like opus, ogg, flac that may not be
+    well-supported by browsers or some APIs.
+    
+    Args:
+        audio_path: Path to the input audio file
+        output_path: Optional output path. If None, creates a temp file.
+        sample_rate: Target sample rate
+        bitrate: Audio bitrate (e.g., '128k')
+    
+    Returns:
+        Tuple of (mp3_path, error_message)
+        - If successful: (mp3_path, None)
+        - If failed: (None, error_message)
+    """
+    try:
+        if output_path is None:
+            temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3", dir=CUSTOM_TEMP_DIR)
+            temp_file.close()
+            output_path = temp_file.name
+        
+        cmd = [
+            "ffmpeg", "-y",  # Overwrite output file
+            "-i", audio_path,
+            "-acodec", "libmp3lame",  # MP3 codec
+            "-ar", str(sample_rate),  # Sample rate
+            "-ac", "1",  # Mono
+            "-b:a", bitrate,  # Audio bitrate
+            output_path
+        ]
+        
+        result = subprocess.run(
+            cmd,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            timeout=300  # 5 minutes timeout
+        )
+        
+        if result.returncode != 0:
+            error_msg = result.stderr.decode('utf-8', errors='ignore')
+            if os.path.exists(output_path):
+                os.unlink(output_path)
+            return None, f"ffmpeg error: {error_msg[:500]}"
+        
+        if not os.path.exists(output_path) or os.path.getsize(output_path) == 0:
+            return None, "Failed to convert audio: output file is empty"
+        
+        return output_path, None
+        
+    except subprocess.TimeoutExpired:
+        return None, "Audio conversion timed out (>5 minutes)"
+    except Exception as e:
+        return None, f"Error converting audio: {str(e)}"
+
+
+def extract_audio_from_video(
+    video_path: str,
+    output_path: Optional[str] = None,
+    sample_rate: int = 16000,
+    output_format: str = "mp3",
+    bitrate: str = "128k"
+) -> Tuple[Optional[str], Optional[str]]:
+    """
+    Extract audio from video file using ffmpeg.
+    
+    Args:
+        video_path: Path to the video file
+        output_path: Optional output path for extracted audio. If None, creates a temp file.
+        sample_rate: Target sample rate for extracted audio
+        output_format: Output audio format ('mp3' or 'wav')
+        bitrate: Audio bitrate for mp3 (e.g., '128k')
+    
+    Returns:
+        Tuple of (audio_path, error_message)
+        - If successful: (audio_path, None)
+        - If failed: (None, error_message)
+    """
+    try:
+        if output_path is None:
+            suffix = f".{output_format}"
+            temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=suffix, dir=CUSTOM_TEMP_DIR)
+            temp_file.close()
+            output_path = temp_file.name
+        
+        # Use ffmpeg to extract audio
+        if output_format == "mp3":
+            cmd = [
+                "ffmpeg", "-y",  # Overwrite output file
+                "-i", video_path,
+                "-vn",  # No video
+                "-acodec", "libmp3lame",  # MP3 codec
+                "-ar", str(sample_rate),  # Sample rate
+                "-ac", "1",  # Mono
+                "-b:a", bitrate,  # Audio bitrate
+                output_path
+            ]
+        else:
+            cmd = [
+                "ffmpeg", "-y",  # Overwrite output file
+                "-i", video_path,
+                "-vn",  # No video
+                "-acodec", "pcm_s16le",  # PCM 16-bit
+                "-ar", str(sample_rate),  # Sample rate
+                "-ac", "1",  # Mono
+                output_path
+            ]
+        
+        result = subprocess.run(
+            cmd,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            timeout=300  # 5 minutes timeout
+        )
+        
+        if result.returncode != 0:
+            error_msg = result.stderr.decode('utf-8', errors='ignore')
+            # Clean up temp file on error
+            if os.path.exists(output_path):
+                os.unlink(output_path)
+            return None, f"ffmpeg error: {error_msg[:500]}"
+        
+        # Verify the output file exists and has content
+        if not os.path.exists(output_path) or os.path.getsize(output_path) == 0:
+            return None, "Failed to extract audio: output file is empty"
+        
+        return output_path, None
+        
+    except subprocess.TimeoutExpired:
+        return None, "Audio extraction timed out (>5 minutes)"
+    except Exception as e:
+        return None, f"Error extracting audio: {str(e)}"
+
+
+def convert_video_to_mp4(
+    video_path: str,
+    output_path: Optional[str] = None,
+    height: int = 480,
+    crf: int = 28,
+    fps: int = 30
+) -> Tuple[Optional[str], Optional[str]]:
+    """
+    Convert video to MP4 format with compression (480p by default).
+    
+    Args:
+        video_path: Path to the input video file (e.g., WebM)
+        output_path: Optional output path. If None, creates a temp file.
+        height: Target video height (width auto-scaled to maintain aspect ratio)
+        crf: Constant Rate Factor for compression (18-28 recommended, higher = smaller file)
+        fps: Target frame rate (default 30fps)
+    
+    Returns:
+        Tuple of (mp4_path, error_message)
+        - If successful: (mp4_path, None)
+        - If failed: (None, error_message)
+    """
+    try:
+        if output_path is None:
+            temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4", dir=CUSTOM_TEMP_DIR)
+            temp_file.close()
+            output_path = temp_file.name
+        
+        # Use ffmpeg to convert to MP4 with H.264 codec
+        # Scale to target height while maintaining aspect ratio (-2 ensures even dimensions)
+        cmd = [
+            "ffmpeg", "-y",  # Overwrite output file
+            "-i", video_path,
+            "-vf", f"scale=-2:{height},fps={fps}",  # Scale to 480p height + set fps
+            "-c:v", "libx264",  # H.264 video codec
+            "-preset", "fast",  # Encoding speed/compression tradeoff
+            "-crf", str(crf),  # Quality (lower = better, 18-28 typical)
+            "-c:a", "aac",  # AAC audio codec
+            "-b:a", "128k",  # Audio bitrate
+            "-movflags", "+faststart",  # Enable streaming
+            output_path
+        ]
+        
+        result = subprocess.run(
+            cmd,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            timeout=600  # 10 minutes timeout
+        )
+        
+        if result.returncode != 0:
+            error_msg = result.stderr.decode('utf-8', errors='ignore')
+            # Clean up temp file on error
+            if os.path.exists(output_path):
+                os.unlink(output_path)
+            return None, f"ffmpeg error: {error_msg[:500]}"
+        
+        # Verify the output file exists and has content
+        if not os.path.exists(output_path) or os.path.getsize(output_path) == 0:
+            return None, "Failed to convert video: output file is empty"
+        
+        return output_path, None
+        
+    except subprocess.TimeoutExpired:
+        return None, "Video conversion timed out (>10 minutes)"
+    except Exception as e:
+        return None, f"Error converting video: {str(e)}"
+
+
+def parse_time_to_seconds(val: Optional[str]) -> Optional[float]:
+    """Parse seconds or hh:mm:ss to float seconds."""
+    if val is None:
+        return None
+    val = val.strip()
+    if not val:
+        return None
+    try:
+        return float(val)
+    except ValueError:
+        pass
+    if ":" in val:
+        parts = val.split(":")
+        if not all(p.strip().replace(".", "", 1).isdigit() for p in parts):
+            return None
+        parts = [float(p) for p in parts]
+        if len(parts) == 3:
+            h, m, s = parts
+        elif len(parts) == 2:
+            h = 0
+            m, s = parts
+        else:
+            return None
+        return h * 3600 + m * 60 + s
+    return None
+
+
+def slice_audio_to_temp(
+    audio_path: str,
+    start_sec: Optional[float],
+    end_sec: Optional[float]
+) -> Tuple[Optional[str], Optional[str]]:
+    """Slice audio to [start_sec, end_sec) and write to a temp WAV file."""
+    try:
+        audio_data, sample_rate = load_audio_ffmpeg(audio_path)
+        n_samples = len(audio_data)
+        full_duration = n_samples / float(sample_rate)
+        start = 0.0 if start_sec is None else max(0.0, start_sec)
+        end = full_duration if end_sec is None else min(full_duration, end_sec)
+        if end <= start:
+            return None, f"Invalid time range: start={start}, end={end}"
+        start_idx = int(start * sample_rate)
+        end_idx = int(end * sample_rate)
+        segment = audio_data[start_idx:end_idx]
+        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".wav", dir=CUSTOM_TEMP_DIR)
+        temp_file.close()
+        # Use 32767.0 instead of 32768.0 to avoid potential overflow
+        segment_int16 = (segment * 32767.0).astype(np.int16)
+        sf.write(temp_file.name, segment_int16, sample_rate, subtype='PCM_16')
+        return temp_file.name, None
+    except Exception as e:
+        return None, f"Error slicing audio: {e}"
+
+
+def clip_and_encode_audio(
+    audio_data: np.ndarray,
+    sr: int,
+    start_time: float,
+    end_time: float,
+    segment_idx: int,
+    use_mp3: bool = True,
+    target_sr: int = 16000,
+    mp3_bitrate: str = "32k"
+) -> Tuple[int, Optional[str], Optional[str]]:
+    """Clip audio segment and encode to base64."""
+    try:
+        start_sample = int(start_time * sr)
+        end_sample = int(end_time * sr)
+        start_sample = max(0, start_sample)
+        end_sample = min(len(audio_data), end_sample)
+        
+        if start_sample >= end_sample:
+            return segment_idx, None, f"Invalid segment range: {start_time}-{end_time}"
+        
+        segment_data = audio_data[start_sample:end_sample]
+        
+        # Resample if needed
+        if sr != target_sr and target_sr < sr:
+            import scipy.signal
+            num_samples = int(len(segment_data) * target_sr / sr)
+            segment_data = scipy.signal.resample(segment_data, num_samples)
+            sr = target_sr
+        
+        # Use 32767.0 instead of 32768.0 to avoid potential overflow
+        segment_data_int16 = (segment_data * 32767.0).astype(np.int16)
+        
+        # Convert to MP3 if pydub is available
+        if use_mp3 and HAS_PYDUB:
+            try:
+                wav_buffer = io.BytesIO()
+                sf.write(wav_buffer, segment_data_int16, sr, format='WAV', subtype='PCM_16')
+                wav_buffer.seek(0)
+                audio_segment = AudioSegment.from_wav(wav_buffer)
+                mp3_buffer = io.BytesIO()
+                audio_segment.export(mp3_buffer, format='mp3', bitrate=mp3_bitrate)
+                mp3_buffer.seek(0)
+                audio_bytes = mp3_buffer.read()
+                audio_base64 = base64.b64encode(audio_bytes).decode('utf-8')
+                return segment_idx, f"data:audio/mpeg;base64,{audio_base64}", None
+            except Exception:
+                pass
+        
+        # Fallback to WAV
+        wav_buffer = io.BytesIO()
+        sf.write(wav_buffer, segment_data_int16, sr, format='WAV', subtype='PCM_16')
+        wav_buffer.seek(0)
+        audio_bytes = wav_buffer.read()
+        audio_base64 = base64.b64encode(audio_bytes).decode('utf-8')
+        return segment_idx, f"data:audio/wav;base64,{audio_base64}", None
+        
+    except Exception as e:
+        return segment_idx, None, f"Error: {str(e)}"
+
+
+def extract_audio_segments(audio_path: str, segments: List[Dict]) -> List[Tuple[str, str, Optional[str]]]:
+    """Extract multiple segments from audio file with parallel processing."""
+    try:
+        audio_data, sr = load_audio_ffmpeg(audio_path)
+        
+        tasks = []
+        use_mp3 = HAS_PYDUB
+        
+        for i, seg in enumerate(segments):
+            start_time = seg.get('Start', seg.get('start', seg.get('Start time', 0)))
+            end_time = seg.get('End', seg.get('end', seg.get('End time', 0)))
+            if start_time is not None and end_time is not None:
+                tasks.append((audio_data, sr, float(start_time), float(end_time), i, use_mp3))
+        
+        results = []
+        max_workers = os.cpu_count() or 4
+        
+        with ThreadPoolExecutor(max_workers=max_workers) as executor:
+            futures = {
+                executor.submit(clip_and_encode_audio, *task): task[4]
+                for task in tasks
+            }
+            for future in as_completed(futures):
+                results.append(future.result())
+        
+        results.sort(key=lambda x: x[0])
+        
+        audio_segments = []
+        for i, (idx, audio_src, error_msg) in enumerate(results):
+            if idx < len(segments):
+                seg = segments[idx]
+                label = f"Segment {idx + 1}"
+                audio_segments.append((label, audio_src, error_msg))
+        
+        return audio_segments
+        
+    except Exception as e:
+        print(f"Error loading audio file: {e}")
+        return []
+
+
+# ============================================================================
+# API Client
+# ============================================================================
+
+class VibeVoiceAPIClient:
+    """Client for VibeVoice vLLM API."""
+    
+    def __init__(self, api_url: str = "http://localhost:8000", model_name: str = None):
+        self.api_url = api_url.rstrip("/")
+        self._model_name = model_name  # User-specified model name (can be None for auto-detect)
+        self._available_models: List[str] = []  # Cached available models
+        self.endpoint = f"{self.api_url}/v1/chat/completions"
+    
+    @property
+    def model_name(self) -> str:
+        """Get the model name (auto-detected if not specified)."""
+        if self._model_name:
+            return self._model_name
+        if self._available_models:
+            return self._available_models[0]
+        return "vibevoice"  # Fallback default
+    
+    @model_name.setter
+    def model_name(self, value: str):
+        """Set the model name."""
+        self._model_name = value
+    
+    def get_available_models_sync(self) -> List[str]:
+        """Fetch available models from vLLM API (synchronous)."""
+        try:
+            response = httpx.get(f"{self.api_url}/v1/models", timeout=10)
+            if response.status_code == 200:
+                data = response.json()
+                models = [m.get('id') for m in data.get('data', []) if m.get('id')]
+                self._available_models = models
+                print(f"📋 Available models: {models}")
+                return models
+            return []
+        except Exception as e:
+            print(f"⚠️ Failed to fetch models: {e}")
+            return []
+    
+    async def get_available_models(self) -> List[str]:
+        """Fetch available models from vLLM API (async)."""
+        try:
+            async with httpx.AsyncClient() as client:
+                response = await client.get(f"{self.api_url}/v1/models", timeout=10)
+                if response.status_code == 200:
+                    data = response.json()
+                    models = [m.get('id') for m in data.get('data', []) if m.get('id')]
+                    self._available_models = models
+                    return models
+            return []
+        except Exception as e:
+            print(f"⚠️ Failed to fetch models: {e}")
+            return []
+    
+    async def check_health(self) -> Tuple[bool, str]:
+        """Check if the API server is healthy."""
+        try:
+            async with httpx.AsyncClient() as client:
+                response = await client.get(f"{self.api_url}/health", timeout=5)
+                if response.status_code == 200:
+                    return True, "API server is healthy"
+                return False, f"API returned status {response.status_code}"
+        except httpx.ConnectError:
+            return False, "Cannot connect to API server"
+        except Exception as e:
+            return False, f"Health check failed: {e}"
+    
+    def check_health_sync(self) -> Tuple[bool, str]:
+        """Synchronous health check for startup."""
+        try:
+            response = httpx.get(f"{self.api_url}/health", timeout=5)
+            if response.status_code == 200:
+                return True, "API server is healthy"
+            return False, f"API returned status {response.status_code}"
+        except httpx.ConnectError:
+            return False, "Cannot connect to API server"
+        except Exception as e:
+            return False, f"Health check failed: {e}"
+    
+    async def transcribe_streaming(
+        self,
+        audio_path: str,
+        max_tokens: int = 4096,
+        temperature: float = 0.0,
+        top_p: float = 1.0,
+        context_info: str = None,
+        timeout: int = 1200,
+    ) -> AsyncGenerator[Tuple[str, Optional[Dict]], None]:
+        """
+        Transcribe audio using streaming API (async version).
+        
+        Yields:
+            Tuple of (accumulated_text, final_result_or_none)
+        """
+        # Get audio duration
+        duration = _get_duration_seconds_ffprobe(audio_path)
+        
+        # Read and encode audio
+        with open(audio_path, "rb") as f:
+            audio_bytes = f.read()
+        audio_b64 = base64.b64encode(audio_bytes).decode("utf-8")
+        
+        # Build prompt
+        show_keys = ["Start", "End", "Speaker", "Content"]
+        prompt_text = (
+            f"This is a {duration:.2f} seconds audio, please transcribe it with these keys: "
+            + ", ".join(show_keys)
+        )
+        
+        # Add context info if provided
+        if context_info and context_info.strip():
+            prompt_text += f"\n\nContext information (hotwords, speaker names, etc.):\n{context_info.strip()}"
+        
+        # Build request payload
+        mime = _guess_mime_type(audio_path)
+        data_url = f"data:{mime};base64,{audio_b64}"
+        
+        payload = {
+            "model": self.model_name,
+            "messages": [
+                {
+                    "role": "system",
+                    "content": "You are a helpful assistant that transcribes audio input into text output in JSON format."
+                },
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "audio_url", "audio_url": {"url": data_url}},
+                        {"type": "text", "text": prompt_text}
+                    ]
+                }
+            ],
+            "max_tokens": max_tokens,
+            "temperature": temperature,
+            "top_p": top_p,
+            "stream": True,
+            "stream_options": {"include_usage": True},  # Enable token statistics
+        }
+        
+        # Send request with streaming using async httpx
+        try:
+            async with httpx.AsyncClient(timeout=httpx.Timeout(timeout)) as client:
+                async with client.stream(
+                    "POST",
+                    self.endpoint,
+                    json=payload,
+                ) as response:
+                    if response.status_code != 200:
+                        error_text = await response.aread()
+                        error_msg = f"API error: {response.status_code} - {error_text.decode()}"
+                        yield error_msg, {"error": error_msg}
+                        return
+                    
+                    accumulated_text = ""
+                    usage_info = None
+                    stopped = False
+                    
+                    # Declare global at function level for proper access
+                    global stop_generation_flag
+                    
+                    # Process streaming lines with periodic stop flag checking
+                    async for line in response.aiter_lines():
+                        # Check stop flag after each line
+                        if stop_generation_flag:
+                            stopped = True
+                            print("[INFO] Stop flag detected, breaking out of stream...")
+                            break
+                        
+                        if line:
+                            if line.startswith("data: "):
+                                json_str = line[6:]
+                                if json_str.strip() == "[DONE]":
+                                    break
+                                
+                                try:
+                                    data = json.loads(json_str)
+                                    
+                                    # Check for usage info (sent in final chunk)
+                                    if 'usage' in data and data['usage']:
+                                        usage_info = data['usage']
+                                    
+                                    if 'choices' in data and data['choices']:
+                                        delta = data['choices'][0].get('delta', {})
+                                        content = delta.get('content', '')
+                                        
+                                        if content:
+                                            # Handle incremental or full text
+                                            if content.startswith(accumulated_text):
+                                                accumulated_text = content
+                                            else:
+                                                accumulated_text += content
+                                            
+                                            yield accumulated_text, None
+                                        
+                                except json.JSONDecodeError:
+                                    pass
+                    
+                    # If stopped, try to close the response to stop receiving more data
+                    if stopped:
+                        try:
+                            await response.aclose()
+                            print("[INFO] Response closed after stop")
+                        except Exception:
+                            pass
+            
+            # Parse final result with partial parsing support
+            segments, parse_warning = self._parse_segments(accumulated_text)
+            if segments is None:
+                segments = []
+            
+            final_result = {
+                "raw_text": accumulated_text,
+                "segments": segments,
+                "duration": duration,
+                "usage": usage_info,  # Include token statistics
+                "stopped": stopped,  # Whether generation was stopped by user
+                "parse_warning": parse_warning,  # Warning if partial parse
+            }
+            
+            yield accumulated_text, final_result
+            
+        except httpx.TimeoutException:
+            yield "Request timed out", {"error": "timeout"}
+        except Exception as e:
+            yield f"Error: {str(e)}", {"error": str(e)}
+
+    def _parse_segments(self, raw_text: str) -> Tuple[Optional[List[Dict]], Optional[str]]:
+        """
+        Parse segments from raw API response.
+        Handles truncated responses by extracting complete segments.
+        
+        Returns:
+            Tuple of (segments_list, warning_message)
+            - If fully successful: (segments, None)
+            - If partially successful: (segments, warning_message)
+            - If failed: (None, error_message)
+        """
+        if not raw_text:
+            return None, "Empty response"
+        
+        # Try to find JSON array in the response
+        text = raw_text.strip()
+        
+        # Try direct parse first
+        try:
+            result = json.loads(text)
+            if isinstance(result, list):
+                return result, None
+            elif isinstance(result, dict) and "segments" in result:
+                return result["segments"], None
+            elif isinstance(result, dict):
+                # Single segment
+                return [result], None
+        except json.JSONDecodeError:
+            pass
+        
+        # Try to extract JSON array from text
+        import re
+        
+        # Try to find array pattern
+        array_match = re.search(r'\[[\s\S]*\]', text)
+        if array_match:
+            try:
+                result = json.loads(array_match.group())
+                if isinstance(result, list):
+                    return result, None
+            except json.JSONDecodeError:
+                pass
+        
+        # Try to find object with segments
+        obj_match = re.search(r'\{[\s\S]*"segments"[\s\S]*\}', text)
+        if obj_match:
+            try:
+                result = json.loads(obj_match.group())
+                if "segments" in result:
+                    return result["segments"], None
+            except json.JSONDecodeError:
+                pass
+        
+        # ===== Handle truncated response =====
+        # Try to parse individual complete segments from truncated array
+        segments = self._parse_truncated_segments(text)
+        if segments:
+            return segments, f"⚠️ Partial parse: {len(segments)} segments recovered from truncated response"
+        
+        return None, "Cannot parse JSON from response"
+
+    def _parse_truncated_segments(self, text: str) -> Optional[List[Dict]]:
+        """
+        Parse complete segments from a truncated JSON array response.
+        This handles cases where the response is cut off mid-segment.
+        
+        Strategy:
+        1. Find all complete JSON objects {...} that look like segments
+        2. Validate each has expected keys (Start, End, Content or Speaker)
+        3. Try to recover incomplete last segment (e.g., repetition truncation)
+        """
+        # Check if text starts with array
+        text = text.strip()
+        if not text.startswith('['):
+            # Try to find array start
+            array_start = text.find('[')
+            if array_start == -1:
+                return None
+            text = text[array_start:]
+        
+        # Find all complete JSON objects
+        # Pattern: {...} that are properly closed
+        segments = []
+        depth = 0
+        obj_start = -1
+        in_string = False
+        escape_next = False
+        
+        for i, char in enumerate(text):
+            if escape_next:
+                escape_next = False
+                continue
+            
+            if char == '\\':
+                escape_next = True
+                continue
+            
+            if char == '"' and not escape_next:
+                in_string = not in_string
+                continue
+            
+            if in_string:
+                continue
+            
+            if char == '{':
+                if depth == 0:
+                    obj_start = i
+                depth += 1
+            elif char == '}':
+                depth -= 1
+                if depth == 0 and obj_start != -1:
+                    # Found a complete object
+                    obj_str = text[obj_start:i+1]
+                    try:
+                        obj = json.loads(obj_str)
+                        # Validate it looks like a segment
+                        if self._is_valid_segment(obj):
+                            segments.append(obj)
+                    except json.JSONDecodeError:
+                        pass
+                    obj_start = -1
+        
+        # Try to recover incomplete last segment (truncated due to repetition)
+        if obj_start != -1:
+            incomplete_text = text[obj_start:]
+            recovered = self._recover_incomplete_segment(incomplete_text)
+            if recovered and self._is_valid_segment(recovered):
+                segments.append(recovered)
+        
+        return segments if segments else None
+
+    def _recover_incomplete_segment(self, incomplete_text: str) -> Optional[Dict]:
+        """
+        Try to recover an incomplete segment that was truncated.
+        Handles cases like repetition where Content is cut off mid-string.
+        
+        Example input:
+        {"Start":198.36,"End":206.86,"Speaker":0,"Content":"I'm not gonna do it, I'm not gonna do it, I'm not...
+        """
+        import re
+        
+        # Try to extract available fields
+        segment = {}
+        
+        # Extract Start
+        start_match = re.search(r'"Start"\s*:\s*([0-9.]+)', incomplete_text)
+        if start_match:
+            segment['Start'] = float(start_match.group(1))
+        
+        # Extract End
+        end_match = re.search(r'"End"\s*:\s*([0-9.]+)', incomplete_text)
+        if end_match:
+            segment['End'] = float(end_match.group(1))
+        
+        # Extract Speaker
+        speaker_match = re.search(r'"Speaker"\s*:\s*([0-9]+)', incomplete_text)
+        if speaker_match:
+            segment['Speaker'] = int(speaker_match.group(1))
+        
+        # Extract Content - handle truncated string
+        content_match = re.search(r'"Content"\s*:\s*"', incomplete_text)
+        if content_match:
+            content_start = content_match.end()
+            # Find the content, may be truncated
+            content_text = incomplete_text[content_start:]
+            
+            # Check for repetition pattern and clean it
+            cleaned_content = self._clean_repetition(content_text)
+            if cleaned_content:
+                segment['Content'] = cleaned_content
+                segment['_truncated'] = True  # Mark as recovered from truncation
+        
+        # Must have at least Start, End to be valid
+        if 'Start' in segment and 'End' in segment:
+            return segment
+        
+        return None
+
+    def _clean_repetition(self, content: str) -> Optional[str]:
+        """
+        Clean content from truncated string.
+        For repetition cases, keep first 500 characters.
+        """
+        # Remove trailing incomplete quote if any
+        content = content.rstrip('"')
+        
+        if not content:
+            return None
+        
+        # Keep first 500 characters for repetition cases
+        if len(content) > 500:
+            return content[:500] + "..."
+        
+        return content
+
+    def _is_valid_segment(self, obj: Dict) -> bool:
+        """
+        Check if a dict looks like a valid ASR segment.
+        Must have Start, End, and either Content or Speaker.
+        """
+        if not isinstance(obj, dict):
+            return False
+        
+        # Check for time boundaries (various possible key names)
+        has_start = any(k in obj for k in ['Start', 'start', 'Start time'])
+        has_end = any(k in obj for k in ['End', 'end', 'End time'])
+        
+        if not (has_start and has_end):
+            return False
+        
+        # Should have content or speaker
+        has_content = any(k in obj for k in ['Content', 'content', 'text'])
+        has_speaker = any(k in obj for k in ['Speaker', 'speaker', 'Speaker ID'])
+        
+        return has_content or has_speaker
+
+
+# ============================================================================
+# Global State
+# ============================================================================
+
+api_client: Optional[VibeVoiceAPIClient] = None
+
+# Global stop flag for generation
+stop_generation_flag = False
+
+# Event to signal stop for async operations
+stop_event: Optional[asyncio.Event] = None
+
+
+# ============================================================================
+# Gradio Interface Functions
+# ============================================================================
+
+async def transcribe_audio(
+    media_input,
+    max_new_tokens: int,
+    temperature: float,
+    top_p: float,
+    do_sample: bool,
+    context_info: str = "",
+    max_video_size_mb: float = DEFAULT_MAX_VIDEO_SIZE_MB
+) -> AsyncGenerator[Tuple[str, str, Optional[str], Optional[str], Optional[str]], None]:
+    """
+    Transcribe audio/video using API and return results (async streaming version).
+    
+    Args:
+        media_input: Audio/Video file path or tuple (sample_rate, audio_data) for microphone
+        max_new_tokens: Maximum tokens to generate
+        temperature: Temperature for sampling
+        top_p: Top-p for nucleus sampling
+        do_sample: Whether to use sampling (affects temperature)
+        context_info: Optional context information
+        max_video_size_mb: Maximum video file size in MB
+    
+    Yields:
+        Tuple of (raw_text, audio_segments_html, srt_content, video_path, vtt_content)
+    """
+    global api_client, stop_generation_flag
+    
+    # Reset stop flag at the start of each transcription
+    stop_generation_flag = False
+    print("[INFO] Stop flag reset at transcribe_audio start")
+    
+    if api_client is None:
+        yield "❌ API client not initialized. Please check API URL.", "", None, None, None
+        return
+    
+    # Check API health (async)
+    healthy, msg = await api_client.check_health()
+    if not healthy:
+        yield f"❌ API server not available: {msg}", "", None, None, None
+        return
+    
+    if media_input is None:
+        yield "❌ Please provide an audio or video input.", "", None, None, None
+        return
+    
+    try:
+        print("[INFO] Transcription requested via API")
+        
+        # Determine audio path and track if input is video
+        audio_path = None
+        original_video_path = None  # Track original video for playback with subtitles
+        temp_file_to_cleanup = None
+        extracted_audio_to_cleanup = None  # Track extracted audio from video
+        is_video_input = False
+        
+        # Handle media input
+        if isinstance(media_input, str):
+            # Check if uploaded file is a video
+            if is_video_file(media_input):
+                is_video_input = True
+                original_video_path = media_input  # Keep video path for subtitle playback
+                video_size = get_file_size_mb(media_input)
+                print(f"[INFO] Uploaded video file size: {video_size:.2f} MB (limit: {max_video_size_mb} MB)")
+                
+                if video_size > max_video_size_mb:
+                    yield f"❌ Video file too large: {video_size:.2f} MB. Maximum allowed: {max_video_size_mb} MB", "", None, None, None
+                    return
+                
+                yield f"🎬 Extracting audio from video ({video_size:.2f} MB)...", "", None, None, None
+                extracted_path, extract_error = extract_audio_from_video(media_input)
+                
+                if extract_error:
+                    yield f"❌ Failed to extract audio from video: {extract_error}", "", None, None, None
+                    return
+                
+                audio_path = extracted_path
+                extracted_audio_to_cleanup = extracted_path
+                print(f"[INFO] Extracted audio from video: {audio_path}")
+            else:
+                # Audio file
+                audio_path = media_input
+                print(f"[INFO] Using uploaded audio file: {audio_path}")
+        elif isinstance(media_input, tuple):
+            # Gradio microphone input: (sample_rate, audio_array)
+            sample_rate, audio_array = media_input
+            audio_array = np.array(audio_array, dtype=np.float32)
+            
+            # Debug: log input audio info
+            print(f"[DEBUG] Microphone input: shape={audio_array.shape}, sr={sample_rate}")
+            print(f"[DEBUG] Microphone audio range: min={audio_array.min():.6f}, max={audio_array.max():.6f}")
+            
+            # Normalize to [-1, 1] range properly
+            max_val = max(abs(audio_array.max()), abs(audio_array.min()))
+            if max_val > 10.0:
+                # Data is likely int16 or similar, normalize it
+                audio_array = audio_array / max_val
+                print(f"[DEBUG] Normalized microphone audio (int format detected), original max_val={max_val}")
+            elif max_val > 1.0:
+                # Float format with slight overflow, just clip it
+                audio_array = np.clip(audio_array, -1.0, 1.0)
+                print(f"[DEBUG] Clipped microphone audio (slight overflow), original max_val={max_val}")
+            
+            # Check for silent audio
+            if max_val == 0:
+                print(f"[WARNING] Microphone audio appears to be completely silent!")
+            
+            temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".wav", dir=CUSTOM_TEMP_DIR)
+            temp_file.close()
+            audio_int16 = (audio_array * 32767.0).astype(np.int16)
+            sf.write(temp_file.name, audio_int16, sample_rate, subtype='PCM_16')
+            audio_path = temp_file.name
+            temp_file_to_cleanup = temp_file.name
+            print(f"[INFO] Saved microphone input to: {audio_path}")
+        
+        # Final check - if we still don't have audio_path, something went wrong
+        if audio_path is None:
+            yield "❌ Invalid audio input format.", "", None, None, None
+            return
+        
+        # Set temperature based on sampling mode
+        actual_temp = temperature if do_sample else 0.0
+        
+        # Start streaming transcription
+        print("[INFO] Starting API transcription (streaming mode)")
+        start_time = time.time()
+        
+        final_result = None
+        token_count = 0
+        accumulated_text = ""  # Track accumulated text for stop case
+        
+        async for text, result in api_client.transcribe_streaming(
+            audio_path=audio_path,
+            max_tokens=max_new_tokens,
+            temperature=actual_temp,
+            top_p=top_p,
+            context_info=context_info,
+        ):
+            # Track accumulated text
+            if text:
+                accumulated_text = text
+            
+            # Check stop flag at higher level too (already declared global at function start)
+            if stop_generation_flag:
+                print("[INFO] Stop flag detected in transcribe_audio, breaking...")
+                # Create a stopped result - parse whatever we have so far
+                stopped_segments, stopped_warning = api_client._parse_segments(accumulated_text) if accumulated_text else ([], None)
+                final_result = {
+                    "raw_text": accumulated_text or "",
+                    "segments": stopped_segments or [],
+                    "duration": 0,
+                    "usage": None,
+                    "stopped": True,
+                    "parse_warning": stopped_warning,
+                }
+                break
+            
+            if result is not None:
+                final_result = result
+            else:
+                # Streaming update - format for readability
+                token_count = len(text.split())  # Rough estimate
+                formatted_text = text.replace('},', '},\n')
+                yield f"🔄 Transcribing... ({token_count} tokens)\n---\n{formatted_text}", "", None, None, None
+        
+        generation_time = time.time() - start_time
+        
+        if final_result is None or "error" in final_result:
+            error_msg = final_result.get("error", "Unknown error") if final_result else "No response"
+            yield f"❌ Transcription failed: {error_msg}", "", None, None, None
+            return
+        
+        # Check if stopped by user
+        was_stopped = final_result.get('stopped', False)
+        
+        # Format final output with token statistics
+        if was_stopped:
+            raw_output = f"--- ⏹️ Transcription Stopped ---\n"
+        else:
+            raw_output = f"--- ✅ Transcription Complete ---\n"
+        raw_output += f"⏱️ Time: {generation_time:.2f}s | 🎵 Audio: {final_result.get('duration', 0):.2f}s\n"
+        
+        # Add token statistics if available
+        usage = final_result.get('usage')
+        if usage:
+            prompt_tokens = usage.get('prompt_tokens', 0)
+            completion_tokens = usage.get('completion_tokens', 0)
+            total_tokens = usage.get('total_tokens', 0)
+            tokens_per_sec = completion_tokens / generation_time if generation_time > 0 else 0
+            raw_output += f"📊 Tokens: {prompt_tokens} (prompt) + {completion_tokens} (completion) = {total_tokens} (total)\n"
+            raw_output += f"⚡ Speed: {tokens_per_sec:.1f} tokens/s\n"
+        
+        # Add parse warning if partial parsing was used
+        parse_warning = final_result.get('parse_warning')
+        if parse_warning:
+            raw_output += f"{parse_warning}\n"
+        
+        raw_output += f"---\n"
+        formatted_raw_text = final_result['raw_text'].replace('},', '},\n')
+        raw_output += formatted_raw_text
+        
+        # Generate audio segments HTML
+        segments = final_result.get('segments', [])
+        audio_segments_html = ""
+        num_segments = len(segments)
+        
+        if segments:
+            # Extract audio clips for each segment
+            audio_clips = extract_audio_segments(audio_path, segments)
+            
+            # Calculate approximate total size
+            total_duration = sum(
+                (float(seg.get('End', seg.get('end', 0))) - float(seg.get('Start', seg.get('start', 0))))
+                for seg in segments
+                if seg.get('End') is not None and seg.get('Start') is not None
+            )
+            approx_size_kb = total_duration * 4  # ~4KB per second at 32kbps
+            
+            # Add CSS for theme-aware styling (matching original demo)
+            theme_css = """
+            
+            """
+            
+            # Build HTML
+            format_info = "MP3 32kbps 16kHz mono" if HAS_PYDUB else "WAV 16kHz"
+            audio_segments_html = theme_css
+            audio_segments_html += "
"
+            audio_segments_html += f"
🔊 Audio Segments ({num_segments} segments)"
+            audio_segments_html += f"📦 ~{approx_size_kb:.0f}KB ({format_info})
"
+            audio_segments_html += "
🎵 Click the play button to listen to each segment directly!
"
+            
+            for i, seg in enumerate(segments):
+                start = seg.get('Start', seg.get('start', seg.get('Start time', None)))
+                end = seg.get('End', seg.get('end', seg.get('End time', None)))
+                speaker = seg.get('Speaker', seg.get('speaker', seg.get('Speaker ID', None)))
+                content = seg.get('Content', seg.get('content', seg.get('text', '')))
+                
+                # Format times nicely
+                start_str = f"{float(start):.2f}" if start is not None else "N/A"
+                end_str = f"{float(end):.2f}" if end is not None else "N/A"
+                speaker_str = str(speaker) if speaker is not None else "N/A"
+                
+                # Get audio clip
+                audio_html = ""
+                error_html = ""
+                if i < len(audio_clips):
+                    _, audio_src, error_msg = audio_clips[i]
+                    if audio_src:
+                        audio_type = 'audio/mp3' if 'audio/mp3' in audio_src or 'audio/mpeg' in audio_src else 'audio/wav'
+                        audio_html = f"""
+                        
+                        """
+                    elif error_msg:
+                        error_html = f"""
+                        

+                            ❌ {error_msg}
+                        

+                        """
+                
+                audio_segments_html += f"""
+                

+                    

+                        
🎤 Speaker {speaker_str}

+                        

+                            ⏱️ {start_str}s - {end_str}s
+                        

+                    

+                    

+                        {content}
+                    

+                    {audio_html}
+                    {error_html}
+                

+                """
+            
+            audio_segments_html += "
"
+        else:
+            audio_segments_html = """
+            
+            

+                
❌ No audio segments available.

+                
This could happen if the model output doesn't contain valid time stamps.

+            

+            """
+        
+        # Cleanup temp files
+        if temp_file_to_cleanup and os.path.exists(temp_file_to_cleanup):
+            try:
+                os.unlink(temp_file_to_cleanup)
+            except Exception:
+                pass
+        
+        if extracted_audio_to_cleanup and os.path.exists(extracted_audio_to_cleanup):
+            try:
+                os.unlink(extracted_audio_to_cleanup)
+            except Exception:
+                pass
+        
+        # Generate SRT and VTT content if we have segments
+        srt_content = None
+        vtt_content = None
+        if segments:
+            srt_content = segments_to_srt(segments)
+            vtt_content = segments_to_vtt(segments)
+        
+        yield raw_output, audio_segments_html, srt_content, original_video_path, vtt_content
+        
+    except Exception as e:
+        print(f"Error during transcription: {e}")
+        print(traceback.format_exc())
+        yield f"❌ Error: {str(e)}", "", None, None, None
+
+
+def create_gradio_interface(api_url: str, model_name: str = None, default_max_tokens: int = 4096, max_video_size_mb: float = DEFAULT_MAX_VIDEO_SIZE_MB):
+    """Create and launch Gradio interface."""
+    global api_client
+    
+    # Initialize API client
+    api_client = VibeVoiceAPIClient(api_url=api_url, model_name=model_name)
+    
+    # Check API health and fetch available models (sync for startup)
+    healthy, health_msg = api_client.check_health_sync()
+    available_models = []
+    if healthy:
+        available_models = api_client.get_available_models_sync()
+        if available_models:
+            # Auto-select first model if not specified
+            if not model_name:
+                print(f"🎯 Auto-selected model: {api_client.model_name}")
+            api_status = f"✅ Connected to API: {api_url} | Model: {api_client.model_name}"
+        else:
+            api_status = f"⚠️ Connected but no models found at: {api_url}"
+    else:
+        api_status = f"⚠️ API not available: {health_msg}"
+    print(api_status)
+    
+    # Custom CSS for button styling
+    custom_css = """
+    #transcribe-btn {
+        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
+        border: none !important;
+        color: white !important;
+    }
+    #transcribe-btn:hover {
+        background: linear-gradient(135deg, #764ba2 0%, #667eea 100%) !important;
+    }
+    #stop-btn {
+        background-color: #dc3545 !important;
+        border-color: #dc3545 !important;
+    }
+    #stop-btn:hover {
+        background-color: #c82333 !important;
+    }
+    /* Fix tab layout on small screens */
+    .tabs > .tab-nav {
+        flex-wrap: nowrap !important;
+        overflow-x: auto !important;
+    }
+    .tabs > .tab-nav > button {
+        white-space: nowrap !important;
+        flex-shrink: 0 !important;
+        font-size: 13px !important;
+        padding: 8px 12px !important;
+    }
+    """
+    
+    with gr.Blocks(title="VibeVoice ASR Demo") as demo:
+        gr.Markdown("# 🎙️ VibeVoice ASR Demo")
+        gr.Markdown("Upload audio/video files or record from microphone to get speech-to-text transcription with speaker diarization.")
+        
+        # Store max video size for use in transcribe function
+        max_video_size_state = gr.State(value=max_video_size_mb)
+        
+        # Hidden slider for max_tokens (use default value from args)
+        max_tokens_slider = gr.Slider(
+            minimum=512,
+            maximum=32768,
+            value=default_max_tokens,
+            step=512,
+            label="Max New Tokens",
+            visible=False
+        )
+        
+        # Define example files
+        # Look for demo files relative to repo root (/app in container)
+        _script_dir = os.path.dirname(os.path.abspath(__file__))
+        _repo_root = os.path.dirname(os.path.dirname(_script_dir))
+        example_dir = os.path.join(_repo_root, "demo", "asr_demo")
+        example_files = {
+            "chat_audio": os.path.join(example_dir, "demo1-chat.mp3"),
+            "chat_video": os.path.join(example_dir, "demo1-chat.mp4"),
+            "song_audio": os.path.join(example_dir, "demo2-song.mp3"),
+            "song_video": os.path.join(example_dir, "demo2-song.mp4"),
+            "hotword": os.path.join(example_dir, "demo3-hotwords.wav"),
+        }
+        
+        with gr.Row():
+            # Left column: Media Input (1/3)
+            with gr.Column(scale=1):
+                # Examples section
+                gr.Markdown("## 🎯 Examples")
+                with gr.Row():
+                    example1_btn = gr.Button("🗣️ Chat", size="sm", scale=1)
+                    example2_btn = gr.Button("🎵 Song", size="sm", scale=1)
+                    example3_btn = gr.Button("📝 Hotword", size="sm", scale=1)
+                
+                # Media input section (combined audio/video)
+                gr.Markdown("## 🎵 Media Input")
+                gr.Markdown(f"*Upload or record audio/video. For video (max {max_video_size_mb} MB), audio will be extracted.*")
+                
+                # Tabs for Upload File and Record
+                with gr.Tabs():
+                    with gr.TabItem("📁 Upload File"):
+                        media_input = gr.File(
+                            label="Upload Audio or Video File",
+                            file_types=list(COMMON_AUDIO_EXTS) + list(COMMON_VIDEO_EXTS),
+                            type="filepath"
+                        )
+                    with gr.TabItem("🎙️ Record Audio"):
+                        audio_record = gr.Audio(
+                            label="Record Audio",
+                            sources=["microphone"],
+                            type="filepath",
+                            interactive=True
+                        )
+                    with gr.TabItem("🎥 Record Video"):
+                        video_record = gr.Video(
+                            label="Record Video (auto-converts to 480p@30fps)",
+                            sources=["webcam"],
+                            include_audio=True,
+                            interactive=True
+                        )
+                
+                # Preview section - expanded by default
+                with gr.Accordion("👁️ Media Preview", open=True):
+                    audio_preview = gr.Audio(
+                        label="Audio Preview",
+                        interactive=False,
+                        visible=False
+                    )
+                    video_preview = gr.Video(
+                        label="Video Preview", 
+                        interactive=False,
+                        visible=False
+                    )
+            
+            # Right column: Context + Sampling + Results (2/3)
+            with gr.Column(scale=2):
+                # Context info and Sampling in one row
+                with gr.Row(equal_height=True):
+                    # Context information section
+                    with gr.Column(scale=1):
+                        gr.Markdown("## 📋 Customized Context")
+                        context_info_input = gr.Textbox(
+                            label="Add your customized terms below for better recognition. ",
+                            placeholder="VibeVoice \nMicrosoft \nAzure ... ",
+                            lines=5,
+                            max_lines=6,
+                            interactive=True,
+                        )
+                    
+                    # Sampling parameters - side by side with Hotwords
+                    with gr.Column(scale=1):
+                        gr.Markdown("## 🎲 Sampling")
+                        do_sample_checkbox = gr.Checkbox(
+                            value=False,
+                            label="Enable Sampling"
+                        )
+                        temperature_slider = gr.Slider(
+                            minimum=0.0,
+                            maximum=2.0,
+                            value=0.0,
+                            step=0.1,
+                            label="Temperature"
+                        )
+                        top_p_slider = gr.Slider(
+                            minimum=0.0,
+                            maximum=1.0,
+                            value=1.0,
+                            step=0.05,
+                            label="Top-p"
+                        )
+                
+                # Transcribe buttons
+                with gr.Row():
+                    transcribe_button = gr.Button("🎯 Transcribe", variant="primary", size="lg", scale=3, elem_id="transcribe-btn")
+                    stop_button = gr.Button("⏹️ Stop", variant="secondary", size="lg", scale=1, elem_id="stop-btn")
+                
+                # Results section
+                gr.Markdown("## 📝 Results")
+                
+                with gr.Tabs():
+                    with gr.TabItem("📝 Raw Output"):
+                        raw_output = gr.Textbox(
+                            label="Raw Transcription Output",
+                            lines=15,
+                            max_lines=30,
+                            interactive=False
+                        )
+                    
+                    with gr.TabItem("🔊 Audio Segments", visible=False) as audio_segments_tab:
+                        audio_segments_output = gr.HTML(
+                            label="Play individual segments to verify accuracy"
+                        )
+                    
+                    with gr.TabItem("🎬 Video with Subtitles", visible=False) as video_subs_tab:
+                        gr.Markdown("*Video playback with generated subtitles (only available for video input)*")
+                        video_with_subs_output = gr.HTML(
+                            label="Video Player with Subtitles"
+                        )
+                    
+                    with gr.TabItem("📥 Download Subtitles", visible=False) as download_subs_tab:
+                        gr.Markdown("*Download generated subtitles in SRT format*")
+                        srt_download = gr.File(
+                            label="Download SRT Subtitle File",
+                            interactive=False
+                        )
+        
+        # Event handlers
+        def async_copy_uploaded_file(file_path: str, save_dir: str = DEFAULT_UPLOAD_SAVE_DIR):
+            """Asynchronously copy uploaded file to save directory with unique filename."""
+            def _copy_file():
+                try:
+                    # Create save directory if not exists
+                    save_path = os.path.join(os.path.dirname(__file__), save_dir)
+                    os.makedirs(save_path, exist_ok=True)
+                    
+                    # Generate unique filename: timestamp_uuid_originalname
+                    original_name = os.path.basename(file_path)
+                    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+                    unique_id = str(uuid.uuid4())[:8]
+                    new_filename = f"{timestamp}_{unique_id}_{original_name}"
+                    dest_path = os.path.join(save_path, new_filename)
+                    
+                    # Copy file
+                    shutil.copy2(file_path, dest_path)
+                    print(f"[INFO] Uploaded file saved to: {dest_path}")
+                except Exception as e:
+                    print(f"[WARNING] Failed to save uploaded file: {e}")
+            
+            # Run copy in background thread
+            thread = threading.Thread(target=_copy_file, daemon=True)
+            thread.start()
+        
+        def update_media_preview(file_path):
+            """Update media preview based on uploaded file type. 
+            Compress video to 480p for preview.
+            Convert unsupported audio formats (opus, ogg, flac) to MP3.
+            """
+            if file_path is None:
+                # Don't clear preview when input is None (e.g., when cleared by example button)
+                return gr.update(), gr.update()
+            
+            # Async copy uploaded file to save directory
+            async_copy_uploaded_file(file_path)
+            
+            ext = os.path.splitext(file_path)[1].lower()
+            if ext in COMMON_VIDEO_EXTS:
+                # Video file - compress to 480p for preview, then show
+                print(f"[INFO] Compressing uploaded video for preview: {file_path}")
+                compressed_path, error = convert_video_to_mp4(file_path, height=480, crf=28, fps=30)
+                if compressed_path:
+                    print(f"[INFO] Video compressed for preview: {compressed_path}")
+                    return gr.update(value=None, visible=False), gr.update(value=compressed_path, visible=True)
+                else:
+                    # Fallback to original if compression fails
+                    print(f"[WARNING] Video compression failed: {error}, using original")
+                    return gr.update(value=None, visible=False), gr.update(value=file_path, visible=True)
+            elif ext in AUDIO_FORMATS_NEED_CONVERSION:
+                # Audio format needs conversion (opus, ogg, flac, etc.) - convert to MP3
+                print(f"[INFO] Converting {ext} audio to MP3 for better compatibility: {file_path}")
+                converted_path, error = convert_audio_to_mp3(file_path)
+                if converted_path:
+                    print(f"[INFO] Audio converted to MP3: {converted_path}")
+                    return gr.update(value=converted_path, visible=True), gr.update(value=None, visible=False)
+                else:
+                    # Fallback to original if conversion fails
+                    print(f"[WARNING] Audio conversion failed: {error}, using original")
+                    return gr.update(value=file_path, visible=True), gr.update(value=None, visible=False)
+            else:
+                # Audio file in supported format - show directly
+                return gr.update(value=file_path, visible=True), gr.update(value=None, visible=False)
+        
+        def update_audio_preview(audio_path):
+            """Update preview when audio is recorded."""
+            if audio_path is None:
+                # Don't clear preview when input is None (e.g., when cleared by example button)
+                return gr.update(), gr.update()
+            return gr.update(value=audio_path, visible=True), gr.update(value=None, visible=False)
+        
+        def update_video_preview(video_path):
+            """Update preview when video is recorded."""
+            if video_path is None:
+                # Don't clear preview when input is None (e.g., when cleared by example button)
+                return gr.update(), gr.update()
+            return gr.update(value=None, visible=False), gr.update(value=video_path, visible=True)
+        
+        # Update preview when file is uploaded
+        media_input.change(
+            fn=update_media_preview,
+            inputs=[media_input],
+            outputs=[audio_preview, video_preview]
+        )
+        
+        # Update preview when audio is recorded
+        audio_record.change(
+            fn=update_audio_preview,
+            inputs=[audio_record],
+            outputs=[audio_preview, video_preview]
+        )
+        
+        # Update preview when video is recorded
+        video_record.change(
+            fn=update_video_preview,
+            inputs=[video_record],
+            outputs=[audio_preview, video_preview]
+        )
+        
+        # Example button handlers - clear upload/record inputs when example is selected
+        def load_example_chat():
+            """Load chat example with video preview, clear other inputs."""
+            video_path = example_files["chat_video"]
+            if os.path.exists(video_path):
+                return (
+                    gr.update(value=None, visible=False),       # audio_preview
+                    gr.update(value=video_path, visible=True),  # video_preview
+                    "",                                          # context_info (no hotwords)
+                    gr.update(value=None),                       # media_input (clear)
+                    gr.update(value=None),                       # audio_record (clear)
+                    gr.update(value=None),                       # video_record (clear)
+                )
+            return gr.update(), gr.update(), "", gr.update(), gr.update(), gr.update()
+        
+        def load_example_song():
+            """Load song example with video preview, clear other inputs."""
+            video_path = example_files["song_video"]
+            if os.path.exists(video_path):
+                return (
+                    gr.update(value=None, visible=False),       # audio_preview
+                    gr.update(value=video_path, visible=True),  # video_preview
+                    "",                                          # context_info (no hotwords)
+                    gr.update(value=None),                       # media_input (clear)
+                    gr.update(value=None),                       # audio_record (clear)
+                    gr.update(value=None),                       # video_record (clear)
+                )
+            return gr.update(), gr.update(), "", gr.update(), gr.update(), gr.update()
+        
+        def load_example_hotword():
+            """Load hotword example with VibeVoice in context, clear other inputs."""
+            audio_path = example_files["hotword"]
+            if os.path.exists(audio_path):
+                return (
+                    gr.update(value=audio_path, visible=True),  # audio_preview
+                    gr.update(value=None, visible=False),       # video_preview
+                    "VibeVoice",                                 # context_info with hotword
+                    gr.update(value=None),                       # media_input (clear)
+                    gr.update(value=None),                       # audio_record (clear)
+                    gr.update(value=None),                       # video_record (clear)
+                )
+            return gr.update(), gr.update(), "VibeVoice", gr.update(), gr.update(), gr.update()
+        
+        example1_btn.click(
+            fn=load_example_chat,
+            inputs=[],
+            outputs=[audio_preview, video_preview, context_info_input, media_input, audio_record, video_record]
+        )
+        
+        example2_btn.click(
+            fn=load_example_song,
+            inputs=[],
+            outputs=[audio_preview, video_preview, context_info_input, media_input, audio_record, video_record]
+        )
+        
+        example3_btn.click(
+            fn=load_example_hotword,
+            inputs=[],
+            outputs=[audio_preview, video_preview, context_info_input, media_input, audio_record, video_record]
+        )
+        
+        def reset_stop_flag():
+            """Reset stop flag before starting transcription."""
+            global stop_generation_flag
+            stop_generation_flag = False
+            print("[INFO] Stop flag reset")
+        
+        def set_stop_flag():
+            """Set stop flag to interrupt generation."""
+            global stop_generation_flag
+            stop_generation_flag = True
+            print("[INFO] Stop flag set - stopping generation...")
+            return "⏹️ Stop requested, waiting for current chunk to complete..."
+        
+        def get_media_input(file_input, audio_rec, video_rec, audio_prev, video_prev):
+            """Get the media input from preview (which shows what will be transcribed).
+            
+            Priority: preview content (video_prev > audio_prev) since that's what user sees.
+            Recorded videos are automatically converted to 480p MP4 to reduce file size.
+            """
+            # Always use preview content - it shows what will be transcribed
+            if video_prev is not None:
+                # Check if it's a recorded video that needs conversion
+                if video_rec is not None and video_prev == video_rec:
+                    print(f"[INFO] Recorded video detected: {video_rec}")
+                    converted_path, error = convert_video_to_mp4(video_rec, height=480, crf=28, fps=30)
+                    if converted_path:
+                        print(f"[INFO] Recorded video converted to 480p@30fps: {converted_path}")
+                        return converted_path
+                    else:
+                        print(f"[WARNING] Failed to convert recorded video: {error}, using original")
+                return video_prev
+            if audio_prev is not None:
+                return audio_prev
+            return None
+        
+        async def transcribe_wrapper(
+            file_input, audio_rec, video_rec, audio_prev, video_prev, max_tokens, temp, top_p, do_sample, context_info, max_video_size
+        ):
+            """Wrapper to handle file/recording input and process results."""
+            media = get_media_input(file_input, audio_rec, video_rec, audio_prev, video_prev)
+            
+            video_html = ""
+            srt_file_path = None
+            
+            async for raw_text, segments_html, srt_content, video_path, vtt_content in transcribe_audio(
+                media, max_tokens, temp, top_p, do_sample, context_info, max_video_size
+            ):
+                # Generate video player HTML with subtitles if video was uploaded
+                if video_path and vtt_content:
+                    # Create a temp VTT file for the video player
+                    vtt_b64 = base64.b64encode(vtt_content.encode('utf-8')).decode('utf-8')
+                    vtt_data_url = f"data:text/vtt;base64,{vtt_b64}"
+                    
+                    # Read video file and create data URL
+                    with open(video_path, 'rb') as f:
+                        video_bytes = f.read()
+                    video_b64 = base64.b64encode(video_bytes).decode('utf-8')
+                    ext = os.path.splitext(video_path)[1].lower()
+                    mime_type = 'video/mp4' if ext == '.mp4' else f'video/{ext[1:]}'
+                    video_data_url = f"data:{mime_type};base64,{video_b64}"
+                    
+                    video_html = f'''
+                    
+                    

+                        
+                    

+                    '''
+                else:
+                    video_html = '''
+                    

+                        
🎬 No video input detected.

+                        
Upload a video file to see playback with subtitles.

+                    

+                    '''
+                
+                # Create SRT file for download if available
+                if srt_content:
+                    srt_temp = tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.srt', encoding='utf-8', dir=CUSTOM_TEMP_DIR)
+                    srt_temp.write(srt_content)
+                    srt_temp.close()
+                    srt_file_path = srt_temp.name
+                
+                # Determine tab visibility based on whether we have final results
+                # Show tabs only when we have actual content (not during streaming)
+                has_segments = segments_html and '
 1, launches N independent vLLM processes behind an nginx
+reverse proxy for optimal throughput (avoids single-process HTTP
+bottleneck of vLLM's built-in DP coordinator).
+
+Usage:
+    python3 start_server.py [--model MODEL_ID] [--port PORT]
+"""
+
+import argparse
+import os
+import signal
+import subprocess
+import sys
+import textwrap
+import time
+
+
+def run_command(cmd: list[str], description: str, shell: bool = False) -> None:
+    """Run a command with logging."""
+    print(f"\n{'='*60}")
+    print(f"  {description}")
+    print(f"{'='*60}\n")
+    if shell:
+        subprocess.run(cmd, shell=True, check=True)
+    else:
+        subprocess.run(cmd, check=True)
+
+
+def install_system_deps() -> None:
+    """Install system dependencies (FFmpeg, etc.)."""
+    run_command(["apt-get", "update"], "Updating package list")
+    run_command(
+        ["apt-get", "install", "-y", "ffmpeg", "libsndfile1"],
+        "Installing FFmpeg and audio libraries"
+    )
+
+
+def install_vibevoice() -> None:
+    """Install VibeVoice Python package."""
+    run_command(
+        [sys.executable, "-m", "pip", "install", "-e", "/app[vllm]"],
+        "Installing VibeVoice with vLLM support"
+    )
+
+
+def download_model(model_id: str) -> str:
+    """Download model from HuggingFace using default cache."""
+    print(f"\n{'='*60}")
+    print(f"  Downloading model: {model_id}")
+    print(f"{'='*60}\n")
+    
+    import warnings
+    from huggingface_hub import snapshot_download
+    
+    # Suppress deprecation warnings from huggingface_hub
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        model_path = snapshot_download(model_id)
+    
+    print(f"\n{'='*60}")
+    print(f"  ✅ Model downloaded successfully!")
+    print(f"  📁 Path: {model_path}")
+    print(f"{'='*60}\n")
+    return model_path
+
+
+def generate_tokenizer(model_path: str) -> None:
+    """Generate tokenizer files for the model."""
+    run_command(
+        [sys.executable, "-m", "vllm_plugin.tools.generate_tokenizer_files", 
+         "--output", model_path],
+        "Generating tokenizer files"
+    )
+
+
+def _build_vllm_cmd(model_path: str, port: int,
+                     tensor_parallel_size: int = 1,
+                     data_parallel_size: int = 1,
+                     max_num_seqs: int = 64,
+                     max_model_len: int = 65536,
+                     gpu_memory_utilization: float = 0.8) -> list[str]:
+    """Build the vllm serve command."""
+    return [
+        "vllm", "serve", model_path,
+        "--served-model-name", "vibevoice",
+        "--trust-remote-code",
+        "--dtype", "bfloat16",
+        "--max-num-seqs", str(max_num_seqs),
+        "--max-model-len", str(max_model_len),
+        "--gpu-memory-utilization", str(gpu_memory_utilization),
+        "--no-enable-prefix-caching",
+        "--enable-chunked-prefill",
+        "--chat-template-content-format", "openai",
+        "--tensor-parallel-size", str(tensor_parallel_size),
+        "--data-parallel-size", str(data_parallel_size),
+        "--allowed-local-media-path", "/app",
+        "--port", str(port),
+    ]
+
+
+def start_vllm_server(model_path: str, port: int,
+                      tensor_parallel_size: int = 1,
+                      data_parallel_size: int = 1,
+                      max_num_seqs: int = 64,
+                      max_model_len: int = 65536,
+                      gpu_memory_utilization: float = 0.8) -> None:
+    """Start a single vLLM server (replaces current process)."""
+    print(f"\n{'='*60}")
+    print(f"  Starting vLLM server on port {port}")
+    print(f"  Tensor Parallel (TP): {tensor_parallel_size}")
+    print(f"  Data Parallel   (DP): {data_parallel_size}")
+    print(f"  Max Num Seqs:         {max_num_seqs}")
+    print(f"  Max Model Len:        {max_model_len}")
+    print(f"  GPU Mem Utilization:  {gpu_memory_utilization}")
+    print(f"{'='*60}\n")
+    
+    vllm_cmd = _build_vllm_cmd(
+        model_path, port,
+        tensor_parallel_size=tensor_parallel_size,
+        data_parallel_size=data_parallel_size,
+        max_num_seqs=max_num_seqs,
+        max_model_len=max_model_len,
+        gpu_memory_utilization=gpu_memory_utilization,
+    )
+    os.execvp("vllm", vllm_cmd)
+
+
+def _install_nginx() -> None:
+    """Install nginx if not already available."""
+    if subprocess.run(["which", "nginx"], capture_output=True).returncode != 0:
+        run_command(["apt-get", "update"], "Updating package list for nginx")
+        run_command(
+            ["apt-get", "install", "-y", "nginx"],
+            "Installing nginx for load balancing"
+        )
+
+
+def _write_nginx_config(frontend_port: int, backend_ports: list[int],
+                        num_workers: int = 0) -> str:
+    """Write nginx config for round-robin load balancing.
+    
+    Args:
+        num_workers: Number of nginx worker processes. 0 = auto (2 × num backends).
+    """
+    if num_workers <= 0:
+        num_workers = len(backend_ports) * 2
+    backends = "\n".join(f"        server 127.0.0.1:{p};" for p in backend_ports)
+    config = textwrap.dedent(f"""\
+        worker_processes {num_workers};
+        worker_rlimit_nofile 65536;
+        error_log /dev/stderr warn;
+        pid /tmp/nginx.pid;
+
+        events {{
+            worker_connections 8192;
+        }}
+
+        http {{
+            access_log off;
+
+            upstream vllm_backends {{
+                least_conn;
+        {backends}
+            }}
+
+            server {{
+                listen {frontend_port};
+                client_max_body_size 200m;
+                client_body_buffer_size 10m;
+                proxy_buffering on;
+                proxy_buffer_size 64k;
+                proxy_buffers 16 64k;
+
+                location / {{
+                    proxy_pass http://vllm_backends;
+                    proxy_read_timeout 600s;
+                    proxy_connect_timeout 10s;
+                    proxy_send_timeout 600s;
+                    proxy_http_version 1.1;
+                    proxy_set_header Connection "";
+                }}
+            }}
+        }}
+    """)
+    config_path = "/tmp/nginx_vllm.conf"
+    with open(config_path, "w") as f:
+        f.write(config)
+    return config_path
+
+
+def start_dp_server(model_path: str, frontend_port: int,
+                    data_parallel_size: int,
+                    tensor_parallel_size: int = 1,
+                    max_num_seqs: int = 64,
+                    max_model_len: int = 65536,
+                    gpu_memory_utilization: float = 0.8) -> None:
+    """Start multiple vLLM workers behind nginx for data parallelism.
+    
+    Launches N independent vLLM processes (one per GPU group) on internal
+    ports, with an nginx reverse proxy on the frontend port for load
+    balancing. This avoids the single-process HTTP bottleneck of vLLM's
+    built-in DP coordinator when handling large audio payloads.
+    """
+    import torch
+    num_gpus = torch.cuda.device_count()
+    gpus_per_replica = tensor_parallel_size
+    total_gpus_needed = data_parallel_size * gpus_per_replica
+    assert num_gpus >= total_gpus_needed, (
+        f"Need {total_gpus_needed} GPUs (dp={data_parallel_size} × tp={tensor_parallel_size}) "
+        f"but only {num_gpus} available"
+    )
+
+    # Auto-tune per-worker env vars based on dp size
+    ffmpeg_concurrency = max(
+        64, int(os.environ.get("VIBEVOICE_FFMPEG_MAX_CONCURRENCY", "64"))
+    )
+    media_threads = max(
+        8, int(os.environ.get("VLLM_MEDIA_LOADING_THREAD_COUNT", "8"))
+    )
+
+    _install_nginx()
+
+    # Assign internal ports: frontend_port + 100, +101, ...
+    backend_ports = [frontend_port + 100 + i for i in range(data_parallel_size)]
+
+    print(f"\n{'='*60}")
+    print(f"  Starting DP server with nginx load balancing")
+    print(f"  Frontend port:     {frontend_port} (nginx)")
+    print(f"  Backend ports:     {backend_ports}")
+    print(f"  Data Parallel:     {data_parallel_size}")
+    print(f"  Tensor Parallel:   {tensor_parallel_size}")
+    print(f"  GPUs per replica:  {gpus_per_replica}")
+    print(f"  Max Num Seqs:      {max_num_seqs}")
+    print(f"  Max Model Len:     {max_model_len}")
+    print(f"  FFmpeg concurrency (per worker): {ffmpeg_concurrency}")
+    print(f"  Media loading threads (per worker): {media_threads}")
+    print(f"{'='*60}\n")
+
+    # Write nginx config
+    nginx_conf = _write_nginx_config(frontend_port, backend_ports)
+
+    # Launch vLLM workers
+    workers: list[subprocess.Popen] = []
+    for rank in range(data_parallel_size):
+        gpu_start = rank * gpus_per_replica
+        gpu_ids = ",".join(str(gpu_start + j) for j in range(gpus_per_replica))
+        port = backend_ports[rank]
+        
+        env = os.environ.copy()
+        env["CUDA_VISIBLE_DEVICES"] = gpu_ids
+        env["VIBEVOICE_FFMPEG_MAX_CONCURRENCY"] = str(ffmpeg_concurrency)
+        env["VLLM_MEDIA_LOADING_THREAD_COUNT"] = str(media_threads)
+
+        vllm_cmd = _build_vllm_cmd(
+            model_path, port,
+            tensor_parallel_size=tensor_parallel_size,
+            data_parallel_size=1,  # Each worker is dp=1
+            max_num_seqs=max_num_seqs,
+            max_model_len=max_model_len,
+            gpu_memory_utilization=gpu_memory_utilization,
+        )
+
+        print(f"  Launching worker rank={rank} on GPU(s) {gpu_ids}, port {port}")
+        proc = subprocess.Popen(vllm_cmd, env=env)
+        workers.append(proc)
+
+    # Start nginx
+    print(f"\n  Starting nginx on port {frontend_port} ...")
+    nginx_proc = subprocess.Popen(
+        ["nginx", "-c", nginx_conf, "-g", "daemon off;"]
+    )
+
+    # Wait for all backends to be ready
+    print("  Waiting for all backends to be ready ...")
+    import urllib.request
+    for port in backend_ports:
+        url = f"http://127.0.0.1:{port}/v1/models"
+        for attempt in range(600):  # up to 10 minutes
+            try:
+                urllib.request.urlopen(url, timeout=2)
+                print(f"    ✅ Backend on port {port} is ready")
+                break
+            except Exception:
+                time.sleep(1)
+        else:
+            print(f"    ❌ Backend on port {port} failed to start")
+
+    print(f"\n{'='*60}")
+    print(f"  ✅ VibeVoice DP server ready on port {frontend_port}")
+    print(f"     {data_parallel_size} replicas behind nginx load balancer")
+    print(f"{'='*60}\n")
+
+    # Handle shutdown: forward signals to all children
+    def _shutdown(signum, frame):
+        print("\nShutting down ...")
+        nginx_proc.terminate()
+        for w in workers:
+            w.terminate()
+        for w in workers:
+            w.wait(timeout=10)
+        nginx_proc.wait(timeout=5)
+        sys.exit(0)
+
+    signal.signal(signal.SIGTERM, _shutdown)
+    signal.signal(signal.SIGINT, _shutdown)
+
+    # Wait for any child to exit (indicates a failure)
+    while True:
+        for i, w in enumerate(workers):
+            ret = w.poll()
+            if ret is not None:
+                print(f"  ❌ Worker {i} exited with code {ret}")
+                _shutdown(None, None)
+        if nginx_proc.poll() is not None:
+            print(f"  ❌ nginx exited with code {nginx_proc.returncode}")
+            _shutdown(None, None)
+        time.sleep(1)
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="VibeVoice vLLM ASR Server - One-Click Deployment",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+    # Start with default settings (single GPU)
+    python3 start_server.py
+
+    # Use custom port
+    python3 start_server.py --port 8080
+
+    # Data parallel: 4 replicas on 4 GPUs (nginx load balancing)
+    python3 start_server.py --dp 4
+
+    # Tensor parallel: split model across 2 GPUs
+    python3 start_server.py --tp 2
+
+    # Skip dependency installation (if already installed)
+    python3 start_server.py --skip-deps
+        """
+    )
+    parser.add_argument(
+        "--model", "-m",
+        default="microsoft/VibeVoice-ASR",
+        help="HuggingFace model ID (default: microsoft/VibeVoice-ASR)"
+    )
+    parser.add_argument(
+        "--port", "-p",
+        type=int,
+        default=8000,
+        help="Server port (default: 8000)"
+    )
+    parser.add_argument(
+        "--skip-deps",
+        action="store_true",
+        help="Skip installing system dependencies"
+    )
+    parser.add_argument(
+        "--skip-tokenizer",
+        action="store_true",
+        help="Skip generating tokenizer files"
+    )
+    parser.add_argument(
+        "--tp", "--tensor-parallel-size",
+        type=int,
+        default=1,
+        dest="tensor_parallel_size",
+        help="Tensor parallel size: split one model across N GPUs (default: 1)"
+    )
+    parser.add_argument(
+        "--dp", "--data-parallel-size",
+        type=int,
+        default=1,
+        dest="data_parallel_size",
+        help="Data parallel size: run N independent model replicas for load balancing (default: 1)"
+    )
+    parser.add_argument(
+        "--max-num-seqs",
+        type=int,
+        default=64,
+        dest="max_num_seqs",
+        help="Maximum number of sequences per batch (default: 64)"
+    )
+    parser.add_argument(
+        "--max-model-len",
+        type=int,
+        default=65536,
+        dest="max_model_len",
+        help="Maximum model context length (default: 65536)"
+    )
+    parser.add_argument(
+        "--gpu-memory-utilization",
+        type=float,
+        default=0.8,
+        dest="gpu_memory_utilization",
+        help="GPU memory utilization fraction (default: 0.8)"
+    )
+    args = parser.parse_args()
+
+    print("\n" + "="*60)
+    print("  VibeVoice vLLM ASR Server - One-Click Deployment")
+    print("="*60)
+
+    # Step 1: Install system dependencies
+    if not args.skip_deps:
+        install_system_deps()
+
+    # Step 2: Install VibeVoice
+    install_vibevoice()
+
+    # Step 3: Download model
+    model_path = download_model(args.model)
+
+    # Step 4: Generate tokenizer files
+    if not args.skip_tokenizer:
+        generate_tokenizer(model_path)
+
+    # Step 5: Start server
+    if args.data_parallel_size > 1:
+        start_dp_server(
+            model_path, args.port,
+            data_parallel_size=args.data_parallel_size,
+            tensor_parallel_size=args.tensor_parallel_size,
+            max_num_seqs=args.max_num_seqs,
+            max_model_len=args.max_model_len,
+            gpu_memory_utilization=args.gpu_memory_utilization,
+        )
+    else:
+        start_vllm_server(
+            model_path, args.port,
+            tensor_parallel_size=args.tensor_parallel_size,
+            data_parallel_size=1,
+            max_num_seqs=args.max_num_seqs,
+            max_model_len=args.max_model_len,
+            gpu_memory_utilization=args.gpu_memory_utilization,
+        )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/vllm_plugin/tests/test_api.py b/vllm_plugin/tests/test_api.py
new file mode 100644
index 0000000..4076c20
--- /dev/null
+++ b/vllm_plugin/tests/test_api.py
@@ -0,0 +1,270 @@
+#!/usr/bin/env python3
+"""
+Test VibeVoice vLLM API with Streaming and Optional Hotwords Support.
+
+This script tests ASR transcription via the vLLM OpenAI-compatible API.
+By default, it runs standard transcription without hotwords.
+
+Optionally, you can provide hotwords (context_info) to improve recognition
+of domain-specific content like proper nouns, technical terms, and speaker names.
+Hotwords are embedded in the prompt as "with extra info: {hotwords}".
+
+Usage:
+    python test_api_with_hotwords.py [audio_path] [--url URL] [--hotwords "word1,word2"]
+    
+Examples:
+    # Standard transcription (no hotwords)
+    python3 test_api.py audio.wav
+    
+    # With hotwords for better recognition of specific terms
+    python3 test_api.py audio.wav --hotwords "Microsoft,Azure,VibeVoice"
+"""
+import requests
+import json
+import base64
+import time
+import sys
+import os
+import subprocess
+import argparse
+
+
+def _guess_mime_type(path: str) -> str:
+    """Guess MIME type from file extension."""
+    ext = os.path.splitext(path)[1].lower()
+    mime_map = {
+        ".wav": "audio/wav",
+        ".mp3": "audio/mpeg",
+        ".m4a": "audio/mp4",
+        ".mp4": "video/mp4",
+        ".flac": "audio/flac",
+        ".ogg": "audio/ogg",
+        ".opus": "audio/ogg",
+    }
+    return mime_map.get(ext, "application/octet-stream")
+
+
+def _get_duration_seconds_ffprobe(path: str) -> float:
+    """Get audio duration using ffprobe."""
+    cmd = [
+        "ffprobe", "-v", "error",
+        "-show_entries", "format=duration",
+        "-of", "default=noprint_wrappers=1:nokey=1",
+        path,
+    ]
+    out = subprocess.check_output(cmd, stderr=subprocess.STDOUT).decode("utf-8").strip()
+    return float(out)
+
+
+def _is_video_file(path: str) -> bool:
+    """Check if the file is a video file that needs audio extraction."""
+    ext = os.path.splitext(path)[1].lower()
+    return ext in (".mp4", ".m4v", ".mov", ".webm", ".avi", ".mkv")
+
+
+def _extract_audio_from_video(video_path: str) -> str:
+    """
+    Extract audio from video file (mp4/mov/webm) to a temporary mp3 file.
+    Returns the path to the extracted audio file.
+    """
+    import tempfile
+    # Create temp file with .mp3 extension
+    fd, audio_path = tempfile.mkstemp(suffix=".mp3")
+    os.close(fd)
+    
+    cmd = [
+        "ffmpeg", "-y", "-i", video_path,
+        "-vn",  # No video
+        "-acodec", "libmp3lame",
+        "-q:a", "2",  # High quality
+        audio_path
+    ]
+    subprocess.run(cmd, check=True, capture_output=True)
+    return audio_path
+
+
+def test_transcription_with_hotwords(
+    audio_path: str,
+    context_info: str = None,
+    base_url: str = "http://localhost:8000",
+):
+    """
+    Test ASR transcription with customized hotwords.
+    
+    Hotwords are embedded in the prompt text as "with extra info: {hotwords}".
+    This helps the model recognize domain-specific terms more accurately.
+    
+    Args:
+        audio_path: Path to the audio file
+        context_info: Hotwords string (e.g., "Microsoft,Azure,VibeVoice")
+        base_url: vLLM server URL
+    """
+    
+    print(f"=" * 70)
+    print(f"Testing Customized Hotwords Support")
+    print(f"=" * 70)
+    print(f"Input file: {audio_path}")
+    print(f"Hotwords: {context_info or '(none)'}")
+    print()
+    
+    # Handle video files: extract audio first
+    temp_audio_path = None
+    actual_audio_path = audio_path
+    if _is_video_file(audio_path):
+        print(f"🎬 Detected video file, extracting audio...")
+        temp_audio_path = _extract_audio_from_video(audio_path)
+        actual_audio_path = temp_audio_path
+        print(f"✅ Audio extracted to: {temp_audio_path}")
+    
+    # Load audio
+    try:
+        duration = _get_duration_seconds_ffprobe(actual_audio_path)
+        print(f"Audio duration: {duration:.2f} seconds")
+        
+        with open(actual_audio_path, "rb") as f:
+            audio_bytes = f.read()
+        
+        audio_b64 = base64.b64encode(audio_bytes).decode("utf-8")
+        print(f"Audio size: {len(audio_bytes)} bytes")
+        
+    except Exception as e:
+        print(f"Error preparing audio: {e}")
+        # Cleanup temp file if created
+        if temp_audio_path and os.path.exists(temp_audio_path):
+            os.remove(temp_audio_path)
+        return
+
+    # Build the request
+    url = f"{base_url}/v1/chat/completions"
+    
+    show_keys = ["Start time", "End time", "Speaker ID", "Content"]
+    
+    # Build prompt with optional hotwords
+    # Hotwords are embedded as "with extra info: {hotwords}" in the prompt
+    if context_info and context_info.strip():
+        prompt_text = (
+            f"This is a {duration:.2f} seconds audio, with extra info: {context_info.strip()}\n\n"
+            f"Please transcribe it with these keys: " + ", ".join(show_keys)
+        )
+        print(f"\n📝 Hotwords embedded in prompt: '{context_info}'")
+    else:
+        prompt_text = (
+            f"This is a {duration:.2f} seconds audio, please transcribe it with these keys: "
+            + ", ".join(show_keys)
+        )
+        print(f"\n📝 No hotwords provided")
+
+    mime = _guess_mime_type(actual_audio_path)
+    data_url = f"data:{mime};base64,{audio_b64}"
+
+    payload = {
+        "model": "vibevoice",
+        "messages": [
+            {
+                "role": "system",
+                "content": "You are a helpful assistant that transcribes audio input into text output in JSON format."
+            },
+            {
+                "role": "user",
+                "content": [
+                    {"type": "audio_url", "audio_url": {"url": data_url}},
+                    {"type": "text", "text": prompt_text}
+                ]
+            }
+        ],
+        "max_tokens": 32768,       
+        "temperature": 0.0,      
+        "stream": True,
+        "top_p": 1.0,
+    }
+    
+    print(f"\n{'=' * 70}")
+    print(f"Sending request to {url}")
+    print(f"{'=' * 70}")
+    
+    t0 = time.time()
+    try:
+        response = requests.post(url, json=payload, stream=True, timeout=12000)
+        
+        if response.status_code == 200:
+            print("\n✅ Response received. Streaming content:\n")
+            print("-" * 50)
+
+            printed = ""
+            for line in response.iter_lines():
+                if line:
+                    decoded_line = line.decode('utf-8')
+                    
+                    if decoded_line.startswith("data: "):
+                        json_str = decoded_line[6:] 
+                        if json_str.strip() == "[DONE]":
+                            print("\n" + "-" * 50)
+                            print("✅ [Finished]")
+                            break
+                        try:
+                            data = json.loads(json_str)
+                            delta = data['choices'][0]['delta']
+                            content = delta.get('content', '')
+                            if content:
+                                if content.startswith(printed):
+                                    to_print = content[len(printed):]
+                                else:
+                                    to_print = content
+                                if to_print:
+                                    print(to_print, end='', flush=True)
+                                    printed += to_print
+                        except json.JSONDecodeError:
+                            pass
+        else:
+            print(f"❌ Error: {response.status_code}")
+            print(response.text)
+            
+    except requests.exceptions.Timeout:
+        print("❌ Request timed out!")
+    except Exception as e:
+        print(f"❌ Error: {e}")
+        
+    elapsed = time.time() - t0
+    print(f"\n{'=' * 70}")
+    print(f"⏱️  Total time elapsed: {elapsed:.2f}s")
+    print(f"📊 RTF (Real-Time Factor): {elapsed / duration:.2f}x")
+    print(f"{'=' * 70}")
+    
+    # Cleanup temp audio file if created
+    if temp_audio_path and os.path.exists(temp_audio_path):
+        os.remove(temp_audio_path)
+        print(f"🗑️  Cleaned up temp file: {temp_audio_path}")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Test VibeVoice vLLM API with Customized Hotwords"
+    )
+    parser.add_argument(
+        "audio_path",
+        help="Path to audio file (wav, mp3, flac, etc.) or video file"
+    )
+    parser.add_argument(
+        "--url",
+        default="http://localhost:8000",
+        help="vLLM server URL (default: http://localhost:8000)"
+    )
+    parser.add_argument(
+        "--hotwords",
+        type=str,
+        default=None,
+        help="Hotwords to improve recognition (e.g., 'Microsoft,Azure,VibeVoice')"
+    )
+    
+    args = parser.parse_args()
+    
+    # Run test
+    test_transcription_with_hotwords(
+        audio_path=args.audio_path,
+        context_info=args.hotwords,
+        base_url=args.url,
+    )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/vllm_plugin/tests/test_api_auto_recover.py b/vllm_plugin/tests/test_api_auto_recover.py
new file mode 100644
index 0000000..482e258
--- /dev/null
+++ b/vllm_plugin/tests/test_api_auto_recover.py
@@ -0,0 +1,638 @@
+#!/usr/bin/env python3
+"""
+Test VibeVoice vLLM API with Streaming, Hotwords, and Auto-Recovery.
+
+This script tests ASR transcription with automatic recovery from repetition loops.
+Supports optional hotwords to improve recognition of domain-specific terms.
+
+Features:
+- Streaming output with real-time repetition detection
+- Auto-recovery when model enters repetition loops
+- Optional hotwords support (embedded in prompt as "with extra info: {hotwords}")
+- Video file support (auto-extracts audio)
+
+Recovery Strategy:
+1. First attempt: greedy decoding (temperature=0, top_p=1.0)
+2. If loop detected: retry with temperature=0.2/0.3/0.4, top_p=0.95
+3. Max 3 retries, truncate to last complete segment boundary
+
+Usage:
+    python test_api_auto_recover.py  [output_path] [--url URL] [--hotwords "word1,word2"] [--debug]
+
+Examples:
+    # Basic usage
+    python3 test_api_auto_recover.py audio.wav
+    
+    # With hotwords
+    python3 test_api_auto_recover.py audio.wav --hotwords "Microsoft,VibeVoice"
+    
+    # Save result to file
+    python3 test_api_auto_recover.py audio.wav result.txt
+    
+    # Debug mode (show recovery info)
+    python3 test_api_auto_recover.py audio.wav --debug
+"""
+import requests
+import json
+import base64
+import time
+import sys
+import os
+import subprocess
+import re
+import argparse
+from collections import Counter
+
+
+def _guess_mime_type(path: str) -> str:
+    ext = os.path.splitext(path)[1].lower()
+    if ext == ".wav":
+        return "audio/wav"
+    if ext in (".mp3",):
+        return "audio/mpeg"
+    if ext in (".m4a",):
+        return "audio/mp4"
+    if ext in (".mp4", ".m4v", ".mov", ".webm"):
+        return "video/mp4"
+    if ext in (".flac",):
+        return "audio/flac"
+    if ext in (".ogg", ".opus"):
+        return "audio/ogg"
+    return "application/octet-stream"
+
+
+def _get_duration_seconds_ffprobe(path: str) -> float:
+    cmd = [
+        "ffprobe", "-v", "error",
+        "-show_entries", "format=duration",
+        "-of", "default=noprint_wrappers=1:nokey=1",
+        path,
+    ]
+    out = subprocess.check_output(cmd, stderr=subprocess.STDOUT).decode("utf-8").strip()
+    return float(out)
+
+
+def _extract_audio_from_video(video_path: str) -> str:
+    """
+    Extract audio from video file (mp4/mov/webm) to a temporary mp3 file.
+    Returns the path to the extracted audio file.
+    """
+    import tempfile
+    # Create temp file with .mp3 extension
+    fd, audio_path = tempfile.mkstemp(suffix=".mp3")
+    os.close(fd)
+    
+    cmd = [
+        "ffmpeg", "-y", "-i", video_path,
+        "-vn",  # No video
+        "-acodec", "libmp3lame",
+        "-q:a", "2",  # High quality
+        audio_path
+    ]
+    subprocess.run(cmd, check=True, capture_output=True)
+    return audio_path
+
+
+def _is_video_file(path: str) -> bool:
+    """Check if the file is a video file that needs audio extraction."""
+    ext = os.path.splitext(path)[1].lower()
+    return ext in (".mp4", ".m4v", ".mov", ".webm", ".avi", ".mkv")
+
+
+def _find_last_segment_boundary(text: str) -> int:
+    """
+    Find the position after the last complete segment boundary (},).
+    Returns -1 if no complete segment found.
+    """
+    # Find last "}, " or "}," pattern (segment separator)
+    pos = text.rfind("},")
+    if pos != -1:
+        return pos + 2  # Include the },
+    return -1
+
+
+def _find_safe_print_boundary(text: str, max_pos: int) -> int:
+    """
+    Find the last complete segment boundary before max_pos.
+    Returns 0 if no complete segment found before max_pos.
+    """
+    search_text = text[:max_pos]
+    pos = search_text.rfind("},")
+    if pos != -1:
+        return pos + 2  # Include the },
+    return 0
+
+
+class RepetitionDetector:
+    """Detect repetition patterns in streaming text."""
+    
+    def __init__(self, 
+                 min_pattern_len: int = 10,      # Minimum chars for a pattern
+                 min_repeats: int = 3,           # Minimum repetitions to trigger
+                 window_size: int = 500):        # Window to check for patterns
+        self.min_pattern_len = min_pattern_len
+        self.min_repeats = min_repeats
+        self.window_size = window_size
+        self.text = ""
+        
+    def add_text(self, new_text: str):
+        """Add new text and return (is_looping, good_text_end_pos)."""
+        self.text += new_text
+        return self._check_repetition()
+    
+    def _check_repetition(self):
+        """Check if the recent text contains repetition loops."""
+        if len(self.text) < self.min_pattern_len * self.min_repeats:
+            return False, len(self.text)
+        
+        # Check the recent window
+        window = self.text[-self.window_size:] if len(self.text) > self.window_size else self.text
+        
+        # Method 1: Check for repeated substrings
+        for pattern_len in range(self.min_pattern_len, len(window) // self.min_repeats + 1):
+            # Get the last pattern_len characters as potential pattern
+            pattern = window[-pattern_len:]
+            
+            # Count how many times this pattern appears at the end
+            count = 0
+            pos = len(window)
+            while pos >= pattern_len:
+                if window[pos - pattern_len:pos] == pattern:
+                    count += 1
+                    pos -= pattern_len
+                else:
+                    break
+            
+            if count >= self.min_repeats:
+                # Found repetition! Calculate where the good text ends
+                repetition_start = len(self.text) - (count * pattern_len)
+                # Keep one instance of the pattern (or none if it's garbage)
+                good_end = repetition_start + pattern_len if self._is_meaningful(pattern) else repetition_start
+                return True, good_end
+        
+        # Method 2: Check for repeated short phrases (like "you're not, you're not")
+        # Look for patterns like "X, X, X" or "X X X"
+        words = window.split()
+        if len(words) >= self.min_repeats * 2:
+            # Check last N words for repetition
+            for phrase_len in range(2, 6):  # 2-5 word phrases
+                if len(words) < phrase_len * self.min_repeats:
+                    continue
+                
+                phrase = " ".join(words[-phrase_len:])
+                count = 0
+                idx = len(words)
+                while idx >= phrase_len:
+                    candidate = " ".join(words[idx - phrase_len:idx])
+                    if candidate == phrase:
+                        count += 1
+                        idx -= phrase_len
+                    else:
+                        break
+                
+                if count >= self.min_repeats:
+                    # Calculate position in original text
+                    repeated_text = (phrase + " ") * count
+                    good_end = len(self.text) - len(repeated_text.rstrip()) + len(phrase)
+                    return True, max(0, good_end)
+        
+        return False, len(self.text)
+    
+    def _is_meaningful(self, pattern: str) -> bool:
+        """Check if pattern is meaningful content (not just garbage)."""
+        # Filter out patterns that are just punctuation, spaces, or very repetitive
+        clean = pattern.strip()
+        if not clean:
+            return False
+        if len(set(clean)) < 3:  # Too few unique characters
+            return False
+        return True
+    
+    def get_good_text(self, end_pos: int) -> str:
+        """Get text up to the specified position."""
+        return self.text[:end_pos]
+    
+    def reset(self, keep_text: str = ""):
+        """Reset detector, optionally keeping some text."""
+        self.text = keep_text
+
+
+def stream_with_recovery(
+    url: str,
+    base_messages: list,
+    audio_data_url: str,
+    prompt_text: str,
+    max_tokens: int = 32768,
+    max_retries: int = 3,
+    timeout: int = 12000,
+    debug: bool = False,
+):
+    """
+    Stream transcription with automatic recovery from repetition loops.
+    
+    Args:
+        url: API endpoint
+        base_messages: Base messages (system + user with audio)
+        audio_data_url: The audio data URL for the request
+        prompt_text: The text prompt
+        max_tokens: Maximum tokens to generate
+        max_retries: Maximum recovery attempts (default 3)
+        timeout: Request timeout
+        debug: If True, show recovery debug info to stderr
+    
+    Recovery strategy:
+        - First attempt: temperature=0.0, top_p=1.0 (greedy)
+        - Recovery: temperature=0.2/0.3/0.4 for retry 1/2/3, top_p=0.95
+        - If has complete segments: use assistant prefix
+        - If no complete segments: restart from scratch
+        - Max 3 retries, if all fail output error message
+    
+    Returns:
+        Final transcription text
+    """
+    import sys as _sys
+    
+    def _log(msg):
+        """Log to stderr only if debug."""
+        if debug:
+            print(msg, file=_sys.stderr)
+    
+    detector = RepetitionDetector(
+        min_pattern_len=10,   # At least 10 chars for a pattern
+        min_repeats=10,       # Must repeat 10+ times
+        window_size=400,      # Check last 400 chars (can detect 10-40 char patterns repeated 10 times)
+    )
+    
+    accumulated_text = ""
+    retry_count = 0
+    user_safe_printed_len = 0  # Track how much we've safely shown to user (at segment boundaries)
+    is_recovery = False  # Whether we're in recovery mode
+    
+    while retry_count <= max_retries:
+        # Build request payload
+        messages = list(base_messages)  # Copy base messages
+        
+        # If we have accumulated text from previous attempt, add it as partial assistant response
+        if accumulated_text:
+            # Add the good content as a partial assistant message
+            # vLLM will continue from here
+            messages.append({
+                "role": "assistant",
+                "content": accumulated_text
+            })
+        
+        # Set sampling parameters based on recovery state
+        if is_recovery:
+            # Recovery: increase temperature each retry to break loops
+            recovery_temp = 0.1 + 0.1 * retry_count  # 0.2, 0.3, 0.4 for retry 1, 2, 3
+            payload = {
+                "model": "vibevoice",
+                "messages": messages,
+                "max_tokens": max_tokens,
+                "temperature": recovery_temp,
+                "top_p": 0.95,
+                "stream": True,
+            }
+            if accumulated_text:
+                _log(f"[RECOVERY #{retry_count}] Continuing from {len(accumulated_text)} chars with temp={recovery_temp}, top_p=0.95")
+            else:
+                _log(f"[RECOVERY #{retry_count}] Restarting from scratch with temp={recovery_temp}, top_p=0.95")
+        else:
+            # First attempt: greedy decoding
+            payload = {
+                "model": "vibevoice",
+                "messages": messages,
+                "max_tokens": max_tokens,
+                "temperature": 0.0,
+                "top_p": 1.0,
+                "stream": True,
+            }
+        
+        try:
+            response = requests.post(url, json=payload, stream=True, timeout=timeout)
+            
+            if response.status_code != 200:
+                _log(f"[ERROR] {response.status_code} - {response.text[:500]}")
+                return accumulated_text
+            
+            new_text = ""
+            printed = ""  # Track what we've already received to handle vLLM duplicates
+            
+            for line in response.iter_lines():
+                if not line:
+                    continue
+                
+                decoded_line = line.decode('utf-8')
+                if not decoded_line.startswith("data: "):
+                    continue
+                
+                json_str = decoded_line[6:]
+                if json_str.strip() == "[DONE]":
+                    # Successfully finished without loops
+                    full_result = accumulated_text + new_text
+                    # Print any remaining content that wasn't printed yet
+                    if len(full_result) > user_safe_printed_len:
+                        remaining = full_result[user_safe_printed_len:]
+                        print(remaining, end='', flush=True)
+                    print()  # Final newline
+                    return full_result
+                
+                try:
+                    data = json.loads(json_str)
+                    delta = data['choices'][0].get('delta', {})
+                    content = delta.get('content', '')
+                    
+                    if content:
+                        # vLLM/OpenAI-compatible streams may emit either
+                        # incremental deltas OR the full accumulated text.
+                        # Only track the newly-added part.
+                        if content.startswith(printed):
+                            to_add = content[len(printed):]
+                        else:
+                            to_add = content
+                        
+                        if to_add:
+                            printed += to_add
+                            new_text += to_add
+                            
+                            # When continuing from prefix, model may add "[" or "[{" at start
+                            # or repeat the ending "}, " from prefix
+                            # We need to handle these to maintain valid JSON array format
+                            if accumulated_text and new_text:
+                                stripped = new_text.lstrip()
+                                # Case 1: Model added "[{" - remove the "["
+                                if stripped.startswith("[{"):
+                                    new_text = stripped[1:]
+                                    _log("[STRIPPED leading '[' from continuation]")
+                                # Case 2: Model added just "[" - remove it
+                                elif stripped.startswith("["):
+                                    new_text = stripped[1:]
+                                    _log("[STRIPPED leading '[' from continuation]")
+                                # Case 3: Model repeated "}," from prefix ending
+                                elif stripped.startswith("},"):
+                                    new_text = stripped[2:]
+                                    _log("[STRIPPED leading '},' from continuation]")
+                                # Case 4: Model repeated "}" from prefix ending
+                                elif stripped.startswith("}") and not stripped.startswith("}]"):
+                                    new_text = stripped[1:]
+                                    _log("[STRIPPED leading '}' from continuation]")
+                                
+                                # Fix malformed JSON: {"2.99,... -> {"Start":2.99,...
+                                # This happens when model skips "Start": key
+                                import re
+                                malformed = re.match(r'^\{"(\d+\.?\d*),', new_text)
+                                if malformed:
+                                    time_val = malformed.group(1)
+                                    new_text = '{"Start":' + time_val + ',' + new_text[malformed.end():]
+                                    _log(f"[FIXED malformed JSON: added Start key]")
+                            
+                            # Check for repetition in the combined text
+                            full_text = accumulated_text + new_text
+                            detector.text = full_text
+                            is_looping, good_end = detector._check_repetition()
+                            
+                            if is_looping:
+                                _log(f"[LOOP DETECTED at char {good_end}]")
+                                
+                                # Use what user has already seen as prefix for retry
+                                # user_safe_printed_len is always at a segment boundary
+                                if user_safe_printed_len > 0:
+                                    accumulated_text = full_text[:user_safe_printed_len]
+                                    _log(f"[RETRY from user-visible content at {user_safe_printed_len}]")
+                                else:
+                                    # No complete segment shown to user yet - restart from scratch
+                                    accumulated_text = ""
+                                    _log(f"[NO CONTENT SHOWN TO USER - restart from scratch]")
+                                
+                                detector.reset(accumulated_text)
+                                is_recovery = True
+                                
+                                if debug:
+                                    print("\n[...recovering...]", end='', flush=True, file=sys.stderr)
+                                
+                                retry_count += 1
+                                
+                                if retry_count > max_retries:
+                                    _log(f"[MAX RETRIES REACHED]")
+                                    print("\n[Error] Transcription failed due to model output anomaly. Please try another audio or contact support.", flush=True)
+                                    return None
+                                
+                                # Break inner loop to retry
+                                break
+                            else:
+                                # No loop detected - stream content to user
+                                # Only print up to (full_text_len - window_size) at segment boundaries
+                                # This ensures user never sees content that might be rolled back
+                                safe_end = max(0, len(full_text) - detector.window_size)
+                                safe_boundary = _find_safe_print_boundary(full_text, safe_end)
+                                
+                                if safe_boundary > user_safe_printed_len:
+                                    # Print new safe content
+                                    to_print = full_text[user_safe_printed_len:safe_boundary]
+                                    print(to_print, end='', flush=True)
+                                    user_safe_printed_len = safe_boundary
+                                
+                except json.JSONDecodeError:
+                    continue
+            else:
+                # Loop completed without break (no repetition detected)
+                full_result = accumulated_text + new_text
+                
+                # Print any remaining content that wasn't printed yet
+                if len(full_result) > user_safe_printed_len:
+                    remaining = full_result[user_safe_printed_len:]
+                    print(remaining, end='', flush=True)
+                
+                print()  # Final newline
+                return full_result
+                
+        except requests.exceptions.Timeout:
+            _log("[TIMEOUT]")
+            print()
+            return accumulated_text
+        except Exception as e:
+            _log(f"[ERROR: {e}]")
+            print()
+            return accumulated_text
+    
+    # All retries exhausted
+    print("\n[Error] Transcription failed due to model output anomaly. Please try another audio or contact support.", flush=True)
+    return None
+
+
+def test_transcription_with_recovery(
+    audio_path: str,
+    output_path: str = None,
+    base_url: str = "http://localhost:8000",
+    hotwords: str = None,
+    debug: bool = False,
+):
+    """
+    Test ASR transcription with auto-recovery from repetition loops.
+    
+    Args:
+        audio_path: Path to the audio file
+        output_path: Optional path to save transcription result
+        base_url: vLLM server URL
+        hotwords: Hotwords string (e.g., "Microsoft,Azure,VibeVoice")
+        debug: Show recovery debug info
+    """
+    
+    print(f"=" * 70)
+    print(f"Testing with Auto-Recovery")
+    print(f"=" * 70)
+    print(f"Input file: {audio_path}")
+    print(f"Hotwords: {hotwords or '(none)'}")
+    print()
+    
+    # Handle video files: extract audio first
+    temp_audio_path = None
+    actual_audio_path = audio_path
+    if _is_video_file(audio_path):
+        print(f"🎬 Detected video file, extracting audio...")
+        temp_audio_path = _extract_audio_from_video(audio_path)
+        actual_audio_path = temp_audio_path
+        print(f"✅ Audio extracted to: {temp_audio_path}")
+    
+    # Load audio
+    try:
+        duration = _get_duration_seconds_ffprobe(actual_audio_path)
+        print(f"Audio duration: {duration:.2f} seconds")
+        
+        with open(actual_audio_path, "rb") as f:
+            audio_bytes = f.read()
+        
+        audio_b64 = base64.b64encode(audio_bytes).decode("utf-8")
+        print(f"Audio size: {len(audio_bytes)} bytes")
+        
+    except Exception as e:
+        print(f"❌ Error preparing audio: {e}")
+        # Cleanup temp file if created
+        if temp_audio_path and os.path.exists(temp_audio_path):
+            os.remove(temp_audio_path)
+        return
+
+    url = f"{base_url}/v1/chat/completions"
+    
+    show_keys = ["Start time", "End time", "Speaker ID", "Content"]
+    
+    # Build prompt with optional hotwords
+    if hotwords and hotwords.strip():
+        prompt_text = (
+            f"This is a {duration:.2f} seconds audio, with extra info: {hotwords.strip()}\n\n"
+            f"Please transcribe it with these keys: " + ", ".join(show_keys)
+        )
+        print(f"\n📝 Hotwords embedded in prompt: '{hotwords}'")
+    else:
+        prompt_text = (
+            f"This is a {duration:.2f} seconds audio, please transcribe it with these keys: "
+            + ", ".join(show_keys)
+        )
+        print(f"\n📝 No hotwords provided")
+
+    mime = _guess_mime_type(actual_audio_path)
+    data_url = f"data:{mime};base64,{audio_b64}"
+
+    # Base messages (without assistant continuation)
+    base_messages = [
+        {
+            "role": "system",
+            "content": "You are a helpful assistant that transcribes audio input into text output in JSON format."
+        },
+        {
+            "role": "user",
+            "content": [
+                {"type": "audio_url", "audio_url": {"url": data_url}},
+                {"type": "text", "text": prompt_text}
+            ]
+        }
+    ]
+    
+    print(f"\n{'=' * 70}")
+    print(f"Sending request to {url}")
+    print(f"{'=' * 70}")
+    
+    t0 = time.time()
+    print("\n✅ Response received. Streaming content:\n")
+    print("-" * 50)
+    
+    result = stream_with_recovery(
+        url=url,
+        base_messages=base_messages,
+        audio_data_url=data_url,
+        prompt_text=prompt_text,
+        max_tokens=32768,
+        max_retries=3,
+        debug=debug,
+    )
+    
+    elapsed = time.time() - t0
+    print("-" * 50)
+    print("✅ [Finished]")
+    print(f"\n{'=' * 70}")
+    print(f"⏱️  Total time elapsed: {elapsed:.2f}s")
+    print(f"{'=' * 70}")
+    
+    if result is None:
+        print("❌ Transcription failed")
+        return
+    
+    print(f"📄 Final output length: {len(result)} chars")
+    
+    # Optionally save result
+    if output_path:
+        with open(output_path, "w", encoding="utf-8") as f:
+            f.write(result)
+        print(f"💾 Result saved to: {output_path}")
+    
+    # Cleanup temp audio file if created
+    if temp_audio_path and os.path.exists(temp_audio_path):
+        os.remove(temp_audio_path)
+        print(f"🗑️  Cleaned up temp file: {temp_audio_path}")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Test VibeVoice vLLM API with auto-recovery from repetition loops"
+    )
+    parser.add_argument(
+        "audio_path",
+        help="Path to audio file (wav, mp3, flac, etc.) or video file"
+    )
+    parser.add_argument(
+        "output_path",
+        nargs="?",
+        default=None,
+        help="Optional path to save transcription result"
+    )
+    parser.add_argument(
+        "--url",
+        default="http://localhost:8000",
+        help="vLLM server URL (default: http://localhost:8000)"
+    )
+    parser.add_argument(
+        "--hotwords",
+        type=str,
+        default=None,
+        help="Hotwords to improve recognition (e.g., 'Microsoft,Azure,VibeVoice')"
+    )
+    parser.add_argument(
+        "--debug",
+        action="store_true",
+        help="Show recovery debug info"
+    )
+    
+    args = parser.parse_args()
+    
+    # Run test
+    test_transcription_with_recovery(
+        audio_path=args.audio_path,
+        output_path=args.output_path,
+        base_url=args.url,
+        hotwords=args.hotwords,
+        debug=args.debug,
+    )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/vllm_plugin/tools/generate_tokenizer_files.py b/vllm_plugin/tools/generate_tokenizer_files.py
new file mode 100644
index 0000000..62de55f
--- /dev/null
+++ b/vllm_plugin/tools/generate_tokenizer_files.py
@@ -0,0 +1,575 @@
+#!/usr/bin/env python3
+"""
+Standalone tool to generate VibeVoice tokenizer files from Qwen2 base.
+
+Downloads base tokenizer from Qwen2 and patches it with VibeVoice-specific
+audio tokens and chat template modifications.
+
+Usage:
+    python generate_tokenizer_files.py --output /path/to/output [--compare /path/to/reference]
+"""
+
+import argparse
+import json
+import os
+import shutil
+import tempfile
+from typing import Optional, Dict, Any
+
+
+# Qwen2.5 extended tokens (151646-151664)
+# These are NOT in base Qwen2-7B but ARE in Qwen2.5 and Qwen2-VL
+# VibeVoice uses some of these for speech: object_ref_start/end, box_start
+QWEN25_EXTENDED_TOKENS = {
+    "<|object_ref_start|>": 151646,  # Used as speech_start_id
+    "<|object_ref_end|>": 151647,    # Used as speech_end_id
+    "<|box_start|>": 151648,         # Used as speech_pad_id
+    "<|box_end|>": 151649,
+    "<|quad_start|>": 151650,
+    "<|quad_end|>": 151651,
+    "<|vision_start|>": 151652,
+    "<|vision_end|>": 151653,
+    "<|vision_pad|>": 151654,
+    "<|image_pad|>": 151655,
+    "<|video_pad|>": 151656,
+    "": 151657,
+    "": 151658,
+    "<|fim_prefix|>": 151659,
+    "<|fim_middle|>": 151660,
+    "<|fim_suffix|>": 151661,
+    "<|fim_pad|>": 151662,
+    "<|repo_name|>": 151663,
+    "<|file_sep|>": 151664,
+}
+
+# VibeVoice-specific audio tokens (IDs follow Qwen2.5's last token 151664)
+VIBEVOICE_AUDIO_TOKENS = {
+    "<|AUDIO|>": 151665,
+    "<|audio_bos|>": 151666,
+    "<|audio_eos|>": 151667,
+}
+
+# All extended tokens (Qwen2.5 + VibeVoice)
+ALL_EXTENDED_TOKENS = {**QWEN25_EXTENDED_TOKENS, **VIBEVOICE_AUDIO_TOKENS}
+
+# Chat template with audio support
+# Key modification: handles part['type'] == 'audio' or 'audio_url' -> '<|AUDIO|>'
+VIBEVOICE_CHAT_TEMPLATE = """{%- if tools %}
+    {{- '<|im_start|>system\\n' }}
+    {%- if messages[0]['role'] == 'system' %}
+        {%- if messages[0]['content'] is string %}
+            {{- messages[0]['content'] }}
+        {%- else %}
+            {%- for part in messages[0]['content'] %}
+                {%- if part['type'] == 'text' %}
+                    {{- part['text'] }}
+                {%- elif part['type'] == 'audio' or part['type'] == 'audio_url' %}
+                    {{- '<|AUDIO|>' }}
+                {%- endif %}
+            {%- endfor %}
+        {%- endif %}
+    {%- else %}
+        {{- 'You are a helpful assistant.' }}
+    {%- endif %}
+    {{- "\\n\\n# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within  XML tags:\\n" }}
+    {%- for tool in tools %}
+        {{- "\\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\\n\\n\\nFor each function call, return a json object with function name and arguments within  XML tags:\\n\\n{\\"name\\": , \\"arguments\\": }\\n<|im_end|>\\n" }}
+{%- else %}
+    {%- if messages[0]['role'] == 'system' %}
+        {{- '<|im_start|>system\\n' }}
+        {%- if messages[0]['content'] is string %}
+            {{- messages[0]['content'] }}
+        {%- else %}
+            {%- for part in messages[0]['content'] %}
+                {%- if part['type'] == 'text' %}
+                    {{- part['text'] }}
+                {%- elif part['type'] == 'audio' or part['type'] == 'audio_url' %}
+                    {{- '<|AUDIO|>' }}
+                {%- endif %}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>\\n' }}
+    {%- else %}
+        {{- '<|im_start|>system\\nYou are a helpful assistant.<|im_end|>\\n' }}
+    {%- endif %}
+{%- endif %}
+{%- for message in messages %}
+    {%- if (message.role == "user") or (message.role == "system" and not loop.first) or (message.role == "assistant" and not message.tool_calls) %}
+        {{- '<|im_start|>' + message.role + '\\n' }}
+        {%- if message['content'] is string %}
+            {{- message['content'] }}
+        {%- else %}
+            {%- for part in message['content'] %}
+                {%- if part['type'] == 'text' %}
+                    {{- part['text'] }}
+                {%- elif part['type'] == 'audio' or part['type'] == 'audio_url' %}
+                    {{- '<|AUDIO|>' }}
+                {%- endif %}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>\\n' }}
+    {%- elif message.role == "assistant" %}
+        {{- '<|im_start|>' + message.role }}
+        {%- if message.content %}
+            {{- '\\n' + message.content }}
+        {%- endif %}
+        {%- for tool_call in message.tool_calls %}
+            {%- if tool_call.function is defined %}
+                {%- set tool_call = tool_call.function %}
+            {%- endif %}
+            {{- '\\n\\n{"name": "' }}
+            {{- tool_call.name }}
+            {{- '", "arguments": ' }}
+            {{- tool_call.arguments | tojson }}
+            {{- '}\\n' }}
+        {%- endfor %}
+        {{- '<|im_end|>\\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != "tool") %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\\n\\n' }}
+        {{- message.content }}
+        {{- '\\n' }}
+        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+            {{- '<|im_end|>\\n' }}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\\n' }}
+{%- endif %}"""
+
+
+# Default to Qwen2.5-7B which has all the extended tokens (151646-151664)
+DEFAULT_QWEN_MODEL = "Qwen/Qwen2.5-7B"
+
+
+def download_qwen_tokenizer_files(output_dir: str, qwen_model: str = DEFAULT_QWEN_MODEL) -> None:
+    """Download base tokenizer files from Qwen2.5 (which includes extended tokens)."""
+    try:
+        from huggingface_hub import hf_hub_download
+    except ImportError:
+        raise ImportError("Please install huggingface_hub: pip install huggingface_hub")
+    
+    files_to_download = [
+        "vocab.json",
+        "merges.txt",
+        "tokenizer.json",
+        "tokenizer_config.json",
+    ]
+    
+    os.makedirs(output_dir, exist_ok=True)
+    
+    for filename in files_to_download:
+        print(f"Downloading {filename} from {qwen_model}...")
+        hf_hub_download(
+            repo_id=qwen_model,
+            filename=filename,
+            local_dir=output_dir,
+            local_dir_use_symlinks=False,
+        )
+
+
+def patch_tokenizer_config(output_dir: str) -> None:
+    """
+    Patch tokenizer_config.json with VibeVoice audio tokens and chat template.
+    """
+    config_path = os.path.join(output_dir, "tokenizer_config.json")
+    
+    with open(config_path, "r", encoding="utf-8") as f:
+        config = json.load(f)
+    
+    # 1. Add ALL extended tokens to added_tokens_decoder (Qwen2.5 + VibeVoice audio)
+    if "added_tokens_decoder" not in config:
+        config["added_tokens_decoder"] = {}
+    
+    for token, token_id in ALL_EXTENDED_TOKENS.items():
+        if str(token_id) not in config["added_tokens_decoder"]:
+            # Determine if token should be marked as "special" 
+            # tool_call tokens are NOT special in Qwen2.5
+            is_special = token not in ("", "", "<|fim_prefix|>", 
+                                       "<|fim_middle|>", "<|fim_suffix|>", "<|fim_pad|>",
+                                       "<|repo_name|>", "<|file_sep|>")
+            config["added_tokens_decoder"][str(token_id)] = {
+                "content": token,
+                "lstrip": False,
+                "normalized": False,
+                "rstrip": False,
+                "single_word": False,
+                "special": is_special,
+            }
+    
+    # 2. Add audio tokens to additional_special_tokens
+    if "additional_special_tokens" not in config:
+        config["additional_special_tokens"] = []
+    
+    for token in VIBEVOICE_AUDIO_TOKENS.keys():
+        if token not in config["additional_special_tokens"]:
+            config["additional_special_tokens"].append(token)
+    
+    # 3. Modify chat_template to support audio
+    # Instead of replacing entirely, we patch the existing template to handle audio
+    chat_template = config.get("chat_template", "")
+    if chat_template and "<|AUDIO|>" not in chat_template:
+        # Insert audio handling into the template
+        # Find patterns like: {%- if part['type'] == 'text' %}
+        # Add after: {%- elif part['type'] == 'audio' or part['type'] == 'audio_url' %}\n                    {{- '<|AUDIO|>' }}
+        audio_handler = """{%- elif part['type'] == 'audio' or part['type'] == 'audio_url' %}
+                    {{- '<|AUDIO|>' }}"""
+        
+        # Pattern to find: after handling 'text' type, before endif
+        import re
+        # Look for the pattern where we handle text type and add audio handling
+        pattern = r"(\{\%- if part\['type'\] == 'text' \%\}\s*\n\s*\{\{- part\['text'\] \}\})"
+        replacement = r"\1\n                " + audio_handler.replace("\n", r"\n")
+        
+        modified_template = re.sub(pattern, replacement, chat_template)
+        
+        if modified_template != chat_template:
+            config["chat_template"] = modified_template
+            print("  - Added audio support to existing chat_template")
+        else:
+            # Fallback: use our predefined template
+            print("  - Warning: Could not patch existing template, using predefined template")
+            config["chat_template"] = VIBEVOICE_CHAT_TEMPLATE
+    
+    # 4. Update model_max_length for long audio support
+    config["model_max_length"] = 131072
+    
+    # 5. Add add_bos_token if not present
+    if "add_bos_token" not in config:
+        config["add_bos_token"] = False
+    
+    # Write back
+    with open(config_path, "w", encoding="utf-8") as f:
+        json.dump(config, f, indent=2, ensure_ascii=False)
+    
+    print(f"Patched {config_path}")
+
+
+def patch_tokenizer_json(output_dir: str) -> None:
+    """
+    Patch tokenizer.json with VibeVoice audio tokens.
+    """
+    tokenizer_path = os.path.join(output_dir, "tokenizer.json")
+    
+    with open(tokenizer_path, "r", encoding="utf-8") as f:
+        tokenizer = json.load(f)
+    
+    # Find existing token IDs to avoid duplicates
+    existing_ids = set()
+    if "added_tokens" in tokenizer:
+        for token_entry in tokenizer["added_tokens"]:
+            existing_ids.add(token_entry.get("id"))
+    
+    # Add ALL extended tokens (Qwen2.5 + VibeVoice audio)
+    for token, token_id in ALL_EXTENDED_TOKENS.items():
+        if token_id not in existing_ids:
+            # Determine if token should be marked as "special"
+            is_special = token not in ("", "", "<|fim_prefix|>", 
+                                       "<|fim_middle|>", "<|fim_suffix|>", "<|fim_pad|>",
+                                       "<|repo_name|>", "<|file_sep|>")
+            tokenizer["added_tokens"].append({
+                "id": token_id,
+                "content": token,
+                "single_word": False,
+                "lstrip": False,
+                "rstrip": False,
+                "normalized": False,
+                "special": is_special,
+            })
+    
+    # Write back
+    with open(tokenizer_path, "w", encoding="utf-8") as f:
+        json.dump(tokenizer, f, indent=2, ensure_ascii=False)
+    
+    print(f"Patched {tokenizer_path}")
+
+
+def generate_added_tokens_json(output_dir: str) -> None:
+    """
+    Generate added_tokens.json from tokenizer_config.json.
+    """
+    config_path = os.path.join(output_dir, "tokenizer_config.json")
+    
+    with open(config_path, "r", encoding="utf-8") as f:
+        config = json.load(f)
+    
+    added_tokens = {}
+    for token_id, token_info in config.get("added_tokens_decoder", {}).items():
+        content = token_info.get("content")
+        if content:
+            added_tokens[content] = int(token_id)
+    
+    output_path = os.path.join(output_dir, "added_tokens.json")
+    with open(output_path, "w", encoding="utf-8") as f:
+        json.dump(added_tokens, f, indent=2, ensure_ascii=False)
+    
+    print(f"Generated {output_path}")
+
+
+def generate_special_tokens_map_json(output_dir: str) -> None:
+    """
+    Generate special_tokens_map.json with VibeVoice special tokens.
+    """
+    # Build the special tokens map
+    special_tokens_map = {
+        "additional_special_tokens": [],
+        "eos_token": "<|endoftext|>",
+        "pad_token": "<|endoftext|>",
+        "unk_token": "<|endoftext|>",
+    }
+    
+    # Add audio tokens as additional_special_tokens
+    for token in VIBEVOICE_AUDIO_TOKENS.keys():
+        special_tokens_map["additional_special_tokens"].append({
+            "content": token,
+            "lstrip": False,
+            "normalized": False,
+            "rstrip": False,
+            "single_word": False,
+        })
+    
+    # Add some commonly used special tokens
+    common_special = ["<|object_ref_start|>", "<|object_ref_end|>", "<|box_start|>"]
+    for token in common_special:
+        special_tokens_map["additional_special_tokens"].append({
+            "content": token,
+            "lstrip": False,
+            "normalized": False,
+            "rstrip": False,
+            "single_word": False,
+        })
+    
+    output_path = os.path.join(output_dir, "special_tokens_map.json")
+    with open(output_path, "w", encoding="utf-8") as f:
+        json.dump(special_tokens_map, f, indent=2, ensure_ascii=False)
+    
+    print(f"Generated {output_path}")
+
+
+def generate_vibevoice_tokenizer_files(output_dir: str, qwen_model: str = DEFAULT_QWEN_MODEL) -> None:
+    """
+    Generate all 6 VibeVoice tokenizer files.
+    
+    Files generated:
+    1. vocab.json - from Qwen2.5 (unchanged)
+    2. merges.txt - from Qwen2.5 (unchanged)
+    3. tokenizer.json - from Qwen2.5 + audio tokens
+    4. tokenizer_config.json - from Qwen2.5 + audio tokens + chat_template
+    5. added_tokens.json - generated from tokenizer_config.json
+    6. special_tokens_map.json - generated with VibeVoice tokens
+    """
+    print(f"=== Generating VibeVoice tokenizer files to {output_dir} ===\n")
+    
+    # Step 1: Download base files from Qwen2
+    download_qwen_tokenizer_files(output_dir, qwen_model)
+    
+    # Step 2: Patch tokenizer_config.json
+    patch_tokenizer_config(output_dir)
+    
+    # Step 3: Patch tokenizer.json
+    patch_tokenizer_json(output_dir)
+    
+    # Step 4: Generate added_tokens.json
+    generate_added_tokens_json(output_dir)
+    
+    # Step 5: Generate special_tokens_map.json
+    generate_special_tokens_map_json(output_dir)
+    
+    print(f"\n✅ All 6 tokenizer files generated in {output_dir}")
+
+
+def compare_json_files(file1: str, file2: str, name: str) -> Dict[str, Any]:
+    """Compare two JSON files and return differences."""
+    result = {
+        "name": name,
+        "identical": False,
+        "differences": [],
+    }
+    
+    if not os.path.exists(file1):
+        result["differences"].append(f"File 1 not found: {file1}")
+        return result
+    
+    if not os.path.exists(file2):
+        result["differences"].append(f"File 2 not found: {file2}")
+        return result
+    
+    with open(file1, "r", encoding="utf-8") as f:
+        data1 = json.load(f)
+    
+    with open(file2, "r", encoding="utf-8") as f:
+        data2 = json.load(f)
+    
+    if data1 == data2:
+        result["identical"] = True
+        return result
+    
+    # Find specific differences
+    def find_diff(d1, d2, path=""):
+        diffs = []
+        if isinstance(d1, dict) and isinstance(d2, dict):
+            all_keys = set(d1.keys()) | set(d2.keys())
+            for k in all_keys:
+                new_path = f"{path}.{k}" if path else k
+                if k not in d1:
+                    diffs.append(f"Missing in generated: {new_path}")
+                elif k not in d2:
+                    diffs.append(f"Extra in generated: {new_path}")
+                else:
+                    diffs.extend(find_diff(d1[k], d2[k], new_path))
+        elif isinstance(d1, list) and isinstance(d2, list):
+            if len(d1) != len(d2):
+                diffs.append(f"{path}: list length differs ({len(d1)} vs {len(d2)})")
+            # For lists, just check if they're equal (detailed diff is complex)
+            if d1 != d2:
+                diffs.append(f"{path}: list content differs")
+        elif d1 != d2:
+            # Truncate long values for readability
+            v1 = str(d1)[:100] + "..." if len(str(d1)) > 100 else str(d1)
+            v2 = str(d2)[:100] + "..." if len(str(d2)) > 100 else str(d2)
+            diffs.append(f"{path}: '{v1}' vs '{v2}'")
+        return diffs
+    
+    result["differences"] = find_diff(data1, data2)
+    return result
+
+
+def compare_text_files(file1: str, file2: str, name: str) -> Dict[str, Any]:
+    """Compare two text files."""
+    result = {
+        "name": name,
+        "identical": False,
+        "differences": [],
+    }
+    
+    if not os.path.exists(file1):
+        result["differences"].append(f"File 1 not found: {file1}")
+        return result
+    
+    if not os.path.exists(file2):
+        result["differences"].append(f"File 2 not found: {file2}")
+        return result
+    
+    with open(file1, "r", encoding="utf-8") as f:
+        content1 = f.read()
+    
+    with open(file2, "r", encoding="utf-8") as f:
+        content2 = f.read()
+    
+    if content1 == content2:
+        result["identical"] = True
+    else:
+        lines1 = content1.splitlines()
+        lines2 = content2.splitlines()
+        result["differences"].append(f"Line count: {len(lines1)} vs {len(lines2)}")
+        
+        # Find first difference
+        for i, (l1, l2) in enumerate(zip(lines1, lines2)):
+            if l1 != l2:
+                result["differences"].append(f"First diff at line {i+1}")
+                break
+    
+    return result
+
+
+def compare_with_reference(generated_dir: str, reference_dir: str) -> None:
+    """Compare generated files with reference files."""
+    print(f"\n=== Comparing generated files with reference ===")
+    print(f"Generated: {generated_dir}")
+    print(f"Reference: {reference_dir}\n")
+    
+    files_to_compare = [
+        ("vocab.json", "json"),
+        ("merges.txt", "text"),
+        ("tokenizer.json", "json"),
+        ("tokenizer_config.json", "json"),
+        ("added_tokens.json", "json"),
+        ("special_tokens_map.json", "json"),
+    ]
+    
+    all_identical = True
+    
+    for filename, file_type in files_to_compare:
+        gen_file = os.path.join(generated_dir, filename)
+        ref_file = os.path.join(reference_dir, filename)
+        
+        if file_type == "json":
+            result = compare_json_files(gen_file, ref_file, filename)
+        else:
+            result = compare_text_files(gen_file, ref_file, filename)
+        
+        if result["identical"]:
+            print(f"✅ {filename}: IDENTICAL")
+        else:
+            print(f"❌ {filename}: DIFFERENT")
+            for diff in result["differences"][:5]:  # Show first 5 differences
+                print(f"   - {diff}")
+            if len(result["differences"]) > 5:
+                print(f"   ... and {len(result['differences']) - 5} more differences")
+            all_identical = False
+    
+    print()
+    if all_identical:
+        print("🎉 All files are identical!")
+    else:
+        print("⚠️  Some files have differences. See details above.")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Generate VibeVoice tokenizer files from Qwen2 base"
+    )
+    parser.add_argument(
+        "--output", "-o",
+        type=str,
+        default=None,
+        help="Output directory for generated files (default: temp directory)"
+    )
+    parser.add_argument(
+        "--compare", "-c",
+        type=str,
+        default=None,
+        help="Reference directory to compare generated files against"
+    )
+    parser.add_argument(
+        "--qwen-model",
+        type=str,
+        default=DEFAULT_QWEN_MODEL,
+        help=f"Qwen model to download base tokenizer from (default: {DEFAULT_QWEN_MODEL})"
+    )
+    
+    args = parser.parse_args()
+    
+    # Determine output directory
+    if args.output:
+        output_dir = args.output
+        cleanup = False
+    else:
+        output_dir = tempfile.mkdtemp(prefix="vibevoice_tokenizer_")
+        cleanup = not args.compare  # Only cleanup if not comparing
+    
+    try:
+        # Generate files
+        generate_vibevoice_tokenizer_files(output_dir, args.qwen_model)
+        
+        # Compare if requested
+        if args.compare:
+            compare_with_reference(output_dir, args.compare)
+        
+        if not args.output:
+            print(f"\nGenerated files are in: {output_dir}")
+            
+    finally:
+        if cleanup and not args.output:
+            print(f"\nCleaning up temporary directory: {output_dir}")
+            shutil.rmtree(output_dir, ignore_errors=True)
+
+
+if __name__ == "__main__":
+    main()