go back to whisperx

handler.py

@@ -7,12 +7,9 @@ import torch
 #     process = subprocess.Popen(['pip', 'install', '--force-reinstall', 'onnxruntime-gpu'], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
 #     stdout, stderr = process.communicate()
 
-from faster_whisper import WhisperModel
-from pathlib import Path
+import whisperx
 import os
 
-# import nvidia.cublas.lib
-# import nvidia.cudnn.lib
 import time
 import json
 import base64
@@ -28,11 +25,6 @@ import logging
 
 logger = logging.getLogger(__name__)
 
-# logger.info(Path(nvidia.cublas.lib.__file__).parent)
-# logger.info(Path(nvidia.cudnn.lib.__file__).parent)
-# os.environ["LD_LIBRARY_PATH"] = ":".join(
-#     [Path(nvidia.cublas.lib.__file__).parent, Path(nvidia.cudnn.lib.__file__).parent]
-# )
 SAMPLE_RATE = 16000
 
 
@@ -196,8 +188,7 @@ class EndpointHandler:
     def __init__(self, path=""):
         # load the model
         device, batch_size, compute_type, whisper_model = whisper_config()
-
-        self.model = WhisperModel(whisper_model, device=device, compute_type=compute_type)
+        self.model = whisperx.load_model(whisper_model, device=device, compute_type=compute_type)
         # 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 ERIC
         # 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 OLD
         logger.info(f"Model {whisper_model} initialized")
@@ -253,13 +244,13 @@ class EndpointHandler:
         with open("/tmp/myfile.tmp", "wb") as w:
             w.write(inputs)
 
-        # audio_nparray = ffmpeg_load_audio('/tmp/myfile.tmp', sr=SAMPLE_RATE, mono=True, out_type=np.float32)
-        # audio_nparray = load_audio("/tmp/myfile.tmp", sr=SAMPLE_RATE)
+        audio_nparray = ffmpeg_load_audio("/tmp/myfile.tmp", sr=SAMPLE_RATE, mono=True, out_type=np.float32)
+        audio_nparray = load_audio("/tmp/myfile.tmp", sr=SAMPLE_RATE)
         # clean up
-        # os.remove("/tmp/myfile.tmp")
+        os.remove("/tmp/myfile.tmp")
 
-        # audio_nparray = ffmpeg_read(inputs, SAMPLE_RATE)
-        # audio_tensor= torch.from_numpy(audio_nparray)
+        audio_nparray = ffmpeg_read(inputs, SAMPLE_RATE)
+        audio_tensor = torch.from_numpy(audio_nparray)
 
         # get the end time
         et = time.time()
@@ -272,19 +263,16 @@ class EndpointHandler:
 
         # 2. transcribe
         logger.info("--------------- STARTING TRANSCRIPTION ------------------------")
-        # transcription = self.model.transcribe(audio_nparray, batch_size=batch_size, language=language)
-        segments, _ = self.model.transcribe("/tmp/myfile.tmp", beam_size=5)
+        transcription = self.model.transcribe(audio_nparray, batch_size=batch_size, language=language)
         if info:
-            # print(transcription["segments"][0:10000])  # before alignment
-            print(segments[0:10000])  # before alignment
-        # logger.info(transcription["segments"][0:10000])
-        logger.info("segments"[0:10000])
+            print(transcription["segments"][0:10000])  # before alignment
+        logger.info(transcription["segments"][0:10000])
 
         try:
-            first_text = segments[0]["text"]
+            first_text = transcription["segments"][0]["text"]
         except:
             logger.warning("No transcription")
-            return {"transcription": segments}
+            return {"transcription": transcription["segments"]}
 
         # get the execution time
         et = time.time()
@@ -338,4 +326,4 @@ class EndpointHandler:
 
         # results_json = json.dumps(results)
         # return {"results": results_json}
-        return {"transcription": [s["text"] for s in segments]}
+        return {"transcription": [s["text"] for s in transcription["segments"]]}