Update gradio_app.py

gradio_app.py

@@ -104,13 +104,9 @@ def separate_dnr_video(video_path):
 
 @spaces.GPU()
 def separate_speakers_video(video_path):
-    # Extract audio
-    video = VideoFileClip(video_path)
-    audio_path = f"/tmp/{uuid.uuid4().hex}_audio.wav"
-    video.audio.write_audiofile(audio_path, fps=TARGET_SR, verbose=False, logger=None)
-
-    # Load and resample
+    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)
 
@@ -118,34 +114,23 @@ def separate_speakers_video(video_path):
         waveform = waveform.unsqueeze(0)
     audio_input = waveform.unsqueeze(0).to(device)
 
-    # Inference
     with torch.no_grad():
         ests_speech = sep_model(audio_input).squeeze(0)
 
-    # Output directory
     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]):
-        audio_np = ests_speech[i].cpu().numpy()
-        if audio_np.ndim == 1:
-            audio_np = audio_np[:, None]  # Ensure shape [samples, 1]
-
-        # Save separated audio
         separated_audio_path = os.path.join(output_dir, f"speaker_{i+1}.wav")
-        sf.write(separated_audio_path, audio_np, TARGET_SR)
-
-        # Combine with original video (no original audio)
-        output_video_path = os.path.join(output_dir, f"speaker_{i+1}_video.mp4")
-        new_audio = AudioFileClip(separated_audio_path)
-        new_video = video.set_audio(new_audio)
-        new_video.write_videofile(output_video_path, audio_codec="aac", verbose=False, logger=None)
+        audio_np = ests_speech[i].cpu().numpy()
+        sf.write(separated_audio_path, audio_np, TARGET_SR, format='WAV', subtype='PCM_16')
 
-        output_videos.append(output_video_path)
+        speaker_video_path = os.path.join(output_dir, f"speaker_{i+1}_video.mp4")
+        final_video = attach_audio_to_video(video, separated_audio_path, speaker_video_path)
+        output_videos.append(final_video)
 
-    # Pad with empty videos if less than MAX_SPEAKERS
     updates = []
     for i in range(MAX_SPEAKERS):
         if i < len(output_videos):
@@ -155,6 +140,7 @@ def separate_speakers_video(video_path):
     return updates
 
 
+
 # --- Gradio UI ---
 with gr.Blocks() as demo:
     gr.Markdown("# TIGER: Time-frequency Interleaved Gain Extraction and Reconstruction for Efficient Speech Separation")