Coverage for wrapper/vibevoice/vibevoice_processor.py: 18%
304 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/vibevoice_processor.py
4import os
5import json
6import re
7import math
9from typing import List
10from typing import Optional
11from typing import Union
12from typing import Dict
13from typing import Any
14from typing import Tuple
15from typing import cast
17import numpy as np
18import torch
20from transformers.tokenization_utils_base import BatchEncoding
21from transformers.tokenization_utils_base import PaddingStrategy
22from transformers.tokenization_utils_base import PreTokenizedInput
23from transformers.tokenization_utils_base import TextInput
24from transformers.tokenization_utils_base import TruncationStrategy
25from transformers.utils import TensorType
26from transformers.utils import logging
28from vibevoice_audio_processor import AudioNormalizer
29from vibevoice_audio_processor import VibeVoiceTokenizerProcessor
31from modular_vibevoice_text_tokenizer import VibeVoiceTextTokenizerFast
33logger = logging.get_logger(__name__)
36class VibeVoiceProcessor:
37 r"""
38 Constructs a VibeVoice processor which wraps a VibeVoice tokenizer and audio processor into a single processor.
39 [`VibeVoiceProcessor`] offers all the functionalities of [`VibeVoiceTokenizer`] and [`VibeVoiceTokenizerProcessor`].
40 See the [`~VibeVoiceProcessor.__call__`] and [`~VibeVoiceProcessor.decode`] for more information.
41 Args:
42 tokenizer (`VibeVoiceTextTokenizer` or `VibeVoiceTextTokenizerFast`):
43 The tokenizer for text processing.
44 audio_processor (`VibeVoiceTokenizerProcessor`):
45 The audio processor for speech processing.
46 speech_tok_compress_ratio (`int`, *optional*, defaults to 3200):
47 The compression ratio for speech tokenization.
48 db_normalize (`bool`, *optional*, defaults to True):
49 Whether to apply decibel normalization to audio inputs.
50 """
52 def __init__(
53 self,
54 tokenizer: Optional[Any] = None,
55 audio_processor: Optional[Any] = None,
56 speech_tok_compress_ratio: int = 3200,
57 db_normalize: bool = True,
58 **kwargs: Any
59 ) -> None:
60 self.tokenizer = tokenizer
61 self.audio_processor = audio_processor
62 self.speech_tok_compress_ratio = speech_tok_compress_ratio
63 self.db_normalize = db_normalize
64 self.audio_normalizer = AudioNormalizer() if db_normalize else None
65 self.system_prompt = (
66 " Transform the text provided by various speakers into speech output, utilizing the "
67 "distinct voice of each respective speaker.\n"
68 )
70 @classmethod
71 def from_pretrained(
72 cls,
73 pretrained_model_name_or_path: Union[str, os.PathLike],
74 **kwargs: Any
75 ) -> "VibeVoiceProcessor":
76 """
77 Instantiate a VibeVoiceProcessor from a pretrained VibeVoice processor.
78 Args:
79 pretrained_model_name_or_path (`str` or `os.PathLike`):
80 This can be either:
81 - a string, the *model id* of a pretrained model
82 - a path to a *directory* containing processor config
83 Returns:
84 [`VibeVoiceProcessor`]: The processor object instantiated from pretrained model.
85 """
86 # Load processor configuration
87 config_path = os.path.join(pretrained_model_name_or_path, "preprocessor_config.json")
88 if os.path.exists(config_path):
89 with open(config_path, 'r') as f:
90 config = json.load(f)
91 else:
92 logger.warning(f"No preprocessor_config.json found at {pretrained_model_name_or_path}, using defaults")
93 config = {
94 "speech_tok_compress_ratio": 3200,
95 "db_normalize": True,
96 }
98 # Extract main processor parameters
99 speech_tok_compress_ratio = config.get("speech_tok_compress_ratio", 3200)
100 db_normalize = config.get("db_normalize", True)
102 # Load tokenizer - try from model path first, then fallback to Qwen
103 language_model_pretrained_name = (
104 config.get("language_model_pretrained_name", None)
105 or kwargs.pop("language_model_pretrained_name", "Qwen/Qwen2.5-1.5B")
106 )
107 logger.info(f"Loading tokenizer from {language_model_pretrained_name}")
108 if 'qwen' in language_model_pretrained_name.lower():
109 tokenizer = VibeVoiceTextTokenizerFast.from_pretrained(
110 language_model_pretrained_name,
111 **kwargs
112 )
113 else:
114 raise ValueError(
115 f"Unsupported tokenizer type for {language_model_pretrained_name}. "
116 "Supported types: Qwen, Llama, Gemma.")
118 # Load audio processor
119 if "audio_processor" in config:
120 # Create audio processor from config
121 audio_config = config["audio_processor"]
122 audio_processor = VibeVoiceTokenizerProcessor(
123 sampling_rate=audio_config.get("sampling_rate", 24000),
124 normalize_audio=audio_config.get("normalize_audio", True),
125 target_dB_FS=audio_config.get("target_dB_FS", -25),
126 eps=audio_config.get("eps", 1e-6),
127 )
128 else:
129 # Create default audio processor
130 audio_processor = VibeVoiceTokenizerProcessor()
132 # Create and return the processor
133 return cls(
134 tokenizer=tokenizer,
135 audio_processor=audio_processor,
136 speech_tok_compress_ratio=speech_tok_compress_ratio,
137 db_normalize=db_normalize,
138 )
140 def save_pretrained(
141 self,
142 save_directory: Union[str, os.PathLike],
143 **kwargs: Dict[str, Any]
144 ) -> None:
145 """
146 Save a processor to a directory, so that it can be re-loaded using the
147 [`~VibeVoiceProcessor.from_pretrained`] class method.
148 Args:
149 save_directory (`str` or `os.PathLike`):
150 Directory where the processor will be saved.
151 """
152 import os
153 import json
155 os.makedirs(save_directory, exist_ok=True)
157 # Save processor configuration
158 processor_config = {
159 "processor_class": "VibeVoiceProcessor",
160 "speech_tok_compress_ratio": self.speech_tok_compress_ratio,
161 "db_normalize": self.db_normalize,
162 "audio_processor": {
163 "feature_extractor_type": "VibeVoiceTokenizerProcessor",
164 "sampling_rate": getattr(self.audio_processor, 'sampling_rate', 24000),
165 "normalize_audio": getattr(self.audio_processor, 'normalize_audio', True),
166 "target_dB_FS": getattr(self.audio_processor, 'target_dB_FS', -25),
167 "eps": getattr(self.audio_processor, 'eps', 1e-6),
168 }
169 }
171 config_path = os.path.join(save_directory, "preprocessor_config.json")
172 with open(config_path, 'w') as f:
173 json.dump(processor_config, f, indent=2)
175 logger.info(f"Processor configuration saved in {config_path}")
177 def __call__(
178 self,
179 text: Optional[Union[str, List[Any], TextInput, PreTokenizedInput]] = None,
180 voice_samples: Optional[Union[List[Union[str, np.ndarray]], List[List[Union[str, np.ndarray]]]]] = None,
181 padding: Union[bool, str, PaddingStrategy] = True,
182 truncation: Union[bool, str, TruncationStrategy] = False,
183 max_length: Optional[int] = None,
184 return_tensors: Optional[Union[str, TensorType]] = None,
185 return_attention_mask: bool = True,
186 **kwargs: Dict[str, Any]
187 ) -> BatchEncoding:
188 """
189 Main method to process one or more podcast scripts with optional voice samples.
190 Args:
191 text (`str`, `List[str]`):
192 The input text(s) to process. Can be:
193 - A single script string
194 - A list of script strings for batch processing
195 - A path to a .json or .txt file
196 - A list of paths
197 voice_samples (`List[Union[str, np.ndarray]]`, `List[List[Union[str, np.ndarray]]]`, *optional*):
198 Voice samples for each script. Can be:
199 - A list of samples for a single script
200 - A list of lists for batch processing
201 padding (`bool`, `str` or `PaddingStrategy`, defaults to `True`):
202 Whether to pad sequences to the same length
203 truncation (`bool`, `str` or `TruncationStrategy`, defaults to `False`):
204 Whether to truncate sequences
205 max_length (`int`, *optional*):
206 Maximum length of the returned sequences
207 return_tensors (`str` or `TensorType`, *optional*):
208 If set, will return tensors of a particular framework
209 return_attention_mask (`bool`, defaults to `True`):
210 Whether to return the attention mask
211 Returns:
212 `BatchEncoding`: A BatchEncoding with the following fields:
213 - **input_ids** -- List of token id sequences or tensor
214 - **attention_mask** -- List of attention masks or tensor
215 - **speech_tensors** -- Padded speech inputs (if voice_samples provided)
216 - **speech_masks** -- Speech masks (if voice_samples provided)
217 - **speech_input_mask** -- Boolean masks indicating speech token positions
218 """
219 # Handle single vs batch input
220 texts: List[Any] = []
221 is_batched = False
222 if isinstance(text, str) or (isinstance(text, list) and len(text) > 0 and not isinstance(text[0], str)):
223 # Single input
224 texts = [text]
225 is_batched = False
226 else:
227 # Batch input
228 texts = cast(List[Any], text)
229 is_batched = True
231 # Handle voice samples
232 voice_samples_list: List[Optional[List[Union[str, np.ndarray]]]] = []
233 if voice_samples is not None:
234 if not is_batched or (isinstance(voice_samples[0], (str, np.ndarray))):
235 # Single set of voice samples
236 voice_samples_list = [cast(List[Union[str, np.ndarray]], voice_samples)]
237 else:
238 # Batch of voice samples
239 voice_samples_list = cast(List[Optional[List[Union[str, np.ndarray]]]], voice_samples)
240 else:
241 voice_samples_list = [None] * len(texts)
243 # Process each input
244 all_encodings = []
245 for text_input, voice_input in zip(texts, voice_samples_list):
246 encoding = self._process_single(text_input, voice_input)
247 all_encodings.append(encoding)
249 # Combine batch
250 batch_encoding = self._batch_encode(
251 all_encodings,
252 padding=padding,
253 truncation=truncation,
254 max_length=max_length,
255 return_tensors=return_tensors,
256 return_attention_mask=return_attention_mask,
257 )
259 return batch_encoding
261 def _process_single(
262 self,
263 text: Union[str, TextInput],
264 voice_samples: Optional[List[Union[str, np.ndarray]]] = None,
265 ) -> Dict[str, Any]:
266 """Process a single podcast script."""
267 # Determine if text is a file path or direct script
268 script = None
269 if isinstance(text, str):
270 # Check if it's a file path
271 if text.endswith('.json') and os.path.exists(text):
272 script = self._convert_json_to_script(text)
273 elif text.endswith('.txt') and os.path.exists(text):
274 script = self._convert_text_to_script(text)
275 else:
276 # Assume it's the script content directly
277 script = text
279 if script is None:
280 raise ValueError(f"Could not process input text: {text}")
282 # Parse the script
283 parsed_lines = self._parse_script(script)
284 all_speakers = list(set(speaker_id for speaker_id, _ in parsed_lines))
286 # Create system prompt
287 # system_tokens = self.tokenizer.encode(self.system_prompt, add_special_tokens=False)
288 if not self.tokenizer:
289 raise ValueError("Tokenizer is not initialized.")
290 system_tokens = self.tokenizer.encode(self.system_prompt)
292 # Process voice samples if provided
293 if voice_samples:
294 voice_tokens, voice_speech_inputs, voice_speech_masks = self._create_voice_prompt(
295 voice_samples[:len(all_speakers)])
296 else:
297 voice_tokens, voice_speech_inputs, voice_speech_masks = [], [], []
299 # Build full token sequence
300 full_tokens = system_tokens + voice_tokens
301 speech_input_mask = [False] * len(system_tokens) + voice_speech_masks
303 # Add text input section
304 full_tokens += self.tokenizer.encode(' Text input:\n', add_special_tokens=False)
305 speech_input_mask += [False] * len(self.tokenizer.encode(' Text input:\n', add_special_tokens=False))
307 for speaker_id, speaker_text in parsed_lines:
308 speaker_text_tokens = self.tokenizer.encode(
309 f" Speaker {speaker_id}:{speaker_text}\n",
310 add_special_tokens=False)
311 full_tokens += speaker_text_tokens
312 speech_input_mask += [False] * len(speaker_text_tokens)
314 # Add speech output section
315 full_tokens += (
316 self.tokenizer.encode(' Speech output:\n', add_special_tokens=False)
317 + [self.tokenizer.speech_start_id]
318 )
319 speech_input_mask += [False] * (len(self.tokenizer.encode(' Speech output:\n', add_special_tokens=False)) + 1)
321 return {
322 "input_ids": full_tokens,
323 "speech_inputs": voice_speech_inputs if voice_speech_inputs else None,
324 "speech_input_mask": speech_input_mask,
325 "parsed_script": parsed_lines,
326 "all_speakers": all_speakers,
327 }
329 def _batch_encode(
330 self,
331 encodings: List[Dict[str, Any]],
332 padding: Union[bool, str, PaddingStrategy] = True,
333 truncation: Union[bool, str, TruncationStrategy] = False,
334 max_length: Optional[int] = None,
335 return_tensors: Optional[Union[str, TensorType]] = None,
336 return_attention_mask: bool = True,
337 ) -> BatchEncoding:
338 """Combine multiple encodings into a batch with padding."""
339 # Extract input_ids and create attention_mask
340 input_ids_list = [enc["input_ids"] for enc in encodings]
341 speech_input_masks_list = [enc["speech_input_mask"] for enc in encodings]
343 # Determine padding strategy
344 if isinstance(padding, bool):
345 padding_strategy = PaddingStrategy.LONGEST if padding else PaddingStrategy.DO_NOT_PAD
346 elif isinstance(padding, str):
347 padding_strategy = PaddingStrategy(padding)
348 else:
349 padding_strategy = padding
351 if not self.tokenizer:
352 raise ValueError("Tokenizer is not initialized.")
354 # Apply padding to input_ids
355 if padding_strategy != PaddingStrategy.DO_NOT_PAD:
356 if padding_strategy == PaddingStrategy.LONGEST:
357 max_len = max(len(ids) for ids in input_ids_list)
358 elif padding_strategy == PaddingStrategy.MAX_LENGTH and max_length is not None:
359 max_len = max_length
360 else:
361 max_len = max(len(ids) for ids in input_ids_list)
363 # Pad sequences
364 padded_input_ids = []
365 padded_attention_masks: List[List[int]] = []
366 padded_speech_input_masks = []
368 for input_ids, speech_mask in zip(input_ids_list, speech_input_masks_list):
369 # Truncate if needed
370 if truncation and len(input_ids) > max_len:
371 input_ids = input_ids[:max_len]
372 speech_mask = speech_mask[:max_len]
374 # Pad
375 padding_length = max_len - len(input_ids)
376 # padded_ids = [self.tokenizer.pad_token_id] * padding_length + input_ids
377 padded_ids = [self.tokenizer.pad_id] * padding_length + input_ids
378 attention_mask = [0] * padding_length + [1] * len(input_ids)
379 padded_speech_mask = [False] * padding_length + speech_mask
381 padded_input_ids.append(padded_ids)
382 padded_attention_masks.append(attention_mask)
383 padded_speech_input_masks.append(padded_speech_mask)
385 attention_masks: Optional[List[List[int]]] = padded_attention_masks
386 input_ids_list = padded_input_ids
387 speech_input_masks_list = padded_speech_input_masks
388 else:
389 # No padding, just create attention masks
390 attention_masks = [[1] * len(ids) for ids in input_ids_list] if return_attention_mask else None
392 # Process speech inputs
393 all_speech_inputs = []
394 has_speech = False
395 for enc in encodings:
396 if enc["speech_inputs"] is not None:
397 all_speech_inputs.extend(enc["speech_inputs"])
398 has_speech = True
400 # Prepare batch encoding
401 batch_encoding = BatchEncoding()
403 # Handle tensor conversion
404 if return_tensors is not None:
405 batch_encoding["input_ids"] = torch.tensor(input_ids_list, dtype=torch.long)
406 if return_attention_mask and attention_masks is not None:
407 batch_encoding["attention_mask"] = torch.tensor(attention_masks, dtype=torch.long)
408 batch_encoding["speech_input_mask"] = torch.tensor(speech_input_masks_list, dtype=torch.bool)
409 else:
410 batch_encoding["input_ids"] = input_ids_list
411 if return_attention_mask and attention_masks is not None:
412 batch_encoding["attention_mask"] = attention_masks
413 batch_encoding["speech_input_mask"] = speech_input_masks_list
415 # Process speech tensors if present
416 if has_speech:
417 speech_dict = self.prepare_speech_inputs(
418 all_speech_inputs,
419 return_tensors=return_tensors,
420 )
421 batch_encoding["speech_tensors"] = speech_dict["padded_speeches"]
422 batch_encoding["speech_masks"] = speech_dict["speech_masks"]
423 else:
424 batch_encoding["speech_tensors"] = None
425 batch_encoding["speech_masks"] = None
427 # Add metadata
428 batch_encoding["parsed_scripts"] = [enc["parsed_script"] for enc in encodings]
429 batch_encoding["all_speakers_list"] = [enc["all_speakers"] for enc in encodings]
431 return batch_encoding
433 def _create_voice_prompt(
434 self,
435 speaker_samples: List[Union[str, np.ndarray]]
436 ) -> Tuple[List[int], List[np.ndarray], List[bool]]:
437 """
438 Create voice prompt tokens and process audio samples.
440 Returns:
441 tuple: (voice_tokens, voice_speech_inputs, voice_speech_masks)
442 """
443 if not self.tokenizer:
444 raise ValueError("Tokenizer is not initialized.")
446 vae_token_id = self.tokenizer.speech_diffusion_id
448 voice_full_tokens = self.tokenizer.encode(' Voice input:\n', add_special_tokens=False)
449 voice_speech_inputs = []
450 voice_speech_masks = [False] * len(voice_full_tokens)
452 for speaker_id, speaker_audio in enumerate(speaker_samples):
453 prefix_tokens = self.tokenizer.encode(f" Speaker {speaker_id}:", add_special_tokens=False)
455 # Process audio
456 if isinstance(speaker_audio, str):
457 # Load audio from file
458 assert self.audio_processor is not None
459 wav = self.audio_processor._load_audio_from_path(speaker_audio)
460 else:
461 wav = np.array(speaker_audio, dtype=np.float32)
463 # Apply normalization if needed
464 if self.db_normalize and self.audio_normalizer:
465 wav = self.audio_normalizer(wav)
467 # Calculate token length based on compression ratio
468 # if speaker_audio.endswith('.pt') or speaker_audio.endswith('.npy'):
469 # vae_tok_len = wav.shape[0]
470 # else:
471 vae_tok_len = math.ceil(wav.shape[0] / self.speech_tok_compress_ratio)
473 # Build tokens and masks
474 speaker_tokens = (
475 prefix_tokens
476 + [self.tokenizer.speech_start_id]
477 + [vae_token_id] * vae_tok_len
478 + [self.tokenizer.speech_end_id]
479 + self.tokenizer.encode('\n', add_special_tokens=False)
480 )
482 vae_input_mask = (
483 [False] * len(prefix_tokens)
484 + [False]
485 + [True] * vae_tok_len
486 + [False]
487 + [False]
488 )
490 voice_full_tokens.extend(speaker_tokens)
491 voice_speech_masks.extend(vae_input_mask)
492 voice_speech_inputs.append(wav)
494 return voice_full_tokens, voice_speech_inputs, voice_speech_masks
496 def prepare_speech_inputs(
497 self,
498 speech_inputs: List[np.ndarray],
499 return_tensors: Optional[Union[str, TensorType]] = None,
500 device: Optional[Union[str, torch.device]] = None,
501 dtype: Optional[torch.dtype] = None,
502 ) -> Dict[str, Any]:
503 """
504 Prepare speech inputs for model consumption.
506 Args:
507 speech_inputs: List of speech arrays
508 return_tensors: Output tensor type
509 device: Device to place tensors on
510 dtype: Data type for tensors
512 Returns:
513 Dictionary with padded_speeches and speech_masks
514 """
515 if not speech_inputs:
516 return {"padded_speeches": None, "speech_masks": None}
518 # Calculate sequence lengths
519 vae_tok_seqlens = [math.ceil(s.shape[0] / self.speech_tok_compress_ratio) for s in speech_inputs]
520 # vae_tok_seqlens = [math.ceil(s.shape[0] / self.speech_tok_compress_ratio)
521 # if s.ndim == 1 else s.shape[0] for s in speech_inputs]
522 max_speech_length = max(s.shape[0] for s in speech_inputs)
524 # Pad speeches
525 if speech_inputs[0].ndim == 1:
526 padded_speeches = np.full((len(speech_inputs), max_speech_length), fill_value=0, dtype=np.float32)
527 else:
528 padded_speeches = np.full(
529 (len(speech_inputs), max_speech_length, speech_inputs[0].shape[-1]),
530 fill_value=0,
531 dtype=np.float32)
532 speech_masks = np.zeros((len(speech_inputs), max(vae_tok_seqlens)), dtype=np.bool_)
534 for i, (speech, vae_tok_length) in enumerate(zip(speech_inputs, vae_tok_seqlens)):
535 padded_speeches[i, :len(speech)] = speech
536 speech_masks[i, :vae_tok_length] = True
538 result: Dict[str, Any] = {
539 "padded_speeches": padded_speeches,
540 "speech_masks": speech_masks,
541 }
543 # Convert to tensors if requested
544 if return_tensors == "pt":
545 result["padded_speeches"] = torch.tensor(padded_speeches, device=device, dtype=dtype or torch.float32)
546 result["speech_masks"] = torch.tensor(speech_masks, device=device, dtype=torch.bool)
548 return result
550 def _convert_json_to_script(self, json_file: str) -> str:
551 """
552 Convert JSON format to script format.
553 Expected JSON format:
554 [
555 {"speaker": "1", "text": "Hello everyone..."},
556 {"speaker": "2", "text": "Great to be here..."}
557 ]
558 """
559 import json
561 with open(json_file, 'r', encoding='utf-8') as f:
562 data = json.load(f)
564 if not isinstance(data, list):
565 raise ValueError("JSON file must contain a list of speaker entries")
567 script_lines = []
568 for item in data:
569 if not isinstance(item, dict):
570 logger.warning(f"Skipping non-dict entry: {item}")
571 continue
573 speaker = item.get('speaker')
574 text = item.get('text')
576 if speaker is None or text is None:
577 logger.warning(f"Skipping entry missing speaker or text: {item}")
578 continue
580 # Ensure speaker ID is valid
581 try:
582 speaker_id = int(speaker)
583 except (ValueError, TypeError):
584 logger.warning(f"Invalid speaker ID: {speaker}, skipping entry")
585 continue
587 # Clean up text
588 text = text.strip()
589 if text:
590 script_lines.append(f"Speaker {speaker_id}: {text}")
592 if not script_lines:
593 raise ValueError("No valid entries found in JSON file")
595 return "\n".join(script_lines)
597 def _convert_text_to_script(self, text_file: str) -> str:
598 """
599 Convert text file to script format.
600 Handles multiple formats:
601 1. Already formatted as "Speaker X: text"
602 2. Plain text (assigns to Speaker 1)
604 Handles edge cases like multiple colons in a line.
605 """
606 with open(text_file, 'r', encoding='utf-8') as f:
607 lines = f.readlines()
609 script_lines = []
610 current_speaker = 1
612 for line in lines:
613 line = line.strip()
614 if not line:
615 continue
617 # Try to parse as "Speaker X: text" format
618 # Use regex to be more robust
619 speaker_match = re.match(r'^Speaker\s+(\d+)\s*:\s*(.*)$', line, re.IGNORECASE)
621 if speaker_match:
622 speaker_id = int(speaker_match.group(1))
623 text = speaker_match.group(2).strip()
624 if text:
625 script_lines.append(f"Speaker {speaker_id}: {text}")
626 else:
627 # Treat as plain text - assign to current speaker
628 script_lines.append(f"Speaker {current_speaker}: {line}")
630 if not script_lines:
631 raise ValueError("No valid content found in text file")
633 return "\n".join(script_lines)
635 def _parse_script(self, script: str) -> List[Tuple[int, str]]:
636 """Parse script into list of (speaker_id, text) tuples."""
637 lines = script.strip().split("\n")
638 parsed_lines = []
639 speaker_ids = []
641 # First pass: parse all lines and collect speaker IDs
642 for line in lines:
643 if not line.strip():
644 continue
646 # Use regex to handle edge cases like multiple colons
647 match = re.match(r'^Speaker\s+(\d+)\s*:\s*(.*)$', line.strip(), re.IGNORECASE)
649 if match:
650 speaker_id = int(match.group(1))
651 text = ' ' + match.group(2).strip()
652 parsed_lines.append((speaker_id, text))
653 speaker_ids.append(speaker_id)
654 else:
655 logger.warning(f"Could not parse line: '{line}'")
657 if not parsed_lines:
658 raise ValueError("No valid speaker lines found in script")
660 # Check if we need to normalize speaker IDs (only if all are > 0)
661 min_speaker_id = min(speaker_ids)
662 if min_speaker_id > 0:
663 # Normalize to start from 0
664 normalized_lines = []
665 for speaker_id, text in parsed_lines:
666 normalized_lines.append((speaker_id - 1, text))
667 return normalized_lines
668 else:
669 # Keep original IDs
670 return parsed_lines
672 def _merge_inputs(self, text_inputs: BatchEncoding, audio_inputs: Dict) -> BatchEncoding:
673 """Merge text and audio inputs into a single BatchEncoding."""
674 # Start with text inputs
675 merged = BatchEncoding(dict(text_inputs))
677 # Add audio-specific fields
678 if "audio" in audio_inputs:
679 merged["speech_inputs"] = audio_inputs["audio"]
680 if "streaming" in audio_inputs:
681 merged["streaming"] = audio_inputs["streaming"]
683 return merged
685 def batch_decode(self, *args: Any, **kwargs: Any) -> Any:
686 """
687 This method forwards all its arguments to VibeVoiceTextTokenizer's [`~PreTrainedTokenizer.batch_decode`].
688 Please refer to the docstring of this method for more information.
689 """
690 if not self.tokenizer:
691 raise ValueError("Tokenizer is not initialized.")
692 return self.tokenizer.batch_decode(*args, **kwargs)
694 def decode(self, *args: Any, **kwargs: Any) -> Any:
695 """
696 This method forwards all its arguments to VibeVoiceTextTokenizer's [`~PreTrainedTokenizer.decode`].
697 Please refer to the docstring of this method for more information.
698 """
699 if not self.tokenizer:
700 raise ValueError("Tokenizer is not initialized.")
701 return self.tokenizer.decode(*args, **kwargs)
703 @property
704 def model_input_names(self) -> List[str]:
705 """
706 Return the list of inputs accepted by the model.
707 """
708 if not self.tokenizer or not self.audio_processor:
709 raise ValueError("Tokenizer or audio processor is not initialized.")
710 tokenizer_input_names = self.tokenizer.model_input_names
711 audio_processor_input_names = self.audio_processor.model_input_names
712 return list(dict.fromkeys(
713 tokenizer_input_names + audio_processor_input_names + ["speech_inputs", "speech_input_mask"]))
715 def save_audio(
716 self,
717 audio: Union[torch.Tensor, np.ndarray, List[Union[torch.Tensor, np.ndarray]]],
718 output_path: str = "output.wav",
719 sampling_rate: Optional[int] = None,
720 normalize: bool = False,
721 batch_prefix: str = "audio_",
722 ) -> str:
723 """
724 Save audio data to a file.
725 Args:
726 audio (Union[torch.Tensor, np.ndarray, List[Union[torch.Tensor, np.ndarray]]]):
727 The audio data to save. Can be a single tensor/array or a list of them.
728 output_path (str, optional): Path to save the audio file. Defaults to "output.wav".
729 sampling_rate (int, optional): Sampling rate for the audio. If None, uses the processor's default.
730 normalize (bool, optional): Whether to normalize the audio before saving. Defaults to False.
731 batch_prefix (str, optional): Prefix for batch audio files. Defaults to "audio_".
732 Returns:
733 str: The path to the saved audio file.
734 """
735 if not self.audio_processor:
736 raise ValueError("Audio processor is not initialized.")
737 return self.audio_processor.save_audio(
738 audio, output_path=output_path, sampling_rate=sampling_rate,
739 normalize=normalize, batch_prefix=batch_prefix)
742__all__ = [
743 "VibeVoiceProcessor",
744]