|
import tqdm |
|
import torch |
|
import random |
|
|
|
import typing as tp |
|
|
|
from torch import nn |
|
from torch.nn import functional as F |
|
from concurrent.futures import ThreadPoolExecutor |
|
|
|
from .demucs import Demucs |
|
from .hdemucs import HDemucs |
|
from .htdemucs import HTDemucs |
|
from .utils import center_trim |
|
|
|
|
|
Model = tp.Union[Demucs, HDemucs, HTDemucs] |
|
|
|
class DummyPoolExecutor: |
|
class DummyResult: |
|
def __init__(self, func, *args, **kwargs): |
|
self.func = func |
|
self.args = args |
|
self.kwargs = kwargs |
|
|
|
def result(self): |
|
return self.func(*self.args, **self.kwargs) |
|
|
|
def __init__(self, workers=0): |
|
pass |
|
|
|
def submit(self, func, *args, **kwargs): |
|
return DummyPoolExecutor.DummyResult(func, *args, **kwargs) |
|
|
|
def __enter__(self): |
|
return self |
|
|
|
def __exit__(self, exc_type, exc_value, exc_tb): |
|
return |
|
|
|
class BagOfModels(nn.Module): |
|
def __init__(self, models: tp.List[Model], weights: tp.Optional[tp.List[tp.List[float]]] = None, segment: tp.Optional[float] = None): |
|
super().__init__() |
|
assert len(models) > 0 |
|
first = models[0] |
|
|
|
for other in models: |
|
assert other.sources == first.sources |
|
assert other.samplerate == first.samplerate |
|
assert other.audio_channels == first.audio_channels |
|
|
|
if segment is not None: other.segment = segment |
|
|
|
self.audio_channels = first.audio_channels |
|
self.samplerate = first.samplerate |
|
self.sources = first.sources |
|
self.models = nn.ModuleList(models) |
|
|
|
if weights is None: weights = [[1.0 for _ in first.sources] for _ in models] |
|
else: |
|
assert len(weights) == len(models) |
|
|
|
for weight in weights: |
|
assert len(weight) == len(first.sources) |
|
|
|
self.weights = weights |
|
|
|
def forward(self, x): |
|
raise NotImplementedError("`apply_model`") |
|
|
|
|
|
class TensorChunk: |
|
def __init__(self, tensor, offset=0, length=None): |
|
total_length = tensor.shape[-1] |
|
assert offset >= 0 |
|
assert offset < total_length |
|
|
|
length = total_length - offset if length is None else min(total_length - offset, length) |
|
|
|
if isinstance(tensor, TensorChunk): |
|
self.tensor = tensor.tensor |
|
self.offset = offset + tensor.offset |
|
else: |
|
self.tensor = tensor |
|
self.offset = offset |
|
|
|
self.length = length |
|
self.device = tensor.device |
|
|
|
@property |
|
def shape(self): |
|
shape = list(self.tensor.shape) |
|
shape[-1] = self.length |
|
|
|
return shape |
|
|
|
def padded(self, target_length): |
|
delta = target_length - self.length |
|
total_length = self.tensor.shape[-1] |
|
assert delta >= 0 |
|
|
|
start = self.offset - delta // 2 |
|
end = start + target_length |
|
|
|
correct_start = max(0, start) |
|
correct_end = min(total_length, end) |
|
|
|
pad_left = correct_start - start |
|
pad_right = end - correct_end |
|
|
|
out = F.pad(self.tensor[..., correct_start:correct_end], (pad_left, pad_right)) |
|
|
|
assert out.shape[-1] == target_length |
|
|
|
return out |
|
|
|
|
|
def tensor_chunk(tensor_or_chunk): |
|
if isinstance(tensor_or_chunk, TensorChunk): return tensor_or_chunk |
|
else: |
|
assert isinstance(tensor_or_chunk, torch.Tensor) |
|
|
|
return TensorChunk(tensor_or_chunk) |
|
|
|
|
|
def apply_model(model, mix, shifts=1, split=True, overlap=0.25, transition_power=1.0, static_shifts=1, set_progress_bar=None, device=None, progress=False, num_workers=0, pool=None): |
|
global fut_length |
|
global bag_num |
|
global prog_bar |
|
|
|
device = mix.device if device is None else torch.device(device) |
|
|
|
if pool is None: pool = ThreadPoolExecutor(num_workers) if num_workers > 0 and device.type == "cpu" else DummyPoolExecutor() |
|
|
|
kwargs = { |
|
"shifts": shifts, |
|
"split": split, |
|
"overlap": overlap, |
|
"transition_power": transition_power, |
|
"progress": progress, |
|
"device": device, |
|
"pool": pool, |
|
"set_progress_bar": set_progress_bar, |
|
"static_shifts": static_shifts, |
|
} |
|
|
|
if isinstance(model, BagOfModels): |
|
estimates = 0 |
|
totals = [0] * len(model.sources) |
|
bag_num = len(model.models) |
|
fut_length = 0 |
|
prog_bar = 0 |
|
current_model = 0 |
|
|
|
for sub_model, weight in zip(model.models, model.weights): |
|
original_model_device = next(iter(sub_model.parameters())).device |
|
sub_model.to(device) |
|
fut_length += fut_length |
|
current_model += 1 |
|
|
|
out = apply_model(sub_model, mix, **kwargs) |
|
sub_model.to(original_model_device) |
|
|
|
for k, inst_weight in enumerate(weight): |
|
out[:, k, :, :] *= inst_weight |
|
totals[k] += inst_weight |
|
|
|
estimates += out |
|
del out |
|
|
|
for k in range(estimates.shape[1]): |
|
estimates[:, k, :, :] /= totals[k] |
|
|
|
return estimates |
|
|
|
model.to(device) |
|
model.eval() |
|
|
|
assert transition_power >= 1 |
|
|
|
batch, channels, length = mix.shape |
|
|
|
if shifts: |
|
kwargs["shifts"] = 0 |
|
max_shift = int(0.5 * model.samplerate) |
|
mix = tensor_chunk(mix) |
|
padded_mix = mix.padded(length + 2 * max_shift) |
|
out = 0 |
|
|
|
for _ in range(shifts): |
|
offset = random.randint(0, max_shift) |
|
shifted = TensorChunk(padded_mix, offset, length + max_shift - offset) |
|
shifted_out = apply_model(model, shifted, **kwargs) |
|
out += shifted_out[..., max_shift - offset :] |
|
|
|
out /= shifts |
|
|
|
return out |
|
elif split: |
|
kwargs["split"] = False |
|
out = torch.zeros(batch, len(model.sources), channels, length, device=mix.device) |
|
sum_weight = torch.zeros(length, device=mix.device) |
|
segment = int(model.samplerate * model.segment) |
|
stride = int((1 - overlap) * segment) |
|
offsets = range(0, length, stride) |
|
|
|
weight = torch.cat([torch.arange(1, segment // 2 + 1, device=device), torch.arange(segment - segment // 2, 0, -1, device=device)]) |
|
assert len(weight) == segment |
|
|
|
weight = (weight / weight.max()) ** transition_power |
|
futures = [] |
|
|
|
for offset in offsets: |
|
chunk = TensorChunk(mix, offset, segment) |
|
future = pool.submit(apply_model, model, chunk, **kwargs) |
|
futures.append((future, offset)) |
|
offset += segment |
|
|
|
if progress: futures = tqdm.tqdm(futures) |
|
|
|
for future, offset in futures: |
|
if set_progress_bar: |
|
fut_length = len(futures) * bag_num * static_shifts |
|
prog_bar += 1 |
|
set_progress_bar(0.1, (0.8 / fut_length * prog_bar)) |
|
|
|
chunk_out = future.result() |
|
chunk_length = chunk_out.shape[-1] |
|
|
|
out[..., offset : offset + segment] += (weight[:chunk_length] * chunk_out).to(mix.device) |
|
sum_weight[offset : offset + segment] += weight[:chunk_length].to(mix.device) |
|
|
|
assert sum_weight.min() > 0 |
|
|
|
out /= sum_weight |
|
|
|
return out |
|
else: |
|
valid_length = model.valid_length(length) if hasattr(model, "valid_length") else length |
|
|
|
mix = tensor_chunk(mix) |
|
padded_mix = mix.padded(valid_length).to(device) |
|
|
|
with torch.no_grad(): |
|
out = model(padded_mix) |
|
|
|
return center_trim(out, length) |
|
|
|
|
|
def demucs_segments(demucs_segment, demucs_model): |
|
if demucs_segment == "Default": |
|
segment = None |
|
|
|
if isinstance(demucs_model, BagOfModels): |
|
if segment is not None: |
|
for sub in demucs_model.models: |
|
sub.segment = segment |
|
else: |
|
if segment is not None: sub.segment = segment |
|
else: |
|
try: |
|
segment = int(demucs_segment) |
|
|
|
if isinstance(demucs_model, BagOfModels): |
|
if segment is not None: |
|
for sub in demucs_model.models: |
|
sub.segment = segment |
|
else: |
|
if segment is not None: sub.segment = segment |
|
except: |
|
segment = None |
|
|
|
if isinstance(demucs_model, BagOfModels): |
|
if segment is not None: |
|
for sub in demucs_model.models: |
|
sub.segment = segment |
|
else: |
|
if segment is not None: sub.segment = segment |
|
|
|
return demucs_model |
