Spaces:
Running
Running
| import torch | |
| import auraloss | |
| import torchaudio | |
| from itertools import chain | |
| import pytorch_lightning as pl | |
| from argparse import ArgumentParser | |
| from typing import Tuple, List, Dict | |
| import deepafx_st.utils as utils | |
| from deepafx_st.utils import DSPMode | |
| from deepafx_st.data.dataset import AudioDataset | |
| from deepafx_st.models.encoder import SpectralEncoder | |
| from deepafx_st.models.controller import StyleTransferController | |
| from deepafx_st.processors.spsa.channel import SPSAChannel | |
| from deepafx_st.processors.spsa.eps_scheduler import EpsilonScheduler | |
| from deepafx_st.processors.proxy.channel import ProxyChannel | |
| from deepafx_st.processors.autodiff.channel import AutodiffChannel | |
| class System(pl.LightningModule): | |
| def __init__( | |
| self, | |
| ext="wav", | |
| dsp_sample_rate=24000, | |
| **kwargs, | |
| ): | |
| super().__init__() | |
| self.save_hyperparameters() | |
| self.eps_scheduler = EpsilonScheduler( | |
| self.hparams.spsa_epsilon, | |
| self.hparams.spsa_patience, | |
| self.hparams.spsa_factor, | |
| self.hparams.spsa_verbose, | |
| ) | |
| self.hparams.dsp_mode = DSPMode.NONE | |
| # first construct the processor, since this will dictate encoder | |
| if self.hparams.processor_model == "spsa": | |
| self.processor = SPSAChannel( | |
| self.hparams.dsp_sample_rate, | |
| self.hparams.spsa_parallel, | |
| self.hparams.batch_size, | |
| ) | |
| elif self.hparams.processor_model == "autodiff": | |
| self.processor = AutodiffChannel(self.hparams.dsp_sample_rate) | |
| elif self.hparams.processor_model == "proxy0": | |
| # print('self.hparams.proxy_ckpts,',self.hparams.proxy_ckpts) | |
| self.hparams.dsp_mode = DSPMode.NONE | |
| self.processor = ProxyChannel( | |
| self.hparams.proxy_ckpts, | |
| self.hparams.freeze_proxies, | |
| self.hparams.dsp_mode, | |
| sample_rate=self.hparams.dsp_sample_rate, | |
| ) | |
| elif self.hparams.processor_model == "proxy1": | |
| # print('self.hparams.proxy_ckpts,',self.hparams.proxy_ckpts) | |
| self.hparams.dsp_mode = DSPMode.INFER | |
| self.processor = ProxyChannel( | |
| self.hparams.proxy_ckpts, | |
| self.hparams.freeze_proxies, | |
| self.hparams.dsp_mode, | |
| sample_rate=self.hparams.dsp_sample_rate, | |
| ) | |
| elif self.hparams.processor_model == "proxy2": | |
| # print('self.hparams.proxy_ckpts,',self.hparams.proxy_ckpts) | |
| self.hparams.dsp_mode = DSPMode.TRAIN_INFER | |
| self.processor = ProxyChannel( | |
| self.hparams.proxy_ckpts, | |
| self.hparams.freeze_proxies, | |
| self.hparams.dsp_mode, | |
| sample_rate=self.hparams.dsp_sample_rate, | |
| ) | |
| elif self.hparams.processor_model == "tcn1": | |
| # self.processor = ConditionalTCN(self.hparams.sample_rate) | |
| self.hparams.dsp_mode = DSPMode.NONE | |
| self.processor = ProxyChannel( | |
| [], | |
| freeze_proxies=False, | |
| dsp_mode=self.hparams.dsp_mode, | |
| tcn_nblocks=self.hparams.tcn_nblocks, | |
| tcn_dilation_growth=self.hparams.tcn_dilation_growth, | |
| tcn_channel_width=self.hparams.tcn_channel_width, | |
| tcn_kernel_size=self.hparams.tcn_kernel_size, | |
| num_tcns=1, | |
| sample_rate=self.hparams.sample_rate, | |
| ) | |
| elif self.hparams.processor_model == "tcn2": | |
| self.hparams.dsp_mode = DSPMode.NONE | |
| self.processor = ProxyChannel( | |
| [], | |
| freeze_proxies=False, | |
| dsp_mode=self.hparams.dsp_mode, | |
| tcn_nblocks=self.hparams.tcn_nblocks, | |
| tcn_dilation_growth=self.hparams.tcn_dilation_growth, | |
| tcn_channel_width=self.hparams.tcn_channel_width, | |
| tcn_kernel_size=self.hparams.tcn_kernel_size, | |
| num_tcns=2, | |
| sample_rate=self.hparams.sample_rate, | |
| ) | |
| else: | |
| raise ValueError(f"Invalid processor_model: {self.hparams.processor_model}") | |
| if self.hparams.encoder_ckpt is not None: | |
| # load encoder weights from a pre-trained system | |
| system = System.load_from_checkpoint(self.hparams.encoder_ckpt) | |
| self.encoder = system.encoder | |
| self.hparams.encoder_embed_dim = system.encoder.embed_dim | |
| else: | |
| self.encoder = SpectralEncoder( | |
| self.processor.num_control_params, | |
| self.hparams.sample_rate, | |
| encoder_model=self.hparams.encoder_model, | |
| embed_dim=self.hparams.encoder_embed_dim, | |
| width_mult=self.hparams.encoder_width_mult, | |
| ) | |
| if self.hparams.encoder_freeze: | |
| for param in self.encoder.parameters(): | |
| param.requires_grad = False | |
| self.controller = StyleTransferController( | |
| self.processor.num_control_params, | |
| self.hparams.encoder_embed_dim, | |
| ) | |
| if len(self.hparams.recon_losses) != len(self.hparams.recon_loss_weights): | |
| raise ValueError("Must supply same number of weights as losses.") | |
| self.recon_losses = torch.nn.ModuleDict() | |
| for recon_loss in self.hparams.recon_losses: | |
| if recon_loss == "mrstft": | |
| self.recon_losses[recon_loss] = auraloss.freq.MultiResolutionSTFTLoss( | |
| fft_sizes=[32, 128, 512, 2048, 8192, 32768], | |
| hop_sizes=[16, 64, 256, 1024, 4096, 16384], | |
| win_lengths=[32, 128, 512, 2048, 8192, 32768], | |
| w_sc=0.0, | |
| w_phs=0.0, | |
| w_lin_mag=1.0, | |
| w_log_mag=1.0, | |
| ) | |
| elif recon_loss == "mrstft-md": | |
| self.recon_losses[recon_loss] = auraloss.freq.MultiResolutionSTFTLoss( | |
| fft_sizes=[128, 512, 2048, 8192], | |
| hop_sizes=[32, 128, 512, 2048], # 1 / 4 | |
| win_lengths=[128, 512, 2048, 8192], | |
| w_sc=0.0, | |
| w_phs=0.0, | |
| w_lin_mag=1.0, | |
| w_log_mag=1.0, | |
| ) | |
| elif recon_loss == "mrstft-sm": | |
| self.recon_losses[recon_loss] = auraloss.freq.MultiResolutionSTFTLoss( | |
| fft_sizes=[512, 2048, 8192], | |
| hop_sizes=[256, 1024, 4096], # 1 / 4 | |
| win_lengths=[512, 2048, 8192], | |
| w_sc=0.0, | |
| w_phs=0.0, | |
| w_lin_mag=1.0, | |
| w_log_mag=1.0, | |
| ) | |
| elif recon_loss == "melfft": | |
| self.recon_losses[recon_loss] = auraloss.freq.MelSTFTLoss( | |
| self.hparams.sample_rate, | |
| fft_size=self.hparams.train_length, | |
| hop_size=self.hparams.train_length // 2, | |
| win_length=self.hparams.train_length, | |
| n_mels=128, | |
| w_sc=0.0, | |
| device="cuda" if self.hparams.gpus > 0 else "cpu", | |
| ) | |
| elif recon_loss == "melstft": | |
| self.recon_losses[recon_loss] = auraloss.freq.MelSTFTLoss( | |
| self.hparams.sample_rate, | |
| device="cuda" if self.hparams.gpus > 0 else "cpu", | |
| ) | |
| elif recon_loss == "l1": | |
| self.recon_losses[recon_loss] = torch.nn.L1Loss() | |
| elif recon_loss == "sisdr": | |
| self.recon_losses[recon_loss] = auraloss.time.SISDRLoss() | |
| else: | |
| raise ValueError( | |
| f"Invalid reconstruction loss: {self.hparams.recon_losses}" | |
| ) | |
| def forward( | |
| self, | |
| x: torch.Tensor, | |
| y: torch.Tensor = None, | |
| e_y: torch.Tensor = None, | |
| z: torch.Tensor = None, | |
| dsp_mode: DSPMode = DSPMode.NONE, | |
| analysis_length: int = 0, | |
| sample_rate: int = 24000, | |
| ): | |
| """Forward pass through the system subnetworks. | |
| Args: | |
| x (tensor): Input audio tensor with shape (batch x 1 x samples) | |
| y (tensor): Target audio tensor with shape (batch x 1 x samples) | |
| e_y (tensor): Target embedding with shape (batch x edim) | |
| z (tensor): Bottleneck latent. | |
| dsp_mode (DSPMode): Mode of operation for the DSP blocks. | |
| analysis_length (optional, int): Only analyze the first N samples. | |
| sample_rate (optional, int): Desired sampling rate for the DSP blocks. | |
| You must supply target audio `y`, `z`, or an embedding for the target `e_y`. | |
| Returns: | |
| y_hat (tensor): Output audio. | |
| p (tensor): | |
| e (tensor): | |
| """ | |
| bs, chs, samp = x.size() | |
| if sample_rate != self.hparams.sample_rate: | |
| x_enc = torchaudio.transforms.Resample( | |
| sample_rate, self.hparams.sample_rate | |
| ).to(x.device)(x) | |
| if y is not None: | |
| y_enc = torchaudio.transforms.Resample( | |
| sample_rate, self.hparams.sample_rate | |
| ).to(x.device)(y) | |
| else: | |
| x_enc = x | |
| y_enc = y | |
| if analysis_length > 0: | |
| x_enc = x_enc[..., :analysis_length] | |
| if y is not None: | |
| y_enc = y_enc[..., :analysis_length] | |
| e_x = self.encoder(x_enc) # generate latent embedding for input | |
| if y is not None: | |
| e_y = self.encoder(y_enc) # generate latent embedding for target | |
| elif e_y is None: | |
| raise RuntimeError("Must supply y, z, or e_y. None supplied.") | |
| # learnable comparision | |
| p = self.controller(e_x, e_y, z=z) | |
| # process audio conditioned on parameters | |
| # if there are multiple channels process them using same parameters | |
| y_hat = torch.zeros(x.shape).type_as(x) | |
| for ch_idx in range(chs): | |
| y_hat_ch = self.processor( | |
| x[:, ch_idx : ch_idx + 1, :], | |
| p, | |
| epsilon=self.eps_scheduler.epsilon, | |
| dsp_mode=dsp_mode, | |
| sample_rate=sample_rate, | |
| ) | |
| y_hat[:, ch_idx : ch_idx + 1, :] = y_hat_ch | |
| return y_hat, p, e_x | |
| def common_paired_step( | |
| self, | |
| batch: Tuple, | |
| batch_idx: int, | |
| optimizer_idx: int = 0, | |
| train: bool = False, | |
| ): | |
| """Model step used for validation and training. | |
| Args: | |
| batch (Tuple[Tensor, Tensor]): Batch items containing input audio (x) and target audio (y). | |
| batch_idx (int): Index of the batch within the current epoch. | |
| optimizer_idx (int): Index of the optimizer, this step is called once for each optimizer. | |
| The firs optimizer corresponds to the generator and the second optimizer, | |
| corresponds to the adversarial loss (when in use). | |
| train (bool): Whether step is called during training (True) or validation (False). | |
| """ | |
| x, y = batch | |
| loss = 0 | |
| dsp_mode = self.hparams.dsp_mode | |
| if train and dsp_mode.INFER.name == DSPMode.INFER.name: | |
| dsp_mode = DSPMode.NONE | |
| # proces input audio through model | |
| if self.hparams.style_transfer: | |
| length = x.shape[-1] | |
| x_A = x[..., : length // 2] | |
| x_B = x[..., length // 2 :] | |
| y_A = y[..., : length // 2] | |
| y_B = y[..., length // 2 :] | |
| if torch.rand(1).sum() > 0.5: | |
| y_ref = y_B | |
| y = y_A | |
| x = x_A | |
| else: | |
| y_ref = y_A | |
| y = y_B | |
| x = x_B | |
| y_hat, p, e = self(x, y=y_ref, dsp_mode=dsp_mode) | |
| else: | |
| y_ref = None | |
| y_hat, p, e = self(x, dsp_mode=dsp_mode) | |
| # compute reconstruction loss terms | |
| for loss_idx, (loss_name, recon_loss_fn) in enumerate( | |
| self.recon_losses.items() | |
| ): | |
| temp_loss = recon_loss_fn(y_hat, y) # reconstruction loss | |
| loss += float(self.hparams.recon_loss_weights[loss_idx]) * temp_loss | |
| self.log( | |
| ("train" if train else "val") + f"_loss/{loss_name}", | |
| temp_loss, | |
| on_step=True, | |
| on_epoch=True, | |
| prog_bar=False, | |
| logger=True, | |
| sync_dist=True, | |
| ) | |
| # log the overall aggregate loss | |
| self.log( | |
| ("train" if train else "val") + "_loss/loss", | |
| loss, | |
| on_step=True, | |
| on_epoch=True, | |
| prog_bar=False, | |
| logger=True, | |
| sync_dist=True, | |
| ) | |
| # store audio data | |
| data_dict = { | |
| "x": x.cpu(), | |
| "y": y.cpu(), | |
| "p": p.cpu(), | |
| "e": e.cpu(), | |
| "y_hat": y_hat.cpu(), | |
| } | |
| if y_ref is not None: | |
| data_dict["y_ref"] = y_ref.cpu() | |
| return loss, data_dict | |
| def training_step(self, batch, batch_idx, optimizer_idx=0): | |
| loss, _ = self.common_paired_step( | |
| batch, | |
| batch_idx, | |
| optimizer_idx, | |
| train=True, | |
| ) | |
| return loss | |
| def training_epoch_end(self, training_step_outputs): | |
| if self.hparams.spsa_schedule and self.hparams.processor_model == "spsa": | |
| self.eps_scheduler.step( | |
| self.trainer.callback_metrics[self.hparams.train_monitor], | |
| ) | |
| def validation_step(self, batch, batch_idx): | |
| loss, data_dict = self.common_paired_step(batch, batch_idx) | |
| return data_dict | |
| def optimizer_step( | |
| self, | |
| epoch, | |
| batch_idx, | |
| optimizer, | |
| optimizer_idx, | |
| optimizer_closure, | |
| on_tpu=False, | |
| using_native_amp=False, | |
| using_lbfgs=False, | |
| ): | |
| if optimizer_idx == 0: | |
| optimizer.step(closure=optimizer_closure) | |
| def configure_optimizers(self): | |
| # we need additional optimizer for the discriminator | |
| optimizers = [] | |
| g_optimizer = torch.optim.Adam( | |
| chain( | |
| self.encoder.parameters(), | |
| self.processor.parameters(), | |
| self.controller.parameters(), | |
| ), | |
| lr=self.hparams.lr, | |
| betas=(0.9, 0.999), | |
| ) | |
| optimizers.append(g_optimizer) | |
| g_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( | |
| g_optimizer, | |
| patience=self.hparams.lr_patience, | |
| verbose=True, | |
| ) | |
| ms1 = int(self.hparams.max_epochs * 0.8) | |
| ms2 = int(self.hparams.max_epochs * 0.95) | |
| print( | |
| "Learning rate schedule:", | |
| f"0 {self.hparams.lr:0.2e} -> ", | |
| f"{ms1} {self.hparams.lr*0.1:0.2e} -> ", | |
| f"{ms2} {self.hparams.lr*0.01:0.2e}", | |
| ) | |
| g_scheduler = torch.optim.lr_scheduler.MultiStepLR( | |
| g_optimizer, | |
| milestones=[ms1, ms2], | |
| gamma=0.1, | |
| ) | |
| lr_schedulers = { | |
| "scheduler": g_scheduler, | |
| } | |
| return optimizers, lr_schedulers | |
| def train_dataloader(self): | |
| train_dataset = AudioDataset( | |
| self.hparams.audio_dir, | |
| subset="train", | |
| train_frac=self.hparams.train_frac, | |
| half=self.hparams.half, | |
| length=self.hparams.train_length, | |
| input_dirs=self.hparams.input_dirs, | |
| random_scale_input=self.hparams.random_scale_input, | |
| random_scale_target=self.hparams.random_scale_target, | |
| buffer_size_gb=self.hparams.buffer_size_gb, | |
| buffer_reload_rate=self.hparams.buffer_reload_rate, | |
| num_examples_per_epoch=self.hparams.train_examples_per_epoch, | |
| augmentations={ | |
| "pitch": {"sr": self.hparams.sample_rate}, | |
| "tempo": {"sr": self.hparams.sample_rate}, | |
| }, | |
| freq_corrupt=self.hparams.freq_corrupt, | |
| drc_corrupt=self.hparams.drc_corrupt, | |
| ext=self.hparams.ext, | |
| ) | |
| g = torch.Generator() | |
| g.manual_seed(0) | |
| return torch.utils.data.DataLoader( | |
| train_dataset, | |
| num_workers=self.hparams.num_workers, | |
| batch_size=self.hparams.batch_size, | |
| worker_init_fn=utils.seed_worker, | |
| generator=g, | |
| pin_memory=True, | |
| persistent_workers=True, | |
| timeout=60, | |
| ) | |
| def val_dataloader(self): | |
| val_dataset = AudioDataset( | |
| self.hparams.audio_dir, | |
| subset="val", | |
| half=self.hparams.half, | |
| train_frac=self.hparams.train_frac, | |
| length=self.hparams.val_length, | |
| input_dirs=self.hparams.input_dirs, | |
| buffer_size_gb=self.hparams.buffer_size_gb, | |
| buffer_reload_rate=self.hparams.buffer_reload_rate, | |
| random_scale_input=self.hparams.random_scale_input, | |
| random_scale_target=self.hparams.random_scale_target, | |
| num_examples_per_epoch=self.hparams.val_examples_per_epoch, | |
| augmentations={}, | |
| freq_corrupt=self.hparams.freq_corrupt, | |
| drc_corrupt=self.hparams.drc_corrupt, | |
| ext=self.hparams.ext, | |
| ) | |
| self.val_dataset = val_dataset | |
| g = torch.Generator() | |
| g.manual_seed(0) | |
| return torch.utils.data.DataLoader( | |
| val_dataset, | |
| num_workers=1, | |
| batch_size=self.hparams.batch_size, | |
| worker_init_fn=utils.seed_worker, | |
| generator=g, | |
| pin_memory=True, | |
| persistent_workers=True, | |
| timeout=60, | |
| ) | |
| def shutdown(self): | |
| del self.processor | |
| # add any model hyperparameters here | |
| def add_model_specific_args(parent_parser): | |
| parser = ArgumentParser(parents=[parent_parser], add_help=False) | |
| # --- Training --- | |
| parser.add_argument("--batch_size", type=int, default=32) | |
| parser.add_argument("--lr", type=float, default=3e-4) | |
| parser.add_argument("--lr_patience", type=int, default=20) | |
| parser.add_argument("--recon_losses", nargs="+", default=["l1"]) | |
| parser.add_argument("--recon_loss_weights", nargs="+", default=[1.0]) | |
| # --- Controller --- | |
| parser.add_argument( | |
| "--processor_model", | |
| type=str, | |
| help="autodiff, spsa, tcn1, tcn2, proxy0, proxy1, proxy2", | |
| ) | |
| parser.add_argument("--controller_hidden_dim", type=int, default=256) | |
| parser.add_argument("--style_transfer", action="store_true") | |
| # --- Encoder --- | |
| parser.add_argument("--encoder_model", type=str, default="mobilenet_v2") | |
| parser.add_argument("--encoder_embed_dim", type=int, default=128) | |
| parser.add_argument("--encoder_width_mult", type=int, default=2) | |
| parser.add_argument("--encoder_ckpt", type=str, default=None) | |
| parser.add_argument("--encoder_freeze", action="store_true", default=False) | |
| # --- TCN --- | |
| parser.add_argument("--tcn_causal", action="store_true") | |
| parser.add_argument("--tcn_nblocks", type=int, default=4) | |
| parser.add_argument("--tcn_dilation_growth", type=int, default=8) | |
| parser.add_argument("--tcn_channel_width", type=int, default=32) | |
| parser.add_argument("--tcn_kernel_size", type=int, default=13) | |
| # --- SPSA --- | |
| parser.add_argument("--plugin_config_file", type=str, default=None) | |
| parser.add_argument("--spsa_epsilon", type=float, default=0.001) | |
| parser.add_argument("--spsa_schedule", action="store_true") | |
| parser.add_argument("--spsa_patience", type=int, default=10) | |
| parser.add_argument("--spsa_verbose", action="store_true") | |
| parser.add_argument("--spsa_factor", type=float, default=0.5) | |
| parser.add_argument("--spsa_parallel", action="store_true") | |
| # --- Proxy ---- | |
| parser.add_argument("--proxy_ckpts", nargs="+") | |
| parser.add_argument("--freeze_proxies", action="store_true", default=False) | |
| parser.add_argument("--use_dsp", action="store_true", default=False) | |
| parser.add_argument("--dsp_mode", choices=DSPMode, type=DSPMode) | |
| # --- Dataset --- | |
| parser.add_argument("--audio_dir", type=str) | |
| parser.add_argument("--ext", type=str, default="wav") | |
| parser.add_argument("--input_dirs", nargs="+") | |
| parser.add_argument("--buffer_reload_rate", type=int, default=1000) | |
| parser.add_argument("--buffer_size_gb", type=float, default=1.0) | |
| parser.add_argument("--sample_rate", type=int, default=24000) | |
| parser.add_argument("--dsp_sample_rate", type=int, default=24000) | |
| parser.add_argument("--shuffle", type=bool, default=True) | |
| parser.add_argument("--random_scale_input", action="store_true") | |
| parser.add_argument("--random_scale_target", action="store_true") | |
| parser.add_argument("--freq_corrupt", action="store_true") | |
| parser.add_argument("--drc_corrupt", action="store_true") | |
| parser.add_argument("--train_length", type=int, default=65536) | |
| parser.add_argument("--train_frac", type=float, default=0.8) | |
| parser.add_argument("--half", action="store_true") | |
| parser.add_argument("--train_examples_per_epoch", type=int, default=10000) | |
| parser.add_argument("--val_length", type=int, default=131072) | |
| parser.add_argument("--val_examples_per_epoch", type=int, default=1000) | |
| parser.add_argument("--num_workers", type=int, default=16) | |
| return parser | |