Update app.py

app.py

@@ -21,29 +21,52 @@ def load_models(model_size):
         model = whisper.load_model(model_size)
         return None, model
 
+from flask import jsonify
+import tempfile
+import os
+import io
+
 def process_audio(audio_file, num_speakers, model_size):
     transcriber, whisper_model = load_models(model_size)
 
-    with tempfile.NamedTemporaryFile(delete=True, suffix=".wav") as tmp:
+    # Ensure audio file is provided
+    if audio_file is None:
+        return jsonify({"error": "Audio file is required"}), 400
+
+    try:
         audio_file.seek(0)  # Reset the file pointer
-        tmp.write(audio_file.read())
-        tmp.flush()  # Ensure all data is written
-        tmp_path = tmp.name
-
-        try:
-            if transcriber:
-                result = transcriber(tmp_path)
-                transcription_text = result['text']
-            elif whisper_model:
-                result = whisper_model.transcribe(tmp_path)
-                transcription_text = result['text']
-
-            diarization_pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization-3.1", use_auth_token=os.getenv('HF_TOKEN'))
-            diarization = diarization_pipeline(tmp_path, min_speakers=num_speakers, max_speakers=5)
-            return transcription_text, diarization.get_timeline().json()
-        except Exception as e:
-            print("Error processing audio file or diarization:", e)
-            return None, None
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
+            tmp.write(audio_file.read())
+            tmp_path = tmp.name
+
+        # Initialize transcription_text
+        transcription_text = None
+
+        if transcriber:
+            result = transcriber(tmp_path)
+            transcription_text = result['text']
+        elif whisper_model:
+            result = whisper_model.transcribe(tmp_path)
+            transcription_text = result['text']
+
+        if transcription_text is None:
+            raise ValueError("No transcription results")
+
+        # Diarization process
+        diarization_pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization-3.1", use_auth_token=HF_TOKEN)
+        diarization = diarization_pipeline(tmp_path, min_speakers=num_speakers, max_speakers=5)
+        
+        os.remove(tmp_path)  # Cleanup the temporary file
+
+        return jsonify({
+            "transcription": transcription_text,
+            "diarization": diarization.get_timeline().json()
+        })
+
+    except Exception as e:
+        os.remove(tmp_path)  # Ensure to cleanup on error
+        return jsonify({"error": f"Error processing audio file: {e}"}), 500
+
 
 def gradio_interface(audio_file, num_speakers, model_size):
     transcription, diarization = process_audio(audio_file, num_speakers, model_size)