Update app.py

app.py

@@ -8,6 +8,7 @@ from transformers import (
     AutoProcessor,
     MusicgenForConditionalGeneration,
 )
+from diffusers import StableDiffusionPipeline
 from scipy.io.wavfile import write
 from pydub import AudioSegment
 from dotenv import load_dotenv
@@ -86,24 +87,36 @@ def generate_voice(script: str, speaker: str = "default"):
 # Music Generation Function
 # ---------------------------------------------------------------------
 @spaces.GPU(duration=300)
-def generate_music(prompt: str, audio_length: int):
+def generate_music(prompt: str, audio_length: int, model_choice: str):
     try:
-        musicgen_model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small")
-        musicgen_processor = AutoProcessor.from_pretrained("facebook/musicgen-small")
+        if model_choice == "Stable Audio Open 1.0":
+            stable_pipeline = StableDiffusionPipeline.from_pretrained("stabilityai/stable-audio-open-1.0")
+            stable_pipeline.to("cuda" if torch.cuda.is_available() else "cpu")
 
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        musicgen_model.to(device)
+            audio = stable_pipeline(prompt, num_inference_steps=50, guidance_scale=7.5)
+            output_path = f"{tempfile.gettempdir()}/stable_generated_music.wav"
+            write(output_path, 44100, audio["sample"].cpu().numpy())
+            return output_path
 
-        inputs = musicgen_processor(text=[prompt], padding=True, return_tensors="pt").to(device)
-        outputs = musicgen_model.generate(**inputs, max_new_tokens=audio_length)
+        elif model_choice == "MusicGen":
+            musicgen_model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small")
+            musicgen_processor = AutoProcessor.from_pretrained("facebook/musicgen-small")
 
-        audio_data = outputs[0, 0].cpu().numpy()
-        normalized_audio = (audio_data / max(abs(audio_data)) * 32767).astype("int16")
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+            musicgen_model.to(device)
 
-        output_path = f"{tempfile.gettempdir()}/generated_music.wav"
-        write(output_path, 44100, normalized_audio)
+            inputs = musicgen_processor(text=[prompt], padding=True, return_tensors="pt").to(device)
+            outputs = musicgen_model.generate(**inputs, max_new_tokens=audio_length)
 
-        return output_path
+            audio_data = outputs[0, 0].cpu().numpy()
+            normalized_audio = (audio_data / max(abs(audio_data)) * 32767).astype("int16")
+
+            output_path = f"{tempfile.gettempdir()}/musicgen_generated_music.wav"
+            write(output_path, 44100, normalized_audio)
+            return output_path
+
+        else:
+            return "Invalid model choice!"
     except Exception as e:
         return f"Error generating music: {e}"
 
@@ -131,7 +144,7 @@ def blend_audio(voice_path: str, music_path: str, ducking: bool):
 # ---------------------------------------------------------------------
 with gr.Blocks() as demo:
     gr.Markdown("""
-        # 🎧 AI Promo Studio 🚀
+        # 🎵 AI Promo Studio 🚀
         Generate scripts, sound design, and music suggestions with ease.
     """)
 
@@ -154,6 +167,7 @@ with gr.Blocks() as demo:
                 outputs=[script_output, sound_design_output, music_suggestion_output],
             )
 
+        # Step 2: Generate Voice
         with gr.Tab("Step 2: Generate Voice"):
             with gr.Row():
                 speaker = gr.Textbox(label="Voice Style (optional)", placeholder="E.g., male, female, or neutral.")
@@ -167,16 +181,22 @@ with gr.Blocks() as demo:
                 outputs=[voice_output],
             )
 
+        # 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)
+                model_choice = gr.Dropdown(
+                    label="Select Music Generation Model",
+                    choices=["Stable Audio Open 1.0", "MusicGen"],
+                    value="Stable Audio Open 1.0"
+                )
 
             generate_music_button = gr.Button("Generate Music")
             music_output = gr.Audio(label="Generated Music", type="filepath")
 
             generate_music_button.click(
-                fn=lambda music_suggestion, audio_length: generate_music(music_suggestion, audio_length),
-                inputs=[music_suggestion_output, audio_length],
+                fn=lambda music_suggestion, audio_length, model_choice: generate_music(music_suggestion, audio_length, model_choice),
+                inputs=[music_suggestion_output, audio_length, model_choice],
                 outputs=[music_output],
             )