Update app.py

app.py

@@ -50,7 +50,6 @@ def get_llama_pipeline(model_id: str, token: str):
     """
     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,
@@ -70,10 +69,8 @@ def get_musicgen_model(model_key: str = "facebook/musicgen-large"):
     """
     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)
@@ -85,7 +82,6 @@ def get_tts_model(model_name: str = "tts_models/en/ljspeech/tacotron2-DDC"):
     """
     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
@@ -97,7 +93,7 @@ def get_tts_model(model_name: str = "tts_models/en/ljspeech/tacotron2-DDC"):
 def generate_script(user_prompt: str, model_id: str, token: str, duration: int):
     """
     Generates a script, sound design suggestions, and music ideas from a user prompt.
-    Returns a tuple of strings: (voice_script, sound_design, music_suggestions).
+    Returns a tuple: (voice_script, sound_design, music_suggestions).
     """
     try:
         text_pipeline = get_llama_pipeline(model_id, token)
@@ -122,12 +118,10 @@ 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()
 
-        # Default placeholders
         voice_script = "No voice-over script found."
         sound_design = "No sound design suggestions found."
         music_suggestions = "No music suggestions found."
 
-        # Extract Voice-Over Script
         if "Voice-Over Script:" in generated_text:
             parts = generated_text.split("Voice-Over Script:")
             voice_script_part = parts[1]
@@ -136,7 +130,6 @@ def generate_script(user_prompt: str, model_id: str, token: str, duration: int):
             else:
                 voice_script = voice_script_part.strip()
 
-        # Extract Sound Design Suggestions
         if "Sound Design Suggestions:" in generated_text:
             parts = generated_text.split("Sound Design Suggestions:")
             sound_design_part = parts[1]
@@ -145,7 +138,6 @@ def generate_script(user_prompt: str, model_id: str, token: str, duration: int):
             else:
                 sound_design = sound_design_part.strip()
 
-        # Extract Music Suggestions
         if "Music Suggestions:" in generated_text:
             parts = generated_text.split("Music Suggestions:")
             music_suggestions = parts[1].strip()
@@ -161,19 +153,17 @@ def generate_script(user_prompt: str, model_id: str, token: str, duration: int):
 @spaces.GPU(duration=100)
 def generate_voice(script: str, tts_model_name: str = "tts_models/en/ljspeech/tacotron2-DDC"):
     """
-    Generates a voice-over from the provided script using the Coqui TTS model.
+    Generates a voice-over from the provided script using Coqui TTS.
     Returns the file path to the generated .wav file.
     """
     try:
         if not script.strip():
             return "Error: No script provided."
-
         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
-
     except Exception as e:
         return f"Error generating voice: {e}"
 
@@ -194,7 +184,6 @@ def generate_music(prompt: str, audio_length: int):
         musicgen_model, musicgen_processor = get_musicgen_model(model_key)
 
         device = "cuda" if torch.cuda.is_available() else "cpu"
-        # Process the input and move each tensor to the proper device
         inputs = musicgen_processor(text=[prompt], padding=True, return_tensors="pt")
         inputs = {k: v.to(device) for k, v in inputs.items()}
 
@@ -211,15 +200,12 @@ def generate_music(prompt: str, audio_length: int):
         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):
     """
     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:
@@ -228,7 +214,6 @@ def blend_audio(voice_path: str, music_path: str, ducking: bool, duck_level: int
 
         voice = AudioSegment.from_wav(voice_path)
         music = AudioSegment.from_wav(music_path)
-
         voice_len = len(voice)
         music_len = len(music)
 
@@ -241,12 +226,7 @@ def blend_audio(voice_path: str, music_path: str, ducking: bool, duck_level: int
         if len(music) > voice_len:
             music = music[:voice_len]
 
-        if ducking:
-            ducked_music = music - duck_level
-            final_audio = ducked_music.overlay(voice)
-        else:
-            final_audio = music.overlay(voice)
-
+        final_audio = music.overlay(voice, gain_during_overlay=-duck_level) if ducking else music.overlay(voice)
         output_path = os.path.join(tempfile.gettempdir(), "blended_output.wav")
         final_audio.export(output_path, format="wav")
         return output_path
@@ -261,19 +241,15 @@ def blend_audio(voice_path: str, music_path: str, ducking: bool, duck_level: int
 def run_agent(user_prompt: str, llama_model_id: str, duration: int, tts_model_name: str, music_length: int, ducking: bool, duck_level: int):
     """
     Runs the full workflow as an agent:
-    1. Generates a script (voice-over, sound design, music suggestions) from a user prompt.
-    2. Synthesizes a voice-over from the generated script.
-    3. Generates a music track based on the music suggestions.
-    4. Blends the voice and music tracks.
-    Returns a tuple with the generated script components, voice file, music file, and final blended audio.
+      1. Generates a script (voice-over, sound design, and music suggestions).
+      2. Synthesizes a voice-over.
+      3. Generates a music track.
+      4. Blends the voice and music.
+    Returns all generated components.
     """
-    # Step 1: Generate Script
     voice_script, sound_design, music_suggestions = generate_script(user_prompt, llama_model_id, HF_TOKEN, duration)
-    # Step 2: Generate Voice-Over
     voice_file = generate_voice(voice_script, tts_model_name)
-    # Step 3: Generate Music
     music_file = generate_music(music_suggestions, music_length)
-    # Step 4: Blend Audio
     blended_file = blend_audio(voice_file, music_file, ducking, duck_level)
     return voice_script, sound_design, music_suggestions, voice_file, music_file, blended_file
 
@@ -281,7 +257,6 @@ def run_agent(user_prompt: str, llama_model_id: str, duration: int, tts_model_na
 # Gradio Interface with Enhanced UI
 # ---------------------------------------------------------------------
 with gr.Blocks(css="""
-    /* Global Styles */
     body {
         background: linear-gradient(135deg, #1d1f21, #3a3d41);
         color: #f0f0f0;
@@ -328,163 +303,88 @@ with gr.Blocks(css="""
 
     gr.Markdown("""
     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.
-    - **Music Production**: Produce custom music tracks with MusicGen.
-    - **Audio Blending**: Seamlessly blend voice and music with options for ducking.
+    - A compelling voice-over script (with sound design and music suggestions),
+    - A natural-sounding voice-over,
+    - Custom music tracks,
+    - And a fully blended audio promo.
     """)
 
     with gr.Tabs():
         # Tab 1: Script Generation
         with gr.Tab("📝 Script Generation"):
             with gr.Row():
-                user_prompt = gr.Textbox(
-                    label="Promo Idea", 
-                    placeholder="E.g., A 30-second promo for a morning show...",
-                    lines=2
-                )
+                user_prompt = gr.Textbox(label="Promo Idea", placeholder="E.g., A 30-second promo for a morning show...", lines=2)
             with gr.Row():
-                llama_model_id = gr.Textbox(
-                    label="LLaMA Model ID", 
-                    value="meta-llama/Meta-Llama-3-8B-Instruct", 
-                    placeholder="Enter a valid Hugging Face model ID"
-                )
-                duration = gr.Slider(
-                    label="Desired Promo Duration (seconds)",
-                    minimum=15, 
-                    maximum=60, 
-                    step=15, 
-                    value=30
-                )
+                llama_model_id = gr.Textbox(label="LLaMA Model ID", value="meta-llama/Meta-Llama-3-8B-Instruct", placeholder="Enter a valid Hugging Face model ID")
+                duration = gr.Slider(label="Promo Duration (seconds)", minimum=15, maximum=60, step=15, value=30)
             generate_script_button = gr.Button("Generate Script", variant="primary")
-            script_output = gr.Textbox(label="Generated Voice-Over Script", lines=5, interactive=False)
+            script_output = gr.Textbox(label="Voice-Over Script", lines=5, interactive=False)
             sound_design_output = gr.Textbox(label="Sound Design Suggestions", lines=3, interactive=False)
             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),
-                inputs=[user_prompt, llama_model_id, duration],
-                outputs=[script_output, sound_design_output, music_suggestion_output],
-            )
+            generate_script_button.click(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])
 
         # Tab 2: Voice Synthesis
         with gr.Tab("🎤 Voice Synthesis"):
             gr.Markdown("Generate a natural-sounding voice-over using Coqui TTS.")
-            selected_tts_model = gr.Dropdown(
-                label="TTS Model",
-                choices=[
-                    "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
-            )
+            selected_tts_model = gr.Dropdown(label="TTS Model",
+                                             choices=["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)
             generate_voice_button = gr.Button("Generate Voice-Over", variant="primary")
             voice_audio_output = gr.Audio(label="Voice-Over (WAV)", type="filepath")
-
-            generate_voice_button.click(
-                fn=lambda script, tts_model: generate_voice(script, tts_model),
-                inputs=[script_output, selected_tts_model],
-                outputs=voice_audio_output,
-            )
+            generate_voice_button.click(fn=lambda script, tts: generate_voice(script, tts),
+                                        inputs=[script_output, selected_tts_model],
+                                        outputs=voice_audio_output)
 
         # Tab 3: Music Production
         with gr.Tab("🎶 Music Production"):
-            gr.Markdown("Generate a custom music track using the **MusicGen Large** model.")
-            audio_length = gr.Slider(
-                label="Music Length (tokens)",
-                minimum=128, 
-                maximum=1024, 
-                step=64, 
-                value=512,
-                info="Increase tokens for longer audio (inference time may vary)."
-            )
+            gr.Markdown("Generate a custom music track using the MusicGen Large model.")
+            audio_length = gr.Slider(label="Music Length (tokens)", minimum=128, maximum=1024, step=64, value=512, info="Increase tokens for longer audio (inference time may vary).")
             generate_music_button = gr.Button("Generate Music", variant="primary")
             music_output = gr.Audio(label="Generated Music (WAV)", type="filepath")
-
-            generate_music_button.click(
-                fn=lambda music_suggestion, length: generate_music(music_suggestion, length),
-                inputs=[music_suggestion_output, audio_length],
-                outputs=[music_output],
-            )
+            generate_music_button.click(fn=lambda sugg, length: generate_music(sugg, length),
+                                        inputs=[music_suggestion_output, audio_length],
+                                        outputs=[music_output])
 
         # Tab 4: Audio Blending
         with gr.Tab("🎚️ Audio Blending"):
-            gr.Markdown("Blend your voice-over and music track. Music will be looped/truncated to match the voice duration. Enable ducking to lower the music during voice segments.")
+            gr.Markdown("Blend your voice-over and music track. Enable ducking to lower the music during voice segments.")
             ducking_checkbox = gr.Checkbox(label="Enable Ducking?", value=True)
-            duck_level_slider = gr.Slider(
-                label="Ducking Level (dB attenuation)", 
-                minimum=0, 
-                maximum=20, 
-                step=1, 
-                value=10
-            )
+            duck_level_slider = gr.Slider(label="Ducking Level (dB attenuation)", minimum=0, maximum=20, step=1, value=10)
             blend_button = gr.Button("Blend Voice + Music", variant="primary")
             blended_output = gr.Audio(label="Final Blended Output (WAV)", type="filepath")
-
-            blend_button.click(
-                fn=blend_audio,
-                inputs=[voice_audio_output, music_output, ducking_checkbox, duck_level_slider],
-                outputs=blended_output
-            )
+            blend_button.click(fn=blend_audio,
+                               inputs=[voice_audio_output, music_output, ducking_checkbox, duck_level_slider],
+                               outputs=blended_output)
 
         # Tab 5: Agent – Full Workflow
         with gr.Tab("🤖 Agent"):
-            gr.Markdown("Let the agent handle everything in one go: generate the script, synthesize voice, produce music, and blend the final ad.")
+            gr.Markdown("Let the agent handle everything in one go: generate script, synthesize voice, produce music, and blend the final ad.")
             with gr.Row():
-                agent_prompt = gr.Textbox(
-                    label="Ad Promo Idea",
-                    placeholder="Enter your ad promo concept...",
-                    lines=2
-                )
+                agent_prompt = gr.Textbox(label="Ad Promo Idea", placeholder="Enter your ad promo concept...", lines=2)
             with gr.Row():
-                agent_llama_model_id = gr.Textbox(
-                    label="LLaMA Model ID", 
-                    value="meta-llama/Meta-Llama-3-8B-Instruct", 
-                    placeholder="Enter a valid Hugging Face model ID"
-                )
-                agent_duration = gr.Slider(
-                    label="Promo Duration (seconds)",
-                    minimum=15, maximum=60, step=15, value=30
-                )
+                agent_llama_model_id = gr.Textbox(label="LLaMA Model ID", value="meta-llama/Meta-Llama-3-8B-Instruct", placeholder="Enter a valid Hugging Face model ID")
+                agent_duration = gr.Slider(label="Promo Duration (seconds)", minimum=15, maximum=60, step=15, value=30)
             with gr.Row():
-                agent_tts_model = gr.Dropdown(
-                    label="TTS Model",
-                    choices=[
-                        "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
-                )
-                agent_music_length = gr.Slider(
-                    label="Music Length (tokens)",
-                    minimum=128, maximum=1024, step=64, value=512
-                )
+                agent_tts_model = gr.Dropdown(label="TTS Model",
+                                              choices=["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)
+                agent_music_length = gr.Slider(label="Music Length (tokens)", minimum=128, maximum=1024, step=64, value=512)
             with gr.Row():
                 agent_ducking = gr.Checkbox(label="Enable Ducking?", value=True)
-                agent_duck_level = gr.Slider(
-                    label="Ducking Level (dB attenuation)",
-                    minimum=0, maximum=20, step=1, value=10
-                )
+                agent_duck_level = gr.Slider(label="Ducking Level (dB attenuation)", minimum=0, maximum=20, step=1, value=10)
             agent_run_button = gr.Button("Run Agent", variant="primary")
-            agent_script_output = gr.Textbox(label="Generated Voice-Over Script", lines=5, interactive=False)
+            agent_script_output = gr.Textbox(label="Voice-Over Script", lines=5, interactive=False)
             agent_sound_output = gr.Textbox(label="Sound Design Suggestions", lines=3, interactive=False)
             agent_music_suggestions_output = gr.Textbox(label="Music Suggestions", lines=3, interactive=False)
             agent_voice_audio = gr.Audio(label="Voice-Over (WAV)", type="filepath")
             agent_music_audio = gr.Audio(label="Generated Music (WAV)", type="filepath")
             agent_blended_audio = gr.Audio(label="Final Blended Output (WAV)", type="filepath")
+            agent_run_button.click(fn=run_agent,
+                                   inputs=[agent_prompt, agent_llama_model_id, agent_duration, agent_tts_model, agent_music_length, agent_ducking, agent_duck_level],
+                                   outputs=[agent_script_output, agent_sound_output, agent_music_suggestions_output, agent_voice_audio, agent_music_audio, agent_blended_audio])
 
-            agent_run_button.click(
-                fn=run_agent,
-                inputs=[agent_prompt, agent_llama_model_id, agent_duration, agent_tts_model, agent_music_length, agent_ducking, agent_duck_level],
-                outputs=[agent_script_output, agent_sound_output, agent_music_suggestions_output, agent_voice_audio, agent_music_audio, agent_blended_audio]
-            )
-
-    # Footer
     gr.Markdown("""
     <div class="footer">
         <hr>