ericmattmann/whisperX-endpoint
3
1import subprocess2import torch3 4# if torch.cuda.is_available():5# process = subprocess.Popen(['pip', 'uninstall', 'onnxruntime'], stdout=subprocess.PIPE, stderr=subprocess.PIPE)6# stdout, stderr = process.communicate()7# process = subprocess.Popen(['pip', 'install', '--force-reinstall', 'onnxruntime-gpu'], stdout=subprocess.PIPE, stderr=subprocess.PIPE)8# stdout, stderr = process.communicate()9 10import whisperx11import os, gc12 13import time14import json15import base6416import numpy as np17 18DEVNULL = open(os.devnull, "w")19 20 21# from transformers.pipelines.audio_utils import ffmpeg_read22from typing import Dict, List, Any23 24import logging25 26logger = logging.getLogger(__name__)27 28SAMPLE_RATE = 1600029 30 31def whisper_config():32 device = "cuda" if torch.cuda.is_available() else "cpu"33 whisper_model = "large-v3"34 batch_size = 48 if device == "cuda" else 135 compute_type = "float16" if device == "cuda" else "int8"36 return device, batch_size, compute_type, whisper_model37 38 39# From https://gist.github.com/kylemcdonald/85d70bf53e207bab377540# load_audio can not detect the input type41def ffmpeg_load_audio(filename, sr=44100, mono=False, normalize=True, in_type=np.int16, out_type=np.float32):42 channels = 1 if mono else 243 format_strings = {44 np.float64: "f64le",45 np.float32: "f32le",46 np.int16: "s16le",47 np.int32: "s32le",48 np.uint32: "u32le",49 }50 format_string = format_strings[in_type]51 command = [52 "ffmpeg",53 "-i",54 filename,55 "-f",56 format_string,57 "-acodec",58 "pcm_" + format_string,59 "-ar",60 str(sr),61 "-ac",62 str(channels),63 "-",64 ]65 p = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=DEVNULL, bufsize=4096)66 bytes_per_sample = np.dtype(in_type).itemsize67 frame_size = bytes_per_sample * channels68 chunk_size = frame_size * sr # read in 1-second chunks69 raw = b""70 with p.stdout as stdout:71 while True:72 data = stdout.read(chunk_size)73 if data:74 raw += data75 else:76 break77 audio = np.fromstring(raw, dtype=in_type).astype(out_type)78 if channels > 1:79 audio = audio.reshape((-1, channels)).transpose()80 if audio.size == 0:81 return audio, sr82 if issubclass(out_type, np.floating):83 if normalize:84 peak = np.abs(audio).max()85 if peak > 0:86 audio /= peak87 elif issubclass(in_type, np.integer):88 audio /= np.iinfo(in_type).max89 return audio90 91 92# FROM HuggingFace93def ffmpeg_read(bpayload: bytes, sampling_rate: int) -> np.array:94 """95 Helper function to read an audio file through ffmpeg.96 """97 ar = f"{sampling_rate}"98 ac = "1"99 format_for_conversion = "f32le"100 ffmpeg_command = [101 "ffmpeg",102 "-i",103 "pipe:0",104 "-ac",105 ac,106 "-ar",107 ar,108 "-f",109 format_for_conversion,110 "-hide_banner",111 "-loglevel",112 "quiet",113 "pipe:1",114 ]115 116 try:117 with subprocess.Popen(ffmpeg_command, stdin=subprocess.PIPE, stdout=subprocess.PIPE) as ffmpeg_process:118 output_stream = ffmpeg_process.communicate(bpayload)119 except FileNotFoundError as error:120 raise ValueError("ffmpeg was not found but is required to load audio files from filename") from error121 out_bytes = output_stream[0]122 audio = np.frombuffer(out_bytes, np.float32)123 if audio.shape[0] == 0:124 raise ValueError(125 "Soundfile is either not in the correct format or is malformed. Ensure that the soundfile has "126 "a valid audio file extension (e.g. wav, flac or mp3) and is not corrupted. If reading from a remote "127 "URL, ensure that the URL is the full address to **download** the audio file."128 )129 return audio130 131 132# FROM whisperX133def load_audio(file: str, sr: int = SAMPLE_RATE):134 """135 Open an audio file and read as mono waveform, resampling as necessary136 137 Parameters138 ----------139 file: str140 The audio file to open141 142 sr: int143 The sample rate to resample the audio if necessary144 145 Returns146 -------147 A NumPy array containing the audio waveform, in float32 dtype.148 """149 try:150 # Launches a subprocess to decode audio while down-mixing and resampling as necessary.151 # Requires the ffmpeg CLI to be installed.152 cmd = [153 "ffmpeg",154 "-nostdin",155 "-threads",156 "0",157 "-i",158 file,159 "-f",160 "s16le",161 "-ac",162 "1",163 "-acodec",164 "pcm_s16le",165 "-ar",166 str(sr),167 "-",168 ]169 out = subprocess.run(cmd, capture_output=True, check=True).stdout170 except subprocess.CalledProcessError as e:171 raise RuntimeError(f"Failed to load audio: {e.stderr.decode()}") from e172 173 return np.frombuffer(out, np.int16).flatten().astype(np.float32) / 32768.0174 175 176def display_gpu_infos():177 if not torch.cuda.is_available():178 return "NO CUDA"179 180 infos = "torch.cuda.current_device(): " + str(torch.cuda.current_device()) + ", "181 infos = infos + "torch.cuda.device(0): " + str(torch.cuda.device(0)) + ", "182 infos = infos + "torch.cuda.device_count(): " + str(torch.cuda.device_count()) + ", "183 infos = infos + "torch.cuda.get_device_name(0): " + str(torch.cuda.get_device_name(0))184 return infos185 186 187class EndpointHandler:188 def __init__(self, path=""):189 # load the model190 device, batch_size, compute_type, whisper_model = whisper_config()191 self.model = whisperx.load_model(whisper_model, device=device, compute_type=compute_type, language="fr")192 # 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 ERIC193 # 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 OLD194 # logger.info(f"Model {whisper_model} initialized")195 196 self.diarize_model = whisperx.DiarizationPipeline(197 "pyannote/speaker-diarization-3.1", use_auth_token="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", device=device198 )199 200 logger.info(f"Model for diarization initialized")201 202 def __call__(self, data: Any) -> Dict[str, str]:203 """204 Args:205 data (:obj:):206 includes the deserialized audio file as bytes207 Return:208 A :obj:`dict`:. base64 encoded image209 """210 # get the start time211 st = time.time()212 213 logger.info("--------------- CONFIGURATION ------------------------")214 device, batch_size, compute_type, whisper_model = whisper_config()215 logger.info(display_gpu_infos())216 217 # 1. process input218 parameters = data.pop("parameters", None)219 options = data.pop("options", None)220 221 # OPTIONS are given as parameters222 info = options.get("info", False)223 transcribe = options.get("transcription", False)224 alignment = options.get("alignment", False)225 diarization = options.get("diarization", False)226 language = parameters.get("language", "fr")227 min_speakers = parameters.get("min_speakers", 2)228 max_speakers = parameters.get("max_speakers", 25)229 230 # for diarization without transcription, the transcription is given as input, so data is now a tuple (inputs, transcription)231 if transcribe:232 inputs_encoded = data.pop("inputs", data)233 elif diarization:234 inputs_encoded, transcription = data.pop("inputs", data)235 236 inputs = base64.b64decode(inputs_encoded)237 logger.info(f"inputs decoded.")238 # make a tmp file239 with open("/tmp/myfile.tmp", "wb") as w:240 w.write(inputs)241 logger.info(f"inputs saved.")242 243 audio_nparray = load_audio("/tmp/myfile.tmp", sr=SAMPLE_RATE)244 logger.info(f"inputs loaded as mono 16kHz.")245 # clean up246 os.remove("/tmp/myfile.tmp")247 logger.info(f"temp file removed.")248 249 et = time.time()250 elapsed_time = et - st251 252 logger.info(f"TIME for audio processing : {elapsed_time:.2f} seconds")253 if info:254 print(f"TIME for audio processing : {elapsed_time:.2f} seconds")255 256 # 2. transcribe257 if transcribe:258 gc.collect()259 torch.cuda.empty_cache()260 logger.info("--------------- STARTING TRANSCRIPTION ------------------------")261 transcription = self.model.transcribe(audio_nparray, batch_size=batch_size, language=language)262 if info:263 print(transcription["segments"][0:10_000]) # before alignment264 else:265 logger.info(transcription["segments"][0:1_000])266 267 try:268 first_text = transcription["segments"][0]["text"]269 except:270 logger.warning("No transcription")271 return {"transcription": transcription["segments"]}272 273 et = time.time()274 elapsed_time = et - st275 st = time.time()276 logger.info(f"TIME for audio transcription : {elapsed_time:.2f} seconds")277 if info:278 print(f"TIME for audio transcription : {elapsed_time:.2f} seconds")279 280 # 3. align281 if alignment:282 gc.collect()283 torch.cuda.empty_cache()284 logger.info("--------------- STARTING ALIGNMENT ------------------------")285 model_a, metadata = whisperx.load_align_model(language_code=transcription["language"], device=device)286 transcription = whisperx.align(287 transcription["segments"], model_a, metadata, audio_nparray, device, return_char_alignments=False288 )289 del model_a290 if info:291 print(transcription["segments"][0:10000])292 else:293 logger.info(transcription["segments"][0:1_000])294 295 et = time.time()296 elapsed_time = et - st297 st = time.time()298 logger.info(f"TIME for alignment : {elapsed_time:.2f} seconds")299 if info:300 print(f"TIME for alignment : {elapsed_time:.2f} seconds")301 302 # 4. Assign speaker labels303 if diarization:304 gc.collect()305 torch.cuda.empty_cache()306 logger.info("--------------- STARTING DIARIZATION ------------------------")307 if not transcription:308 logger.warning("No transcription to diarize")309 # add min/max number of speakers if known310 diarize_segments = self.diarize_model(audio_nparray, min_speakers=min_speakers, max_speakers=max_speakers)311 if info:312 print(diarize_segments)313 else:314 logger.info(diarize_segments)315 316 transcription = whisperx.assign_word_speakers(diarize_segments, transcription)317 318 et = time.time()319 elapsed_time = et - st320 st = time.time()321 logger.info(f"TIME for audio diarization : {elapsed_time:.2f} seconds")322 if info:323 print(f"TIME for audio diarization : {elapsed_time:.2f} seconds")324 325 # results_json = json.dumps(results)326 # return {"results": results_json}327 # return {"transcription": [s["text"] for s in transcription["segments"]]}328 gc.collect()329 torch.cuda.empty_cache()330 return transcription331