Update gradio_app.py

gradio_app.py

@@ -32,6 +32,32 @@ def attach_audio_to_video(original_video, audio_path, out_path):
     new_video.write_videofile(out_path, audio_codec='aac', verbose=False, logger=None)
     return out_path
 
+
+def separate_speakers_core(audio_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", session_id)
+    os.makedirs(output_dir, exist_ok=True)
+
+    output_files = []
+    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)
+        output_files.append(path)
+
+    return output_files
+
+
 @spaces.GPU()
 def separate_dnr(audio_file):
     audio, sr = torchaudio.load(audio_file)
@@ -58,27 +84,7 @@ def separate_dnr(audio_file):
 
 @spaces.GPU()
 def separate_speakers(audio_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", session_id)
-    os.makedirs(output_dir, exist_ok=True)
-
-    output_files = []
-    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)
-        output_files.append(path)
-
+    output_files = separate_speakers_core(audio_path)
     updates = []
     for i in range(MAX_SPEAKERS):
         if i < len(output_files):
@@ -102,40 +108,16 @@ def separate_dnr_video(video_path):
 
     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_files = []
-    for i in range(ests_speech.shape[0]):
-        separated_audio_path = os.path.join(output_dir, f"speaker_{i+1}.wav")
-        mono_audio = ests_speech[i].cpu().unsqueeze(0)  # Shape: [1, time]
-        torchaudio.save(separated_audio_path, mono_audio.contiguous(), TARGET_SR, format="wav", encoding="PCM_S")  # safest combo
-
-        # Attach audio back to video
-        out_video_path = os.path.join(output_dir, f"speaker_{i+1}.mp4")
-        attach_audio_to_video(video, separated_audio_path, out_video_path)
-        output_files.append(out_video_path)
-
-    return output_files + [None] * (MAX_SPEAKERS - len(output_files))
-
+def separate_speakers(audio_path):
+    output_files = separate_speakers_core(audio_path)
+    updates = []
+    for i in range(MAX_SPEAKERS):
+        if i < len(output_files):
+            updates.append(gr.update(value=output_files[i], visible=True, label=f"Speaker {i+1}"))
+        else:
+            updates.append(gr.update(value=None, visible=False))
+    return updates
 
 
 # --- Gradio UI ---