Update app.py

app.py

@@ -1,6 +1,8 @@
-import gradio as gr
 import os
+import uuid
 import torch
+import numpy as np
+import gradio as gr
 from transformers import (
     AutoTokenizer,
     AutoModelForCausalLM,
@@ -29,6 +31,7 @@ MUSICGEN_MODELS = {}
 TTS_MODELS = {}
 
 def get_llama_pipeline(model_id: str, token: str):
+    """Load and cache the LLaMA text-generation pipeline."""
     if model_id in LLAMA_PIPELINES:
         return LLAMA_PIPELINES[model_id]
 
@@ -45,6 +48,7 @@ def get_llama_pipeline(model_id: str, token: str):
     return text_pipeline
 
 def get_musicgen_model(model_key: str = "facebook/musicgen-large"):
+    """Load and cache the MusicGen model and processor."""
     if model_key in MUSICGEN_MODELS:
         return MUSICGEN_MODELS[model_key]
 
@@ -56,6 +60,7 @@ def get_musicgen_model(model_key: str = "facebook/musicgen-large"):
     return model, processor
 
 def get_tts_model(model_name: str = "tts_models/en/ljspeech/tacotron2-DDC"):
+    """Load and cache the TTS model."""
     if model_name in TTS_MODELS:
         return TTS_MODELS[model_name]
     tts_model = TTS(model_name)
@@ -67,12 +72,21 @@ def get_tts_model(model_name: str = "tts_models/en/ljspeech/tacotron2-DDC"):
 # -----------------------------------------------------------
 @spaces.GPU(duration=100)
 def generate_script(user_prompt: str, model_id: str, token: str, duration: int):
+    """
+    Generate a professional promo script including a voice-over script,
+    sound design suggestions, and music recommendations.
+    """
     try:
         text_pipeline = get_llama_pipeline(model_id, token)
-        system_prompt = f"""You are a professional audio producer creating {duration}-second content. Generate:
-1. Voice script (clear and concise)
-2. Sound design suggestions (specific effects)
-3. Music style recommendations (genre, tempo)"""
+        # Updated prompt to instruct the model to output sections with explicit headers.
+        system_prompt = (
+            f"You are a professional audio producer creating {duration}-second content. "
+            "Please generate the following three sections exactly as shown:\n\n"
+            "Voice-Over Script: [A clear and concise script for the voiceover.]\n"
+            "Sound Design Suggestions: [Specific ideas, effects, and ambience recommendations.]\n"
+            "Music Suggestions: [Recommendations for music style, genre, and tempo.]\n\n"
+            "Make sure each section starts with its header exactly."
+        )
         
         full_prompt = f"{system_prompt}\nClient brief: {user_prompt}\nOutput:"
         
@@ -87,7 +101,7 @@ def generate_script(user_prompt: str, model_id: str, token: str, duration: int):
 
         generated_text = result[0]["generated_text"].split("Output:")[-1].strip()
         
-        # Parse sections
+        # Parse the output into the three expected sections.
         sections = {
             "Voice-Over Script:": "",
             "Sound Design Suggestions:": "",
@@ -99,7 +113,9 @@ def generate_script(user_prompt: str, model_id: str, token: str, duration: int):
             for section in sections:
                 if section in line:
                     current_section = section
+                    # Remove header from the line.
                     line = line.replace(section, '').strip()
+                    break
             if current_section:
                 sections[current_section] += line + '\n'
         
@@ -114,19 +130,41 @@ 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):
+    """
+    Generate full voice-over audio from the provided script using a TTS model.
+    """
     try:
         if not script.strip():
             return None
         tts_model = get_tts_model(tts_model_name)
-        output_path = f"{tempfile.gettempdir()}/voice_temp.wav"
+        # Create a unique temporary file name for the output.
+        output_path = os.path.join(tempfile.gettempdir(), f"voice_{uuid.uuid4().hex}.wav")
         tts_model.tts_to_file(text=script, file_path=output_path)
         return output_path
     except Exception as e:
         print(f"Voice generation error: {e}")
         return None
 
+@spaces.GPU(duration=100)
+def generate_voice_preview(script: str, tts_model_name: str):
+    """
+    Generate a short preview of the voice-over by taking the first 100 words.
+    """
+    try:
+        if not script.strip():
+            return None
+        words = script.split()
+        preview_text = ' '.join(words[:100]) if len(words) > 100 else script
+        return generate_voice(preview_text, tts_model_name)
+    except Exception as e:
+        print(f"Voice preview error: {e}")
+        return None
+
 @spaces.GPU(duration=100)
 def generate_music(prompt: str, audio_length: int):
+    """
+    Generate music audio from a text prompt using the MusicGen model.
+    """
     try:
         model, processor = get_musicgen_model()
         device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -135,9 +173,15 @@ def generate_music(prompt: str, audio_length: int):
         with torch.inference_mode():
             outputs = model.generate(**inputs, max_new_tokens=audio_length)
         
+        # Assuming outputs[0, 0] holds the generated audio waveform.
         audio_data = outputs[0, 0].cpu().numpy()
-        normalized_audio = (audio_data / max(abs(audio_data)) * 32767).astype("int16")
-        output_path = f"{tempfile.gettempdir()}/music_temp.wav"
+        # Prevent division by zero during normalization.
+        max_val = np.max(np.abs(audio_data))
+        if max_val == 0:
+            normalized_audio = audio_data.astype("int16")
+        else:
+            normalized_audio = (audio_data / max_val * 32767).astype("int16")
+        output_path = os.path.join(tempfile.gettempdir(), f"music_{uuid.uuid4().hex}.wav")
         write(output_path, 44100, normalized_audio)
         return output_path
     except Exception as e:
@@ -146,24 +190,27 @@ def generate_music(prompt: str, audio_length: int):
 
 @spaces.GPU(duration=100)
 def blend_audio(voice_path: str, music_path: str, ducking: bool, duck_level: int):
+    """
+    Blend the generated voice and music audio files.
+    If ducking is enabled, lower the music volume during the voice segments.
+    """
     try:
         voice = AudioSegment.from_wav(voice_path)
         music = AudioSegment.from_wav(music_path)
         
-        # Adjust music length
+        # Loop the music track if it's shorter than the voice track.
         if len(music) < len(voice):
             loops_needed = (len(voice) // len(music)) + 1
             music = music * loops_needed
         music = music[:len(voice)]
         
-        # Ducking effect
         if ducking:
             ducked_music = music - duck_level
             final_audio = ducked_music.overlay(voice)
         else:
             final_audio = music.overlay(voice)
         
-        output_path = f"{tempfile.gettempdir()}/final_mix.wav"
+        output_path = os.path.join(tempfile.gettempdir(), f"final_mix_{uuid.uuid4().hex}.wav")
         final_audio.export(output_path, format="wav")
         return output_path
     except Exception as e:
@@ -268,7 +315,7 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Default()) as demo:
         
         # Main Workflow Tabs
         with gr.Tabs(elem_classes="tab-nav"):
-            # Script Generation
+            # Script Generation Tab
             with gr.Tab("📝 Script Design", elem_classes="tab-button"):
                 with gr.Row(equal_height=False):
                     with gr.Column(scale=2):
@@ -301,7 +348,7 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Default()) as demo:
                         sound_design_output = gr.Textbox(label="Sound Design", lines=3)
                         music_suggestion_output = gr.Textbox(label="Music Style", lines=3)
 
-            # Voice Production
+            # Voice Production Tab
             with gr.Tab("🎙️ Voice Production", elem_classes="tab-button"):
                 with gr.Row():
                     with gr.Column(scale=1):
@@ -327,7 +374,7 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Default()) as demo:
                             waveform_options={"show_controls": True}
                         )
 
-            # Music Production
+            # Music Production Tab
             with gr.Tab("🎵 Music Design", elem_classes="tab-button"):
                 with gr.Row():
                     with gr.Column(scale=1):
@@ -349,7 +396,7 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Default()) as demo:
                             waveform_options={"show_controls": True}
                         )
 
-            # Final Mix
+            # Final Mix Tab
             with gr.Tab("🔊 Final Mix", elem_classes="tab-button"):
                 with gr.Row():
                     with gr.Column(scale=1):
@@ -375,7 +422,7 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Default()) as demo:
                             waveform_options={"show_controls": True}
                         )
 
-        # Footer
+        # Footer Section
         with gr.Column(elem_classes="output-card"):
             gr.Markdown("""
             <div style="text-align: center; padding: 1.5em 0;">
@@ -391,13 +438,26 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Default()) as demo:
             </p>
             """)
 
+    # -----------------------------------------------------------
     # Event Handling
+    # -----------------------------------------------------------
+    # Hidden textbox for HF_TOKEN (its value is set via the environment variable).
+    hf_token_hidden = gr.Textbox(value=HF_TOKEN, visible=False)
+    
     generate_btn.click(
         generate_script,
-        inputs=[user_prompt, llama_model_id, gr.Textbox(HF_TOKEN, visible=False), duration],
+        inputs=[user_prompt, llama_model_id, hf_token_hidden, duration],
         outputs=[script_output, sound_design_output, music_suggestion_output]
     )
     
+    # Voice preview: generates a trimmed version of the script.
+    voice_preview_btn.click(
+        generate_voice_preview,
+        inputs=[script_output, tts_model],
+        outputs=voice_audio
+    )
+    
+    # Full voice generation using the complete script.
     voice_generate_btn.click(
         generate_voice,
         inputs=[script_output, tts_model],
@@ -417,4 +477,4 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Default()) as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch(debug=True)
+    demo.launch(debug=True)