| |
| |
| """Triton Python backend: the release's `load_audio` resampling (torchaudio 2.1.1's sinc resampler, as copied unchanged |
| into anuvani_diar.nn.resample, 2 CPU threads) for non-16 kHz input.""" |
| import numpy as np |
| import triton_python_backend_utils as pb_utils |
|
|
|
|
| class TritonPythonModel: |
| def initialize(self, args): |
| import torch |
| from anuvani_diar.nn.resample import resample |
| self.torch, self.resample = torch, resample |
| torch.set_num_threads(2) |
|
|
| def execute(self, requests): |
| t, resample = self.torch, self.resample |
| responses = [] |
| for req in requests: |
| try: |
| audio = pb_utils.get_input_tensor_by_name(req, "AUDIO").as_numpy() |
| sr = int(pb_utils.get_input_tensor_by_name(req, "SR").as_numpy()[0]) |
| t.set_num_threads(2) |
| out = resample(t.from_numpy(np.ascontiguousarray(audio)).unsqueeze(0), sr, 16000).squeeze(0).numpy().astype(np.float32) |
| responses.append(pb_utils.InferenceResponse(output_tensors=[pb_utils.Tensor("OUT", out)])) |
| except Exception as e: |
| responses.append(pb_utils.InferenceResponse(output_tensors=[], error=pb_utils.TritonError(f"{type(e).__name__}: {e}"))) |
| return responses |
|
|
