Coverage for wrapper/vibevoice/vibevoice_audio_processor.py: 14%
175 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1# mypy: ignore-errors
2# Copy from https://github.com/microsoft/VibeVoice/blob/main/vibevoice/processor/audio_utils.py
3# Copy from https://github.com/microsoft/VibeVoice/blob/main/vibevoice/processor/vibevoice_tokenizer_processor.py
5"""
6Processor class for VibeVoice models.
7"""
9import os
10from typing import List, Optional, Union, Dict, Any, cast
12import numpy as np
13import torch
15from transformers.feature_extraction_utils import FeatureExtractionMixin
16from transformers.utils import logging
18logger = logging.get_logger(__name__)
21class AudioNormalizer:
22 """
23 Audio normalization class for VibeVoice tokenizer.
25 This class provides audio normalization to ensure consistent input levels
26 for the VibeVoice tokenizer while maintaining audio quality.
27 """
29 def __init__(self, target_dB_FS: float = -25, eps: float = 1e-6):
30 """
31 Initialize the audio normalizer.
32 Args:
33 target_dB_FS (float): Target dB FS level for the audio. Default: -25
34 eps (float): Small value to avoid division by zero. Default: 1e-6
35 """
36 self.target_dB_FS = target_dB_FS
37 self.eps = eps
39 def tailor_dB_FS(self, audio: np.ndarray) -> tuple:
40 """
41 Adjust the audio to the target dB FS level.
42 Args:
43 audio (np.ndarray): Input audio signal
44 Returns:
45 tuple: (normalized_audio, rms, scalar)
46 """
47 rms = np.sqrt(np.mean(audio**2))
48 scalar = 10 ** (self.target_dB_FS / 20) / (rms + self.eps)
49 normalized_audio = audio * scalar
50 return normalized_audio, rms, scalar
52 def avoid_clipping(self, audio: np.ndarray, scalar: Optional[float] = None) -> tuple:
53 """
54 Avoid clipping by scaling down if necessary.
55 Args:
56 audio (np.ndarray): Input audio signal
57 scalar (float, optional): Explicit scaling factor
58 Returns:
59 tuple: (normalized_audio, scalar)
60 """
61 if scalar is None:
62 max_val = np.max(np.abs(audio))
63 if max_val > 1.0:
64 scalar = max_val + self.eps
65 else:
66 scalar = 1.0
68 return audio / scalar, scalar
70 def __call__(self, audio: np.ndarray) -> np.ndarray:
71 """
72 Normalize the audio by adjusting to target dB FS and avoiding clipping.
73 Args:
74 audio (np.ndarray): Input audio signal
75 Returns:
76 np.ndarray: Normalized audio signal
77 """
78 # First adjust to target dB FS
79 audio, _, _ = self.tailor_dB_FS(audio)
80 # Then avoid clipping
81 audio, _ = self.avoid_clipping(audio)
82 return audio
85# Change from ProcessorMixin to FeatureExtractionMixin which is designed for single components
86class VibeVoiceTokenizerProcessor(FeatureExtractionMixin):
87 """
88 Processor for VibeVoice acoustic tokenizer models.
90 This processor handles audio preprocessing for VibeVoice models, including:
91 - Audio format conversion (stereo to mono)
92 - Optional audio normalization
93 - Streaming support for infinite-length audio
94 Args:
95 sampling_rate (int, optional): Expected sampling rate. Defaults to 24000.
96 normalize_audio (bool, optional): Whether to normalize audio. Defaults to True.
97 target_dB_FS (float, optional): Target dB FS for normalization. Defaults to -25.
98 eps (float, optional): Small value for numerical stability. Defaults to 1e-6.
99 """
100 model_input_names = ["input_features"]
102 def __init__(
103 self,
104 sampling_rate: int = 24000,
105 normalize_audio: bool = True,
106 target_dB_FS: float = -25,
107 eps: float = 1e-6,
108 **kwargs: Any,
109 ) -> None:
110 super().__init__(**kwargs)
112 self.sampling_rate = sampling_rate
113 self.normalize_audio = normalize_audio
115 # Initialize audio normalizer if needed
116 self.normalizer: Optional[AudioNormalizer]
117 if self.normalize_audio:
118 self.normalizer = AudioNormalizer(target_dB_FS=target_dB_FS, eps=eps)
119 else:
120 self.normalizer = None
122 # Save config
123 self.feature_extractor_dict = {
124 "sampling_rate": sampling_rate,
125 "normalize_audio": normalize_audio,
126 "target_dB_FS": target_dB_FS,
127 "eps": eps,
128 }
130 def _ensure_mono(self, audio: np.ndarray) -> np.ndarray:
131 """
132 Convert stereo audio to mono if needed.
133 Args:
134 audio (np.ndarray): Input audio array
135 Returns:
136 np.ndarray: Mono audio array
137 """
138 if len(audio.shape) == 1:
139 return audio
140 elif len(audio.shape) == 2:
141 if audio.shape[0] == 2: # (2, time)
142 return np.mean(audio, axis=0)
143 elif audio.shape[1] == 2: # (time, 2)
144 return np.mean(audio, axis=1)
145 else:
146 # If one dimension is 1, squeeze it
147 if audio.shape[0] == 1:
148 return audio.squeeze(0)
149 elif audio.shape[1] == 1:
150 return audio.squeeze(1)
151 else:
152 raise ValueError(f"Unexpected audio shape: {audio.shape}")
153 else:
154 raise ValueError(f"Audio should be 1D or 2D, got shape: {audio.shape}")
156 def _process_single_audio(self, audio: Union[np.ndarray, List[float]]) -> np.ndarray:
157 """
158 Process a single audio array.
159 Args:
160 audio: Single audio input
161 Returns:
162 np.ndarray: Processed audio
163 """
164 # Convert to numpy array
165 if not isinstance(audio, np.ndarray):
166 audio = np.array(audio, dtype=np.float32)
167 else:
168 audio = audio.astype(np.float32)
170 # Ensure mono
171 audio = self._ensure_mono(audio)
173 # Normalize if requested
174 if self.normalize_audio and self.normalizer is not None:
175 audio = self.normalizer(audio)
177 return audio
179 def __call__(
180 self,
181 audio: Optional[Union[str, np.ndarray, List[float], List[np.ndarray], List[List[float]], List[str]]] = None,
182 sampling_rate: Optional[int] = None,
183 return_tensors: Optional[str] = None,
184 **kwargs: Any,
185 ) -> Dict[str, Any]:
186 """
187 Process audio for VibeVoice models.
188 Args:
189 audio: Audio input(s) to process. Can be:
190 - str: Path to audio file
191 - np.ndarray: Audio array
192 - List[float]: Audio as list of floats
193 - List[np.ndarray]: Batch of audio arrays
194 - List[str]: Batch of audio file paths
195 sampling_rate (int, optional): Sampling rate of the input audio
196 return_tensors (str, optional): Return format ('pt' for PyTorch, 'np' for NumPy)
197 Returns:
198 dict: Processed audio inputs with keys:
199 - input_features: Audio tensor(s) ready for the model
200 """
201 if audio is None:
202 raise ValueError("Audio input is required")
204 # Validate sampling rate
205 if sampling_rate is not None and sampling_rate != self.sampling_rate:
206 logger.warning(
207 f"Input sampling rate ({sampling_rate}) differs from expected "
208 f"sampling rate ({self.sampling_rate}). Please resample your audio."
209 )
211 # Handle different input types
212 if isinstance(audio, str):
213 # Single audio file path
214 audio = self._load_audio_from_path(audio)
215 is_batched = False
216 elif isinstance(audio, list):
217 if len(audio) == 0:
218 raise ValueError("Empty audio list provided")
220 # Check if it's a list of file paths
221 if all(isinstance(item, str) for item in audio):
222 # Batch of audio file paths
223 audio = [self._load_audio_from_path(path) for path in audio if isinstance(path, str)]
224 is_batched = True
225 else:
226 # Check if it's batched audio arrays
227 is_batched = isinstance(audio[0], (np.ndarray, list))
228 else:
229 # Single audio array or list
230 is_batched = False
232 # Process audio
233 if is_batched:
234 processed_audio = [self._process_single_audio(cast(Union[np.ndarray, List[float]], a)) for a in audio]
235 else:
236 processed_audio = [self._process_single_audio(cast(Union[np.ndarray, List[float]], audio))]
238 # Convert to tensors if requested
239 input_features: Union[torch.Tensor, np.ndarray, List[np.ndarray]]
240 if return_tensors == "pt":
241 if len(processed_audio) == 1:
242 # Create a proper batch dimension (B, T)
243 input_features = torch.from_numpy(processed_audio[0]).unsqueeze(0).unsqueeze(1)
244 else:
245 # For batched input with different lengths, create a batch properly
246 input_features = torch.stack([torch.from_numpy(a) for a in processed_audio]).unsqueeze(1)
247 elif return_tensors == "np":
248 if len(processed_audio) == 1:
249 input_features = processed_audio[0][np.newaxis, np.newaxis, :]
250 else:
251 input_features = np.stack(processed_audio)[:, np.newaxis, :]
252 else:
253 input_features = processed_audio[0] if len(processed_audio) == 1 else processed_audio
255 outputs = {
256 "audio": input_features, # Use "audio" instead of "input_features"
257 }
259 return outputs
261 def _load_audio_from_path(self, audio_path: str) -> np.ndarray:
262 """
263 Load audio from file path.
264 Args:
265 audio_path (str): Path to audio file
266 Returns:
267 np.ndarray: Loaded audio array
268 """
269 # Get file extension to determine loading method
270 file_ext = os.path.splitext(audio_path)[1].lower()
272 if file_ext in ['.wav', '.mp3', '.flac', '.m4a', '.ogg']:
273 # Audio file - use librosa
274 import librosa
275 audio_array, sr = librosa.load(
276 audio_path,
277 sr=self.sampling_rate,
278 mono=True
279 )
280 return audio_array
281 elif file_ext == '.pt':
282 # PyTorch tensor file
283 audio_tensor = torch.load(audio_path, map_location='cpu', weights_only=True).squeeze()
284 if isinstance(audio_tensor, torch.Tensor):
285 audio_array = audio_tensor.numpy()
286 else:
287 audio_array = np.array(audio_tensor)
288 return audio_array.astype(np.float32)
289 elif file_ext == '.npy':
290 # NumPy file
291 audio_array = np.load(audio_path)
292 return audio_array.astype(np.float32)
293 else:
294 raise ValueError(
295 f"Unsupported file format: {file_ext}. "
296 f"Supported formats: .wav, .mp3, .flac, .m4a, .ogg, .pt, .npy, .npz"
297 )
299 def preprocess_audio(
300 self,
301 audio_path_or_array: Union[str, np.ndarray],
302 normalize: Optional[bool] = None,
303 ) -> np.ndarray:
304 """
305 Convenience method to preprocess audio from file path or array.
306 This method is kept for backward compatibility but __call__ is recommended.
307 Args:
308 audio_path_or_array: Path to audio file or numpy array
309 normalize: Whether to normalize (overrides default setting)
310 Returns:
311 np.ndarray: Preprocessed audio array
312 """
313 if isinstance(audio_path_or_array, str):
314 audio_array = self._load_audio_from_path(audio_path_or_array)
315 else:
316 audio_array = np.array(audio_path_or_array, dtype=np.float32)
318 # Override normalization setting if specified
319 original_normalize = self.normalize_audio
320 if normalize is not None:
321 self.normalize_audio = normalize
323 try:
324 processed = self._process_single_audio(audio_array)
325 finally:
326 # Restore original setting
327 self.normalize_audio = original_normalize
329 return processed
331 # Override to_dict method for configuration saving
332 def to_dict(self) -> Dict[str, Any]:
333 """
334 Convert the object to a dict containing all attributes needed for serialization.
335 """
336 return self.feature_extractor_dict
338 def save_audio(
339 self,
340 audio: Union[torch.Tensor, np.ndarray, List[Union[torch.Tensor, np.ndarray]]],
341 output_path: str = "output.wav",
342 sampling_rate: Optional[int] = None,
343 normalize: bool = False,
344 batch_prefix: str = "audio_",
345 ) -> List[str]:
346 """
347 Save audio data to WAV file(s).
348 Args:
349 audio: Audio data to save. Can be:
350 - torch.Tensor: PyTorch tensor with shape (B, C, T) or (B, T) or (T)
351 - np.ndarray: NumPy array with shape (B, C, T) or (B, T) or (T)
352 - List of tensors or arrays
353 output_path: Path where to save the audio. If saving multiple files,
354 this is treated as a directory and individual files will be saved inside.
355 sampling_rate: Sampling rate for the saved audio. Defaults to the processor's rate.
356 normalize: Whether to normalize audio before saving.
357 batch_prefix: Prefix for batch files when saving multiple audios.
358 Returns:
359 List[str]: Paths to the saved audio files.
360 """
361 if sampling_rate is None:
362 sampling_rate = self.sampling_rate
364 try:
365 import soundfile as sf
366 except ImportError:
367 raise ImportError(
368 "soundfile is required to save audio files. "
369 "Install it with: pip install soundfile"
370 )
372 # Ensure audio is in the right format
373 audio_np: Union[np.ndarray, List[Any]]
374 if isinstance(audio, torch.Tensor):
375 # Convert PyTorch tensor to numpy
376 audio_np = audio.float().detach().cpu().numpy()
377 elif isinstance(audio, np.ndarray):
378 audio_np = audio
379 elif isinstance(audio, list):
380 # Handle list of tensors or arrays
381 if all(isinstance(a, torch.Tensor) for a in audio):
382 audio_np = [cast(torch.Tensor, a).float().detach().cpu().numpy() for a in audio]
383 else:
384 audio_np = audio
385 else:
386 raise ValueError(f"Unsupported audio type: {type(audio)}")
388 saved_paths = []
390 # Handle based on shape or type
391 if isinstance(audio_np, list):
392 # Multiple separate audios to save
393 output_dir = output_path
395 # Ensure output directory exists
396 os.makedirs(output_dir, exist_ok=True)
398 # Save each audio
399 for i, audio_item in enumerate(audio_np):
400 audio_item = self._prepare_audio_for_save(audio_item, normalize)
401 file_path = os.path.join(output_dir, f"{batch_prefix}{i}.wav")
402 sf.write(file_path, audio_item, sampling_rate)
403 saved_paths.append(file_path)
404 else:
405 # Handle different dimensions
406 if len(audio_np.shape) >= 3: # (B, C, T) or similar
407 # Get batch size
408 batch_size = audio_np.shape[0]
410 if batch_size > 1:
411 # Multiple audios in a batch
412 output_dir = output_path
414 # Ensure output directory exists
415 os.makedirs(output_dir, exist_ok=True)
417 # Save each audio in the batch
418 for i in range(batch_size):
419 # Extract single audio and remove channel dim if present
420 single_audio = audio_np[i]
421 if len(single_audio.shape) > 1:
422 if single_audio.shape[0] == 1: # (1, T)
423 single_audio = single_audio.squeeze(0)
425 single_audio = self._prepare_audio_for_save(single_audio, normalize)
426 file_path = os.path.join(output_dir, f"{batch_prefix}{i}.wav")
427 sf.write(file_path, single_audio, sampling_rate)
428 saved_paths.append(file_path)
429 else:
430 # Single audio with batch and channel dims
431 audio_item = audio_np.squeeze() # Remove batch and channel dimensions
432 audio_item = self._prepare_audio_for_save(audio_item, normalize)
433 sf.write(output_path, audio_item, sampling_rate)
434 saved_paths.append(output_path)
435 else:
436 # Single audio without batch dimension
437 audio_item = self._prepare_audio_for_save(audio_np, normalize)
438 sf.write(output_path, audio_item, sampling_rate)
439 saved_paths.append(output_path)
441 return saved_paths
443 def _prepare_audio_for_save(self, audio: np.ndarray, normalize: bool) -> np.ndarray:
444 """
445 Prepare audio for saving by ensuring it's the right shape and optionally normalizing.
447 Args:
448 audio: Audio data as numpy array
449 normalize: Whether to normalize audio
451 Returns:
452 np.ndarray: Processed audio ready for saving
453 """
454 # Ensure right dimensionality
455 if len(audio.shape) > 1 and audio.shape[0] == 1: # (1, T)
456 audio = audio.squeeze(0)
458 # Normalize if requested
459 if normalize:
460 max_val = np.abs(audio).max()
461 if max_val > 0:
462 audio = audio / max_val
464 return audio
467__all__ = ["VibeVoiceTokenizerProcessor", "AudioNormalizer"]