Update app.py

app.py

@@ -3,6 +3,8 @@ import re
 import torch
 import tempfile
 import logging
+import math
+from typing import Tuple, Union, Any
 from scipy.io.wavfile import write
 from pydub import AudioSegment
 from dotenv import load_dotenv
@@ -30,37 +32,53 @@ from packaging import version
 # ---------------------------------------------------------------------
 # Setup Logging and Environment Variables
 # ---------------------------------------------------------------------
-logging.basicConfig(level=logging.INFO)
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 load_dotenv()
 HF_TOKEN = os.getenv("HF_TOKEN")
+if not HF_TOKEN:
+    logging.warning("HF_TOKEN is not set in your environment. Some model downloads might fail.")
 
 # ---------------------------------------------------------------------
 # Global Model Caches
 # ---------------------------------------------------------------------
-LLAMA_PIPELINES = {}
-MUSICGEN_MODELS = {}
-TTS_MODELS = {}
-SOUND_DESIGN_PIPELINES = {}
+LLAMA_PIPELINES: dict[str, Any] = {}
+MUSICGEN_MODELS: dict[str, Any] = {}
+TTS_MODELS: dict[str, Any] = {}
+SOUND_DESIGN_PIPELINES: dict[str, Any] = {}
 
 # ---------------------------------------------------------------------
-# Utility Function
+# Utility Functions
 # ---------------------------------------------------------------------
 def clean_text(text: str) -> str:
     """
-    Removes undesired characters (e.g., asterisks) that might not be recognized by the model's vocabulary.
+    Remove undesired characters that may not be recognized by the model.
+    
+    Args:
+        text (str): Input text to be cleaned.
+    
+    Returns:
+        str: Cleaned text.
     """
     return re.sub(r'\*', '', text)
 
 # ---------------------------------------------------------------------
 # Model Helper Functions
 # ---------------------------------------------------------------------
-def get_llama_pipeline(model_id: str, token: str):
+def get_llama_pipeline(model_id: str, token: str) -> Any:
     """
-    Returns a cached LLaMA text-generation pipeline if available; otherwise, loads and caches it.
+    Returns a cached LLaMA text-generation pipeline or loads a new one.
+    
+    Args:
+        model_id (str): Hugging Face model ID.
+        token (str): Hugging Face token.
+        
+    Returns:
+        Any: A Hugging Face text-generation pipeline.
     """
     if model_id in LLAMA_PIPELINES:
         return LLAMA_PIPELINES[model_id]
 
+    logging.info(f"Loading LLaMA model from {model_id}...")
     tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=token)
     model = AutoModelForCausalLM.from_pretrained(
         model_id,
@@ -73,14 +91,20 @@ def get_llama_pipeline(model_id: str, token: str):
     LLAMA_PIPELINES[model_id] = text_pipeline
     return text_pipeline
 
-def get_musicgen_model(model_key: str = "facebook/musicgen-large"):
+def get_musicgen_model(model_key: str = "facebook/musicgen-large") -> Tuple[Any, Any]:
     """
-    Returns a cached MusicGen model and processor if available; otherwise, loads and caches them.
-    Uses the 'large' variant for higher quality outputs.
+    Returns a cached MusicGen model and processor, or loads new ones.
+    
+    Args:
+        model_key (str): Hugging Face model key (default is 'facebook/musicgen-large').
+    
+    Returns:
+        Tuple[Any, Any]: The MusicGen model and its processor.
     """
     if model_key in MUSICGEN_MODELS:
         return MUSICGEN_MODELS[model_key]
 
+    logging.info(f"Loading MusicGen model from {model_key}...")
     model = MusicgenForConditionalGeneration.from_pretrained(model_key)
     processor = AutoProcessor.from_pretrained(model_key)
     device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -88,30 +112,51 @@ def get_musicgen_model(model_key: str = "facebook/musicgen-large"):
     MUSICGEN_MODELS[model_key] = (model, processor)
     return model, processor
 
-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 and caches it.
+    Returns a cached TTS model or loads a new one.
+    
+    Args:
+        model_name (str): Identifier for the TTS model.
+    
+    Returns:
+        TTS: A Coqui TTS model.
     """
     if model_name in TTS_MODELS:
         return TTS_MODELS[model_name]
 
+    logging.info(f"Loading TTS model: {model_name}...")
     tts_model = TTS(model_name)
     TTS_MODELS[model_name] = tts_model
     return tts_model
 
-def get_sound_design_pipeline(model_name: str, token: str):
+def get_sound_design_pipeline(model_name: str, token: str) -> Any:
     """
-    Returns a cached DiffusionPipeline for sound design if available;
-    otherwise, it loads and caches the pipeline.
+    Returns a cached DiffusionPipeline for sound design, or loads a new one.
+    Raises an error if diffusers version is less than 0.21.0.
+    
+    Args:
+        model_name (str): The model name to load.
+        token (str): Hugging Face token.
     
-    NOTE: AudioLDM2Pipeline is available only in diffusers>=0.21.0.
-    Since your requirements fix diffusers==0.20.2, this function will raise an error.
+    Returns:
+        Any: A DiffusionPipeline for sound design.
+    
+    Raises:
+        ValueError: If diffusers version is lower than 0.21.0.
     """
     if version.parse(diffusers.__version__) < version.parse("0.21.0"):
         raise ValueError("AudioLDM2 requires diffusers>=0.21.0. Please upgrade your diffusers package.")
+    
     if model_name in SOUND_DESIGN_PIPELINES:
         return SOUND_DESIGN_PIPELINES[model_name]
-    pipe = DiffusionPipeline.from_pretrained(model_name, pipeline_class=AudioLDMPipeline, use_auth_token=token)
+    
+    logging.info(f"Loading sound design pipeline from {model_name}...")
+    pipe = DiffusionPipeline.from_pretrained(
+        model_name,
+        pipeline_class=AudioLDMPipeline,
+        use_auth_token=token
+    )
     SOUND_DESIGN_PIPELINES[model_name] = pipe
     return pipe
 
@@ -119,14 +164,21 @@ def get_sound_design_pipeline(model_name: str, token: str):
 # 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 voice-over script, sound design suggestions, and music ideas from a user prompt.
-    Returns a tuple: (voice_script, sound_design, music_suggestions).
+    Generates a voice-over script, sound design suggestions, and music ideas based on the user prompt.
+    
+    Args:
+        user_prompt (str): The user-provided concept.
+        model_id (str): The LLaMA model ID.
+        token (str): Hugging Face token.
+        duration (int): The desired duration in seconds.
+    
+    Returns:
+        Tuple[str, str, str]: Voice-over script, sound design suggestions, and 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"
@@ -148,6 +200,7 @@ def generate_script(user_prompt: str, model_id: str, token: str, duration: int):
         if "Output:" in generated_text:
             generated_text = generated_text.split("Output:")[-1].strip()
 
+        # Extract sections using regex
         pattern = r"Voice-Over Script:\s*(.*?)\s*Sound Design Suggestions:\s*(.*?)\s*Music Suggestions:\s*(.*)"
         match = re.search(pattern, generated_text, re.DOTALL)
         if match:
@@ -167,10 +220,16 @@ def generate_script(user_prompt: str, model_id: str, token: str, duration: int):
 # Voice-Over Generation Function
 # ---------------------------------------------------------------------
 @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, Any]:
     """
-    Generates a voice-over audio file from the provided script using Coqui TTS.
-    Returns the file path to the generated .wav file.
+    Generates a voice-over audio file from a script using Coqui TTS.
+    
+    Args:
+        script (str): The voice-over script.
+        tts_model_name (str): The TTS model name.
+    
+    Returns:
+        Union[str, Any]: The file path to the generated .wav file or an error message.
     """
     try:
         if not script.strip():
@@ -178,7 +237,6 @@ def generate_voice(script: str, tts_model_name: str = "tts_models/en/ljspeech/ta
 
         cleaned_script = clean_text(script)
         tts_model = get_tts_model(tts_model_name)
-
         output_path = os.path.join(tempfile.gettempdir(), "voice_over.wav")
         tts_model.tts_to_file(text=cleaned_script, file_path=output_path)
         return output_path
@@ -191,10 +249,16 @@ def generate_voice(script: str, tts_model_name: str = "tts_models/en/ljspeech/ta
 # 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, Any]:
     """
-    Generates a music track from the 'facebook/musicgen-large' model based on the prompt.
-    Returns the file path to the generated .wav file.
+    Generates a music track using the MusicGen model based on the prompt.
+    
+    Args:
+        prompt (str): Music suggestion prompt.
+        audio_length (int): Number of tokens determining audio length.
+    
+    Returns:
+        Union[str, Any]: The file path to the generated .wav file or an error message.
     """
     try:
         if not prompt.strip():
@@ -209,11 +273,10 @@ def generate_music(prompt: str, audio_length: int):
             outputs = musicgen_model.generate(**inputs, max_new_tokens=audio_length)
 
         audio_data = outputs[0, 0].cpu().numpy()
-        normalized_audio = (audio_data / max(abs(audio_data)) * 32767).astype("int16")
-
+        # Normalize audio data to 16-bit integer range
+        normalized_audio = (audio_data / np.max(np.abs(audio_data)) * 32767).astype("int16")
         output_path = os.path.join(tempfile.gettempdir(), "musicgen_large_generated_music.wav")
         write(output_path, 44100, normalized_audio)
-
         return output_path
 
     except Exception as e:
@@ -224,21 +287,23 @@ def generate_music(prompt: str, audio_length: int):
 # Sound Design Generation Function
 # ---------------------------------------------------------------------
 @spaces.GPU(duration=200)
-def generate_sound_design(prompt: str):
+def generate_sound_design(prompt: str) -> Union[str, Any]:
     """
-    Generates a sound design audio file based on the provided prompt using AudioLDM 2.
-    Returns the file path to the generated .wav file.
+    Generates a sound design audio file using AudioLDM 2 based on the prompt.
+    
+    Args:
+        prompt (str): Sound design prompt.
+    
+    Returns:
+        Union[str, Any]: The file path to the generated .wav file or an error message.
     """
     try:
         if not prompt.strip():
             return "Error: No sound design suggestion provided."
         
         pipe = get_sound_design_pipeline("cvssp/audioldm2", HF_TOKEN)
-        
-        # Generate audio from the prompt; assumes the pipeline returns a dict with key 'audios'
-        result = pipe(prompt)
+        result = pipe(prompt)  # Expected to return a dict with key 'audios'
         audio_samples = result["audios"][0]
-        
         normalized_audio = (audio_samples / np.max(np.abs(audio_samples)) * 32767).astype("int16")
         output_path = os.path.join(tempfile.gettempdir(), "sound_design_generated.wav")
         write(output_path, 44100, normalized_audio)
@@ -249,19 +314,27 @@ def generate_sound_design(prompt: str):
         return f"Error generating sound design: {e}"
 
 # ---------------------------------------------------------------------
-# Audio Blending with Duration Sync & Ducking (Voice + Sound Design + Music)
+# Audio Blending Function
 # ---------------------------------------------------------------------
 @spaces.GPU(duration=100)
-def blend_audio(voice_path: str, sound_effect_path: str, music_path: str, ducking: bool, duck_level: int = 10):
+def blend_audio(voice_path: str, sound_effect_path: str, music_path: str, ducking: bool, duck_level: int = 10) -> Union[str, Any]:
     """
-    Blends three audio files (voice, sound design/sound effect, and music):
-      - Loops music and sound design if shorter than the voice track.
-      - Trims both to match the voice duration.
-      - Applies ducking to lower music and sound design volumes during voice segments if enabled.
-    Returns the file path to the blended .wav file.
+    Blends three audio files (voice, sound design, and music) by:
+      - Looping/trimming music and sound design to match voice duration.
+      - Optionally applying ducking to background tracks.
+      - Overlaying the voice on top of the background.
+    
+    Args:
+        voice_path (str): Path to the voice audio file.
+        sound_effect_path (str): Path to the sound design audio file.
+        music_path (str): Path to the music audio file.
+        ducking (bool): Whether to apply ducking.
+        duck_level (int): Amount of attenuation in dB.
+    
+    Returns:
+        Union[str, Any]: The file path to the blended .wav file or an error message.
     """
     try:
-        # Verify input files exist
         for path in [voice_path, sound_effect_path, music_path]:
             if not os.path.isfile(path):
                 return f"Error: Missing audio file for {path}"
@@ -270,34 +343,30 @@ def blend_audio(voice_path: str, sound_effect_path: str, music_path: str, duckin
         voice = AudioSegment.from_wav(voice_path)
         music = AudioSegment.from_wav(music_path)
         sound_effect = AudioSegment.from_wav(sound_effect_path)
-
         voice_len = len(voice)  # duration in milliseconds
 
-        # Loop or trim music to match voice duration
+        # Loop or trim music to match voice duration using pydub multiplication
         if len(music) < voice_len:
-            looped_music = AudioSegment.empty()
-            while len(looped_music) < voice_len:
-                looped_music += music
-            music = looped_music
-        music = music[:voice_len]
+            repeats = math.ceil(voice_len / len(music))
+            music = (music * repeats)[:voice_len]
+        else:
+            music = music[:voice_len]
 
-        # Loop or trim sound effect to match voice duration
+        # Loop or trim sound design to match voice duration
         if len(sound_effect) < voice_len:
-            looped_effect = AudioSegment.empty()
-            while len(looped_effect) < voice_len:
-                looped_effect += sound_effect
-            sound_effect = looped_effect
-        sound_effect = sound_effect[:voice_len]
+            repeats = math.ceil(voice_len / len(sound_effect))
+            sound_effect = (sound_effect * repeats)[:voice_len]
+        else:
+            sound_effect = sound_effect[:voice_len]
 
-        # Apply ducking to background tracks if enabled
+        # Apply ducking if enabled
         if ducking:
             music = music - duck_level
             sound_effect = sound_effect - duck_level
 
-        # Combine music and sound effect into a background track
+        # Overlay music and sound effect for background
         background = music.overlay(sound_effect)
-
-        # Overlay voice on top of the background
+        # Overlay voice on top of background
         final_audio = background.overlay(voice)
 
         output_path = os.path.join(tempfile.gettempdir(), "blended_output.wav")
@@ -360,25 +429,13 @@ with gr.Blocks(css="""
     gr.Markdown("""
     **Welcome to Ai Ads Promo!**
 
-    This simple, easy-to-use app helps you create amazing audio ads in just a few steps. Here’s how it works:
-
-    1. **Script Generation:** 
-       - Share your idea and let our AI craft a clear and engaging voice-over script, along with sound design and music suggestions.
-    2. **Voice Synthesis:** 
-       - Convert your script into a natural-sounding voice-over using advanced text-to-speech technology.
-    3. **Music Production:** 
-       - Generate a custom music track that perfectly fits your ad.
-    4. **Sound Design:** 
-       - Generate creative sound effects based on our sound design suggestions.
-    5. **Audio Blending:** 
-       - Combine your voice-over, sound effects, and music seamlessly. Enable ducking to lower background audio during voice segments.
-
-    **Benefits:**
-    - **Easy to Use:** Designed for everyone – no technical skills required.
-    - **Fast Results:** Quickly produce professional-sounding audio ads.
-    - **All-In-One:** Everything you need in one convenient app.
-
-    Get started now and create your perfect audio ad with Ai Ads Promo!
+    This app helps you create amazing audio ads in just a few steps:
+    
+    1. **Script Generation:** Provide your idea and get a voice-over script, sound design, and music suggestions.
+    2. **Voice Synthesis:** Convert the script into natural-sounding speech.
+    3. **Music Production:** Generate a custom music track.
+    4. **Sound Design:** Create creative sound effects.
+    5. **Audio Blending:** Seamlessly blend voice, music, and sound design (with optional ducking).
     """)
 
     with gr.Tabs():
@@ -458,7 +515,7 @@ with gr.Blocks(css="""
 
         # Step 4: Sound Design Generation
         with gr.Tab("🎧 Sound Design Generation"):
-            gr.Markdown("Generate a creative sound design track based on the sound design suggestions from the script.")
+            gr.Markdown("Generate a creative sound design track based on the script's suggestions.")
             generate_sound_design_button = gr.Button("Generate Sound Design", variant="primary")
             sound_design_audio_output = gr.Audio(label="Generated Sound Design (WAV)", type="filepath")
             
@@ -470,7 +527,7 @@ with gr.Blocks(css="""
 
         # Step 5: Audio Blending (Voice + Sound Design + Music)
         with gr.Tab("🎚️ Audio Blending"):
-            gr.Markdown("Blend your voice-over, sound design, and music track. The background audio (music and sound design) can be ducked during voice segments.")
+            gr.Markdown("Blend your voice-over, sound design, and music track. Enable ducking to lower background audio during voice segments.")
             ducking_checkbox = gr.Checkbox(label="Enable Ducking?", value=True)
             duck_level_slider = gr.Slider(
                 label="Ducking Level (dB attenuation)", 
@@ -488,7 +545,7 @@ with gr.Blocks(css="""
                 outputs=blended_output
             )
 
-    # Footer
+    # Footer and Visitor Badge
     gr.Markdown("""
     <div class="footer">
         <hr>
@@ -497,8 +554,6 @@ with gr.Blocks(css="""
         <small>Ai Ads Promo &copy; 2025</small>
     </div>
     """)
-    
-    # Visitor Badge
     gr.HTML("""
     <div style="text-align: center; margin-top: 1rem;">
         <a href="https://visitorbadge.io/status?path=https%3A%2F%2Fhuggingface.co%2Fspaces%2FBils%2Fradiogold">
@@ -507,4 +562,5 @@ with gr.Blocks(css="""
     </div>
     """)
 
-demo.launch(debug=True)
+if __name__ == "__main__":
+    demo.launch(debug=True)