Update app.py

app.py

@@ -3,11 +3,27 @@ import spaces
 import json
 import re
 from gradio_client import Client
+from moviepy.editor import VideoFileClip
+from moviepy.audio.AudioClip import AudioClip
 
-kosmos2_client = Client("https://ydshieh-kosmos-2.hf.space/")
+def extract_audio(video_in):
+    input_video = video_in
+    output_audio = 'audio.wav'
+    
+    # Open the video file and extract the audio
+    video_clip = VideoFileClip(input_video)
+    audio_clip = video_clip.audio
+    
+    # Save the audio as a .wav file
+    audio_clip.write_audiofile(output_audio, fps=44100)  # Use 44100 Hz as the sample rate for .wav files  
+    print("Audio extraction complete.")
 
-def get_caption(image_in):
+    return 'audio.wav'
 
+
+
+def get_caption(image_in):
+    kosmos2_client = Client("https://ydshieh-kosmos-2.hf.space/")
     kosmos2_result = kosmos2_client.predict(
         image_in,	# str (filepath or URL to image) in 'Test Image' Image component
         "Detailed",	# str in 'Description Type' Radio component
@@ -77,6 +93,22 @@ def get_magnet(prompt):
     )
     print(result)
     return result[1]
+
+def get_audioldm(prompt):
+    client = Client("https://haoheliu-audioldm2-text2audio-text2music.hf.space/")
+    result = client.predict(
+        prompt,	# str in 'Input text' Textbox component
+        "Low quality.",	# str in 'Negative prompt' Textbox component
+        10,	# int | float (numeric value between 5 and 15) in 'Duration (seconds)' Slider component
+        3.5,	# int | float (numeric value between 0 and 7) in 'Guidance scale' Slider component
+        45,	# int | float in 'Seed' Number component
+        3,	# int | float (numeric value between 1 and 5) in 'Number waveforms to generate' Slider component
+        fn_index=1
+    )
+    print(result)
+    audio_result = extract_audio(result)
+    return audio_result
+
     
 import re
 import torch
@@ -112,15 +144,19 @@ def get_musical_prompt(user_prompt):
     print(f"SUGGESTED Musical prompt: {cleaned_text}")
     return cleaned_text.lstrip("\n")
 
-def infer(image_in):
+def infer(image_in, chosen_model):
     gr.Info("Getting image caption with Kosmos2...")
     user_prompt = get_caption(image_in)
     
     gr.Info("Building a musical prompt according to the image caption ...")
     musical_prompt = get_musical_prompt(user_prompt)
 
-    gr.Info("Now calling MAGNet for music ...")
-    music_o = get_magnet(musical_prompt)
+    if chosen_model == "MAGNet" :
+        gr.Info("Now calling MAGNet for music...")
+        music_o = get_magnet(musical_prompt)
+    elif chosen_model == "AudioLDM-2" :
+        gr.Info("Now calling AudioLDM-2 for music...")
+        music_o = get_magnet(musical_prompt)
     
     return musical_prompt, music_o
 
@@ -149,10 +185,18 @@ with gr.Blocks(css=css) as demo:
                     type = "filepath",
                     elem_id = "image-in"
                 )
+                chosen_model = gr.Radio(
+                    label = "Choose a model",
+                    choices = [
+                        "MAGNet",
+                        "AudioLDM-2"
+                    ],
+                    value = "MAGNet"
+                )
                 submit_btn = gr.Button("Make music from my pic !")
             with gr.Column():
                 caption = gr.Textbox(
-                    label = "Musical prompt",
+                    label = "Inspirational musical prompt",
                     max_lines = 3
                 )
                 result = gr.Audio(
@@ -161,16 +205,16 @@ with gr.Blocks(css=css) as demo:
         with gr.Column():
             gr.Examples(
                 examples = [
-                    ["examples/monalisa.png"],
-                    ["examples/santa.png"],
-                    ["examples/ocean_poet.jpeg"],
-                    ["examples/winter_hiking.png"],
-                    ["examples/teatime.jpeg"],
-                    ["examples/news_experts.jpeg"],
-                    ["examples/chicken_adobo.jpeg"]
+                    ["examples/monalisa.png", "MAGNet"],
+                    ["examples/santa.png", "MAGNet"],
+                    ["examples/ocean_poet.jpeg", "MAGNet"],
+                    ["examples/winter_hiking.png", "MAGNet"],
+                    ["examples/teatime.jpeg", "MAGNet"],
+                    ["examples/news_experts.jpeg", "MAGNet"],
+                    ["examples/chicken_adobo.jpeg", "MAGNet"]
                 ],
                 fn = infer,
-                inputs = [image_in],
+                inputs = [image_in, chosen_model],
                 outputs = [caption, result],
                 cache_examples = False
             )
@@ -178,7 +222,8 @@ with gr.Blocks(css=css) as demo:
     submit_btn.click(
         fn = infer,
         inputs = [
-            image_in
+            image_in,
+            chosen_model
         ],
         outputs =[
             caption,
@@ -186,4 +231,4 @@ with gr.Blocks(css=css) as demo:
         ]
     )
 
-demo.queue().launch(show_api=False)
+demo.queue(max_size=16).launch(show_api=False)