import os import time import numpy as np import gradio as gr import librosa import soundfile as sf import torch import traceback import threading from spaces import GPU from datetime import datetime 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_paths: dict, device: str = "cuda", inference_steps: int = 5): """ model_paths: dict like {"VibeVoice-1.5B": "microsoft/VibeVoice-1.5B", "VibeVoice-1.1B": "microsoft/VibeVoice-1.1B"} """ self.model_paths = model_paths self.device = device self.inference_steps = inference_steps self.is_generating = False # Multi-model holders self.models = {} # name -> model self.processors = {} # name -> processor self.current_model_name = None self.available_voices = {} self.load_models() # load all on CPU self.setup_voice_presets() self.load_example_scripts() def load_models(self): print("Loading processors and models on CPU...") for name, path in self.model_paths.items(): print(f" - {name} from {path}") proc = VibeVoiceProcessor.from_pretrained(path) mdl = VibeVoiceForConditionalGenerationInference.from_pretrained( path, torch_dtype=torch.bfloat16 ) # Keep on CPU initially self.processors[name] = proc self.models[name] = mdl # choose default self.current_model_name = next(iter(self.models)) print(f"Default model is {self.current_model_name}") def _place_model(self, target_name: str): """ Move the selected model to CUDA and push all others back to CPU. """ for name, mdl in self.models.items(): if name == target_name: self.models[name] = mdl.to(self.device) else: self.models[name] = mdl.to("cpu") self.current_model_name = target_name print(f"Model {target_name} is now on {self.device}. Others moved to CPU.") def setup_voice_presets(self): voices_dir = os.path.join(os.path.dirname(__file__), "voices") if not os.path.exists(voices_dir): print(f"Warning: Voices directory not found at {voices_dir}") return wav_files = [f for f in os.listdir(voices_dir) if f.lower().endswith(('.wav', '.mp3', '.flac', '.ogg', '.m4a', '.aac'))] for wav_file in wav_files: name = os.path.splitext(wav_file)[0] self.available_voices[name] = os.path.join(voices_dir, wav_file) print(f"Voices loaded: {list(self.available_voices.keys())}") def read_audio(self, audio_path: str, target_sr: int = 24000) -> np.ndarray: 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([]) @GPU(duration=120) @torch.inference_mode() def generate_podcast(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, model_name: str = None): """ Generates a podcast as a single audio file from a script and saves it. Non-streaming. """ try: # pick model model_name = model_name or self.current_model_name if model_name not in self.models: raise gr.Error(f"Unknown model: {model_name}") # place models on devices self._place_model(model_name) model = self.models[model_name] processor = self.processors[model_name] print(f"Using model {model_name} on {self.device}") model.eval() model.set_ddpm_inference_steps(num_steps=self.inference_steps) self.is_generating = True if not script.strip(): raise gr.Error("Error: Please provide a script.") script = script.replace("’", "'") if not 1 <= num_speakers <= 4: raise gr.Error("Error: Number of speakers must be between 1 and 4.") selected_speakers = [speaker_1, speaker_2, speaker_3, speaker_4][:num_speakers] for i, speaker_name in enumerate(selected_speakers): if not speaker_name or speaker_name not in self.available_voices: raise gr.Error(f"Error: Please select a valid speaker for Speaker {i+1}.") log = f"🎙️ Generating podcast with {num_speakers} speakers\n" log += f"🧠 Model: {model_name}\n" log += f"📊 Parameters: CFG Scale={cfg_scale}\n" log += f"🎭 Speakers: {', '.join(selected_speakers)}\n" 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: 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" lines = script.strip().split('\n') formatted_script_lines = [] for line in lines: line = line.strip() if not line: continue if line.startswith('Speaker ') and ':' in line: formatted_script_lines.append(line) else: 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" log += "🔄 Processing with VibeVoice...\n" inputs = processor( text=[formatted_script], voice_samples=[voice_samples], padding=True, return_tensors="pt", return_attention_mask=True, ) start_time = time.time() outputs = model.generate( **inputs, max_new_tokens=None, cfg_scale=cfg_scale, tokenizer=processor.tokenizer, generation_config={'do_sample': False}, verbose=False, ) generation_time = time.time() - start_time if hasattr(outputs, 'speech_outputs') and outputs.speech_outputs[0] is not None: audio_tensor = outputs.speech_outputs[0] audio = audio_tensor.cpu().float().numpy() else: raise gr.Error("❌ Error: No audio was generated by the model. Please try again.") if audio.ndim > 1: audio = audio.squeeze() sample_rate = 24000 output_dir = "outputs" os.makedirs(output_dir, exist_ok=True) timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") file_path = os.path.join(output_dir, f"podcast_{timestamp}.wav") sf.write(file_path, audio, sample_rate) print(f"💾 Podcast saved to {file_path}") total_duration = len(audio) / sample_rate log += f"⏱️ Generation completed in {generation_time:.2f} seconds\n" log += f"🎵 Final audio duration: {total_duration:.2f} seconds\n" log += f"✅ Successfully saved podcast to: {file_path}\n" self.is_generating = False return (sample_rate, audio), log except gr.Error as e: self.is_generating = False error_msg = f"❌ Input Error: {str(e)}" print(error_msg) return None, error_msg except Exception as e: self.is_generating = False error_msg = f"❌ An unexpected error occurred: {str(e)}" print(error_msg) traceback.print_exc() return None, error_msg @staticmethod def _infer_num_speakers_from_script(script: str) -> int: """ Infer number of speakers by counting distinct 'Speaker X:' tags in the script. Robust to 0- or 1-indexed labels and repeated turns. Falls back to 1 if none found. """ import re ids = re.findall(r'(?mi)^\s*Speaker\s+(\d+)\s*:', script) return len({int(x) for x in ids}) if ids else 1 def load_example_scripts(self): examples_dir = os.path.join(os.path.dirname(__file__), "text_examples") self.example_scripts = [] if not os.path.exists(examples_dir): return txt_files = sorted( [f for f in os.listdir(examples_dir) if f.lower().endswith('.txt')] ) for txt_file in txt_files: try: with open(os.path.join(examples_dir, txt_file), 'r', encoding='utf-8') as f: script_content = f.read().strip() if script_content: num_speakers = self._infer_num_speakers_from_script(script_content) self.example_scripts.append([num_speakers, script_content]) except Exception as e: print(f"Error loading {txt_file}: {e}") def convert_to_16_bit_wav(data): if torch.is_tensor(data): data = data.detach().cpu().numpy() data = np.array(data) if np.max(np.abs(data)) > 1.0: data = data / np.max(np.abs(data)) return (data * 32767).astype(np.int16) def create_demo_interface(demo_instance: VibeVoiceDemo): custom_css = """ """ 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: gr.HTML("""

🎙️ Vibe Podcasting

Generating Long-form Multi-speaker AI Podcast with VibeVoice

""") with gr.Row(): with gr.Column(scale=1, elem_classes="settings-card"): gr.Markdown("### 🎛️ Podcast Settings") # NEW - model dropdown model_dropdown = gr.Dropdown( choices=list(demo_instance.models.keys()), value=demo_instance.current_model_name, label="Model", ) num_speakers = gr.Slider( minimum=1, maximum=4, value=2, step=1, label="Number of Speakers", elem_classes="slider-container" ) gr.Markdown("### 🎭 Speaker Selection") available_speaker_names = list(demo_instance.available_voices.keys()) 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), elem_classes="speaker-item" ) speaker_selections.append(speaker) gr.Markdown("### ⚙️ Advanced 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)", elem_classes="slider-container" ) 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...", lines=12, max_lines=20, elem_classes="script-input" ) with gr.Row(): random_example_btn = gr.Button( "🎲 Random Example", size="lg", variant="secondary", elem_classes="random-btn", scale=1 ) generate_btn = gr.Button( "🚀 Generate Podcast", size="lg", variant="primary", elem_classes="generate-btn", scale=2 ) gr.Markdown("### 🎵 Generated Podcast") complete_audio_output = gr.Audio( label="Complete Podcast (Download)", type="numpy", elem_classes="audio-output complete-audio-section", autoplay=False, show_download_button=True, visible=True ) log_output = gr.Textbox( label="Generation Log", lines=8, max_lines=15, interactive=False, elem_classes="log-output" ) def update_speaker_visibility(num_speakers): return [gr.update(visible=(i < num_speakers)) for i in range(4)] num_speakers.change( fn=update_speaker_visibility, inputs=[num_speakers], outputs=speaker_selections ) def generate_podcast_wrapper(model_choice, num_speakers, script, *speakers_and_params): try: speakers = speakers_and_params[:4] cfg_scale_val = speakers_and_params[4] audio, log = demo_instance.generate_podcast( 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_val, model_name=model_choice ) return audio, log except Exception as e: traceback.print_exc() return None, f"❌ Error: {str(e)}" generate_btn.click( fn=generate_podcast_wrapper, inputs=[model_dropdown, num_speakers, script_input] + speaker_selections + [cfg_scale], outputs=[complete_audio_output, log_output], queue=True ) def load_random_example(): import random examples = getattr(demo_instance, "example_scripts", []) if not examples: examples = [ [2, "Speaker 0: Welcome to our AI podcast demo!\nSpeaker 1: Thanks, excited to be here!"] ] num_speakers_value, script_value = random.choice(examples) return num_speakers_value, script_value random_example_btn.click( fn=load_random_example, inputs=[], outputs=[num_speakers, script_input], queue=False ) gr.Markdown("### 📚 Example Scripts") examples = getattr(demo_instance, "example_scripts", []) or [ [1, "Speaker 1: Welcome to our AI podcast demo. This is a sample script."] ] gr.Examples( examples=examples, inputs=[num_speakers, script_input], label="Try these example scripts:" ) return interface def run_demo( model_paths: dict = None, device: str = "cuda", inference_steps: int = 5, share: bool = True, ): """ model_paths default includes two entries. Replace paths as needed. """ if model_paths is None: model_paths = { "VibeVoice-Large": "microsoft/VibeVoice-Large", "VibeVoice-1.5B": "microsoft/VibeVoice-1.5B" } set_seed(42) demo_instance = VibeVoiceDemo(model_paths, device, inference_steps) interface = create_demo_interface(demo_instance) interface.queue().launch( share=share, server_name="0.0.0.0" if share else "127.0.0.1", show_error=True, show_api=False ) if __name__ == "__main__": run_demo()