483 lines
No EOL
17 KiB
Python
483 lines
No EOL
17 KiB
Python
"""
|
|
Processor class for VibeVoice models.
|
|
"""
|
|
|
|
import os
|
|
import json
|
|
import warnings
|
|
from typing import List, Optional, Union, Dict, Any
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from transformers.feature_extraction_utils import FeatureExtractionMixin
|
|
from transformers.utils import logging
|
|
|
|
logger = logging.get_logger(__name__)
|
|
|
|
|
|
class AudioNormalizer:
|
|
"""
|
|
Audio normalization class for VibeVoice tokenizer.
|
|
|
|
This class provides audio normalization to ensure consistent input levels
|
|
for the VibeVoice tokenizer while maintaining audio quality.
|
|
"""
|
|
|
|
def __init__(self, target_dB_FS: float = -25, eps: float = 1e-6):
|
|
"""
|
|
Initialize the audio normalizer.
|
|
|
|
Args:
|
|
target_dB_FS (float): Target dB FS level for the audio. Default: -25
|
|
eps (float): Small value to avoid division by zero. Default: 1e-6
|
|
"""
|
|
self.target_dB_FS = target_dB_FS
|
|
self.eps = eps
|
|
|
|
def tailor_dB_FS(self, audio: np.ndarray) -> tuple:
|
|
"""
|
|
Adjust the audio to the target dB FS level.
|
|
|
|
Args:
|
|
audio (np.ndarray): Input audio signal
|
|
|
|
Returns:
|
|
tuple: (normalized_audio, rms, scalar)
|
|
"""
|
|
rms = np.sqrt(np.mean(audio**2))
|
|
scalar = 10 ** (self.target_dB_FS / 20) / (rms + self.eps)
|
|
normalized_audio = audio * scalar
|
|
return normalized_audio, rms, scalar
|
|
|
|
def avoid_clipping(self, audio: np.ndarray, scalar: Optional[float] = None) -> tuple:
|
|
"""
|
|
Avoid clipping by scaling down if necessary.
|
|
|
|
Args:
|
|
audio (np.ndarray): Input audio signal
|
|
scalar (float, optional): Explicit scaling factor
|
|
|
|
Returns:
|
|
tuple: (normalized_audio, scalar)
|
|
"""
|
|
if scalar is None:
|
|
max_val = np.max(np.abs(audio))
|
|
if max_val > 1.0:
|
|
scalar = max_val + self.eps
|
|
else:
|
|
scalar = 1.0
|
|
|
|
return audio / scalar, scalar
|
|
|
|
def __call__(self, audio: np.ndarray) -> np.ndarray:
|
|
"""
|
|
Normalize the audio by adjusting to target dB FS and avoiding clipping.
|
|
|
|
Args:
|
|
audio (np.ndarray): Input audio signal
|
|
|
|
Returns:
|
|
np.ndarray: Normalized audio signal
|
|
"""
|
|
# First adjust to target dB FS
|
|
audio, _, _ = self.tailor_dB_FS(audio)
|
|
# Then avoid clipping
|
|
audio, _ = self.avoid_clipping(audio)
|
|
return audio
|
|
|
|
|
|
# Change from ProcessorMixin to FeatureExtractionMixin which is designed for single components
|
|
class VibeVoiceTokenizerProcessor(FeatureExtractionMixin):
|
|
"""
|
|
Processor for VibeVoice acoustic tokenizer models.
|
|
|
|
This processor handles audio preprocessing for VibeVoice models, including:
|
|
- Audio format conversion (stereo to mono)
|
|
- Optional audio normalization
|
|
- Streaming support for infinite-length audio
|
|
|
|
Args:
|
|
sampling_rate (int, optional): Expected sampling rate. Defaults to 24000.
|
|
normalize_audio (bool, optional): Whether to normalize audio. Defaults to True.
|
|
target_dB_FS (float, optional): Target dB FS for normalization. Defaults to -25.
|
|
eps (float, optional): Small value for numerical stability. Defaults to 1e-6.
|
|
"""
|
|
model_input_names = ["input_features"]
|
|
|
|
def __init__(
|
|
self,
|
|
sampling_rate: int = 24000,
|
|
normalize_audio: bool = True,
|
|
target_dB_FS: float = -25,
|
|
eps: float = 1e-6,
|
|
**kwargs,
|
|
):
|
|
super().__init__(**kwargs)
|
|
|
|
self.sampling_rate = sampling_rate
|
|
self.normalize_audio = normalize_audio
|
|
|
|
# Initialize audio normalizer if needed
|
|
if self.normalize_audio:
|
|
self.normalizer = AudioNormalizer(target_dB_FS=target_dB_FS, eps=eps)
|
|
else:
|
|
self.normalizer = None
|
|
|
|
# Save config
|
|
self.feature_extractor_dict = {
|
|
"sampling_rate": sampling_rate,
|
|
"normalize_audio": normalize_audio,
|
|
"target_dB_FS": target_dB_FS,
|
|
"eps": eps,
|
|
}
|
|
|
|
def _ensure_mono(self, audio: np.ndarray) -> np.ndarray:
|
|
"""
|
|
Convert stereo audio to mono if needed.
|
|
|
|
Args:
|
|
audio (np.ndarray): Input audio array
|
|
|
|
Returns:
|
|
np.ndarray: Mono audio array
|
|
"""
|
|
if len(audio.shape) == 1:
|
|
return audio
|
|
elif len(audio.shape) == 2:
|
|
if audio.shape[0] == 2: # (2, time)
|
|
return np.mean(audio, axis=0)
|
|
elif audio.shape[1] == 2: # (time, 2)
|
|
return np.mean(audio, axis=1)
|
|
else:
|
|
# If one dimension is 1, squeeze it
|
|
if audio.shape[0] == 1:
|
|
return audio.squeeze(0)
|
|
elif audio.shape[1] == 1:
|
|
return audio.squeeze(1)
|
|
else:
|
|
raise ValueError(f"Unexpected audio shape: {audio.shape}")
|
|
else:
|
|
raise ValueError(f"Audio should be 1D or 2D, got shape: {audio.shape}")
|
|
|
|
def _process_single_audio(self, audio: Union[np.ndarray, List[float]]) -> np.ndarray:
|
|
"""
|
|
Process a single audio array.
|
|
|
|
Args:
|
|
audio: Single audio input
|
|
|
|
Returns:
|
|
np.ndarray: Processed audio
|
|
"""
|
|
# Convert to numpy array
|
|
if not isinstance(audio, np.ndarray):
|
|
audio = np.array(audio, dtype=np.float32)
|
|
else:
|
|
audio = audio.astype(np.float32)
|
|
|
|
# Ensure mono
|
|
audio = self._ensure_mono(audio)
|
|
|
|
# Normalize if requested
|
|
if self.normalize_audio and self.normalizer is not None:
|
|
audio = self.normalizer(audio)
|
|
|
|
return audio
|
|
|
|
def __call__(
|
|
self,
|
|
audio: Union[str, np.ndarray, List[float], List[np.ndarray], List[List[float]], List[str]] = None,
|
|
sampling_rate: Optional[int] = None,
|
|
return_tensors: Optional[str] = None,
|
|
**kwargs,
|
|
):
|
|
"""
|
|
Process audio for VibeVoice models.
|
|
|
|
Args:
|
|
audio: Audio input(s) to process. Can be:
|
|
- str: Path to audio file
|
|
- np.ndarray: Audio array
|
|
- List[float]: Audio as list of floats
|
|
- List[np.ndarray]: Batch of audio arrays
|
|
- List[str]: Batch of audio file paths
|
|
sampling_rate (int, optional): Sampling rate of the input audio
|
|
return_tensors (str, optional): Return format ('pt' for PyTorch, 'np' for NumPy)
|
|
|
|
Returns:
|
|
dict: Processed audio inputs with keys:
|
|
- input_features: Audio tensor(s) ready for the model
|
|
"""
|
|
if audio is None:
|
|
raise ValueError("Audio input is required")
|
|
|
|
# Validate sampling rate
|
|
if sampling_rate is not None and sampling_rate != self.sampling_rate:
|
|
logger.warning(
|
|
f"Input sampling rate ({sampling_rate}) differs from expected "
|
|
f"sampling rate ({self.sampling_rate}). Please resample your audio."
|
|
)
|
|
|
|
# Handle different input types
|
|
if isinstance(audio, str):
|
|
# Single audio file path
|
|
audio = self._load_audio_from_path(audio)
|
|
is_batched = False
|
|
elif isinstance(audio, list):
|
|
if len(audio) == 0:
|
|
raise ValueError("Empty audio list provided")
|
|
|
|
# Check if it's a list of file paths
|
|
if all(isinstance(item, str) for item in audio):
|
|
# Batch of audio file paths
|
|
audio = [self._load_audio_from_path(path) for path in audio]
|
|
is_batched = True
|
|
else:
|
|
# Check if it's batched audio arrays
|
|
is_batched = isinstance(audio[0], (np.ndarray, list))
|
|
else:
|
|
# Single audio array or list
|
|
is_batched = False
|
|
|
|
# Process audio
|
|
if is_batched:
|
|
processed_audio = [self._process_single_audio(a) for a in audio]
|
|
else:
|
|
processed_audio = [self._process_single_audio(audio)]
|
|
|
|
# Convert to tensors if requested
|
|
if return_tensors == "pt":
|
|
if len(processed_audio) == 1:
|
|
# Create a proper batch dimension (B, T)
|
|
input_features = torch.from_numpy(processed_audio[0]).unsqueeze(0).unsqueeze(1)
|
|
else:
|
|
# For batched input with different lengths, create a batch properly
|
|
input_features = torch.stack([torch.from_numpy(a) for a in processed_audio]).unsqueeze(1)
|
|
elif return_tensors == "np":
|
|
if len(processed_audio) == 1:
|
|
input_features = processed_audio[0][np.newaxis, np.newaxis, :]
|
|
else:
|
|
input_features = np.stack(processed_audio)[:, np.newaxis, :]
|
|
else:
|
|
input_features = processed_audio[0] if len(processed_audio) == 1 else processed_audio
|
|
|
|
outputs = {
|
|
"audio": input_features, # Use "audio" instead of "input_features"
|
|
}
|
|
|
|
return outputs
|
|
|
|
def _load_audio_from_path(self, audio_path: str) -> np.ndarray:
|
|
"""
|
|
Load audio from file path.
|
|
|
|
Args:
|
|
audio_path (str): Path to audio file
|
|
|
|
Returns:
|
|
np.ndarray: Loaded audio array
|
|
"""
|
|
# Get file extension to determine loading method
|
|
file_ext = os.path.splitext(audio_path)[1].lower()
|
|
|
|
if file_ext in ['.wav', '.mp3', '.flac', '.m4a', '.ogg']:
|
|
# Audio file - use librosa
|
|
import librosa
|
|
audio_array, sr = librosa.load(
|
|
audio_path,
|
|
sr=self.sampling_rate,
|
|
mono=True
|
|
)
|
|
return audio_array
|
|
elif file_ext == '.pt':
|
|
# PyTorch tensor file
|
|
audio_tensor = torch.load(audio_path, map_location='cpu').squeeze()
|
|
if isinstance(audio_tensor, torch.Tensor):
|
|
audio_array = audio_tensor.numpy()
|
|
else:
|
|
audio_array = np.array(audio_tensor)
|
|
return audio_array.astype(np.float32)
|
|
elif file_ext == '.npy':
|
|
# NumPy file
|
|
audio_array = np.load(audio_path)
|
|
return audio_array.astype(np.float32)
|
|
else:
|
|
raise ValueError(
|
|
f"Unsupported file format: {file_ext}. "
|
|
f"Supported formats: .wav, .mp3, .flac, .m4a, .ogg, .pt, .npy, .npz"
|
|
)
|
|
|
|
def preprocess_audio(
|
|
self,
|
|
audio_path_or_array: Union[str, np.ndarray],
|
|
normalize: Optional[bool] = None,
|
|
) -> np.ndarray:
|
|
"""
|
|
Convenience method to preprocess audio from file path or array.
|
|
This method is kept for backward compatibility but __call__ is recommended.
|
|
|
|
Args:
|
|
audio_path_or_array: Path to audio file or numpy array
|
|
normalize: Whether to normalize (overrides default setting)
|
|
|
|
Returns:
|
|
np.ndarray: Preprocessed audio array
|
|
"""
|
|
if isinstance(audio_path_or_array, str):
|
|
audio_array = self._load_audio_from_path(audio_path_or_array)
|
|
else:
|
|
audio_array = np.array(audio_path_or_array, dtype=np.float32)
|
|
|
|
# Override normalization setting if specified
|
|
original_normalize = self.normalize_audio
|
|
if normalize is not None:
|
|
self.normalize_audio = normalize
|
|
|
|
try:
|
|
processed = self._process_single_audio(audio_array)
|
|
finally:
|
|
# Restore original setting
|
|
self.normalize_audio = original_normalize
|
|
|
|
return processed
|
|
|
|
# Override to_dict method for configuration saving
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
"""
|
|
Convert the object to a dict containing all attributes needed for serialization.
|
|
"""
|
|
return self.feature_extractor_dict
|
|
|
|
def save_audio(
|
|
self,
|
|
audio: Union[torch.Tensor, np.ndarray, List[Union[torch.Tensor, np.ndarray]]],
|
|
output_path: str = "output.wav",
|
|
sampling_rate: Optional[int] = None,
|
|
normalize: bool = False,
|
|
batch_prefix: str = "audio_",
|
|
):
|
|
"""
|
|
Save audio data to WAV file(s).
|
|
|
|
Args:
|
|
audio: Audio data to save. Can be:
|
|
- torch.Tensor: PyTorch tensor with shape (B, C, T) or (B, T) or (T)
|
|
- np.ndarray: NumPy array with shape (B, C, T) or (B, T) or (T)
|
|
- List of tensors or arrays
|
|
output_path: Path where to save the audio. If saving multiple files,
|
|
this is treated as a directory and individual files will be saved inside.
|
|
sampling_rate: Sampling rate for the saved audio. Defaults to the processor's rate.
|
|
normalize: Whether to normalize audio before saving.
|
|
batch_prefix: Prefix for batch files when saving multiple audios.
|
|
|
|
Returns:
|
|
List[str]: Paths to the saved audio files.
|
|
"""
|
|
if sampling_rate is None:
|
|
sampling_rate = self.sampling_rate
|
|
|
|
try:
|
|
import soundfile as sf
|
|
except ImportError:
|
|
raise ImportError(
|
|
"soundfile is required to save audio files. "
|
|
"Install it with: pip install soundfile"
|
|
)
|
|
|
|
# Ensure audio is in the right format
|
|
if isinstance(audio, torch.Tensor):
|
|
# Convert PyTorch tensor to numpy
|
|
audio_np = audio.float().detach().cpu().numpy()
|
|
elif isinstance(audio, np.ndarray):
|
|
audio_np = audio
|
|
elif isinstance(audio, list):
|
|
# Handle list of tensors or arrays
|
|
if all(isinstance(a, torch.Tensor) for a in audio):
|
|
audio_np = [a.float().detach().cpu().numpy() for a in audio]
|
|
else:
|
|
audio_np = audio
|
|
else:
|
|
raise ValueError(f"Unsupported audio type: {type(audio)}")
|
|
|
|
saved_paths = []
|
|
|
|
# Handle based on shape or type
|
|
if isinstance(audio_np, list):
|
|
# Multiple separate audios to save
|
|
output_dir = output_path
|
|
|
|
# Ensure output directory exists
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
# Save each audio
|
|
for i, audio_item in enumerate(audio_np):
|
|
audio_item = self._prepare_audio_for_save(audio_item, normalize)
|
|
file_path = os.path.join(output_dir, f"{batch_prefix}{i}.wav")
|
|
sf.write(file_path, audio_item, sampling_rate)
|
|
saved_paths.append(file_path)
|
|
|
|
else:
|
|
# Handle different dimensions
|
|
if len(audio_np.shape) >= 3: # (B, C, T) or similar
|
|
# Get batch size
|
|
batch_size = audio_np.shape[0]
|
|
|
|
if batch_size > 1:
|
|
# Multiple audios in a batch
|
|
output_dir = output_path
|
|
|
|
# Ensure output directory exists
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
# Save each audio in the batch
|
|
for i in range(batch_size):
|
|
# Extract single audio and remove channel dim if present
|
|
single_audio = audio_np[i]
|
|
if len(single_audio.shape) > 1:
|
|
if single_audio.shape[0] == 1: # (1, T)
|
|
single_audio = single_audio.squeeze(0)
|
|
|
|
single_audio = self._prepare_audio_for_save(single_audio, normalize)
|
|
file_path = os.path.join(output_dir, f"{batch_prefix}{i}.wav")
|
|
sf.write(file_path, single_audio, sampling_rate)
|
|
saved_paths.append(file_path)
|
|
else:
|
|
# Single audio with batch and channel dims
|
|
audio_item = audio_np.squeeze() # Remove batch and channel dimensions
|
|
audio_item = self._prepare_audio_for_save(audio_item, normalize)
|
|
sf.write(output_path, audio_item, sampling_rate)
|
|
saved_paths.append(output_path)
|
|
else:
|
|
# Single audio without batch dimension
|
|
audio_item = self._prepare_audio_for_save(audio_np, normalize)
|
|
sf.write(output_path, audio_item, sampling_rate)
|
|
saved_paths.append(output_path)
|
|
|
|
return saved_paths
|
|
|
|
def _prepare_audio_for_save(self, audio: np.ndarray, normalize: bool) -> np.ndarray:
|
|
"""
|
|
Prepare audio for saving by ensuring it's the right shape and optionally normalizing.
|
|
|
|
Args:
|
|
audio: Audio data as numpy array
|
|
normalize: Whether to normalize audio
|
|
|
|
Returns:
|
|
np.ndarray: Processed audio ready for saving
|
|
"""
|
|
# Ensure right dimensionality
|
|
if len(audio.shape) > 1 and audio.shape[0] == 1: # (1, T)
|
|
audio = audio.squeeze(0)
|
|
|
|
# Normalize if requested
|
|
if normalize:
|
|
max_val = np.abs(audio).max()
|
|
if max_val > 0:
|
|
audio = audio / max_val
|
|
|
|
return audio
|
|
|
|
|
|
__all__ = ["VibeVoiceTokenizerProcessor", "AudioNormalizer"] |