1178 lines
46 KiB
Python
1178 lines
46 KiB
Python
"""
|
||
VibeVoice Gradio Demo - High-Quality Dialogue Generation Interface with Streaming Support
|
||
"""
|
||
|
||
import argparse
|
||
import json
|
||
import os
|
||
import sys
|
||
import tempfile
|
||
import time
|
||
from pathlib import Path
|
||
from typing import List, Dict, Any, Iterator
|
||
from datetime import datetime
|
||
import threading
|
||
import numpy as np
|
||
import gradio as gr
|
||
import librosa
|
||
import soundfile as sf
|
||
import torch
|
||
import os
|
||
import traceback
|
||
|
||
from vibevoice.modular.configuration_vibevoice import VibeVoiceConfig
|
||
from vibevoice.modular.modeling_vibevoice_inference import VibeVoiceForConditionalGenerationInference
|
||
from vibevoice.processor.vibevoice_processor import VibeVoiceProcessor
|
||
from vibevoice.modular.streamer import AudioStreamer
|
||
from transformers.utils import logging
|
||
from transformers import set_seed
|
||
|
||
logging.set_verbosity_info()
|
||
logger = logging.get_logger(__name__)
|
||
|
||
|
||
class VibeVoiceDemo:
|
||
def __init__(self, model_path: str, device: str = "cuda", inference_steps: int = 5):
|
||
"""Initialize the VibeVoice demo with model loading."""
|
||
self.model_path = model_path
|
||
self.device = device
|
||
self.inference_steps = inference_steps
|
||
self.is_generating = False # Track generation state
|
||
self.stop_generation = False # Flag to stop generation
|
||
self.current_streamer = None # Track current audio streamer
|
||
self.load_model()
|
||
self.setup_voice_presets()
|
||
self.load_example_scripts() # Load example scripts
|
||
|
||
def load_model(self):
|
||
"""Load the VibeVoice model and processor."""
|
||
print(f"Loading processor & model from {self.model_path}")
|
||
|
||
# Load processor
|
||
self.processor = VibeVoiceProcessor.from_pretrained(
|
||
self.model_path,
|
||
)
|
||
|
||
# Load model
|
||
self.model = VibeVoiceForConditionalGenerationInference.from_pretrained(
|
||
self.model_path,
|
||
torch_dtype=torch.bfloat16,
|
||
device_map='cuda',
|
||
attn_implementation="flash_attention_2",
|
||
)
|
||
self.model.eval()
|
||
|
||
# Use SDE solver by default
|
||
self.model.model.noise_scheduler = self.model.model.noise_scheduler.from_config(
|
||
self.model.model.noise_scheduler.config,
|
||
algorithm_type='sde-dpmsolver++',
|
||
beta_schedule='squaredcos_cap_v2'
|
||
)
|
||
self.model.set_ddpm_inference_steps(num_steps=self.inference_steps)
|
||
|
||
if hasattr(self.model.model, 'language_model'):
|
||
print(f"Language model attention: {self.model.model.language_model.config._attn_implementation}")
|
||
|
||
def setup_voice_presets(self):
|
||
"""Setup voice presets by scanning the voices directory."""
|
||
voices_dir = os.path.join(os.path.dirname(__file__), "voices")
|
||
|
||
# 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 WAV files in the voices directory
|
||
self.voice_presets = {}
|
||
|
||
# Get all .wav files in the voices directory
|
||
wav_files = [f for f in os.listdir(voices_dir)
|
||
if f.lower().endswith(('.wav', '.mp3', '.flac', '.ogg', '.m4a', '.aac')) and os.path.isfile(os.path.join(voices_dir, f))]
|
||
|
||
# Create dictionary with filename (without extension) as key
|
||
for wav_file in wav_files:
|
||
# Remove .wav extension to get the name
|
||
name = os.path.splitext(wav_file)[0]
|
||
# Create full path
|
||
full_path = os.path.join(voices_dir, wav_file)
|
||
self.voice_presets[name] = full_path
|
||
|
||
# Sort the voice presets alphabetically by name for better UI
|
||
self.voice_presets = dict(sorted(self.voice_presets.items()))
|
||
|
||
# Filter out voices that don't exist (this is now redundant but kept for safety)
|
||
self.available_voices = {
|
||
name: path for name, path in self.voice_presets.items()
|
||
if os.path.exists(path)
|
||
}
|
||
|
||
if not self.available_voices:
|
||
raise gr.Error("No voice presets found. Please add .wav files to the demo/voices directory.")
|
||
|
||
print(f"Found {len(self.available_voices)} voice files in {voices_dir}")
|
||
print(f"Available voices: {', '.join(self.available_voices.keys())}")
|
||
|
||
def read_audio(self, audio_path: str, target_sr: int = 24000) -> np.ndarray:
|
||
"""Read and preprocess audio file."""
|
||
try:
|
||
wav, sr = sf.read(audio_path)
|
||
if len(wav.shape) > 1:
|
||
wav = np.mean(wav, axis=1)
|
||
if sr != target_sr:
|
||
wav = librosa.resample(wav, orig_sr=sr, target_sr=target_sr)
|
||
return wav
|
||
except Exception as e:
|
||
print(f"Error reading audio {audio_path}: {e}")
|
||
return np.array([])
|
||
|
||
def generate_podcast_streaming(self,
|
||
num_speakers: int,
|
||
script: str,
|
||
speaker_1: str = None,
|
||
speaker_2: str = None,
|
||
speaker_3: str = None,
|
||
speaker_4: str = None,
|
||
cfg_scale: float = 1.3) -> Iterator[tuple]:
|
||
try:
|
||
|
||
# Reset stop flag and set generating state
|
||
self.stop_generation = False
|
||
self.is_generating = True
|
||
|
||
# Validate inputs
|
||
if not script.strip():
|
||
self.is_generating = False
|
||
raise gr.Error("Error: Please provide a script.")
|
||
|
||
# Defend against common mistake
|
||
script = script.replace("’", "'")
|
||
|
||
if num_speakers < 1 or num_speakers > 4:
|
||
self.is_generating = False
|
||
raise gr.Error("Error: Number of speakers must be between 1 and 4.")
|
||
|
||
# Collect selected speakers
|
||
selected_speakers = [speaker_1, speaker_2, speaker_3, speaker_4][:num_speakers]
|
||
|
||
# Validate speaker selections
|
||
for i, speaker in enumerate(selected_speakers):
|
||
if not speaker or speaker not in self.available_voices:
|
||
self.is_generating = False
|
||
raise gr.Error(f"Error: Please select a valid speaker for Speaker {i+1}.")
|
||
|
||
# Build initial log
|
||
log = f"🎙️ Generating podcast with {num_speakers} speakers\n"
|
||
log += f"📊 Parameters: CFG Scale={cfg_scale}, Inference Steps={self.inference_steps}\n"
|
||
log += f"🎭 Speakers: {', '.join(selected_speakers)}\n"
|
||
|
||
# Check for stop signal
|
||
if self.stop_generation:
|
||
self.is_generating = False
|
||
yield None, "🛑 Generation stopped by user", gr.update(visible=False)
|
||
return
|
||
|
||
# Load voice samples
|
||
voice_samples = []
|
||
for speaker_name in selected_speakers:
|
||
audio_path = self.available_voices[speaker_name]
|
||
audio_data = self.read_audio(audio_path)
|
||
if len(audio_data) == 0:
|
||
self.is_generating = False
|
||
raise gr.Error(f"Error: Failed to load audio for {speaker_name}")
|
||
voice_samples.append(audio_data)
|
||
|
||
# log += f"✅ Loaded {len(voice_samples)} voice samples\n"
|
||
|
||
# Check for stop signal
|
||
if self.stop_generation:
|
||
self.is_generating = False
|
||
yield None, "🛑 Generation stopped by user", gr.update(visible=False)
|
||
return
|
||
|
||
# Parse script to assign speaker ID's
|
||
lines = script.strip().split('\n')
|
||
formatted_script_lines = []
|
||
|
||
for line in lines:
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
|
||
# Check if line already has speaker format
|
||
if line.startswith('Speaker ') and ':' in line:
|
||
formatted_script_lines.append(line)
|
||
else:
|
||
# Auto-assign to speakers in rotation
|
||
speaker_id = len(formatted_script_lines) % num_speakers
|
||
formatted_script_lines.append(f"Speaker {speaker_id}: {line}")
|
||
|
||
formatted_script = '\n'.join(formatted_script_lines)
|
||
log += f"📝 Formatted script with {len(formatted_script_lines)} turns\n\n"
|
||
log += "🔄 Processing with VibeVoice (streaming mode)...\n"
|
||
|
||
# Check for stop signal before processing
|
||
if self.stop_generation:
|
||
self.is_generating = False
|
||
yield None, "🛑 Generation stopped by user", gr.update(visible=False)
|
||
return
|
||
|
||
start_time = time.time()
|
||
|
||
inputs = self.processor(
|
||
text=[formatted_script],
|
||
voice_samples=[voice_samples],
|
||
padding=True,
|
||
return_tensors="pt",
|
||
return_attention_mask=True,
|
||
)
|
||
|
||
# Create audio streamer
|
||
audio_streamer = AudioStreamer(
|
||
batch_size=1,
|
||
stop_signal=None,
|
||
timeout=None
|
||
)
|
||
|
||
# Store current streamer for potential stopping
|
||
self.current_streamer = audio_streamer
|
||
|
||
# Start generation in a separate thread
|
||
generation_thread = threading.Thread(
|
||
target=self._generate_with_streamer,
|
||
args=(inputs, cfg_scale, audio_streamer)
|
||
)
|
||
generation_thread.start()
|
||
|
||
# Wait for generation to actually start producing audio
|
||
time.sleep(1) # Reduced from 3 to 1 second
|
||
|
||
# Check for stop signal after thread start
|
||
if self.stop_generation:
|
||
audio_streamer.end()
|
||
generation_thread.join(timeout=5.0) # Wait up to 5 seconds for thread to finish
|
||
self.is_generating = False
|
||
yield None, "🛑 Generation stopped by user", gr.update(visible=False)
|
||
return
|
||
|
||
# Collect audio chunks as they arrive
|
||
sample_rate = 24000
|
||
all_audio_chunks = [] # For final statistics
|
||
pending_chunks = [] # Buffer for accumulating small chunks
|
||
chunk_count = 0
|
||
last_yield_time = time.time()
|
||
min_yield_interval = 15 # Yield every 15 seconds
|
||
min_chunk_size = sample_rate * 30 # At least 2 seconds of audio
|
||
|
||
# Get the stream for the first (and only) sample
|
||
audio_stream = audio_streamer.get_stream(0)
|
||
|
||
has_yielded_audio = False
|
||
has_received_chunks = False # Track if we received any chunks at all
|
||
|
||
for audio_chunk in audio_stream:
|
||
# Check for stop signal in the streaming loop
|
||
if self.stop_generation:
|
||
audio_streamer.end()
|
||
break
|
||
|
||
chunk_count += 1
|
||
has_received_chunks = True # Mark that we received at least one chunk
|
||
|
||
# Convert tensor to numpy
|
||
if torch.is_tensor(audio_chunk):
|
||
# Convert bfloat16 to float32 first, then to numpy
|
||
if audio_chunk.dtype == torch.bfloat16:
|
||
audio_chunk = audio_chunk.float()
|
||
audio_np = audio_chunk.cpu().numpy().astype(np.float32)
|
||
else:
|
||
audio_np = np.array(audio_chunk, dtype=np.float32)
|
||
|
||
# Ensure audio is 1D and properly normalized
|
||
if len(audio_np.shape) > 1:
|
||
audio_np = audio_np.squeeze()
|
||
|
||
# Convert to 16-bit for Gradio
|
||
audio_16bit = convert_to_16_bit_wav(audio_np)
|
||
|
||
# Store for final statistics
|
||
all_audio_chunks.append(audio_16bit)
|
||
|
||
# Add to pending chunks buffer
|
||
pending_chunks.append(audio_16bit)
|
||
|
||
# Calculate pending audio size
|
||
pending_audio_size = sum(len(chunk) for chunk in pending_chunks)
|
||
current_time = time.time()
|
||
time_since_last_yield = current_time - last_yield_time
|
||
|
||
# Decide whether to yield
|
||
should_yield = False
|
||
if not has_yielded_audio and pending_audio_size >= min_chunk_size:
|
||
# First yield: wait for minimum chunk size
|
||
should_yield = True
|
||
has_yielded_audio = True
|
||
elif has_yielded_audio and (pending_audio_size >= min_chunk_size or time_since_last_yield >= min_yield_interval):
|
||
# Subsequent yields: either enough audio or enough time has passed
|
||
should_yield = True
|
||
|
||
if should_yield and pending_chunks:
|
||
# Concatenate and yield only the new audio chunks
|
||
new_audio = np.concatenate(pending_chunks)
|
||
new_duration = len(new_audio) / sample_rate
|
||
total_duration = sum(len(chunk) for chunk in all_audio_chunks) / sample_rate
|
||
|
||
log_update = log + f"🎵 Streaming: {total_duration:.1f}s generated (chunk {chunk_count})\n"
|
||
|
||
# Yield streaming audio chunk and keep complete_audio as None during streaming
|
||
yield (sample_rate, new_audio), None, log_update, gr.update(visible=True)
|
||
|
||
# Clear pending chunks after yielding
|
||
pending_chunks = []
|
||
last_yield_time = current_time
|
||
|
||
# Yield any remaining chunks
|
||
if pending_chunks:
|
||
final_new_audio = np.concatenate(pending_chunks)
|
||
total_duration = sum(len(chunk) for chunk in all_audio_chunks) / sample_rate
|
||
log_update = log + f"🎵 Streaming final chunk: {total_duration:.1f}s total\n"
|
||
yield (sample_rate, final_new_audio), None, log_update, gr.update(visible=True)
|
||
has_yielded_audio = True # Mark that we yielded audio
|
||
|
||
# Wait for generation to complete (with timeout to prevent hanging)
|
||
generation_thread.join(timeout=5.0) # Increased timeout to 5 seconds
|
||
|
||
# If thread is still alive after timeout, force end
|
||
if generation_thread.is_alive():
|
||
print("Warning: Generation thread did not complete within timeout")
|
||
audio_streamer.end()
|
||
generation_thread.join(timeout=5.0)
|
||
|
||
# Clean up
|
||
self.current_streamer = None
|
||
self.is_generating = False
|
||
|
||
generation_time = time.time() - start_time
|
||
|
||
# Check if stopped by user
|
||
if self.stop_generation:
|
||
yield None, None, "🛑 Generation stopped by user", gr.update(visible=False)
|
||
return
|
||
|
||
# Debug logging
|
||
# print(f"Debug: has_received_chunks={has_received_chunks}, chunk_count={chunk_count}, all_audio_chunks length={len(all_audio_chunks)}")
|
||
|
||
# Check if we received any chunks but didn't yield audio
|
||
if has_received_chunks and not has_yielded_audio and all_audio_chunks:
|
||
# We have chunks but didn't meet the yield criteria, yield them now
|
||
complete_audio = np.concatenate(all_audio_chunks)
|
||
final_duration = len(complete_audio) / sample_rate
|
||
|
||
final_log = log + f"⏱️ Generation completed in {generation_time:.2f} seconds\n"
|
||
final_log += f"🎵 Final audio duration: {final_duration:.2f} seconds\n"
|
||
final_log += f"📊 Total chunks: {chunk_count}\n"
|
||
final_log += "✨ Generation successful! Complete audio is ready.\n"
|
||
final_log += "💡 Not satisfied? You can regenerate or adjust the CFG scale for different results."
|
||
|
||
# Yield the complete audio
|
||
yield None, (sample_rate, complete_audio), final_log, gr.update(visible=False)
|
||
return
|
||
|
||
if not has_received_chunks:
|
||
error_log = log + f"\n❌ Error: No audio chunks were received from the model. Generation time: {generation_time:.2f}s"
|
||
yield None, None, error_log, gr.update(visible=False)
|
||
return
|
||
|
||
if not has_yielded_audio:
|
||
error_log = log + f"\n❌ Error: Audio was generated but not streamed. Chunk count: {chunk_count}"
|
||
yield None, None, error_log, gr.update(visible=False)
|
||
return
|
||
|
||
# Prepare the complete audio
|
||
if all_audio_chunks:
|
||
complete_audio = np.concatenate(all_audio_chunks)
|
||
final_duration = len(complete_audio) / sample_rate
|
||
|
||
final_log = log + f"⏱️ Generation completed in {generation_time:.2f} seconds\n"
|
||
final_log += f"🎵 Final audio duration: {final_duration:.2f} seconds\n"
|
||
final_log += f"📊 Total chunks: {chunk_count}\n"
|
||
final_log += "✨ Generation successful! Complete audio is ready in the 'Complete Audio' tab.\n"
|
||
final_log += "💡 Not satisfied? You can regenerate or adjust the CFG scale for different results."
|
||
|
||
# Final yield: Clear streaming audio and provide complete audio
|
||
yield None, (sample_rate, complete_audio), final_log, gr.update(visible=False)
|
||
else:
|
||
final_log = log + "❌ No audio was generated."
|
||
yield None, None, final_log, gr.update(visible=False)
|
||
|
||
except gr.Error as e:
|
||
# Handle Gradio-specific errors (like input validation)
|
||
self.is_generating = False
|
||
self.current_streamer = None
|
||
error_msg = f"❌ Input Error: {str(e)}"
|
||
print(error_msg)
|
||
yield None, None, error_msg, gr.update(visible=False)
|
||
|
||
except Exception as e:
|
||
self.is_generating = False
|
||
self.current_streamer = None
|
||
error_msg = f"❌ An unexpected error occurred: {str(e)}"
|
||
print(error_msg)
|
||
import traceback
|
||
traceback.print_exc()
|
||
yield None, None, error_msg, gr.update(visible=False)
|
||
|
||
def _generate_with_streamer(self, inputs, cfg_scale, audio_streamer):
|
||
"""Helper method to run generation with streamer in a separate thread."""
|
||
try:
|
||
# Check for stop signal before starting generation
|
||
if self.stop_generation:
|
||
audio_streamer.end()
|
||
return
|
||
|
||
# Define a stop check function that can be called from generate
|
||
def check_stop_generation():
|
||
return self.stop_generation
|
||
|
||
outputs = self.model.generate(
|
||
**inputs,
|
||
max_new_tokens=None,
|
||
cfg_scale=cfg_scale,
|
||
tokenizer=self.processor.tokenizer,
|
||
generation_config={
|
||
'do_sample': False,
|
||
},
|
||
audio_streamer=audio_streamer,
|
||
stop_check_fn=check_stop_generation, # Pass the stop check function
|
||
verbose=False, # Disable verbose in streaming mode
|
||
refresh_negative=True,
|
||
)
|
||
|
||
except Exception as e:
|
||
print(f"Error in generation thread: {e}")
|
||
traceback.print_exc()
|
||
# Make sure to end the stream on error
|
||
audio_streamer.end()
|
||
|
||
def stop_audio_generation(self):
|
||
"""Stop the current audio generation process."""
|
||
self.stop_generation = True
|
||
if self.current_streamer is not None:
|
||
try:
|
||
self.current_streamer.end()
|
||
except Exception as e:
|
||
print(f"Error stopping streamer: {e}")
|
||
print("🛑 Audio generation stop requested")
|
||
|
||
def load_example_scripts(self):
|
||
"""Load example scripts from the text_examples directory."""
|
||
examples_dir = os.path.join(os.path.dirname(__file__), "text_examples")
|
||
self.example_scripts = []
|
||
|
||
# Check if text_examples directory exists
|
||
if not os.path.exists(examples_dir):
|
||
print(f"Warning: text_examples directory not found at {examples_dir}")
|
||
return
|
||
|
||
# Get all .txt files in the text_examples directory
|
||
txt_files = sorted([f for f in os.listdir(examples_dir)
|
||
if f.lower().endswith('.txt') and os.path.isfile(os.path.join(examples_dir, f))])
|
||
|
||
for txt_file in txt_files:
|
||
file_path = os.path.join(examples_dir, txt_file)
|
||
|
||
import re
|
||
# Check if filename contains a time pattern like "45min", "90min", etc.
|
||
time_pattern = re.search(r'(\d+)min', txt_file.lower())
|
||
if time_pattern:
|
||
minutes = int(time_pattern.group(1))
|
||
if minutes > 15:
|
||
print(f"Skipping {txt_file}: duration {minutes} minutes exceeds 15-minute limit")
|
||
continue
|
||
|
||
try:
|
||
with open(file_path, 'r', encoding='utf-8') as f:
|
||
script_content = f.read().strip()
|
||
|
||
# Remove empty lines and lines with only whitespace
|
||
script_content = '\n'.join(line for line in script_content.split('\n') if line.strip())
|
||
|
||
if not script_content:
|
||
continue
|
||
|
||
# Parse the script to determine number of speakers
|
||
num_speakers = self._get_num_speakers_from_script(script_content)
|
||
|
||
# Add to examples list as [num_speakers, script_content]
|
||
self.example_scripts.append([num_speakers, script_content])
|
||
print(f"Loaded example: {txt_file} with {num_speakers} speakers")
|
||
|
||
except Exception as e:
|
||
print(f"Error loading example script {txt_file}: {e}")
|
||
|
||
if self.example_scripts:
|
||
print(f"Successfully loaded {len(self.example_scripts)} example scripts")
|
||
else:
|
||
print("No example scripts were loaded")
|
||
|
||
def _get_num_speakers_from_script(self, script: str) -> int:
|
||
"""Determine the number of unique speakers in a script."""
|
||
import re
|
||
speakers = set()
|
||
|
||
lines = script.strip().split('\n')
|
||
for line in lines:
|
||
# Use regex to find speaker patterns
|
||
match = re.match(r'^Speaker\s+(\d+)\s*:', line.strip(), re.IGNORECASE)
|
||
if match:
|
||
speaker_id = int(match.group(1))
|
||
speakers.add(speaker_id)
|
||
|
||
# If no speakers found, default to 1
|
||
if not speakers:
|
||
return 1
|
||
|
||
# Return the maximum speaker ID + 1 (assuming 0-based indexing)
|
||
# or the count of unique speakers if they're 1-based
|
||
max_speaker = max(speakers)
|
||
min_speaker = min(speakers)
|
||
|
||
if min_speaker == 0:
|
||
return max_speaker + 1
|
||
else:
|
||
# Assume 1-based indexing, return the count
|
||
return len(speakers)
|
||
|
||
|
||
def create_demo_interface(demo_instance: VibeVoiceDemo):
|
||
"""Create the Gradio interface with streaming support."""
|
||
|
||
# Custom CSS for high-end aesthetics with lighter theme
|
||
custom_css = """
|
||
/* Modern light theme with gradients */
|
||
.gradio-container {
|
||
background: linear-gradient(135deg, #f8fafc 0%, #e2e8f0 100%);
|
||
font-family: 'SF Pro Display', -apple-system, BlinkMacSystemFont, sans-serif;
|
||
}
|
||
|
||
/* Header styling */
|
||
.main-header {
|
||
background: linear-gradient(90deg, #667eea 0%, #764ba2 100%);
|
||
padding: 2rem;
|
||
border-radius: 20px;
|
||
margin-bottom: 2rem;
|
||
text-align: center;
|
||
box-shadow: 0 10px 40px rgba(102, 126, 234, 0.3);
|
||
}
|
||
|
||
.main-header h1 {
|
||
color: white;
|
||
font-size: 2.5rem;
|
||
font-weight: 700;
|
||
margin: 0;
|
||
text-shadow: 0 2px 4px rgba(0,0,0,0.3);
|
||
}
|
||
|
||
.main-header p {
|
||
color: rgba(255,255,255,0.9);
|
||
font-size: 1.1rem;
|
||
margin: 0.5rem 0 0 0;
|
||
}
|
||
|
||
/* Card styling */
|
||
.settings-card, .generation-card {
|
||
background: rgba(255, 255, 255, 0.8);
|
||
backdrop-filter: blur(10px);
|
||
border: 1px solid rgba(226, 232, 240, 0.8);
|
||
border-radius: 16px;
|
||
padding: 1.5rem;
|
||
margin-bottom: 1rem;
|
||
box-shadow: 0 8px 32px rgba(0, 0, 0, 0.1);
|
||
}
|
||
|
||
/* Speaker selection styling */
|
||
.speaker-grid {
|
||
display: grid;
|
||
gap: 1rem;
|
||
margin-bottom: 1rem;
|
||
}
|
||
|
||
.speaker-item {
|
||
background: linear-gradient(135deg, #e2e8f0 0%, #cbd5e1 100%);
|
||
border: 1px solid rgba(148, 163, 184, 0.4);
|
||
border-radius: 12px;
|
||
padding: 1rem;
|
||
color: #374151;
|
||
font-weight: 500;
|
||
}
|
||
|
||
/* Streaming indicator */
|
||
.streaming-indicator {
|
||
display: inline-block;
|
||
width: 10px;
|
||
height: 10px;
|
||
background: #22c55e;
|
||
border-radius: 50%;
|
||
margin-right: 8px;
|
||
animation: pulse 1.5s infinite;
|
||
}
|
||
|
||
@keyframes pulse {
|
||
0% { opacity: 1; transform: scale(1); }
|
||
50% { opacity: 0.5; transform: scale(1.1); }
|
||
100% { opacity: 1; transform: scale(1); }
|
||
}
|
||
|
||
/* Queue status styling */
|
||
.queue-status {
|
||
background: linear-gradient(135deg, #f0f9ff 0%, #e0f2fe 100%);
|
||
border: 1px solid rgba(14, 165, 233, 0.3);
|
||
border-radius: 8px;
|
||
padding: 0.75rem;
|
||
margin: 0.5rem 0;
|
||
text-align: center;
|
||
font-size: 0.9rem;
|
||
color: #0369a1;
|
||
}
|
||
|
||
.generate-btn {
|
||
background: linear-gradient(135deg, #059669 0%, #0d9488 100%);
|
||
border: none;
|
||
border-radius: 12px;
|
||
padding: 1rem 2rem;
|
||
color: white;
|
||
font-weight: 600;
|
||
font-size: 1.1rem;
|
||
box-shadow: 0 4px 20px rgba(5, 150, 105, 0.4);
|
||
transition: all 0.3s ease;
|
||
}
|
||
|
||
.generate-btn:hover {
|
||
transform: translateY(-2px);
|
||
box-shadow: 0 6px 25px rgba(5, 150, 105, 0.6);
|
||
}
|
||
|
||
.stop-btn {
|
||
background: linear-gradient(135deg, #ef4444 0%, #dc2626 100%);
|
||
border: none;
|
||
border-radius: 12px;
|
||
padding: 1rem 2rem;
|
||
color: white;
|
||
font-weight: 600;
|
||
font-size: 1.1rem;
|
||
box-shadow: 0 4px 20px rgba(239, 68, 68, 0.4);
|
||
transition: all 0.3s ease;
|
||
}
|
||
|
||
.stop-btn:hover {
|
||
transform: translateY(-2px);
|
||
box-shadow: 0 6px 25px rgba(239, 68, 68, 0.6);
|
||
}
|
||
|
||
/* Audio player styling */
|
||
.audio-output {
|
||
background: linear-gradient(135deg, #f1f5f9 0%, #e2e8f0 100%);
|
||
border-radius: 16px;
|
||
padding: 1.5rem;
|
||
border: 1px solid rgba(148, 163, 184, 0.3);
|
||
}
|
||
|
||
.complete-audio-section {
|
||
margin-top: 1rem;
|
||
padding: 1rem;
|
||
background: linear-gradient(135deg, #f0fdf4 0%, #dcfce7 100%);
|
||
border: 1px solid rgba(34, 197, 94, 0.3);
|
||
border-radius: 12px;
|
||
}
|
||
|
||
/* Text areas */
|
||
.script-input, .log-output {
|
||
background: rgba(255, 255, 255, 0.9) !important;
|
||
border: 1px solid rgba(148, 163, 184, 0.4) !important;
|
||
border-radius: 12px !important;
|
||
color: #1e293b !important;
|
||
font-family: 'JetBrains Mono', monospace !important;
|
||
}
|
||
|
||
.script-input::placeholder {
|
||
color: #64748b !important;
|
||
}
|
||
|
||
/* Sliders */
|
||
.slider-container {
|
||
background: rgba(248, 250, 252, 0.8);
|
||
border: 1px solid rgba(226, 232, 240, 0.6);
|
||
border-radius: 8px;
|
||
padding: 1rem;
|
||
margin: 0.5rem 0;
|
||
}
|
||
|
||
/* Labels and text */
|
||
.gradio-container label {
|
||
color: #374151 !important;
|
||
font-weight: 600 !important;
|
||
}
|
||
|
||
.gradio-container .markdown {
|
||
color: #1f2937 !important;
|
||
}
|
||
|
||
/* Responsive design */
|
||
@media (max-width: 768px) {
|
||
.main-header h1 { font-size: 2rem; }
|
||
.settings-card, .generation-card { padding: 1rem; }
|
||
}
|
||
|
||
/* Random example button styling - more subtle professional color */
|
||
.random-btn {
|
||
background: linear-gradient(135deg, #64748b 0%, #475569 100%);
|
||
border: none;
|
||
border-radius: 12px;
|
||
padding: 1rem 1.5rem;
|
||
color: white;
|
||
font-weight: 600;
|
||
font-size: 1rem;
|
||
box-shadow: 0 4px 20px rgba(100, 116, 139, 0.3);
|
||
transition: all 0.3s ease;
|
||
display: inline-flex;
|
||
align-items: center;
|
||
gap: 0.5rem;
|
||
}
|
||
|
||
.random-btn:hover {
|
||
transform: translateY(-2px);
|
||
box-shadow: 0 6px 25px rgba(100, 116, 139, 0.4);
|
||
background: linear-gradient(135deg, #475569 0%, #334155 100%);
|
||
}
|
||
"""
|
||
|
||
with gr.Blocks(
|
||
title="VibeVoice - AI Podcast Generator",
|
||
css=custom_css,
|
||
theme=gr.themes.Soft(
|
||
primary_hue="blue",
|
||
secondary_hue="purple",
|
||
neutral_hue="slate",
|
||
)
|
||
) as interface:
|
||
|
||
# Header
|
||
gr.HTML("""
|
||
<div class="main-header">
|
||
<h1>🎙️ Vibe Podcasting </h1>
|
||
<p>Generating Long-form Multi-speaker AI Podcast with VibeVoice</p>
|
||
</div>
|
||
""")
|
||
|
||
with gr.Row():
|
||
# Left column - Settings
|
||
with gr.Column(scale=1, elem_classes="settings-card"):
|
||
gr.Markdown("### 🎛️ **Podcast Settings**")
|
||
|
||
# Number of speakers
|
||
num_speakers = gr.Slider(
|
||
minimum=1,
|
||
maximum=4,
|
||
value=2,
|
||
step=1,
|
||
label="Number of Speakers",
|
||
elem_classes="slider-container"
|
||
)
|
||
|
||
# Speaker selection
|
||
gr.Markdown("### 🎭 **Speaker Selection**")
|
||
|
||
available_speaker_names = list(demo_instance.available_voices.keys())
|
||
# default_speakers = available_speaker_names[:4] if len(available_speaker_names) >= 4 else available_speaker_names
|
||
default_speakers = ['en-Alice_woman', 'en-Carter_man', 'en-Frank_man', 'en-Maya_woman']
|
||
|
||
speaker_selections = []
|
||
for i in range(4):
|
||
default_value = default_speakers[i] if i < len(default_speakers) else None
|
||
speaker = gr.Dropdown(
|
||
choices=available_speaker_names,
|
||
value=default_value,
|
||
label=f"Speaker {i+1}",
|
||
visible=(i < 2), # Initially show only first 2 speakers
|
||
elem_classes="speaker-item"
|
||
)
|
||
speaker_selections.append(speaker)
|
||
|
||
# Advanced settings
|
||
gr.Markdown("### ⚙️ **Advanced Settings**")
|
||
|
||
# Sampling parameters (contains all generation settings)
|
||
with gr.Accordion("Generation Parameters", open=False):
|
||
cfg_scale = gr.Slider(
|
||
minimum=1.0,
|
||
maximum=2.0,
|
||
value=1.3,
|
||
step=0.05,
|
||
label="CFG Scale (Guidance Strength)",
|
||
# info="Higher values increase adherence to text",
|
||
elem_classes="slider-container"
|
||
)
|
||
|
||
# Right column - Generation
|
||
with gr.Column(scale=2, elem_classes="generation-card"):
|
||
gr.Markdown("### 📝 **Script Input**")
|
||
|
||
script_input = gr.Textbox(
|
||
label="Conversation Script",
|
||
placeholder="""Enter your podcast script here. You can format it as:
|
||
|
||
Speaker 0: Welcome to our podcast today!
|
||
Speaker 1: Thanks for having me. I'm excited to discuss...
|
||
|
||
Or paste text directly and it will auto-assign speakers.""",
|
||
lines=12,
|
||
max_lines=20,
|
||
elem_classes="script-input"
|
||
)
|
||
|
||
# Button row with Random Example on the left and Generate on the right
|
||
with gr.Row():
|
||
# Random example button (now on the left)
|
||
random_example_btn = gr.Button(
|
||
"🎲 Random Example",
|
||
size="lg",
|
||
variant="secondary",
|
||
elem_classes="random-btn",
|
||
scale=1 # Smaller width
|
||
)
|
||
|
||
# Generate button (now on the right)
|
||
generate_btn = gr.Button(
|
||
"🚀 Generate Podcast",
|
||
size="lg",
|
||
variant="primary",
|
||
elem_classes="generate-btn",
|
||
scale=2 # Wider than random button
|
||
)
|
||
|
||
# Stop button
|
||
stop_btn = gr.Button(
|
||
"🛑 Stop Generation",
|
||
size="lg",
|
||
variant="stop",
|
||
elem_classes="stop-btn",
|
||
visible=False
|
||
)
|
||
|
||
# Streaming status indicator
|
||
streaming_status = gr.HTML(
|
||
value="""
|
||
<div style="background: linear-gradient(135deg, #dcfce7 0%, #bbf7d0 100%);
|
||
border: 1px solid rgba(34, 197, 94, 0.3);
|
||
border-radius: 8px;
|
||
padding: 0.75rem;
|
||
margin: 0.5rem 0;
|
||
text-align: center;
|
||
font-size: 0.9rem;
|
||
color: #166534;">
|
||
<span class="streaming-indicator"></span>
|
||
<strong>LIVE STREAMING</strong> - Audio is being generated in real-time
|
||
</div>
|
||
""",
|
||
visible=False,
|
||
elem_id="streaming-status"
|
||
)
|
||
|
||
# Output section
|
||
gr.Markdown("### 🎵 **Generated Podcast**")
|
||
|
||
# Streaming audio output (outside of tabs for simpler handling)
|
||
audio_output = gr.Audio(
|
||
label="Streaming Audio (Real-time)",
|
||
type="numpy",
|
||
elem_classes="audio-output",
|
||
streaming=True, # Enable streaming mode
|
||
autoplay=True,
|
||
show_download_button=False, # Explicitly show download button
|
||
visible=True
|
||
)
|
||
|
||
# Complete audio output (non-streaming)
|
||
complete_audio_output = gr.Audio(
|
||
label="Complete Podcast (Download after generation)",
|
||
type="numpy",
|
||
elem_classes="audio-output complete-audio-section",
|
||
streaming=False, # Non-streaming mode
|
||
autoplay=False,
|
||
show_download_button=True, # Explicitly show download button
|
||
visible=False # Initially hidden, shown when audio is ready
|
||
)
|
||
|
||
gr.Markdown("""
|
||
*💡 **Streaming**: Audio plays as it's being generated (may have slight pauses)
|
||
*💡 **Complete Audio**: Will appear below after generation finishes*
|
||
""")
|
||
|
||
# Generation log
|
||
log_output = gr.Textbox(
|
||
label="Generation Log",
|
||
lines=8,
|
||
max_lines=15,
|
||
interactive=False,
|
||
elem_classes="log-output"
|
||
)
|
||
|
||
def update_speaker_visibility(num_speakers):
|
||
updates = []
|
||
for i in range(4):
|
||
updates.append(gr.update(visible=(i < num_speakers)))
|
||
return updates
|
||
|
||
num_speakers.change(
|
||
fn=update_speaker_visibility,
|
||
inputs=[num_speakers],
|
||
outputs=speaker_selections
|
||
)
|
||
|
||
# Main generation function with streaming
|
||
def generate_podcast_wrapper(num_speakers, script, *speakers_and_params):
|
||
"""Wrapper function to handle the streaming generation call."""
|
||
try:
|
||
# Extract speakers and parameters
|
||
speakers = speakers_and_params[:4] # First 4 are speaker selections
|
||
cfg_scale = speakers_and_params[4] # CFG scale
|
||
|
||
# Clear outputs and reset visibility at start
|
||
yield None, gr.update(value=None, visible=False), "🎙️ Starting generation...", gr.update(visible=True), gr.update(visible=False), gr.update(visible=True)
|
||
|
||
# The generator will yield multiple times
|
||
final_log = "Starting generation..."
|
||
|
||
for streaming_audio, complete_audio, log, streaming_visible in demo_instance.generate_podcast_streaming(
|
||
num_speakers=int(num_speakers),
|
||
script=script,
|
||
speaker_1=speakers[0],
|
||
speaker_2=speakers[1],
|
||
speaker_3=speakers[2],
|
||
speaker_4=speakers[3],
|
||
cfg_scale=cfg_scale
|
||
):
|
||
final_log = log
|
||
|
||
# Check if we have complete audio (final yield)
|
||
if complete_audio is not None:
|
||
# Final state: clear streaming, show complete audio
|
||
yield None, gr.update(value=complete_audio, visible=True), log, gr.update(visible=False), gr.update(visible=True), gr.update(visible=False)
|
||
else:
|
||
# Streaming state: update streaming audio only
|
||
if streaming_audio is not None:
|
||
yield streaming_audio, gr.update(visible=False), log, streaming_visible, gr.update(visible=False), gr.update(visible=True)
|
||
else:
|
||
# No new audio, just update status
|
||
yield None, gr.update(visible=False), log, streaming_visible, gr.update(visible=False), gr.update(visible=True)
|
||
|
||
except Exception as e:
|
||
error_msg = f"❌ A critical error occurred in the wrapper: {str(e)}"
|
||
print(error_msg)
|
||
import traceback
|
||
traceback.print_exc()
|
||
# Reset button states on error
|
||
yield None, gr.update(value=None, visible=False), error_msg, gr.update(visible=False), gr.update(visible=True), gr.update(visible=False)
|
||
|
||
def stop_generation_handler():
|
||
"""Handle stopping generation."""
|
||
demo_instance.stop_audio_generation()
|
||
# Return values for: log_output, streaming_status, generate_btn, stop_btn
|
||
return "🛑 Generation stopped.", gr.update(visible=False), gr.update(visible=True), gr.update(visible=False)
|
||
|
||
# Add a clear audio function
|
||
def clear_audio_outputs():
|
||
"""Clear both audio outputs before starting new generation."""
|
||
return None, gr.update(value=None, visible=False)
|
||
|
||
# Connect generation button with streaming outputs
|
||
generate_btn.click(
|
||
fn=clear_audio_outputs,
|
||
inputs=[],
|
||
outputs=[audio_output, complete_audio_output],
|
||
queue=False
|
||
).then(
|
||
fn=generate_podcast_wrapper,
|
||
inputs=[num_speakers, script_input] + speaker_selections + [cfg_scale],
|
||
outputs=[audio_output, complete_audio_output, log_output, streaming_status, generate_btn, stop_btn],
|
||
queue=True # Enable Gradio's built-in queue
|
||
)
|
||
|
||
# Connect stop button
|
||
stop_btn.click(
|
||
fn=stop_generation_handler,
|
||
inputs=[],
|
||
outputs=[log_output, streaming_status, generate_btn, stop_btn],
|
||
queue=False # Don't queue stop requests
|
||
).then(
|
||
# Clear both audio outputs after stopping
|
||
fn=lambda: (None, None),
|
||
inputs=[],
|
||
outputs=[audio_output, complete_audio_output],
|
||
queue=False
|
||
)
|
||
|
||
# Function to randomly select an example
|
||
def load_random_example():
|
||
"""Randomly select and load an example script."""
|
||
import random
|
||
|
||
# Get available examples
|
||
if hasattr(demo_instance, 'example_scripts') and demo_instance.example_scripts:
|
||
example_scripts = demo_instance.example_scripts
|
||
else:
|
||
# Fallback to default
|
||
example_scripts = [
|
||
[2, "Speaker 0: Welcome to our AI podcast demonstration!\nSpeaker 1: Thanks for having me. This is exciting!"]
|
||
]
|
||
|
||
# Randomly select one
|
||
if example_scripts:
|
||
selected = random.choice(example_scripts)
|
||
num_speakers_value = selected[0]
|
||
script_value = selected[1]
|
||
|
||
# Return the values to update the UI
|
||
return num_speakers_value, script_value
|
||
|
||
# Default values if no examples
|
||
return 2, ""
|
||
|
||
# Connect random example button
|
||
random_example_btn.click(
|
||
fn=load_random_example,
|
||
inputs=[],
|
||
outputs=[num_speakers, script_input],
|
||
queue=False # Don't queue this simple operation
|
||
)
|
||
|
||
# Add usage tips
|
||
gr.Markdown("""
|
||
### 💡 **Usage Tips**
|
||
|
||
- Click **🚀 Generate Podcast** to start audio generation
|
||
- **Live Streaming** tab shows audio as it's generated (may have slight pauses)
|
||
- **Complete Audio** tab provides the full, uninterrupted podcast after generation
|
||
- During generation, you can click **🛑 Stop Generation** to interrupt the process
|
||
- The streaming indicator shows real-time generation progress
|
||
""")
|
||
|
||
# Add example scripts
|
||
gr.Markdown("### 📚 **Example Scripts**")
|
||
|
||
# Use dynamically loaded examples if available, otherwise provide a default
|
||
if hasattr(demo_instance, 'example_scripts') and demo_instance.example_scripts:
|
||
example_scripts = demo_instance.example_scripts
|
||
else:
|
||
# Fallback to a simple default example if no scripts loaded
|
||
example_scripts = [
|
||
[1, "Speaker 1: Welcome to our AI podcast demonstration! This is a sample script showing how VibeVoice can generate natural-sounding speech."]
|
||
]
|
||
|
||
gr.Examples(
|
||
examples=example_scripts,
|
||
inputs=[num_speakers, script_input],
|
||
label="Try these example scripts:"
|
||
)
|
||
|
||
return interface
|
||
|
||
|
||
def convert_to_16_bit_wav(data):
|
||
# Check if data is a tensor and move to cpu
|
||
if torch.is_tensor(data):
|
||
data = data.detach().cpu().numpy()
|
||
|
||
# Ensure data is numpy array
|
||
data = np.array(data)
|
||
|
||
# Normalize to range [-1, 1] if it's not already
|
||
if np.max(np.abs(data)) > 1.0:
|
||
data = data / np.max(np.abs(data))
|
||
|
||
# Scale to 16-bit integer range
|
||
data = (data * 32767).astype(np.int16)
|
||
return data
|
||
|
||
|
||
def parse_args():
|
||
parser = argparse.ArgumentParser(description="VibeVoice Gradio Demo")
|
||
parser.add_argument(
|
||
"--model_path",
|
||
type=str,
|
||
default="/tmp/vibevoice-model",
|
||
help="Path to the VibeVoice model directory",
|
||
)
|
||
parser.add_argument(
|
||
"--device",
|
||
type=str,
|
||
default="cuda" if torch.cuda.is_available() else "cpu",
|
||
help="Device for inference",
|
||
)
|
||
parser.add_argument(
|
||
"--inference_steps",
|
||
type=int,
|
||
default=10,
|
||
help="Number of inference steps for DDPM (not exposed to users)",
|
||
)
|
||
parser.add_argument(
|
||
"--share",
|
||
action="store_true",
|
||
help="Share the demo publicly via Gradio",
|
||
)
|
||
parser.add_argument(
|
||
"--port",
|
||
type=int,
|
||
default=7860,
|
||
help="Port to run the demo on",
|
||
)
|
||
|
||
return parser.parse_args()
|
||
|
||
|
||
def main():
|
||
"""Main function to run the demo."""
|
||
args = parse_args()
|
||
|
||
set_seed(42) # Set a fixed seed for reproducibility
|
||
|
||
print("🎙️ Initializing VibeVoice Demo with Streaming Support...")
|
||
|
||
# Initialize demo instance
|
||
demo_instance = VibeVoiceDemo(
|
||
model_path=args.model_path,
|
||
device=args.device,
|
||
inference_steps=args.inference_steps
|
||
)
|
||
|
||
# Create interface
|
||
interface = create_demo_interface(demo_instance)
|
||
|
||
print(f"🚀 Launching demo on port {args.port}")
|
||
print(f"📁 Model path: {args.model_path}")
|
||
print(f"🎭 Available voices: {len(demo_instance.available_voices)}")
|
||
print(f"🔴 Streaming mode: ENABLED")
|
||
print(f"🔒 Session isolation: ENABLED")
|
||
|
||
# Launch the interface
|
||
try:
|
||
interface.queue(
|
||
max_size=20, # Maximum queue size
|
||
default_concurrency_limit=1 # Process one request at a time
|
||
).launch(
|
||
share=args.share,
|
||
# server_port=args.port,
|
||
server_name="0.0.0.0" if args.share else "127.0.0.1",
|
||
show_error=True,
|
||
show_api=False # Hide API docs for cleaner interface
|
||
)
|
||
except KeyboardInterrupt:
|
||
print("\n🛑 Shutting down gracefully...")
|
||
except Exception as e:
|
||
print(f"❌ Server error: {e}")
|
||
raise
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|