146 lines
4.6 KiB
Python
146 lines
4.6 KiB
Python
import subprocess
|
|
import torch as th
|
|
import torchaudio as ta
|
|
|
|
from pathlib import Path
|
|
from typing import Optional, Callable, Dict, Tuple, Union
|
|
|
|
from .apply import apply_model, _replace_dict
|
|
from .audio import AudioFile, convert_audio
|
|
from .pretrained import get_model
|
|
|
|
|
|
class LoadAudioError(Exception):
|
|
pass
|
|
|
|
|
|
class LoadModelError(Exception):
|
|
pass
|
|
|
|
|
|
class _NotProvided:
|
|
pass
|
|
|
|
|
|
NotProvided = _NotProvided()
|
|
|
|
|
|
class Separator:
|
|
def __init__(
|
|
self,
|
|
model: str = "htdemucs",
|
|
repo: Optional[Path] = None,
|
|
device: str = "cuda" if th.cuda.is_available() else "cpu",
|
|
shifts: int = 1,
|
|
overlap: float = 0.25,
|
|
split: bool = True,
|
|
segment: Optional[int] = None,
|
|
jobs: int = 0,
|
|
progress: bool = False,
|
|
callback: Optional[Callable[[dict], None]] = None,
|
|
callback_arg: Optional[dict] = None,
|
|
):
|
|
|
|
self._name = model
|
|
self._repo = repo
|
|
self._load_model()
|
|
self.update_parameter(device=device, shifts=shifts, overlap=overlap, split=split,
|
|
segment=segment, jobs=jobs, progress=progress, callback=callback,
|
|
callback_arg=callback_arg)
|
|
|
|
def update_parameter(
|
|
self,
|
|
device: Union[str, _NotProvided] = NotProvided,
|
|
shifts: Union[int, _NotProvided] = NotProvided,
|
|
overlap: Union[float, _NotProvided] = NotProvided,
|
|
split: Union[bool, _NotProvided] = NotProvided,
|
|
segment: Optional[Union[int, _NotProvided]] = NotProvided,
|
|
jobs: Union[int, _NotProvided] = NotProvided,
|
|
progress: Union[bool, _NotProvided] = NotProvided,
|
|
callback: Optional[
|
|
Union[Callable[[dict], None], _NotProvided]
|
|
] = NotProvided,
|
|
callback_arg: Optional[Union[dict, _NotProvided]] = NotProvided,
|
|
):
|
|
|
|
if not isinstance(device, _NotProvided):
|
|
self._device = device
|
|
if not isinstance(shifts, _NotProvided):
|
|
self._shifts = shifts
|
|
if not isinstance(overlap, _NotProvided):
|
|
self._overlap = overlap
|
|
if not isinstance(split, _NotProvided):
|
|
self._split = split
|
|
if not isinstance(segment, _NotProvided):
|
|
self._segment = segment
|
|
if not isinstance(jobs, _NotProvided):
|
|
self._jobs = jobs
|
|
if not isinstance(progress, _NotProvided):
|
|
self._progress = progress
|
|
if not isinstance(callback, _NotProvided):
|
|
self._callback = callback
|
|
if not isinstance(callback_arg, _NotProvided):
|
|
self._callback_arg = callback_arg
|
|
|
|
def _load_model(self):
|
|
self._model = get_model(name=self._name, repo=self._repo)
|
|
if self._model is None:
|
|
raise LoadModelError("Failed to load model")
|
|
self._audio_channels = self._model.audio_channels
|
|
self._samplerate = self._model.samplerate
|
|
|
|
|
|
def _load_audio(self, track: Dict):
|
|
if track is not None:
|
|
wav = track["waveform"].squeeze()
|
|
sr = track["sample_rate"]
|
|
|
|
wav = convert_audio(wav, sr, self._samplerate, self._audio_channels)
|
|
|
|
return wav
|
|
|
|
def separate_tensor(
|
|
self, wav: th.Tensor, sr: Optional[int] = None
|
|
) -> Tuple[th.Tensor, Dict[str, th.Tensor]]:
|
|
if sr is not None and sr != self.samplerate:
|
|
wav = convert_audio(wav, sr, self._samplerate, self._audio_channels)
|
|
ref = wav.mean(0)
|
|
wav -= ref.mean()
|
|
wav /= ref.std() + 1e-8
|
|
out = apply_model(
|
|
self._model,
|
|
wav[None],
|
|
segment=self._segment,
|
|
shifts=self._shifts,
|
|
split=self._split,
|
|
overlap=self._overlap,
|
|
device=self._device,
|
|
num_workers=self._jobs,
|
|
callback=self._callback,
|
|
callback_arg=_replace_dict(
|
|
self._callback_arg, ("audio_length", wav.shape[1])
|
|
),
|
|
progress=self._progress,
|
|
)
|
|
if out is None:
|
|
raise KeyboardInterrupt
|
|
out *= ref.std() + 1e-8
|
|
out += ref.mean()
|
|
wav *= ref.std() + 1e-8
|
|
wav += ref.mean()
|
|
return (wav, dict(zip(self._model.sources, out[0])))
|
|
|
|
def separate_audio_file(self, track:Dict):
|
|
return self.separate_tensor(self._load_audio(track), self.samplerate)
|
|
|
|
@property
|
|
def samplerate(self):
|
|
return self._samplerate
|
|
|
|
@property
|
|
def audio_channels(self):
|
|
return self._audio_channels
|
|
|
|
@property
|
|
def model(self):
|
|
return self._model |