import torch
from torch import Tensor, nn
import pytorch_lightning as pl
from einops import rearrange
import wandb
from audio_diffusion_pytorch import DiffusionModel
from auraloss.time import SISDRLoss
from auraloss.freq import MultiResolutionSTFTLoss, STFTLoss
from torch.nn import L1Loss
from frechet_audio_distance import FrechetAudioDistance
import numpy as np
from umx.openunmix.model import OpenUnmix, Separator
from torchaudio.models import HDemucs
class FADLoss(torch.nn.Module):
def __init__(self, sample_rate: float):
super().__init__()
self.fad = FrechetAudioDistance(
use_pca=False, use_activation=False, verbose=False
)
self.fad.model = self.fad.model.to("cpu")
self.sr = sample_rate
def forward(self, audio_background, audio_eval):
embds_background = []
embds_eval = []
for sample in audio_background:
embd = self.fad.model.forward(sample.T.cpu().detach().numpy(), self.sr)
embds_background.append(embd.cpu().detach().numpy())
for sample in audio_eval:
embd = self.fad.model.forward(sample.T.cpu().detach().numpy(), self.sr)
embds_eval.append(embd.cpu().detach().numpy())
embds_background = np.concatenate(embds_background, axis=0)
embds_eval = np.concatenate(embds_eval, axis=0)
mu_background, sigma_background = self.fad.calculate_embd_statistics(
embds_background
)
mu_eval, sigma_eval = self.fad.calculate_embd_statistics(embds_eval)
fad_score = self.fad.calculate_frechet_distance(
mu_background, sigma_background, mu_eval, sigma_eval
)
return fad_score
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
self.metrics = torch.nn.ModuleDict(
{
"SISDR": SISDRLoss(),
"STFT": STFTLoss(),
"FAD": FADLoss(sample_rate=sample_rate),
}
)
# Log first batch metrics input vs output only once
self.log_first_metrics = True
self.log_train_audio = True
@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")
return loss
def common_step(self, batch, batch_idx, mode: str = "train"):
loss, output = self.model(batch)
self.log(f"{mode}_loss", loss)
x, y, label = batch
# Metric logging
for metric in self.metrics:
# SISDR returns negative values, so negate them
if metric == "SISDR":
negate = -1
else:
negate = 1
self.log(
f"{mode}_{metric}",
negate * self.metrics[metric](output, y),
on_step=False,
on_epoch=True,
logger=True,
prog_bar=True,
sync_dist=True,
)
return loss
def on_train_batch_start(self, batch, batch_idx):
if self.log_train_audio:
x, y, label = batch
# Concat samples together for easier viewing in dashboard
input_samples = rearrange(x, "b c t -> c (b t)").unsqueeze(0)
target_samples = rearrange(y, "b c t -> c (b t)").unsqueeze(0)
log_wandb_audio_batch(
logger=self.logger,
id="input_effected_audio",
samples=input_samples.cpu(),
sampling_rate=self.sample_rate,
caption="Training Data",
)
log_wandb_audio_batch(
logger=self.logger,
id="target_audio",
samples=target_samples.cpu(),
sampling_rate=self.sample_rate,
caption="Target Data",
)
self.log_train_audio = False
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
# Log Input Metrics
for metric in self.metrics:
# SISDR returns negative values, so negate them
if metric == "SISDR":
negate = -1
else:
negate = 1
self.log(
f"Input_{metric}",
negate * self.metrics[metric](x, target),
on_step=False,
on_epoch=True,
logger=True,
prog_bar=True,
sync_dist=True,
)
self.model.eval()
with torch.no_grad():
y = self.model.sample(x)
# Concat samples together for easier viewing in dashboard
# 2 seconds of silence between each sample
silence = torch.zeros_like(x)
silence = silence[:, : self.sample_rate * 2]
concat_samples = torch.cat([y, silence, x, silence, target], dim=-1)
log_wandb_audio_batch(
logger=self.logger,
id="prediction_input_target",
samples=concat_samples.cpu(),
sampling_rate=self.sample_rate,
caption=f"Epoch {self.current_epoch}",
)
self.log_next = False
self.model.train()
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 = 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, sep_out
def sample(self, x: Tensor) -> Tensor:
return self.separator(x).squeeze(1)
class DemucsModel(torch.nn.Module):
def __init__(self, sample_rate, **kwargs) -> None:
super().__init__()
self.model = HDemucs(**kwargs)
self.num_bins = kwargs["nfft"] // 2 + 1
self.loss_fn = MultiResolutionSTFTLoss(
n_bins=self.num_bins, sample_rate=sample_rate
)
def forward(self, batch):
x, target, label = batch
output = self.model(x).squeeze(1)
loss = self.loss_fn(output, target)
return loss, output
def sample(self, x: Tensor) -> Tensor:
return self.model(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
sampled_out = self.model.sample(x)
return self.model(x), sampled_out
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)