Add VibeVoice streaming model (#35)

2025-12-04 22:00:30 -08:00
parent 187cf62036
commit 864a4a7ed2
18 changed files with 2041 additions and 5 deletions
@@ -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!