app.py

import os
import time
import numpy as np
import gradio as gr
import librosa
import soundfile as sf
import torch
import traceback
from spaces import GPU
from datetime import datetime

from vibevoice.modular.modeling_vibevoice_inference import VibeVoiceForConditionalGenerationInference
from vibevoice.processor.vibevoice_processor import VibeVoiceProcessor
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):
        self.model_path = model_path
        self.device = device
        self.inference_steps = inference_steps
        self.is_generating = False
        self.processor = None
        self.model = None
        self.available_voices = {}
        self.load_model()
        self.setup_voice_presets()
        self.load_example_scripts()

    def load_model(self):
        print(f"Loading processor & model from {self.model_path}")
        self.processor = VibeVoiceProcessor.from_pretrained(self.model_path)
        self.model = VibeVoiceForConditionalGenerationInference.from_pretrained(
            self.model_path,
            torch_dtype=torch.bfloat16,
            device_map=self.device
        )
        self.model.eval()
        self.model.set_ddpm_inference_steps(num_steps=self.inference_steps)

    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
    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):
        """Final audio generation only (no streaming)."""
        self.is_generating = True

        if not script.strip():
            raise gr.Error("Please provide a script.")

        if num_speakers < 1 or num_speakers > 4:
            raise gr.Error("Number of speakers must be 1–4.")

        selected = [speaker_1, speaker_2, speaker_3, speaker_4][:num_speakers]
        for i, sp in enumerate(selected):
            if not sp or sp not in self.available_voices:
                raise gr.Error(f"Invalid speaker {i+1} selection.")

        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.")

        # format script
        lines = script.strip().split("\n")
        formatted = []
        for i, line in enumerate(lines):
            line = line.strip()
            if not line:
                continue
            if line.startswith("Speaker "):
                formatted.append(line)
            else:
                sp_id = i % num_speakers
                formatted.append(f"Speaker {sp_id}: {line}")
        formatted_script = "\n".join(formatted)

        # processor input
        inputs = self.processor(
            text=[formatted_script],
            voice_samples=[voice_samples],
            padding=True,
            return_tensors="pt"
        )

        start = time.time()
        outputs = self.model.generate(
            **inputs,
            cfg_scale=cfg_scale,
            tokenizer=self.processor.tokenizer,
            verbose=False
        )

        # --- handle model output ---
        if hasattr(outputs, "audio"):
            audio = outputs.audio
        elif hasattr(outputs, "audios"):
            audio = outputs.audios[0]
        else:
            raise gr.Error("Model did not return audio in expected format.")

        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)

        # --- save automatically to disk ---
        os.makedirs("outputs", exist_ok=True)
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        file_path = os.path.join("outputs", f"podcast_{timestamp}.wav")
        sf.write(file_path, audio16, sample_rate)
        print(f"💾 Saved podcast to {file_path}")

        total_dur = len(audio16) / sample_rate
        log = f"✅ Generation complete in {time.time()-start:.1f}s, {total_dur:.1f}s audio\nSaved to {file_path}"

        self.is_generating = False
        return (sample_rate, audio16), log

    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:
                    self.example_scripts.append([1, 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):
    with gr.Blocks(
        title="VibeVoice - AI Podcast Generator",
        theme=gr.themes.Soft(primary_hue="blue", secondary_hue="purple")
    ) as interface:

        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
        )
        log_output = gr.Textbox(label="Log", interactive=False, lines=5)

        def generate_podcast_wrapper(num_speakers, script, *speakers_and_params):
            try:
                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],
                    speaker_2=speakers[1],
                    speaker_3=speakers[2],
                    speaker_4=speakers[3],
                    cfg_scale=cfg_scale
                )
                return audio, log
            except Exception as e:
                traceback.print_exc()
                return None, f"❌ Error: {str(e)}"

        generate_btn.click(
            fn=generate_podcast_wrapper,
            inputs=[num_speakers, script_input] + speaker_selections + [cfg_scale],
            outputs=[audio_output, log_output]
        )

        return interface


def run_demo(
    model_path: str = "microsoft/VibeVoice-1.5B",
    device: str = "cuda",
    inference_steps: int = 5,
    share: bool = True,
):
    set_seed(42)
    demo_instance = VibeVoiceDemo(model_path, 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()