Update app.py

app.py

@@ -8,6 +8,8 @@ import os
 import warnings
 from transformers import logging
 import math
+import json
+from pyannote.audio import Pipeline
 
 # Suppress warnings
 warnings.filterwarnings("ignore")
@@ -58,7 +60,12 @@ def detect_language(audio_path):
     
     return max(langs, key=lambda x: x.prob).lang
 
-def transcribe_audio_stream(audio, model_name):
+def diarize_audio(wav_audio):
+    pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization")
+    diarization = pipeline(wav_audio)
+    return diarization
+
+def transcribe_audio_stream(audio, model_name, diarization):
     wav_audio = convert_audio_to_wav(audio)
     speech, rate = librosa.load(wav_audio, sr=16000)
     duration = len(speech) / rate
@@ -69,6 +76,7 @@ def transcribe_audio_stream(audio, model_name):
         
         chunk_duration = 30  # seconds
         
+        transcriptions = []
         for i in range(0, int(duration), chunk_duration):
             end = min(i + chunk_duration, duration)
             chunk = speech[int(i * rate):int(end * rate)]
@@ -78,19 +86,38 @@ def transcribe_audio_stream(audio, model_name):
             transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
             
             progress = min(100, (end / duration) * 100)
-            yield transcription, progress
+            timestamp = i
+            transcriptions.append((timestamp, transcription))
+            yield transcriptions, progress
     else:
         transcriber = pipeline("automatic-speech-recognition", model=model_name)
         
         chunk_duration = 10  # seconds
         
+        transcriptions = []
         for i in range(0, int(duration), chunk_duration):
             end = min(i + chunk_duration, duration)
             chunk = speech[int(i * rate):int(end * rate)]
             result = transcriber(chunk)
             
             progress = min(100, (end / duration) * 100)
-            yield result["text"], progress
+            timestamp = i
+            transcriptions.append((timestamp, result["text"]))
+            yield transcriptions, progress
+
+    # Merge diarization results with transcription
+    speaker_transcriptions = []
+    for segment in diarization.itertracks(yield_label=True):
+        start, end, speaker = segment
+        start_time = start / rate
+        end_time = end / rate
+        text_segment = ""
+        for ts, text in transcriptions:
+            if start_time <= ts <= end_time:
+                text_segment += text + " "
+        speaker_transcriptions.append((start_time, end_time, speaker, text_segment.strip()))
+
+    return speaker_transcriptions
 
 def detect_and_select_model(audio):
     wav_audio = convert_audio_to_wav(audio)
@@ -98,24 +125,38 @@ def detect_and_select_model(audio):
     model_options = MODELS.get(language, MODELS["en"])
     return language, model_options
 
+def save_transcription(transcriptions, file_format):
+    if file_format == "txt":
+        with open("transcription.txt", "w") as f:
+            for start, end, speaker, text in transcriptions:
+                f.write(f"[{start}-{end}] {speaker}: {text}\n")
+        return "transcription.txt"
+    elif file_format == "json":
+        with open("transcription.json", "w") as f:
+            json.dump(transcriptions, f)
+        return "transcription.json"
+
 def combined_interface(audio):
     try:
         language, model_options = detect_and_select_model(audio)
         selected_model = model_options[0]
         
-        yield language, model_options, selected_model, "", 0, "Initializing..."
+        yield language, model_options, selected_model, [], 0, "Initializing..."
         
-        full_transcription = ""
-        for partial_transcription, progress in transcribe_audio_stream(audio, selected_model):
-            full_transcription += partial_transcription + " "
+        wav_audio = convert_audio_to_wav(audio)
+        diarization = diarize_audio(wav_audio)
+        transcriptions = []
+        for partial_transcriptions, progress in transcribe_audio_stream(audio, selected_model, diarization):
+            transcriptions = partial_transcriptions
+            transcriptions_text = "\n".join([f"[{start}-{end}] {speaker}: {text}" for start, end, speaker, text in transcriptions])
             progress_int = math.floor(progress)
             status = f"Transcribing... {progress_int}% complete"
-            yield language, model_options, selected_model, full_transcription.strip(), progress_int, status
+            yield language, model_options, selected_model, transcriptions_text, progress_int, status
         
         # Clean up temporary files
         os.remove("converted_audio.wav")
         
-        yield language, model_options, selected_model, full_transcription.strip(), 100, "Transcription complete!"
+        yield language, model_options, selected_model, transcriptions_text, 100, "Transcription complete!"
         
     except Exception as e:
         yield str(e), [], "", "An error occurred during processing.", 0, "Error"
@@ -129,12 +170,14 @@ iface = gr.Interface(
         gr.Textbox(label="Selected Model"),
         gr.Textbox(label="Transcription", lines=10),
         gr.Slider(minimum=0, maximum=100, label="Progress", interactive=False),
-        gr.Textbox(label="Status")
+        gr.Textbox(label="Status"),
+        gr.File(label="Download Transcription (TXT)", type="file", interactive=True, value="transcription.txt"),
+        gr.File(label="Download Transcription (JSON)", type="file", interactive=True, value="transcription.json")
     ],
-    title="Multilingual Audio Transcriber with Real-time Display and Progress Indicator",
-    description="Upload an audio file to detect the language, select the transcription model, and get the transcription in real-time. Optimized for Spanish, English, and Portuguese.",
+    title="Multilingual Audio Transcriber with Real-time Display, Timestamps, and Speaker Diarization",
+    description="Upload an audio file to detect the language, select the transcription model, and get the transcription with timestamps and speaker labels in real-time. Download the transcription as TXT or JSON. Optimized for Spanish, English, and Portuguese.",
     live=True
 )
 
 if __name__ == "__main__":
-    iface.queue().launch()
+    iface.queue().launch()