Update app.py

app.py

@@ -1,4 +1,4 @@
-# import os
+import os
 import re
 import torch
 import tempfile
@@ -39,7 +39,6 @@ def clean_text(text: str) -> str:
     """
     Removes undesired characters (e.g., asterisks) that might not be recognized by the model's vocabulary.
     """
-    # Remove all asterisks. You can add more cleaning steps here as needed.
     return re.sub(r'\*', '', text)
 
 # ---------------------------------------------------------------------
@@ -64,7 +63,6 @@ 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"):
     """
     Returns a cached MusicGen model if available; otherwise, loads it.
@@ -81,7 +79,6 @@ 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"):
     """
     Returns a cached TTS model if available; otherwise, loads it.
@@ -93,7 +90,6 @@ def get_tts_model(model_name: str = "tts_models/en/ljspeech/tacotron2-DDC"):
     TTS_MODELS[model_name] = tts_model
     return tts_model
 
-
 # ---------------------------------------------------------------------
 # Script Generation Function
 # ---------------------------------------------------------------------
@@ -105,7 +101,6 @@ def generate_script(user_prompt: str, model_id: str, token: str, duration: int):
     """
     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: "
@@ -132,7 +127,7 @@ def generate_script(user_prompt: str, model_id: str, token: str, duration: int):
         sound_design = "No sound design suggestions found."
         music_suggestions = "No music suggestions found."
 
-        # Voice-Over Script
+        # Extract Voice-Over Script
         if "Voice-Over Script:" in generated_text:
             parts = generated_text.split("Voice-Over Script:")
             voice_script_part = parts[1]
@@ -141,7 +136,7 @@ def generate_script(user_prompt: str, model_id: str, token: str, duration: int):
             else:
                 voice_script = voice_script_part.strip()
 
-        # Sound Design
+        # Extract Sound Design Suggestions
         if "Sound Design Suggestions:" in generated_text:
             parts = generated_text.split("Sound Design Suggestions:")
             sound_design_part = parts[1]
@@ -150,7 +145,7 @@ def generate_script(user_prompt: str, model_id: str, token: str, duration: int):
             else:
                 sound_design = sound_design_part.strip()
 
-        # Music Suggestions
+        # Extract Music Suggestions
         if "Music Suggestions:" in generated_text:
             parts = generated_text.split("Music Suggestions:")
             music_suggestions = parts[1].strip()
@@ -160,7 +155,6 @@ def generate_script(user_prompt: str, model_id: str, token: str, duration: int):
     except Exception as e:
         return f"Error generating script: {e}", "", ""
 
-
 # ---------------------------------------------------------------------
 # Voice-Over Generation Function
 # ---------------------------------------------------------------------
@@ -174,12 +168,8 @@ def generate_voice(script: str, tts_model_name: str = "tts_models/en/ljspeech/ta
         if not script.strip():
             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)
-
-        # 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)
         return output_path
@@ -187,7 +177,6 @@ def generate_voice(script: str, tts_model_name: str = "tts_models/en/ljspeech/ta
     except Exception as e:
         return f"Error generating voice: {e}"
 
-
 # ---------------------------------------------------------------------
 # Music Generation Function
 # ---------------------------------------------------------------------
@@ -205,23 +194,22 @@ 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"
-        inputs = musicgen_processor(text=[prompt], padding=True, return_tensors="pt").to(device)
+        # 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()}
 
         with torch.inference_mode():
             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")
-
         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:
         return f"Error generating music: {e}"
 
-
 # ---------------------------------------------------------------------
 # Audio Blending with Duration Sync & Ducking
 # ---------------------------------------------------------------------
@@ -241,17 +229,15 @@ 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)  # in milliseconds
-        music_len = len(music)  # in milliseconds
+        voice_len = len(voice)
+        music_len = len(music)
 
-        # Loop music if it's shorter than the voice
         if music_len < voice_len:
             looped_music = AudioSegment.empty()
             while len(looped_music) < voice_len:
                 looped_music += music
             music = looped_music
 
-        # Trim music if it's longer than the voice
         if len(music) > voice_len:
             music = music[:voice_len]
 
@@ -268,6 +254,28 @@ def blend_audio(voice_path: str, music_path: str, ducking: bool, duck_level: int
     except Exception as e:
         return f"Error blending audio: {e}"
 
+# ---------------------------------------------------------------------
+# Agent Function: Orchestrate the Full Workflow
+# ---------------------------------------------------------------------
+@spaces.GPU(duration=400)
+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.
+    """
+    # 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
 
 # ---------------------------------------------------------------------
 # Gradio Interface with Enhanced UI
@@ -314,12 +322,12 @@ with gr.Blocks(css="""
     # Custom Header
     with gr.Row(elem_classes="header"):
         gr.Markdown("""
-        <h1>🎧 AI Ads Promo</h1>
-        <p>Your all-in-one AI solution for crafting engaging audio ads. <br><em>Demo MVP</em></p>
+        <h1>🎧 AI Promo Studio</h1>
+        <p>Your all-in-one AI solution for crafting engaging audio promos.</p>
         """)
 
     gr.Markdown("""
-    Welcome to **AI Ads Promo (Demo MVP)**! This platform leverages state-of-the-art AI models to help you generate:
+    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.
@@ -328,7 +336,7 @@ with gr.Blocks(css="""
     """)
 
     with gr.Tabs():
-        # Step 1: Generate Script
+        # Tab 1: Script Generation
         with gr.Tab("📝 Script Generation"):
             with gr.Row():
                 user_prompt = gr.Textbox(
@@ -360,7 +368,7 @@ with gr.Blocks(css="""
                 outputs=[script_output, sound_design_output, music_suggestion_output],
             )
 
-        # Step 2: Generate Voice
+        # 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(
@@ -382,7 +390,7 @@ with gr.Blocks(css="""
                 outputs=voice_audio_output,
             )
 
-        # Step 3: Generate Music
+        # 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(
@@ -402,7 +410,7 @@ with gr.Blocks(css="""
                 outputs=[music_output],
             )
 
-        # Step 4: Blend Audio
+        # 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.")
             ducking_checkbox = gr.Checkbox(label="Enable Ducking?", value=True)
@@ -422,17 +430,70 @@ with gr.Blocks(css="""
                 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.")
+            with gr.Row():
+                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
+                )
+            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
+                )
+            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_run_button = gr.Button("Run Agent", variant="primary")
+            agent_script_output = gr.Textbox(label="Generated 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]
+            )
+
     # Footer
     gr.Markdown("""
     <div class="footer">
         <hr>
         Created with ❤️ by <a href="https://bilsimaging.com" target="_blank" style="color: #88aaff;">bilsimaging.com</a>
         <br>
-        <small>AI Ads Promo (Demo MVP) &copy; 2025</small>
+        <small>AI Promo Studio &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">