vibevoice_integration.py

"""
Integration code for VibeVoice-PodcastCreator to use ACE-Music-Generator

Add this to your VibeVoice space to generate background music on demand.
"""

from gradio_client import Client
import numpy as np
from scipy.io import wavfile
import tempfile
import os


class MusicGenerator:
    """
    Music generator client for VibeVoice podcast creation
    """
    
    def __init__(self, space_name="ACloudCenter/ACE-Music-Generator"):
        """Initialize connection to music generator space"""
        try:
            self.client = Client(space_name)
            self.connected = True
        except:
            print("Warning: Could not connect to music generator. Music features disabled.")
            self.connected = False
    
    def generate_intro_music(self, duration=10):
        """Generate intro music for podcast"""
        if not self.connected:
            return None
            
        return self._generate(
            duration=duration,
            tags="podcast intro, upbeat, electronic, professional, 120 bpm, energetic",
            lyrics="[instrumental]"
        )
    
    def generate_outro_music(self, duration=10):
        """Generate outro music for podcast"""
        if not self.connected:
            return None
            
        return self._generate(
            duration=duration,
            tags="podcast outro, calm, ambient, soft, 80 bpm, fade out, peaceful",
            lyrics="[instrumental]"
        )
    
    def generate_background_music(self, style="ambient", duration=30):
        """
        Generate background music for podcast segments
        
        Styles:
        - ambient: Soft background music
        - news: Professional news-style background
        - dramatic: Intense, dramatic music
        - tech: Futuristic tech music
        - chill: Relaxed lofi music
        """
        styles = {
            "ambient": "ambient, soft, background, minimal, 70 bpm, atmospheric",
            "news": "news, professional, subtle, electronic, 90 bpm, serious",
            "dramatic": "dramatic, orchestral, cinematic, 100 bpm, intense",
            "tech": "electronic, futuristic, synth, 110 bpm, innovative",
            "chill": "lofi, relaxed, warm, 75 bpm, cozy, mellow"
        }
        
        if not self.connected:
            return None
            
        tags = styles.get(style, styles["ambient"])
        return self._generate(duration=duration, tags=tags, lyrics="[instrumental]")
    
    def generate_commercial_jingle(self, duration=5):
        """Generate a short commercial jingle"""
        if not self.connected:
            return None
            
        return self._generate(
            duration=duration,
            tags="jingle, commercial, catchy, upbeat, 140 bpm, happy, memorable",
            lyrics="[instrumental]"
        )
    
    def _generate(self, duration, tags, lyrics):
        """Internal method to generate music"""
        try:
            result = self.client.predict(
                duration,
                tags,
                lyrics,
                60,  # infer_steps
                15.0,  # guidance_scale
                api_name="/generate"
            )
            return result
        except Exception as e:
            print(f"Error generating music: {e}")
            return None
    
    def mix_with_podcast(self, podcast_audio_path, music_path, music_volume=0.2):
        """
        Mix background music with podcast audio
        
        Args:
            podcast_audio_path: Path to podcast audio file
            music_path: Path to music file
            music_volume: Volume of music (0-1, lower = quieter background)
        
        Returns:
            mixed_audio_path: Path to mixed audio file
        """
        try:
            # Load audio files
            podcast_rate, podcast_data = wavfile.read(podcast_audio_path)
            music_rate, music_data = wavfile.read(music_path)
            
            # Ensure same sample rate
            if podcast_rate != music_rate:
                # Simple resampling (you might want to use librosa for better quality)
                music_data = np.interp(
                    np.linspace(0, len(music_data), int(len(music_data) * podcast_rate / music_rate)),
                    np.arange(len(music_data)),
                    music_data
                )
            
            # Match lengths
            if len(music_data) < len(podcast_data):
                # Loop music if it's shorter
                music_data = np.tile(music_data, (len(podcast_data) // len(music_data) + 1))
            music_data = music_data[:len(podcast_data)]
            
            # Mix audio
            mixed = podcast_data + (music_data * music_volume)
            
            # Normalize to prevent clipping
            mixed = np.clip(mixed, -32768, 32767).astype(np.int16)
            
            # Save mixed audio
            output_path = tempfile.mktemp(suffix=".wav")
            wavfile.write(output_path, podcast_rate, mixed)
            
            return output_path
            
        except Exception as e:
            print(f"Error mixing audio: {e}")
            return podcast_audio_path  # Return original if mixing fails


# Example usage in VibeVoice generator
def enhance_podcast_with_music(podcast_generator):
    """
    Example of how to add this to your existing podcast generator
    """
    
    # Initialize music generator
    music_gen = MusicGenerator()
    
    # Your existing podcast generation code
    # podcast_audio = podcast_generator.generate_podcast(...)
    
    # Generate intro music
    intro_music = music_gen.generate_intro_music(duration=5)
    
    # Generate background music for main content
    background_music = music_gen.generate_background_music(
        style="ambient",
        duration=60  # Adjust based on your podcast length
    )
    
    # Generate outro music
    outro_music = music_gen.generate_outro_music(duration=5)
    
    # Mix background music with podcast (optional)
    # if background_music and podcast_audio:
    #     mixed_audio = music_gen.mix_with_podcast(
    #         podcast_audio,
    #         background_music,
    #         music_volume=0.1  # Keep it quiet in background
    #     )
    
    return {
        "intro": intro_music,
        "background": background_music,
        "outro": outro_music
    }


# Quick function to add to your VibeVoice app.py
def add_music_generation_to_vibevoice():
    """
    Add this to your VibeVoice app.py to integrate music generation
    """
    
    # In your create_demo() function, add:
    """
    # Add music generator
    music_gen = MusicGenerator()
    
    # Add checkbox for music generation
    with gr.Row():
        add_intro_music = gr.Checkbox(label="Add Intro Music", value=False)
        add_outro_music = gr.Checkbox(label="Add Outro Music", value=False)
        add_background_music = gr.Checkbox(label="Add Background Music", value=False)
        background_style = gr.Dropdown(
            choices=["ambient", "news", "dramatic", "tech", "chill"],
            value="ambient",
            label="Background Music Style"
        )
    
    # In your generation function:
    def generate_with_music(..., add_intro, add_outro, add_background, bg_style):
        # Your existing generation code
        podcast_audio = generate_podcast(...)
        
        # Add music if requested
        if add_intro:
            intro = music_gen.generate_intro_music(5)
            # Concatenate intro with podcast
        
        if add_background:
            bg_music = music_gen.generate_background_music(bg_style, duration=60)
            # Mix with podcast audio
        
        if add_outro:
            outro = music_gen.generate_outro_music(5)
            # Concatenate outro
        
        return final_audio
    """
    pass


if __name__ == "__main__":
    # Test the music generator
    print("Testing music generator...")
    music_gen = MusicGenerator()
    
    print("Generating intro music...")
    intro = music_gen.generate_intro_music(duration=5)
    if intro:
        print(f"Intro music saved to: {intro}")
    
    print("Generating background music...")
    background = music_gen.generate_background_music(style="ambient", duration=10)
    if background:
        print(f"Background music saved to: {background}")
    
    print("Done!")