Update handler.py

handler.py

@@ -109,12 +109,13 @@ class EndpointHandler():
             stream = torch.cuda.Stream()
             with torch.cuda.stream(stream):
                 try:
-                    diarize_segments = self.diarize_model(
+                    diarize_segments, _embeddings = self.diarize_model(
                         _audio,
                         min_speakers=parameters.min_speakers,
-                        max_speakers=parameters.max_speakers
+                        max_speakers=parameters.max_speakers,
+                        return_embeddings=True
                     )
-                    return diarize_segments
+                    return diarize_segments, _embeddings
                 except RuntimeError as e:
                     logger.error(f"Diarization inference error: {str(e)}")
                     raise HTTPException(status_code=400, detail=f"Diarization inference error: {str(e)}")
@@ -134,14 +135,13 @@ class EndpointHandler():
             diarization_future = executor.submit(run_diarization, audio)
 
             _result = align_future.result()
-            diarization_output = diarization_future.result()
+            diarization_output, embeddings = diarization_future.result()
         result = []
         if diarization_output is not None and _result:
             result = assign_word_speakers(
                 diarization_output,
                 _result,
             )
-        embeddings = diarization_output.speaker_embeddings
         # Final cleanup
         del diarization_output, segments, audio
         gc.collect()