diff --git a/README.md b/README.md index cd3b819..f747a42 100644 --- a/README.md +++ b/README.md @@ -18,6 +18,7 @@ SISU (Super Ingenious Sound Upscaler) is a project that uses GANs (Generative Ad 1. **Set Up**: - Make sure you have Python installed (version 3.8 or higher). - Install needed packages: `pip install -r requirements.txt` + - Install current version of PyTorch (CUDA/ROCm/What ever your device supports) 2. **Prepare Audio Data**: - Put your audio files in the `dataset/good` folder. diff --git a/data.py b/data.py index ac69730..bc7574f 100644 --- a/data.py +++ b/data.py @@ -4,22 +4,20 @@ import torch import torchaudio import os import random - import torchaudio.transforms as T import AudioUtils class AudioDataset(Dataset): - #audio_sample_rates = [8000, 11025, 16000, 22050] audio_sample_rates = [11025] + MAX_LENGTH = 44100 # Define your desired maximum length here - def __init__(self, input_dir): + def __init__(self, input_dir, device): self.input_files = [os.path.join(root, f) for root, _, files in os.walk(input_dir) for f in files if f.endswith('.wav')] - + self.device = device def __len__(self): return len(self.input_files) - def __getitem__(self, idx): # Load high-quality audio high_quality_audio, original_sample_rate = torchaudio.load(self.input_files[idx], normalize=True) @@ -32,4 +30,24 @@ class AudioDataset(Dataset): resample_transform_high = torchaudio.transforms.Resample(mangled_sample_rate, original_sample_rate) low_quality_audio = resample_transform_high(low_quality_audio) - return (AudioUtils.stereo_tensor_to_mono(high_quality_audio), original_sample_rate), (AudioUtils.stereo_tensor_to_mono(low_quality_audio), mangled_sample_rate) + high_quality_audio = AudioUtils.stereo_tensor_to_mono(high_quality_audio) + low_quality_audio = AudioUtils.stereo_tensor_to_mono(low_quality_audio) + + # Pad or truncate high-quality audio + if high_quality_audio.shape[1] < self.MAX_LENGTH: + padding = self.MAX_LENGTH - high_quality_audio.shape[1] + high_quality_audio = F.pad(high_quality_audio, (0, padding)) + elif high_quality_audio.shape[1] > self.MAX_LENGTH: + high_quality_audio = high_quality_audio[:, :self.MAX_LENGTH] + + # Pad or truncate low-quality audio + if low_quality_audio.shape[1] < self.MAX_LENGTH: + padding = self.MAX_LENGTH - low_quality_audio.shape[1] + low_quality_audio = F.pad(low_quality_audio, (0, padding)) + elif low_quality_audio.shape[1] > self.MAX_LENGTH: + low_quality_audio = low_quality_audio[:, :self.MAX_LENGTH] + + high_quality_audio = high_quality_audio.to(self.device) + low_quality_audio = low_quality_audio.to(self.device) + + return (high_quality_audio, original_sample_rate), (low_quality_audio, mangled_sample_rate) diff --git a/discriminator.py b/discriminator.py index af29f5d..dfd0126 100644 --- a/discriminator.py +++ b/discriminator.py @@ -2,40 +2,62 @@ import torch import torch.nn as nn import torch.nn.utils as utils -def discriminator_block(in_channels, out_channels, kernel_size=3, stride=1, dilation=1): +def discriminator_block(in_channels, out_channels, kernel_size=3, stride=1, dilation=1, spectral_norm=True, use_instance_norm=True): padding = (kernel_size // 2) * dilation - return nn.Sequential( - utils.spectral_norm(nn.Conv1d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, dilation=dilation, padding=padding)), - nn.BatchNorm1d(out_channels), - nn.LeakyReLU(0.2, inplace=True) # Changed activation to LeakyReLU + conv_layer = nn.Conv1d( + in_channels, + out_channels, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + padding=padding ) -class SISUDiscriminator(nn.Module): - def __init__(self): - super(SISUDiscriminator, self).__init__() - layers = 32 # Increased base layer count - self.model = nn.Sequential( - # Initial Convolution - discriminator_block(1, layers, kernel_size=7, stride=2, dilation=1), # Downsample + if spectral_norm: + conv_layer = utils.spectral_norm(conv_layer) - # Core Discriminator Blocks with varied kernels and dilations - discriminator_block(layers, layers * 2, kernel_size=5, stride=2, dilation=1), # Downsample - discriminator_block(layers * 2, layers * 2, kernel_size=3, dilation=2), - discriminator_block(layers * 2, layers * 4, kernel_size=5, dilation=4), - discriminator_block(layers * 4, layers * 4, kernel_size=3, dilation=8), - discriminator_block(layers * 4, layers * 8, kernel_size=5, dilation=16), - discriminator_block(layers * 8, layers * 8, kernel_size=3, dilation=8), - discriminator_block(layers * 8, layers * 4, kernel_size=5, dilation=4), - discriminator_block(layers * 4, layers * 2, kernel_size=3, dilation=2), - discriminator_block(layers * 2, layers, kernel_size=5, dilation=1), - # Final Convolution - discriminator_block(layers, 1, kernel_size=3, stride=1), + layers = [conv_layer] + layers.append(nn.LeakyReLU(0.2, inplace=True)) + + if use_instance_norm: + layers.append(nn.InstanceNorm1d(out_channels)) + + return nn.Sequential(*layers) + +class AttentionBlock(nn.Module): + def __init__(self, channels): + super(AttentionBlock, self).__init__() + self.attention = nn.Sequential( + nn.Conv1d(channels, channels // 4, kernel_size=1), + nn.ReLU(inplace=True), + nn.Conv1d(channels // 4, channels, kernel_size=1), + nn.Sigmoid() ) + + def forward(self, x): + attention_weights = self.attention(x) + return x * attention_weights + +class SISUDiscriminator(nn.Module): + def __init__(self, base_channels=16): + super(SISUDiscriminator, self).__init__() + layers = base_channels + self.model = nn.Sequential( + discriminator_block(1, layers, kernel_size=7, stride=1, spectral_norm=True, use_instance_norm=False), + discriminator_block(layers, layers * 2, kernel_size=5, stride=2, spectral_norm=True, use_instance_norm=True), + discriminator_block(layers * 2, layers * 4, kernel_size=5, stride=1, dilation=2, spectral_norm=True, use_instance_norm=True), + AttentionBlock(layers * 4), + discriminator_block(layers * 4, layers * 8, kernel_size=5, stride=1, dilation=4, spectral_norm=True, use_instance_norm=True), + discriminator_block(layers * 8, layers * 4, kernel_size=5, stride=2, spectral_norm=True, use_instance_norm=True), + discriminator_block(layers * 4, layers * 2, kernel_size=3, stride=1, spectral_norm=True, use_instance_norm=True), + discriminator_block(layers * 2, layers, kernel_size=3, stride=1, spectral_norm=True, use_instance_norm=True), + discriminator_block(layers, 1, kernel_size=3, stride=1, spectral_norm=False, use_instance_norm=False) + ) + self.global_avg_pool = nn.AdaptiveAvgPool1d(1) def forward(self, x): - # Gaussian noise is not necessary here for discriminator as it is already implicit in the training process x = self.model(x) x = self.global_avg_pool(x) - x = x.view(-1, 1) + x = x.view(x.size(0), -1) return x diff --git a/file_utils.py b/file_utils.py new file mode 100644 index 0000000..a723688 --- /dev/null +++ b/file_utils.py @@ -0,0 +1,28 @@ +import json + +filepath = "my_data.json" + +def write_data(filepath, data): + try: + with open(filepath, 'w') as f: + json.dump(data, f, indent=4) # Use indent for pretty formatting + print(f"Data written to '{filepath}'") + except Exception as e: + print(f"Error writing to file: {e}") + + +def read_data(filepath): + try: + with open(filepath, 'r') as f: + data = json.load(f) + print(f"Data read from '{filepath}'") + return data + except FileNotFoundError: + print(f"File not found: {filepath}") + return None + except json.JSONDecodeError: + print(f"Error decoding JSON from file: {filepath}") + return None + except Exception as e: + print(f"Error reading from file: {e}") + return None diff --git a/generator.py b/generator.py index 6ea267d..a53feb7 100644 --- a/generator.py +++ b/generator.py @@ -1,39 +1,74 @@ +import torch import torch.nn as nn def conv_block(in_channels, out_channels, kernel_size=3, dilation=1): return nn.Sequential( - nn.Conv1d(in_channels, out_channels, kernel_size=kernel_size, dilation=dilation, padding=(kernel_size // 2) * dilation), - nn.BatchNorm1d(out_channels), + nn.Conv1d( + in_channels, + out_channels, + kernel_size=kernel_size, + dilation=dilation, + padding=(kernel_size // 2) * dilation + ), + nn.InstanceNorm1d(out_channels), nn.PReLU() ) -class SISUGenerator(nn.Module): - def __init__(self): - super(SISUGenerator, self).__init__() - layer = 32 # Increased base layer count - self.conv1 = nn.Sequential( - nn.Conv1d(1, layer, kernel_size=7, padding=3), - nn.BatchNorm1d(layer), - nn.PReLU(), - ) - self.conv_blocks = nn.Sequential( - conv_block(layer, layer, kernel_size=3, dilation=1), # Local details - conv_block(layer, layer*2, kernel_size=5, dilation=2), # Local Context - conv_block(layer*2, layer*2, kernel_size=3, dilation=4), # Wider context - conv_block(layer*2, layer*4, kernel_size=7, dilation=8), # Longer range dependencies - conv_block(layer*4, layer*4, kernel_size=3, dilation=16), # Longer range dependencies - conv_block(layer*4, layer*2, kernel_size=5, dilation=8), # Wider context - conv_block(layer*2, layer*2, kernel_size=3, dilation=4), # Wider context - conv_block(layer*2, layer, kernel_size=5, dilation=2), # Local Context - conv_block(layer, layer, kernel_size=3, dilation=1), # Local details - ) - self.final_layer = nn.Sequential( - nn.Conv1d(layer, 1, kernel_size=3, padding=1), +class AttentionBlock(nn.Module): + """ + Simple Channel Attention Block. Learns to weight channels based on their importance. + """ + def __init__(self, channels): + super(AttentionBlock, self).__init__() + self.attention = nn.Sequential( + nn.Conv1d(channels, channels // 4, kernel_size=1), + nn.ReLU(inplace=True), + nn.Conv1d(channels // 4, channels, kernel_size=1), + nn.Sigmoid() ) + def forward(self, x): + attention_weights = self.attention(x) + return x * attention_weights + +class ResidualInResidualBlock(nn.Module): + def __init__(self, channels, num_convs=3): + super(ResidualInResidualBlock, self).__init__() + + self.conv_layers = nn.Sequential( + *[conv_block(channels, channels) for _ in range(num_convs)] + ) + + self.attention = AttentionBlock(channels) + def forward(self, x): residual = x - x = self.conv1(x) - x = self.conv_blocks(x) - x = self.final_layer(x) + x = self.conv_layers(x) + x = self.attention(x) return x + residual + +class SISUGenerator(nn.Module): + def __init__(self, channels=16, num_rirb=4, alpha=1.0): + super(SISUGenerator, self).__init__() + self.alpha = alpha + + self.conv1 = nn.Sequential( + nn.Conv1d(1, channels, kernel_size=7, padding=3), + nn.InstanceNorm1d(channels), + nn.PReLU(), + ) + + self.rir_blocks = nn.Sequential( + *[ResidualInResidualBlock(channels) for _ in range(num_rirb)] + ) + + self.final_layer = nn.Conv1d(channels, 1, kernel_size=3, padding=1) + + def forward(self, x): + residual_input = x + x = self.conv1(x) + x_rirb_out = self.rir_blocks(x) + learned_residual = self.final_layer(x_rirb_out) + output = residual_input + self.alpha * learned_residual + + return output diff --git a/requirements.txt b/requirements.txt index 5cb5df1..21f6bef 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,11 +4,9 @@ Jinja2==3.1.4 MarkupSafe==2.1.5 mpmath==1.3.0 networkx==3.4.2 -numpy==2.2.1 -pytorch-triton-rocm==3.2.0+git0d4682f0 +numpy==2.2.3 +pillow==11.0.0 setuptools==70.2.0 -sympy==1.13.1 -torch==2.6.0.dev20241222+rocm6.2.4 -torchaudio==2.6.0.dev20241222+rocm6.2.4 +sympy==1.13.3 tqdm==4.67.1 typing_extensions==4.12.2 diff --git a/training.py b/training.py index e114817..01ea749 100644 --- a/training.py +++ b/training.py @@ -10,6 +10,8 @@ import argparse import math +import os + from torch.utils.data import random_split from torch.utils.data import DataLoader @@ -18,42 +20,10 @@ from data import AudioDataset from generator import SISUGenerator from discriminator import SISUDiscriminator -def perceptual_loss(y_true, y_pred): - return torch.mean((y_true - y_pred) ** 2) +from training_utils import discriminator_train, generator_train +import file_utils as Data -def discriminator_train(high_quality, low_quality, real_labels, fake_labels): - optimizer_d.zero_grad() - - # Forward pass for real samples - discriminator_decision_from_real = discriminator(high_quality[0]) - d_loss_real = criterion_d(discriminator_decision_from_real, real_labels) - - # Forward pass for fake samples (from generator output) - generator_output = generator(low_quality[0]) - discriminator_decision_from_fake = discriminator(generator_output.detach()) - d_loss_fake = criterion_d(discriminator_decision_from_fake, fake_labels) - - # Combine real and fake losses - d_loss = (d_loss_real + d_loss_fake) / 2.0 - - # Backward pass and optimization - d_loss.backward() - nn.utils.clip_grad_norm_(discriminator.parameters(), max_norm=1.0) # Gradient Clipping - optimizer_d.step() - - return d_loss - -def generator_train(low_quality, real_labels): - optimizer_g.zero_grad() - - # Forward pass for fake samples (from generator output) - generator_output = generator(low_quality[0]) - discriminator_decision = discriminator(generator_output) - g_loss = criterion_g(discriminator_decision, real_labels) - - g_loss.backward() - optimizer_g.step() - return generator_output +import torchaudio.transforms as T # Init script argument parser parser = argparse.ArgumentParser(description="Training script") @@ -61,47 +31,78 @@ parser.add_argument("--generator", type=str, default=None, help="Path to the generator model file") parser.add_argument("--discriminator", type=str, default=None, help="Path to the discriminator model file") +parser.add_argument("--device", type=str, default="cpu", help="Select device") +parser.add_argument("--epoch", type=int, default=0, help="Current epoch for model versioning") +parser.add_argument("--debug", action="store_true", help="Print debug logs") +parser.add_argument("--continue_training", action="store_true", help="Continue training using temp_generator and temp_discriminator models") args = parser.parse_args() -# Check for CUDA availability -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +device = torch.device(args.device if torch.cuda.is_available() else "cpu") print(f"Using device: {device}") +# Parameters +sample_rate = 44100 +n_fft = 2048 +hop_length = 256 +win_length = n_fft +n_mels = 128 +n_mfcc = 20 # If using MFCC + +mfcc_transform = T.MFCC( + sample_rate, + n_mfcc, + melkwargs = {'n_fft': n_fft, 'hop_length': hop_length} +).to(device) + +mel_transform = T.MelSpectrogram( + sample_rate=sample_rate, n_fft=n_fft, hop_length=hop_length, + win_length=win_length, n_mels=n_mels, power=1.0 # Magnitude Mel +).to(device) + +stft_transform = T.Spectrogram( + n_fft=n_fft, win_length=win_length, hop_length=hop_length +).to(device) + +debug = args.debug + # Initialize dataset and dataloader dataset_dir = './dataset/good' -dataset = AudioDataset(dataset_dir) - -# ========= MULTIPLE ========= - -# dataset_size = len(dataset) -# train_size = int(dataset_size * .9) -# val_size = int(dataset_size-train_size) - -#train_dataset, val_dataset = random_split(dataset, [train_size, val_size]) - -# train_data_loader = DataLoader(train_dataset, batch_size=1, shuffle=True) -# val_data_loader = DataLoader(val_dataset, batch_size=1, shuffle=True) +dataset = AudioDataset(dataset_dir, device) +models_dir = "models" +os.makedirs(models_dir, exist_ok=True) +audio_output_dir = "output" +os.makedirs(audio_output_dir, exist_ok=True) # ========= SINGLE ========= -train_data_loader = DataLoader(dataset, batch_size=1, shuffle=True) +train_data_loader = DataLoader(dataset, batch_size=64, shuffle=True) + + +# ========= MODELS ========= -# Initialize models and move them to device generator = SISUGenerator() discriminator = SISUDiscriminator() -if args.generator is not None: - generator.load_state_dict(torch.load(args.generator, weights_only=True)) -if args.discriminator is not None: - discriminator.load_state_dict(torch.load(args.discriminator, weights_only=True)) +epoch: int = args.epoch +epoch_from_file = Data.read_data(f"{models_dir}/epoch_data.json") + +if args.continue_training: + generator.load_state_dict(torch.load(f"{models_dir}/temp_generator.pt", map_location=device, weights_only=True)) + discriminator.load_state_dict(torch.load(f"{models_dir}/temp_generator.pt", map_location=device, weights_only=True)) + epoch = epoch_from_file["epoch"] + 1 +else: + if args.generator is not None: + generator.load_state_dict(torch.load(args.generator, map_location=device, weights_only=True)) + if args.discriminator is not None: + discriminator.load_state_dict(torch.load(args.discriminator, map_location=device, weights_only=True)) generator = generator.to(device) discriminator = discriminator.to(device) # Loss -criterion_g = nn.MSELoss() -criterion_d = nn.BCELoss() +criterion_g = nn.BCEWithLogitsLoss() +criterion_d = nn.BCEWithLogitsLoss() # Optimizers optimizer_g = optim.Adam(generator.parameters(), lr=0.0001, betas=(0.5, 0.999)) @@ -112,31 +113,6 @@ scheduler_g = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer_g, mode='min' scheduler_d = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer_d, mode='min', factor=0.5, patience=5) def start_training(): - - # Training loop - - # ========= DISCRIMINATOR PRE-TRAINING ========= - # discriminator_epochs = 1 - # for discriminator_epoch in range(discriminator_epochs): - - # # ========= TRAINING ========= - # for high_quality_clip, low_quality_clip in tqdm.tqdm(train_data_loader, desc=f"Epoch {discriminator_epoch+1}/{discriminator_epochs}"): - # high_quality_sample = high_quality_clip[0].to(device) - # low_quality_sample = low_quality_clip[0].to(device) - - # scale = high_quality_clip[0].shape[2]/low_quality_clip[0].shape[2] - - # # ========= LABELS ========= - # batch_size = high_quality_sample.size(0) - # real_labels = torch.ones(batch_size, 1).to(device) - # fake_labels = torch.zeros(batch_size, 1).to(device) - - # # ========= DISCRIMINATOR ========= - # discriminator.train() - # discriminator_train(high_quality_sample, low_quality_sample, scale, real_labels, fake_labels) - - # torch.save(discriminator.state_dict(), "models/discriminator-single-shot-pre-train.pt") - generator_epochs = 5000 for generator_epoch in range(generator_epochs): low_quality_audio = (torch.empty((1)), 1) @@ -146,10 +122,10 @@ def start_training(): times_correct = 0 # ========= TRAINING ========= - for high_quality_clip, low_quality_clip in tqdm.tqdm(train_data_loader, desc=f"Epoch {generator_epoch+1}/{generator_epochs}"): + for high_quality_clip, low_quality_clip in tqdm.tqdm(train_data_loader, desc=f"Training epoch {generator_epoch+1}/{generator_epochs}, Current epoch {epoch+1}"): # for high_quality_clip, low_quality_clip in train_data_loader: - high_quality_sample = (high_quality_clip[0].to(device), high_quality_clip[1]) - low_quality_sample = (low_quality_clip[0].to(device), low_quality_clip[1]) + high_quality_sample = (high_quality_clip[0], high_quality_clip[1]) + low_quality_sample = (low_quality_clip[0], low_quality_clip[1]) # ========= LABELS ========= batch_size = high_quality_clip[0].size(0) @@ -158,32 +134,61 @@ def start_training(): # ========= DISCRIMINATOR ========= discriminator.train() - discriminator_train(high_quality_sample, low_quality_sample, real_labels, fake_labels) + d_loss = discriminator_train( + high_quality_sample, + low_quality_sample, + real_labels, + fake_labels, + discriminator, + generator, + criterion_d, + optimizer_d + ) # ========= GENERATOR ========= generator.train() - generator_output = generator_train(low_quality_sample, real_labels) + generator_output, combined_loss, adversarial_loss, mel_l1_tensor, log_stft_l1_tensor, mfcc_l_tensor = generator_train( + low_quality_sample, + high_quality_sample, + real_labels, + generator, + discriminator, + criterion_d, + optimizer_g, + device, + mel_transform, + stft_transform, + mfcc_transform + ) + + if debug: + print(f"D_LOSS: {d_loss.item():.4f}, COMBINED_LOSS: {combined_loss.item():.4f}, ADVERSARIAL_LOSS: {adversarial_loss.item():.4f}, MEL_L1_LOSS: {mel_l1_tensor.item():.4f}, LOG_STFT_L1_LOSS: {log_stft_l1_tensor.item():.4f}, MFCC_LOSS: {mfcc_l_tensor.item():.4f}") + scheduler_d.step(d_loss.detach()) + scheduler_g.step(adversarial_loss.detach()) # ========= SAVE LATEST AUDIO ========= - high_quality_audio = high_quality_clip - low_quality_audio = low_quality_clip - ai_enhanced_audio = (generator_output, high_quality_clip[1]) + high_quality_audio = (high_quality_clip[0][0], high_quality_clip[1][0]) + low_quality_audio = (low_quality_clip[0][0], low_quality_clip[1][0]) + ai_enhanced_audio = (generator_output[0], high_quality_clip[1][0]) - #metric = snr(high_quality_audio[0].to(device), ai_enhanced_audio[0]) - #print(f"Generator metric {metric}!") - #scheduler_g.step(metric) + new_epoch = generator_epoch+epoch - if generator_epoch % 10 == 0: - print(f"Saved epoch {generator_epoch}!") - torchaudio.save(f"./output/epoch-{generator_epoch}-audio-crap.wav", low_quality_audio[0][0].cpu(), high_quality_audio[1]) # <-- Because audio clip was resampled in data.py from original to crap and to original again. - torchaudio.save(f"./output/epoch-{generator_epoch}-audio-ai.wav", ai_enhanced_audio[0][0].cpu(), ai_enhanced_audio[1]) - torchaudio.save(f"./output/epoch-{generator_epoch}-audio-orig.wav", high_quality_audio[0][0].cpu(), high_quality_audio[1]) + if generator_epoch % 25 == 0: + print(f"Saved epoch {new_epoch}!") + torchaudio.save(f"{audio_output_dir}/epoch-{new_epoch}-audio-crap.wav", low_quality_audio[0].cpu().detach(), high_quality_audio[1]) # <-- Because audio clip was resampled in data.py from original to crap and to original again. + torchaudio.save(f"{audio_output_dir}/epoch-{new_epoch}-audio-ai.wav", ai_enhanced_audio[0].cpu().detach(), ai_enhanced_audio[1]) + torchaudio.save(f"{audio_output_dir}/epoch-{new_epoch}-audio-orig.wav", high_quality_audio[0].cpu().detach(), high_quality_audio[1]) - torch.save(discriminator.state_dict(), f"models/current-epoch-discriminator.pt") - torch.save(generator.state_dict(), f"models/current-epoch-generator.pt") + #if debug: + # print(generator.state_dict().keys()) + # print(discriminator.state_dict().keys()) + torch.save(discriminator.state_dict(), f"{models_dir}/temp_discriminator.pt") + torch.save(generator.state_dict(), f"{models_dir}/temp_generator.pt") + Data.write_data(f"{models_dir}/epoch_data.json", {"epoch": new_epoch}) - torch.save(discriminator.state_dict(), "models/epoch-5000-discriminator.pt") - torch.save(generator.state_dict(), "models/epoch-5000-generator.pt") + + torch.save(discriminator, "models/epoch-5000-discriminator.pt") + torch.save(generator, "models/epoch-5000-generator.pt") print("Training complete!") start_training() diff --git a/training_utils.py b/training_utils.py new file mode 100644 index 0000000..6f26f58 --- /dev/null +++ b/training_utils.py @@ -0,0 +1,144 @@ +import torch +import torch.nn as nn +import torch.optim as optim + +import torchaudio +import torchaudio.transforms as T + +def gpu_mfcc_loss(mfcc_transform, y_true, y_pred): + mfccs_true = mfcc_transform(y_true) + mfccs_pred = mfcc_transform(y_pred) + + min_len = min(mfccs_true.shape[2], mfccs_pred.shape[2]) + mfccs_true = mfccs_true[:, :, :min_len] + mfccs_pred = mfccs_pred[:, :, :min_len] + + loss = torch.mean((mfccs_true - mfccs_pred)**2) + return loss + +def mel_spectrogram_l1_loss(mel_transform: T.MelSpectrogram, y_true: torch.Tensor, y_pred: torch.Tensor) -> torch.Tensor: + mel_spec_true = mel_transform(y_true) + mel_spec_pred = mel_transform(y_pred) + + # Ensure same time dimension length (due to potential framing differences) + min_len = min(mel_spec_true.shape[-1], mel_spec_pred.shape[-1]) + mel_spec_true = mel_spec_true[..., :min_len] + mel_spec_pred = mel_spec_pred[..., :min_len] + + # L1 Loss (Mean Absolute Error) + loss = torch.mean(torch.abs(mel_spec_true - mel_spec_pred)) + return loss + +def mel_spectrogram_l2_loss(mel_transform: T.MelSpectrogram, y_true: torch.Tensor, y_pred: torch.Tensor) -> torch.Tensor: + mel_spec_true = mel_transform(y_true) + mel_spec_pred = mel_transform(y_pred) + + min_len = min(mel_spec_true.shape[-1], mel_spec_pred.shape[-1]) + mel_spec_true = mel_spec_true[..., :min_len] + mel_spec_pred = mel_spec_pred[..., :min_len] + + loss = torch.mean((mel_spec_true - mel_spec_pred)**2) + return loss + +def log_stft_magnitude_loss(stft_transform: T.Spectrogram, y_true: torch.Tensor, y_pred: torch.Tensor, eps: float = 1e-7) -> torch.Tensor: + stft_mag_true = stft_transform(y_true) + stft_mag_pred = stft_transform(y_pred) + + min_len = min(stft_mag_true.shape[-1], stft_mag_pred.shape[-1]) + stft_mag_true = stft_mag_true[..., :min_len] + stft_mag_pred = stft_mag_pred[..., :min_len] + + loss = torch.mean(torch.abs(torch.log(stft_mag_true + eps) - torch.log(stft_mag_pred + eps))) + return loss + +def spectral_convergence_loss(stft_transform: T.Spectrogram, y_true: torch.Tensor, y_pred: torch.Tensor, eps: float = 1e-7) -> torch.Tensor: + stft_mag_true = stft_transform(y_true) + stft_mag_pred = stft_transform(y_pred) + + min_len = min(stft_mag_true.shape[-1], stft_mag_pred.shape[-1]) + stft_mag_true = stft_mag_true[..., :min_len] + stft_mag_pred = stft_mag_pred[..., :min_len] + + norm_true = torch.linalg.norm(stft_mag_true, ord='fro', dim=(-2, -1)) + norm_diff = torch.linalg.norm(stft_mag_true - stft_mag_pred, ord='fro', dim=(-2, -1)) + + loss = torch.mean(norm_diff / (norm_true + eps)) + return loss + +def discriminator_train(high_quality, low_quality, real_labels, fake_labels, discriminator, generator, criterion, optimizer): + optimizer.zero_grad() + + # Forward pass for real samples + discriminator_decision_from_real = discriminator(high_quality[0]) + d_loss_real = criterion(discriminator_decision_from_real, real_labels) + + with torch.no_grad(): + generator_output = generator(low_quality[0]) + discriminator_decision_from_fake = discriminator(generator_output) + d_loss_fake = criterion(discriminator_decision_from_fake, fake_labels.expand_as(discriminator_decision_from_fake)) + + d_loss = (d_loss_real + d_loss_fake) / 2.0 + + d_loss.backward() + # Optional: Gradient Clipping (can be helpful) + # nn.utils.clip_grad_norm_(discriminator.parameters(), max_norm=1.0) # Gradient Clipping + optimizer.step() + + return d_loss + +def generator_train( + low_quality, + high_quality, + real_labels, + generator, + discriminator, + adv_criterion, + g_optimizer, + device, + mel_transform: T.MelSpectrogram, + stft_transform: T.Spectrogram, + mfcc_transform: T.MFCC, + lambda_adv: float = 1.0, + lambda_mel_l1: float = 10.0, + lambda_log_stft: float = 1.0, + lambda_mfcc: float = 1.0 +): + g_optimizer.zero_grad() + + generator_output = generator(low_quality[0]) + + discriminator_decision = discriminator(generator_output) + adversarial_loss = adv_criterion(discriminator_decision, real_labels.expand_as(discriminator_decision)) + + mel_l1 = 0.0 + log_stft_l1 = 0.0 + mfcc_l = 0.0 + + # Calculate Mel L1 Loss if weight is positive + if lambda_mel_l1 > 0: + mel_l1 = mel_spectrogram_l1_loss(mel_transform, high_quality[0], generator_output) + + # Calculate Log STFT L1 Loss if weight is positive + if lambda_log_stft > 0: + log_stft_l1 = log_stft_magnitude_loss(stft_transform, high_quality[0], generator_output) + + # Calculate MFCC Loss if weight is positive + if lambda_mfcc > 0: + mfcc_l = gpu_mfcc_loss(mfcc_transform, high_quality[0], generator_output) + + mel_l1_tensor = torch.tensor(mel_l1, device=device) if isinstance(mel_l1, float) else mel_l1 + log_stft_l1_tensor = torch.tensor(log_stft_l1, device=device) if isinstance(log_stft_l1, float) else log_stft_l1 + mfcc_l_tensor = torch.tensor(mfcc_l, device=device) if isinstance(mfcc_l, float) else mfcc_l + + combined_loss = (lambda_adv * adversarial_loss) + \ + (lambda_mel_l1 * mel_l1_tensor) + \ + (lambda_log_stft * log_stft_l1_tensor) + \ + (lambda_mfcc * mfcc_l_tensor) + + combined_loss.backward() + # Optional: Gradient Clipping + # nn.utils.clip_grad_norm_(generator.parameters(), max_norm=1.0) + g_optimizer.step() + + # 6. Return values for logging + return generator_output, combined_loss, adversarial_loss, mel_l1_tensor, log_stft_l1_tensor, mfcc_l_tensor