Update app.py

app.py

@@ -16,86 +16,174 @@ import spaces
 from TTS.api import TTS
 from TTS.utils.synthesizer import Synthesizer
 
-# Load environment variables
+# ---------------------------------------------------------------------
+# Load Environment Variables
+# ---------------------------------------------------------------------
 load_dotenv()
-hf_token = os.getenv("HF_TOKEN")
+HF_TOKEN = os.getenv("HF_TOKEN")
+
+# ---------------------------------------------------------------------
+# Global Model Caches
+# ---------------------------------------------------------------------
+# We store models/pipelines in global variables for reuse,
+# so they are only loaded once.
+LLAMA_PIPELINES = {}
+MUSICGEN_MODELS = {}
+
+# ---------------------------------------------------------------------
+# Helper Functions
+# ---------------------------------------------------------------------
+def get_llama_pipeline(model_id: str, token: str):
+    """
+    Returns a cached LLaMA pipeline if available; otherwise, loads it.
+    This significantly reduces loading time for repeated calls.
+    """
+    if model_id in LLAMA_PIPELINES:
+        return LLAMA_PIPELINES[model_id]
+
+    # Load new pipeline and store in cache
+    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-medium"):
+    """
+    Returns a cached MusicGen model if available; otherwise, loads it.
+    """
+    if model_key in MUSICGEN_MODELS:
+        return MUSICGEN_MODELS[model_key]
+
+    # Load new MusicGen model and store in cache
+    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
+
 
 # ---------------------------------------------------------------------
 # Script Generation Function
 # ---------------------------------------------------------------------
 @spaces.GPU(duration=100)
 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).
+    """
     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,
-        )
-        llama_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
+        text_pipeline = get_llama_pipeline(model_id, token)
 
         # System prompt with clear structure instructions
         system_prompt = (
-            f"You are an expert radio imaging producer specializing in sound design and music. "
+            "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: "
-            f"1. A concise voice-over script. Prefix this section with 'Voice-Over Script:'.\n"
-            f"2. Suggestions for sound design. Prefix this section with 'Sound Design Suggestions:'.\n"
-            f"3. Music styles or track recommendations. Prefix this section with 'Music Suggestions:'."
+            "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:"
-        result = llama_pipeline(combined_prompt, max_new_tokens=300, do_sample=True, temperature=0.8)
 
-        # Parsing output
-        generated_text = result[0]["generated_text"].split("Output:")[-1].strip()
+        # Use inference mode for efficient forward passes
+        with torch.inference_mode():
+            result = text_pipeline(
+                combined_prompt,
+                max_new_tokens=300,
+                do_sample=True,
+                temperature=0.8
+            )
+
+        # LLaMA pipeline returns a list of dicts with "generated_text"
+        generated_text = result[0]["generated_text"]
+
+        # Basic parsing to isolate everything after "Output:"
+        # (in case the model repeated your system prompt).
+        if "Output:" in generated_text:
+            generated_text = generated_text.split("Output:")[-1].strip()
+
+        # Extract sections based on known prefixes
+        voice_script = "No voice-over script found."
+        sound_design = "No sound design suggestions found."
+        music_suggestions = "No music suggestions found."
+
+        if "Voice-Over Script:" in generated_text:
+            parts = generated_text.split("Voice-Over Script:")
+            if len(parts) > 1:
+                # Everything after "Voice-Over Script:" up until next prefix
+                voice_script_part = parts[1]
+                voice_script = voice_script_part.split("Sound Design Suggestions:")[0].strip() \
+                    if "Sound Design Suggestions:" in voice_script_part else voice_script_part.strip()
 
-        # Extract sections based on prefixes
-        voice_script = generated_text.split("Voice-Over Script:")[1].split("Sound Design Suggestions:")[0].strip() if "Voice-Over Script:" in generated_text else "No voice-over script found."
-        sound_design = generated_text.split("Sound Design Suggestions:")[1].split("Music Suggestions:")[0].strip() if "Sound Design Suggestions:" in generated_text else "No sound design suggestions found."
-        music_suggestions = generated_text.split("Music Suggestions:")[1].strip() if "Music Suggestions:" in generated_text else "No music suggestions found."
+        if "Sound Design Suggestions:" in generated_text:
+            parts = generated_text.split("Sound Design Suggestions:")
+            if len(parts) > 1:
+                sound_design_part = parts[1]
+                sound_design = sound_design_part.split("Music Suggestions:")[0].strip() \
+                    if "Music Suggestions:" in sound_design_part else sound_design_part.strip()
+
+        if "Music Suggestions:" in generated_text:
+            parts = generated_text.split("Music Suggestions:")
+            if len(parts) > 1:
+                music_suggestions = parts[1].strip()
 
         return voice_script, sound_design, music_suggestions
+
     except Exception as e:
         return f"Error generating script: {e}", "", ""
 
+
 # ---------------------------------------------------------------------
 # Voice-Over Generation Function (Inactive)
 # ---------------------------------------------------------------------
 @spaces.GPU(duration=100)
 def generate_voice(script: str, speaker: str = "default"):
+    """
+    Placeholder for future voice-over generation functionality.
+    """
     try:
-        # Placeholder for inactive state
         return "Voice-over generation is currently inactive."
     except Exception as e:
         return f"Error: {e}"
 
+
 # ---------------------------------------------------------------------
-# Music Generation Function (facebook/musicgen-medium)
+# Music Generation Function
 # ---------------------------------------------------------------------
 @spaces.GPU(duration=100)
 def generate_music(prompt: str, audio_length: int):
+    """
+    Generates music from the 'facebook/musicgen-medium' model based on the prompt.
+    Returns the file path to the generated .wav file.
+    """
     try:
-        # Load facebook/musicgen-medium model
-        musicgen_model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-medium")
-        musicgen_processor = AutoProcessor.from_pretrained("facebook/musicgen-medium")
+        model_key = "facebook/musicgen-medium"
+        musicgen_model, musicgen_processor = get_musicgen_model(model_key)
 
-        # Move the model to the appropriate device (CUDA or CPU)
         device = "cuda" if torch.cuda.is_available() else "cpu"
-        musicgen_model.to(device)
-
-        # Prepare inputs
+        # Prepare input
         inputs = musicgen_processor(text=[prompt], padding=True, return_tensors="pt").to(device)
 
-        # Generate music
-        outputs = musicgen_model.generate(**inputs, max_new_tokens=audio_length)
+        # Generate music within inference mode
+        with torch.inference_mode():
+            outputs = musicgen_model.generate(**inputs, max_new_tokens=audio_length)
 
-        # Process audio data
         audio_data = outputs[0, 0].cpu().numpy()
+        # Normalize audio to int16 format
         normalized_audio = (audio_data / max(abs(audio_data)) * 32767).astype("int16")
 
-        # Save generated music to a file
+        # Save generated music to a temp file
         output_path = f"{tempfile.gettempdir()}/musicgen_medium_generated_music.wav"
         write(output_path, 44100, normalized_audio)
 
@@ -106,11 +194,13 @@ def generate_music(prompt: str, audio_length: int):
 
 
 # ---------------------------------------------------------------------
-# Audio Blending Function with Ducking (Inactive)
+# Audio Blending Function (Inactive)
 # ---------------------------------------------------------------------
 def blend_audio(voice_path: str, music_path: str, ducking: bool):
+    """
+    Placeholder for future audio blending functionality with optional ducking.
+    """
     try:
-        # Placeholder for inactive state
         return "Audio blending functionality is currently inactive."
     except Exception as e:
         return f"Error: {e}"
@@ -121,73 +211,92 @@ def blend_audio(voice_path: str, music_path: str, ducking: bool):
 # ---------------------------------------------------------------------
 with gr.Blocks() as demo:
     gr.Markdown("""
-        # 🎧 AI Promo Studio 🚀  
-        Welcome to **AI Promo Studio**, your one-stop solution for creating stunning and professional radio promos with ease!  
-        Whether you're a sound designer, radio producer, or content creator, our AI-driven tools, powered by advanced LLM Llama models, empower you to bring your vision to life in just a few steps.  
+    # 🎧 AI Promo Studio 🚀  
+    Welcome to **AI Promo Studio**, your one-stop solution for creating stunning and professional radio promos with ease!  
+    Whether you're a sound designer, radio producer, or content creator, our AI-driven tools, powered by advanced LLM Llama models, empower you to bring your vision to life in just a few steps.  
     """)
 
     with gr.Tabs():
         # Step 1: Generate Script
         with gr.Tab("Step 1: Generate Script"):
             with gr.Row():
-                user_prompt = gr.Textbox(label="Promo Idea", placeholder="E.g., A 30-second promo for a morning show.")
-                llama_model_id = gr.Textbox(label="Llama Model ID", value="meta-llama/Meta-Llama-3-8B-Instruct")
-                duration = gr.Slider(label="Duration (seconds)", minimum=15, maximum=60, step=15, value=30)
+                user_prompt = gr.Textbox(
+                    label="Promo Idea", 
+                    placeholder="E.g., A 30-second promo for a morning show...",
+                    lines=2
+                )
+                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
+                )
 
             generate_script_button = gr.Button("Generate Script")
-            script_output = gr.Textbox(label="Generated Voice-Over Script", lines=5)
-            sound_design_output = gr.Textbox(label="Sound Design Suggestions", lines=3)
-            music_suggestion_output = gr.Textbox(label="Music Suggestions", lines=3)
+            script_output = gr.Textbox(label="Generated 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, duration: generate_script(user_prompt, model_id, hf_token, duration),
+                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
+        # Step 2: Generate Voice (Inactive)
         with gr.Tab("Step 2: Generate Voice"):
             gr.Markdown("""
-                **Note:** Voice-over generation is currently inactive. 
-                This feature will be available in future updates!
+            **Note:** Voice-over generation is currently inactive. 
+            This feature will be available in future updates!
             """)
-           
 
         # Step 3: Generate Music
         with gr.Tab("Step 3: Generate Music"):
             with gr.Row():
-                audio_length = gr.Slider(label="Music Length (tokens)", minimum=128, maximum=1024, step=64, value=512)
-
+                audio_length = gr.Slider(
+                    label="Music Length (tokens)",
+                    minimum=128, 
+                    maximum=1024, 
+                    step=64, 
+                    value=512,
+                    info="Increase tokens for longer audio, but be mindful of inference time."
+                )
             generate_music_button = gr.Button("Generate Music")
-            music_output = gr.Audio(label="Generated Music", type="filepath")
+            music_output = gr.Audio(label="Generated Music (WAV)", type="filepath")
 
             generate_music_button.click(
-                fn=lambda music_suggestion, audio_length: generate_music(music_suggestion, audio_length),
+                fn=lambda music_suggestion, length: generate_music(music_suggestion, length),
                 inputs=[music_suggestion_output, audio_length],
                 outputs=[music_output],
             )
 
-        # Step 4: Blend Audio
+        # Step 4: Blend Audio (Inactive)
         with gr.Tab("Step 4: Blend Audio"):
             gr.Markdown("""
-                **Note:** Audio blending functionality is currently inactive. 
-                This feature will be available in future updates!
+            **Note:** Audio blending functionality is currently inactive. 
+            This feature will be available in future updates!
             """)
-            
 
+    # Footer / Credits
     gr.Markdown("""
-        <hr>
-        <p style="text-align: center; font-size: 0.9em;">
-            Created with ❤️ by <a href="https://bilsimaging.com" target="_blank">bilsimaging.com</a>
-        </p>
+    <hr>
+    <p style="text-align: center; font-size: 0.9em;">
+        Created with ❤️ by <a href="https://bilsimaging.com" target="_blank">bilsimaging.com</a>
+    </p>
     """)
     
-    # Add visitor badge HTML
+    # Visitor Badge
     gr.HTML("""
-        <a href="https://visitorbadge.io/status?path=https%3A%2F%2Fhuggingface.co%2Fspaces%2FBils%2Fradiogold">
-            <img src="https://api.visitorbadge.io/api/visitors?path=https%3A%2F%2Fhuggingface.co%2Fspaces%2FBils%2Fradiogold&countColor=%23263759" />
-        </a>
+    <a href="https://visitorbadge.io/status?path=https%3A%2F%2Fhuggingface.co%2Fspaces%2FBils%2Fradiogold">
+        <img src="https://api.visitorbadge.io/api/visitors?path=https%3A%2F%2Fhuggingface.co%2Fspaces%2FBils%2Fradiogold&countColor=%23263759" />
+    </a>
     """)
 
+# Launch the Gradio app
 demo.launch(debug=True)
-