Add VibeVoice ASR (#48)

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).

vibevoice/modular/__init__.py

@@ -12,6 +12,12 @@ from .configuration_vibevoice import (
     VibeVoiceAcousticTokenizerConfig,
     VibeVoiceSemanticTokenizerConfig,
     VibeVoiceDiffusionHeadConfig,
+    VibeVoiceASRConfig,
+)
+from .modeling_vibevoice_asr import (
+    VibeVoiceASRPreTrainedModel,
+    VibeVoiceASRModel,
+    VibeVoiceASRForConditionalGeneration,
 )
 from .modeling_vibevoice import (
     VibeVoicePreTrainedModel,
@@ -67,10 +73,15 @@ __all__ = [
     "VibeVoiceAcousticTokenizerConfig",
     "VibeVoiceSemanticTokenizerConfig",
     "VibeVoiceDiffusionHeadConfig",
+    "VibeVoiceASRConfig",
     # Multi-speaker models
     "VibeVoicePreTrainedModel",
     "VibeVoiceModel",
     "VibeVoiceForConditionalGenerationInference",
+    # ASR models
+    "VibeVoiceASRPreTrainedModel",
+    "VibeVoiceASRModel",
+    "VibeVoiceASRForConditionalGeneration",
     # Streaming configs
     "VibeVoiceStreamingConfig",
     # Streaming models

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",
     "VibeVoiceSemanticTokenizerConfig",
     "VibeVoiceDiffusionHeadConfig",
-    "VibeVoiceConfig"
+    "VibeVoiceConfig",
+    "VibeVoiceASRConfig"
 ]

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/processor/__init__.py

@@ -9,10 +9,12 @@ This module provides processors for preparing inputs for VibeVoice models:
 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"]