Update app.py

app.py

@@ -5,9 +5,25 @@ import json
 import re
 from moviepy.editor import VideoFileClip
 from moviepy.audio.AudioClip import AudioClip
+import cv2
 
 hf_token = os.environ.get("HF_TKN")
 
+def extract_firstframe():
+    vidcap = cv2.VideoCapture('yourvideo.mp4') # replace yourvideo.mp4 with actual filename of your video
+    success,image = vidcap.read()
+    count = 0
+    while success:
+        if count == 0:
+            cv2.imwrite("first_frame.jpg", image)     # save first extracted frame as jpg file named first_frame.jpg
+        else:
+            break   # exit loop after saving first frame
+        success,image = vidcap.read()
+        print ('Read a new frame: ', success)
+        count += 1
+    print ("Done extracted first frame!")
+    return "first_frame.jpg"
+
 def extract_audio(video_in):
     input_video = video_in
     output_audio = 'audio.wav'
@@ -136,21 +152,27 @@ def get_tango(prompt):
     print(result)
     return result
 
-def infer(image_in, chosen_model):
+def blend_vsfx(video_in, audio_result):
+    audioClip = AudioFileClip(audio_result)
+    clip = VideoFileClip(video_in)
+    final_clip = clip.set_audio(audioClip)
+    final_clip.write_videofile('final_video_with_sound.mp4')
+    return "final_video_with_sound.mp4"
+
+def infer(video_in, chosen_model):
+    image_in = extract_firstframe(video_in)
     caption = get_caption(image_in)
     if chosen_model == "MAGNet" :
-        magnet_result = get_magnet(caption)
-        return magnet_result
+        audio_result = get_magnet(caption)
     elif chosen_model == "AudioLDM-2" : 
-        audioldm_result = get_audioldm(caption)
-        return audioldm_result
+        audio_result = get_audioldm(caption)
     elif chosen_model == "AudioGen" :
-        audiogen_result = get_audiogen(caption)
-        return audiogen_result
+        audio_result = get_audiogen(caption)
     elif chosen_model == "Tango" :
-        tango_result = get_tango(caption)
-        return tango_result
-
+        audio_result = get_tango(caption)
+    
+    final_res = blend_vsfx(video_in, audio_result)
+    return audio_result, final_res
 css="""
 #col-container{
     margin: 0 auto;
@@ -162,25 +184,28 @@ with gr.Blocks(css=css) as demo:
     with gr.Column(elem_id="col-container"):
         gr.HTML("""
         <h2 style="text-align: center;">
-            Image to SFX
+            Video to SoundFX
         </h2>
         <p style="text-align: center;">
-            Compare sound effects generation models from image caption.
+            Get sound effectsfor from video while comparing models from image caption.
         </p>
         """)
+
+        with gr.Row():
         
-        with gr.Column():
-            image_in = gr.Image(sources=["upload"], type="filepath", label="Image input", value="oiseau.png")
-            with gr.Row():
-                chosen_model = gr.Dropdown(label="Choose a model", choices=["MAGNet", "AudioLDM-2", "AudioGen", "Tango"], value="AudioLDM-2")
-                submit_btn = gr.Button("Submit")
-        with gr.Column():
-            audio_o = gr.Audio(label="Audio output")
+            with gr.Column():
+                video_in = gr.Video(sources=["upload"], type="filepath", label="Video input")
+                with gr.Row():
+                    chosen_model = gr.Dropdown(label="Choose a model", choices=["MAGNet", "AudioLDM-2", "AudioGen", "Tango"], value="AudioLDM-2")
+                    submit_btn = gr.Button("Submit")
+            with gr.Column():
+                audio_o = gr.Audio(label="Audio output")
+                video_o  gr.Video(label="Video with soundFX")
     
     submit_btn.click(
         fn=infer,
-        inputs=[image_in, chosen_model],
-        outputs=[audio_o],
+        inputs=[video_in_in, chosen_model],
+        outputs=[audio_o, video_o],
         concurrency_limit = 2
     )