Update app.py

app.py

@@ -234,7 +234,25 @@ def infer(image_in, chosen_model):
         gr.Info("Now calling Mustango for music...")
         music_o = get_mustango(musical_prompt)
     
-    return musical_prompt, music_o
+    return musical_prompt, gr.update(visible=True), music_o
+
+def retry(chosen_model, caption):
+    musical_prompt = caption
+
+    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)
+    elif chosen_model == "Riffusion" :
+        gr.Info("Now calling Riffusion for music...")
+        music_o = get_riffusion(musical_prompt)
+    elif chosen_model == "Mustango" :
+        gr.Info("Now calling Mustango for music...")
+        music_o = get_mustango(musical_prompt)
+
+    return music_o
 
 demo_title = "Image to Music V2"
 description = "Get music from a picture"
@@ -278,8 +296,10 @@ with gr.Blocks(css=css) as demo:
             with gr.Column():
                 caption = gr.Textbox(
                     label = "Inspirational musical prompt",
-                    max_lines = 3
+                    max_lines = 3,
+                    interactive = False
                 )
+                retry_btn = gr.Button("Retry with edited prompt", visible=False)
                 result = gr.Audio(
                     label = "Music"
                 )
@@ -305,6 +325,12 @@ with gr.Blocks(css=css) as demo:
         outputs = check_status,
         queue = False
     )
+
+    retry_btn.click(
+        fn = retry,
+        inputs = [chosen_model, caption],
+        outputs = [result]
+    )
     
     submit_btn.click(
         fn = infer,
@@ -314,6 +340,7 @@ with gr.Blocks(css=css) as demo:
         ],
         outputs =[
             caption,
+            retry_btn
             result
         ]
     )