Update app.py

app.py

@@ -9,18 +9,20 @@ from transformers import (
     MusicgenForConditionalGeneration,
 )
 from scipy.io.wavfile import write
+from TTS.api import TTS
 import tempfile
 from dotenv import load_dotenv
 import spaces
 
+# Load environment variables
 load_dotenv()
 hf_token = os.getenv("HF_TOKEN")
 
 # ---------------------------------------------------------------------
 # Load Llama 3 Pipeline with Zero GPU (Encapsulated)
 # ---------------------------------------------------------------------
-@spaces.GPU(duration=300)  # GPU allocation for 300 seconds
-def generate_script(user_prompt: str, model_id: str, token: str):
+@spaces.GPU(duration=300)
+def generate_script(user_prompt: str, duration: int, model_id: str, token: str):
     try:
         tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=token)
         model = AutoModelForCausalLM.from_pretrained(
@@ -34,7 +36,7 @@ def generate_script(user_prompt: str, model_id: str, token: str):
 
         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"Generate a concise, creative promo script for a {duration}-second ad, focusing on auditory elements and musical appeal."
         )
 
         combined_prompt = f"{system_prompt}\nUser concept: {user_prompt}\nRefined script:"
@@ -43,7 +45,6 @@ def generate_script(user_prompt: str, model_id: str, token: str):
     except Exception as e:
         return f"Error generating script: {e}"
 
-
 # ---------------------------------------------------------------------
 # Load MusicGen Model (Encapsulated)
 # ---------------------------------------------------------------------
@@ -53,49 +54,60 @@ def generate_audio(prompt: str, audio_length: int):
         musicgen_model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small")
         musicgen_processor = AutoProcessor.from_pretrained("facebook/musicgen-small")
 
-        # Ensure everything is on the same device (GPU or CPU)
         device = "cuda" if torch.cuda.is_available() else "cpu"
         musicgen_model.to(device)
 
         inputs = musicgen_processor(text=[prompt], padding=True, return_tensors="pt").to(device)
         outputs = musicgen_model.generate(**inputs, max_new_tokens=audio_length)
 
-        # Move outputs to CPU for further processing
         audio_data = outputs[0, 0].cpu().numpy()
-
-        # Normalize and save the audio file
         normalized_audio = (audio_data / max(abs(audio_data)) * 32767).astype("int16")
+
         output_path = f"{tempfile.gettempdir()}/generated_audio.wav"
         write(output_path, musicgen_model.config.audio_encoder.sampling_rate, normalized_audio)
-
         return output_path
     except Exception as e:
         return f"Error generating audio: {e}"
 
-
 # ---------------------------------------------------------------------
-# Gradio Interface Functions
+# Generate Voice-Over with Coqui XTTS-v2
 # ---------------------------------------------------------------------
-def interface_generate_script(user_prompt, llama_model_id):
-    return generate_script(user_prompt, llama_model_id, hf_token)
+@spaces.GPU(duration=300)
+def generate_voice(script: str, reference_audio: str, language: str):
+    try:
+        tts = TTS("tts_models/multilingual/multi-dataset/xtts_v2", gpu=torch.cuda.is_available())
+        output_path = f"{tempfile.gettempdir()}/voice_over.wav"
+        tts.tts_to_file(
+            text=script,
+            file_path=output_path,
+            speaker_wav=reference_audio,
+            language=language,
+        )
+        return output_path
+    except Exception as e:
+        return f"Error generating voice-over: {e}"
 
+# ---------------------------------------------------------------------
+# Interface Functions
+# ---------------------------------------------------------------------
+def interface_generate_script(user_prompt, duration, llama_model_id):
+    return generate_script(user_prompt, duration, llama_model_id, hf_token)
 
 def interface_generate_audio(script, audio_length):
     return generate_audio(script, audio_length)
 
+def interface_generate_voice(script, reference_audio, language):
+    return generate_voice(script, reference_audio, language)
 
 # ---------------------------------------------------------------------
 # Interface
 # ---------------------------------------------------------------------
 with gr.Blocks() as demo:
-    # Header
-    gr.Markdown(
-        """
-        # 🎧 AI-Powered Radio Imaging Studio 🚀  
-        ### Create stunning **radio promos** with **Llama 3** and **MusicGen**  
-        🔥 **Zero GPU** integration for efficiency and ease!  
-        """
-    )
+    gr.Markdown("""
+        # 🎧 All-in-One Radio Promo Studio 🚀  
+        ### Create professional scripts, soundscapes, and voice-overs in minutes!  
+        🔥 Powered by **Llama 3**, **MusicGen**, and **XTTS-v2**
+    """)
 
     # Script Generation Section
     gr.Markdown("## ✍️ Step 1: Generate Your Promo Script")
@@ -103,62 +115,72 @@ with gr.Blocks() as demo:
         user_prompt = gr.Textbox(
             label="🎤 Enter Promo Idea",
             placeholder="E.g., A 15-second energetic jingle for a morning talk show.",
-            lines=2,
-            info="Describe your promo idea clearly to generate a creative script."
+            lines=2
+        )
+        duration = gr.Dropdown(
+            label="⏳ Duration",
+            choices=["15", "30", "60"],
+            value="15",
+            info="Choose the duration of the promo (in seconds)."
         )
         llama_model_id = gr.Textbox(
-            label="🎿 Llama 3 Model ID",
-            value="meta-llama/Meta-Llama-3-8B-Instruct",
-            info="Enter the Hugging Face model ID for Llama 3."
+            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 Promo Script",
-        lines=4,
-        interactive=False,
-        info="Your generated promo script will appear here."
-    )
+    script_output = gr.Textbox(label="🖌️ Generated Promo Script", lines=4, interactive=False)
 
     # Audio Generation Section
-    gr.Markdown("## 🎵 Step 2: Generate Audio from Your Script")
+    gr.Markdown("## 🎵 Step 2: Generate Background Music")
     with gr.Row():
         audio_length = gr.Slider(
             label="🎶 Audio Length (tokens)",
             minimum=128,
             maximum=1024,
             step=64,
-            value=512,
-            info="Select the desired audio token length."  
+            value=512
         )
     generate_audio_button = gr.Button("Generate Audio 🎶")
-    audio_output = gr.Audio(
-        label="🎵 Generated Audio File",
-        type="filepath",
-        interactive=False
-    )
+    audio_output = gr.Audio(label="🎵 Generated Audio", type="filepath")
+
+    # Voice-Over Section
+    gr.Markdown("## 🎙️ Step 3: Generate Voice-Over")
+    with gr.Row():
+        reference_audio = gr.Audio(
+            label="🎤 Upload Reference Voice (6 seconds)",
+            type="filepath"
+        )
+        language = gr.Dropdown(
+            label="🌍 Language",
+            choices=["en", "es", "fr", "de", "it"],
+            value="en"
+        )
+    generate_voice_button = gr.Button("Generate Voice-Over 🎤")
+    voice_output = gr.Audio(label="🔊 Generated Voice-Over", type="filepath")
 
     # Footer
-    gr.Markdown(
-        """
+    gr.Markdown("""
         <br><hr>
         <p style="text-align: center; font-size: 0.9em;">
             Created with ❤️ by <a href="https://bilsimaging.com" target="_blank">bilsimaging.com</a>
         </p>
-        """,
-        elem_id="footer"
-    )
+    """)
 
     # Button Actions
     generate_script_button.click(
         fn=interface_generate_script,
-        inputs=[user_prompt, llama_model_id],
-        outputs=script_output,
+        inputs=[user_prompt, duration, llama_model_id],
+        outputs=script_output
     )
-
     generate_audio_button.click(
         fn=interface_generate_audio,
         inputs=[script_output, audio_length],
-        outputs=audio_output,
+        outputs=audio_output
+    )
+    generate_voice_button.click(
+        fn=interface_generate_voice,
+        inputs=[script_output, reference_audio, language],
+        outputs=voice_output
     )
 
 # ---------------------------------------------------------------------