Initial ptl model and training script for umx

.gitignore

@@ -4,4 +4,6 @@ wandb/
 *.egg-info/
 data/
 .DS_Store
-__pycache__/
+__pycache__/
+lightning_logs/
+RemFX/

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "umx"]
+	path = umx
+	url = https://github.com/sigsep/open-unmix-pytorch

models.py

@@ -1,44 +1,103 @@
-from audio_diffusion_pytorch import AudioDiffusionModel
 import torch
 from torch import Tensor
 import pytorch_lightning as pl
 from einops import rearrange
 import wandb
+from audio_diffusion_pytorch import AudioDiffusionModel
+
+import sys
+
+sys.path.append("/Users/matthewrice/Developer/remfx/umx/")
+from umx.openunmix.model import OpenUnmix, Separator
+
 
 SAMPLE_RATE = 22050  # From audio-diffusion-pytorch
 
 
-class TCNWrapper(pl.LightningModule):
-    def __init__(self):
+class OpenUnmixModel(pl.LightningModule):
+    def __init__(
+        self,
+        n_fft: int = 2048,
+        hop_length: int = 512,
+        alpha: float = 0.3,
+    ):
         super().__init__()
-        self.model = AudioDiffusionModel(in_channels=1)
+        self.model = OpenUnmix(
+            nb_channels=1,
+            nb_bins=n_fft // 2 + 1,
+        )
+        self.n_fft = n_fft
+        self.hop_length = hop_length
+        self.alpha = alpha
+        window = torch.hann_window(n_fft)
+        self.register_buffer("window", window)
 
     def forward(self, x: torch.Tensor):
         return self.model(x)
 
     def training_step(self, batch, batch_idx):
-        loss = self.common_step(batch, batch_idx, mode="train")
+        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="val")
+        loss, Y = self.common_step(batch, batch_idx, mode="val")
+        return loss, Y
 
     def common_step(self, batch, batch_idx, mode: str = "train"):
         x, target, label = batch
-        loss = self(x)
+        X = spectrogram(x, self.window, self.n_fft, self.hop_length, self.alpha)
+        Y = self(X)
+        Y_hat = spectrogram(
+            target, self.window, self.n_fft, self.hop_length, self.alpha
+        )
+        loss = torch.nn.functional.mse_loss(Y, Y_hat)
         self.log(f"{mode}_loss", loss, on_step=True, on_epoch=True)
-        return loss
+        return loss, Y
 
     def configure_optimizers(self):
         return torch.optim.Adam(
             self.parameters(), lr=1e-4, betas=(0.95, 0.999), eps=1e-6, weight_decay=1e-3
         )
 
+    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
+            s = Separator(
+                target_models={"other": self.model},
+                nb_channels=1,
+                sample_rate=SAMPLE_RATE,
+                n_fft=self.n_fft,
+                n_hop=self.hop_length,
+            )
+            outputs = s(x).squeeze(1)
+            log_wandb_audio_batch(
+                id="sample",
+                samples=x,
+                sampling_rate=SAMPLE_RATE,
+                caption=f"Epoch {self.current_epoch}",
+            )
+            log_wandb_audio_batch(
+                id="prediction",
+                samples=outputs,
+                sampling_rate=SAMPLE_RATE,
+                caption=f"Epoch {self.current_epoch}",
+            )
+            log_wandb_audio_batch(
+                id="target",
+                samples=target,
+                sampling_rate=SAMPLE_RATE,
+                caption=f"Epoch {self.current_epoch}",
+            )
+            self.log_next = False
+
 
-class AudioDiffusionWrapper(pl.LightningModule):
-    def __init__(self):
+class DiffusionGenerationModel(pl.LightningModule):
+    def __init__(self, model: torch.nn.Module):
         super().__init__()
-        self.model = AudioDiffusionModel(in_channels=1)
+        self.model = model
 
     def forward(self, x: torch.Tensor):
         return self.model(x)
@@ -77,10 +136,8 @@ class AudioDiffusionWrapper(pl.LightningModule):
     def log_sample(self, batch, num_steps=10):
         # Get start diffusion noise
         noise = torch.randn(batch.shape, device=self.device)
-        sampled = self.model.sample(
-            noise=noise, num_steps=num_steps  # Suggested range: 2-50
-        )
-        self.log_wandb_audio_batch(
+        sampled = self.sample(noise=noise, num_steps=num_steps)  # Suggested range: 2-50
+        log_wandb_audio_batch(
             id="sample",
             samples=sampled,
             sampling_rate=SAMPLE_RATE,
@@ -96,10 +153,34 @@ def log_wandb_audio_batch(
     for idx in range(num_items):
         wandb.log(
             {
-                f"sample_{idx}_{id}": wandb.Audio(
+                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)

train.py

@@ -3,17 +3,18 @@ import pytorch_lightning as pl
 import torch
 from torch.utils.data import DataLoader
 from datasets import GuitarFXDataset
-from models import AudioDiffusionWrapper
+from models import DiffusionGenerationModel, OpenUnmixModel
+
 
 SAMPLE_RATE = 22050
 TRAIN_SPLIT = 0.8
 
 
 def main():
-    # wandb_logger = WandbLogger(project="RemFX", save_dir="./")
-    trainer = pl.Trainer()  # logger=wandb_logger)
+    wandb_logger = WandbLogger(project="RemFX", save_dir="./")
+    trainer = pl.Trainer(logger=wandb_logger, max_epochs=10)
     guitfx = GuitarFXDataset(
-        root="/Users/matthewrice/mir_datasets/egfxset",
+        root="/Users/matthewrice/Developer/remfx/data/egfx",
         sample_rate=SAMPLE_RATE,
         effect_type=["Phaser"],
     )
@@ -24,7 +25,10 @@ def main():
     )
     train = DataLoader(train_dataset, batch_size=2)
     val = DataLoader(val_dataset, batch_size=2)
-    model = AudioDiffusionWrapper()
+
+    # model = DiffusionGenerationModel()
+    model = OpenUnmixModel()
+
     trainer.fit(model=model, train_dataloaders=train, val_dataloaders=val)
 
 

umx

@@ -0,0 +1 @@
+Subproject commit 05fd4d8a0e3e50e308579052d762a342647c3408