Update README.md

.gitignore

@@ -175,3 +175,4 @@ tags
 /checkpoints/
 exp
 .gradio/
+MSFT-VV

README.md

@@ -34,6 +34,8 @@ Fine-tuning is now supported, which is incredibly powerful. You can adapt VibeVo
 
 ## Updates
 
+- **[2026-06-12]** [Unofficial voice cloning](https://huggingface.co/mohammed-bahumaish/vibevoice-realtime-0.5b-with-encoder) for Streaming/Realtime model released by community member
+- **[2025-12-04]** Added support for VibeVoice-Streaming-0.5B model for real-time TTS!
 - **[2025-09-05]** Microsoft repo restored (without code) with statement about responsible AI use.
 - **[2025-09-04]** Community backup created after Microsoft removed original repo and models.
 - **[2025-08-26]** The [VibeVoice-7B](https://huggingface.co/vibevoice/VibeVoice-7B) model weights are open-sourced!
@@ -47,10 +49,16 @@ Fine-tuning is now supported, which is incredibly powerful. You can adapt VibeVo
 
 ## Model Zoo
 
-| Model | Context Length | Generation Length |  Weight |
+| Model | Context Length | Generation Length | Speakers | Weight |
-|-------|----------------|----------|----------|
+|-------|----------------|-------------------|----------|--------|
-| VibeVoice-1.5B | 64K | ~90 min | [HF link](https://huggingface.co/vibevoice/VibeVoice-1.5B) |
+| VibeVoice-Streaming-0.5B | 8K | Real-time | 1 | [HF link](https://huggingface.co/microsoft/VibeVoice-Realtime-0.5B) |
-| VibeVoice-Large| 32K | ~45 min | [HF link](https://huggingface.co/vibevoice/VibeVoice-7B) |
+| VibeVoice-1.5B | 64K | ~90 min | Up to 4 | [HF link](https://huggingface.co/vibevoice/VibeVoice-1.5B) |
+| VibeVoice-Large (7B) | 32K | ~45 min | Up to 4 | [HF link](https://huggingface.co/vibevoice/VibeVoice-7B) |
+
+### Model Comparison
+
+- **VibeVoice-Streaming-0.5B**: Optimized for **real-time** low-latency TTS. Single speaker only. Uses pre-computed voice embeddings (.pt files) for fast inference. Best for live applications. [Unofficial voice cloning](https://huggingface.co/mohammed-bahumaish/vibevoice-realtime-0.5b-with-encoder)
+- **VibeVoice-1.5B/7B**: Full-featured models for **long-form multi-speaker** content like podcasts. Support up to 4 speakers with voice cloning from audio samples.
 
 ## Installation
 
@@ -82,7 +90,7 @@ python demo/gradio_demo.py --model_path vibevoice/VibeVoice-1.5B --share
 # use the in-app "Disable voice cloning" setting (Advanced Settings) to skip speaker conditioning
 ```
 
-**Option 2: Inference from files directly**
+**Option 2: Inference from files directly (Multi-Speaker 1.5B/7B)**
 
 ```bash
 # We provide some LLM generated example scripts under demo/text_examples/ for demo
@@ -99,6 +107,30 @@ python demo/inference_from_file.py --model_path vibevoice/VibeVoice-7B --txt_pat
 python demo/inference_from_file.py --model_path vibevoice/VibeVoice-7B --txt_path demo/text_examples/1p_abs.txt --speaker_names Alice --disable_prefill
 ```
 
+**Option 3: Streaming Model (0.5B) - Real-time TTS**
+
+The streaming model uses pre-computed voice embeddings for low-latency generation:
+
+```bash
+# Basic usage with streaming model
+python demo/streaming_inference_from_file.py \
+    --model_path microsoft/VibeVoice-Realtime-0.5B \
+    --txt_path demo/text_examples/1p_vibevoice.txt \
+    --speaker_name Emma
+
+# Available voice presets: Carter, Davis, Emma, Frank, Grace, Mike (English), Samuel (Indian English)
+# Adjust CFG scale and DDPM steps for quality/speed tradeoff
+python demo/streaming_inference_from_file.py \
+    --model_path microsoft/VibeVoice-Realtime-0.5B \
+    --txt_path demo/text_examples/1p_vibevoice.txt \
+    --speaker_name Mike \
+    --cfg_scale 1.5 \
+    --ddpm_steps 5
+```
+
+Voice presets are stored as `.pt` files in `demo/voices/streaming_model/`. These contain pre-computed KV cache embeddings for fast inference. Voice cloning is not supported for now.
+
+NOTE: If you get the warning `Some weights of VibeVoiceStreamingForConditionalGenerationInference were not initialized from the model checkpoint` when loading, this is expected. This is because voice cloning capabilities have been removed from the model.
 
 ## [Finetuning](./FINETUNING.md)
 

demo/VibeVoice_colab.ipynb

@@ -7,7 +7,7 @@
         "id": "view-in-github"
       },
       "source": [
-        "<a href=\"https://colab.research.google.com/github/microsoft/VibeVoice/blob/main/demo/VibeVoice_colab.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+        "<a href=\"https://colab.research.google.com/github/vibevoice-community/VibeVoice/blob/main/demo/VibeVoice_colab.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
       ]
     },
     {

demo/gradio_demo.py

@@ -19,6 +19,7 @@ import soundfile as sf
 import torch
 import os
 import traceback
+import re
 
 from vibevoice.modular.configuration_vibevoice import VibeVoiceConfig
 from vibevoice.modular.modeling_vibevoice_inference import VibeVoiceForConditionalGenerationInference
@@ -210,6 +211,8 @@ class VibeVoiceDemo:
                                  speaker_3: str = None,
                                  speaker_4: str = None,
                                  cfg_scale: float = 1.3,
+                                 inference_steps: Optional[int] = None,
+                                 seed: Optional[int] = None,
                                  disable_voice_cloning: bool = False) -> Iterator[tuple]:
         try:
             
@@ -240,9 +243,37 @@ class VibeVoiceDemo:
             
             voice_cloning_enabled = not disable_voice_cloning
 
+            # Resolve per-run parameters
+            resolved_inference_steps = int(inference_steps) if inference_steps is not None else int(self.inference_steps)
+            resolved_seed: Optional[int]
+            if seed is None:
+                resolved_seed = None
+            else:
+                # Gradio Number may come through as float
+                resolved_seed = int(seed)
+                if resolved_seed < 0:
+                    resolved_seed = None
+
+            def _sanitize_filename_part(value: str, max_len: int = 60) -> str:
+                value = (value or "").strip()
+                value = re.sub(r"\s+", "_", value)
+                value = re.sub(r"[^A-Za-z0-9_\-\+\.]+", "", value)
+                return value[:max_len] if len(value) > max_len else value
+
+            def _write_complete_wav(audio_int16: np.ndarray, sample_rate: int = 24000) -> str:
+                ts = datetime.now().strftime("%Y%m%d%H%M%S")
+                speakers_part = _sanitize_filename_part("-".join(selected_speakers) if selected_speakers else "speaker")
+                cfg_part = _sanitize_filename_part(f"{cfg_scale:.2f}")
+                seed_part = str(resolved_seed) if resolved_seed is not None else "rand"
+                # Requested pattern: yyyyMMddhhmmss_[speaker]_[cfg_scale]_[inference_step][seed].wav
+                basename = f"{ts}_{speakers_part}_{cfg_part}_{resolved_inference_steps}_{seed_part}.wav"
+                out_path = os.path.join(tempfile.gettempdir(), basename)
+                sf.write(out_path, np.asarray(audio_int16, dtype=np.int16), sample_rate, subtype="PCM_16")
+                return out_path
+
             # Build initial log
             log = f"🎙️ Generating podcast with {num_speakers} speakers\n"
-            log += f"📊 Parameters: CFG Scale={cfg_scale}, Inference Steps={self.inference_steps}\n"
+            log += f"📊 Parameters: CFG Scale={cfg_scale}, Inference Steps={resolved_inference_steps}, Seed={(resolved_seed if resolved_seed is not None else 'random')}\n"
             log += f"🎭 Speakers: {', '.join(selected_speakers)}\n"
             log += f"🔊 Voice cloning: {'Enabled' if voice_cloning_enabled else 'Disabled'}\n"
             if self.loaded_adapter_root:
@@ -331,7 +362,7 @@ class VibeVoiceDemo:
             # Start generation in a separate thread
             generation_thread = threading.Thread(
                 target=self._generate_with_streamer,
-                args=(inputs, cfg_scale, audio_streamer, voice_cloning_enabled)
+                args=(inputs, cfg_scale, audio_streamer, voice_cloning_enabled, resolved_inference_steps, resolved_seed, target_device)
             )
             generation_thread.start()
             
@@ -465,8 +496,9 @@ class VibeVoiceDemo:
                 final_log += "✨ Generation successful! Complete audio is ready.\n"
                 final_log += "💡 Not satisfied? You can regenerate or adjust the CFG scale for different results."
                 
-                # Yield the complete audio
+                # Yield the complete audio as a filepath so download uses our filename
-                yield None, (sample_rate, complete_audio), final_log, gr.update(visible=False)
+                complete_wav_path = _write_complete_wav(complete_audio, sample_rate=sample_rate)
+                yield None, complete_wav_path, final_log, gr.update(visible=False)
                 return
             
             if not has_received_chunks:
@@ -490,8 +522,9 @@ class VibeVoiceDemo:
                 final_log += "✨ Generation successful! Complete audio is ready in the 'Complete Audio' tab.\n"
                 final_log += "💡 Not satisfied? You can regenerate or adjust the CFG scale for different results."
                 
-                # Final yield: Clear streaming audio and provide complete audio
+                # Final yield: Clear streaming audio and provide complete audio as filepath
-                yield None, (sample_rate, complete_audio), final_log, gr.update(visible=False)
+                complete_wav_path = _write_complete_wav(complete_audio, sample_rate=sample_rate)
+                yield None, complete_wav_path, final_log, gr.update(visible=False)
             else:
                 final_log = log + "❌ No audio was generated."
                 yield None, None, final_log, gr.update(visible=False)
@@ -513,17 +546,44 @@ class VibeVoiceDemo:
             traceback.print_exc()
             yield None, None, error_msg, gr.update(visible=False)
     
-    def _generate_with_streamer(self, inputs, cfg_scale, audio_streamer, voice_cloning_enabled: bool):
+    def _generate_with_streamer(
+        self,
+        inputs,
+        cfg_scale,
+        audio_streamer,
+        voice_cloning_enabled: bool,
+        inference_steps: int,
+        seed: Optional[int],
+        target_device: str,
+    ):
         """Helper method to run generation with streamer in a separate thread."""
         try:
             # Check for stop signal before starting generation
             if self.stop_generation:
                 audio_streamer.end()
                 return
+
+            # Apply per-run DDPM steps
+            try:
+                self.model.set_ddpm_inference_steps(num_steps=int(inference_steps))
+            except Exception as e:
+                print(f"Warning: failed to set inference steps ({inference_steps}): {e}")
                 
             # Define a stop check function that can be called from generate
             def check_stop_generation():
                 return self.stop_generation
+
+            generator = None
+            if seed is not None:
+                try:
+                    if target_device == "cuda":
+                        generator = torch.Generator(device="cuda")
+                    else:
+                        generator = torch.Generator()
+                    generator.manual_seed(int(seed))
+                except Exception as e:
+                    print(f"Warning: failed to create seeded generator (seed={seed}, device={target_device}): {e}")
+                    generator = None
                 
             outputs = self.model.generate(
                 **inputs,
@@ -533,6 +593,7 @@ class VibeVoiceDemo:
                 generation_config={
                     'do_sample': False,
                 },
+                generator=generator,
                 audio_streamer=audio_streamer,
                 stop_check_fn=check_stop_generation,  # Pass the stop check function
                 verbose=False,  # Disable verbose in streaming mode
@@ -639,227 +700,17 @@ class VibeVoiceDemo:
 def create_demo_interface(demo_instance: VibeVoiceDemo):
     """Create the Gradio interface with streaming support."""
     
-    # Custom CSS for high-end aesthetics with lighter theme
+    with gr.Blocks() as interface:
-    custom_css = """
-    /* Modern light theme with gradients */
-    .gradio-container {
-        background: linear-gradient(135deg, #f8fafc 0%, #e2e8f0 100%);
-        font-family: 'SF Pro Display', -apple-system, BlinkMacSystemFont, sans-serif;
-    }
-    
-    /* Header styling */
-    .main-header {
-        background: linear-gradient(90deg, #667eea 0%, #764ba2 100%);
-        padding: 2rem;
-        border-radius: 20px;
-        margin-bottom: 2rem;
-        text-align: center;
-        box-shadow: 0 10px 40px rgba(102, 126, 234, 0.3);
-    }
-    
-    .main-header h1 {
-        color: white;
-        font-size: 2.5rem;
-        font-weight: 700;
-        margin: 0;
-        text-shadow: 0 2px 4px rgba(0,0,0,0.3);
-    }
-    
-    .main-header p {
-        color: rgba(255,255,255,0.9);
-        font-size: 1.1rem;
-        margin: 0.5rem 0 0 0;
-    }
-    
-    /* Card styling */
-    .settings-card, .generation-card {
-        background: rgba(255, 255, 255, 0.8);
-        backdrop-filter: blur(10px);
-        border: 1px solid rgba(226, 232, 240, 0.8);
-        border-radius: 16px;
-        padding: 1.5rem;
-        margin-bottom: 1rem;
-        box-shadow: 0 8px 32px rgba(0, 0, 0, 0.1);
-    }
-    
-    /* Speaker selection styling */
-    .speaker-grid {
-        display: grid;
-        gap: 1rem;
-        margin-bottom: 1rem;
-    }
-    
-    .speaker-item {
-        background: linear-gradient(135deg, #e2e8f0 0%, #cbd5e1 100%);
-        border: 1px solid rgba(148, 163, 184, 0.4);
-        border-radius: 12px;
-        padding: 1rem;
-        color: #374151;
-        font-weight: 500;
-    }
-    
-    /* Streaming indicator */
-    .streaming-indicator {
-        display: inline-block;
-        width: 10px;
-        height: 10px;
-        background: #22c55e;
-        border-radius: 50%;
-        margin-right: 8px;
-        animation: pulse 1.5s infinite;
-    }
-    
-    @keyframes pulse {
-        0% { opacity: 1; transform: scale(1); }
-        50% { opacity: 0.5; transform: scale(1.1); }
-        100% { opacity: 1; transform: scale(1); }
-    }
-    
-    /* Queue status styling */
-    .queue-status {
-        background: linear-gradient(135deg, #f0f9ff 0%, #e0f2fe 100%);
-        border: 1px solid rgba(14, 165, 233, 0.3);
-        border-radius: 8px;
-        padding: 0.75rem;
-        margin: 0.5rem 0;
-        text-align: center;
-        font-size: 0.9rem;
-        color: #0369a1;
-    }
-    
-    .generate-btn {
-        background: linear-gradient(135deg, #059669 0%, #0d9488 100%);
-        border: none;
-        border-radius: 12px;
-        padding: 1rem 2rem;
-        color: white;
-        font-weight: 600;
-        font-size: 1.1rem;
-        box-shadow: 0 4px 20px rgba(5, 150, 105, 0.4);
-        transition: all 0.3s ease;
-    }
-    
-    .generate-btn:hover {
-        transform: translateY(-2px);
-        box-shadow: 0 6px 25px rgba(5, 150, 105, 0.6);
-    }
-    
-    .stop-btn {
-        background: linear-gradient(135deg, #ef4444 0%, #dc2626 100%);
-        border: none;
-        border-radius: 12px;
-        padding: 1rem 2rem;
-        color: white;
-        font-weight: 600;
-        font-size: 1.1rem;
-        box-shadow: 0 4px 20px rgba(239, 68, 68, 0.4);
-        transition: all 0.3s ease;
-    }
-    
-    .stop-btn:hover {
-        transform: translateY(-2px);
-        box-shadow: 0 6px 25px rgba(239, 68, 68, 0.6);
-    }
-    
-    /* Audio player styling */
-    .audio-output {
-        background: linear-gradient(135deg, #f1f5f9 0%, #e2e8f0 100%);
-        border-radius: 16px;
-        padding: 1.5rem;
-        border: 1px solid rgba(148, 163, 184, 0.3);
-    }
-    
-    .complete-audio-section {
-        margin-top: 1rem;
-        padding: 1rem;
-        background: linear-gradient(135deg, #f0fdf4 0%, #dcfce7 100%);
-        border: 1px solid rgba(34, 197, 94, 0.3);
-        border-radius: 12px;
-    }
-    
-    /* Text areas */
-    .script-input, .log-output {
-        background: rgba(255, 255, 255, 0.9) !important;
-        border: 1px solid rgba(148, 163, 184, 0.4) !important;
-        border-radius: 12px !important;
-        color: #1e293b !important;
-        font-family: 'JetBrains Mono', monospace !important;
-    }
-    
-    .script-input::placeholder {
-        color: #64748b !important;
-    }
-    
-    /* Sliders */
-    .slider-container {
-        background: rgba(248, 250, 252, 0.8);
-        border: 1px solid rgba(226, 232, 240, 0.6);
-        border-radius: 8px;
-        padding: 1rem;
-        margin: 0.5rem 0;
-    }
-    
-    /* Labels and text */
-    .gradio-container label {
-        color: #374151 !important;
-        font-weight: 600 !important;
-    }
-    
-    .gradio-container .markdown {
-        color: #1f2937 !important;
-    }
-    
-    /* Responsive design */
-    @media (max-width: 768px) {
-        .main-header h1 { font-size: 2rem; }
-        .settings-card, .generation-card { padding: 1rem; }
-    }
-    
-    /* Random example button styling - more subtle professional color */
-    .random-btn {
-        background: linear-gradient(135deg, #64748b 0%, #475569 100%);
-        border: none;
-        border-radius: 12px;
-        padding: 1rem 1.5rem;
-        color: white;
-        font-weight: 600;
-        font-size: 1rem;
-        box-shadow: 0 4px 20px rgba(100, 116, 139, 0.3);
-        transition: all 0.3s ease;
-        display: inline-flex;
-        align-items: center;
-        gap: 0.5rem;
-    }
-    
-    .random-btn:hover {
-        transform: translateY(-2px);
-        box-shadow: 0 6px 25px rgba(100, 116, 139, 0.4);
-        background: linear-gradient(135deg, #475569 0%, #334155 100%);
-    }
-    """
-    
-    with gr.Blocks(
-        title="VibeVoice - AI Podcast Generator",
-        css=custom_css,
-        theme=gr.themes.Soft(
-            primary_hue="blue",
-            secondary_hue="purple",
-            neutral_hue="slate",
-        )
-    ) as interface:
         
         # Header
         gr.HTML("""
-        <div class="main-header">
+        # VibeVoice
-            <h1>🎙️ Vibe Podcasting </h1>
-            <p>Generating Long-form Multi-speaker AI Podcast with VibeVoice</p>
-        </div>
         """)
         
         with gr.Row():
             # Left column - Settings
             with gr.Column(scale=1, elem_classes="settings-card"):
-                gr.Markdown("### 🎛️ **Podcast Settings**")
+                gr.Markdown("### **Podcast Settings**")
                 
                 # Number of speakers
                 num_speakers = gr.Slider(
@@ -897,13 +748,26 @@ def create_demo_interface(demo_instance: VibeVoiceDemo):
                 with gr.Accordion("Generation Parameters", open=False):
                     cfg_scale = gr.Slider(
                         minimum=1.0,
-                        maximum=2.0,
+                        maximum=12.0,
                         value=1.3,
                         step=0.05,
                         label="CFG Scale (Guidance Strength)",
                         # info="Higher values increase adherence to text",
                         elem_classes="slider-container"
                     )
+                    inference_steps = gr.Slider(
+                        minimum=1,
+                        maximum=50,
+                        value=int(getattr(demo_instance, "inference_steps", 10)),
+                        step=1,
+                        label="Inference Steps",
+                        elem_classes="slider-container",
+                    )
+                    seed = gr.Number(
+                        value=42,
+                        precision=0,
+                        label="Seed (-1 = random)",
+                    )
                     disable_voice_cloning = gr.Checkbox(
                         value=False,
                         label="Disable voice cloning (skip conditioning voice prompts)",
@@ -992,7 +856,7 @@ Or paste text directly and it will auto-assign speakers.""",
                 # Complete audio output (non-streaming)
                 complete_audio_output = gr.Audio(
                     label="Complete Podcast (Download after generation)",
-                    type="numpy",
+                    type="filepath",
                     elem_classes="audio-output complete-audio-section",
                     streaming=False,  # Non-streaming mode
                     autoplay=False,
@@ -1027,7 +891,7 @@ Or paste text directly and it will auto-assign speakers.""",
         )
         
         # Main generation function with streaming
-        def generate_podcast_wrapper(num_speakers, script, speaker_1, speaker_2, speaker_3, speaker_4, cfg_scale, disable_voice_cloning):
+        def generate_podcast_wrapper(num_speakers, script, speaker_1, speaker_2, speaker_3, speaker_4, cfg_scale, inference_steps, seed, disable_voice_cloning):
             """Wrapper function to handle the streaming generation call."""
             try:
                 speakers = [speaker_1, speaker_2, speaker_3, speaker_4]
@@ -1046,6 +910,8 @@ Or paste text directly and it will auto-assign speakers.""",
                     speaker_3=speakers[2],
                     speaker_4=speakers[3],
                     cfg_scale=cfg_scale,
+                    inference_steps=inference_steps,
+                    seed=seed,
                     disable_voice_cloning=disable_voice_cloning
                 ):
                     final_log = log
@@ -1094,7 +960,7 @@ Or paste text directly and it will auto-assign speakers.""",
             queue=False
         ).then(
             fn=generate_podcast_wrapper,
-            inputs=[num_speakers, script_input] + speaker_selections + [cfg_scale, disable_voice_cloning],
+            inputs=[num_speakers, script_input] + speaker_selections + [cfg_scale, inference_steps, seed, disable_voice_cloning],
             outputs=[audio_output, complete_audio_output, log_output, streaming_status, generate_btn, stop_btn],
             queue=True  # Enable Gradio's built-in queue
         )

demo/streaming_inference_from_file.py

@@ -0,0 +1,377 @@
+"""
+VibeVoice Streaming Model Inference from File
+
+This script demonstrates how to use the VibeVoice-Streaming-0.5B model
+for text-to-speech synthesis from a text file.
+
+The streaming model uses pre-computed voice embeddings (.pt files) for
+low-latency generation, making it suitable for real-time applications.
+
+Usage:
+    python demo/streaming_inference_from_file.py \
+        --model_path microsoft/VibeVoice-Realtime-0.5B \
+        --txt_path demo/text_examples/1p_vibevoice.txt \
+        --speaker_name Emma \
+        --output_dir ./outputs
+
+Available voice presets (in demo/voices/streaming_model/):
+    - Carter, Davis, Emma, Frank, Grace, Mike (English)
+    - Samuel (Indian English)
+"""
+
+import argparse
+import os
+import re
+import traceback
+from typing import List, Tuple, Union, Dict, Any
+import time
+import torch
+import copy
+
+from vibevoice.modular.modeling_vibevoice_streaming_inference import VibeVoiceStreamingForConditionalGenerationInference
+from vibevoice.processor.vibevoice_streaming_processor import VibeVoiceStreamingProcessor
+from transformers.utils import logging
+
+logging.set_verbosity_info()
+logger = logging.get_logger(__name__)
+
+
+class VoiceMapper:
+    """Maps speaker names to voice file paths for the streaming model."""
+
+    def __init__(self):
+        self.setup_voice_presets()
+
+        # Create name aliases without prefixes/suffixes for easier matching
+        new_dict = {}
+        for name, path in self.voice_presets.items():
+            # Remove gender suffix (e.g., "en-Emma_woman" -> "Emma")
+            if '_' in name:
+                name = name.split('_')[0]
+
+            # Remove language prefix (e.g., "en-Emma" -> "Emma")
+            if '-' in name:
+                name = name.split('-')[-1]
+
+            new_dict[name] = path
+        self.voice_presets.update(new_dict)
+
+    def setup_voice_presets(self):
+        """Setup voice presets by scanning the voices/streaming_model 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 .pt files in the voices directory
+        self.voice_presets = {}
+
+        pt_files = [f for f in os.listdir(voices_dir)
+                    if f.lower().endswith('.pt') and os.path.isfile(os.path.join(voices_dir, f))]
+
+        # Create dictionary with filename (without extension) as key
+        for pt_file in pt_files:
+            name = os.path.splitext(pt_file)[0]
+            full_path = os.path.join(voices_dir, pt_file)
+            self.voice_presets[name] = full_path
+
+        # Sort alphabetically for consistent ordering
+        self.voice_presets = dict(sorted(self.voice_presets.items()))
+
+        # Filter out voices that don't exist (safety check)
+        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
+        if speaker_name in self.voice_presets:
+            return self.voice_presets[speaker_name]
+
+        # Try partial matching (case insensitive)
+        speaker_lower = speaker_name.lower()
+        for preset_name, path in self.voice_presets.items():
+            if preset_name.lower() in speaker_lower or speaker_lower in preset_name.lower():
+                return path
+
+        # Default to first voice if no match found
+        if self.voice_presets:
+            default_voice = list(self.voice_presets.values())[0]
+            print(f"Warning: No voice preset found for '{speaker_name}', using default voice")
+            return default_voice
+
+        raise ValueError(f"No voice presets available. Please check the voices/streaming_model directory.")
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="VibeVoice Streaming Model (0.5B) - Text-to-Speech from File",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+    # Basic usage with default settings
+    python demo/streaming_inference_from_file.py --txt_path demo/text_examples/1p_vibevoice.txt
+
+    # Specify a different voice
+    python demo/streaming_inference_from_file.py --txt_path demo/text_examples/1p_vibevoice.txt --speaker_name Emma
+
+    # Use local model path
+    python demo/streaming_inference_from_file.py --model_path /path/to/model --txt_path input.txt
+        """
+    )
+    parser.add_argument(
+        "--model_path",
+        type=str,
+        default="microsoft/VibeVoice-Realtime-0.5B",
+        help="Path to the HuggingFace model directory or model ID",
+    )
+    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="Emma",
+        help="Speaker name for voice selection (e.g., Emma, Carter, Mike)",
+    )
+    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)",
+    )
+    parser.add_argument(
+        "--ddpm_steps",
+        type=int,
+        default=5,
+        help="Number of DDPM inference steps (default: 5, lower is faster)",
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=None,
+        help="Random seed for reproducibility",
+    )
+
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    # Set random seed if specified
+    if args.seed is not None:
+        torch.manual_seed(args.seed)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(args.seed)
+        print(f"Using random seed: {args.seed}")
+
+    # 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
+
+    # Normalize quotes
+    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 based on device
+    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,
+            )
+            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. Note that flash_attention_2 is recommended for best 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=args.ddpm_steps)
+
+    if hasattr(model.model, 'language_model'):
+        print(f"Language model attention: {model.model.language_model.config._attn_implementation}")
+
+    # Load voice preset
+    target_device = args.device if args.device != "cpu" else "cpu"
+    voice_sample = voice_mapper.get_voice_path(args.speaker_name)
+    print(f"Loading voice preset: {voice_sample}")
+    all_prefilled_outputs = torch.load(voice_sample, map_location=target_device, weights_only=False)
+
+    # 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}, ddpm_steps: {args.ddpm_steps}")
+
+    # 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:
+        sample_rate = 24000  # 24kHz sample rate
+        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")
+        audio_duration = 0
+        rtf = float('inf')
+
+    # Calculate token metrics
+    input_tokens = inputs['tts_text_ids'].shape[1]
+    output_tokens = outputs.sequences.shape[1]
+    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
+    txt_filename = os.path.splitext(os.path.basename(args.txt_path))[0]
+    output_path = os.path.join(args.output_dir, f"{txt_filename}_streaming_generated.wav")
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    processor.save_audio(
+        outputs.speech_outputs[0],
+        output_path=output_path,
+    )
+    print(f"Saved output to {output_path}")
+
+    # Print summary
+    print("\n" + "=" * 50)
+    print("GENERATION SUMMARY (Streaming 0.5B)")
+    print("=" * 50)
+    print(f"Model: {args.model_path}")
+    print(f"Input file: {args.txt_path}")
+    print(f"Output file: {output_path}")
+    print(f"Speaker: {args.speaker_name}")
+    print(f"CFG Scale: {args.cfg_scale}")
+    print(f"DDPM Steps: {args.ddpm_steps}")
+    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()

demo/text_examples/1p_vibevoice.txt

@@ -0,0 +1,7 @@
+Hello and welcome! Today I'm going to tell you about VibeVoice, a groundbreaking text-to-speech model developed for generating natural, expressive speech.
+
+VibeVoice uses an innovative architecture combining autoregressive language modeling with diffusion-based audio generation. This allows it to produce remarkably natural-sounding speech with proper intonation and rhythm.
+
+The streaming version you're hearing right now is optimized for real-time applications, with low latency that makes it perfect for interactive use cases like virtual assistants and live applications.
+
+Thank you for listening to this demonstration!

demo/vibevoice_asr_inference_from_file.py

@@ -0,0 +1,554 @@
+#!/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
+
+
+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 = "flash_attention_2"
+    ):
+        """
+        Initialize the ASR batch inference pipeline.
+        
+        Args:
+            model_path: Path to the pretrained model
+            device: Device to run inference on
+            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 "cpu",
+        choices=["cuda", "cpu", "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="flash_attention_2",
+        help="Attention implementation to use (default: flash_attention_2)"
+    )
+    
+    args = parser.parse_args()
+    
+    # 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:
+        import glob
+        for ext in ["*.wav", "*.mp3", "*.flac", "*.mp4", "*.m4a", "*.webm"]:
+            audio_files.extend(glob.glob(os.path.join(args.audio_dir, ext)))
+    
+    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
+    asr = VibeVoiceASRBatchInference(
+        model_path=args.model_path,
+        device=args.device,
+        dtype=torch.bfloat16 if args.device != "cpu" else torch.float32,
+        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()

docs/vibevoice-asr.md

@@ -0,0 +1,72 @@
+# 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**.
+
+## 🔥 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*.
+
+**Demo:** [VibeVoice-ASR-Demo](https://aka.ms/vibevoice-asr)
+
+## 🏗️ Model Architecture
+
+<p align="center">
+  <img src="../Figures/VibeVoice_ASR_archi.png" alt="VibeVoice ASR Architecture" width="80%">
+</p>
+
+## Evaluation
+<p align="center">
+  <img src="../Figures/DER.jpg" alt="DER" width="50%">
+  <img src="../Figures/cpWER.jpg" alt="cpWER" width="50%">
+  <img src="../Figures/tcpWER.jpg" alt="tcpWER" width="50%">
+</p>
+
+## Installation
+We recommend to use 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 .[asr]
+```
+
+## 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 a audio path here] 
+```
+
+
+## 📄 License
+
+This project is licensed under the [MIT License](../LICENSE).

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "vibevoice"
-version = "0.0.1"
+version = "0.1.0"
 authors = [
   { name="vibevoice team", email="vibepod@microsoft.com" },
 ]
@@ -31,11 +31,17 @@ dependencies = [
     "librosa",
     "ml-collections",
     "absl-py",
-    "gradio",
+    "gradio==5.50.0",
     "av",
     "aiortc"
 ]
 
+[project.optional-dependencies]
+# Optional dependencies for streaming model web service
+streaming-web = [
+    "fastapi",
+    "uvicorn[standard]",
+]
 
 [project.urls]
 "Homepage" = "https://github.com/microsoft/VibeVoice"

vibevoice/modular/__init__.py

@@ -0,0 +1,107 @@
+"""
+VibeVoice Modular Components
+
+This module provides the core model architectures for VibeVoice:
+- Multi-speaker models (1.5B, 7B) for high-quality multi-speaker TTS
+- Streaming model (0.5B) for real-time low-latency TTS
+"""
+
+# Multi-speaker model components
+from .configuration_vibevoice import (
+    VibeVoiceConfig,
+    VibeVoiceAcousticTokenizerConfig,
+    VibeVoiceSemanticTokenizerConfig,
+    VibeVoiceDiffusionHeadConfig,
+    VibeVoiceASRConfig,
+)
+from .modeling_vibevoice_asr import (
+    VibeVoiceASRPreTrainedModel,
+    VibeVoiceASRModel,
+    VibeVoiceASRForConditionalGeneration,
+)
+from .modeling_vibevoice import (
+    VibeVoicePreTrainedModel,
+    VibeVoiceModel,
+)
+from .modeling_vibevoice_inference import (
+    VibeVoiceForConditionalGenerationInference,
+)
+
+# Streaming model components (0.5B)
+from .configuration_vibevoice_streaming import (
+    VibeVoiceStreamingConfig,
+)
+from .modeling_vibevoice_streaming import (
+    VibeVoiceStreamingPreTrainedModel,
+    VibeVoiceStreamingModel,
+    BinaryClassifier,
+    SpeechConnector,
+)
+from .modeling_vibevoice_streaming_inference import (
+    VibeVoiceStreamingForConditionalGenerationInference,
+    VibeVoiceGenerationOutput,
+    TTS_TEXT_WINDOW_SIZE,
+    TTS_SPEECH_WINDOW_SIZE,
+)
+
+# Shared components
+from .modular_vibevoice_tokenizer import (
+    VibeVoiceTokenizerStreamingCache,
+    VibeVoiceAcousticTokenizerModel,
+    VibeVoiceSemanticTokenizerModel,
+)
+from .modular_vibevoice_text_tokenizer import (
+    VibeVoiceTextTokenizer,
+    VibeVoiceTextTokenizerFast,
+)
+from .modular_vibevoice_diffusion_head import (
+    VibeVoiceDiffusionHead,
+)
+from .streamer import (
+    AudioStreamer,
+    AsyncAudioStreamer,
+)
+
+# LoRA support
+from .lora_loading import (
+    load_lora_assets,
+)
+
+__all__ = [
+    # Multi-speaker configs
+    "VibeVoiceConfig",
+    "VibeVoiceAcousticTokenizerConfig",
+    "VibeVoiceSemanticTokenizerConfig",
+    "VibeVoiceDiffusionHeadConfig",
+    "VibeVoiceASRConfig",
+    # Multi-speaker models
+    "VibeVoicePreTrainedModel",
+    "VibeVoiceModel",
+    "VibeVoiceForConditionalGenerationInference",
+    # ASR models
+    "VibeVoiceASRPreTrainedModel",
+    "VibeVoiceASRModel",
+    "VibeVoiceASRForConditionalGeneration",
+    # Streaming configs
+    "VibeVoiceStreamingConfig",
+    # Streaming models
+    "VibeVoiceStreamingPreTrainedModel",
+    "VibeVoiceStreamingModel",
+    "VibeVoiceStreamingForConditionalGenerationInference",
+    "VibeVoiceGenerationOutput",
+    "BinaryClassifier",
+    "SpeechConnector",
+    "TTS_TEXT_WINDOW_SIZE",
+    "TTS_SPEECH_WINDOW_SIZE",
+    # Shared components
+    "VibeVoiceTokenizerStreamingCache",
+    "VibeVoiceAcousticTokenizerModel",
+    "VibeVoiceSemanticTokenizerModel",
+    "VibeVoiceTextTokenizer",
+    "VibeVoiceTextTokenizerFast",
+    "VibeVoiceDiffusionHead",
+    "AudioStreamer",
+    "AsyncAudioStreamer",
+    # LoRA
+    "load_lora_assets",
+]

vibevoice/modular/configuration_vibevoice.py

@@ -240,9 +240,112 @@ class VibeVoiceConfig(PretrainedConfig):
 
         super().__init__(**kwargs)
 
+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 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", 
+    "VibeVoiceAcousticTokenizerConfig",
-    "VibeVoiceSemanticTokenizerConfig", 
+    "VibeVoiceSemanticTokenizerConfig",
-    "VibeVoiceDiffusionHeadConfig", 
+    "VibeVoiceDiffusionHeadConfig",
-    "VibeVoiceConfig"
+    "VibeVoiceConfig",
+    "VibeVoiceASRConfig"
 ]

vibevoice/modular/configuration_vibevoice_streaming.py

@@ -0,0 +1,96 @@
+""" VibeVoice Streaming model configuration"""
+
+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
+
+logger = logging.get_logger(__name__)
+
+
+class VibeVoiceStreamingConfig(PretrainedConfig):
+    """
+    Configuration class for the VibeVoice Streaming model (0.5B).
+
+    The streaming model differs from the multi-speaker model:
+    - No semantic tokenizer (only acoustic)
+    - Split language model: lower layers for text encoding, upper layers for TTS
+    - Optimized for low-latency real-time generation
+
+    Args:
+        acoustic_tokenizer_config: Configuration for the acoustic tokenizer
+        decoder_config: Configuration for the Qwen2 language model backbone
+        diffusion_head_config: Configuration for the diffusion prediction head
+        tts_backbone_num_hidden_layers: Number of upper transformer layers used for TTS (default: 20)
+    """
+    model_type = "vibevoice_streaming"
+    is_composition = True
+    sub_configs = {
+        "acoustic_tokenizer_config": VibeVoiceAcousticTokenizerConfig,
+        "decoder_config": Qwen2Config,
+        "diffusion_head_config": VibeVoiceDiffusionHeadConfig,
+    }
+    # 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_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
+            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)
+
+__all__ = [
+    "VibeVoiceStreamingConfig"
+]

vibevoice/modular/lora_loading.py

@@ -74,6 +74,7 @@ def _load_diffusion_head(
     diff_dir = adapter_root / "diffusion_head"
     adapter_config = diff_dir / "adapter_config.json"
     adapter_model = diff_dir / "adapter_model.bin"
+    adapter_model_safetensors = diff_dir / "adapter_model.safetensors"
 
     try:
         from peft import PeftModel
@@ -82,7 +83,7 @@ def _load_diffusion_head(
             "peft is required to load diffusion head adapters but is not installed"
         ) from exc
 
-    if adapter_config.exists() and adapter_model.exists():
+    if adapter_config.exists() and (adapter_model.exists() or adapter_model_safetensors.exists()):
         logger.info(f"Loading diffusion head LoRA from {diff_dir}")
         shim = _DiffusionHeadForwardShim(model.model.prediction_head)
         peft_head = PeftModel.from_pretrained(shim, diff_dir)
@@ -113,7 +114,8 @@ def _load_language_model(
 ) -> None:
     config_file = adapter_root / "adapter_config.json"
     bin_file = adapter_root / "adapter_model.bin"
-    if not (config_file.exists() and bin_file.exists()):
+    safe_tensors_file = adapter_root / "adapter_model.safetensors"
+    if not (config_file.exists() and (bin_file.exists() or safe_tensors_file.exists())):
         return
 
     try:

vibevoice/modular/modeling_vibevoice_asr.py

@@ -0,0 +1,520 @@
+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 self.config.use_cache
+
+        # 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,
+            )
+            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()
+            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",
+]

vibevoice/modular/modeling_vibevoice_inference.py

@@ -466,11 +466,13 @@ class VibeVoiceForConditionalGenerationInference(VibeVoicePreTrainedModel, Gener
             model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
             if is_prefill:
                 # we process the speech inputs only during the first generation step
-                prefill_inputs = {
+                prefill_inputs = {}
-                    "speech_tensors": speech_tensors.to(device=device),
+                if speech_tensors is not None:
-                    "speech_masks": speech_masks.to(device),
+                    prefill_inputs["speech_tensors"] = speech_tensors.to(device=device)
-                    "speech_input_mask": speech_input_mask.to(device),
+                if speech_masks is not None:
-                }
+                    prefill_inputs["speech_masks"] = speech_masks.to(device)
+                if speech_input_mask is not None:
+                    prefill_inputs["speech_input_mask"] = speech_input_mask.to(device)
                 is_prefill = False
             else:
                 _ = model_inputs.pop('inputs_embeds', None)

vibevoice/modular/modeling_vibevoice_streaming.py

@@ -0,0 +1,216 @@
+"""
+VibeVoice Streaming Model Architecture (0.5B)
+
+This module implements the streaming-optimized version of VibeVoice for real-time TTS.
+Key differences from the multi-speaker model:
+- No semantic tokenizer (only acoustic)
+- Split language model architecture: lower layers for text, upper layers for TTS
+- Optimized for low-latency generation
+"""
+
+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):
+    """Binary classifier for end-of-speech detection in streaming TTS."""
+
+    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):
+    """Connector module that projects speech latents to language model hidden dimension."""
+
+    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
+
+
+class VibeVoiceStreamingPreTrainedModel(PreTrainedModel):
+    """Base class for VibeVoice Streaming models."""
+
+    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_()
+
+
+class VibeVoiceStreamingModel(VibeVoiceStreamingPreTrainedModel):
+    """
+    VibeVoice Streaming Model for real-time TTS.
+
+    The model uses a split architecture:
+    - language_model: Lower transformer layers for text encoding
+    - tts_language_model: Upper transformer layers for TTS generation
+
+    This separation enables streaming text input and low-latency speech output.
+    """
+
+    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",
+    "BinaryClassifier",
+    "SpeechConnector",
+]

vibevoice/modular/modeling_vibevoice_streaming_inference.py

@@ -0,0 +1,762 @@
+"""
+VibeVoice Streaming Inference Model (0.5B)
+
+This module implements the inference engine for real-time streaming TTS.
+Key features:
+- Window-based text/speech interleaving for streaming
+- Binary EOS classifier for end-of-speech detection
+- Classifier-free guidance for speech quality
+- Audio streaming support
+"""
+
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple, Union, Callable
+from tqdm import tqdm
+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"]
+
+# Window sizes for streaming generation
+TTS_TEXT_WINDOW_SIZE = 5      # Process 5 text tokens at a time
+TTS_SPEECH_WINDOW_SIZE = 6    # Generate 6 speech frames per window
+
+
+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.
+
+    Mainly for the case num_new_tokens > 1 (e.g. a whole text window):
+      - past_key_values: take from current outputs
+      - attention_mask: append num_new_tokens ones
+      - cache_position: advance by creating a range for all new positions
+    """
+
+    # update past_key_values keeping its naming used in model code
+    model_kwargs["past_key_values"] = getattr(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
+    )
+
+    model_kwargs["cache_position"] = torch.arange(model_kwargs["cache_position"][-1] + 1, model_kwargs["cache_position"][-1] + num_new_tokens + 1).to(model_kwargs["cache_position"].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):
+    """
+    VibeVoice Streaming model for conditional generation inference.
+
+    This class handles the full streaming TTS pipeline:
+    1. Prefill text prompt with cached voice embeddings
+    2. Window-based text/speech interleaving
+    3. Diffusion-based speech token generation
+    4. Binary EOS detection for natural stopping
+    """
+
+    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 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,
+        )
+
+    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
+        self._prepare_cache_for_generation(generation_config, model_kwargs, None, 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]:
+        """
+        Generate speech from text using streaming TTS.
+
+        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:
+            inputs: Input tensor (usually from processor)
+            generation_config: Generation configuration
+            audio_streamer: If provided, emits audio chunks during generation
+            tts_text_ids: Full text tokens to stream in windows
+            cfg_scale: Classifier-free guidance scale for speech diffusion (default: 1.0)
+            return_speech: If False, skips audio decode concatenation
+            stop_check_fn: External early-stop hook (returns True to halt)
+            **kwargs: Additional arguments including tokenizer, all_prefilled_outputs, etc.
+
+        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
+
+            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):
+        """
+        Sample speech tokens using diffusion with classifier-free guidance.
+
+        Args:
+            condition: Positive conditioning from TTS LM hidden states
+            neg_condition: Negative conditioning for CFG
+            cfg_scale: Classifier-free guidance scale (higher = more adherence to text)
+
+        Returns:
+            Generated speech latents
+        """
+        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",
+    "VibeVoiceGenerationOutput",
+    "VibeVoiceCausalLMOutputWithPast",
+    "TTS_TEXT_WINDOW_SIZE",
+    "TTS_SPEECH_WINDOW_SIZE",
+]

vibevoice/processor/__init__.py

@@ -0,0 +1,20 @@
+"""
+VibeVoice Processors
+
+This module provides processors for preparing inputs for VibeVoice models:
+- VibeVoiceProcessor: For multi-speaker models (1.5B, 7B)
+- VibeVoiceStreamingProcessor: For streaming model (0.5B)
+"""
+
+from .vibevoice_processor import VibeVoiceProcessor
+from .vibevoice_streaming_processor import VibeVoiceStreamingProcessor
+from .vibevoice_tokenizer_processor import VibeVoiceTokenizerProcessor, AudioNormalizer
+from .vibevoice_asr_processor import VibeVoiceASRProcessor
+
+__all__ = [
+    "VibeVoiceProcessor",
+    "VibeVoiceStreamingProcessor",
+    "VibeVoiceTokenizerProcessor",
+    "AudioNormalizer",
+    "VibeVoiceASRProcessor",
+]

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
+        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=320,
+        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"]

vibevoice/processor/vibevoice_streaming_processor.py

@@ -0,0 +1,426 @@
+"""
+VibeVoice Streaming Processor
+
+This processor handles input preparation for the streaming 0.5B model,
+including text tokenization and cached voice prompt handling.
+"""
+
+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.
+
+    The streaming processor is designed for the 0.5B real-time model that uses pre-computed voice embeddings
+    (.pt files) instead of live audio conditioning.
+
+    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 fallback to Qwen
+        language_model_pretrained_name = config.get("language_model_pretrained_name", None) or kwargs.pop("language_model_pretrained_name", "Qwen/Qwen2.5-0.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.")
+
+        # 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 streaming model uses pre-computed voice embeddings (cached prompts) loaded from .pt files
+        instead of processing audio on-the-fly.
+
+        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.
+                Load this from a .pt file using torch.load().
+            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]
+        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",
+]