Fix folder structure

.gitignore

@@ -6,7 +6,6 @@ data/
 .DS_Store
 __pycache__/
 lightning_logs/
-RemFX/
 outputs/
 logs/
 .vscode/

remfx/models.py

@@ -0,0 +1,192 @@
+import torch
+from torch import Tensor, nn
+import pytorch_lightning as pl
+from einops import rearrange
+import wandb
+from audio_diffusion_pytorch import DiffusionModel
+import auraloss
+
+from umx.openunmix.model import OpenUnmix, Separator
+
+
+class RemFXModel(pl.LightningModule):
+    def __init__(
+        self,
+        lr: float,
+        lr_beta1: float,
+        lr_beta2: float,
+        lr_eps: float,
+        lr_weight_decay: float,
+        sample_rate: float,
+        network: nn.Module,
+    ):
+        super().__init__()
+        self.lr = lr
+        self.lr_beta1 = lr_beta1
+        self.lr_beta2 = lr_beta2
+        self.lr_eps = lr_eps
+        self.lr_weight_decay = lr_weight_decay
+        self.sample_rate = sample_rate
+        self.model = network
+
+    @property
+    def device(self):
+        return next(self.model.parameters()).device
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.AdamW(
+            list(self.model.parameters()),
+            lr=self.lr,
+            betas=(self.lr_beta1, self.lr_beta2),
+            eps=self.lr_eps,
+            weight_decay=self.lr_weight_decay,
+        )
+        return optimizer
+
+    def training_step(self, batch, batch_idx):
+        loss = self.common_step(batch, batch_idx, mode="train")
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        loss = self.common_step(batch, batch_idx, mode="valid")
+
+    def common_step(self, batch, batch_idx, mode: str = "train"):
+        loss = self.model(batch)
+        self.log(f"{mode}_loss", loss)
+        return loss
+
+    def on_validation_epoch_start(self):
+        self.log_next = True
+
+    def on_validation_batch_start(self, batch, batch_idx, dataloader_idx):
+        if self.log_next:
+            x, target, label = batch
+            y = self.model.sample(x)
+            log_wandb_audio_batch(
+                logger=self.logger,
+                id="sample",
+                samples=x.cpu(),
+                sampling_rate=self.sample_rate,
+                caption=f"Epoch {self.current_epoch}",
+            )
+            log_wandb_audio_batch(
+                logger=self.logger,
+                id="prediction",
+                samples=y.cpu(),
+                sampling_rate=self.sample_rate,
+                caption=f"Epoch {self.current_epoch}",
+            )
+            log_wandb_audio_batch(
+                logger=self.logger,
+                id="target",
+                samples=target.cpu(),
+                sampling_rate=self.sample_rate,
+                caption=f"Epoch {self.current_epoch}",
+            )
+            self.log_next = False
+
+
+class OpenUnmixModel(torch.nn.Module):
+    def __init__(
+        self,
+        n_fft: int = 2048,
+        hop_length: int = 512,
+        n_channels: int = 1,
+        alpha: float = 0.3,
+        sample_rate: int = 22050,
+    ):
+        super().__init__()
+        self.n_channels = n_channels
+        self.n_fft = n_fft
+        self.hop_length = hop_length
+        self.alpha = alpha
+        window = torch.hann_window(n_fft)
+        self.register_buffer("window", window)
+
+        self.num_bins = self.n_fft // 2 + 1
+        self.sample_rate = sample_rate
+        self.model = OpenUnmix(
+            nb_channels=self.n_channels,
+            nb_bins=self.num_bins,
+        )
+        self.separator = Separator(
+            target_models={"other": self.model},
+            nb_channels=self.n_channels,
+            sample_rate=self.sample_rate,
+            n_fft=self.n_fft,
+            n_hop=self.hop_length,
+        )
+        self.loss_fn = auraloss.freq.MultiResolutionSTFTLoss(
+            n_bins=self.num_bins, sample_rate=self.sample_rate
+        )
+
+    def forward(self, batch):
+        x, target, label = batch
+        X = spectrogram(x, self.window, self.n_fft, self.hop_length, self.alpha)
+        Y = self.model(X)
+        sep_out = self.separator(x).squeeze(1)
+        loss = self.loss_fn(sep_out, target)
+
+        return loss
+
+    def sample(self, x: Tensor) -> Tensor:
+        return self.separator(x).squeeze(1)
+
+
+class DiffusionGenerationModel(nn.Module):
+    def __init__(self, n_channels: int = 1):
+        super().__init__()
+        self.model = DiffusionModel(in_channels=n_channels)
+
+    def forward(self, batch):
+        x, target, label = batch
+        return self.model(x)
+
+    def sample(self, x: Tensor, num_steps: int = 10) -> Tensor:
+        noise = torch.randn(x.shape).to(x)
+        return self.model.sample(noise, num_steps=num_steps)
+
+
+def log_wandb_audio_batch(
+    logger: pl.loggers.WandbLogger,
+    id: str,
+    samples: Tensor,
+    sampling_rate: int,
+    caption: str = "",
+):
+    num_items = samples.shape[0]
+    samples = rearrange(samples, "b c t -> b t c")
+    for idx in range(num_items):
+        logger.experiment.log(
+            {
+                f"{id}_{idx}": wandb.Audio(
+                    samples[idx].cpu().numpy(),
+                    caption=caption,
+                    sample_rate=sampling_rate,
+                )
+            }
+        )
+
+
+def spectrogram(
+    x: torch.Tensor,
+    window: torch.Tensor,
+    n_fft: int,
+    hop_length: int,
+    alpha: float,
+) -> torch.Tensor:
+    bs, chs, samp = x.size()
+    x = x.view(bs * chs, -1)  # move channels onto batch dim
+
+    X = torch.stft(
+        x,
+        n_fft=n_fft,
+        hop_length=hop_length,
+        window=window,
+        return_complex=True,
+    )
+
+    # move channels back
+    X = X.view(bs, chs, X.shape[-2], X.shape[-1])
+
+    return torch.pow(X.abs() + 1e-8, alpha)

remfx/utils.py

@@ -0,0 +1,71 @@
+import logging
+from typing import List
+import pytorch_lightning as pl
+from omegaconf import DictConfig
+from pytorch_lightning.utilities import rank_zero_only
+
+
+def get_logger(name=__name__) -> logging.Logger:
+    """Initializes multi-GPU-friendly python command line logger."""
+
+    logger = logging.getLogger(name)
+
+    # this ensures all logging levels get marked with the rank zero decorator
+    # otherwise logs would get multiplied for each GPU process in multi-GPU setup
+    for level in (
+        "debug",
+        "info",
+        "warning",
+        "error",
+        "exception",
+        "fatal",
+        "critical",
+    ):
+        setattr(logger, level, rank_zero_only(getattr(logger, level)))
+
+    return logger
+
+
+log = get_logger(__name__)
+
+
+@rank_zero_only
+def log_hyperparameters(
+    config: DictConfig,
+    model: pl.LightningModule,
+    datamodule: pl.LightningDataModule,
+    trainer: pl.Trainer,
+    callbacks: List[pl.Callback],
+    logger: pl.loggers.logger.Logger,
+) -> None:
+    """Controls which config parts are saved by Lightning loggers.
+    Additionaly saves:
+    - number of model parameters
+    """
+
+    if not trainer.logger:
+        return
+
+    hparams = {}
+
+    # choose which parts of hydra config will be saved to loggers
+    hparams["model"] = config["model"]
+
+    # save number of model parameters
+    hparams["model/params/total"] = sum(p.numel() for p in model.parameters())
+    hparams["model/params/trainable"] = sum(
+        p.numel() for p in model.parameters() if p.requires_grad
+    )
+    hparams["model/params/non_trainable"] = sum(
+        p.numel() for p in model.parameters() if not p.requires_grad
+    )
+
+    hparams["datamodule"] = config["datamodule"]
+    hparams["trainer"] = config["trainer"]
+
+    if "seed" in config:
+        hparams["seed"] = config["seed"]
+    if "callbacks" in config:
+        hparams["callbacks"] = config["callbacks"]
+
+    logger.experiment.config.update(hparams)