File size: 1,082 Bytes
dd35111 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 | import torch, torchaudio
from diffsynth import load_model, ModelConfig
from diffsynth.models.demucs import HTDemucs
class AudioTrackSeparator(torch.nn.Module):
def __init__(self, torch_dtype=torch.float32, device="cuda", model_config=ModelConfig(model_id="DiffSynth-Studio/Demucs-Repackage", origin_file_pattern="model.safetensors")):
super().__init__()
model_config.download_if_necessary()
self.model = load_model(HTDemucs, model_config.path, torch_dtype=torch_dtype, device=device)
@torch.no_grad()
def __call__(self, audio, target_sample_rate=48000, **kwargs):
if isinstance(audio, str):
audio, sample_rate = torchaudio.load(audio)
else:
audio, sample_rate = audio
audio = audio.to(dtype=next(iter(self.model.parameters())).dtype, device=next(iter(self.model.parameters())).device)
vocals = self.model.extract_track(audio, sample_rate)
if target_sample_rate != 44100:
vocals = torchaudio.functional.resample(vocals, 44100, target_sample_rate)
return vocals
|