Update app.py

app.py

@@ -179,8 +179,27 @@ def infer(video_in, chosen_model):
     elif chosen_model == "Tango" :
         audio_result = get_tango(caption)
     
+    final_res = blend_vsfx(video_in, audio_result)
+    return gr.update(value=caption, interactive=True), gr.update(visible=True), audio_result, final_res
+
+def retry(edited_prompt, video_in, chosen_model):
+    image_in = extract_firstframe(video_in)
+    caption = edited_prompt
+    if chosen_model == "MAGNet" :
+        audio_result = get_magnet(caption)
+    elif chosen_model == "AudioLDM-2" : 
+        audio_result = get_audioldm(caption)
+    elif chosen_model == "AudioGen" :
+        audio_result = get_audiogen(caption)
+    elif chosen_model == "Tango" :
+        audio_result = get_tango(caption)
+    
     final_res = blend_vsfx(video_in, audio_result)
     return audio_result, final_res
+
+def refresh():
+    return gr.update(value="", interactive=False), gr.update(visible=False)
+
 css="""
 #col-container{
     margin: 0 auto;
@@ -205,8 +224,10 @@ with gr.Blocks(css=css) as demo:
                 video_in = gr.Video(sources=["upload"], label="Video input")
                 with gr.Row():
                     chosen_model = gr.Dropdown(label="Choose a model", choices=["MAGNet", "AudioLDM-2", "AudioGen", "Tango"], value="Tango")
-                    submit_btn = gr.Button("Submit")
+                    submit_btn = gr.Button("Submit", scale=0)
             with gr.Column():
+                caption_o = gr.Textbox(label="Scene caption", interactive=False)
+                retry_btn = gr.Button("Retry with edited prompt", visible=False)
                 audio_o = gr.Audio(label="Audio output")
                 video_o = gr.Video(label="Video with soundFX")
 
@@ -217,15 +238,26 @@ with gr.Blocks(css=css) as demo:
                 ["examples/chinese-new-year-dragon.mp4", "Tango"], 
                 ["examples/big-sur.mp4", "AudioLDM-2"]
             ],
-            fn = infer,
-            inputs = [video_in, chosen_model],
-            outputs = [audio_o, video_o],
-            cache_examples = False
+            inputs = [video_in, chosen_model]
         )
+
+    video_in.change(
+        fn = refresh,
+        inputs = None,
+        outputs = [caption_o, retry_btn],
+        queue = False
+    )
     
     submit_btn.click(
         fn=infer,
         inputs=[video_in, chosen_model],
+        outputs=[caption_o, retry_btn, audio_o, video_o],
+        concurrency_limit = 2
+    )
+
+    retry_btn.click(
+        fn=retry,
+        inputs=[caption_o, video_in, chosen_model],
         outputs=[audio_o, video_o],
         concurrency_limit = 2
     )