Spaces:
Running
Running
| import torch | |
| from deepafx_st.processors.proxy.proxy_system import ProxySystem | |
| from deepafx_st.utils import DSPMode | |
| class ProxyChannel(torch.nn.Module): | |
| def __init__( | |
| self, | |
| proxy_system_ckpts: list, | |
| freeze_proxies: bool = True, | |
| dsp_mode: DSPMode = DSPMode.NONE, | |
| num_tcns: int = 2, | |
| tcn_nblocks: int = 4, | |
| tcn_dilation_growth: int = 8, | |
| tcn_channel_width: int = 64, | |
| tcn_kernel_size: int = 13, | |
| sample_rate: int = 24000, | |
| ): | |
| super().__init__() | |
| self.freeze_proxies = freeze_proxies | |
| self.dsp_mode = dsp_mode | |
| self.num_tcns = num_tcns | |
| # load the proxies | |
| self.proxies = torch.nn.ModuleList() | |
| self.num_control_params = 0 | |
| self.ports = [] | |
| for proxy_system_ckpt in proxy_system_ckpts: | |
| proxy = ProxySystem.load_from_checkpoint(proxy_system_ckpt) | |
| # freeze model parameters | |
| if freeze_proxies: | |
| for param in proxy.parameters(): | |
| param.requires_grad = False | |
| self.proxies.append(proxy) | |
| if proxy.hparams.processor == "channel": | |
| self.ports = proxy.processor.ports | |
| else: | |
| self.ports.append(proxy.processor.ports) | |
| self.num_control_params += proxy.processor.num_control_params | |
| if len(proxy_system_ckpts) == 0: | |
| if self.num_tcns == 2: | |
| peq_proxy = ProxySystem( | |
| processor="peq", | |
| output_gain=False, | |
| nblocks=tcn_nblocks, | |
| dilation_growth=tcn_dilation_growth, | |
| kernel_size=tcn_kernel_size, | |
| channel_width=tcn_channel_width, | |
| sample_rate=sample_rate, | |
| ) | |
| self.proxies.append(peq_proxy) | |
| self.ports.append(peq_proxy.processor.ports) | |
| self.num_control_params += peq_proxy.processor.num_control_params | |
| comp_proxy = ProxySystem( | |
| processor="comp", | |
| output_gain=True, | |
| nblocks=tcn_nblocks, | |
| dilation_growth=tcn_dilation_growth, | |
| kernel_size=tcn_kernel_size, | |
| channel_width=tcn_channel_width, | |
| sample_rate=sample_rate, | |
| ) | |
| self.proxies.append(comp_proxy) | |
| self.ports.append(comp_proxy.processor.ports) | |
| self.num_control_params += comp_proxy.processor.num_control_params | |
| elif self.num_tcns == 1: | |
| channel_proxy = ProxySystem( | |
| processor="channel", | |
| output_gain=True, | |
| nblocks=tcn_nblocks, | |
| dilation_growth=tcn_dilation_growth, | |
| kernel_size=tcn_kernel_size, | |
| channel_width=tcn_channel_width, | |
| sample_rate=sample_rate, | |
| ) | |
| self.proxies.append(channel_proxy) | |
| for port_list in channel_proxy.processor.ports: | |
| self.ports.append(port_list) | |
| self.num_control_params += channel_proxy.processor.num_control_params | |
| else: | |
| raise ValueError(f"num_tcns must be <= 2. Asked for {self.num_tcns}.") | |
| def forward( | |
| self, | |
| x: torch.Tensor, | |
| p: torch.Tensor, | |
| dsp_mode: DSPMode = DSPMode.NONE, | |
| sample_rate: int = 24000, | |
| **kwargs, | |
| ): | |
| # loop over the proxies and pass parameters | |
| stop_idx = 0 | |
| for proxy in self.proxies: | |
| start_idx = stop_idx | |
| stop_idx += proxy.processor.num_control_params | |
| p_subset = p[:, start_idx:stop_idx] | |
| if dsp_mode.name == DSPMode.NONE.name: | |
| x = proxy( | |
| x, | |
| p_subset, | |
| use_dsp=False, | |
| ) | |
| elif dsp_mode.name == DSPMode.INFER.name: | |
| x = proxy( | |
| x, | |
| p_subset, | |
| use_dsp=True, | |
| sample_rate=sample_rate, | |
| ) | |
| elif dsp_mode.name == DSPMode.TRAIN_INFER.name: | |
| # Mimic gumbel softmax implementation to replace grads similar to | |
| # https://gist.github.com/yzh119/fd2146d2aeb329d067568a493b20172f | |
| x_hard = proxy( | |
| x, | |
| p_subset, | |
| use_dsp=True, | |
| sample_rate=sample_rate, | |
| ) | |
| x = proxy( | |
| x, | |
| p_subset, | |
| use_dsp=False, | |
| sample_rate=sample_rate, | |
| ) | |
| x = (x_hard - x).detach() + x | |
| else: | |
| assert 0, "invalid dsp model for proxy" | |
| return x | |