Update app.py

app.py

@@ -1,9 +1,7 @@
 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
@@ -21,149 +19,102 @@ 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
-
 # ---------------------------------------------------------------------
-# Configuration & Logging Setup
+# Load Environment Variables
 # ---------------------------------------------------------------------
 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 Functions
+# Utility Function: Clean Text
 # ---------------------------------------------------------------------
 def clean_text(text: str) -> str:
     """
-    Clean text by removing undesired characters.
-    
-    Args:
-        text (str): Input text to be cleaned.
-        
-    Returns:
-        str: Cleaned text.
+    Removes undesired characters (e.g., asterisks) that might not be recognized by the model's vocabulary.
     """
-    # Remove all asterisks. Additional cleaning rules can be added.
+    # Remove all asterisks. You can add more cleaning steps here as needed.
     return re.sub(r'\*', '', text)
 
 # ---------------------------------------------------------------------
-# Model Loading Helper Functions
+# Helper Functions
 # ---------------------------------------------------------------------
-def get_llama_pipeline(model_id: str, token: str) -> pipeline:
+def get_llama_pipeline(model_id: str, token: str):
     """
-    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.
+    Returns a cached LLaMA pipeline if available; otherwise, loads it.
     """
     if model_id in LLAMA_PIPELINES:
         return LLAMA_PIPELINES[model_id]
 
-    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
+    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
 
-def get_musicgen_model(model_key: str = "facebook/musicgen-large") -> Tuple[MusicgenForConditionalGeneration, AutoProcessor]:
+
+def get_musicgen_model(model_key: str = "facebook/musicgen-large"):
     """
-    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)
+    Returns a cached MusicGen model if available; otherwise, loads it.
+    Uses the 'large' variant for higher quality outputs.
     """
     if model_key in MUSICGEN_MODELS:
         return MUSICGEN_MODELS[model_key]
 
-    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
+    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
 
-def get_tts_model(model_name: str = "tts_models/en/ljspeech/tacotron2-DDC") -> TTS:
+
+def get_tts_model(model_name: str = "tts_models/en/ljspeech/tacotron2-DDC"):
     """
-    Load and cache the TTS model.
-    
-    Args:
-        model_name (str): Name of the TTS model.
-    
-    Returns:
-        TTS: TTS model instance.
+    Returns a cached TTS model if available; otherwise, loads it.
     """
     if model_name in TTS_MODELS:
         return TTS_MODELS[model_name]
 
-    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
+    tts_model = TTS(model_name)
+    TTS_MODELS[model_name] = tts_model
+    return tts_model
+
 
 # ---------------------------------------------------------------------
 # Script Generation Function
 # ---------------------------------------------------------------------
 @spaces.GPU(duration=100)
-def generate_script(user_prompt: str, model_id: str, token: str, duration: int) -> Tuple[str, str, str]:
+def generate_script(user_prompt: str, model_id: str, token: str, duration: int):
     """
-    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)
+    Generates a script, sound design suggestions, and music ideas from a user prompt.
+    Returns a tuple of strings: (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:\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:'"
+            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:'."
         )
         combined_prompt = f"{system_prompt}\nUser concept: {user_prompt}\nOutput:"
-        
+
         with torch.inference_mode():
             result = text_pipeline(
                 combined_prompt,
@@ -173,130 +124,86 @@ 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()
 
-        # Initialize default outputs
+        # Default placeholders
         voice_script = "No voice-over script found."
         sound_design = "No sound design suggestions found."
         music_suggestions = "No music suggestions found."
 
-        # Parse generated text based on expected prefixes
+        # Voice-Over Script
         if "Voice-Over Script:" in generated_text:
-            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()
+            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()
             else:
-                voice_script = voice_section.strip()
+                voice_script = voice_script_part.strip()
 
+        # Sound Design
         if "Sound Design Suggestions:" in generated_text:
-            sound_section = generated_text.split("Sound Design Suggestions:")[1]
-            if "Music Suggestions:" in sound_section:
-                sound_design = sound_section.split("Music Suggestions:")[0].strip()
+            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()
             else:
-                sound_design = sound_section.strip()
+                sound_design = sound_design_part.strip()
 
+        # Music Suggestions
         if "Music Suggestions:" in generated_text:
-            music_suggestions = generated_text.split("Music Suggestions:")[-1].strip()
+            parts = generated_text.split("Music Suggestions:")
+            music_suggestions = parts[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 Functions
+# Voice-Over Generation Function
 # ---------------------------------------------------------------------
 @spaces.GPU(duration=100)
-def generate_voice(script: str, tts_model_name: str = "tts_models/en/ljspeech/tacotron2-DDC") -> Union[str, None]:
+def generate_voice(script: str, tts_model_name: str = "tts_models/en/ljspeech/tacotron2-DDC"):
     """
-    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.
+    Generates a voice-over from the provided script using the Coqui TTS model.
+    Returns the file path to the generated .wav file.
     """
     try:
         if not script.strip():
-            raise ValueError("No script provided.")
+            return "Error: No script provided."
+
+        # Clean the script to remove special characters (e.g., asterisks) that may produce warnings
         cleaned_script = clean_text(script)
+
         tts_model = get_tts_model(tts_model_name)
-        output_path = os.path.join(tempfile.gettempdir(), "voice_over_coqui.wav")
+
+        # Generate and save voice
+        output_path = os.path.join(tempfile.gettempdir(), "voice_over.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) -> Union[str, None]:
+def generate_music(prompt: str, audio_length: int):
     """
-    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.
+    Generates music from the 'facebook/musicgen-large' model based on the prompt.
+    Returns the file path to the generated .wav file.
     """
     try:
         if not prompt.strip():
-            raise ValueError("No music suggestion provided.")
+            return "Error: 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)
 
@@ -305,47 +212,48 @@ def generate_music(prompt: str, audio_length: int) -> Union[str, None]:
 
         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_generated_music.wav")
+
+        output_path = os.path.join(tempfile.gettempdir(), "musicgen_large_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 Function
+# Audio Blending with Duration Sync & Ducking
 # ---------------------------------------------------------------------
 @spaces.GPU(duration=100)
-def blend_audio(voice_path: str, music_path: str, ducking: bool, duck_level: int = 10) -> Union[str, None]:
+def blend_audio(voice_path: str, music_path: str, ducking: bool, duck_level: int = 10):
     """
-    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.
+    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.
     """
     try:
-        if not (os.path.isfile(voice_path) and os.path.isfile(music_path)):
-            raise FileNotFoundError("Missing audio files for blending.")
+        if not os.path.isfile(voice_path) or not os.path.isfile(music_path):
+            return "Error: Missing audio files for blending."
 
         voice = AudioSegment.from_wav(voice_path)
         music = AudioSegment.from_wav(music_path)
-        voice_duration = len(voice)
 
-        if len(music) < voice_duration:
+        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:
             looped_music = AudioSegment.empty()
-            while len(looped_music) < voice_duration:
+            while len(looped_music) < voice_len:
                 looped_music += music
             music = looped_music
-        else:
-            music = music[:voice_duration]
+
+        # Trim music if it's longer than the voice
+        if len(music) > voice_len:
+            music = music[:voice_len]
 
         if ducking:
             ducked_music = music - duck_level
@@ -355,13 +263,12 @@ 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
 # ---------------------------------------------------------------------
@@ -415,7 +322,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 or Kokoro TTS.
+    - **Voice Synthesis**: Create natural-sounding voice-overs using Coqui TTS.
     - **Music Production**: Produce custom music tracks with MusicGen.
     - **Audio Blending**: Seamlessly blend voice and music with options for ducking.
     """)
@@ -448,26 +355,20 @@ with gr.Blocks(css="""
             music_suggestion_output = gr.Textbox(label="Music Suggestions", lines=3, interactive=False)
 
             generate_script_button.click(
-                fn=lambda prompt, model, dur: generate_script(prompt, model, HF_TOKEN, dur),
+                fn=lambda user_prompt, model_id, dur: generate_script(user_prompt, model_id, 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. Choose your TTS engine below:")
-            voice_engine = gr.Dropdown(
-                label="TTS Engine",
-                choices=["Coqui TTS", "Kokoro TTS"],
-                value="Coqui TTS",
-                multiselect=False
-            )
+            gr.Markdown("Generate a natural-sounding voice-over using Coqui TTS.")
             selected_tts_model = gr.Dropdown(
-                label="TTS Model / Voice Option",
+                label="TTS Model",
                 choices=[
-                    "tts_models/en/ljspeech/tacotron2-DDC",  # Coqui TTS option
-                    "tts_models/en/ljspeech/vits",            # Coqui TTS option
-                    "af_heart"                                # Kokoro TTS voice option
+                    "tts_models/en/ljspeech/tacotron2-DDC",  
+                    "tts_models/en/ljspeech/vits", 
+                    "tts_models/en/sam/tacotron-DDC", 
                 ],
                 value="tts_models/en/ljspeech/tacotron2-DDC",
                 multiselect=False
@@ -475,18 +376,9 @@ 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=generate_voice_combined,
-                inputs=[script_output, voice_engine, selected_tts_model],
+                fn=lambda script, tts_model: generate_voice(script, tts_model),
+                inputs=[script_output, selected_tts_model],
                 outputs=voice_audio_output,
             )
 
@@ -505,7 +397,7 @@ with gr.Blocks(css="""
             music_output = gr.Audio(label="Generated Music (WAV)", type="filepath")
 
             generate_music_button.click(
-                fn=lambda prompt, length: generate_music(prompt, length),
+                fn=lambda music_suggestion, length: generate_music(music_suggestion, length),
                 inputs=[music_suggestion_output, audio_length],
                 outputs=[music_output],
             )