deepafx_st/system.py

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
    @staticmethod
    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