File size: 2,428 Bytes
bc3e180 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 |
import torch
import torch.nn as nn
from src.model.nn.blocks import FMBlock, AMBlock, ModBlock
from src.utils.ddsp import scale_function, remove_above_nyquist, upsample
from src.utils.ddsp import remove_above_nyquist_mode
from src.utils.ddsp import harmonic_synth, amp_to_impulse_response, fft_convolve
from src.utils.ddsp import modal_synth
from src.utils.ddsp import resample
import math
class DMSP(nn.Module):
def __init__(self,
embed_dim, hidden_size, n_features,
n_modes, n_bands, sampling_rate, block_size,
):
super().__init__()
self.n_modes = n_modes
self.freq_modulator = FMBlock(n_modes, embed_dim, n_features)
self.coef_modulator = AMBlock(n_modes, embed_dim, n_features)
self.proj_noise = nn.Linear(n_features*embed_dim, n_bands)
self.register_buffer("sampling_rate", torch.tensor(sampling_rate))
self.register_buffer("block_size", torch.tensor(block_size))
def forward(self, hidden, mode_freq, mode_coef, times, alpha, omega, lengths):
''' hidden : (Bs, 1, hidden_size)
mode_freq : (Bs, Nt, n_modes)
mode_coef : (Bs, 1, n_modes)
times : (Bs, Nt, 1)
'''
freq_m = self.freq_modulator(mode_freq, hidden, alpha, omega)
coef_m = self.coef_modulator(mode_coef, hidden, times)
#==============================
# harmonic part
#==============================
freqs = freq_m / (2*math.pi) * self.sampling_rate
coef_m = remove_above_nyquist_mode(coef_m, freqs, self.sampling_rate) # (Bs, Nt, n_modes)
freq_s = upsample(freq_m, self.block_size).narrow(1,0,lengths)
coef_s = upsample(coef_m, self.block_size).narrow(1,0,lengths)
harmonic = modal_synth(freq_s, coef_s, self.sampling_rate)
#==============================
# noise part
#==============================
param = scale_function(self.proj_noise(hidden) - 5)
impulse = amp_to_impulse_response(param, self.block_size)
noise = torch.rand(
impulse.shape[0],
impulse.shape[1],
self.block_size,
).to(impulse) * 2 - 1
noise = fft_convolve(noise, impulse).contiguous()
noise = noise.reshape(noise.shape[0], -1, 1).narrow(1,0,lengths)
signal = harmonic + noise
return signal.squeeze(-1), freq_m, coef_m
|