| import base64 |
| import io |
| from typing import Any |
|
|
| import numpy as np |
| import soundfile as sf |
| import torch |
| import torchaudio |
|
|
|
|
| |
| if not hasattr(torchaudio, "list_audio_backends"): |
| torchaudio.list_audio_backends = lambda: ["soundfile"] |
|
|
| from speechbrain.inference.separation import SepformerSeparation |
|
|
|
|
| TARGET_SAMPLE_RATE = 8000 |
|
|
|
|
| class EndpointHandler: |
| def __init__(self, path: str = ""): |
| device = "cuda" if torch.cuda.is_available() else "cpu" |
| self.model = SepformerSeparation.from_hparams( |
| source=path or ".", |
| savedir=path or ".", |
| run_opts={"device": device}, |
| ) |
|
|
| def __call__(self, data: Any) -> dict: |
| audio_bytes = self._extract_audio_bytes(data) |
| waveform, sample_rate = self._load_audio(audio_bytes) |
|
|
| with torch.no_grad(): |
| est_sources = self.model.separate_batch(waveform.unsqueeze(0)) |
|
|
| est_sources = est_sources.squeeze(0).detach().cpu() |
| if est_sources.ndim == 1: |
| est_sources = est_sources.unsqueeze(-1) |
|
|
| outputs = [] |
| for idx in range(est_sources.shape[-1]): |
| source = est_sources[:, idx].numpy() |
| buffer = io.BytesIO() |
| sf.write(buffer, source, TARGET_SAMPLE_RATE, format="WAV") |
| outputs.append( |
| { |
| "speaker": idx, |
| "audio_base64": base64.b64encode(buffer.getvalue()).decode("utf-8"), |
| "sample_rate": TARGET_SAMPLE_RATE, |
| "mime_type": "audio/wav", |
| } |
| ) |
|
|
| return { |
| "num_speakers": len(outputs), |
| "sources": outputs, |
| } |
|
|
| def _extract_audio_bytes(self, data: Any) -> bytes: |
| if isinstance(data, (bytes, bytearray)): |
| return bytes(data) |
|
|
| if isinstance(data, dict): |
| payload = data.get("inputs", data) |
|
|
| if isinstance(payload, (bytes, bytearray)): |
| return bytes(payload) |
|
|
| if isinstance(payload, str): |
| return self._decode_base64_audio(payload) |
|
|
| if isinstance(payload, dict): |
| for key in ("audio", "audio_base64", "data"): |
| value = payload.get(key) |
| if isinstance(value, str): |
| return self._decode_base64_audio(value) |
|
|
| raise ValueError("Unsupported request format. Send raw audio bytes or a JSON body with base64 audio.") |
|
|
| def _decode_base64_audio(self, value: str) -> bytes: |
| if "," in value and value.startswith("data:"): |
| value = value.split(",", 1)[1] |
| return base64.b64decode(value) |
|
|
| def _load_audio(self, audio_bytes: bytes) -> tuple[torch.Tensor, int]: |
| waveform, sample_rate = sf.read(io.BytesIO(audio_bytes), dtype="float32", always_2d=True) |
| waveform = torch.from_numpy(waveform.T) |
|
|
| if waveform.shape[0] > 1: |
| waveform = waveform.mean(dim=0, keepdim=True) |
|
|
| if sample_rate != TARGET_SAMPLE_RATE: |
| resampler = torchaudio.transforms.Resample(sample_rate, TARGET_SAMPLE_RATE) |
| waveform = resampler(waveform) |
|
|
| return waveform.squeeze(0), TARGET_SAMPLE_RATE |
|
|
