Update app.py

app.py

@@ -10,8 +10,9 @@ from transformers import (
 )
 from scipy.io.wavfile import write
 from pydub import AudioSegment
-from dotenv import load_dotenv
+from pydub.playback import play
 import tempfile
+from dotenv import load_dotenv
 import spaces
 
 # Load environment variables
@@ -19,10 +20,10 @@ load_dotenv()
 hf_token = os.getenv("HF_TOKEN")
 
 # ---------------------------------------------------------------------
-# Generate Script
+# Script Generation Function
 # ---------------------------------------------------------------------
 @spaces.GPU(duration=300)
-def generate_script(user_prompt: str, model_id: str, token: str):
+def generate_script(user_prompt: str, model_id: str, token: str, duration: int):
     try:
         tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=token)
         model = AutoModelForCausalLM.from_pretrained(
@@ -35,21 +36,45 @@ def generate_script(user_prompt: str, model_id: str, token: str):
         llama_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
         system_prompt = (
-            "You are an expert radio imaging producer specializing in sound design and music. "
-            "Take the user's concept and craft a concise, creative promo script with a strong focus on auditory elements and musical appeal."
+            f"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, craft a concise, engaging promo script. "
+            f"Ensure the script fits within the time limit and suggest a matching music style that complements the theme."
         )
 
-        combined_prompt = f"{system_prompt}\nUser concept: {user_prompt}\nRefined script:"
+        combined_prompt = f"{system_prompt}\nUser concept: {user_prompt}\nRefined script and music suggestion:"
         result = llama_pipeline(combined_prompt, max_new_tokens=200, do_sample=True, temperature=0.9)
-        return result[0]["generated_text"].split("Refined script:")[-1].strip()
+
+        generated_text = result[0]["generated_text"].split("Refined script and music suggestion:")[-1].strip()
+        script, music_suggestion = generated_text.split("Music Suggestion:")
+        return script.strip(), music_suggestion.strip()
+    except Exception as e:
+        return f"Error generating script: {e}", None
+
+# ---------------------------------------------------------------------
+# Voice-Over Generation Function
+# ---------------------------------------------------------------------
+@spaces.GPU(duration=300)
+def generate_voice(script: str, speaker: str):
+    try:
+        # Replace with your chosen TTS model
+        tts_model = "coqui/XTTS-v2"
+        processor = AutoProcessor.from_pretrained(tts_model)
+        model = AutoModelForCausalLM.from_pretrained(tts_model)
+
+        inputs = processor(script, return_tensors="pt")
+        speech = model.generate(**inputs)
+
+        output_path = f"{tempfile.gettempdir()}/generated_voice.wav"
+        write(output_path, 22050, speech.cpu().numpy())
+        return output_path
     except Exception as e:
-        return f"Error generating script: {e}"
+        return f"Error generating voice-over: {e}"
 
 # ---------------------------------------------------------------------
-# Generate Music
+# Music Generation Function
 # ---------------------------------------------------------------------
 @spaces.GPU(duration=300)
-def generate_audio(prompt: str, audio_length: int):
+def generate_music(prompt: str, audio_length: int):
     try:
         musicgen_model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small")
         musicgen_processor = AutoProcessor.from_pretrained("facebook/musicgen-small")
@@ -62,116 +87,81 @@ def generate_audio(prompt: str, audio_length: int):
 
         audio_data = outputs[0, 0].cpu().numpy()
         normalized_audio = (audio_data / max(abs(audio_data)) * 32767).astype("int16")
+
         output_path = f"{tempfile.gettempdir()}/generated_music.wav"
-        write(output_path, musicgen_model.config.audio_encoder.sampling_rate, normalized_audio)
+        write(output_path, 44100, normalized_audio)
 
         return output_path
     except Exception as e:
-        return f"Error generating audio: {e}"
+        return f"Error generating music: {e}"
 
 # ---------------------------------------------------------------------
-# Generate Voice-Over (TTS)
+# Audio Blending Function with Ducking
 # ---------------------------------------------------------------------
-@spaces.GPU(duration=300)
-def generate_voice(script: str, language: str):
+def blend_audio(voice_path: str, music_path: str, ducking: bool):
     try:
-        tts_model = pipeline("text-to-speech", model="coqui/XTTS-v2")
-        tts_output = tts_model(script, language=language)
-
-        voice_path = f"{tempfile.gettempdir()}/generated_voice.wav"
-        with open(voice_path, "wb") as f:
-            f.write(tts_output["audio"])
-
-        return voice_path
-    except Exception as e:
-        return f"Error generating voice-over: {e}"
-
-# ---------------------------------------------------------------------
-# Mix Audio with Ducking Option
-# ---------------------------------------------------------------------
-def mix_audio(voice_file, music_file, output_file, ducking: bool):
-    try:
-        voice = AudioSegment.from_file(voice_file)
-        music = AudioSegment.from_file(music_file)
+        voice = AudioSegment.from_file(voice_path)
+        music = AudioSegment.from_file(music_path)
 
         if ducking:
-            music = music - 10  # Lower the volume of the music
-        combined = music.overlay(voice, position=0)
+            music = music - 10  # Lower music volume for ducking
 
-        combined.export(output_file, format="wav")
-        return output_file
+        combined = music.overlay(voice)
+        output_path = f"{tempfile.gettempdir()}/final_promo.wav"
+        combined.export(output_path, format="wav")
+
+        return output_path
     except Exception as e:
-        return f"Error mixing audio: {e}"
+        return f"Error blending audio: {e}"
 
 # ---------------------------------------------------------------------
-# Gradio Interface Functions
+# Gradio Interface
 # ---------------------------------------------------------------------
-def interface_generate_script(user_prompt, llama_model_id):
-    return generate_script(user_prompt, llama_model_id, hf_token)
+def process_all(user_prompt, llama_model_id, duration, audio_length, speaker, ducking):
+    script, music_suggestion = generate_script(user_prompt, llama_model_id, hf_token, duration)
+    if "Error" in script:
+        return script, None
 
-def interface_generate_audio(script, audio_length):
-    return generate_audio(script, audio_length)
+    voice_path = generate_voice(script, speaker)
+    if "Error" in voice_path:
+        return voice_path, None
 
-def interface_generate_voice(script, language):
-    return generate_voice(script, language)
+    music_path = generate_music(music_suggestion, audio_length)
+    if "Error" in music_path:
+        return music_path, None
 
-def interface_mix_audio(voice_file, music_file, ducking):
-    output_file = f"{tempfile.gettempdir()}/final_promo.wav"
-    return mix_audio(voice_file, music_file, output_file, ducking)
+    final_audio = blend_audio(voice_path, music_path, ducking)
+    return f"Script:\n{script}\n\nMusic Suggestion:\n{music_suggestion}", final_audio
 
-# ---------------------------------------------------------------------
-# Interface
-# ---------------------------------------------------------------------
 with gr.Blocks() as demo:
-    gr.Markdown(
-        """
-        # 🎙️ AI Radio Promo Maker 🚀  
-        ### Your one-stop solution for **scripts**, **voice-overs**, and **music**!  
-        🔥 **Zero GPU** integration powered by **Hugging Face** models.  
-        """
-    )
-
-    # Step 1: Generate Script
-    gr.Markdown("## ✍️ Step 1: Generate Your Promo Script")
-    with gr.Row():
-        user_prompt = gr.Textbox(label="Enter Promo Idea", placeholder="E.g., A 15-second energetic jingle.", lines=2)
-        llama_model_id = gr.Textbox(label="Llama 3 Model ID", value="meta-llama/Meta-Llama-3-8B-Instruct")
-    generate_script_button = gr.Button("Generate Script")
-    script_output = gr.Textbox(label="Generated Script", lines=4, interactive=False)
-
-    # Step 2: Generate Voice-Over
-    gr.Markdown("## 🎤 Step 2: Generate Voice-Over")
-    with gr.Row():
-        language = gr.Dropdown(label="Select Language", choices=["en", "es", "fr", "de"], value="en")
-    generate_voice_button = gr.Button("Generate Voice")
-    voice_output = gr.Audio(label="Generated Voice", type="filepath", interactive=False)
+    gr.Markdown("""
+        # 🎧 AI Promo Studio with Script, Voice, Music, and Mixing 🚀  
+        Generate fully mixed promos effortlessly with AI-driven tools for radio and media!
+    """)
 
-    # Step 3: Generate Music
-    gr.Markdown("## 🎵 Step 3: Generate Background Music")
     with gr.Row():
-        audio_length = gr.Slider(label="Audio Length (tokens)", minimum=128, maximum=1024, step=64, value=512)
-    generate_audio_button = gr.Button("Generate Music")
-    audio_output = gr.Audio(label="Generated Music", type="filepath", interactive=False)
+        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)
+        audio_length = gr.Slider(label="Music Length (tokens)", minimum=128, maximum=1024, step=64, value=512)
+        speaker = gr.Textbox(label="Voice Style (optional)", placeholder="E.g., male, female, or neutral.")
+        ducking = gr.Checkbox(label="Enable Ducking", value=True)
+
+    generate_button = gr.Button("Generate Full Promo")
+    script_output = gr.Textbox(label="Generated Script and Music Suggestion")
+    audio_output = gr.Audio(label="Final Promo Audio", type="filepath")
+
+    generate_button.click(
+        fn=process_all,
+        inputs=[user_prompt, llama_model_id, duration, audio_length, speaker, ducking],
+        outputs=[script_output, audio_output],
+    )
 
-    # Step 4: Mix Audio
-    gr.Markdown("## 🎶 Step 4: Mix Audio")
-    with gr.Row():
-        ducking = gr.Checkbox(label="Enable Ducking (lower background music volume)", value=True)
-    mix_audio_button = gr.Button("Mix Audio")
-    final_output = gr.Audio(label="Final Promo Audio", type="filepath", interactive=False)
-
-    # Button Actions
-    generate_script_button.click(interface_generate_script, inputs=[user_prompt, llama_model_id], outputs=script_output)
-    generate_voice_button.click(interface_generate_voice, inputs=[script_output, language], outputs=voice_output)
-    generate_audio_button.click(interface_generate_audio, inputs=[script_output, audio_length], outputs=audio_output)
-    mix_audio_button.click(interface_mix_audio, inputs=[voice_output, audio_output, ducking], outputs=final_output)
-
-    gr.Markdown(
-        """
+    gr.Markdown("""
         <hr>
-        <p style="text-align: center;">Created with ❤️ by <a href="https://bilsimaging.com" target="_blank">bilsimaging.com</a></p>
-        """
-    )
+        <p style="text-align: center; font-size: 0.9em;">
+            Created with ❤️ by <a href="https://bilsimaging.com" target="_blank">bilsimaging.com</a>
+        </p>
+    """)
 
-# Launch App
 demo.launch(debug=True)