add video process tabs

gradio_app.py

@@ -6,22 +6,32 @@ import torchaudio.transforms as T
 import soundfile as sf
 import gradio as gr
 import spaces
+from moviepy.editor import VideoFileClip, AudioFileClip, CompositeAudioClip
 import look2hear.models
 
-# Setup device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # Load models
-dnr_model = look2hear.models.TIGERDNR.from_pretrained("JusperLee/TIGER-DnR", cache_dir="cache")
-dnr_model.to(device).eval()
-
-sep_model = look2hear.models.TIGER.from_pretrained("JusperLee/TIGER-speech", cache_dir="cache")
-sep_model.to(device).eval()
+dnr_model = look2hear.models.TIGERDNR.from_pretrained("JusperLee/TIGER-DnR", cache_dir="cache").to(device).eval()
+sep_model = look2hear.models.TIGER.from_pretrained("JusperLee/TIGER-speech", cache_dir="cache").to(device).eval()
 
 TARGET_SR = 16000
 MAX_SPEAKERS = 4
 
-# --- DnR Function ---
+def extract_audio_from_video(video_path):
+    video = VideoFileClip(video_path)
+    session_id = uuid.uuid4().hex[:8]
+    audio_path = f"temp_audio/{session_id}.wav"
+    os.makedirs("temp_audio", exist_ok=True)
+    video.audio.write_audiofile(audio_path, fps=44100, verbose=False, logger=None)
+    return audio_path, video
+
+def attach_audio_to_video(original_video, audio_path, out_path):
+    new_audio = AudioFileClip(audio_path)
+    new_video = original_video.set_audio(new_audio)
+    new_video.write_videofile(out_path, audio_codec='aac', verbose=False, logger=None)
+    return out_path
+
 @spaces.GPU()
 def separate_dnr(audio_file):
     audio, sr = torchaudio.load(audio_file)
@@ -30,22 +40,22 @@ def separate_dnr(audio_file):
     with torch.no_grad():
         dialog, effect, music = dnr_model(audio[None])
 
-    # Unique output folder
     session_id = uuid.uuid4().hex[:8]
     output_dir = os.path.join("output_dnr", session_id)
     os.makedirs(output_dir, exist_ok=True)
 
-    dialog_path = os.path.join(output_dir, "dialog.wav")
-    effect_path = os.path.join(output_dir, "effect.wav")
-    music_path = os.path.join(output_dir, "music.wav")
+    paths = {
+        "dialog": os.path.join(output_dir, "dialog.wav"),
+        "effect": os.path.join(output_dir, "effect.wav"),
+        "music": os.path.join(output_dir, "music.wav"),
+    }
 
-    torchaudio.save(dialog_path, dialog.cpu(), sr)
-    torchaudio.save(effect_path, effect.cpu(), sr)
-    torchaudio.save(music_path, music.cpu(), sr)
+    torchaudio.save(paths["dialog"], dialog.cpu(), sr)
+    torchaudio.save(paths["effect"], effect.cpu(), sr)
+    torchaudio.save(paths["music"], music.cpu(), sr)
 
-    return dialog_path, effect_path, music_path
+    return paths["dialog"], paths["effect"], paths["music"]
 
-# --- Speaker Separation Function ---
 @spaces.GPU()
 def separate_speakers(audio_path):
     waveform, original_sr = torchaudio.load(audio_path)
@@ -57,11 +67,8 @@ def separate_speakers(audio_path):
     audio_input = waveform.unsqueeze(0).to(device)
 
     with torch.no_grad():
-        ests_speech = sep_model(audio_input)
-
-    ests_speech = ests_speech.squeeze(0)
+        ests_speech = sep_model(audio_input).squeeze(0)
 
-    # Unique output folder
     session_id = uuid.uuid4().hex[:8]
     output_dir = os.path.join("output_sep", session_id)
     os.makedirs(output_dir, exist_ok=True)
@@ -69,8 +76,7 @@ def separate_speakers(audio_path):
     output_files = []
     for i in range(ests_speech.shape[0]):
         path = os.path.join(output_dir, f"speaker_{i+1}.wav")
-        audio_np = ests_speech[i].cpu().numpy()
-        sf.write(path, audio_np.T, TARGET_SR)  # Transpose only if shape is [T, C], usually not needed
+        sf.write(path, ests_speech[i].cpu().numpy(), TARGET_SR)
         output_files.append(path)
 
     updates = []
@@ -81,7 +87,57 @@ def separate_speakers(audio_path):
             updates.append(gr.update(value=None, visible=False))
     return updates
 
-# --- Gradio App ---
+@spaces.GPU()
+def separate_dnr_video(video_path):
+    audio_path, video = extract_audio_from_video(video_path)
+    dialog_path, effect_path, music_path = separate_dnr(audio_path)
+
+    session_id = uuid.uuid4().hex[:8]
+    output_dir = os.path.join("output_dnr_video", session_id)
+    os.makedirs(output_dir, exist_ok=True)
+
+    dialog_video = attach_audio_to_video(video, dialog_path, os.path.join(output_dir, "dialog_video.mp4"))
+    effect_video = attach_audio_to_video(video, effect_path, os.path.join(output_dir, "effect_video.mp4"))
+    music_video = attach_audio_to_video(video, music_path, os.path.join(output_dir, "music_video.mp4"))
+
+    return dialog_video, effect_video, music_video
+
+@spaces.GPU()
+def separate_speakers_video(video_path):
+    audio_path, video = extract_audio_from_video(video_path)
+
+    waveform, original_sr = torchaudio.load(audio_path)
+    if original_sr != TARGET_SR:
+        waveform = T.Resample(orig_freq=original_sr, new_freq=TARGET_SR)(waveform)
+
+    if waveform.dim() == 1:
+        waveform = waveform.unsqueeze(0)
+    audio_input = waveform.unsqueeze(0).to(device)
+
+    with torch.no_grad():
+        ests_speech = sep_model(audio_input).squeeze(0)
+
+    session_id = uuid.uuid4().hex[:8]
+    output_dir = os.path.join("output_sep_video", session_id)
+    os.makedirs(output_dir, exist_ok=True)
+
+    output_videos = []
+    for i in range(ests_speech.shape[0]):
+        path = os.path.join(output_dir, f"speaker_{i+1}.wav")
+        sf.write(path, ests_speech[i].cpu().numpy(), TARGET_SR)
+        video_path = os.path.join(output_dir, f"speaker_{i+1}_video.mp4")
+        attach_audio_to_video(video, path, video_path)
+        output_videos.append(video_path)
+
+    updates = []
+    for i in range(MAX_SPEAKERS):
+        if i < len(output_videos):
+            updates.append(gr.update(value=output_videos[i], visible=True, label=f"Speaker {i+1}"))
+        else:
+            updates.append(gr.update(value=None, visible=False))
+    return updates
+
+# --- Gradio UI ---
 with gr.Blocks() as demo:
     gr.Markdown("# TIGER: Time-frequency Interleaved Gain Extraction and Reconstruction for Efficient Speech Separation")
     gr.Markdown("TIGER is a lightweight model for speech separation which effectively extracts key acoustic features through frequency band-split, multi-scale and full-frequency-frame modeling.")
@@ -97,51 +153,39 @@ with gr.Blocks() as demo:
     			</a>	
             </div>
         """)
-    with gr.Tabs():
-        # --- Tab 1: DnR ---
-        with gr.Tab("Dialog/Effects/Music Separation (DnR)"):
-            gr.Markdown("### Separate Dialog, Effects, and Music from Mixed Audio")
-
-            dnr_input = gr.Audio(type="filepath", label="Upload Audio File")
-            dnr_button = gr.Button("Separate Audio")
 
+    with gr.Tabs():
+        with gr.Tab("Audio DnR"):
+            dnr_input = gr.Audio(type="filepath", label="Upload Audio")
+            dnr_btn = gr.Button("Separate")
             gr.Examples(
                 examples = ["./test/test_mixture_466.wav"],
                 inputs = dnr_input
             )
+            dnr_output = [gr.Audio(label=l) for l in ["Dialog", "Effects", "Music"]]
+            dnr_btn.click(separate_dnr, inputs=dnr_input, outputs=dnr_output)
 
-            dnr_output_dialog = gr.Audio(label="Dialog", type="filepath")
-            dnr_output_effect = gr.Audio(label="Effects", type="filepath")
-            dnr_output_music = gr.Audio(label="Music", type="filepath")
-
-            dnr_button.click(
-                fn=separate_dnr,
-                inputs=dnr_input,
-                outputs=[dnr_output_dialog, dnr_output_effect, dnr_output_music]
-            )
-
-        # --- Tab 2: Speaker Separation ---
-        with gr.Tab("Speaker Separation"):
-            gr.Markdown("### Separate Individual Speakers from Mixed Speech")
-
+        with gr.Tab("Audio Speaker Separation"):
             sep_input = gr.Audio(type="filepath", label="Upload Speech Audio")
-            sep_button = gr.Button("Separate Speakers")
-
+            sep_btn = gr.Button("Separate Speakers")
             gr.Examples(
                 examples = ["./test/mix.wav"],
                 inputs = sep_input
             )
+            sep_outputs = [gr.Audio(label=f"Speaker {i+1}", visible=(i==0)) for i in range(MAX_SPEAKERS)]
+            sep_btn.click(separate_speakers, inputs=sep_input, outputs=sep_outputs)
 
-            gr.Markdown("#### Separated Speakers")
-            sep_outputs = []
-            for i in range(MAX_SPEAKERS):
-                sep_outputs.append(gr.Audio(label=f"Speaker {i+1}", visible=(i == 0), interactive=False))
+        with gr.Tab("Video DnR"):
+            vdnr_input = gr.Video(label="Upload Video")
+            vdnr_btn = gr.Button("Separate Audio Tracks")
+            vdnr_output = [gr.Video(label=l) for l in ["Dialog Video", "Effects Video", "Music Video"]]
+            vdnr_btn.click(separate_dnr_video, inputs=vdnr_input, outputs=vdnr_output)
 
-            sep_button.click(
-                fn=separate_speakers,
-                inputs=sep_input,
-                outputs=sep_outputs
-            )
+        with gr.Tab("Video Speaker Separation"):
+            vsep_input = gr.Video(label="Upload Video")
+            vsep_btn = gr.Button("Separate Speakers")
+            vsep_outputs = [gr.Video(label=f"Speaker {i+1}", visible=(i==0)) for i in range(MAX_SPEAKERS)]
+            vsep_btn.click(separate_speakers_video, inputs=vsep_input, outputs=vsep_outputs)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()