code/VibeVoice/vllm_plugin/inputs.py

"""Audio input mapper for vLLM multimodal pipeline.

This module handles audio data loading and preprocessing for VibeVoice ASR inference.
It converts various audio input formats (path, bytes, numpy array) into tensors
that can be processed by the VibeVoice model.
"""
import os
import logging
import torch
import numpy as np
from typing import Union, List
from vllm.multimodal.inputs import MultiModalInputs
from vibevoice.processor.audio_utils import load_audio_use_ffmpeg, load_audio_bytes_use_ffmpeg, AudioNormalizer

logger = logging.getLogger(__name__)

# Maximum audio duration in seconds. Default 3660s (61 minutes) matches the
# model's designed capacity. Override via environment variable to guard against
# OOM on GPUs with less VRAM. See https://github.com/microsoft/VibeVoice/issues/210
_MAX_AUDIO_DURATION = float(os.environ.get("VIBEVOICE_MAX_AUDIO_DURATION", "3660"))


def load_audio(audio_path: str, target_sr: int = 24000) -> np.ndarray:
    """Load and normalize audio from file path.
    
    Args:
        audio_path: Path to audio file
        target_sr: Target sample rate (default 24kHz for VibeVoice)
        
    Returns:
        Normalized audio waveform as numpy array
    """
    # Load with FFmpeg (handles various formats)
    audio, sr = load_audio_use_ffmpeg(audio_path, resample=True, target_sr=target_sr)
    
    # Normalize audio
    normalizer = AudioNormalizer()
    audio = normalizer(audio)
    
    return audio


def vibevoice_audio_input_mapper(ctx, data: Union[str, bytes, np.ndarray, List[str]]) -> MultiModalInputs:
    """Map audio input data to vLLM MultiModalInputs format.
    
    This function is registered as the input mapper for VibeVoice audio processing.
    It handles multiple input formats and converts them to normalized tensors.
    
    Args:
        ctx: vLLM context (unused)
        data: Audio data in one of these formats:
            - str: Path to audio file
            - bytes: Raw audio bytes (any format FFmpeg supports)
            - np.ndarray: Pre-loaded audio waveform
            - List[str]: List of audio paths (only first is used)
            
    Returns:
        MultiModalInputs containing:
            - audio: Audio tensor (float32)
            - audio_length: Length of audio in samples
    """
    # Handle list input (take first item)
    if isinstance(data, list):
        data = data[0]

    audio_waveform = None
    
    if isinstance(data, str):
        # Load from file path
        audio_waveform = load_audio(data)
        
    elif isinstance(data, bytes):
        # Decode bytes directly via ffmpeg stdin pipe to avoid temp-file IO
        audio_waveform, _sr = load_audio_bytes_use_ffmpeg(data, resample=True, target_sr=24000)
        normalizer = AudioNormalizer()
        audio_waveform = normalizer(audio_waveform)
                
    elif isinstance(data, np.ndarray):
        # Already loaded numpy array
        audio_waveform = data
    else:
        raise ValueError(f"Unsupported audio data type: {type(data)}")

    # Validate audio duration before tensor conversion to catch OOM early
    duration_sec = len(audio_waveform) / 24000
    if duration_sec > _MAX_AUDIO_DURATION:
        raise ValueError(
            f"Audio duration ({duration_sec:.1f}s) exceeds the configured "
            f"limit ({_MAX_AUDIO_DURATION:.0f}s). Set the "
            f"VIBEVOICE_MAX_AUDIO_DURATION environment variable to adjust "
            f"this limit, or use shorter audio."
        )

    # Convert to tensor
    audio_tensor = torch.from_numpy(audio_waveform).float()
    audio_length = audio_tensor.shape[0]
    
    return MultiModalInputs({
        "audio": audio_tensor,
        "audio_length": audio_length
    })