Update app.py

app.py

@@ -2,6 +2,10 @@ import gradio as gr
 from pyannote.audio import Pipeline
 import torch
 import os
+import numpy as np
+from pydub import AudioSegment
+import io
+import zipfile
 
 hf_token = os.getenv("HF_TOKEN")
 
@@ -9,26 +13,59 @@ hf_token = os.getenv("HF_TOKEN")
 pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization-3.1", use_auth_token=hf_token)
 pipeline.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
 
-def diarize(audio):
-    diarization = pipeline({"waveform": audio, "sample_rate": 16000})
-    speaker1_segments = []
-    speaker2_segments = []
-    for segment, _, speaker in diarization.itertracks(yield_label=True):
-        if speaker == 'SPEAKER_1':
-            speaker1_segments.append((segment.start, segment.end))
-        elif speaker == 'SPEAKER_2':
-            speaker2_segments.append((segment.start, segment.end))
-    return speaker1_segments, speaker2_segments
+def diarize_and_split(audio, sr):
+    # Convert to mono if stereo
+    if len(audio.shape) > 1:
+        audio = np.mean(audio, axis=1)
+    
+    # Perform diarization
+    diarization = pipeline({"waveform": torch.from_numpy(audio), "sample_rate": sr})
+    
+    # Create an AudioSegment from the numpy array
+    audio_segment = AudioSegment(
+        audio.tobytes(), 
+        frame_rate=sr,
+        sample_width=audio.dtype.itemsize, 
+        channels=1
+    )
+    
+    speaker_segments = {}
+    
+    for turn, _, speaker in diarization.itertracks(yield_label=True):
+        start_ms = int(turn.start * 1000)
+        end_ms = int(turn.end * 1000)
+        segment = audio_segment[start_ms:end_ms]
+        
+        if speaker not in speaker_segments:
+            speaker_segments[speaker] = []
+        speaker_segments[speaker].append(segment)
+    
+    # Create zip files for each speaker
+    zip_files = {}
+    for speaker, segments in speaker_segments.items():
+        zip_buffer = io.BytesIO()
+        with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
+            for i, segment in enumerate(segments):
+                segment_buffer = io.BytesIO()
+                segment.export(segment_buffer, format="wav")
+                zip_file.writestr(f"{speaker}_segment_{i}.wav", segment_buffer.getvalue())
+        
+        zip_buffer.seek(0)
+        zip_files[f"{speaker}.zip"] = zip_buffer.getvalue()
+    
+    return zip_files
 
-interface = gr.Interface(
-    fn=diarize,
+def process_audio(audio):
+    sr, audio_data = audio
+    zip_files = diarize_and_split(audio_data, sr)
+    return list(zip_files.values())
+
+iface = gr.Interface(
+    fn=process_audio,
     inputs=gr.Audio(type="numpy"),
-    outputs=[
-        gr.Textbox(label="Speaker 1 Segments (start, end)"),
-        gr.Textbox(label="Speaker 2 Segments (start, end)")
-    ],
-    title="Speaker Diarization",
-    description="Upload an audio file and get the segments where each speaker talks."
+    outputs=[gr.File(label="Speaker Zip Files") for _ in range(10)],  # Assuming max 10 speakers
+    title="Speaker Diarization and Audio Splitting",
+    description="Upload an audio file to split it into separate files for each speaker."
 )
 
-interface.launch()
+iface.launch()