Add VibeVoice streaming model (#35)
This commit is contained in:
@@ -175,3 +175,4 @@ tags
|
|||||||
/checkpoints/
|
/checkpoints/
|
||||||
exp
|
exp
|
||||||
.gradio/
|
.gradio/
|
||||||
|
MSFT-VV
|
||||||
|
|||||||
@@ -34,6 +34,7 @@ Fine-tuning is now supported, which is incredibly powerful. You can adapt VibeVo
|
|||||||
|
|
||||||
## Updates
|
## Updates
|
||||||
|
|
||||||
|
- **[2025-12-04]** Added support for VibeVoice-Streaming-0.5B model for real-time TTS!
|
||||||
- **[2025-09-05]** Microsoft repo restored (without code) with statement about responsible AI use.
|
- **[2025-09-05]** Microsoft repo restored (without code) with statement about responsible AI use.
|
||||||
- **[2025-09-04]** Community backup created after Microsoft removed original repo and models.
|
- **[2025-09-04]** Community backup created after Microsoft removed original repo and models.
|
||||||
- **[2025-08-26]** The [VibeVoice-7B](https://huggingface.co/vibevoice/VibeVoice-7B) model weights are open-sourced!
|
- **[2025-08-26]** The [VibeVoice-7B](https://huggingface.co/vibevoice/VibeVoice-7B) model weights are open-sourced!
|
||||||
@@ -47,10 +48,16 @@ Fine-tuning is now supported, which is incredibly powerful. You can adapt VibeVo
|
|||||||
|
|
||||||
## Model Zoo
|
## Model Zoo
|
||||||
|
|
||||||
| Model | Context Length | Generation Length | Weight |
|
| Model | Context Length | Generation Length | Speakers | Weight |
|
||||||
|-------|----------------|----------|----------|
|
|-------|----------------|-------------------|----------|--------|
|
||||||
| VibeVoice-1.5B | 64K | ~90 min | [HF link](https://huggingface.co/vibevoice/VibeVoice-1.5B) |
|
| VibeVoice-Streaming-0.5B | 8K | Real-time | 1 | [HF link](https://huggingface.co/microsoft/VibeVoice-Realtime-0.5B) |
|
||||||
| VibeVoice-Large| 32K | ~45 min | [HF link](https://huggingface.co/vibevoice/VibeVoice-7B) |
|
| VibeVoice-1.5B | 64K | ~90 min | Up to 4 | [HF link](https://huggingface.co/vibevoice/VibeVoice-1.5B) |
|
||||||
|
| VibeVoice-Large (7B) | 32K | ~45 min | Up to 4 | [HF link](https://huggingface.co/vibevoice/VibeVoice-7B) |
|
||||||
|
|
||||||
|
### Model Comparison
|
||||||
|
|
||||||
|
- **VibeVoice-Streaming-0.5B**: Optimized for **real-time** low-latency TTS. Single speaker only. Uses pre-computed voice embeddings (.pt files) for fast inference. Best for live applications.
|
||||||
|
- **VibeVoice-1.5B/7B**: Full-featured models for **long-form multi-speaker** content like podcasts. Support up to 4 speakers with voice cloning from audio samples.
|
||||||
|
|
||||||
## Installation
|
## Installation
|
||||||
|
|
||||||
@@ -82,7 +89,7 @@ python demo/gradio_demo.py --model_path vibevoice/VibeVoice-1.5B --share
|
|||||||
# use the in-app "Disable voice cloning" setting (Advanced Settings) to skip speaker conditioning
|
# use the in-app "Disable voice cloning" setting (Advanced Settings) to skip speaker conditioning
|
||||||
```
|
```
|
||||||
|
|
||||||
**Option 2: Inference from files directly**
|
**Option 2: Inference from files directly (Multi-Speaker 1.5B/7B)**
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
# We provide some LLM generated example scripts under demo/text_examples/ for demo
|
# We provide some LLM generated example scripts under demo/text_examples/ for demo
|
||||||
@@ -99,6 +106,30 @@ python demo/inference_from_file.py --model_path vibevoice/VibeVoice-7B --txt_pat
|
|||||||
python demo/inference_from_file.py --model_path vibevoice/VibeVoice-7B --txt_path demo/text_examples/1p_abs.txt --speaker_names Alice --disable_prefill
|
python demo/inference_from_file.py --model_path vibevoice/VibeVoice-7B --txt_path demo/text_examples/1p_abs.txt --speaker_names Alice --disable_prefill
|
||||||
```
|
```
|
||||||
|
|
||||||
|
**Option 3: Streaming Model (0.5B) - Real-time TTS**
|
||||||
|
|
||||||
|
The streaming model uses pre-computed voice embeddings for low-latency generation:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# Basic usage with streaming model
|
||||||
|
python demo/streaming_inference_from_file.py \
|
||||||
|
--model_path microsoft/VibeVoice-Realtime-0.5B \
|
||||||
|
--txt_path demo/text_examples/1p_vibevoice.txt \
|
||||||
|
--speaker_name Emma
|
||||||
|
|
||||||
|
# Available voice presets: Carter, Davis, Emma, Frank, Grace, Mike (English), Samuel (Indian English)
|
||||||
|
# Adjust CFG scale and DDPM steps for quality/speed tradeoff
|
||||||
|
python demo/streaming_inference_from_file.py \
|
||||||
|
--model_path microsoft/VibeVoice-Realtime-0.5B \
|
||||||
|
--txt_path demo/text_examples/1p_vibevoice.txt \
|
||||||
|
--speaker_name Mike \
|
||||||
|
--cfg_scale 1.5 \
|
||||||
|
--ddpm_steps 5
|
||||||
|
```
|
||||||
|
|
||||||
|
Voice presets are stored as `.pt` files in `demo/voices/streaming_model/`. These contain pre-computed KV cache embeddings for fast inference. Voice cloning is not supported for now.
|
||||||
|
|
||||||
|
NOTE: If you get the warning `Some weights of VibeVoiceStreamingForConditionalGenerationInference were not initialized from the model checkpoint` when loading, this is expected. This is because voice cloning capabilities have been removed from the model.
|
||||||
|
|
||||||
## [Finetuning](./FINETUNING.md)
|
## [Finetuning](./FINETUNING.md)
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,377 @@
|
|||||||
|
"""
|
||||||
|
VibeVoice Streaming Model Inference from File
|
||||||
|
|
||||||
|
This script demonstrates how to use the VibeVoice-Streaming-0.5B model
|
||||||
|
for text-to-speech synthesis from a text file.
|
||||||
|
|
||||||
|
The streaming model uses pre-computed voice embeddings (.pt files) for
|
||||||
|
low-latency generation, making it suitable for real-time applications.
|
||||||
|
|
||||||
|
Usage:
|
||||||
|
python demo/streaming_inference_from_file.py \
|
||||||
|
--model_path microsoft/VibeVoice-Realtime-0.5B \
|
||||||
|
--txt_path demo/text_examples/1p_vibevoice.txt \
|
||||||
|
--speaker_name Emma \
|
||||||
|
--output_dir ./outputs
|
||||||
|
|
||||||
|
Available voice presets (in demo/voices/streaming_model/):
|
||||||
|
- Carter, Davis, Emma, Frank, Grace, Mike (English)
|
||||||
|
- Samuel (Indian English)
|
||||||
|
"""
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import os
|
||||||
|
import re
|
||||||
|
import traceback
|
||||||
|
from typing import List, Tuple, Union, Dict, Any
|
||||||
|
import time
|
||||||
|
import torch
|
||||||
|
import copy
|
||||||
|
|
||||||
|
from vibevoice.modular.modeling_vibevoice_streaming_inference import VibeVoiceStreamingForConditionalGenerationInference
|
||||||
|
from vibevoice.processor.vibevoice_streaming_processor import VibeVoiceStreamingProcessor
|
||||||
|
from transformers.utils import logging
|
||||||
|
|
||||||
|
logging.set_verbosity_info()
|
||||||
|
logger = logging.get_logger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class VoiceMapper:
|
||||||
|
"""Maps speaker names to voice file paths for the streaming model."""
|
||||||
|
|
||||||
|
def __init__(self):
|
||||||
|
self.setup_voice_presets()
|
||||||
|
|
||||||
|
# Create name aliases without prefixes/suffixes for easier matching
|
||||||
|
new_dict = {}
|
||||||
|
for name, path in self.voice_presets.items():
|
||||||
|
# Remove gender suffix (e.g., "en-Emma_woman" -> "Emma")
|
||||||
|
if '_' in name:
|
||||||
|
name = name.split('_')[0]
|
||||||
|
|
||||||
|
# Remove language prefix (e.g., "en-Emma" -> "Emma")
|
||||||
|
if '-' in name:
|
||||||
|
name = name.split('-')[-1]
|
||||||
|
|
||||||
|
new_dict[name] = path
|
||||||
|
self.voice_presets.update(new_dict)
|
||||||
|
|
||||||
|
def setup_voice_presets(self):
|
||||||
|
"""Setup voice presets by scanning the voices/streaming_model directory."""
|
||||||
|
voices_dir = os.path.join(os.path.dirname(__file__), "voices/streaming_model")
|
||||||
|
|
||||||
|
# Check if voices directory exists
|
||||||
|
if not os.path.exists(voices_dir):
|
||||||
|
print(f"Warning: Voices directory not found at {voices_dir}")
|
||||||
|
self.voice_presets = {}
|
||||||
|
self.available_voices = {}
|
||||||
|
return
|
||||||
|
|
||||||
|
# Scan for all .pt files in the voices directory
|
||||||
|
self.voice_presets = {}
|
||||||
|
|
||||||
|
pt_files = [f for f in os.listdir(voices_dir)
|
||||||
|
if f.lower().endswith('.pt') and os.path.isfile(os.path.join(voices_dir, f))]
|
||||||
|
|
||||||
|
# Create dictionary with filename (without extension) as key
|
||||||
|
for pt_file in pt_files:
|
||||||
|
name = os.path.splitext(pt_file)[0]
|
||||||
|
full_path = os.path.join(voices_dir, pt_file)
|
||||||
|
self.voice_presets[name] = full_path
|
||||||
|
|
||||||
|
# Sort alphabetically for consistent ordering
|
||||||
|
self.voice_presets = dict(sorted(self.voice_presets.items()))
|
||||||
|
|
||||||
|
# Filter out voices that don't exist (safety check)
|
||||||
|
self.available_voices = {
|
||||||
|
name: path for name, path in self.voice_presets.items()
|
||||||
|
if os.path.exists(path)
|
||||||
|
}
|
||||||
|
|
||||||
|
print(f"Found {len(self.available_voices)} voice files in {voices_dir}")
|
||||||
|
print(f"Available voices: {', '.join(self.available_voices.keys())}")
|
||||||
|
|
||||||
|
def get_voice_path(self, speaker_name: str) -> str:
|
||||||
|
"""Get voice file path for a given speaker name."""
|
||||||
|
# First try exact match
|
||||||
|
if speaker_name in self.voice_presets:
|
||||||
|
return self.voice_presets[speaker_name]
|
||||||
|
|
||||||
|
# Try partial matching (case insensitive)
|
||||||
|
speaker_lower = speaker_name.lower()
|
||||||
|
for preset_name, path in self.voice_presets.items():
|
||||||
|
if preset_name.lower() in speaker_lower or speaker_lower in preset_name.lower():
|
||||||
|
return path
|
||||||
|
|
||||||
|
# Default to first voice if no match found
|
||||||
|
if self.voice_presets:
|
||||||
|
default_voice = list(self.voice_presets.values())[0]
|
||||||
|
print(f"Warning: No voice preset found for '{speaker_name}', using default voice")
|
||||||
|
return default_voice
|
||||||
|
|
||||||
|
raise ValueError(f"No voice presets available. Please check the voices/streaming_model directory.")
|
||||||
|
|
||||||
|
|
||||||
|
def parse_args():
|
||||||
|
parser = argparse.ArgumentParser(
|
||||||
|
description="VibeVoice Streaming Model (0.5B) - Text-to-Speech from File",
|
||||||
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||||
|
epilog="""
|
||||||
|
Examples:
|
||||||
|
# Basic usage with default settings
|
||||||
|
python demo/streaming_inference_from_file.py --txt_path demo/text_examples/1p_vibevoice.txt
|
||||||
|
|
||||||
|
# Specify a different voice
|
||||||
|
python demo/streaming_inference_from_file.py --txt_path demo/text_examples/1p_vibevoice.txt --speaker_name Emma
|
||||||
|
|
||||||
|
# Use local model path
|
||||||
|
python demo/streaming_inference_from_file.py --model_path /path/to/model --txt_path input.txt
|
||||||
|
"""
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--model_path",
|
||||||
|
type=str,
|
||||||
|
default="microsoft/VibeVoice-Realtime-0.5B",
|
||||||
|
help="Path to the HuggingFace model directory or model ID",
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--txt_path",
|
||||||
|
type=str,
|
||||||
|
default="demo/text_examples/1p_vibevoice.txt",
|
||||||
|
help="Path to the txt file containing the script",
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--speaker_name",
|
||||||
|
type=str,
|
||||||
|
default="Emma",
|
||||||
|
help="Speaker name for voice selection (e.g., Emma, Carter, Mike)",
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--output_dir",
|
||||||
|
type=str,
|
||||||
|
default="./outputs",
|
||||||
|
help="Directory to save output audio files",
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--device",
|
||||||
|
type=str,
|
||||||
|
default=("cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")),
|
||||||
|
help="Device for inference: cuda | mps | cpu",
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--cfg_scale",
|
||||||
|
type=float,
|
||||||
|
default=1.5,
|
||||||
|
help="CFG (Classifier-Free Guidance) scale for generation (default: 1.5)",
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--ddpm_steps",
|
||||||
|
type=int,
|
||||||
|
default=5,
|
||||||
|
help="Number of DDPM inference steps (default: 5, lower is faster)",
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--seed",
|
||||||
|
type=int,
|
||||||
|
default=None,
|
||||||
|
help="Random seed for reproducibility",
|
||||||
|
)
|
||||||
|
|
||||||
|
return parser.parse_args()
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
args = parse_args()
|
||||||
|
|
||||||
|
# Set random seed if specified
|
||||||
|
if args.seed is not None:
|
||||||
|
torch.manual_seed(args.seed)
|
||||||
|
if torch.cuda.is_available():
|
||||||
|
torch.cuda.manual_seed_all(args.seed)
|
||||||
|
print(f"Using random seed: {args.seed}")
|
||||||
|
|
||||||
|
# Normalize potential 'mpx' typo to 'mps'
|
||||||
|
if args.device.lower() == "mpx":
|
||||||
|
print("Note: device 'mpx' detected, treating it as 'mps'.")
|
||||||
|
args.device = "mps"
|
||||||
|
|
||||||
|
# Validate mps availability if requested
|
||||||
|
if args.device == "mps" and not torch.backends.mps.is_available():
|
||||||
|
print("Warning: MPS not available. Falling back to CPU.")
|
||||||
|
args.device = "cpu"
|
||||||
|
|
||||||
|
print(f"Using device: {args.device}")
|
||||||
|
|
||||||
|
# Initialize voice mapper
|
||||||
|
voice_mapper = VoiceMapper()
|
||||||
|
|
||||||
|
# Check if txt file exists
|
||||||
|
if not os.path.exists(args.txt_path):
|
||||||
|
print(f"Error: txt file not found: {args.txt_path}")
|
||||||
|
return
|
||||||
|
|
||||||
|
# Read and parse txt file
|
||||||
|
print(f"Reading script from: {args.txt_path}")
|
||||||
|
with open(args.txt_path, 'r', encoding='utf-8') as f:
|
||||||
|
scripts = f.read().strip()
|
||||||
|
|
||||||
|
if not scripts:
|
||||||
|
print("Error: No valid scripts found in the txt file")
|
||||||
|
return
|
||||||
|
|
||||||
|
# Normalize quotes
|
||||||
|
full_script = scripts.replace("'", "'").replace('"', '"').replace('"', '"')
|
||||||
|
|
||||||
|
print(f"Loading processor & model from {args.model_path}")
|
||||||
|
processor = VibeVoiceStreamingProcessor.from_pretrained(args.model_path)
|
||||||
|
|
||||||
|
# Decide dtype & attention implementation based on device
|
||||||
|
if args.device == "mps":
|
||||||
|
load_dtype = torch.float32 # MPS requires float32
|
||||||
|
attn_impl_primary = "sdpa" # flash_attention_2 not supported on MPS
|
||||||
|
elif args.device == "cuda":
|
||||||
|
load_dtype = torch.bfloat16
|
||||||
|
attn_impl_primary = "flash_attention_2"
|
||||||
|
else: # cpu
|
||||||
|
load_dtype = torch.float32
|
||||||
|
attn_impl_primary = "sdpa"
|
||||||
|
|
||||||
|
print(f"Using device: {args.device}, torch_dtype: {load_dtype}, attn_implementation: {attn_impl_primary}")
|
||||||
|
|
||||||
|
# Load model with device-specific logic
|
||||||
|
try:
|
||||||
|
if args.device == "mps":
|
||||||
|
model = VibeVoiceStreamingForConditionalGenerationInference.from_pretrained(
|
||||||
|
args.model_path,
|
||||||
|
torch_dtype=load_dtype,
|
||||||
|
attn_implementation=attn_impl_primary,
|
||||||
|
device_map=None,
|
||||||
|
)
|
||||||
|
model.to("mps")
|
||||||
|
elif args.device == "cuda":
|
||||||
|
model = VibeVoiceStreamingForConditionalGenerationInference.from_pretrained(
|
||||||
|
args.model_path,
|
||||||
|
torch_dtype=load_dtype,
|
||||||
|
device_map="cuda",
|
||||||
|
attn_implementation=attn_impl_primary,
|
||||||
|
)
|
||||||
|
else: # cpu
|
||||||
|
model = VibeVoiceStreamingForConditionalGenerationInference.from_pretrained(
|
||||||
|
args.model_path,
|
||||||
|
torch_dtype=load_dtype,
|
||||||
|
device_map="cpu",
|
||||||
|
attn_implementation=attn_impl_primary,
|
||||||
|
)
|
||||||
|
except Exception as e:
|
||||||
|
if attn_impl_primary == 'flash_attention_2':
|
||||||
|
print(f"[ERROR] : {type(e).__name__}: {e}")
|
||||||
|
print(traceback.format_exc())
|
||||||
|
print("Error loading the model. Trying to use SDPA. Note that flash_attention_2 is recommended for best quality.")
|
||||||
|
model = VibeVoiceStreamingForConditionalGenerationInference.from_pretrained(
|
||||||
|
args.model_path,
|
||||||
|
torch_dtype=load_dtype,
|
||||||
|
device_map=(args.device if args.device in ("cuda", "cpu") else None),
|
||||||
|
attn_implementation='sdpa'
|
||||||
|
)
|
||||||
|
if args.device == "mps":
|
||||||
|
model.to("mps")
|
||||||
|
else:
|
||||||
|
raise e
|
||||||
|
|
||||||
|
model.eval()
|
||||||
|
model.set_ddpm_inference_steps(num_steps=args.ddpm_steps)
|
||||||
|
|
||||||
|
if hasattr(model.model, 'language_model'):
|
||||||
|
print(f"Language model attention: {model.model.language_model.config._attn_implementation}")
|
||||||
|
|
||||||
|
# Load voice preset
|
||||||
|
target_device = args.device if args.device != "cpu" else "cpu"
|
||||||
|
voice_sample = voice_mapper.get_voice_path(args.speaker_name)
|
||||||
|
print(f"Loading voice preset: {voice_sample}")
|
||||||
|
all_prefilled_outputs = torch.load(voice_sample, map_location=target_device, weights_only=False)
|
||||||
|
|
||||||
|
# Prepare inputs for the model
|
||||||
|
inputs = processor.process_input_with_cached_prompt(
|
||||||
|
text=full_script,
|
||||||
|
cached_prompt=all_prefilled_outputs,
|
||||||
|
padding=True,
|
||||||
|
return_tensors="pt",
|
||||||
|
return_attention_mask=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Move tensors to target device
|
||||||
|
for k, v in inputs.items():
|
||||||
|
if torch.is_tensor(v):
|
||||||
|
inputs[k] = v.to(target_device)
|
||||||
|
|
||||||
|
print(f"Starting generation with cfg_scale: {args.cfg_scale}, ddpm_steps: {args.ddpm_steps}")
|
||||||
|
|
||||||
|
# Generate audio
|
||||||
|
start_time = time.time()
|
||||||
|
outputs = model.generate(
|
||||||
|
**inputs,
|
||||||
|
max_new_tokens=None,
|
||||||
|
cfg_scale=args.cfg_scale,
|
||||||
|
tokenizer=processor.tokenizer,
|
||||||
|
generation_config={'do_sample': False},
|
||||||
|
verbose=True,
|
||||||
|
all_prefilled_outputs=copy.deepcopy(all_prefilled_outputs) if all_prefilled_outputs is not None else None,
|
||||||
|
)
|
||||||
|
generation_time = time.time() - start_time
|
||||||
|
print(f"Generation time: {generation_time:.2f} seconds")
|
||||||
|
|
||||||
|
# Calculate audio duration and additional metrics
|
||||||
|
if outputs.speech_outputs and outputs.speech_outputs[0] is not None:
|
||||||
|
sample_rate = 24000 # 24kHz sample rate
|
||||||
|
audio_samples = outputs.speech_outputs[0].shape[-1] if len(outputs.speech_outputs[0].shape) > 0 else len(outputs.speech_outputs[0])
|
||||||
|
audio_duration = audio_samples / sample_rate
|
||||||
|
rtf = generation_time / audio_duration if audio_duration > 0 else float('inf')
|
||||||
|
|
||||||
|
print(f"Generated audio duration: {audio_duration:.2f} seconds")
|
||||||
|
print(f"RTF (Real Time Factor): {rtf:.2f}x")
|
||||||
|
else:
|
||||||
|
print("No audio output generated")
|
||||||
|
audio_duration = 0
|
||||||
|
rtf = float('inf')
|
||||||
|
|
||||||
|
# Calculate token metrics
|
||||||
|
input_tokens = inputs['tts_text_ids'].shape[1]
|
||||||
|
output_tokens = outputs.sequences.shape[1]
|
||||||
|
generated_tokens = output_tokens - input_tokens - all_prefilled_outputs['tts_lm']['last_hidden_state'].size(1)
|
||||||
|
|
||||||
|
print(f"Prefilling text tokens: {input_tokens}")
|
||||||
|
print(f"Generated speech tokens: {generated_tokens}")
|
||||||
|
print(f"Total tokens: {output_tokens}")
|
||||||
|
|
||||||
|
# Save output
|
||||||
|
txt_filename = os.path.splitext(os.path.basename(args.txt_path))[0]
|
||||||
|
output_path = os.path.join(args.output_dir, f"{txt_filename}_streaming_generated.wav")
|
||||||
|
os.makedirs(args.output_dir, exist_ok=True)
|
||||||
|
|
||||||
|
processor.save_audio(
|
||||||
|
outputs.speech_outputs[0],
|
||||||
|
output_path=output_path,
|
||||||
|
)
|
||||||
|
print(f"Saved output to {output_path}")
|
||||||
|
|
||||||
|
# Print summary
|
||||||
|
print("\n" + "=" * 50)
|
||||||
|
print("GENERATION SUMMARY (Streaming 0.5B)")
|
||||||
|
print("=" * 50)
|
||||||
|
print(f"Model: {args.model_path}")
|
||||||
|
print(f"Input file: {args.txt_path}")
|
||||||
|
print(f"Output file: {output_path}")
|
||||||
|
print(f"Speaker: {args.speaker_name}")
|
||||||
|
print(f"CFG Scale: {args.cfg_scale}")
|
||||||
|
print(f"DDPM Steps: {args.ddpm_steps}")
|
||||||
|
print(f"Prefilling text tokens: {input_tokens}")
|
||||||
|
print(f"Generated speech tokens: {generated_tokens}")
|
||||||
|
print(f"Total tokens: {output_tokens}")
|
||||||
|
print(f"Generation time: {generation_time:.2f} seconds")
|
||||||
|
print(f"Audio duration: {audio_duration:.2f} seconds")
|
||||||
|
print(f"RTF (Real Time Factor): {rtf:.2f}x")
|
||||||
|
print("=" * 50)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
@@ -0,0 +1,7 @@
|
|||||||
|
Hello and welcome! Today I'm going to tell you about VibeVoice, a groundbreaking text-to-speech model developed for generating natural, expressive speech.
|
||||||
|
|
||||||
|
VibeVoice uses an innovative architecture combining autoregressive language modeling with diffusion-based audio generation. This allows it to produce remarkably natural-sounding speech with proper intonation and rhythm.
|
||||||
|
|
||||||
|
The streaming version you're hearing right now is optimized for real-time applications, with low latency that makes it perfect for interactive use cases like virtual assistants and live applications.
|
||||||
|
|
||||||
|
Thank you for listening to this demonstration!
|
||||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -36,6 +36,12 @@ dependencies = [
|
|||||||
"aiortc"
|
"aiortc"
|
||||||
]
|
]
|
||||||
|
|
||||||
|
[project.optional-dependencies]
|
||||||
|
# Optional dependencies for streaming model web service
|
||||||
|
streaming-web = [
|
||||||
|
"fastapi",
|
||||||
|
"uvicorn[standard]",
|
||||||
|
]
|
||||||
|
|
||||||
[project.urls]
|
[project.urls]
|
||||||
"Homepage" = "https://github.com/microsoft/VibeVoice"
|
"Homepage" = "https://github.com/microsoft/VibeVoice"
|
||||||
|
|||||||
@@ -0,0 +1,96 @@
|
|||||||
|
"""
|
||||||
|
VibeVoice Modular Components
|
||||||
|
|
||||||
|
This module provides the core model architectures for VibeVoice:
|
||||||
|
- Multi-speaker models (1.5B, 7B) for high-quality multi-speaker TTS
|
||||||
|
- Streaming model (0.5B) for real-time low-latency TTS
|
||||||
|
"""
|
||||||
|
|
||||||
|
# Multi-speaker model components
|
||||||
|
from .configuration_vibevoice import (
|
||||||
|
VibeVoiceConfig,
|
||||||
|
VibeVoiceAcousticTokenizerConfig,
|
||||||
|
VibeVoiceSemanticTokenizerConfig,
|
||||||
|
VibeVoiceDiffusionHeadConfig,
|
||||||
|
)
|
||||||
|
from .modeling_vibevoice import (
|
||||||
|
VibeVoicePreTrainedModel,
|
||||||
|
VibeVoiceModel,
|
||||||
|
)
|
||||||
|
from .modeling_vibevoice_inference import (
|
||||||
|
VibeVoiceForConditionalGenerationInference,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Streaming model components (0.5B)
|
||||||
|
from .configuration_vibevoice_streaming import (
|
||||||
|
VibeVoiceStreamingConfig,
|
||||||
|
)
|
||||||
|
from .modeling_vibevoice_streaming import (
|
||||||
|
VibeVoiceStreamingPreTrainedModel,
|
||||||
|
VibeVoiceStreamingModel,
|
||||||
|
BinaryClassifier,
|
||||||
|
SpeechConnector,
|
||||||
|
)
|
||||||
|
from .modeling_vibevoice_streaming_inference import (
|
||||||
|
VibeVoiceStreamingForConditionalGenerationInference,
|
||||||
|
VibeVoiceGenerationOutput,
|
||||||
|
TTS_TEXT_WINDOW_SIZE,
|
||||||
|
TTS_SPEECH_WINDOW_SIZE,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Shared components
|
||||||
|
from .modular_vibevoice_tokenizer import (
|
||||||
|
VibeVoiceTokenizerStreamingCache,
|
||||||
|
VibeVoiceAcousticTokenizerModel,
|
||||||
|
VibeVoiceSemanticTokenizerModel,
|
||||||
|
)
|
||||||
|
from .modular_vibevoice_text_tokenizer import (
|
||||||
|
VibeVoiceTextTokenizer,
|
||||||
|
VibeVoiceTextTokenizerFast,
|
||||||
|
)
|
||||||
|
from .modular_vibevoice_diffusion_head import (
|
||||||
|
VibeVoiceDiffusionHead,
|
||||||
|
)
|
||||||
|
from .streamer import (
|
||||||
|
AudioStreamer,
|
||||||
|
AsyncAudioStreamer,
|
||||||
|
)
|
||||||
|
|
||||||
|
# LoRA support
|
||||||
|
from .lora_loading import (
|
||||||
|
load_lora_assets,
|
||||||
|
)
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
# Multi-speaker configs
|
||||||
|
"VibeVoiceConfig",
|
||||||
|
"VibeVoiceAcousticTokenizerConfig",
|
||||||
|
"VibeVoiceSemanticTokenizerConfig",
|
||||||
|
"VibeVoiceDiffusionHeadConfig",
|
||||||
|
# Multi-speaker models
|
||||||
|
"VibeVoicePreTrainedModel",
|
||||||
|
"VibeVoiceModel",
|
||||||
|
"VibeVoiceForConditionalGenerationInference",
|
||||||
|
# Streaming configs
|
||||||
|
"VibeVoiceStreamingConfig",
|
||||||
|
# Streaming models
|
||||||
|
"VibeVoiceStreamingPreTrainedModel",
|
||||||
|
"VibeVoiceStreamingModel",
|
||||||
|
"VibeVoiceStreamingForConditionalGenerationInference",
|
||||||
|
"VibeVoiceGenerationOutput",
|
||||||
|
"BinaryClassifier",
|
||||||
|
"SpeechConnector",
|
||||||
|
"TTS_TEXT_WINDOW_SIZE",
|
||||||
|
"TTS_SPEECH_WINDOW_SIZE",
|
||||||
|
# Shared components
|
||||||
|
"VibeVoiceTokenizerStreamingCache",
|
||||||
|
"VibeVoiceAcousticTokenizerModel",
|
||||||
|
"VibeVoiceSemanticTokenizerModel",
|
||||||
|
"VibeVoiceTextTokenizer",
|
||||||
|
"VibeVoiceTextTokenizerFast",
|
||||||
|
"VibeVoiceDiffusionHead",
|
||||||
|
"AudioStreamer",
|
||||||
|
"AsyncAudioStreamer",
|
||||||
|
# LoRA
|
||||||
|
"load_lora_assets",
|
||||||
|
]
|
||||||
|
|||||||
@@ -0,0 +1,96 @@
|
|||||||
|
""" VibeVoice Streaming model configuration"""
|
||||||
|
|
||||||
|
from transformers.configuration_utils import PretrainedConfig
|
||||||
|
from transformers.utils import logging
|
||||||
|
|
||||||
|
from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
|
||||||
|
|
||||||
|
from .configuration_vibevoice import VibeVoiceAcousticTokenizerConfig, VibeVoiceDiffusionHeadConfig
|
||||||
|
|
||||||
|
logger = logging.get_logger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class VibeVoiceStreamingConfig(PretrainedConfig):
|
||||||
|
"""
|
||||||
|
Configuration class for the VibeVoice Streaming model (0.5B).
|
||||||
|
|
||||||
|
The streaming model differs from the multi-speaker model:
|
||||||
|
- No semantic tokenizer (only acoustic)
|
||||||
|
- Split language model: lower layers for text encoding, upper layers for TTS
|
||||||
|
- Optimized for low-latency real-time generation
|
||||||
|
|
||||||
|
Args:
|
||||||
|
acoustic_tokenizer_config: Configuration for the acoustic tokenizer
|
||||||
|
decoder_config: Configuration for the Qwen2 language model backbone
|
||||||
|
diffusion_head_config: Configuration for the diffusion prediction head
|
||||||
|
tts_backbone_num_hidden_layers: Number of upper transformer layers used for TTS (default: 20)
|
||||||
|
"""
|
||||||
|
model_type = "vibevoice_streaming"
|
||||||
|
is_composition = True
|
||||||
|
sub_configs = {
|
||||||
|
"acoustic_tokenizer_config": VibeVoiceAcousticTokenizerConfig,
|
||||||
|
"decoder_config": Qwen2Config,
|
||||||
|
"diffusion_head_config": VibeVoiceDiffusionHeadConfig,
|
||||||
|
}
|
||||||
|
# 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,
|
||||||
|
decoder_config=None,
|
||||||
|
diffusion_head_config=None,
|
||||||
|
tts_backbone_num_hidden_layers=20,
|
||||||
|
**kwargs
|
||||||
|
):
|
||||||
|
|
||||||
|
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 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
|
||||||
|
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
|
||||||
|
|
||||||
|
if diffusion_head_config is None:
|
||||||
|
self.diffusion_head_config = self.sub_configs["diffusion_head_config"]()
|
||||||
|
elif isinstance(diffusion_head_config, dict):
|
||||||
|
diffusion_head_config["model_type"] = "vibevoice_diffusion_head"
|
||||||
|
self.diffusion_head_config = self.sub_configs["diffusion_head_config"](**diffusion_head_config)
|
||||||
|
elif isinstance(diffusion_head_config, VibeVoiceDiffusionHeadConfig):
|
||||||
|
# If an instance of the config class is provided
|
||||||
|
self.diffusion_head_config = diffusion_head_config
|
||||||
|
|
||||||
|
# other parameters
|
||||||
|
self.acoustic_vae_dim = getattr(self.acoustic_tokenizer_config, 'vae_dim', 64)
|
||||||
|
# The decoder of the model is divided into two components. The lower Transformer layers are only used for encoding text, while the upper Transformer layers are used for encoding text and generating speech. `tts_backbone_num_hidden_layers` indicates the number of upper layers used for TTS.
|
||||||
|
self.tts_backbone_num_hidden_layers = tts_backbone_num_hidden_layers
|
||||||
|
|
||||||
|
super().__init__(**kwargs)
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"VibeVoiceStreamingConfig"
|
||||||
|
]
|
||||||
@@ -0,0 +1,216 @@
|
|||||||
|
"""
|
||||||
|
VibeVoice Streaming Model Architecture (0.5B)
|
||||||
|
|
||||||
|
This module implements the streaming-optimized version of VibeVoice for real-time TTS.
|
||||||
|
Key differences from the multi-speaker model:
|
||||||
|
- No semantic tokenizer (only acoustic)
|
||||||
|
- Split language model architecture: lower layers for text, upper layers for TTS
|
||||||
|
- Optimized for low-latency generation
|
||||||
|
"""
|
||||||
|
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from typing import Dict, List, Optional, Tuple, Union, Callable
|
||||||
|
from tqdm import tqdm
|
||||||
|
import copy
|
||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
import torch.nn.functional as F
|
||||||
|
import torch.distributed as dist
|
||||||
|
|
||||||
|
from transformers.models.auto import AutoModel, AutoModelForCausalLM
|
||||||
|
|
||||||
|
from transformers.activations import ACT2FN
|
||||||
|
from transformers.modeling_outputs import CausalLMOutput, BaseModelOutputWithPast, ModelOutput
|
||||||
|
from transformers.models.llama.modeling_llama import LlamaRMSNorm
|
||||||
|
from transformers import modeling_utils
|
||||||
|
from transformers.modeling_utils import PreTrainedModel
|
||||||
|
from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
|
||||||
|
from transformers.utils import logging
|
||||||
|
|
||||||
|
from .modular_vibevoice_diffusion_head import VibeVoiceDiffusionHead
|
||||||
|
from vibevoice.schedule.dpm_solver import DPMSolverMultistepScheduler
|
||||||
|
|
||||||
|
from .configuration_vibevoice_streaming import VibeVoiceStreamingConfig
|
||||||
|
|
||||||
|
|
||||||
|
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"]
|
||||||
|
|
||||||
|
|
||||||
|
class BinaryClassifier(nn.Module):
|
||||||
|
"""Binary classifier for end-of-speech detection in streaming TTS."""
|
||||||
|
|
||||||
|
def __init__(self, hidden_size):
|
||||||
|
super(BinaryClassifier, self).__init__()
|
||||||
|
self.fc1 = nn.Linear(hidden_size, hidden_size)
|
||||||
|
self.fc2 = nn.Linear(hidden_size, 1)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
x = torch.relu(self.fc1(x))
|
||||||
|
x = self.fc2(x)
|
||||||
|
return x
|
||||||
|
|
||||||
|
|
||||||
|
class SpeechConnector(nn.Module):
|
||||||
|
"""Connector module that projects speech latents to language model hidden dimension."""
|
||||||
|
|
||||||
|
def __init__(self, input_dim, output_dim):
|
||||||
|
super().__init__()
|
||||||
|
self.fc1 = nn.Linear(input_dim, output_dim)
|
||||||
|
self.norm = LlamaRMSNorm(output_dim, eps=1e-6)
|
||||||
|
self.fc2 = nn.Linear(output_dim, output_dim)
|
||||||
|
|
||||||
|
def forward(self, features, **kwargs):
|
||||||
|
x = self.fc1(features)
|
||||||
|
x = self.norm(x)
|
||||||
|
x = self.fc2(x)
|
||||||
|
return x
|
||||||
|
|
||||||
|
|
||||||
|
class VibeVoiceStreamingPreTrainedModel(PreTrainedModel):
|
||||||
|
"""Base class for VibeVoice Streaming models."""
|
||||||
|
|
||||||
|
config_class = VibeVoiceStreamingConfig
|
||||||
|
base_model_prefix = "model"
|
||||||
|
supports_gradient_checkpointing = True
|
||||||
|
_skip_keys_device_placement = "past_key_values"
|
||||||
|
_supports_cache_class = 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):
|
||||||
|
if isinstance(module, VibeVoiceDiffusionHead):
|
||||||
|
module.initialize_weights()
|
||||||
|
return
|
||||||
|
|
||||||
|
# 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_()
|
||||||
|
|
||||||
|
|
||||||
|
class VibeVoiceStreamingModel(VibeVoiceStreamingPreTrainedModel):
|
||||||
|
"""
|
||||||
|
VibeVoice Streaming Model for real-time TTS.
|
||||||
|
|
||||||
|
The model uses a split architecture:
|
||||||
|
- language_model: Lower transformer layers for text encoding
|
||||||
|
- tts_language_model: Upper transformer layers for TTS generation
|
||||||
|
|
||||||
|
This separation enables streaming text input and low-latency speech output.
|
||||||
|
"""
|
||||||
|
|
||||||
|
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.
|
||||||
|
# The lower Transformer layers are only used for encoding text, while the upper Transformer layers are used for encoding text and generating speech.
|
||||||
|
# To keep the code clean, we constructs two language models.
|
||||||
|
# The final norm layer of the first language_model is set to identity and will not be used in inference.
|
||||||
|
lm_config = copy.deepcopy(config.decoder_config)
|
||||||
|
lm_backbone_num_hidden_layers = getattr(lm_config, 'num_hidden_layers', 24) - config.tts_backbone_num_hidden_layers
|
||||||
|
lm_config.num_hidden_layers = lm_backbone_num_hidden_layers
|
||||||
|
self.language_model = AutoModel.from_config(lm_config)
|
||||||
|
self.language_model.norm = nn.Identity()
|
||||||
|
|
||||||
|
# We only need the Transformer layers here. Note that embed_tokens in tts_language_model is unused
|
||||||
|
tts_lm_config = copy.deepcopy(lm_config)
|
||||||
|
tts_lm_config.num_hidden_layers = config.tts_backbone_num_hidden_layers
|
||||||
|
self.tts_language_model = AutoModel.from_config(tts_lm_config)
|
||||||
|
|
||||||
|
# Marks the text that needs to be spoken by the TTS model.
|
||||||
|
self.tts_input_types = nn.Embedding(num_embeddings=2, embedding_dim=config.decoder_config.hidden_size)
|
||||||
|
|
||||||
|
# Initialize speech components if needed
|
||||||
|
self.acoustic_tokenizer = AutoModel.from_config(config.acoustic_tokenizer_config).to(dtype)
|
||||||
|
self.acoustic_connector = SpeechConnector(config.acoustic_vae_dim, lm_config.hidden_size).to(dtype)
|
||||||
|
|
||||||
|
# Register scaling factors as buffers - use 1D tensors for FSDP compatibility
|
||||||
|
self.register_buffer('speech_scaling_factor', torch.tensor(float('nan')))
|
||||||
|
self.register_buffer('speech_bias_factor', torch.tensor(float('nan')))
|
||||||
|
|
||||||
|
# Initialize prediction head for speech generation
|
||||||
|
self.prediction_head = AutoModel.from_config(config.diffusion_head_config).to(dtype)
|
||||||
|
|
||||||
|
# Initialize noise scheduler
|
||||||
|
self.noise_scheduler = DPMSolverMultistepScheduler(
|
||||||
|
num_train_timesteps=config.diffusion_head_config.ddpm_num_steps,
|
||||||
|
beta_schedule=config.diffusion_head_config.ddpm_beta_schedule,
|
||||||
|
prediction_type=config.diffusion_head_config.prediction_type
|
||||||
|
)
|
||||||
|
|
||||||
|
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):
|
||||||
|
"""Set the speech tokenizers used for encoding and decoding speech."""
|
||||||
|
self.acoustic_tokenizer = acoustic_tokenizer
|
||||||
|
|
||||||
|
# Reset the encoder to evaluation mode
|
||||||
|
if self.acoustic_tokenizer is not None:
|
||||||
|
self.acoustic_tokenizer.eval()
|
||||||
|
|
||||||
|
def forward(self, *args, **kwargs):
|
||||||
|
"""
|
||||||
|
Intentionally not implemented.
|
||||||
|
|
||||||
|
This streaming model is split into two explicit submodules:
|
||||||
|
- `language_model` for plain text processing (lower layers).
|
||||||
|
- `tts_language_model` for TTS-related upper layers.
|
||||||
|
|
||||||
|
We deliberately avoid a unified `forward` to prevent accidental calls
|
||||||
|
that mix responsibilities.
|
||||||
|
|
||||||
|
To use the model:
|
||||||
|
- Call `self.language_model(...)` for text embeddings / hidden states.
|
||||||
|
- Call `self.tts_language_model(...)` for the TTS portion.
|
||||||
|
- Use the dedicated inference class for combined generation logic.
|
||||||
|
"""
|
||||||
|
raise RuntimeError(
|
||||||
|
"VibeVoiceStreamingModel.forward is intentionally disabled. "
|
||||||
|
"Use `model.language_model(...)` or `model.tts_language_model(...)` instead."
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
AutoModel.register(VibeVoiceStreamingConfig, VibeVoiceStreamingModel)
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"VibeVoiceStreamingPreTrainedModel",
|
||||||
|
"VibeVoiceStreamingModel",
|
||||||
|
"BinaryClassifier",
|
||||||
|
"SpeechConnector",
|
||||||
|
]
|
||||||
@@ -0,0 +1,762 @@
|
|||||||
|
"""
|
||||||
|
VibeVoice Streaming Inference Model (0.5B)
|
||||||
|
|
||||||
|
This module implements the inference engine for real-time streaming TTS.
|
||||||
|
Key features:
|
||||||
|
- Window-based text/speech interleaving for streaming
|
||||||
|
- Binary EOS classifier for end-of-speech detection
|
||||||
|
- Classifier-free guidance for speech quality
|
||||||
|
- Audio streaming support
|
||||||
|
"""
|
||||||
|
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from typing import Any, Dict, List, Optional, Tuple, Union, Callable
|
||||||
|
from tqdm import tqdm
|
||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
|
||||||
|
from transformers.models.auto import AutoModel, AutoModelForCausalLM
|
||||||
|
|
||||||
|
from transformers.generation import GenerationMixin, GenerationConfig, LogitsProcessor, LogitsProcessorList, StoppingCriteriaList
|
||||||
|
from transformers.modeling_outputs import BaseModelOutputWithPast, ModelOutput
|
||||||
|
from transformers import modeling_utils
|
||||||
|
from transformers.modeling_utils import PreTrainedModel
|
||||||
|
from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
|
||||||
|
from transformers.utils import logging
|
||||||
|
|
||||||
|
from .modular_vibevoice_tokenizer import VibeVoiceTokenizerStreamingCache
|
||||||
|
from .modular_vibevoice_diffusion_head import VibeVoiceDiffusionHead
|
||||||
|
from vibevoice.schedule.dpm_solver import DPMSolverMultistepScheduler
|
||||||
|
from .configuration_vibevoice_streaming import VibeVoiceStreamingConfig
|
||||||
|
from .modular_vibevoice_text_tokenizer import VibeVoiceTextTokenizer, VibeVoiceTextTokenizerFast
|
||||||
|
from .modeling_vibevoice_streaming import VibeVoiceStreamingPreTrainedModel, VibeVoiceStreamingModel, BinaryClassifier
|
||||||
|
from .streamer import AudioStreamer, AsyncAudioStreamer
|
||||||
|
|
||||||
|
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"]
|
||||||
|
|
||||||
|
# Window sizes for streaming generation
|
||||||
|
TTS_TEXT_WINDOW_SIZE = 5 # Process 5 text tokens at a time
|
||||||
|
TTS_SPEECH_WINDOW_SIZE = 6 # Generate 6 speech frames per window
|
||||||
|
|
||||||
|
|
||||||
|
def _update_model_kwargs_for_generation(
|
||||||
|
outputs: ModelOutput,
|
||||||
|
model_kwargs: Dict[str, Any],
|
||||||
|
num_new_tokens: int = 1,
|
||||||
|
) -> Dict[str, Any]:
|
||||||
|
"""
|
||||||
|
Update model_kwargs after adding new tokens.
|
||||||
|
|
||||||
|
Mainly for the case num_new_tokens > 1 (e.g. a whole text window):
|
||||||
|
- past_key_values: take from current outputs
|
||||||
|
- attention_mask: append num_new_tokens ones
|
||||||
|
- cache_position: advance by creating a range for all new positions
|
||||||
|
"""
|
||||||
|
|
||||||
|
# update past_key_values keeping its naming used in model code
|
||||||
|
model_kwargs["past_key_values"] = getattr(outputs, "past_key_values")
|
||||||
|
|
||||||
|
attention_mask = model_kwargs["attention_mask"]
|
||||||
|
model_kwargs["attention_mask"] = torch.cat(
|
||||||
|
[attention_mask, attention_mask.new_ones((attention_mask.shape[0], num_new_tokens))], dim=-1
|
||||||
|
)
|
||||||
|
|
||||||
|
model_kwargs["cache_position"] = torch.arange(model_kwargs["cache_position"][-1] + 1, model_kwargs["cache_position"][-1] + num_new_tokens + 1).to(model_kwargs["cache_position"].device)
|
||||||
|
|
||||||
|
return model_kwargs
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class VibeVoiceCausalLMOutputWithPast(BaseModelOutputWithPast):
|
||||||
|
logits: Optional[torch.FloatTensor] = None
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class VibeVoiceGenerationOutput(ModelOutput):
|
||||||
|
"""
|
||||||
|
Output type for VibeVoice generation.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
sequences (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
|
||||||
|
The generated sequences.
|
||||||
|
speech_outputs (`List[torch.FloatTensor]`, *optional*):
|
||||||
|
List of generated speech waveforms or latents for each speech segment.
|
||||||
|
"""
|
||||||
|
sequences: torch.LongTensor = None
|
||||||
|
speech_outputs: Optional[List[torch.FloatTensor]] = None
|
||||||
|
reach_max_step_sample: Optional[torch.BoolTensor] = None
|
||||||
|
|
||||||
|
|
||||||
|
class VibeVoiceStreamingForConditionalGenerationInference(VibeVoiceStreamingPreTrainedModel, GenerationMixin):
|
||||||
|
"""
|
||||||
|
VibeVoice Streaming model for conditional generation inference.
|
||||||
|
|
||||||
|
This class handles the full streaming TTS pipeline:
|
||||||
|
1. Prefill text prompt with cached voice embeddings
|
||||||
|
2. Window-based text/speech interleaving
|
||||||
|
3. Diffusion-based speech token generation
|
||||||
|
4. Binary EOS detection for natural stopping
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, config):
|
||||||
|
super().__init__(config)
|
||||||
|
|
||||||
|
# Initialize the base model
|
||||||
|
self.model = VibeVoiceStreamingModel(config)
|
||||||
|
|
||||||
|
# TTS generation EOS classifier
|
||||||
|
self.tts_eos_classifier = BinaryClassifier(config.decoder_config.hidden_size)
|
||||||
|
|
||||||
|
# inference configuration
|
||||||
|
self.ddpm_inference_steps = config.diffusion_head_config.ddpm_num_inference_steps
|
||||||
|
|
||||||
|
# Initialize weights and apply final processing
|
||||||
|
self.post_init()
|
||||||
|
|
||||||
|
@property
|
||||||
|
def noise_scheduler(self):
|
||||||
|
return self.model.noise_scheduler
|
||||||
|
|
||||||
|
@property
|
||||||
|
def prediction_head(self):
|
||||||
|
return self.model.prediction_head
|
||||||
|
|
||||||
|
@property
|
||||||
|
def speech_scaling_factor(self):
|
||||||
|
return self.model.speech_scaling_factor
|
||||||
|
|
||||||
|
@property
|
||||||
|
def speech_bias_factor(self):
|
||||||
|
return self.model.speech_bias_factor
|
||||||
|
|
||||||
|
@property
|
||||||
|
def acoustic_tokenizer(self):
|
||||||
|
return self.model.acoustic_tokenizer
|
||||||
|
|
||||||
|
@property
|
||||||
|
def acoustic_connector(self):
|
||||||
|
return self.model.acoustic_connector
|
||||||
|
|
||||||
|
def tie_weights(self):
|
||||||
|
"""
|
||||||
|
Tie the weights between the input embeddings and the output embeddings.
|
||||||
|
"""
|
||||||
|
# Tie lm_head.weight to language_model.embed_tokens.weight
|
||||||
|
if not getattr(self.config, 'tie_word_embeddings', False):
|
||||||
|
return
|
||||||
|
|
||||||
|
if hasattr(self, 'lm_head') and hasattr(self.model.language_model, 'embed_tokens'):
|
||||||
|
self.lm_head.weight = self.model.language_model.embed_tokens.weight
|
||||||
|
|
||||||
|
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):
|
||||||
|
"""
|
||||||
|
This model does not define an `lm_head` (vocabulary projection).
|
||||||
|
"""
|
||||||
|
return None
|
||||||
|
|
||||||
|
def set_output_embeddings(self, new_embeddings):
|
||||||
|
"""
|
||||||
|
No-op because there is no `lm_head`. Provided only to satisfy optional API calls.
|
||||||
|
To enable, first create `self.lm_head` then allow assignment.
|
||||||
|
"""
|
||||||
|
raise RuntimeError("Output embeddings (lm_head) are not defined for this model. "
|
||||||
|
"Create one before calling set_output_embeddings if needed.")
|
||||||
|
|
||||||
|
def set_speech_tokenizers(self, acoustic_tokenizer=None):
|
||||||
|
"""Set the speech tokenizers used for encoding and decoding speech."""
|
||||||
|
self.model.set_speech_tokenizers(acoustic_tokenizer)
|
||||||
|
|
||||||
|
def set_ddpm_inference_steps(self, num_steps=None):
|
||||||
|
self.ddpm_inference_steps = num_steps or self.config.diffusion_head_config.ddpm_num_inference_steps
|
||||||
|
|
||||||
|
def forward_lm(
|
||||||
|
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,
|
||||||
|
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,
|
||||||
|
**kwargs,
|
||||||
|
) -> Union[Tuple, BaseModelOutputWithPast]:
|
||||||
|
"""
|
||||||
|
Single pass of the base text LM.
|
||||||
|
|
||||||
|
- Builds embeddings if `inputs_embeds` not provided.
|
||||||
|
- Uses (and returns) `past_key_values` when `use_cache=True`.
|
||||||
|
- No loss / no lm_head / no speech logic.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
input_ids: (B, S) token ids.
|
||||||
|
attention_mask: (B, S) mask.
|
||||||
|
past_key_values: cache from previous steps.
|
||||||
|
cache_position: positions for cached tokens.
|
||||||
|
labels: unsupported (will raise).
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
BaseModelOutputWithPast with `last_hidden_state` and `past_key_values`.
|
||||||
|
"""
|
||||||
|
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
||||||
|
|
||||||
|
# Get embeddings
|
||||||
|
if inputs_embeds is None:
|
||||||
|
inputs_embeds = self.model.get_input_embeddings()(input_ids)
|
||||||
|
|
||||||
|
outputs = self.model.language_model(
|
||||||
|
inputs_embeds=inputs_embeds,
|
||||||
|
attention_mask=attention_mask,
|
||||||
|
position_ids=position_ids,
|
||||||
|
past_key_values=past_key_values,
|
||||||
|
use_cache=use_cache,
|
||||||
|
output_attentions=output_attentions,
|
||||||
|
output_hidden_states=output_hidden_states,
|
||||||
|
return_dict=return_dict,
|
||||||
|
cache_position=cache_position,
|
||||||
|
**kwargs,
|
||||||
|
)
|
||||||
|
|
||||||
|
hidden_states = outputs[0] if not return_dict else outputs.last_hidden_state
|
||||||
|
|
||||||
|
if labels is not None:
|
||||||
|
raise NotImplementedError("Loss computation is not implemented in this version.")
|
||||||
|
|
||||||
|
return BaseModelOutputWithPast(
|
||||||
|
past_key_values=outputs.past_key_values,
|
||||||
|
last_hidden_state=hidden_states,
|
||||||
|
attentions=outputs.attentions,
|
||||||
|
)
|
||||||
|
|
||||||
|
def forward_tts_lm(
|
||||||
|
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,
|
||||||
|
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,
|
||||||
|
lm_last_hidden_state: Optional[torch.FloatTensor] = None,
|
||||||
|
tts_text_masks: Optional[torch.BoolTensor] = None,
|
||||||
|
**kwargs,
|
||||||
|
) -> Union[Tuple, VibeVoiceCausalLMOutputWithPast]:
|
||||||
|
"""
|
||||||
|
Single pass of the TTS LM.
|
||||||
|
|
||||||
|
- Overwrites tail embeddings with `lm_last_hidden_state`.
|
||||||
|
- Adds type embedding via `tts_text_masks` (1=text, 0=speech).
|
||||||
|
- Predicts EOS from last hidden state (binary classifier).
|
||||||
|
- No loss / no full acoustic decoding here.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
input_ids: (B, S) token ids.
|
||||||
|
attention_mask: (B, S) mask.
|
||||||
|
lm_last_hidden_state: (B, K, H) hidden states to splice into the tail.
|
||||||
|
tts_text_masks: (B, 1) mask marking current position as text(1)/speech(0).
|
||||||
|
past_key_values: cache from previous TTS steps.
|
||||||
|
cache_position: positions for cached tokens.
|
||||||
|
labels: unsupported (will raise).
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
VibeVoiceCausalLMOutputWithPast with `logits` (EOS), `last_hidden_state`, `past_key_values`.
|
||||||
|
"""
|
||||||
|
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
||||||
|
|
||||||
|
# Get embeddings
|
||||||
|
if inputs_embeds is None:
|
||||||
|
# Will be replaced with lm_last_hidden_state
|
||||||
|
inputs_embeds = self.model.get_input_embeddings()(input_ids)
|
||||||
|
|
||||||
|
# Replace the last part of inputs_embeds with lm_last_hidden_state
|
||||||
|
start_idx = inputs_embeds.shape[1] - lm_last_hidden_state.shape[1]
|
||||||
|
inputs_embeds[:, start_idx:, :] = lm_last_hidden_state
|
||||||
|
|
||||||
|
# Adds type embedding via `tts_text_masks`.
|
||||||
|
inputs_embeds = inputs_embeds + self.model.tts_input_types(tts_text_masks.long())
|
||||||
|
|
||||||
|
outputs = self.model.tts_language_model(
|
||||||
|
inputs_embeds=inputs_embeds,
|
||||||
|
attention_mask=attention_mask,
|
||||||
|
position_ids=position_ids,
|
||||||
|
past_key_values=past_key_values,
|
||||||
|
use_cache=use_cache,
|
||||||
|
output_attentions=output_attentions,
|
||||||
|
output_hidden_states=output_hidden_states,
|
||||||
|
return_dict=return_dict,
|
||||||
|
cache_position=cache_position,
|
||||||
|
**kwargs,
|
||||||
|
)
|
||||||
|
|
||||||
|
hidden_states = outputs[0] if not return_dict else outputs.last_hidden_state
|
||||||
|
logits = self.tts_eos_classifier(hidden_states[:, -1, :])
|
||||||
|
|
||||||
|
if labels is not None:
|
||||||
|
raise NotImplementedError("Loss computation is not implemented in this version.")
|
||||||
|
|
||||||
|
return VibeVoiceCausalLMOutputWithPast(
|
||||||
|
logits=logits,
|
||||||
|
past_key_values=outputs.past_key_values,
|
||||||
|
last_hidden_state=hidden_states,
|
||||||
|
attentions=outputs.attentions,
|
||||||
|
)
|
||||||
|
|
||||||
|
def forward(self, *args, **kwargs):
|
||||||
|
"""
|
||||||
|
Unified forward is intentionally disabled.
|
||||||
|
|
||||||
|
Reasons:
|
||||||
|
1. The inference pipeline is staged: base text LM, then TTS LM, plus streaming & diffusion handled in `generate`.
|
||||||
|
2. A monolithic call would hide required sequencing (prefill, window stepping, speech diffusion sampling).
|
||||||
|
|
||||||
|
Use instead:
|
||||||
|
- self.forward_lm(...) for a base text LM step (prefill or incremental).
|
||||||
|
- self.forward_tts_lm(...) for a single TTS LM step (needs LM hidden states).
|
||||||
|
- self.generate(...) for full streaming (text + speech + diffusion + audio assembly).
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
RuntimeError: Always (by design).
|
||||||
|
"""
|
||||||
|
raise RuntimeError(
|
||||||
|
"Unified forward is disabled. Use `forward_lm`, `forward_tts_lm`, or `generate` instead."
|
||||||
|
)
|
||||||
|
|
||||||
|
def _build_generate_config_model_kwargs(self, generation_config, inputs, tokenizer, return_processors=False, **kwargs):
|
||||||
|
if generation_config is None:
|
||||||
|
generation_config = GenerationConfig(
|
||||||
|
bos_token_id=tokenizer.bos_token_id,
|
||||||
|
eos_token_id=tokenizer.eos_token_id,
|
||||||
|
pad_token_id=tokenizer.pad_token_id
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
generation_config = GenerationConfig(
|
||||||
|
**generation_config,
|
||||||
|
bos_token_id=tokenizer.bos_token_id,
|
||||||
|
eos_token_id=tokenizer.eos_token_id,
|
||||||
|
pad_token_id=tokenizer.pad_token_id
|
||||||
|
)
|
||||||
|
|
||||||
|
generation_config, model_kwargs = self._prepare_generation_config(
|
||||||
|
generation_config,
|
||||||
|
True,
|
||||||
|
speech_start_id=tokenizer.speech_start_id,
|
||||||
|
speech_end_id=tokenizer.speech_end_id,
|
||||||
|
speech_diffusion_id=tokenizer.speech_diffusion_id,
|
||||||
|
**kwargs
|
||||||
|
)
|
||||||
|
generation_config.speech_start_id = tokenizer.speech_start_id
|
||||||
|
generation_config.speech_end_id = tokenizer.speech_end_id
|
||||||
|
generation_config.speech_diffusion_id = tokenizer.speech_diffusion_id
|
||||||
|
|
||||||
|
inputs_tensor, model_input_name, model_kwargs = self._prepare_model_inputs(inputs, generation_config.bos_token_id, model_kwargs)
|
||||||
|
batch_size = inputs_tensor.shape[0]
|
||||||
|
device = self.device
|
||||||
|
|
||||||
|
self._prepare_special_tokens(generation_config, True, device=device)
|
||||||
|
generation_config.use_cache = True
|
||||||
|
model_kwargs["use_cache"] = generation_config.use_cache
|
||||||
|
input_ids = inputs_tensor.to(self.device)
|
||||||
|
|
||||||
|
input_ids_length = input_ids.shape[1]
|
||||||
|
has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None
|
||||||
|
has_default_min_length = kwargs.get("min_length") is None and generation_config.min_length is not None
|
||||||
|
generation_config = self._prepare_generated_length(
|
||||||
|
generation_config=generation_config,
|
||||||
|
has_default_max_length=has_default_max_length,
|
||||||
|
has_default_min_length=has_default_min_length,
|
||||||
|
model_input_name=model_input_name,
|
||||||
|
inputs_tensor=inputs_tensor,
|
||||||
|
input_ids_length=input_ids_length,
|
||||||
|
)
|
||||||
|
|
||||||
|
max_cache_length = generation_config.max_length - 1
|
||||||
|
self._prepare_cache_for_generation(generation_config, model_kwargs, None, batch_size, max_cache_length, device)
|
||||||
|
model_kwargs['cache_position'] = torch.arange(input_ids_length, device=device, dtype=torch.long)
|
||||||
|
for k, v in model_kwargs.items():
|
||||||
|
if isinstance(v, torch.Tensor):
|
||||||
|
model_kwargs[k] = v.to(device=device)
|
||||||
|
|
||||||
|
if return_processors:
|
||||||
|
logits_processor = self._get_logits_processor(
|
||||||
|
generation_config=generation_config,
|
||||||
|
input_ids_seq_length=input_ids_length,
|
||||||
|
encoder_input_ids=inputs_tensor,
|
||||||
|
prefix_allowed_tokens_fn=None,
|
||||||
|
logits_processor=LogitsProcessorList(),
|
||||||
|
device=inputs_tensor.device,
|
||||||
|
model_kwargs=model_kwargs,
|
||||||
|
)
|
||||||
|
|
||||||
|
stopping_criteria = self._get_stopping_criteria(generation_config=generation_config, stopping_criteria=StoppingCriteriaList())
|
||||||
|
|
||||||
|
return generation_config, model_kwargs, input_ids, logits_processor, stopping_criteria
|
||||||
|
else:
|
||||||
|
return generation_config, model_kwargs, input_ids
|
||||||
|
|
||||||
|
@torch.no_grad()
|
||||||
|
def generate(
|
||||||
|
self,
|
||||||
|
inputs: Optional[torch.Tensor] = None,
|
||||||
|
generation_config: Optional[GenerationConfig] = None,
|
||||||
|
logits_processor: Optional[LogitsProcessorList] = None,
|
||||||
|
stopping_criteria: Optional[StoppingCriteriaList] = None,
|
||||||
|
prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None,
|
||||||
|
synced_gpus: Optional[bool] = None,
|
||||||
|
assistant_model: Optional["PreTrainedModel"] = None,
|
||||||
|
audio_streamer: Optional[Union[AudioStreamer, AsyncAudioStreamer]] = None,
|
||||||
|
negative_prompt_ids: Optional[torch.Tensor] = None,
|
||||||
|
negative_prompt_attention_mask: Optional[torch.Tensor] = None,
|
||||||
|
speech_tensors: Optional[torch.FloatTensor] = None,
|
||||||
|
speech_masks: Optional[torch.BoolTensor] = None,
|
||||||
|
speech_input_mask: Optional[torch.BoolTensor] = None,
|
||||||
|
tts_text_ids: Optional[torch.LongTensor] = None,
|
||||||
|
return_speech: bool = True,
|
||||||
|
cfg_scale: float = 1.0,
|
||||||
|
stop_check_fn: Optional[Callable[[], bool]] = None,
|
||||||
|
**kwargs,
|
||||||
|
) -> Union[torch.LongTensor, VibeVoiceGenerationOutput]:
|
||||||
|
"""
|
||||||
|
Generate speech from text using streaming TTS.
|
||||||
|
|
||||||
|
Text is fed in small windows (dynamic slicing of `tts_text_ids`), which enables streaming text input:
|
||||||
|
you don't need the full text upfront. After each text window, a loop samples several speech latents
|
||||||
|
(diffusion). The interleaved text encoding + speech generation enables streaming text input and
|
||||||
|
realtime speech output.
|
||||||
|
|
||||||
|
The function only supports batch size = 1 currently.
|
||||||
|
|
||||||
|
- Windowed text prefill → incremental LM + TTS LM updates.
|
||||||
|
- Interleave speech token diffusion sampling (`sample_speech_tokens`).
|
||||||
|
- Stops on EOS (binary classifier) or max length / external `stop_check_fn`.
|
||||||
|
- Returns final token `sequences` and (optionally) concatenated speech audio.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
inputs: Input tensor (usually from processor)
|
||||||
|
generation_config: Generation configuration
|
||||||
|
audio_streamer: If provided, emits audio chunks during generation
|
||||||
|
tts_text_ids: Full text tokens to stream in windows
|
||||||
|
cfg_scale: Classifier-free guidance scale for speech diffusion (default: 1.0)
|
||||||
|
return_speech: If False, skips audio decode concatenation
|
||||||
|
stop_check_fn: External early-stop hook (returns True to halt)
|
||||||
|
**kwargs: Additional arguments including tokenizer, all_prefilled_outputs, etc.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
VibeVoiceGenerationOutput with:
|
||||||
|
- sequences: final token ids
|
||||||
|
- speech_outputs: list of concatenated audio tensors (or None)
|
||||||
|
- reach_max_step_sample: flags for samples stopped by max length
|
||||||
|
"""
|
||||||
|
# 1. Handle `generation_config` and kwargs that might update it, and validate the `.generate()` call
|
||||||
|
tokenizer = kwargs.pop("tokenizer", None)
|
||||||
|
neg_text_input_id = tokenizer.convert_tokens_to_ids("<|image_pad|>")
|
||||||
|
|
||||||
|
tts_lm_input_ids = kwargs.pop("tts_lm_input_ids", None)
|
||||||
|
tts_lm_attention_mask = kwargs.pop("tts_lm_attention_mask", None)
|
||||||
|
# all_prefilled_outputs: cached prefilled prompt outputs for lm, tts_lm, neg_lm, neg_tts_lm
|
||||||
|
all_prefilled_outputs = kwargs.pop("all_prefilled_outputs", None)
|
||||||
|
tts_text_ids = tts_text_ids.to(self.device)
|
||||||
|
|
||||||
|
if kwargs.get('max_new_tokens', None) is None:
|
||||||
|
kwargs['max_new_tokens'] = self.config.decoder_config.max_position_embeddings - tts_lm_input_ids.shape[-1]
|
||||||
|
|
||||||
|
generation_config, model_kwargs, input_ids, logits_processor, stopping_criteria = self._build_generate_config_model_kwargs(
|
||||||
|
generation_config, inputs, tokenizer, return_processors=True, **kwargs
|
||||||
|
)
|
||||||
|
|
||||||
|
negative_kwargs = {
|
||||||
|
'input_ids': torch.full((kwargs['input_ids'].shape[0], 1), neg_text_input_id, dtype=torch.long, device=kwargs['input_ids'].device),
|
||||||
|
'attention_mask': torch.ones((kwargs['input_ids'].shape[0], 1), dtype=torch.long, device=kwargs['input_ids'].device),
|
||||||
|
'max_new_tokens': kwargs.get('max_new_tokens', 100)
|
||||||
|
}
|
||||||
|
negative_generation_config, negative_model_kwargs, negative_input_ids = self._build_generate_config_model_kwargs(
|
||||||
|
None, None, tokenizer, return_processors=False, **negative_kwargs
|
||||||
|
)
|
||||||
|
|
||||||
|
tts_lm_kwargs = {
|
||||||
|
'input_ids': tts_lm_input_ids,
|
||||||
|
'attention_mask': tts_lm_attention_mask,
|
||||||
|
'max_new_tokens': kwargs.get('max_new_tokens', 100)
|
||||||
|
}
|
||||||
|
tts_lm_generation_config, tts_lm_model_kwargs, tts_lm_input_ids = self._build_generate_config_model_kwargs(
|
||||||
|
None, None, tokenizer, return_processors=False, **tts_lm_kwargs
|
||||||
|
)
|
||||||
|
|
||||||
|
tts_lm_negative_kwargs = {
|
||||||
|
'input_ids': torch.full((kwargs['input_ids'].shape[0], 1), neg_text_input_id, dtype=torch.long, device=kwargs['input_ids'].device),
|
||||||
|
'attention_mask': torch.ones((kwargs['input_ids'].shape[0], 1), dtype=torch.long, device=kwargs['input_ids'].device),
|
||||||
|
'max_new_tokens': kwargs.get('max_new_tokens', 100)
|
||||||
|
}
|
||||||
|
tts_lm_negative_generation_config, tts_lm_negative_model_kwargs, tts_lm_negative_input_ids = self._build_generate_config_model_kwargs(
|
||||||
|
None, None, tokenizer, return_processors=False, **tts_lm_negative_kwargs
|
||||||
|
)
|
||||||
|
|
||||||
|
acoustic_cache = VibeVoiceTokenizerStreamingCache()
|
||||||
|
batch_size = input_ids.shape[0]
|
||||||
|
assert batch_size == 1, "Currently only supports batch size == 1"
|
||||||
|
device = input_ids.device
|
||||||
|
finished_tags = torch.zeros(batch_size, dtype=torch.bool, device=device)
|
||||||
|
verbose = kwargs.get("verbose", False)
|
||||||
|
|
||||||
|
# Initialize audio chunks storage for each sample
|
||||||
|
audio_chunks = [[] for _ in range(batch_size)]
|
||||||
|
tts_text_window_index = 0
|
||||||
|
reach_max_step_sample = torch.zeros(batch_size, dtype=torch.bool, device=device)
|
||||||
|
first_text_window_size = TTS_TEXT_WINDOW_SIZE if tts_text_ids.shape[1] >= TTS_TEXT_WINDOW_SIZE else tts_text_ids.shape[1]
|
||||||
|
|
||||||
|
outputs = all_prefilled_outputs["lm"]
|
||||||
|
tts_lm_outputs = all_prefilled_outputs["tts_lm"]
|
||||||
|
negative_outputs = all_prefilled_outputs["neg_lm"]
|
||||||
|
tts_lm_negative_outputs = all_prefilled_outputs["neg_tts_lm"]
|
||||||
|
|
||||||
|
model_kwargs = _update_model_kwargs_for_generation(
|
||||||
|
outputs, model_kwargs, num_new_tokens=first_text_window_size,
|
||||||
|
)
|
||||||
|
tts_lm_model_kwargs = _update_model_kwargs_for_generation(
|
||||||
|
tts_lm_outputs, tts_lm_model_kwargs, num_new_tokens=first_text_window_size,
|
||||||
|
)
|
||||||
|
negative_model_kwargs = self._update_model_kwargs_for_generation(
|
||||||
|
negative_outputs, negative_model_kwargs, is_encoder_decoder=False,
|
||||||
|
)
|
||||||
|
tts_lm_negative_model_kwargs = self._update_model_kwargs_for_generation(
|
||||||
|
tts_lm_negative_outputs, tts_lm_negative_model_kwargs, is_encoder_decoder=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
step = tts_lm_input_ids.shape[1]
|
||||||
|
total_generated_speech_tokens = 0
|
||||||
|
total_prefilled_text_tokens = 0
|
||||||
|
if kwargs.get("show_progress_bar", True):
|
||||||
|
progress_bar = tqdm(
|
||||||
|
total=tts_lm_generation_config.max_length,
|
||||||
|
desc=f"Prefilled {step} tokens, current step ({step} / {tts_lm_generation_config.max_length})",
|
||||||
|
initial=step,
|
||||||
|
leave=False
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
progress_bar = None
|
||||||
|
|
||||||
|
while True:
|
||||||
|
# Check for external stop signal
|
||||||
|
if stop_check_fn is not None and stop_check_fn():
|
||||||
|
if verbose:
|
||||||
|
print(f"Generation stopped externally at step {step + 1}")
|
||||||
|
# End the audio streamer if it exists
|
||||||
|
if audio_streamer is not None:
|
||||||
|
audio_streamer.end()
|
||||||
|
break
|
||||||
|
|
||||||
|
if finished_tags.all():
|
||||||
|
if hasattr(progress_bar, 'set_description'):
|
||||||
|
progress_bar.set_description("Generation complete")
|
||||||
|
break
|
||||||
|
|
||||||
|
cur_input_tts_text_ids = tts_text_ids[:, tts_text_window_index*TTS_TEXT_WINDOW_SIZE:(tts_text_window_index+1)*TTS_TEXT_WINDOW_SIZE]
|
||||||
|
next_text_window_size = tts_text_ids[:, (tts_text_window_index+1)*TTS_TEXT_WINDOW_SIZE:(tts_text_window_index+2)*TTS_TEXT_WINDOW_SIZE].shape[1]
|
||||||
|
tts_text_window_index += 1
|
||||||
|
|
||||||
|
if cur_input_tts_text_ids.shape[1] > 0:
|
||||||
|
input_ids = torch.cat([input_ids, cur_input_tts_text_ids], dim=-1)
|
||||||
|
tts_lm_input_ids = torch.cat([tts_lm_input_ids, cur_input_tts_text_ids], dim=-1)
|
||||||
|
|
||||||
|
if tts_lm_input_ids.shape[1] > tts_lm_generation_config.max_length:
|
||||||
|
if verbose:
|
||||||
|
print(f"Reached maximum generation length {generation_config.max_length}, stopped it.")
|
||||||
|
reached_samples = torch.arange(batch_size, device=device)[~finished_tags]
|
||||||
|
if reached_samples.numel() > 0:
|
||||||
|
reach_max_step_sample[reached_samples] = True
|
||||||
|
break
|
||||||
|
|
||||||
|
step += cur_input_tts_text_ids.shape[1]
|
||||||
|
total_prefilled_text_tokens += cur_input_tts_text_ids.shape[1]
|
||||||
|
if progress_bar is not None:
|
||||||
|
progress_bar.update(cur_input_tts_text_ids.shape[1])
|
||||||
|
progress_bar.set_description(f"Prefilled {total_prefilled_text_tokens} text tokens, generated {total_generated_speech_tokens} speech tokens, current step ({step} / {tts_lm_generation_config.max_length})")
|
||||||
|
|
||||||
|
model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
|
||||||
|
# Forward pass through the model
|
||||||
|
outputs = self.forward_lm(
|
||||||
|
**model_inputs, return_dict=True, output_attentions=False, output_hidden_states=False,
|
||||||
|
)
|
||||||
|
model_kwargs = _update_model_kwargs_for_generation(
|
||||||
|
outputs, model_kwargs, num_new_tokens=next_text_window_size,
|
||||||
|
)
|
||||||
|
|
||||||
|
tts_lm_model_inputs = self.prepare_inputs_for_generation(tts_lm_input_ids, **tts_lm_model_kwargs)
|
||||||
|
tts_lm_additional_inputs = {
|
||||||
|
"tts_text_masks": torch.ones_like(tts_lm_input_ids[:, -1:]),
|
||||||
|
"lm_last_hidden_state": outputs.last_hidden_state,
|
||||||
|
}
|
||||||
|
# Forward pass through the model
|
||||||
|
tts_lm_outputs = self.forward_tts_lm(
|
||||||
|
**tts_lm_model_inputs, **tts_lm_additional_inputs, return_dict=True, output_attentions=False, output_hidden_states=False,
|
||||||
|
)
|
||||||
|
tts_lm_model_kwargs = self._update_model_kwargs_for_generation(
|
||||||
|
tts_lm_outputs, tts_lm_model_kwargs, is_encoder_decoder=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
diffusion_indices = torch.LongTensor([0])
|
||||||
|
for cur_speech_index in range(TTS_SPEECH_WINDOW_SIZE):
|
||||||
|
positive_condition = tts_lm_outputs.last_hidden_state[diffusion_indices, -1, :]
|
||||||
|
negative_condition = tts_lm_negative_outputs.last_hidden_state[diffusion_indices, -1, :]
|
||||||
|
|
||||||
|
speech_latent = self.sample_speech_tokens(
|
||||||
|
positive_condition,
|
||||||
|
negative_condition,
|
||||||
|
cfg_scale=cfg_scale,
|
||||||
|
).unsqueeze(1)
|
||||||
|
|
||||||
|
# Decode acoustic latent to audio using acoustic streaming cache
|
||||||
|
scaled_latent = speech_latent / self.model.speech_scaling_factor.to(speech_latent.device) - self.model.speech_bias_factor.to(speech_latent.device)
|
||||||
|
audio_chunk = self.model.acoustic_tokenizer.decode(
|
||||||
|
scaled_latent.to(self.model.acoustic_tokenizer.device),
|
||||||
|
cache=acoustic_cache, # Use acoustic-specific cache
|
||||||
|
sample_indices=diffusion_indices.to(self.model.acoustic_tokenizer.device),
|
||||||
|
use_cache=True,
|
||||||
|
debug=False
|
||||||
|
)
|
||||||
|
|
||||||
|
# Store audio chunks for each sample
|
||||||
|
for i, sample_idx in enumerate(diffusion_indices):
|
||||||
|
idx = sample_idx.item()
|
||||||
|
# Only append audio chunk if the sample is not finished
|
||||||
|
if not finished_tags[idx]:
|
||||||
|
audio_chunks[idx].append(audio_chunk[i])
|
||||||
|
|
||||||
|
# Add streaming support here
|
||||||
|
if audio_streamer is not None:
|
||||||
|
# Stream the audio chunks immediately
|
||||||
|
audio_streamer.put(audio_chunk, diffusion_indices)
|
||||||
|
|
||||||
|
acoustic_embed = self.model.acoustic_connector(speech_latent)
|
||||||
|
tts_lm_input_ids = torch.cat([tts_lm_input_ids, torch.ones_like(tts_lm_input_ids[:, -1:])], dim=-1)
|
||||||
|
|
||||||
|
if tts_lm_input_ids.shape[1] > tts_lm_generation_config.max_length:
|
||||||
|
break
|
||||||
|
|
||||||
|
step += 1
|
||||||
|
total_generated_speech_tokens += 1
|
||||||
|
if progress_bar is not None:
|
||||||
|
progress_bar.update(1)
|
||||||
|
progress_bar.set_description(f"Prefilled {total_prefilled_text_tokens} text tokens, generated {total_generated_speech_tokens} speech tokens, current step ({step} / {tts_lm_generation_config.max_length})")
|
||||||
|
|
||||||
|
tts_lm_model_inputs = self.prepare_inputs_for_generation(tts_lm_input_ids, **tts_lm_model_kwargs)
|
||||||
|
tts_lm_additional_inputs = {
|
||||||
|
"tts_text_masks": torch.zeros_like(tts_lm_input_ids[:, -1:]),
|
||||||
|
"lm_last_hidden_state": acoustic_embed,
|
||||||
|
}
|
||||||
|
# Forward pass through the model
|
||||||
|
tts_lm_outputs = self.forward_tts_lm(
|
||||||
|
**tts_lm_model_inputs, **tts_lm_additional_inputs, return_dict=True, output_attentions=False, output_hidden_states=False,
|
||||||
|
)
|
||||||
|
if cur_speech_index == TTS_SPEECH_WINDOW_SIZE - 1 and next_text_window_size > 0:
|
||||||
|
tts_lm_model_kwargs = _update_model_kwargs_for_generation(
|
||||||
|
tts_lm_outputs, tts_lm_model_kwargs, num_new_tokens=next_text_window_size,
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
tts_lm_model_kwargs = self._update_model_kwargs_for_generation(
|
||||||
|
tts_lm_outputs, tts_lm_model_kwargs, is_encoder_decoder=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
tts_lm_negative_input_ids = torch.cat([tts_lm_negative_input_ids, torch.ones_like(tts_lm_input_ids[:, -1:])], dim=-1)
|
||||||
|
tts_lm_negative_model_inputs = self.prepare_inputs_for_generation(tts_lm_negative_input_ids, **tts_lm_negative_model_kwargs)
|
||||||
|
# Forward negative pass through the model
|
||||||
|
tts_lm_negative_additional_inputs = {
|
||||||
|
"tts_text_masks": torch.zeros_like(tts_lm_negative_input_ids[:, -1:]),
|
||||||
|
"lm_last_hidden_state": acoustic_embed,
|
||||||
|
}
|
||||||
|
tts_lm_negative_outputs = self.forward_tts_lm(
|
||||||
|
**tts_lm_negative_model_inputs, **tts_lm_negative_additional_inputs, return_dict=True, output_attentions=False, output_hidden_states=False,
|
||||||
|
)
|
||||||
|
tts_lm_negative_model_kwargs = self._update_model_kwargs_for_generation(
|
||||||
|
tts_lm_negative_outputs, tts_lm_negative_model_kwargs, is_encoder_decoder=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
tts_eos_logits = torch.sigmoid(self.tts_eos_classifier(tts_lm_outputs.last_hidden_state[diffusion_indices, -1, :]))
|
||||||
|
if tts_eos_logits[0].item() > 0.5:
|
||||||
|
# If EOS token is predicted, we can stop generation for this sample
|
||||||
|
finished_tags[diffusion_indices] = True
|
||||||
|
if audio_streamer is not None:
|
||||||
|
audio_streamer.end(diffusion_indices)
|
||||||
|
|
||||||
|
if tts_lm_input_ids.shape[1] > tts_lm_generation_config.max_length:
|
||||||
|
if verbose:
|
||||||
|
print(f"Reached maximum generation length {tts_lm_generation_config.max_length}, stopped it.")
|
||||||
|
reached_samples = torch.arange(batch_size, device=device)[~finished_tags]
|
||||||
|
if reached_samples.numel() > 0:
|
||||||
|
reach_max_step_sample[reached_samples] = True
|
||||||
|
break
|
||||||
|
|
||||||
|
if audio_streamer is not None:
|
||||||
|
audio_streamer.end()
|
||||||
|
|
||||||
|
# Concatenate audio chunks for each sample
|
||||||
|
final_audio_outputs = []
|
||||||
|
for sample_chunks in audio_chunks:
|
||||||
|
if sample_chunks:
|
||||||
|
# Concatenate all chunks along the time dimension (assumed to be the last dimension)
|
||||||
|
concatenated_audio = torch.cat(sample_chunks, dim=-1)
|
||||||
|
final_audio_outputs.append(concatenated_audio)
|
||||||
|
else:
|
||||||
|
# If no audio was generated for this sample, append None
|
||||||
|
final_audio_outputs.append(None)
|
||||||
|
|
||||||
|
if reach_max_step_sample is not None and reach_max_step_sample.any():
|
||||||
|
print(f"Reached maximum generation length {tts_lm_generation_config.max_length}, stopped it.")
|
||||||
|
|
||||||
|
return VibeVoiceGenerationOutput(
|
||||||
|
sequences=tts_lm_input_ids,
|
||||||
|
speech_outputs=final_audio_outputs if return_speech else None,
|
||||||
|
reach_max_step_sample=reach_max_step_sample,
|
||||||
|
)
|
||||||
|
|
||||||
|
@torch.no_grad()
|
||||||
|
def sample_speech_tokens(self, condition, neg_condition, cfg_scale=3.0):
|
||||||
|
"""
|
||||||
|
Sample speech tokens using diffusion with classifier-free guidance.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
condition: Positive conditioning from TTS LM hidden states
|
||||||
|
neg_condition: Negative conditioning for CFG
|
||||||
|
cfg_scale: Classifier-free guidance scale (higher = more adherence to text)
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Generated speech latents
|
||||||
|
"""
|
||||||
|
self.model.noise_scheduler.set_timesteps(self.ddpm_inference_steps)
|
||||||
|
condition = torch.cat([condition, neg_condition], dim=0).to(self.model.prediction_head.device)
|
||||||
|
speech = torch.randn(condition.shape[0], self.config.acoustic_vae_dim).to(condition)
|
||||||
|
for t in self.model.noise_scheduler.timesteps:
|
||||||
|
half = speech[: len(speech) // 2]
|
||||||
|
combined = torch.cat([half, half], dim=0)
|
||||||
|
eps = self.model.prediction_head(combined, t.repeat(combined.shape[0]).to(combined), condition=condition)
|
||||||
|
cond_eps, uncond_eps = torch.split(eps, len(eps) // 2, dim=0)
|
||||||
|
half_eps = uncond_eps + cfg_scale * (cond_eps - uncond_eps)
|
||||||
|
eps = torch.cat([half_eps, half_eps], dim=0)
|
||||||
|
speech = self.model.noise_scheduler.step(eps, t, speech).prev_sample
|
||||||
|
return speech[: len(speech) // 2]
|
||||||
|
|
||||||
|
|
||||||
|
AutoModelForCausalLM.register(VibeVoiceStreamingConfig, VibeVoiceStreamingForConditionalGenerationInference)
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"VibeVoiceStreamingForConditionalGenerationInference",
|
||||||
|
"VibeVoiceGenerationOutput",
|
||||||
|
"VibeVoiceCausalLMOutputWithPast",
|
||||||
|
"TTS_TEXT_WINDOW_SIZE",
|
||||||
|
"TTS_SPEECH_WINDOW_SIZE",
|
||||||
|
]
|
||||||
@@ -0,0 +1,18 @@
|
|||||||
|
"""
|
||||||
|
VibeVoice Processors
|
||||||
|
|
||||||
|
This module provides processors for preparing inputs for VibeVoice models:
|
||||||
|
- VibeVoiceProcessor: For multi-speaker models (1.5B, 7B)
|
||||||
|
- VibeVoiceStreamingProcessor: For streaming model (0.5B)
|
||||||
|
"""
|
||||||
|
|
||||||
|
from .vibevoice_processor import VibeVoiceProcessor
|
||||||
|
from .vibevoice_streaming_processor import VibeVoiceStreamingProcessor
|
||||||
|
from .vibevoice_tokenizer_processor import VibeVoiceTokenizerProcessor, AudioNormalizer
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"VibeVoiceProcessor",
|
||||||
|
"VibeVoiceStreamingProcessor",
|
||||||
|
"VibeVoiceTokenizerProcessor",
|
||||||
|
"AudioNormalizer",
|
||||||
|
]
|
||||||
|
|||||||
@@ -0,0 +1,426 @@
|
|||||||
|
"""
|
||||||
|
VibeVoice Streaming Processor
|
||||||
|
|
||||||
|
This processor handles input preparation for the streaming 0.5B model,
|
||||||
|
including text tokenization and cached voice prompt handling.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import math
|
||||||
|
import warnings
|
||||||
|
from typing import List, Optional, Union, Dict, Any, Tuple
|
||||||
|
import os
|
||||||
|
import re
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import torch
|
||||||
|
|
||||||
|
from transformers.tokenization_utils_base import BatchEncoding, PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
|
||||||
|
from transformers.utils import TensorType, logging
|
||||||
|
from .vibevoice_tokenizer_processor import AudioNormalizer
|
||||||
|
|
||||||
|
logger = logging.get_logger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class VibeVoiceStreamingProcessor:
|
||||||
|
r"""
|
||||||
|
Constructs a VibeVoice Streaming processor which wraps a VibeVoice tokenizer and audio processor into a single processor.
|
||||||
|
|
||||||
|
The streaming processor is designed for the 0.5B real-time model that uses pre-computed voice embeddings
|
||||||
|
(.pt files) instead of live audio conditioning.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
tokenizer (`VibeVoiceTextTokenizer` or `VibeVoiceTextTokenizerFast`):
|
||||||
|
The tokenizer for text processing.
|
||||||
|
audio_processor (`VibeVoiceTokenizerProcessor`):
|
||||||
|
The audio processor for speech processing.
|
||||||
|
speech_tok_compress_ratio (`int`, *optional*, defaults to 3200):
|
||||||
|
The compression ratio for speech tokenization.
|
||||||
|
db_normalize (`bool`, *optional*, defaults to True):
|
||||||
|
Whether to apply decibel normalization to audio inputs.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, tokenizer=None, audio_processor=None, speech_tok_compress_ratio=3200, db_normalize=True, **kwargs):
|
||||||
|
self.tokenizer = tokenizer
|
||||||
|
self.audio_processor = audio_processor
|
||||||
|
self.speech_tok_compress_ratio = speech_tok_compress_ratio
|
||||||
|
self.db_normalize = db_normalize
|
||||||
|
self.audio_normalizer = AudioNormalizer() if db_normalize else None
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
|
||||||
|
"""
|
||||||
|
Instantiate a VibeVoiceStreamingProcessor from a pretrained VibeVoice Streaming processor.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
pretrained_model_name_or_path (`str` or `os.PathLike`):
|
||||||
|
This can be either:
|
||||||
|
- a string, the *model id* of a pretrained model
|
||||||
|
- a path to a *directory* containing processor config
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
[`VibeVoiceStreamingProcessor`]: The processor object instantiated from pretrained model.
|
||||||
|
"""
|
||||||
|
import os
|
||||||
|
import json
|
||||||
|
from transformers.utils import cached_file
|
||||||
|
from .vibevoice_tokenizer_processor import VibeVoiceTokenizerProcessor
|
||||||
|
from vibevoice.modular.modular_vibevoice_text_tokenizer import (
|
||||||
|
VibeVoiceTextTokenizer,
|
||||||
|
VibeVoiceTextTokenizerFast
|
||||||
|
)
|
||||||
|
|
||||||
|
# Try to load from local path first, then from HF hub
|
||||||
|
config_path = os.path.join(pretrained_model_name_or_path, "preprocessor_config.json")
|
||||||
|
config = None
|
||||||
|
|
||||||
|
if os.path.exists(config_path):
|
||||||
|
# Local path exists
|
||||||
|
with open(config_path, 'r') as f:
|
||||||
|
config = json.load(f)
|
||||||
|
else:
|
||||||
|
# Try to load from HF hub
|
||||||
|
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 from {pretrained_model_name_or_path}: {e}")
|
||||||
|
logger.warning("Using default configuration")
|
||||||
|
config = {
|
||||||
|
"speech_tok_compress_ratio": 3200,
|
||||||
|
"db_normalize": True,
|
||||||
|
}
|
||||||
|
|
||||||
|
# Extract main processor parameters
|
||||||
|
speech_tok_compress_ratio = config.get("speech_tok_compress_ratio", 3200)
|
||||||
|
db_normalize = config.get("db_normalize", True)
|
||||||
|
|
||||||
|
# Load tokenizer - try from model path first, then fallback to Qwen
|
||||||
|
language_model_pretrained_name = config.get("language_model_pretrained_name", None) or kwargs.pop("language_model_pretrained_name", "Qwen/Qwen2.5-0.5B")
|
||||||
|
logger.info(f"Loading tokenizer from {language_model_pretrained_name}")
|
||||||
|
if 'qwen' in language_model_pretrained_name.lower():
|
||||||
|
tokenizer = VibeVoiceTextTokenizerFast.from_pretrained(
|
||||||
|
language_model_pretrained_name,
|
||||||
|
**kwargs
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
raise ValueError(f"Unsupported tokenizer type for {language_model_pretrained_name}. Supported types: Qwen.")
|
||||||
|
|
||||||
|
# Load audio processor
|
||||||
|
if "audio_processor" in config:
|
||||||
|
# Create audio processor from config
|
||||||
|
audio_config = config["audio_processor"]
|
||||||
|
audio_processor = VibeVoiceTokenizerProcessor(
|
||||||
|
sampling_rate=audio_config.get("sampling_rate", 24000),
|
||||||
|
normalize_audio=audio_config.get("normalize_audio", True),
|
||||||
|
target_dB_FS=audio_config.get("target_dB_FS", -25),
|
||||||
|
eps=audio_config.get("eps", 1e-6),
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
# Create default audio processor
|
||||||
|
audio_processor = VibeVoiceTokenizerProcessor()
|
||||||
|
|
||||||
|
# Create and return the processor
|
||||||
|
return cls(
|
||||||
|
tokenizer=tokenizer,
|
||||||
|
audio_processor=audio_processor,
|
||||||
|
speech_tok_compress_ratio=speech_tok_compress_ratio,
|
||||||
|
db_normalize=db_normalize,
|
||||||
|
)
|
||||||
|
|
||||||
|
def save_pretrained(self, save_directory: Union[str, os.PathLike], **kwargs):
|
||||||
|
"""
|
||||||
|
Save a processor to a directory, so that it can be re-loaded using the
|
||||||
|
[`~VibeVoiceStreamingProcessor.from_pretrained`] class method.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
save_directory (`str` or `os.PathLike`):
|
||||||
|
Directory where the processor will be saved.
|
||||||
|
"""
|
||||||
|
import os
|
||||||
|
import json
|
||||||
|
|
||||||
|
os.makedirs(save_directory, exist_ok=True)
|
||||||
|
|
||||||
|
# Save processor configuration
|
||||||
|
processor_config = {
|
||||||
|
"processor_class": "VibeVoiceStreamingProcessor",
|
||||||
|
"speech_tok_compress_ratio": self.speech_tok_compress_ratio,
|
||||||
|
"db_normalize": self.db_normalize,
|
||||||
|
"audio_processor": {
|
||||||
|
"feature_extractor_type": "VibeVoiceTokenizerProcessor",
|
||||||
|
"sampling_rate": getattr(self.audio_processor, 'sampling_rate', 24000),
|
||||||
|
"normalize_audio": getattr(self.audio_processor, 'normalize_audio', True),
|
||||||
|
"target_dB_FS": getattr(self.audio_processor, 'target_dB_FS', -25),
|
||||||
|
"eps": getattr(self.audio_processor, '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) -> BatchEncoding:
|
||||||
|
"""
|
||||||
|
Note:
|
||||||
|
This method is intentionally not implemented in the streaming processor.
|
||||||
|
Use `process_input_with_cached_prompt` for streaming use cases.
|
||||||
|
"""
|
||||||
|
raise NotImplementedError(
|
||||||
|
"VibeVoiceStreamingProcessor.__call__ is not implemented. "
|
||||||
|
"Use process_input_with_cached_prompt for streaming inputs."
|
||||||
|
)
|
||||||
|
|
||||||
|
def process_input_with_cached_prompt(
|
||||||
|
self,
|
||||||
|
text: Optional[str] = None,
|
||||||
|
cached_prompt: Optional[Dict[str, Any]] = None,
|
||||||
|
padding: Union[bool, str, PaddingStrategy] = True,
|
||||||
|
truncation: Union[bool, str, TruncationStrategy] = False,
|
||||||
|
max_length: Optional[int] = None,
|
||||||
|
return_tensors: Optional[Union[str, TensorType]] = None,
|
||||||
|
return_attention_mask: bool = True,
|
||||||
|
**kwargs,
|
||||||
|
) -> BatchEncoding:
|
||||||
|
"""
|
||||||
|
Main method to process one text script based on cached prompt.
|
||||||
|
|
||||||
|
The streaming model uses pre-computed voice embeddings (cached prompts) loaded from .pt files
|
||||||
|
instead of processing audio on-the-fly.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
text (`str`):
|
||||||
|
The input text to process.
|
||||||
|
cached_prompt (`Dict[str, Any]`, *optional*):
|
||||||
|
The cached prompt to use for processing. It contains the kv cache of the voice prompt.
|
||||||
|
Load this from a .pt file using torch.load().
|
||||||
|
padding (`bool`, `str` or `PaddingStrategy`, defaults to `True`):
|
||||||
|
Whether to pad sequences to the same length
|
||||||
|
truncation (`bool`, `str` or `TruncationStrategy`, defaults to `False`):
|
||||||
|
Whether to truncate sequences
|
||||||
|
max_length (`int`, *optional*):
|
||||||
|
Maximum length of the returned sequences
|
||||||
|
return_tensors (`str` or `TensorType`, *optional*):
|
||||||
|
If set, will return tensors of a particular framework
|
||||||
|
return_attention_mask (`bool`, defaults to `True`):
|
||||||
|
Whether to return the attention mask
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
`BatchEncoding`: A BatchEncoding with the following fields:
|
||||||
|
- **input_ids** -- List of token id sequences or tensor
|
||||||
|
- **attention_mask** -- List of attention masks or tensor
|
||||||
|
- **tts_lm_input_ids** -- List of token id sequences or tensor used for TTS LM
|
||||||
|
- **tts_lm_attention_mask** -- List of attention masks or tensor used for TTS LM
|
||||||
|
- **tts_text_ids** -- List of token id sequences or tensor for TTS text input
|
||||||
|
- **speech_tensors** -- Padded speech inputs (if voice_samples provided)
|
||||||
|
- **speech_masks** -- Speech masks (if voice_samples provided)
|
||||||
|
- **speech_input_mask** -- Boolean masks indicating speech token positions
|
||||||
|
"""
|
||||||
|
# Only support single example
|
||||||
|
texts = [text]
|
||||||
|
cached_prompts = [cached_prompt]
|
||||||
|
is_batched = False
|
||||||
|
|
||||||
|
# Process each input
|
||||||
|
all_encodings = []
|
||||||
|
for text_input, cached_prompt_input in zip(texts, cached_prompts):
|
||||||
|
script_tokens = self.tokenizer.encode(text_input.strip() + "\n", add_special_tokens=False)
|
||||||
|
input_id_length = cached_prompt_input['lm']['last_hidden_state'].size(1)
|
||||||
|
tts_lm_input_id_length = cached_prompt_input['tts_lm']['last_hidden_state'].size(1)
|
||||||
|
|
||||||
|
# pseudo input ids and masks
|
||||||
|
input_ids = [self.tokenizer.pad_id] * input_id_length
|
||||||
|
tts_lm_input_ids = [self.tokenizer.pad_id] * tts_lm_input_id_length
|
||||||
|
speech_input_mask = [False] * tts_lm_input_id_length
|
||||||
|
|
||||||
|
encoding = {
|
||||||
|
"input_ids": input_ids,
|
||||||
|
"tts_lm_input_ids": tts_lm_input_ids,
|
||||||
|
"tts_text_ids": script_tokens,
|
||||||
|
"speech_inputs": None,
|
||||||
|
"speech_input_mask": speech_input_mask,
|
||||||
|
}
|
||||||
|
all_encodings.append(encoding)
|
||||||
|
|
||||||
|
# Combine batch
|
||||||
|
batch_encoding = self._batch_encode(
|
||||||
|
all_encodings,
|
||||||
|
padding=padding,
|
||||||
|
truncation=truncation,
|
||||||
|
max_length=max_length,
|
||||||
|
return_tensors=return_tensors,
|
||||||
|
return_attention_mask=return_attention_mask,
|
||||||
|
)
|
||||||
|
|
||||||
|
return batch_encoding
|
||||||
|
|
||||||
|
def _batch_encode(
|
||||||
|
self,
|
||||||
|
encodings: List[Dict[str, Any]],
|
||||||
|
padding: Union[bool, str, PaddingStrategy] = True,
|
||||||
|
truncation: Union[bool, str, TruncationStrategy] = False,
|
||||||
|
max_length: Optional[int] = None,
|
||||||
|
return_tensors: Optional[Union[str, TensorType]] = None,
|
||||||
|
return_attention_mask: bool = True,
|
||||||
|
) -> BatchEncoding:
|
||||||
|
"""Combine multiple encodings into a batch with padding."""
|
||||||
|
# Extract input_ids and create attention_mask
|
||||||
|
input_ids_list = [enc["input_ids"] for enc in encodings]
|
||||||
|
tts_lm_input_ids_list = [enc["tts_lm_input_ids"] for enc in encodings]
|
||||||
|
tts_text_ids_list = [enc["tts_text_ids"] for enc in encodings]
|
||||||
|
speech_input_masks_list = [enc["speech_input_mask"] for enc in encodings]
|
||||||
|
|
||||||
|
attention_masks = [[1] * len(ids) for ids in input_ids_list] if return_attention_mask else None
|
||||||
|
tts_lm_attention_masks = [[1] * len(ids) for ids in tts_lm_input_ids_list] if return_attention_mask else None
|
||||||
|
|
||||||
|
# Process speech inputs
|
||||||
|
all_speech_inputs = []
|
||||||
|
has_speech = False
|
||||||
|
for enc in encodings:
|
||||||
|
if enc["speech_inputs"] is not None:
|
||||||
|
all_speech_inputs.extend(enc["speech_inputs"])
|
||||||
|
has_speech = True
|
||||||
|
|
||||||
|
# Prepare batch encoding
|
||||||
|
batch_encoding = BatchEncoding()
|
||||||
|
|
||||||
|
# Handle tensor conversion
|
||||||
|
if return_tensors is not None:
|
||||||
|
batch_encoding["input_ids"] = torch.tensor(input_ids_list, dtype=torch.long)
|
||||||
|
batch_encoding["tts_lm_input_ids"] = torch.tensor(tts_lm_input_ids_list, dtype=torch.long)
|
||||||
|
batch_encoding["tts_text_ids"] = torch.tensor(tts_text_ids_list, dtype=torch.long)
|
||||||
|
|
||||||
|
if return_attention_mask and attention_masks is not None:
|
||||||
|
batch_encoding["attention_mask"] = torch.tensor(attention_masks, dtype=torch.long)
|
||||||
|
batch_encoding["tts_lm_attention_mask"] = torch.tensor(tts_lm_attention_masks, dtype=torch.long)
|
||||||
|
|
||||||
|
batch_encoding["speech_input_mask"] = torch.tensor(speech_input_masks_list, dtype=torch.bool)
|
||||||
|
else:
|
||||||
|
batch_encoding["input_ids"] = input_ids_list
|
||||||
|
batch_encoding["tts_lm_input_ids"] = tts_lm_input_ids_list
|
||||||
|
batch_encoding["tts_text_ids"] = tts_text_ids_list
|
||||||
|
if return_attention_mask and attention_masks is not None:
|
||||||
|
batch_encoding["attention_mask"] = attention_masks
|
||||||
|
batch_encoding["tts_lm_attention_mask"] = tts_lm_attention_masks
|
||||||
|
batch_encoding["speech_input_mask"] = speech_input_masks_list
|
||||||
|
|
||||||
|
# Process speech tensors if present
|
||||||
|
if has_speech:
|
||||||
|
speech_dict = self.prepare_speech_inputs(
|
||||||
|
all_speech_inputs,
|
||||||
|
return_tensors=return_tensors,
|
||||||
|
)
|
||||||
|
batch_encoding["speech_tensors"] = speech_dict["padded_speeches"]
|
||||||
|
batch_encoding["speech_masks"] = speech_dict["speech_masks"]
|
||||||
|
else:
|
||||||
|
batch_encoding["speech_tensors"] = None
|
||||||
|
batch_encoding["speech_masks"] = None
|
||||||
|
|
||||||
|
return batch_encoding
|
||||||
|
|
||||||
|
def prepare_speech_inputs(
|
||||||
|
self,
|
||||||
|
speech_inputs: List[np.ndarray],
|
||||||
|
return_tensors: Optional[Union[str, TensorType]] = None,
|
||||||
|
device: Optional[Union[str, torch.device]] = None,
|
||||||
|
dtype: Optional[torch.dtype] = None,
|
||||||
|
) -> Dict[str, Any]:
|
||||||
|
"""
|
||||||
|
Prepare speech inputs for model consumption.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
speech_inputs: List of speech arrays
|
||||||
|
return_tensors: Output tensor type
|
||||||
|
device: Device to place tensors on
|
||||||
|
dtype: Data type for tensors
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Dictionary with padded_speeches and speech_masks
|
||||||
|
"""
|
||||||
|
if not speech_inputs:
|
||||||
|
return {"padded_speeches": None, "speech_masks": None}
|
||||||
|
|
||||||
|
# Calculate sequence lengths
|
||||||
|
vae_tok_seqlens = [math.ceil(s.shape[0] / self.speech_tok_compress_ratio) for s in speech_inputs]
|
||||||
|
max_speech_length = max(s.shape[0] for s in speech_inputs)
|
||||||
|
|
||||||
|
# Pad speeches
|
||||||
|
if speech_inputs[0].ndim == 1:
|
||||||
|
padded_speeches = np.full((len(speech_inputs), max_speech_length), fill_value=0, dtype=np.float32)
|
||||||
|
else:
|
||||||
|
padded_speeches = np.full((len(speech_inputs), max_speech_length, speech_inputs[0].shape[-1]), fill_value=0, dtype=np.float32)
|
||||||
|
speech_masks = np.zeros((len(speech_inputs), max(vae_tok_seqlens)), dtype=np.bool_)
|
||||||
|
|
||||||
|
for i, (speech, vae_tok_length) in enumerate(zip(speech_inputs, vae_tok_seqlens)):
|
||||||
|
padded_speeches[i, :len(speech)] = speech
|
||||||
|
speech_masks[i, :vae_tok_length] = True
|
||||||
|
|
||||||
|
result = {
|
||||||
|
"padded_speeches": padded_speeches,
|
||||||
|
"speech_masks": speech_masks,
|
||||||
|
}
|
||||||
|
|
||||||
|
# Convert to tensors if requested
|
||||||
|
if return_tensors == "pt":
|
||||||
|
result["padded_speeches"] = torch.tensor(padded_speeches, device=device, dtype=dtype or torch.float32)
|
||||||
|
result["speech_masks"] = torch.tensor(speech_masks, device=device, dtype=torch.bool)
|
||||||
|
|
||||||
|
return result
|
||||||
|
|
||||||
|
def batch_decode(self, *args, **kwargs):
|
||||||
|
"""
|
||||||
|
This method forwards all its arguments to VibeVoiceTextTokenizer's [`~PreTrainedTokenizer.batch_decode`].
|
||||||
|
Please refer to the docstring of this method for more information.
|
||||||
|
"""
|
||||||
|
return self.tokenizer.batch_decode(*args, **kwargs)
|
||||||
|
|
||||||
|
def decode(self, *args, **kwargs):
|
||||||
|
"""
|
||||||
|
This method forwards all its arguments to VibeVoiceTextTokenizer's [`~PreTrainedTokenizer.decode`].
|
||||||
|
Please refer to the docstring of this method for more information.
|
||||||
|
"""
|
||||||
|
return self.tokenizer.decode(*args, **kwargs)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def model_input_names(self):
|
||||||
|
"""
|
||||||
|
Return the list of inputs accepted by the model.
|
||||||
|
"""
|
||||||
|
tokenizer_input_names = self.tokenizer.model_input_names
|
||||||
|
audio_processor_input_names = self.audio_processor.model_input_names
|
||||||
|
return list(dict.fromkeys(tokenizer_input_names + audio_processor_input_names + ["speech_inputs", "speech_input_mask"]))
|
||||||
|
|
||||||
|
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_",
|
||||||
|
) -> str:
|
||||||
|
"""
|
||||||
|
Save audio data to a file.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
audio (Union[torch.Tensor, np.ndarray, List[Union[torch.Tensor, np.ndarray]]]):
|
||||||
|
The audio data to save. Can be a single tensor/array or a list of them.
|
||||||
|
output_path (str, optional): Path to save the audio file. Defaults to "output.wav".
|
||||||
|
sampling_rate (int, optional): Sampling rate for the audio. If None, uses the processor's default.
|
||||||
|
normalize (bool, optional): Whether to normalize the audio before saving. Defaults to False.
|
||||||
|
batch_prefix (str, optional): Prefix for batch audio files. Defaults to "audio_".
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
str: The path to the saved audio file.
|
||||||
|
"""
|
||||||
|
return self.audio_processor.save_audio(audio, output_path=output_path, sampling_rate=sampling_rate, normalize=normalize, batch_prefix=batch_prefix)
|
||||||
|
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"VibeVoiceStreamingProcessor",
|
||||||
|
]
|
||||||
Reference in New Issue
Block a user