Spaces:
Running on Zero
Running on Zero
| import math | |
| import torch | |
| import torch as th | |
| from einops import rearrange | |
| from torch import nn | |
| from torch.nn import functional as F | |
| import fractions | |
| import torchaudio | |
| def spectro(x, n_fft, hop_length): | |
| *other, length = x.shape | |
| x = x.reshape(-1, length) | |
| z = th.stft(x, n_fft, hop_length, window=th.hann_window(n_fft).to(x), | |
| win_length=n_fft, normalized=True, center=True, | |
| return_complex=True, pad_mode="reflect") | |
| _, freqs, frame = z.shape | |
| return z.view(*other, freqs, frame) | |
| def ispectro(z, hop_length, length): | |
| *other, freqs, frames = z.shape | |
| n_fft = 2 * freqs - 2 | |
| z = z.view(-1, freqs, frames) | |
| x = th.istft(z, n_fft, hop_length, window=th.hann_window(n_fft).to(z.real), | |
| win_length=n_fft, normalized=True, length=length, center=True) | |
| return x.view(*other, x.shape[-1]) | |
| class LayerScale(nn.Module): | |
| def __init__(self, channels, init=0, channel_last=False): | |
| super().__init__() | |
| self.channel_last = channel_last | |
| self.scale = nn.Parameter(torch.full((channels,), float(init))) | |
| def forward(self, x): | |
| return self.scale * x if self.channel_last else self.scale[:, None] * x | |
| class DConv(nn.Module): | |
| def __init__(self, channels, compress=8, depth=2, init=1e-3, kernel=3): | |
| super().__init__() | |
| hidden = int(channels / compress) | |
| self.layers = nn.ModuleList() | |
| for d in range(depth): | |
| dilation = 2 ** d | |
| padding = dilation * (kernel // 2) | |
| self.layers.append(nn.Sequential( | |
| nn.Conv1d(channels, hidden, kernel, dilation=dilation, padding=padding), | |
| nn.GroupNorm(1, hidden), nn.GELU(), | |
| nn.Conv1d(hidden, 2 * channels, 1), | |
| nn.GroupNorm(1, 2 * channels), nn.GLU(1), | |
| LayerScale(channels, init), | |
| )) | |
| def forward(self, x): | |
| for layer in self.layers: | |
| x = x + layer(x) | |
| return x | |
| class ScaledEmbedding(nn.Module): | |
| def __init__(self, num_embeddings, embedding_dim, scale=10.0): | |
| super().__init__() | |
| self.embedding = nn.Embedding(num_embeddings, embedding_dim) | |
| weight = torch.cumsum(self.embedding.weight.data, dim=0) | |
| weight = weight / torch.arange(1, num_embeddings + 1).to(weight).sqrt()[:, None] | |
| self.embedding.weight.data[:] = weight / scale | |
| self.scale = scale | |
| def forward(self, x): | |
| return self.embedding(x) * self.scale | |
| class HEncLayer(nn.Module): | |
| def __init__(self, chin, chout, kernel_size=8, stride=4, freq=True, context=0, | |
| dconv_kw={}): | |
| super().__init__() | |
| self.freq = freq | |
| self.stride = stride | |
| pad = kernel_size // 4 | |
| klass = nn.Conv2d if freq else nn.Conv1d | |
| if freq: | |
| kernel_size, stride, pad = [kernel_size, 1], [stride, 1], [pad, 0] | |
| self.conv = klass(chin, chout, kernel_size, stride, pad) | |
| self.rewrite = klass(chout, 2 * chout, 1 + 2 * context, 1, context) | |
| self.dconv = DConv(chout, **dconv_kw) | |
| def forward(self, x): | |
| if not self.freq: | |
| if x.dim() == 4: | |
| B, C, Fr, T = x.shape | |
| x = x.view(B, -1, T) | |
| le = x.shape[-1] | |
| if le % self.stride: | |
| x = F.pad(x, (0, self.stride - (le % self.stride))) | |
| y = F.gelu(self.conv(x)) | |
| if self.freq: | |
| B, C, Fr, T = y.shape | |
| y = y.permute(0, 2, 1, 3).reshape(-1, C, T) | |
| y = self.dconv(y).view(B, Fr, C, T).permute(0, 2, 1, 3) | |
| else: | |
| y = self.dconv(y) | |
| return F.glu(self.rewrite(y), dim=1) | |
| class HDecLayer(nn.Module): | |
| def __init__(self, chin, chout, last=False, kernel_size=8, stride=4, freq=True, | |
| context=1, dconv_kw={}): | |
| super().__init__() | |
| pad = kernel_size // 4 | |
| self.pad = pad | |
| self.last = last | |
| self.freq = freq | |
| self.chin = chin | |
| klass = nn.Conv2d if freq else nn.Conv1d | |
| klass_tr = nn.ConvTranspose2d if freq else nn.ConvTranspose1d | |
| ks, st = ([kernel_size, 1], [stride, 1]) if freq else (kernel_size, stride) | |
| self.conv_tr = klass_tr(chin, chout, ks, st) | |
| self.rewrite = klass(chin, 2 * chin, 1 + 2 * context, 1, context) | |
| self.dconv = DConv(chin, **dconv_kw) | |
| def forward(self, x, skip, length): | |
| if self.freq and x.dim() == 3: | |
| B, C, T = x.shape | |
| x = x.view(B, self.chin, -1, T) | |
| x = x + skip | |
| y = F.glu(self.rewrite(x), dim=1) | |
| if self.freq: | |
| B, C, Fr, T = y.shape | |
| y = y.permute(0, 2, 1, 3).reshape(-1, C, T) | |
| y = self.dconv(y).view(B, Fr, C, T).permute(0, 2, 1, 3) | |
| else: | |
| y = self.dconv(y) | |
| z = self.conv_tr(y) | |
| if self.freq: | |
| z = z[..., self.pad: -self.pad, :] | |
| else: | |
| z = z[..., self.pad: self.pad + length] | |
| if not self.last: | |
| z = F.gelu(z) | |
| return z, y | |
| def create_sin_embedding(length, dim, device, max_period=10000): | |
| pos = torch.arange(length, device=device).view(-1, 1, 1) | |
| half_dim = dim // 2 | |
| adim = torch.arange(half_dim, device=device).view(1, 1, -1) | |
| phase = pos / (max_period ** (adim / (half_dim - 1))) | |
| return torch.cat([torch.cos(phase), torch.sin(phase)], dim=-1) | |
| def create_2d_sin_embedding(d_model, height, width, device, max_period=10000): | |
| pe = torch.zeros(d_model, height, width) | |
| d_model = d_model // 2 | |
| div_term = torch.exp(torch.arange(0.0, d_model, 2) * -(math.log(max_period) / d_model)) | |
| pos_w = torch.arange(0.0, width).unsqueeze(1) | |
| pos_h = torch.arange(0.0, height).unsqueeze(1) | |
| pe[0:d_model:2, :, :] = torch.sin(pos_w * div_term).transpose(0, 1).unsqueeze(1).repeat(1, height, 1) | |
| pe[1:d_model:2, :, :] = torch.cos(pos_w * div_term).transpose(0, 1).unsqueeze(1).repeat(1, height, 1) | |
| pe[d_model::2, :, :] = torch.sin(pos_h * div_term).transpose(0, 1).unsqueeze(2).repeat(1, 1, width) | |
| pe[d_model + 1:: 2, :, :] = torch.cos(pos_h * div_term).transpose(0, 1).unsqueeze(2).repeat(1, 1, width) | |
| return pe[None, :].to(device) | |
| class MyGroupNorm(nn.GroupNorm): | |
| def forward(self, x): | |
| return super().forward(x.transpose(1, 2)).transpose(1, 2) | |
| class MyTransformerEncoderLayer(nn.TransformerEncoderLayer): | |
| def __init__(self, d_model, nhead, dim_feedforward, dropout, activation): | |
| super().__init__(d_model=d_model, nhead=nhead, dim_feedforward=dim_feedforward, | |
| dropout=dropout, activation=activation, batch_first=True, | |
| norm_first=True) | |
| self.norm_out = MyGroupNorm(1, d_model) | |
| self.gamma_1 = LayerScale(d_model, 1e-4, True) | |
| self.gamma_2 = LayerScale(d_model, 1e-4, True) | |
| def forward(self, x, **unused): | |
| x = x + self.gamma_1(self._sa_block(self.norm1(x), None, None)) | |
| x = x + self.gamma_2(self._ff_block(self.norm2(x))) | |
| return self.norm_out(x) | |
| class CrossTransformerEncoderLayer(nn.Module): | |
| def __init__(self, d_model, nhead, dim_feedforward, dropout, activation): | |
| super().__init__() | |
| self.cross_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=True) | |
| self.linear1 = nn.Linear(d_model, dim_feedforward) | |
| self.dropout = nn.Dropout(dropout) | |
| self.linear2 = nn.Linear(dim_feedforward, d_model) | |
| self.norm1 = nn.LayerNorm(d_model) | |
| self.norm2 = nn.LayerNorm(d_model) | |
| self.norm3 = nn.LayerNorm(d_model) | |
| self.norm_out = MyGroupNorm(1, d_model) | |
| self.gamma_1 = LayerScale(d_model, 1e-4, True) | |
| self.gamma_2 = LayerScale(d_model, 1e-4, True) | |
| self.dropout1 = nn.Dropout(dropout) | |
| self.dropout2 = nn.Dropout(dropout) | |
| self.activation = activation | |
| def forward(self, q, k): | |
| x = q + self.gamma_1(self.dropout1( | |
| self.cross_attn(self.norm1(q), self.norm2(k), self.norm2(k), need_weights=False)[0])) | |
| x = x + self.gamma_2(self._ff_block(self.norm3(x))) | |
| return self.norm_out(x) | |
| def _ff_block(self, x): | |
| return self.dropout2(self.linear2(self.dropout(self.activation(self.linear1(x))))) | |
| class CrossTransformerEncoder(nn.Module): | |
| def __init__(self, dim, hidden_scale=4.0, num_heads=8, num_layers=5, dropout=0.0, | |
| max_period=10000.0, gelu=True): | |
| super().__init__() | |
| hidden_dim = int(dim * hidden_scale) | |
| self.num_layers = num_layers | |
| self.max_period = max_period | |
| self.norm_in = nn.LayerNorm(dim) | |
| self.norm_in_t = nn.LayerNorm(dim) | |
| activation = F.gelu if gelu else F.relu | |
| self.layers = nn.ModuleList() | |
| self.layers_t = nn.ModuleList() | |
| common = dict(d_model=dim, nhead=num_heads, dim_feedforward=hidden_dim, | |
| dropout=dropout, activation=activation) | |
| for idx in range(num_layers): | |
| klass = MyTransformerEncoderLayer if idx % 2 == 0 else CrossTransformerEncoderLayer | |
| self.layers.append(klass(**common)) | |
| self.layers_t.append(klass(**common)) | |
| def forward(self, x, xt): | |
| B, C, Fr, T1 = x.shape | |
| pos_2d = rearrange(create_2d_sin_embedding(C, Fr, T1, x.device, self.max_period), | |
| "b c fr t1 -> b (t1 fr) c") | |
| x = (self.norm_in(rearrange(x, "b c fr t1 -> b (t1 fr) c")) + pos_2d).to(x.dtype) | |
| B, C, T2 = xt.shape | |
| pos = rearrange(create_sin_embedding(T2, C, x.device, self.max_period), "t2 b c -> b t2 c") | |
| xt = (self.norm_in_t(rearrange(xt, "b c t2 -> b t2 c")) + pos).to(xt.dtype) | |
| for idx in range(self.num_layers): | |
| if idx % 2 == 0: | |
| x = self.layers[idx](x) | |
| xt = self.layers_t[idx](xt) | |
| else: | |
| old_x = x | |
| x = self.layers[idx](x, xt) | |
| xt = self.layers_t[idx](xt, old_x) | |
| x = rearrange(x, "b (t1 fr) c -> b c fr t1", t1=T1) | |
| xt = rearrange(xt, "b t2 c -> b c t2") | |
| return x, xt | |
| class HTDemucs(nn.Module): | |
| def __init__(self, sources=['drums', 'bass', 'other', 'vocals'], audio_channels=2, channels=48, growth=2, nfft=4096, cac=True, | |
| depth=4, freq_emb=0.2, emb_scale=10, kernel_size=8, time_stride=2, stride=4, | |
| context=1, context_enc=0, dconv_depth=2, dconv_comp=8, dconv_init=1e-3, | |
| bottom_channels=512, t_layers=5, t_hidden_scale=4.0, t_heads=8, t_dropout=0.0, | |
| t_max_period=10000.0, t_gelu=True, samplerate=44100, segment=fractions.Fraction(39, 5), **unused): | |
| super().__init__() | |
| self.audio_channels = audio_channels | |
| self.sources = sources | |
| self.depth = depth | |
| self.bottom_channels = bottom_channels | |
| self.samplerate = samplerate | |
| self.segment = segment | |
| self.nfft = nfft | |
| self.hop_length = nfft // 4 | |
| self.encoder = nn.ModuleList() | |
| self.decoder = nn.ModuleList() | |
| self.tencoder = nn.ModuleList() | |
| self.tdecoder = nn.ModuleList() | |
| chin = audio_channels | |
| chin_z = chin * 2 | |
| chout = chout_z = channels | |
| freqs = nfft // 2 | |
| for index in range(depth): | |
| dconv_kw = {"depth": dconv_depth, "compress": dconv_comp, "init": dconv_init} | |
| kw = {"kernel_size": kernel_size, "stride": stride, "dconv_kw": dconv_kw} | |
| kwt = dict(kw, freq=False) | |
| self.encoder.append(HEncLayer(chin_z, chout_z, context=context_enc, **kw)) | |
| self.tencoder.append(HEncLayer(chin, chout, freq=False, context=context_enc, **kw)) | |
| if index == 0: | |
| chin = self.audio_channels * len(self.sources) | |
| chin_z = chin * 2 | |
| self.decoder.insert(0, HDecLayer(chout_z, chin_z, last=index == 0, context=context, **kw)) | |
| self.tdecoder.insert(0, HDecLayer(chout, chin, last=index == 0, context=context, **kwt)) | |
| chin, chin_z = chout, chout_z | |
| chout = int(growth * chout) | |
| chout_z = int(growth * chout_z) | |
| freqs //= stride | |
| if index == 0: | |
| self.freq_emb = ScaledEmbedding(freqs, chin_z, scale=emb_scale) | |
| self.freq_emb_scale = freq_emb | |
| transformer_channels = channels * growth ** (depth - 1) | |
| self.channel_upsampler = nn.Conv1d(transformer_channels, bottom_channels, 1) | |
| self.channel_downsampler = nn.Conv1d(bottom_channels, transformer_channels, 1) | |
| self.channel_upsampler_t = nn.Conv1d(transformer_channels, bottom_channels, 1) | |
| self.channel_downsampler_t = nn.Conv1d(bottom_channels, transformer_channels, 1) | |
| self.crosstransformer = CrossTransformerEncoder( | |
| dim=bottom_channels, hidden_scale=t_hidden_scale, num_heads=t_heads, | |
| num_layers=t_layers, dropout=t_dropout, max_period=t_max_period, gelu=t_gelu) | |
| def _spec(self, x): | |
| hl, nfft = self.hop_length, self.nfft | |
| le = int(math.ceil(x.shape[-1] / hl)) | |
| pad = hl // 2 * 3 | |
| x = F.pad(x, (pad, pad + le * hl - x.shape[-1]), mode="reflect") | |
| z = spectro(x.float(), nfft, hl)[..., :-1, :] | |
| return z[..., 2: 2 + le] | |
| def _ispec(self, z, length): | |
| hl = self.hop_length | |
| z = F.pad(z, (0, 0, 0, 1)) | |
| z = F.pad(z, (2, 2)) | |
| pad = hl // 2 * 3 | |
| le = hl * int(math.ceil(length / hl)) + 2 * pad | |
| return ispectro(z, hl, le)[..., pad: pad + length] | |
| def _magnitude(self, z): | |
| # cac: 复数实虚部当作额外通道 | |
| B, C, Fr, T = z.shape | |
| return torch.view_as_real(z).permute(0, 1, 4, 2, 3).reshape(B, C * 2, Fr, T) | |
| def _mask(self, m): | |
| # cac: m 即完整复数谱(通道形式), 还原成复数 | |
| B, S, C, Fr, T = m.shape | |
| out = m.view(B, S, -1, 2, Fr, T).permute(0, 1, 2, 4, 5, 3) | |
| return torch.view_as_complex(out.contiguous()) | |
| def forward(self, mix): | |
| length = mix.shape[-1] | |
| z = self._spec(mix) | |
| x = self._magnitude(z).to(device=mix.device, dtype=mix.dtype) | |
| B, C, Fq, T = x.shape | |
| mean = x.mean(dim=(1, 2, 3), keepdim=True) | |
| std = x.std(dim=(1, 2, 3), keepdim=True) | |
| x = (x - mean) / (1e-5 + std) | |
| xt = mix | |
| meant = xt.mean(dim=(1, 2), keepdim=True) | |
| stdt = xt.std(dim=(1, 2), keepdim=True) | |
| xt = (xt - meant) / (1e-5 + stdt) | |
| saved, saved_t, lengths, lengths_t = [], [], [], [] | |
| for idx, encode in enumerate(self.encoder): | |
| lengths.append(x.shape[-1]) | |
| lengths_t.append(xt.shape[-1]) | |
| xt = self.tencoder[idx](xt) | |
| saved_t.append(xt) | |
| x = encode(x) | |
| if idx == 0: | |
| frs = torch.arange(x.shape[-2], device=x.device) | |
| emb = self.freq_emb(frs).t()[None, :, :, None].expand_as(x) | |
| x = x + self.freq_emb_scale * emb | |
| saved.append(x) | |
| b, c, f, t = x.shape | |
| x = rearrange(x, "b c f t-> b c (f t)") | |
| x = self.channel_upsampler(x) | |
| x = rearrange(x, "b c (f t)-> b c f t", f=f) | |
| xt = self.channel_upsampler_t(xt) | |
| x, xt = self.crosstransformer(x, xt) | |
| x = rearrange(x, "b c f t-> b c (f t)") | |
| x = self.channel_downsampler(x) | |
| x = rearrange(x, "b c (f t)-> b c f t", f=f) | |
| xt = self.channel_downsampler_t(xt) | |
| for idx, decode in enumerate(self.decoder): | |
| x, pre = decode(x, saved.pop(-1), lengths.pop(-1)) | |
| xt, _ = self.tdecoder[idx](xt, saved_t.pop(-1), lengths_t.pop(-1)) | |
| S = len(self.sources) | |
| x = x.view(B, S, -1, Fq, T) * std[:, None] + mean[:, None] | |
| x = self._ispec(self._mask(x.float()), length).to(mix.dtype) | |
| xt = xt.view(B, S, -1, length) * stdt[:, None] + meant[:, None] | |
| return xt + x | |
| def extract_track(self, audio, sample_rate, track="vocals"): | |
| if sample_rate != 44100: | |
| audio = torchaudio.functional.resample(audio, sample_rate, 44100) | |
| ref = audio.mean(0) | |
| audio = (audio - ref.mean()) / (ref.std() + 1e-8) | |
| out = apply_model(self, audio[None], shifts=1, split=True, overlap=0.25, | |
| device=audio.device, progress=True)[0] | |
| out = out * (ref.std() + 1e-8) + ref.mean() | |
| out = out / max(1.01 * out.abs().max(), 1) | |
| out = out.clamp_(-1, 1).cpu() | |
| out = out[self.sources.index(track)] | |
| return out | |
| def center_trim(tensor, length): | |
| delta = tensor.size(-1) - length | |
| if delta: | |
| tensor = tensor[..., delta // 2: -(delta - delta // 2)] | |
| return tensor | |
| class TensorChunk: | |
| def __init__(self, tensor, offset=0, length=None): | |
| total = tensor.shape[-1] | |
| if isinstance(tensor, TensorChunk): | |
| self.tensor, self.offset = tensor.tensor, offset + tensor.offset | |
| else: | |
| self.tensor, self.offset = tensor, offset | |
| self.length = total - offset if length is None else min(total - offset, length) | |
| self.device = tensor.device | |
| 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 = self.tensor.shape[-1] | |
| start = self.offset - delta // 2 | |
| end = start + target_length | |
| correct_start, correct_end = max(0, start), min(total, end) | |
| out = F.pad(self.tensor[..., correct_start:correct_end], | |
| (correct_start - start, end - correct_end)) | |
| return out | |
| def apply_model(model, mix, shifts=1, split=True, overlap=0.25, device=None, segment=None, | |
| progress=False): | |
| device = th.device(device) if device is not None else mix.device | |
| batch, channels, length = mix.shape | |
| if shifts: | |
| max_shift = int(0.5 * model.samplerate) | |
| mix = mix if isinstance(mix, TensorChunk) else TensorChunk(mix) | |
| padded_mix = mix.padded(length + 2 * max_shift) | |
| out = 0.0 | |
| for _ in range(shifts): | |
| offset = int(th.randint(0, max_shift, (1,)).item()) | |
| shifted = TensorChunk(padded_mix, offset, length + max_shift - offset) | |
| shifted_out = apply_model(model, shifted, shifts=0, split=split, overlap=overlap, | |
| device=device, segment=segment, progress=progress) | |
| out += shifted_out[..., max_shift - offset:] | |
| return out / shifts | |
| if split: | |
| out = th.zeros(batch, len(model.sources), channels, length, device=mix.device) | |
| sum_weight = th.zeros(length, device=mix.device) | |
| if segment is None: | |
| segment = model.segment | |
| segment_length = int(model.samplerate * segment) | |
| stride = int((1 - overlap) * segment_length) | |
| offsets = range(0, length, stride) | |
| scale = float(format(stride / model.samplerate, ".2f")) | |
| weight = th.cat([th.arange(1, segment_length // 2 + 1, device=device), | |
| th.arange(segment_length - segment_length // 2, 0, -1, device=device)]) | |
| weight = weight / weight.max() | |
| iterator = offsets | |
| for offset in iterator: | |
| chunk = TensorChunk(mix, offset, segment_length) | |
| chunk_out = apply_model(model, chunk, shifts=0, split=False, device=device, segment=segment) | |
| cl = chunk_out.shape[-1] | |
| out[..., offset:offset + segment_length] += (weight[:cl] * chunk_out).to(mix.device) | |
| sum_weight[offset:offset + segment_length] += weight[:cl].to(mix.device) | |
| return out / sum_weight | |
| valid_length = int(segment * model.samplerate) | |
| mix = mix if isinstance(mix, TensorChunk) else TensorChunk(mix) | |
| padded_mix = mix.padded(valid_length).to(device) | |
| out = model(padded_mix) | |
| return center_trim(out, length) | |