Update app.py

app.py

@@ -1,7 +1,9 @@
 import os
 import re
+import logging
 import torch
 import tempfile
+from typing import Tuple, Union
 from scipy.io.wavfile import write
 from pydub import AudioSegment
 from dotenv import load_dotenv
@@ -19,102 +21,149 @@ from transformers import (
 # Coqui TTS
 from TTS.api import TTS
 
+# Kokoro TTS (ensure these are installed)
+# pip install -q kokoro>=0.8.2 soundfile
+# apt-get -qq -y install espeak-ng > /dev/null 2>&1
+from kokoro import KPipeline
+import soundfile as sf
+
 # ---------------------------------------------------------------------
-# Load Environment Variables
+# Configuration & Logging Setup
 # ---------------------------------------------------------------------
 load_dotenv()
 HF_TOKEN = os.getenv("HF_TOKEN")
+if not HF_TOKEN:
+    logging.warning("HF_TOKEN environment variable not set!")
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
 
-# ---------------------------------------------------------------------
 # Global Model Caches
-# ---------------------------------------------------------------------
 LLAMA_PIPELINES = {}
 MUSICGEN_MODELS = {}
 TTS_MODELS = {}
 
 # ---------------------------------------------------------------------
-# Utility Function: Clean Text
+# Utility Functions
 # ---------------------------------------------------------------------
 def clean_text(text: str) -> str:
     """
-    Removes undesired characters (e.g., asterisks) that might not be recognized by the model's vocabulary.
+    Clean text by removing undesired characters.
+    
+    Args:
+        text (str): Input text to be cleaned.
+        
+    Returns:
+        str: Cleaned text.
     """
-    # Remove all asterisks. You can add more cleaning steps here as needed.
+    # Remove all asterisks. Additional cleaning rules can be added.
     return re.sub(r'\*', '', text)
 
 # ---------------------------------------------------------------------
-# Helper Functions
+# Model Loading Helper Functions
 # ---------------------------------------------------------------------
-def get_llama_pipeline(model_id: str, token: str):
+def get_llama_pipeline(model_id: str, token: str) -> pipeline:
     """
-    Returns a cached LLaMA pipeline if available; otherwise, loads it.
+    Load and cache the LLaMA text-generation pipeline.
+    
+    Args:
+        model_id (str): Hugging Face model identifier.
+        token (str): Hugging Face authentication token.
+    
+    Returns:
+        pipeline: Text-generation pipeline instance.
     """
     if model_id in LLAMA_PIPELINES:
         return LLAMA_PIPELINES[model_id]
 
-    tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=token)
-    model = AutoModelForCausalLM.from_pretrained(
-        model_id,
-        use_auth_token=token,
-        torch_dtype=torch.float16,
-        device_map="auto",
-        trust_remote_code=True,
-    )
-    text_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
-    LLAMA_PIPELINES[model_id] = text_pipeline
-    return text_pipeline
-
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=token)
+        model = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            use_auth_token=token,
+            torch_dtype=torch.float16,
+            device_map="auto",
+            trust_remote_code=True,
+        )
+        text_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
+        LLAMA_PIPELINES[model_id] = text_pipeline
+        return text_pipeline
+    except Exception as e:
+        logging.error(f"Error loading LLaMA pipeline: {e}")
+        raise
 
-def get_musicgen_model(model_key: str = "facebook/musicgen-large"):
+def get_musicgen_model(model_key: str = "facebook/musicgen-large") -> Tuple[MusicgenForConditionalGeneration, AutoProcessor]:
     """
-    Returns a cached MusicGen model if available; otherwise, loads it.
-    Uses the 'large' variant for higher quality outputs.
+    Load and cache the MusicGen model and its processor.
+    
+    Args:
+        model_key (str): Model key (default uses 'facebook/musicgen-large').
+    
+    Returns:
+        tuple: (MusicGen model, processor)
     """
     if model_key in MUSICGEN_MODELS:
         return MUSICGEN_MODELS[model_key]
 
-    model = MusicgenForConditionalGeneration.from_pretrained(model_key)
-    processor = AutoProcessor.from_pretrained(model_key)
-
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    model.to(device)
-    MUSICGEN_MODELS[model_key] = (model, processor)
-    return model, processor
-
+    try:
+        model = MusicgenForConditionalGeneration.from_pretrained(model_key)
+        processor = AutoProcessor.from_pretrained(model_key)
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        model.to(device)
+        MUSICGEN_MODELS[model_key] = (model, processor)
+        return model, processor
+    except Exception as e:
+        logging.error(f"Error loading MusicGen model: {e}")
+        raise
 
-def get_tts_model(model_name: str = "tts_models/en/ljspeech/tacotron2-DDC"):
+def get_tts_model(model_name: str = "tts_models/en/ljspeech/tacotron2-DDC") -> TTS:
     """
-    Returns a cached TTS model if available; otherwise, loads it.
+    Load and cache the TTS model.
+    
+    Args:
+        model_name (str): Name of the TTS model.
+    
+    Returns:
+        TTS: TTS model instance.
     """
     if model_name in TTS_MODELS:
         return TTS_MODELS[model_name]
 
-    tts_model = TTS(model_name)
-    TTS_MODELS[model_name] = tts_model
-    return tts_model
-
+    try:
+        tts_model = TTS(model_name)
+        TTS_MODELS[model_name] = tts_model
+        return tts_model
+    except Exception as e:
+        logging.error(f"Error loading TTS model: {e}")
+        raise
 
 # ---------------------------------------------------------------------
 # Script Generation Function
 # ---------------------------------------------------------------------
 @spaces.GPU(duration=100)
-def generate_script(user_prompt: str, model_id: str, token: str, duration: int):
+def generate_script(user_prompt: str, model_id: str, token: str, duration: int) -> Tuple[str, str, str]:
     """
-    Generates a script, sound design suggestions, and music ideas from a user prompt.
-    Returns a tuple of strings: (voice_script, sound_design, music_suggestions).
+    Generate a script, sound design suggestions, and music ideas from a user prompt.
+    
+    Args:
+        user_prompt (str): The user's creative input.
+        model_id (str): Hugging Face model identifier for LLaMA.
+        token (str): Hugging Face authentication token.
+        duration (int): Desired duration of the promo in seconds.
+    
+    Returns:
+        tuple: (voice_script, sound_design, music_suggestions)
     """
     try:
         text_pipeline = get_llama_pipeline(model_id, token)
-
         system_prompt = (
             "You are an expert radio imaging producer specializing in sound design and music. "
-            f"Based on the user's concept and the selected duration of {duration} seconds, produce the following: "
-            "1. A concise voice-over script. Prefix this section with 'Voice-Over Script:'.\n"
-            "2. Suggestions for sound design. Prefix this section with 'Sound Design Suggestions:'.\n"
-            "3. Music styles or track recommendations. Prefix this section with 'Music Suggestions:'."
+            f"Based on the user's concept and the selected duration of {duration} seconds, produce the following:\n"
+            "1. A concise voice-over script. Prefix this section with 'Voice-Over Script:'\n"
+            "2. Suggestions for sound design. Prefix this section with 'Sound Design Suggestions:'\n"
+            "3. Music styles or track recommendations. Prefix this section with 'Music Suggestions:'"
         )
         combined_prompt = f"{system_prompt}\nUser concept: {user_prompt}\nOutput:"
-
+        
         with torch.inference_mode():
             result = text_pipeline(
                 combined_prompt,
@@ -124,86 +173,130 @@ def generate_script(user_prompt: str, model_id: str, token: str, duration: int):
             )
 
         generated_text = result[0]["generated_text"]
+        # Remove everything before the 'Output:' marker if present
         if "Output:" in generated_text:
             generated_text = generated_text.split("Output:")[-1].strip()
 
-        # Default placeholders
+        # Initialize default outputs
         voice_script = "No voice-over script found."
         sound_design = "No sound design suggestions found."
         music_suggestions = "No music suggestions found."
 
-        # Voice-Over Script
+        # Parse generated text based on expected prefixes
         if "Voice-Over Script:" in generated_text:
-            parts = generated_text.split("Voice-Over Script:")
-            voice_script_part = parts[1]
-            if "Sound Design Suggestions:" in voice_script_part:
-                voice_script = voice_script_part.split("Sound Design Suggestions:")[0].strip()
+            voice_section = generated_text.split("Voice-Over Script:")[1]
+            if "Sound Design Suggestions:" in voice_section:
+                voice_script = voice_section.split("Sound Design Suggestions:")[0].strip()
             else:
-                voice_script = voice_script_part.strip()
+                voice_script = voice_section.strip()
 
-        # Sound Design
         if "Sound Design Suggestions:" in generated_text:
-            parts = generated_text.split("Sound Design Suggestions:")
-            sound_design_part = parts[1]
-            if "Music Suggestions:" in sound_design_part:
-                sound_design = sound_design_part.split("Music Suggestions:")[0].strip()
+            sound_section = generated_text.split("Sound Design Suggestions:")[1]
+            if "Music Suggestions:" in sound_section:
+                sound_design = sound_section.split("Music Suggestions:")[0].strip()
             else:
-                sound_design = sound_design_part.strip()
+                sound_design = sound_section.strip()
 
-        # Music Suggestions
         if "Music Suggestions:" in generated_text:
-            parts = generated_text.split("Music Suggestions:")
-            music_suggestions = parts[1].strip()
+            music_suggestions = generated_text.split("Music Suggestions:")[-1].strip()
 
         return voice_script, sound_design, music_suggestions
 
     except Exception as e:
+        logging.error(f"Error in generate_script: {e}")
         return f"Error generating script: {e}", "", ""
 
-
 # ---------------------------------------------------------------------
-# Voice-Over Generation Function
+# Voice-Over Generation Functions
 # ---------------------------------------------------------------------
 @spaces.GPU(duration=100)
-def generate_voice(script: str, tts_model_name: str = "tts_models/en/ljspeech/tacotron2-DDC"):
+def generate_voice(script: str, tts_model_name: str = "tts_models/en/ljspeech/tacotron2-DDC") -> Union[str, None]:
     """
-    Generates a voice-over from the provided script using the Coqui TTS model.
-    Returns the file path to the generated .wav file.
+    Generate a voice-over audio file using Coqui TTS from the provided script.
+    
+    Args:
+        script (str): The voice-over script.
+        tts_model_name (str): TTS model identifier.
+    
+    Returns:
+        str: File path to the generated .wav file or an error message.
     """
     try:
         if not script.strip():
-            return "Error: No script provided."
-
-        # Clean the script to remove special characters (e.g., asterisks) that may produce warnings
+            raise ValueError("No script provided.")
         cleaned_script = clean_text(script)
-
         tts_model = get_tts_model(tts_model_name)
-
-        # Generate and save voice
-        output_path = os.path.join(tempfile.gettempdir(), "voice_over.wav")
+        output_path = os.path.join(tempfile.gettempdir(), "voice_over_coqui.wav")
         tts_model.tts_to_file(text=cleaned_script, file_path=output_path)
+        logging.info(f"Coqui voice-over generated at {output_path}")
         return output_path
 
     except Exception as e:
+        logging.error(f"Error in generate_voice (Coqui TTS): {e}")
         return f"Error generating voice: {e}"
 
+@spaces.GPU(duration=100)
+def generate_voice_kokoro(script: str, lang_code: str = 'a', voice: str = 'af_heart', speed: float = 1.0) -> Union[str, None]:
+    """
+    Generate a voice-over audio file using the Kokoro TTS model.
+    
+    Args:
+        script (str): The text to synthesize.
+        lang_code (str): Language code ('a' for American English, etc.).
+        voice (str): Specific voice style.
+        speed (float): Speech speed.
+    
+    Returns:
+        str: File path to the generated WAV file or an error message.
+    """
+    try:
+        # Initialize the Kokoro pipeline
+        kp = KPipeline(lang_code=lang_code)
+        audio_segments = []
+        generator = kp(script, voice=voice, speed=speed, split_pattern=r'\n+')
+        for i, (gs, ps, audio) in enumerate(generator):
+            audio_segments.append(audio)
+
+        # Join audio segments using pydub
+        combined = AudioSegment.empty()
+        for seg in audio_segments:
+            segment = AudioSegment(
+                seg.tobytes(),
+                frame_rate=24000,
+                sample_width=seg.dtype.itemsize,
+                channels=1
+            )
+            combined += segment
+
+        output_path = os.path.join(tempfile.gettempdir(), "voice_over_kokoro.wav")
+        combined.export(output_path, format="wav")
+        logging.info(f"Kokoro voice-over generated at {output_path}")
+        return output_path
+
+    except Exception as e:
+        logging.error(f"Error in generate_voice_kokoro: {e}")
+        return f"Error generating Kokoro voice: {e}"
 
 # ---------------------------------------------------------------------
 # Music Generation Function
 # ---------------------------------------------------------------------
 @spaces.GPU(duration=200)
-def generate_music(prompt: str, audio_length: int):
+def generate_music(prompt: str, audio_length: int) -> Union[str, None]:
     """
-    Generates music from the 'facebook/musicgen-large' model based on the prompt.
-    Returns the file path to the generated .wav file.
+    Generate music based on the prompt using MusicGen.
+    
+    Args:
+        prompt (str): Music prompt or style suggestion.
+        audio_length (int): Length parameter (number of tokens).
+    
+    Returns:
+        str: File path to the generated .wav file or an error message.
     """
     try:
         if not prompt.strip():
-            return "Error: No music suggestion provided."
-
+            raise ValueError("No music suggestion provided.")
         model_key = "facebook/musicgen-large"
         musicgen_model, musicgen_processor = get_musicgen_model(model_key)
-
         device = "cuda" if torch.cuda.is_available() else "cpu"
         inputs = musicgen_processor(text=[prompt], padding=True, return_tensors="pt").to(device)
 
@@ -212,48 +305,47 @@ def generate_music(prompt: str, audio_length: int):
 
         audio_data = outputs[0, 0].cpu().numpy()
         normalized_audio = (audio_data / max(abs(audio_data)) * 32767).astype("int16")
-
-        output_path = os.path.join(tempfile.gettempdir(), "musicgen_large_generated_music.wav")
+        output_path = os.path.join(tempfile.gettempdir(), "musicgen_generated_music.wav")
         write(output_path, 44100, normalized_audio)
-
+        logging.info(f"Music generated at {output_path}")
         return output_path
 
     except Exception as e:
+        logging.error(f"Error in generate_music: {e}")
         return f"Error generating music: {e}"
 
-
 # ---------------------------------------------------------------------
-# Audio Blending with Duration Sync & Ducking
+# Audio Blending Function
 # ---------------------------------------------------------------------
 @spaces.GPU(duration=100)
-def blend_audio(voice_path: str, music_path: str, ducking: bool, duck_level: int = 10):
+def blend_audio(voice_path: str, music_path: str, ducking: bool, duck_level: int = 10) -> Union[str, None]:
     """
-    Blends two audio files (voice and music).
-    1. If music < voice, loops the music until it meets/exceeds the voice duration.
-    2. If music > voice, trims music to the voice duration.
-    3. If ducking=True, the music is attenuated by 'duck_level' dB while the voice is playing.
-    Returns the file path to the blended .wav file.
+    Blend voice and music audio files with optional ducking.
+    
+    Args:
+        voice_path (str): File path to the voice audio.
+        music_path (str): File path to the music audio.
+        ducking (bool): If True, attenuate music during voice segments.
+        duck_level (int): Attenuation level in dB.
+    
+    Returns:
+        str: File path to the blended .wav file or an error message.
     """
     try:
-        if not os.path.isfile(voice_path) or not os.path.isfile(music_path):
-            return "Error: Missing audio files for blending."
+        if not (os.path.isfile(voice_path) and os.path.isfile(music_path)):
+            raise FileNotFoundError("Missing audio files for blending.")
 
         voice = AudioSegment.from_wav(voice_path)
         music = AudioSegment.from_wav(music_path)
+        voice_duration = len(voice)
 
-        voice_len = len(voice)  # in milliseconds
-        music_len = len(music)  # in milliseconds
-
-        # Loop music if it's shorter than the voice
-        if music_len < voice_len:
+        if len(music) < voice_duration:
             looped_music = AudioSegment.empty()
-            while len(looped_music) < voice_len:
+            while len(looped_music) < voice_duration:
                 looped_music += music
             music = looped_music
-
-        # Trim music if it's longer than the voice
-        if len(music) > voice_len:
-            music = music[:voice_len]
+        else:
+            music = music[:voice_duration]
 
         if ducking:
             ducked_music = music - duck_level
@@ -263,12 +355,13 @@ def blend_audio(voice_path: str, music_path: str, ducking: bool, duck_level: int
 
         output_path = os.path.join(tempfile.gettempdir(), "blended_output.wav")
         final_audio.export(output_path, format="wav")
+        logging.info(f"Audio blended at {output_path}")
         return output_path
 
     except Exception as e:
+        logging.error(f"Error in blend_audio: {e}")
         return f"Error blending audio: {e}"
 
-
 # ---------------------------------------------------------------------
 # Gradio Interface with Enhanced UI
 # ---------------------------------------------------------------------
@@ -322,7 +415,7 @@ with gr.Blocks(css="""
     Welcome to **AI Promo Studio**! This platform leverages state-of-the-art AI models to help you generate:
     
     - **Script**: Generate a compelling voice-over script with LLaMA.
-    - **Voice Synthesis**: Create natural-sounding voice-overs using Coqui TTS.
+    - **Voice Synthesis**: Create natural-sounding voice-overs using Coqui TTS or Kokoro TTS.
     - **Music Production**: Produce custom music tracks with MusicGen.
     - **Audio Blending**: Seamlessly blend voice and music with options for ducking.
     """)
@@ -355,20 +448,26 @@ with gr.Blocks(css="""
             music_suggestion_output = gr.Textbox(label="Music Suggestions", lines=3, interactive=False)
 
             generate_script_button.click(
-                fn=lambda user_prompt, model_id, dur: generate_script(user_prompt, model_id, HF_TOKEN, dur),
+                fn=lambda prompt, model, dur: generate_script(prompt, model, HF_TOKEN, dur),
                 inputs=[user_prompt, llama_model_id, duration],
                 outputs=[script_output, sound_design_output, music_suggestion_output],
             )
 
         # Step 2: Generate Voice
         with gr.Tab("🎤 Voice Synthesis"):
-            gr.Markdown("Generate a natural-sounding voice-over using Coqui TTS.")
+            gr.Markdown("Generate a natural-sounding voice-over. Choose your TTS engine below:")
+            voice_engine = gr.Dropdown(
+                label="TTS Engine",
+                choices=["Coqui TTS", "Kokoro TTS"],
+                value="Coqui TTS",
+                multiselect=False
+            )
             selected_tts_model = gr.Dropdown(
-                label="TTS Model",
+                label="TTS Model / Voice Option",
                 choices=[
-                    "tts_models/en/ljspeech/tacotron2-DDC",  
-                    "tts_models/en/ljspeech/vits", 
-                    "tts_models/en/sam/tacotron-DDC", 
+                    "tts_models/en/ljspeech/tacotron2-DDC",  # Coqui TTS option
+                    "tts_models/en/ljspeech/vits",            # Coqui TTS option
+                    "af_heart"                                # Kokoro TTS voice option
                 ],
                 value="tts_models/en/ljspeech/tacotron2-DDC",
                 multiselect=False
@@ -376,9 +475,18 @@ with gr.Blocks(css="""
             generate_voice_button = gr.Button("Generate Voice-Over", variant="primary")
             voice_audio_output = gr.Audio(label="Voice-Over (WAV)", type="filepath")
 
+            def generate_voice_combined(script, engine, model_choice):
+                if engine == "Coqui TTS":
+                    return generate_voice(script, model_choice)
+                elif engine == "Kokoro TTS":
+                    # For Kokoro, pass the voice option (e.g., "af_heart") and default language code ('a')
+                    return generate_voice_kokoro(script, lang_code='a', voice=model_choice, speed=1.0)
+                else:
+                    return "Error: Unknown TTS engine."
+
             generate_voice_button.click(
-                fn=lambda script, tts_model: generate_voice(script, tts_model),
-                inputs=[script_output, selected_tts_model],
+                fn=generate_voice_combined,
+                inputs=[script_output, voice_engine, selected_tts_model],
                 outputs=voice_audio_output,
             )
 
@@ -397,7 +505,7 @@ with gr.Blocks(css="""
             music_output = gr.Audio(label="Generated Music (WAV)", type="filepath")
 
             generate_music_button.click(
-                fn=lambda music_suggestion, length: generate_music(music_suggestion, length),
+                fn=lambda prompt, length: generate_music(prompt, length),
                 inputs=[music_suggestion_output, audio_length],
                 outputs=[music_output],
             )