Update app.py

app.py

@@ -1,124 +1,61 @@
-"""
-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
-import spaces
+from spaces import GPU
 
 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__)
 
-# import os
-# os.environ["FLASH_ATTENTION_2"] = "0"
-
 
 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.is_generating = False
+        self.processor = None
+        self.model = None
+        self.available_voices = {}
         self.load_model()
         self.setup_voice_presets()
-        self.load_example_scripts()  # Load example scripts
-        
+        self.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.processor = VibeVoiceProcessor.from_pretrained(self.model_path)
         self.model = VibeVoiceForConditionalGenerationInference.from_pretrained(
             self.model_path,
             torch_dtype=torch.bfloat16,
-            device_map='cuda'
+            device_map=self.device
         )
         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
+        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:
-            # 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())}")
-    
+            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:
-        """Read and preprocess audio file."""
         try:
             wav, sr = sf.read(audio_path)
             if len(wav.shape) > 1:
@@ -129,14 +66,13 @@ class VibeVoiceDemo:
         except Exception as e:
             print(f"Error reading audio {audio_path}: {e}")
             return np.array([])
-    
-    @spaces.GPU
+
+    @GPU
     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):
-        """Single GPU function for full generation (streaming + final)."""
-        self.stop_generation = False
+        """Final audio generation only (no streaming)."""
         self.is_generating = True
 
         if not script.strip():
@@ -150,7 +86,6 @@ class VibeVoiceDemo:
             if not sp or sp not in self.available_voices:
                 raise gr.Error(f"Invalid speaker {i+1} selection.")
 
-        # load voices
         voice_samples = [self.read_audio(self.available_voices[sp]) for sp in selected]
         if any(len(v) == 0 for v in voice_samples):
             raise gr.Error("Failed to load one or more voice samples.")
@@ -177,539 +112,103 @@ class VibeVoiceDemo:
             return_tensors="pt"
         )
 
-        # === direct generation with streamer ===
-        from vibevoice import AudioStreamer, convert_to_16_bit_wav
-        audio_streamer = AudioStreamer(batch_size=1)
         start = time.time()
         outputs = self.model.generate(
             **inputs,
             cfg_scale=cfg_scale,
             tokenizer=self.processor.tokenizer,
-            audio_streamer=audio_streamer,
             verbose=False
         )
 
-        sample_rate = 24000
-        audio_stream = audio_streamer.get_stream(0)
-        all_chunks, pending = [], []
-        min_chunk_size = sample_rate * 2
-        last_yield = time.time()
-
-        for chunk in audio_stream:
-            if torch.is_tensor(chunk):
-                chunk = chunk.float().cpu().numpy()
-            if chunk.ndim > 1:
-                chunk = chunk.squeeze()
-            chunk16 = convert_to_16_bit_wav(chunk)
-            all_chunks.append(chunk16)
-            pending.append(chunk16)
-            if sum(len(c) for c in pending) >= min_chunk_size or (time.time() - last_yield) > 5:
-                new_audio = np.concatenate(pending)
-                yield (sample_rate, new_audio), None, f"Streaming {len(all_chunks)} chunks..."
-                pending = []
-                last_yield = time.time()
-
-        if all_chunks:
-            complete = np.concatenate(all_chunks)
-            total_dur = len(complete) / sample_rate
-            log = f"✅ Generation complete in {time.time()-start:.1f}s, {total_dur:.1f}s audio"
-            yield None, (sample_rate, complete), log
+        # Extract audio
+        if isinstance(outputs, dict) and "audio" in outputs:
+            audio = outputs["audio"]
         else:
-            yield None, None, "❌ No audio generated."
+            audio = outputs
+
+        if torch.is_tensor(audio):
+            audio = audio.float().cpu().numpy()
+        if audio.ndim > 1:
+            audio = audio.squeeze()
+
+        sample_rate = 24000
+        audio16 = convert_to_16_bit_wav(audio)
+
+        total_dur = len(audio16) / sample_rate
+        log = f"✅ Generation complete in {time.time()-start:.1f}s, {total_dur:.1f}s audio"
 
         self.is_generating = False
-    
+        return (sample_rate, audio16), log
 
-    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))])
-        
+        txt_files = sorted([f for f in os.listdir(examples_dir)
+                            if f.lower().endswith('.txt')])
         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:
+                with open(os.path.join(examples_dir, txt_file), '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")
-                
+                if script_content:
+                    self.example_scripts.append([1, script_content])
             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)
-    
+                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):
-    """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",
-        )
+        theme=gr.themes.Soft(primary_hue="blue", secondary_hue="purple")
     ) 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
+
+        gr.Markdown("## 🎙️ VibeVoice Podcast Generator (Final Audio Only)")
+
+        num_speakers = gr.Slider(1, 4, value=2, step=1, label="Number of Speakers")
+        available_speaker_names = list(demo_instance.available_voices.keys())
+        default_speakers = available_speaker_names[:4]
+
+        speaker_selections = []
+        for i in range(4):
+            speaker = gr.Dropdown(
+                choices=available_speaker_names,
+                value=default_speakers[i] if i < len(default_speakers) else None,
+                label=f"Speaker {i+1}",
+                visible=(i < 2)
+            )
+            speaker_selections.append(speaker)
+
+        cfg_scale = gr.Slider(1.0, 2.0, value=1.3, step=0.05, label="CFG Scale")
+
+        script_input = gr.Textbox(
+            label="Podcast Script",
+            placeholder="Enter your script here...",
+            lines=10
+        )
+
+        generate_btn = gr.Button("🚀 Generate Podcast")
+        audio_output = gr.Audio(
+            label="Generated Podcast (Download)",
+            type="numpy",
+            show_download_button=True
         )
-        
-        # Main generation function with streaming
+        log_output = gr.Textbox(label="Log", interactive=False, lines=5)
+
         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(
+                speakers = speakers_and_params[:4]
+                cfg_scale = speakers_and_params[4]
+                audio, log = demo_instance.generate_podcast(
                     num_speakers=int(num_speakers),
                     script=script,
                     speaker_1=speakers[0],
@@ -717,280 +216,37 @@ Or paste text directly and it will auto-assign speakers.""",
                     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)
-
+                )
+                return audio, log
             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
+                return None, f"❌ Error: {str(e)}"
+
         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
+            outputs=[audio_output, log_output]
         )
-        
-        # 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
 
+        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
 
 def run_demo(
     model_path: str = "microsoft/VibeVoice-1.5B",
     device: str = "cuda",
     inference_steps: int = 5,
     share: bool = True,
-    ) -> None:
-    """
-    Run the VibeVoice demo without any command-line arguments.
-    - share=True exposes the app publicly via a Gradio share link.
-    - Default Gradio port (7860) is used automatically.
-    - Errors are shown to help with debugging.
-    """
+):
     set_seed(42)
-
-    print("🎙️ Initializing VibeVoice Demo with Streaming Support...")
-
-    # Initialize demo instance
-    demo_instance = VibeVoiceDemo(
-        model_path=model_path,
-        device=device,
-        inference_steps=inference_steps
-    )
-
-    # Build UI
+    demo_instance = VibeVoiceDemo(model_path, device, inference_steps)
     interface = create_demo_interface(demo_instance)
-
-    # Info
-    print("🚀 Launching demo")
-    print(f"📁 Model path: {model_path}")
-    print(f"🎭 Available voices: {len(getattr(demo_instance, 'available_voices', []))}")
-    print(f"🔴 Streaming mode: ENABLED")
-    print(f"🔒 Session isolation: ENABLED")
-
-    # Launch (no server_port specified → default 7860)
-    try:
-        interface.queue(
-            max_size=20,
-            default_concurrency_limit=1
-        ).launch(
-            share=share,
-            server_name="0.0.0.0" if share else "127.0.0.1",
-            show_error=True,     # show full tracebacks (debug-friendly)
-            show_api=False       # cleaner interface
-        )
-    except KeyboardInterrupt:
-        print("\n🛑 Shutting down gracefully...")
-    except Exception as e:
-        print(f"❌ Server error: {e}")
-        raise
+    interface.queue().launch(
+        share=share,
+        server_name="0.0.0.0" if share else "127.0.0.1",
+        show_error=True,
+        show_api=False
+    )
 
 
-# Run automatically when this file is executed (no CLI needed)
 if __name__ == "__main__":
     run_demo()
-
-