Update handler

handler.py

@@ -23,7 +23,7 @@ SAMPLE_RATE = 16000
 def whisper_config():
     device = "cuda" if torch.cuda.is_available() else "cpu"
     whisper_model = "large-v2"
-    batch_size = 16  # reduce if low on GPU mem
+    batch_size = 8  # reduce if low on GPU mem, 16 initailly
     # change to "int8" if low on GPU mem (may reduce accuracy)
     compute_type = "float16" if device == "cuda" else "int8"
     return device, batch_size, compute_type, whisper_model
@@ -158,6 +158,15 @@ def load_audio(file: str, sr: int = SAMPLE_RATE):
     return np.frombuffer(out, np.int16).flatten().astype(np.float32) / 32768.0
 
 
+def display_gpu_infos():
+    if not torch.cuda.is_available():
+        return "NO CUDA"
+    
+    infos = "torch.cuda.current_device(): " + torch.cuda.current_device() + "\n"
+    infos = infos + "torch.cuda.device(0): " +  torch.cuda.device(0) + "\n"
+    infos = infos + "torch.cuda.device_count(): " + torch.cuda.device_count() + "\n"
+    infos = infos + "torch.cuda.get_device_name(0): " + torch.cuda.get_device_name(0) + "\n"
+    return infos
 
 class EndpointHandler():
     def __init__(self, path=""):
@@ -187,6 +196,9 @@ class EndpointHandler():
                 logger.info(f"key: {x}, value: {data[x]}    ")
                 print(f"key: {x}, value: {data[x]}    ")
             
+            logger.info("--------------- CUDA ------------------------")
+            logger.info(display_gpu_infos())
+            
             # 1. process input
             inputs_encoded = data.pop("inputs", data)
             parameters = data.pop("parameters", None)
@@ -212,11 +224,13 @@ class EndpointHandler():
 
             # 2. transcribe
             device, batch_size, compute_type, whisper_model = whisper_config()
+            logger.info("--------------- STARTING TRANSCRIPTION ------------------------")
             transcription = self.model.transcribe(audio_nparray, batch_size=batch_size,language=language)
 
             logger.info(transcription["segments"])
             
             # 3. align
+            logger.info("--------------- STARTING ALIGNMENT ------------------------")
             # model_a, metadata = whisperx.load_align_model(
             #     language_code=result["language"], device=device)
             # transcription = whisperx.align(
@@ -225,6 +239,7 @@ class EndpointHandler():
             # print(transcription["segments"])
             
             # 4. Assign speaker labels
+            logger.info("--------------- STARTING DIARIZATION ------------------------")
             # add min/max number of speakers if known
             diarize_segments = self.diarize_model(audio_nparray)
             logger.info(diarize_segments)
@@ -234,6 +249,13 @@ class EndpointHandler():
             logger.info(diarized_transcription["segments"])  # segments are now assigned speaker IDs
             results.append({"diarized_transcription": diarized_transcription["segments"]})
 
+            if torch.cuda.is_available():
+                logger.info("--------------- GPU ------------------------")
+                logger.info(display_gpu_infos())
+                torch.cuda.empty_cache()
+                logger.info("--------------- GPU AFTER empty_cache ------------------------")
+                logger.info(display_gpu_infos())
+            
             # results_json = json.dumps(results)
             # return {"results": results_json}
             return results