Update app.py

app.py

@@ -65,18 +65,19 @@ def diarize_audio(wav_audio):
     diarization = pipeline(wav_audio)
     return diarization
 
-def transcribe_audio_stream(audio, model_name, diarization):
+def transcribe_audio_stream(audio, model_name):
     wav_audio = convert_audio_to_wav(audio)
     speech, rate = librosa.load(wav_audio, sr=16000)
     duration = len(speech) / rate
     
+    transcriptions = []
+    
     if "whisper" in model_name:
         processor = WhisperProcessor.from_pretrained(model_name)
         model = WhisperForConditionalGeneration.from_pretrained(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)]
@@ -94,7 +95,6 @@ def transcribe_audio_stream(audio, model_name, diarization):
         
         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)]
@@ -105,7 +105,7 @@ def transcribe_audio_stream(audio, model_name, diarization):
             transcriptions.append((timestamp, result["text"]))
             yield transcriptions, progress
 
-    # Merge diarization results with transcription
+def merge_diarization_with_transcription(transcriptions, diarization, rate):
     speaker_transcriptions = []
     for segment in diarization.itertracks(yield_label=True):
         start, end, speaker = segment
@@ -116,7 +116,6 @@ def transcribe_audio_stream(audio, model_name, diarization):
             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):
@@ -127,39 +126,48 @@ def detect_and_select_model(audio):
 
 def save_transcription(transcriptions, file_format):
     if file_format == "txt":
-        with open("transcription.txt", "w") as f:
+        file_path = "/tmp/transcription.txt"
+        with open(file_path, "w") as f:
             for start, end, speaker, text in transcriptions:
-                f.write(f"[{start}-{end}] {speaker}: {text}\n")
-        return "transcription.txt"
+                f.write(f"[{start:.2f}-{end:.2f}] {speaker}: {text}\n")
+        return file_path
     elif file_format == "json":
-        with open("transcription.json", "w") as f:
+        file_path = "/tmp/transcription.json"
+        with open(file_path, "w") as f:
             json.dump(transcriptions, f)
-        return "transcription.json"
+        return file_path
 
 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...", None, None
         
         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):
+        
+        for partial_transcriptions, progress in transcribe_audio_stream(audio, selected_model):
             transcriptions = partial_transcriptions
-            transcriptions_text = "\n".join([f"[{start}-{end}] {speaker}: {text}" for start, end, speaker, text in transcriptions])
+            transcriptions_text = "\n".join([f"[{start}-{end}] {text}" for start, end, text in transcriptions])
             progress_int = math.floor(progress)
             status = f"Transcribing... {progress_int}% complete"
-            yield language, model_options, selected_model, transcriptions_text, progress_int, status
+            yield language, model_options, selected_model, transcriptions_text, progress_int, status, None, None
+        
+        rate = librosa.get_samplerate(wav_audio)
+        speaker_transcriptions = merge_diarization_with_transcription(transcriptions, diarization, rate)
+        transcriptions_text = "\n".join([f"[{start:.2f}-{end:.2f}] {speaker}: {text}" for start, end, speaker, text in speaker_transcriptions])
         
-        # Clean up temporary files
-        os.remove("converted_audio.wav")
+        txt_file_path = save_transcription(speaker_transcriptions, "txt")
+        json_file_path = save_transcription(speaker_transcriptions, "json")
+
+        os.remove(wav_audio)
         
-        yield language, model_options, selected_model, transcriptions_text, 100, "Transcription complete!"
+        yield language, model_options, selected_model, transcriptions_text, 100, "Transcription complete!", txt_file_path, json_file_path
         
     except Exception as e:
-        yield str(e), [], "", "An error occurred during processing.", 0, "Error"
+        yield str(e), [], "", "An error occurred during processing.", 0, "Error", None, None
 
 iface = gr.Interface(
     fn=combined_interface,
@@ -171,8 +179,8 @@ iface = gr.Interface(
         gr.Textbox(label="Transcription", lines=10),
         gr.Slider(minimum=0, maximum=100, label="Progress", interactive=False),
         gr.Textbox(label="Status"),
-        gr.File(label="Download Transcription (TXT)", type="filepath", interactive=True, value="transcription.txt"),
-        gr.File(label="Download Transcription (JSON)", type="filepath", interactive=True, value="transcription.json")
+        gr.File(label="Download Transcription (TXT)", type="filepath"),
+        gr.File(label="Download Transcription (JSON)", type="filepath")
     ],
     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.",