import os import sys import torch import spaces import numpy as np import soundfile as sf import librosa import logging import gradio as gr import tempfile from typing import Dict, Optional, List # --- 1. Setup Environment --- # Add the project root to the Python path to allow importing local modules project_root = os.path.dirname(os.path.abspath(__file__)) if project_root not in sys.path: sys.path.insert(0, project_root) # Configure logging to see VibeVoice messages logging.basicConfig(level=logging.INFO, format='[%(name)s] %(message)s') logger = logging.getLogger("VibeVoiceGradio") # Mock ComfyUI's folder_paths module for model caching class MockFolderPaths: def get_folder_paths(self, folder_name): if folder_name == "checkpoints": models_dir = os.path.join(project_root, "models") os.makedirs(models_dir, exist_ok=True) return [models_dir] return [] sys.modules['folder_paths'] = MockFolderPaths() # Import the node class after setting up the environment # We use MultiSpeakerNode as it can handle single-speaker text too. from nodes.multi_speaker_node import VibeVoiceMultipleSpeakersNode # --- 2. Load Model Globally --- logger.info("Initializing VibeVoice node...") # We use the multi-speaker node as it can handle single-speaker cases gracefully. # This instance will hold the model in memory for all Gradio calls. vibevoice_node = VibeVoiceMultipleSpeakersNode() try: logger.info("Loading VibeVoice-Large model. This may take a while on the first run...") # Pre-load the model into the node instance. vibevoice_node.load_model( model_name='VibeVoice-Large', model_path='aoi-ot/VibeVoice-Large', attention_type='auto' ) logger.info("VibeVoice-Large model loaded successfully!") except Exception as e: logger.error(f"Failed to load the model: {e}") logger.error("Please ensure you have an internet connection for the first run and sufficient VRAM.") sys.exit(1) # --- 3. Helper Functions --- def load_audio_for_node(file_path: Optional[str]) -> Optional[Dict]: """Loads an audio file from a path and formats it for the VibeVoice node.""" if file_path is None: return None try: waveform, sr = librosa.load(file_path, sr=24000, mono=True) waveform_tensor = torch.from_numpy(waveform).float().unsqueeze(0).unsqueeze(0) return {"waveform": waveform_tensor, "sample_rate": 24000} except Exception as e: logger.error(f"Failed to load audio file {file_path}: {e}") return None def save_audio_to_tempfile(audio_dict: Dict) -> Optional[str]: """Saves the node's audio output to a temporary WAV file for Gradio.""" if not audio_dict or "waveform" not in audio_dict: logger.error("Invalid audio dictionary received from node.") return None waveform_tensor = audio_dict["waveform"] sample_rate = audio_dict["sample_rate"] waveform_np = waveform_tensor.squeeze().cpu().numpy() # Create a temporary file with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmpfile: sf.write(tmpfile.name, waveform_np, sample_rate) return tmpfile.name # --- 4. Gradio Core Logic --- @spaces.GPU def generate_speech_gradio( text: str, speaker1_audio_path: Optional[str], speaker2_audio_path: Optional[str], speaker3_audio_path: Optional[str], speaker4_audio_path: Optional[str], seed: int, diffusion_steps: int, cfg_scale: float, use_sampling: bool, temperature: float, top_p: float, progress=gr.Progress(track_tqdm=True) ): """The main function that Gradio will call to generate speech.""" if not text or not text.strip(): raise gr.Error("Please provide some text to generate.") progress(0, desc="Processing audio inputs...") logger.info("Processing user inputs...") # Load uploaded voices speaker_voices = [ load_audio_for_node(speaker1_audio_path), load_audio_for_node(speaker2_audio_path), load_audio_for_node(speaker3_audio_path), load_audio_for_node(speaker4_audio_path), ] progress(0.2, desc="Generating speech... (this can take a moment)") logger.info("Calling VibeVoice model to generate speech...") try: # Call the generate_speech method on our globally loaded node audio_output_tuple = vibevoice_node.generate_speech( text=text, model='VibeVoice-Large', attention_type='auto', free_memory_after_generate=False, # Keep model in memory for next call diffusion_steps=int(diffusion_steps), seed=int(seed), cfg_scale=cfg_scale, use_sampling=use_sampling, speaker1_voice=speaker_voices[0], speaker2_voice=speaker_voices[1], speaker3_voice=speaker_voices[2], speaker4_voice=speaker_voices[3], temperature=temperature, top_p=top_p ) except Exception as e: logger.error(f"Error during speech generation: {e}") raise gr.Error(f"An error occurred during generation: {e}") progress(0.9, desc="Saving audio file...") logger.info("Generation complete. Saving audio output.") # Save the output to a temporary file for Gradio to serve output_audio_path = save_audio_to_tempfile(audio_output_tuple[0]) if output_audio_path is None: raise gr.Error("Failed to process the generated audio.") return output_audio_path # --- 5. Gradio UI Layout --- with gr.Blocks(theme=gr.themes.Soft()) as demo: gr.Markdown( "# VibeVoice Text-to-Speech Demo\n" "Generate multi-speaker conversations with optional voice cloning using Microsoft's VibeVoice-Large model." ) with gr.Row(): with gr.Column(scale=2): text_input = gr.Textbox( label="Text Input", placeholder=( "Enter text using speaker tags like [1]:, [2]:, etc.\n\n" "[1]: Hello, I'm the first speaker.\n" "[2]: Hi there, I'm the second! How are you?\n" "[1]: I'm doing great, thanks for asking!" ), lines=8, max_lines=20 ) with gr.Accordion("Upload Speaker Voices (Optional)", open=False): gr.Markdown("Upload a short audio clip (3-30 seconds, clear audio) for each speaker you want to clone.") with gr.Row(): speaker1_audio = gr.Audio(label="Speaker 1 Voice", type="filepath") speaker2_audio = gr.Audio(label="Speaker 2 Voice", type="filepath") with gr.Row(): speaker3_audio = gr.Audio(label="Speaker 3 Voice", type="filepath") speaker4_audio = gr.Audio(label="Speaker 4 Voice", type="filepath") with gr.Accordion("Advanced Options", open=False): seed = gr.Slider(label="Seed", minimum=0, maximum=2**32-1, step=1, value=42, interactive=True) diffusion_steps = gr.Slider(label="Diffusion Steps", minimum=5, maximum=100, step=1, value=20, interactive=True, info="More steps = better quality, but slower.") cfg_scale = gr.Slider(label="CFG Scale", minimum=0.5, maximum=3.5, step=0.05, value=1.3, interactive=True, info="Guidance scale.") use_sampling = gr.Checkbox(label="Use Sampling", value=False, interactive=True, info="Enable for more varied, less deterministic output.") temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=2.0, step=0.05, value=0.95, interactive=True, info="Only used when sampling is enabled.") top_p = gr.Slider(label="Top P", minimum=0.1, maximum=1.0, step=0.05, value=0.95, interactive=True, info="Only used when sampling is enabled.") with gr.Column(scale=1): generate_button = gr.Button("Generate Speech", variant="primary") audio_output = gr.Audio(label="Generated Speech", type="filepath", interactive=False) inputs = [ text_input, speaker1_audio, speaker2_audio, speaker3_audio, speaker4_audio, seed, diffusion_steps, cfg_scale, use_sampling, temperature, top_p ] generate_button.click( fn=generate_speech_gradio, inputs=inputs, outputs=audio_output ) if __name__ == "__main__": # Launch the Gradio app demo.launch(share=True) # Add share=True to create a public link: demo.launch(share=True)