From 0790a0d3dab77d071ba5d8c3bc7f68b7845a11c4 Mon Sep 17 00:00:00 2001 From: NikkeDoy Date: Sat, 25 Jan 2025 16:48:10 +0200 Subject: [PATCH 01/14] :alembic: | Experimenting with smaller architecture. --- discriminator.py | 5 +---- generator.py | 5 +---- 2 files changed, 2 insertions(+), 8 deletions(-) diff --git a/discriminator.py b/discriminator.py index af29f5d..b800eda 100644 --- a/discriminator.py +++ b/discriminator.py @@ -20,12 +20,9 @@ class SISUDiscriminator(nn.Module): # 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 * 8, layers * 4, kernel_size=3, dilation=8), discriminator_block(layers * 4, layers * 2, kernel_size=3, dilation=2), discriminator_block(layers * 2, layers, kernel_size=5, dilation=1), # Final Convolution diff --git a/generator.py b/generator.py index 6ea267d..2446275 100644 --- a/generator.py +++ b/generator.py @@ -18,12 +18,9 @@ class SISUGenerator(nn.Module): ) 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, layer*4, kernel_size=5, dilation=2), # Local Context 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 ) -- 2.43.0 From fb7b624c877d6d7d9016b75f04c04e3b39c69761 Mon Sep 17 00:00:00 2001 From: NikkeDoy Date: Mon, 10 Feb 2025 12:44:42 +0200 Subject: [PATCH 02/14] :alembic: | Experimenting with very small model. --- discriminator.py | 9 ++++----- generator.py | 8 ++++---- training.py | 2 +- 3 files changed, 9 insertions(+), 10 deletions(-) diff --git a/discriminator.py b/discriminator.py index b800eda..b1d82e1 100644 --- a/discriminator.py +++ b/discriminator.py @@ -13,7 +13,7 @@ def discriminator_block(in_channels, out_channels, kernel_size=3, stride=1, dila class SISUDiscriminator(nn.Module): def __init__(self): super(SISUDiscriminator, self).__init__() - layers = 32 # Increased base layer count + layers = 4 # Increased base layer count self.model = nn.Sequential( # Initial Convolution discriminator_block(1, layers, kernel_size=7, stride=2, dilation=1), # Downsample @@ -21,10 +21,9 @@ class SISUDiscriminator(nn.Module): # 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 * 4, kernel_size=5, dilation=4), - discriminator_block(layers * 4, layers * 8, kernel_size=5, dilation=16), - discriminator_block(layers * 8, layers * 4, kernel_size=3, dilation=8), - discriminator_block(layers * 4, layers * 2, kernel_size=3, dilation=2), - discriminator_block(layers * 2, layers, kernel_size=5, dilation=1), + discriminator_block(layers * 4, layers * 4, kernel_size=5, dilation=16), + discriminator_block(layers * 4, layers * 2, kernel_size=3, dilation=8), + discriminator_block(layers * 2, layers, kernel_size=3, dilation=1), # Final Convolution discriminator_block(layers, 1, kernel_size=3, stride=1), ) diff --git a/generator.py b/generator.py index 2446275..03fa279 100644 --- a/generator.py +++ b/generator.py @@ -10,7 +10,7 @@ def conv_block(in_channels, out_channels, kernel_size=3, dilation=1): class SISUGenerator(nn.Module): def __init__(self): super(SISUGenerator, self).__init__() - layer = 32 # Increased base layer count + layer = 4 # Increased base layer count self.conv1 = nn.Sequential( nn.Conv1d(1, layer, kernel_size=7, padding=3), nn.BatchNorm1d(layer), @@ -18,9 +18,9 @@ class SISUGenerator(nn.Module): ) self.conv_blocks = nn.Sequential( conv_block(layer, layer, kernel_size=3, dilation=1), # Local details - conv_block(layer, layer*4, kernel_size=5, dilation=2), # Local Context - 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, layer*2, kernel_size=5, dilation=2), # Local Context + conv_block(layer*2, layer*2, kernel_size=3, dilation=16), # Longer range dependencies + conv_block(layer*2, layer*2, kernel_size=5, dilation=8), # 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 ) diff --git a/training.py b/training.py index e114817..6ee7116 100644 --- a/training.py +++ b/training.py @@ -85,7 +85,7 @@ dataset = AudioDataset(dataset_dir) # ========= SINGLE ========= -train_data_loader = DataLoader(dataset, batch_size=1, shuffle=True) +train_data_loader = DataLoader(dataset, batch_size=16, shuffle=True) # Initialize models and move them to device generator = SISUGenerator() -- 2.43.0 From 741dcce7b45af6738c4b8a5df33cb6d3770fccc7 Mon Sep 17 00:00:00 2001 From: NikkeDoy Date: Sun, 23 Feb 2025 13:52:01 +0200 Subject: [PATCH 03/14] :alembic: | Increase discriminator size and implement mfcc_loss for generator. --- discriminator.py | 25 ++++++------- training.py | 96 ++++++++++++++++++++++++++---------------------- 2 files changed, 65 insertions(+), 56 deletions(-) diff --git a/discriminator.py b/discriminator.py index b1d82e1..d090372 100644 --- a/discriminator.py +++ b/discriminator.py @@ -6,8 +6,8 @@ def discriminator_block(in_channels, out_channels, kernel_size=3, stride=1, dila 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 + nn.LeakyReLU(0.2, inplace=True), + nn.BatchNorm1d(out_channels) ) class SISUDiscriminator(nn.Module): @@ -15,17 +15,16 @@ class SISUDiscriminator(nn.Module): super(SISUDiscriminator, self).__init__() layers = 4 # Increased base layer count self.model = nn.Sequential( - # Initial Convolution - discriminator_block(1, layers, kernel_size=7, stride=2, dilation=1), # Downsample - - # 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 * 4, kernel_size=5, dilation=4), - discriminator_block(layers * 4, layers * 4, kernel_size=5, dilation=16), - discriminator_block(layers * 4, layers * 2, kernel_size=3, dilation=8), - discriminator_block(layers * 2, layers, kernel_size=3, dilation=1), - # Final Convolution - discriminator_block(layers, 1, kernel_size=3, stride=1), + discriminator_block(1, layers, kernel_size=7, stride=2), # Initial downsampling + discriminator_block(layers, layers * 2, kernel_size=5, stride=2), # Downsampling + discriminator_block(layers * 2, layers * 4, kernel_size=5, dilation=2), # Increased dilation + discriminator_block(layers * 4, layers * 4, kernel_size=5, dilation=4), # Increased dilation + discriminator_block(layers * 4, layers * 8, kernel_size=5, dilation=8), # Deeper layer! + discriminator_block(layers * 8, layers * 8, kernel_size=5, dilation=1), # Deeper layer! + discriminator_block(layers * 8, layers * 4, kernel_size=3, dilation=2), # Reduced dilation + discriminator_block(layers * 4, layers * 2, kernel_size=3, dilation=1), + discriminator_block(layers * 2, layers, kernel_size=3, stride=1), # Final convolution + discriminator_block(layers, 1, kernel_size=3, stride=1) ) self.global_avg_pool = nn.AdaptiveAvgPool1d(1) diff --git a/training.py b/training.py index 6ee7116..3992829 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,8 +20,26 @@ 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) +import librosa + +def mfcc_loss(y_true, y_pred, sr): + # 1. Ensure sr is a NumPy scalar (not a Tensor) + if isinstance(sr, torch.Tensor): # Check if it's a Tensor + sr = sr.item() # Extract the value as a Python number + + # 2. Convert y_true and y_pred to NumPy arrays + y_true_np = y_true.cpu().detach().numpy()[0] # .cpu() is crucial! + y_pred_np = y_pred.cpu().detach().numpy()[0] + + + mfccs_true = librosa.feature.mfcc(y=y_true_np, sr=sr, n_mfcc=20) + mfccs_pred = librosa.feature.mfcc(y=y_pred_np, sr=sr, n_mfcc=20) + + # 3. Convert MFCCs back to PyTorch tensors and ensure correct device + mfccs_true = torch.tensor(mfccs_true, device=y_true.device, dtype=torch.float32) + mfccs_pred = torch.tensor(mfccs_pred, device=y_pred.device, dtype=torch.float32) + + return torch.mean((mfccs_true - mfccs_pred)**2) def discriminator_train(high_quality, low_quality, real_labels, fake_labels): optimizer_d.zero_grad() @@ -43,17 +63,23 @@ def discriminator_train(high_quality, low_quality, real_labels, fake_labels): return d_loss -def generator_train(low_quality, real_labels): +def generator_train(low_quality, high_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() + mfcc_l = mfcc_loss(high_quality[0], generator_output, high_quality[1]) + + discriminator_decision = discriminator(generator_output) + adversarial_loss = criterion_g(discriminator_decision, real_labels) + + combined_loss = adversarial_loss + 0.5 * mfcc_l + + combined_loss.backward() optimizer_g.step() - return generator_output + + return (generator_output, combined_loss, adversarial_loss, mfcc_l) # Init script argument parser parser = argparse.ArgumentParser(description="Training script") @@ -61,6 +87,7 @@ 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("--verbose", action="store_true", help="Increase output verbosity") args = parser.parse_args() @@ -68,6 +95,8 @@ args = parser.parse_args() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Using device: {device}") +debug = args.verbose + # Initialize dataset and dataloader dataset_dir = './dataset/good' dataset = AudioDataset(dataset_dir) @@ -85,7 +114,7 @@ dataset = AudioDataset(dataset_dir) # ========= SINGLE ========= -train_data_loader = DataLoader(dataset, batch_size=16, shuffle=True) +train_data_loader = DataLoader(dataset, batch_size=1, shuffle=True) # Initialize models and move them to device generator = SISUGenerator() @@ -111,32 +140,10 @@ optimizer_d = optim.Adam(discriminator.parameters(), lr=0.0001, betas=(0.5, 0.99 scheduler_g = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer_g, mode='min', factor=0.5, patience=5) scheduler_d = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer_d, mode='min', factor=0.5, patience=5) +models_dir = "models" +os.makedirs(models_dir, exist_ok=True) + 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) @@ -158,32 +165,35 @@ 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) # ========= GENERATOR ========= generator.train() - generator_output = generator_train(low_quality_sample, real_labels) + generator_output, combined_loss, adversarial_loss, mfcc_l = generator_train(low_quality_sample, high_quality_sample, real_labels) + + if debug: + print(d_loss, combined_loss, adversarial_loss, mfcc_l) + scheduler_d.step(d_loss) + scheduler_g.step(combined_loss) # ========= SAVE LATEST AUDIO ========= high_quality_audio = high_quality_clip low_quality_audio = low_quality_clip ai_enhanced_audio = (generator_output, high_quality_clip[1]) - #metric = snr(high_quality_audio[0].to(device), ai_enhanced_audio[0]) - #print(f"Generator metric {metric}!") - #scheduler_g.step(metric) - 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]) - torch.save(discriminator.state_dict(), f"models/current-epoch-discriminator.pt") - torch.save(generator.state_dict(), f"models/current-epoch-generator.pt") + torch.save(discriminator.state_dict(), f"{models_dir}/discriminator_epoch_{generator_epoch}.pt") + torch.save(generator.state_dict(), f"{models_dir}/generator_epoch_{generator_epoch}.pt") + torch.save(discriminator, f"{models_dir}/discriminator_epoch_{generator_epoch}_full.pt") + torch.save(generator, f"{models_dir}/generator_epoch_{generator_epoch}_full.pt") - 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() -- 2.43.0 From 8332b0df2daa97561a78fb95012dc0bc349ffb41 Mon Sep 17 00:00:00 2001 From: NikkeDoy Date: Wed, 26 Feb 2025 19:36:43 +0200 Subject: [PATCH 04/14] :sparkles: | Added ability to set epoch. --- training.py | 66 ++++++++++++++++++++++++++++++++++++++--------------- 1 file changed, 48 insertions(+), 18 deletions(-) diff --git a/training.py b/training.py index 3992829..710dd65 100644 --- a/training.py +++ b/training.py @@ -21,21 +21,45 @@ from generator import SISUGenerator from discriminator import SISUDiscriminator import librosa +import numpy as np def mfcc_loss(y_true, y_pred, sr): - # 1. Ensure sr is a NumPy scalar (not a Tensor) - if isinstance(sr, torch.Tensor): # Check if it's a Tensor - sr = sr.item() # Extract the value as a Python number + """Calculates MFCC loss between two audio signals. - # 2. Convert y_true and y_pred to NumPy arrays - y_true_np = y_true.cpu().detach().numpy()[0] # .cpu() is crucial! + Args: + y_true: Target audio signal (PyTorch tensor). + y_pred: Predicted audio signal (PyTorch tensor). + sr: Sample rate (NumPy scalar). + + Returns: + MFCC loss (PyTorch tensor). + """ + + # 1. Ensure sr is a NumPy scalar (not a Tensor) + if isinstance(sr, torch.Tensor): + sr = sr.item() + + # 2. Convert y_true and y_pred to NumPy arrays (and detach from graph) + y_true_np = y_true.cpu().detach().numpy()[0] # .cpu() and .detach() are crucial! y_pred_np = y_pred.cpu().detach().numpy()[0] + # 3. Dynamically calculate n_fft based on signal length + signal_length = min(y_true_np.shape[0], y_pred_np.shape[0]) # Use shortest signal length + n_fft = min(2048, 2**int(np.log2(signal_length))) # Power of 2, up to 2048 - mfccs_true = librosa.feature.mfcc(y=y_true_np, sr=sr, n_mfcc=20) - mfccs_pred = librosa.feature.mfcc(y=y_pred_np, sr=sr, n_mfcc=20) + # 4. Calculate MFCCs using adjusted n_fft + mfccs_true = librosa.feature.mfcc(y=y_true_np, sr=sr, n_fft=n_fft, n_mfcc=20) + mfccs_pred = librosa.feature.mfcc(y=y_pred_np, sr=sr, n_fft=n_fft, n_mfcc=20) - # 3. Convert MFCCs back to PyTorch tensors and ensure correct device + # 5. Truncate MFCCs to the same length (important!) + len_true = mfccs_true.shape[1] + len_pred = mfccs_pred.shape[1] + min_len = min(len_true, len_pred) + + mfccs_true = mfccs_true[:, :min_len] + mfccs_pred = mfccs_pred[:, :min_len] + + # 6. Convert MFCCs back to PyTorch tensors and ensure correct device mfccs_true = torch.tensor(mfccs_true, device=y_true.device, dtype=torch.float32) mfccs_pred = torch.tensor(mfccs_pred, device=y_pred.device, dtype=torch.float32) @@ -87,6 +111,7 @@ 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("--epoch", type=int, default=0, help="Current epoch for model versioning") parser.add_argument("--verbose", action="store_true", help="Increase output verbosity") args = parser.parse_args() @@ -120,6 +145,8 @@ train_data_loader = DataLoader(dataset, batch_size=1, shuffle=True) generator = SISUGenerator() discriminator = SISUDiscriminator() +epoch: int = args.epoch + if args.generator is not None: generator.load_state_dict(torch.load(args.generator, weights_only=True)) if args.discriminator is not None: @@ -153,7 +180,7 @@ 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]) @@ -181,16 +208,19 @@ def start_training(): low_quality_audio = low_quality_clip ai_enhanced_audio = (generator_output, high_quality_clip[1]) - 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]) + new_epoch = generator_epoch+epoch - torch.save(discriminator.state_dict(), f"{models_dir}/discriminator_epoch_{generator_epoch}.pt") - torch.save(generator.state_dict(), f"{models_dir}/generator_epoch_{generator_epoch}.pt") - torch.save(discriminator, f"{models_dir}/discriminator_epoch_{generator_epoch}_full.pt") - torch.save(generator, f"{models_dir}/generator_epoch_{generator_epoch}_full.pt") + if generator_epoch % 10 == 0: + print(f"Saved epoch {new_epoch}!") + torchaudio.save(f"./output/epoch-{new_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-{new_epoch}-audio-ai.wav", ai_enhanced_audio[0][0].cpu(), ai_enhanced_audio[1]) + torchaudio.save(f"./output/epoch-{new_epoch}-audio-orig.wav", high_quality_audio[0][0].cpu(), high_quality_audio[1]) + + if debug: + print(generator.state_dict().keys()) + print(discriminator.state_dict().keys()) + torch.save(discriminator.state_dict(), f"{models_dir}/discriminator_epoch_{new_epoch}.pt") + torch.save(generator.state_dict(), f"{models_dir}/generator_epoch_{new_epoch}.pt") torch.save(discriminator, "models/epoch-5000-discriminator.pt") torch.save(generator, "models/epoch-5000-generator.pt") -- 2.43.0 From 416500f7fc3d7fdf4e677d1d1acf375bc885e463 Mon Sep 17 00:00:00 2001 From: NikkeDoy Date: Wed, 26 Feb 2025 20:15:30 +0200 Subject: [PATCH 05/14] :heavy_minus_sign: | Removed/Updated dependencies. --- data.py | 8 +++- requirements.txt | 10 ++--- training.py | 95 +++++++++++++++++------------------------------- 3 files changed, 44 insertions(+), 69 deletions(-) diff --git a/data.py b/data.py index ac69730..2f05581 100644 --- a/data.py +++ b/data.py @@ -12,8 +12,9 @@ class AudioDataset(Dataset): #audio_sample_rates = [8000, 11025, 16000, 22050] audio_sample_rates = [11025] - 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): @@ -32,4 +33,7 @@ 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).to(self.device) + low_quality_audio = AudioUtils.stereo_tensor_to_mono(low_quality_audio).to(self.device) + + return (high_quality_audio, original_sample_rate), (low_quality_audio, mangled_sample_rate) diff --git a/requirements.txt b/requirements.txt index 5cb5df1..eacfc3b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,11 +4,11 @@ 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 +pytorch-triton-rocm==3.2.0+git4b3bb1f8 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 +torch==2.7.0.dev20250226+rocm6.3 +torchaudio==2.6.0.dev20250226+rocm6.3 tqdm==4.67.1 typing_extensions==4.12.2 diff --git a/training.py b/training.py index 710dd65..bf60c5c 100644 --- a/training.py +++ b/training.py @@ -20,49 +20,35 @@ from data import AudioDataset from generator import SISUGenerator from discriminator import SISUDiscriminator -import librosa -import numpy as np +import torchaudio.transforms as T -def mfcc_loss(y_true, y_pred, sr): - """Calculates MFCC loss between two audio signals. +# Init script argument parser +parser = argparse.ArgumentParser(description="Training script") +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("--verbose", action="store_true", help="Increase output verbosity") - Args: - y_true: Target audio signal (PyTorch tensor). - y_pred: Predicted audio signal (PyTorch tensor). - sr: Sample rate (NumPy scalar). +args = parser.parse_args() - Returns: - MFCC loss (PyTorch tensor). - """ +device = torch.device(args.device if torch.cuda.is_available() else "cpu") +print(f"Using device: {device}") - # 1. Ensure sr is a NumPy scalar (not a Tensor) - if isinstance(sr, torch.Tensor): - sr = sr.item() - - # 2. Convert y_true and y_pred to NumPy arrays (and detach from graph) - y_true_np = y_true.cpu().detach().numpy()[0] # .cpu() and .detach() are crucial! - y_pred_np = y_pred.cpu().detach().numpy()[0] - - # 3. Dynamically calculate n_fft based on signal length - signal_length = min(y_true_np.shape[0], y_pred_np.shape[0]) # Use shortest signal length - n_fft = min(2048, 2**int(np.log2(signal_length))) # Power of 2, up to 2048 - - # 4. Calculate MFCCs using adjusted n_fft - mfccs_true = librosa.feature.mfcc(y=y_true_np, sr=sr, n_fft=n_fft, n_mfcc=20) - mfccs_pred = librosa.feature.mfcc(y=y_pred_np, sr=sr, n_fft=n_fft, n_mfcc=20) - - # 5. Truncate MFCCs to the same length (important!) - len_true = mfccs_true.shape[1] - len_pred = mfccs_pred.shape[1] - min_len = min(len_true, len_pred) - - mfccs_true = mfccs_true[:, :min_len] - mfccs_pred = mfccs_pred[:, :min_len] - - # 6. Convert MFCCs back to PyTorch tensors and ensure correct device - mfccs_true = torch.tensor(mfccs_true, device=y_true.device, dtype=torch.float32) - mfccs_pred = torch.tensor(mfccs_pred, device=y_pred.device, dtype=torch.float32) +mfcc_transform = T.MFCC( + sample_rate=16000, # Adjust to your sample rate + n_mfcc=20, + melkwargs={'n_fft': 2048, 'hop_length': 512} # adjust n_fft and hop_length to your needs. +).to(device) +def gpu_mfcc_loss(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] return torch.mean((mfccs_true - mfccs_pred)**2) def discriminator_train(high_quality, low_quality, real_labels, fake_labels): @@ -93,7 +79,7 @@ def generator_train(low_quality, high_quality, real_labels): # Forward pass for fake samples (from generator output) generator_output = generator(low_quality[0]) - mfcc_l = mfcc_loss(high_quality[0], generator_output, high_quality[1]) + mfcc_l = gpu_mfcc_loss(high_quality[0], generator_output) discriminator_decision = discriminator(generator_output) adversarial_loss = criterion_g(discriminator_decision, real_labels) @@ -105,26 +91,11 @@ def generator_train(low_quality, high_quality, real_labels): return (generator_output, combined_loss, adversarial_loss, mfcc_l) -# Init script argument parser -parser = argparse.ArgumentParser(description="Training script") -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("--epoch", type=int, default=0, help="Current epoch for model versioning") -parser.add_argument("--verbose", action="store_true", help="Increase output verbosity") - -args = parser.parse_args() - -# Check for CUDA availability -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -print(f"Using device: {device}") - debug = args.verbose # Initialize dataset and dataloader dataset_dir = './dataset/good' -dataset = AudioDataset(dataset_dir) +dataset = AudioDataset(dataset_dir, device) # ========= MULTIPLE ========= @@ -147,14 +118,14 @@ discriminator = SISUDiscriminator() epoch: int = args.epoch -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)) - generator = generator.to(device) discriminator = discriminator.to(device) +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)) + # Loss criterion_g = nn.MSELoss() criterion_d = nn.BCELoss() @@ -182,8 +153,8 @@ def start_training(): # ========= TRAINING ========= 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) -- 2.43.0 From 7e1c7e935a0e43a9696f424f9757c6ad344dc2c9 Mon Sep 17 00:00:00 2001 From: NikkeDoy Date: Sat, 15 Mar 2025 18:01:19 +0200 Subject: [PATCH 06/14] :albemic: | Experimenting with other model layouts. --- data.py | 26 +++++++++++++++++----- discriminator.py | 57 ++++++++++++++++++++++++++++++++---------------- generator.py | 44 +++++++++++++++++++++++++------------ training.py | 27 +++++++---------------- 4 files changed, 96 insertions(+), 58 deletions(-) diff --git a/data.py b/data.py index 2f05581..9ca5ee5 100644 --- a/data.py +++ b/data.py @@ -4,23 +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 = 88200 # Define your desired maximum length here 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) @@ -33,7 +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) - high_quality_audio = AudioUtils.stereo_tensor_to_mono(high_quality_audio).to(self.device) - low_quality_audio = AudioUtils.stereo_tensor_to_mono(low_quality_audio).to(self.device) + 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 d090372..b1ec6eb 100644 --- a/discriminator.py +++ b/discriminator.py @@ -2,35 +2,54 @@ 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): padding = (kernel_size // 2) * dilation + conv_layer = nn.Conv1d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, dilation=dilation, padding=padding) + if spectral_norm: + conv_layer = utils.spectral_norm(conv_layer) return nn.Sequential( - utils.spectral_norm(nn.Conv1d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, dilation=dilation, padding=padding)), + conv_layer, nn.LeakyReLU(0.2, inplace=True), nn.BatchNorm1d(out_channels) ) -class SISUDiscriminator(nn.Module): - def __init__(self): - super(SISUDiscriminator, self).__init__() - layers = 4 # Increased base layer count - self.model = nn.Sequential( - discriminator_block(1, layers, kernel_size=7, stride=2), # Initial downsampling - discriminator_block(layers, layers * 2, kernel_size=5, stride=2), # Downsampling - discriminator_block(layers * 2, layers * 4, kernel_size=5, dilation=2), # Increased dilation - discriminator_block(layers * 4, layers * 4, kernel_size=5, dilation=4), # Increased dilation - discriminator_block(layers * 4, layers * 8, kernel_size=5, dilation=8), # Deeper layer! - discriminator_block(layers * 8, layers * 8, kernel_size=5, dilation=1), # Deeper layer! - discriminator_block(layers * 8, layers * 4, kernel_size=3, dilation=2), # Reduced dilation - discriminator_block(layers * 4, layers * 2, kernel_size=3, dilation=1), - discriminator_block(layers * 2, layers, kernel_size=3, stride=1), # Final convolution - discriminator_block(layers, 1, kernel_size=3, stride=1) +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(), + nn.Conv1d(channels // 4, channels, kernel_size=1), + nn.Sigmoid() ) - 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 + attention_weights = self.attention(x) + return x * attention_weights + +class SISUDiscriminator(nn.Module): + def __init__(self, layers=4): #Increased base layer count + super(SISUDiscriminator, self).__init__() + self.model = nn.Sequential( + discriminator_block(1, layers, kernel_size=7, stride=4), #Aggressive downsampling + discriminator_block(layers, layers * 2, kernel_size=5, stride=2), + discriminator_block(layers * 2, layers * 4, kernel_size=5, dilation=2), + discriminator_block(layers * 4, layers * 8, kernel_size=5, dilation=4), + AttentionBlock(layers * 8), #Added attention + discriminator_block(layers * 8, layers * 16, kernel_size=5, dilation=8), + discriminator_block(layers * 16, layers * 16, kernel_size=3, dilation=1), + discriminator_block(layers * 16, layers * 8, kernel_size=3, dilation=2), + discriminator_block(layers * 8, layers * 4, kernel_size=3, dilation=1), + discriminator_block(layers * 4, layers * 2, kernel_size=3, stride=1), + discriminator_block(layers * 2, layers, kernel_size=3, stride=1), + discriminator_block(layers, 1, kernel_size=3, stride=1, spectral_norm=False) #last layer no spectral norm. + ) + self.global_avg_pool = nn.AdaptiveAvgPool1d(1) + self.sigmoid = nn.Sigmoid() + + def forward(self, x): x = self.model(x) x = self.global_avg_pool(x) x = x.view(-1, 1) + x = self.sigmoid(x) return x diff --git a/generator.py b/generator.py index 03fa279..950530a 100644 --- a/generator.py +++ b/generator.py @@ -7,30 +7,46 @@ def conv_block(in_channels, out_channels, kernel_size=3, dilation=1): nn.PReLU() ) +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(), + 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.conv_layers(x) + x = self.attention(x) + return x + residual + class SISUGenerator(nn.Module): - def __init__(self): + def __init__(self, layer=4, num_rirb=4): #increased base layer and rirb amounts super(SISUGenerator, self).__init__() - layer = 4 # 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=16), # Longer range dependencies - conv_block(layer*2, layer*2, kernel_size=5, dilation=8), # 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), - ) + self.rir_blocks = nn.Sequential(*[ResidualInResidualBlock(layer) for _ in range(num_rirb)]) + self.final_layer = nn.Conv1d(layer, 1, kernel_size=3, padding=1) def forward(self, x): residual = x x = self.conv1(x) - x = self.conv_blocks(x) + x = self.rir_blocks(x) x = self.final_layer(x) return x + residual diff --git a/training.py b/training.py index bf60c5c..50743be 100644 --- a/training.py +++ b/training.py @@ -38,7 +38,7 @@ device = torch.device(args.device if torch.cuda.is_available() else "cpu") print(f"Using device: {device}") mfcc_transform = T.MFCC( - sample_rate=16000, # Adjust to your sample rate + sample_rate=44100, # Adjust to your sample rate n_mfcc=20, melkwargs={'n_fft': 2048, 'hop_length': 512} # adjust n_fft and hop_length to your needs. ).to(device) @@ -97,20 +97,9 @@ debug = args.verbose dataset_dir = './dataset/good' dataset = AudioDataset(dataset_dir, device) -# ========= 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) - # ========= SINGLE ========= -train_data_loader = DataLoader(dataset, batch_size=1, shuffle=True) +train_data_loader = DataLoader(dataset, batch_size=128, shuffle=True) # Initialize models and move them to device generator = SISUGenerator() @@ -175,17 +164,17 @@ def start_training(): scheduler_g.step(combined_loss) # ========= 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]) new_epoch = generator_epoch+epoch if generator_epoch % 10 == 0: print(f"Saved epoch {new_epoch}!") - torchaudio.save(f"./output/epoch-{new_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-{new_epoch}-audio-ai.wav", ai_enhanced_audio[0][0].cpu(), ai_enhanced_audio[1]) - torchaudio.save(f"./output/epoch-{new_epoch}-audio-orig.wav", high_quality_audio[0][0].cpu(), high_quality_audio[1]) + torchaudio.save(f"./output/epoch-{new_epoch}-audio-crap.wav", low_quality_audio[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-{new_epoch}-audio-ai.wav", ai_enhanced_audio[0].cpu(), ai_enhanced_audio[1]) + torchaudio.save(f"./output/epoch-{new_epoch}-audio-orig.wav", high_quality_audio[0].cpu(), high_quality_audio[1]) if debug: print(generator.state_dict().keys()) -- 2.43.0 From 54338e55a99b3ab5620b36d7b594ca5aa2b29039 Mon Sep 17 00:00:00 2001 From: NikkeDoy Date: Tue, 25 Mar 2025 19:50:51 +0200 Subject: [PATCH 07/14] :albemic: | Tests. --- data.py | 2 +- discriminator.py | 23 +++++++++++++---------- generator.py | 2 +- training.py | 11 ++++++----- 4 files changed, 21 insertions(+), 17 deletions(-) diff --git a/data.py b/data.py index 9ca5ee5..bc7574f 100644 --- a/data.py +++ b/data.py @@ -9,7 +9,7 @@ import AudioUtils class AudioDataset(Dataset): audio_sample_rates = [11025] - MAX_LENGTH = 88200 # Define your desired maximum length here + MAX_LENGTH = 44100 # Define your desired maximum length here 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')] diff --git a/discriminator.py b/discriminator.py index b1ec6eb..1608199 100644 --- a/discriminator.py +++ b/discriminator.py @@ -28,19 +28,22 @@ class AttentionBlock(nn.Module): return x * attention_weights class SISUDiscriminator(nn.Module): - def __init__(self, layers=4): #Increased base layer count + def __init__(self, layers=64): #Increased base layer count super(SISUDiscriminator, self).__init__() self.model = nn.Sequential( - discriminator_block(1, layers, kernel_size=7, stride=4), #Aggressive downsampling + discriminator_block(1, layers, kernel_size=3, stride=1), #Aggressive downsampling discriminator_block(layers, layers * 2, kernel_size=5, stride=2), - discriminator_block(layers * 2, layers * 4, kernel_size=5, dilation=2), - discriminator_block(layers * 4, layers * 8, kernel_size=5, dilation=4), - AttentionBlock(layers * 8), #Added attention - discriminator_block(layers * 8, layers * 16, kernel_size=5, dilation=8), - discriminator_block(layers * 16, layers * 16, kernel_size=3, dilation=1), - discriminator_block(layers * 16, layers * 8, kernel_size=3, dilation=2), - discriminator_block(layers * 8, layers * 4, kernel_size=3, dilation=1), - discriminator_block(layers * 4, layers * 2, kernel_size=3, stride=1), + discriminator_block(layers * 2, layers * 4, kernel_size=5, dilation=4), + + #AttentionBlock(layers * 4), #Added attention + + #discriminator_block(layers * 4, layers * 8, kernel_size=5, dilation=4), + #AttentionBlock(layers * 8), #Added attention + #discriminator_block(layers * 8, layers * 16, kernel_size=5, dilation=8), + #discriminator_block(layers * 16, layers * 16, kernel_size=3, dilation=1), + #discriminator_block(layers * 16, layers * 8, kernel_size=3, dilation=2), + #discriminator_block(layers * 8, layers * 4, kernel_size=3, dilation=1), + discriminator_block(layers * 4, layers * 2, kernel_size=5, stride=2), discriminator_block(layers * 2, layers, kernel_size=3, stride=1), discriminator_block(layers, 1, kernel_size=3, stride=1, spectral_norm=False) #last layer no spectral norm. ) diff --git a/generator.py b/generator.py index 950530a..04ac5b4 100644 --- a/generator.py +++ b/generator.py @@ -34,7 +34,7 @@ class ResidualInResidualBlock(nn.Module): return x + residual class SISUGenerator(nn.Module): - def __init__(self, layer=4, num_rirb=4): #increased base layer and rirb amounts + def __init__(self, layer=64, num_rirb=4): #increased base layer and rirb amounts super(SISUGenerator, self).__init__() self.conv1 = nn.Sequential( nn.Conv1d(1, layer, kernel_size=7, padding=3), diff --git a/training.py b/training.py index 50743be..63fc5b8 100644 --- a/training.py +++ b/training.py @@ -38,9 +38,9 @@ device = torch.device(args.device if torch.cuda.is_available() else "cpu") print(f"Using device: {device}") mfcc_transform = T.MFCC( - sample_rate=44100, # Adjust to your sample rate + sample_rate=44100, n_mfcc=20, - melkwargs={'n_fft': 2048, 'hop_length': 512} # adjust n_fft and hop_length to your needs. + melkwargs={'n_fft': 2048, 'hop_length': 256} ).to(device) def gpu_mfcc_loss(y_true, y_pred): @@ -49,7 +49,8 @@ def gpu_mfcc_loss(y_true, 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] - return torch.mean((mfccs_true - mfccs_pred)**2) + loss = torch.mean((mfccs_true - mfccs_pred)**2) + return loss def discriminator_train(high_quality, low_quality, real_labels, fake_labels): optimizer_d.zero_grad() @@ -99,7 +100,7 @@ dataset = AudioDataset(dataset_dir, device) # ========= SINGLE ========= -train_data_loader = DataLoader(dataset, batch_size=128, shuffle=True) +train_data_loader = DataLoader(dataset, batch_size=16, shuffle=True) # Initialize models and move them to device generator = SISUGenerator() @@ -161,7 +162,7 @@ def start_training(): if debug: print(d_loss, combined_loss, adversarial_loss, mfcc_l) scheduler_d.step(d_loss) - scheduler_g.step(combined_loss) + #scheduler_g.step(combined_loss) # ========= SAVE LATEST AUDIO ========= high_quality_audio = (high_quality_clip[0][0], high_quality_clip[1][0]) -- 2.43.0 From f928d8c2cf3540c1f22f13f70715863220766eb9 Mon Sep 17 00:00:00 2001 From: NikkeDoy Date: Tue, 25 Mar 2025 21:51:29 +0200 Subject: [PATCH 08/14] :albemic: | More tests. --- discriminator.py | 2 +- generator.py | 2 +- training.py | 22 ++++++++++++---------- 3 files changed, 14 insertions(+), 12 deletions(-) diff --git a/discriminator.py b/discriminator.py index 1608199..58b95f0 100644 --- a/discriminator.py +++ b/discriminator.py @@ -28,7 +28,7 @@ class AttentionBlock(nn.Module): return x * attention_weights class SISUDiscriminator(nn.Module): - def __init__(self, layers=64): #Increased base layer count + def __init__(self, layers=4): #Increased base layer count super(SISUDiscriminator, self).__init__() self.model = nn.Sequential( discriminator_block(1, layers, kernel_size=3, stride=1), #Aggressive downsampling diff --git a/generator.py b/generator.py index 04ac5b4..950530a 100644 --- a/generator.py +++ b/generator.py @@ -34,7 +34,7 @@ class ResidualInResidualBlock(nn.Module): return x + residual class SISUGenerator(nn.Module): - def __init__(self, layer=64, num_rirb=4): #increased base layer and rirb amounts + def __init__(self, layer=4, num_rirb=4): #increased base layer and rirb amounts super(SISUGenerator, self).__init__() self.conv1 = nn.Sequential( nn.Conv1d(1, layer, kernel_size=7, padding=3), diff --git a/training.py b/training.py index 63fc5b8..814fcda 100644 --- a/training.py +++ b/training.py @@ -30,7 +30,7 @@ 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("--verbose", action="store_true", help="Increase output verbosity") +parser.add_argument("--debug", action="store_true", help="Print debug logs") args = parser.parse_args() @@ -80,19 +80,20 @@ def generator_train(low_quality, high_quality, real_labels): # Forward pass for fake samples (from generator output) generator_output = generator(low_quality[0]) - mfcc_l = gpu_mfcc_loss(high_quality[0], generator_output) + #mfcc_l = gpu_mfcc_loss(high_quality[0], generator_output) discriminator_decision = discriminator(generator_output) adversarial_loss = criterion_g(discriminator_decision, real_labels) - combined_loss = adversarial_loss + 0.5 * mfcc_l + #combined_loss = adversarial_loss + 0.5 * mfcc_l - combined_loss.backward() + adversarial_loss.backward() optimizer_g.step() - return (generator_output, combined_loss, adversarial_loss, mfcc_l) + #return (generator_output, combined_loss, adversarial_loss, mfcc_l) + return (generator_output, adversarial_loss) -debug = args.verbose +debug = args.debug # Initialize dataset and dataloader dataset_dir = './dataset/good' @@ -100,7 +101,7 @@ dataset = AudioDataset(dataset_dir, device) # ========= SINGLE ========= -train_data_loader = DataLoader(dataset, batch_size=16, shuffle=True) +train_data_loader = DataLoader(dataset, batch_size=256, shuffle=True) # Initialize models and move them to device generator = SISUGenerator() @@ -157,12 +158,13 @@ def start_training(): # ========= GENERATOR ========= generator.train() - generator_output, combined_loss, adversarial_loss, mfcc_l = generator_train(low_quality_sample, high_quality_sample, real_labels) + #generator_output, combined_loss, adversarial_loss, mfcc_l = generator_train(low_quality_sample, high_quality_sample, real_labels) + generator_output, adversarial_loss = generator_train(low_quality_sample, high_quality_sample, real_labels) if debug: - print(d_loss, combined_loss, adversarial_loss, mfcc_l) + print(d_loss, adversarial_loss) scheduler_d.step(d_loss) - #scheduler_g.step(combined_loss) + scheduler_g.step(adversarial_loss) # ========= SAVE LATEST AUDIO ========= high_quality_audio = (high_quality_clip[0][0], high_quality_clip[1][0]) -- 2.43.0 From 9394bc6c5a1f25733a0bc5007148d12dbf11a218 Mon Sep 17 00:00:00 2001 From: NikkeDoy Date: Sun, 6 Apr 2025 00:05:43 +0300 Subject: [PATCH 09/14] :albemic: | Fat architecture. Hopefully better results. --- README.md | 1 + discriminator.py | 59 ++++++++++++++++++++++++++---------------------- generator.py | 48 ++++++++++++++++++++++++++++----------- requirements.txt | 4 +--- training.py | 2 +- 5 files changed, 70 insertions(+), 44 deletions(-) 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/discriminator.py b/discriminator.py index 58b95f0..777abf2 100644 --- a/discriminator.py +++ b/discriminator.py @@ -2,23 +2,34 @@ 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, spectral_norm=True): +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 - conv_layer = nn.Conv1d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, dilation=dilation, padding=padding) + conv_layer = nn.Conv1d( + in_channels, + out_channels, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + padding=padding + ) + if spectral_norm: conv_layer = utils.spectral_norm(conv_layer) - return nn.Sequential( - conv_layer, - nn.LeakyReLU(0.2, inplace=True), - nn.BatchNorm1d(out_channels) - ) + + 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(), + nn.ReLU(inplace=True), nn.Conv1d(channels // 4, channels, kernel_size=1), nn.Sigmoid() ) @@ -28,31 +39,25 @@ class AttentionBlock(nn.Module): return x * attention_weights class SISUDiscriminator(nn.Module): - def __init__(self, layers=4): #Increased base layer count + def __init__(self, base_channels=64): super(SISUDiscriminator, self).__init__() + layers = base_channels self.model = nn.Sequential( - discriminator_block(1, layers, kernel_size=3, stride=1), #Aggressive downsampling - discriminator_block(layers, layers * 2, kernel_size=5, stride=2), - discriminator_block(layers * 2, layers * 4, kernel_size=5, dilation=4), - - #AttentionBlock(layers * 4), #Added attention - - #discriminator_block(layers * 4, layers * 8, kernel_size=5, dilation=4), - #AttentionBlock(layers * 8), #Added attention - #discriminator_block(layers * 8, layers * 16, kernel_size=5, dilation=8), - #discriminator_block(layers * 16, layers * 16, kernel_size=3, dilation=1), - #discriminator_block(layers * 16, layers * 8, kernel_size=3, dilation=2), - #discriminator_block(layers * 8, layers * 4, kernel_size=3, dilation=1), - discriminator_block(layers * 4, layers * 2, kernel_size=5, stride=2), - discriminator_block(layers * 2, layers, kernel_size=3, stride=1), - discriminator_block(layers, 1, kernel_size=3, stride=1, spectral_norm=False) #last layer no spectral norm. + 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) - self.sigmoid = nn.Sigmoid() def forward(self, x): x = self.model(x) x = self.global_avg_pool(x) - x = x.view(-1, 1) - x = self.sigmoid(x) + x = x.view(x.size(0), -1) return x diff --git a/generator.py b/generator.py index 950530a..cd4d48c 100644 --- a/generator.py +++ b/generator.py @@ -1,18 +1,28 @@ +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 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(), + nn.ReLU(inplace=True), nn.Conv1d(channels // 4, channels, kernel_size=1), nn.Sigmoid() ) @@ -24,7 +34,11 @@ class AttentionBlock(nn.Module): 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.conv_layers = nn.Sequential( + *[conv_block(channels, channels) for _ in range(num_convs)] + ) + self.attention = AttentionBlock(channels) def forward(self, x): @@ -34,19 +48,27 @@ class ResidualInResidualBlock(nn.Module): return x + residual class SISUGenerator(nn.Module): - def __init__(self, layer=4, num_rirb=4): #increased base layer and rirb amounts + def __init__(self, channels=64, num_rirb=8, alpha=1.0): super(SISUGenerator, self).__init__() + self.alpha = alpha + self.conv1 = nn.Sequential( - nn.Conv1d(1, layer, kernel_size=7, padding=3), - nn.BatchNorm1d(layer), + nn.Conv1d(1, channels, kernel_size=7, padding=3), + nn.InstanceNorm1d(channels), nn.PReLU(), ) - self.rir_blocks = nn.Sequential(*[ResidualInResidualBlock(layer) for _ in range(num_rirb)]) - self.final_layer = nn.Conv1d(layer, 1, kernel_size=3, padding=1) + + 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 = x + residual_input = x x = self.conv1(x) - x = self.rir_blocks(x) - x = self.final_layer(x) - return x + residual + 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 eacfc3b..21f6bef 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,10 +5,8 @@ MarkupSafe==2.1.5 mpmath==1.3.0 networkx==3.4.2 numpy==2.2.3 -pytorch-triton-rocm==3.2.0+git4b3bb1f8 +pillow==11.0.0 setuptools==70.2.0 sympy==1.13.3 -torch==2.7.0.dev20250226+rocm6.3 -torchaudio==2.6.0.dev20250226+rocm6.3 tqdm==4.67.1 typing_extensions==4.12.2 diff --git a/training.py b/training.py index 814fcda..380f738 100644 --- a/training.py +++ b/training.py @@ -101,7 +101,7 @@ dataset = AudioDataset(dataset_dir, device) # ========= SINGLE ========= -train_data_loader = DataLoader(dataset, batch_size=256, shuffle=True) +train_data_loader = DataLoader(dataset, batch_size=8, shuffle=True) # Initialize models and move them to device generator = SISUGenerator() -- 2.43.0 From fbcd5803b84d660dd6f8b4a0cab04e2929134d0f Mon Sep 17 00:00:00 2001 From: NikkeDoy Date: Mon, 7 Apr 2025 02:14:06 +0300 Subject: [PATCH 10/14] :bug: | Fixed training on CPU and NVIDIA hardware. --- training.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/training.py b/training.py index 380f738..c050b9c 100644 --- a/training.py +++ b/training.py @@ -119,7 +119,7 @@ if args.discriminator is not None: # Loss criterion_g = nn.MSELoss() -criterion_d = nn.BCELoss() +criterion_d = nn.BCEWithLogitsLoss() # Optimizers optimizer_g = optim.Adam(generator.parameters(), lr=0.0001, betas=(0.5, 0.999)) -- 2.43.0 From 3936b6c1600a62b18a7e38371707c47f4dfcfaf4 Mon Sep 17 00:00:00 2001 From: nsiltala <144348410+nsiltala@users.noreply.github.com> Date: Mon, 7 Apr 2025 14:49:07 +0300 Subject: [PATCH 11/14] :bug: | Fixed NVIDIA training... again. --- training.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/training.py b/training.py index c050b9c..47982bf 100644 --- a/training.py +++ b/training.py @@ -101,7 +101,7 @@ dataset = AudioDataset(dataset_dir, device) # ========= SINGLE ========= -train_data_loader = DataLoader(dataset, batch_size=8, shuffle=True) +train_data_loader = DataLoader(dataset, batch_size=12, shuffle=True) # Initialize models and move them to device generator = SISUGenerator() @@ -118,7 +118,7 @@ if args.discriminator is not None: discriminator.load_state_dict(torch.load(args.discriminator, map_location=device, weights_only=True)) # Loss -criterion_g = nn.MSELoss() +criterion_g = nn.BCEWithLogitsLoss() criterion_d = nn.BCEWithLogitsLoss() # Optimizers @@ -163,8 +163,8 @@ def start_training(): if debug: print(d_loss, adversarial_loss) - scheduler_d.step(d_loss) - scheduler_g.step(adversarial_loss) + scheduler_d.step(d_loss.detach()) + scheduler_g.step(adversarial_loss.detach()) # ========= SAVE LATEST AUDIO ========= high_quality_audio = (high_quality_clip[0][0], high_quality_clip[1][0]) @@ -175,9 +175,9 @@ def start_training(): if generator_epoch % 10 == 0: print(f"Saved epoch {new_epoch}!") - torchaudio.save(f"./output/epoch-{new_epoch}-audio-crap.wav", low_quality_audio[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-{new_epoch}-audio-ai.wav", ai_enhanced_audio[0].cpu(), ai_enhanced_audio[1]) - torchaudio.save(f"./output/epoch-{new_epoch}-audio-orig.wav", high_quality_audio[0].cpu(), high_quality_audio[1]) + torchaudio.save(f"./output/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"./output/epoch-{new_epoch}-audio-ai.wav", ai_enhanced_audio[0].cpu().detach(), ai_enhanced_audio[1]) + torchaudio.save(f"./output/epoch-{new_epoch}-audio-orig.wav", high_quality_audio[0].cpu().detach(), high_quality_audio[1]) if debug: print(generator.state_dict().keys()) -- 2.43.0 From b6d16e4f11582078b40af2a05d60d37487ee0090 Mon Sep 17 00:00:00 2001 From: NikkeDoy Date: Mon, 14 Apr 2025 17:51:34 +0300 Subject: [PATCH 12/14] :recycle: | Restructured procject code. --- file_utils.py | 28 ++++++++++ training.py | 127 ++++++++++++++++++++-------------------------- training_utils.py | 55 ++++++++++++++++++++ 3 files changed, 137 insertions(+), 73 deletions(-) create mode 100644 file_utils.py create mode 100644 training_utils.py 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/training.py b/training.py index 47982bf..17843e0 100644 --- a/training.py +++ b/training.py @@ -20,6 +20,9 @@ from data import AudioDataset from generator import SISUGenerator from discriminator import SISUDiscriminator +from training_utils import discriminator_train, generator_train +import file_utils as Data + import torchaudio.transforms as T # Init script argument parser @@ -31,92 +34,55 @@ parser.add_argument("--discriminator", type=str, default=None, 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", type=bool, default=False, help="Continue training using temp_generator and temp_discriminator models") args = parser.parse_args() device = torch.device(args.device if torch.cuda.is_available() else "cpu") print(f"Using device: {device}") -mfcc_transform = T.MFCC( - sample_rate=44100, - n_mfcc=20, - melkwargs={'n_fft': 2048, 'hop_length': 256} -).to(device) - -def gpu_mfcc_loss(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 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, high_quality, real_labels): - optimizer_g.zero_grad() - - # Forward pass for fake samples (from generator output) - generator_output = generator(low_quality[0]) - - #mfcc_l = gpu_mfcc_loss(high_quality[0], generator_output) - - discriminator_decision = discriminator(generator_output) - adversarial_loss = criterion_g(discriminator_decision, real_labels) - - #combined_loss = adversarial_loss + 0.5 * mfcc_l - - adversarial_loss.backward() - optimizer_g.step() - - #return (generator_output, combined_loss, adversarial_loss, mfcc_l) - return (generator_output, adversarial_loss) +# mfcc_transform = T.MFCC( +# sample_rate=44100, +# n_mfcc=20, +# melkwargs={'n_fft': 2048, 'hop_length': 256} +# ).to(device) debug = args.debug # Initialize dataset and dataloader dataset_dir = './dataset/good' 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=12, shuffle=True) -# Initialize models and move them to device + +# ========= MODELS ========= + generator = SISUGenerator() discriminator = SISUDiscriminator() 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) -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)) - # Loss criterion_g = nn.BCEWithLogitsLoss() criterion_d = nn.BCEWithLogitsLoss() @@ -129,9 +95,6 @@ optimizer_d = optim.Adam(discriminator.parameters(), lr=0.0001, betas=(0.5, 0.99 scheduler_g = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer_g, mode='min', factor=0.5, patience=5) scheduler_d = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer_d, mode='min', factor=0.5, patience=5) -models_dir = "models" -os.makedirs(models_dir, exist_ok=True) - def start_training(): generator_epochs = 5000 for generator_epoch in range(generator_epochs): @@ -154,12 +117,28 @@ def start_training(): # ========= DISCRIMINATOR ========= discriminator.train() - d_loss = 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, combined_loss, adversarial_loss, mfcc_l = generator_train(low_quality_sample, high_quality_sample, real_labels) - generator_output, adversarial_loss = generator_train(low_quality_sample, high_quality_sample, real_labels) + generator_output, adversarial_loss = generator_train( + low_quality_sample, + high_quality_sample, + real_labels, + generator, + discriminator, + criterion_g, + optimizer_g + ) if debug: print(d_loss, adversarial_loss) @@ -173,17 +152,19 @@ def start_training(): new_epoch = generator_epoch+epoch - if generator_epoch % 10 == 0: + if generator_epoch % 25 == 0: print(f"Saved epoch {new_epoch}!") - torchaudio.save(f"./output/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"./output/epoch-{new_epoch}-audio-ai.wav", ai_enhanced_audio[0].cpu().detach(), ai_enhanced_audio[1]) - torchaudio.save(f"./output/epoch-{new_epoch}-audio-orig.wav", high_quality_audio[0].cpu().detach(), high_quality_audio[1]) + 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]) if debug: print(generator.state_dict().keys()) print(discriminator.state_dict().keys()) - torch.save(discriminator.state_dict(), f"{models_dir}/discriminator_epoch_{new_epoch}.pt") - torch.save(generator.state_dict(), f"{models_dir}/generator_epoch_{new_epoch}.pt") + 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, "models/epoch-5000-discriminator.pt") torch.save(generator, "models/epoch-5000-generator.pt") diff --git a/training_utils.py b/training_utils.py new file mode 100644 index 0000000..a1d2c19 --- /dev/null +++ b/training_utils.py @@ -0,0 +1,55 @@ +import torch +import torch.nn as nn +import torch.optim as optim + +import torchaudio + +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 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) + + # 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(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.step() + + return d_loss + +def generator_train(low_quality, high_quality, real_labels, generator, discriminator, criterion, optimizer): + optimizer.zero_grad() + + # Forward pass for fake samples (from generator output) + generator_output = generator(low_quality[0]) + + #mfcc_l = gpu_mfcc_loss(high_quality[0], generator_output) + + discriminator_decision = discriminator(generator_output) + adversarial_loss = criterion(discriminator_decision, real_labels) + + #combined_loss = adversarial_loss + 0.5 * mfcc_l + + adversarial_loss.backward() + optimizer.step() + + #return (generator_output, combined_loss, adversarial_loss, mfcc_l) + return (generator_output, adversarial_loss) -- 2.43.0 From c04b072de6d156cb88f2c2652a33c04e80c68a31 Mon Sep 17 00:00:00 2001 From: NikkeDoy Date: Wed, 16 Apr 2025 17:08:13 +0300 Subject: [PATCH 13/14] :sparkles: | Added smarter ways that would've been needed from the begining. --- training.py | 33 ++++++++--- training_utils.py | 139 +++++++++++++++++++++++++++++++++++++++++----- 2 files changed, 148 insertions(+), 24 deletions(-) diff --git a/training.py b/training.py index 17843e0..db7cb86 100644 --- a/training.py +++ b/training.py @@ -41,11 +41,24 @@ args = parser.parse_args() device = torch.device(args.device if torch.cuda.is_available() else "cpu") print(f"Using device: {device}") -# mfcc_transform = T.MFCC( -# sample_rate=44100, -# n_mfcc=20, -# melkwargs={'n_fft': 2048, 'hop_length': 256} -# ).to(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) debug = args.debug @@ -130,18 +143,20 @@ def start_training(): # ========= GENERATOR ========= generator.train() - generator_output, adversarial_loss = generator_train( + generator_output, combined_loss, adversarial_loss, mel_l1_tensor = generator_train( low_quality_sample, high_quality_sample, real_labels, generator, discriminator, - criterion_g, - optimizer_g + criterion_d, + optimizer_g, + device, + mel_transform ) if debug: - print(d_loss, adversarial_loss) + print(combined_loss, adversarial_loss, mel_l1_tensor) scheduler_d.step(d_loss.detach()) scheduler_g.step(adversarial_loss.detach()) diff --git a/training_utils.py b/training_utils.py index a1d2c19..be402d9 100644 --- a/training_utils.py +++ b/training_utils.py @@ -3,16 +3,73 @@ 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] + + # L2 Loss (Mean Squared Error) + 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] + + # Log Magnitude L1 Loss + 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] + + # Calculate Frobenius norms and the loss + # Ensure norms are calculated over frequency and time dims ([..., freq, time]) + 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)) + + # Average loss over the batch + 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() @@ -21,35 +78,87 @@ def discriminator_train(high_quality, low_quality, real_labels, fake_labels, dis d_loss_real = criterion(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(discriminator_decision_from_fake, fake_labels) + with torch.no_grad(): # Detach generator output within no_grad context + generator_output = generator(low_quality[0]) + discriminator_decision_from_fake = discriminator(generator_output) # No need to detach again if inside no_grad + d_loss_fake = criterion(discriminator_decision_from_fake, fake_labels.expand_as(discriminator_decision_from_fake)) # 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 + # 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, criterion, optimizer): - optimizer.zero_grad() +def generator_train( + low_quality, + high_quality, + real_labels, + generator, + discriminator, + adv_criterion, # Criterion for adversarial loss (e.g., BCEWithLogitsLoss) + g_optimizer, + device, + # --- Pass necessary transforms and loss weights --- + mel_transform: T.MelSpectrogram, # Example: Pass Mel transform + # stft_transform: T.Spectrogram, # Pass STFT transform if using STFT losses + # mfcc_transform: T.MFCC, # Pass MFCC transform if using MFCC loss + lambda_adv: float = 1.0, # Weight for adversarial loss + lambda_mel_l1: float = 10.0, # Example: Weight for Mel L1 loss + # lambda_log_stft: float = 0.0, # Set weights > 0 for losses you want to use + # lambda_mfcc: float = 0.0 +): + g_optimizer.zero_grad() - # Forward pass for fake samples (from generator output) + # 1. Generate high-quality audio from low-quality input generator_output = generator(low_quality[0]) - #mfcc_l = gpu_mfcc_loss(high_quality[0], generator_output) - + # 2. Calculate Adversarial Loss (Generator tries to fool discriminator) discriminator_decision = discriminator(generator_output) - adversarial_loss = criterion(discriminator_decision, real_labels) + # Generator wants discriminator to output "real" labels for its fakes + adversarial_loss = adv_criterion(discriminator_decision, real_labels.expand_as(discriminator_decision)) - #combined_loss = adversarial_loss + 0.5 * mfcc_l + # 3. Calculate Reconstruction/Spectrogram Loss(es) + # --- Choose and calculate the losses you want to include --- + mel_l1 = 0.0 + # log_stft_l1 = 0.0 + # mfcc_l = 0.0 - adversarial_loss.backward() - optimizer.step() + # 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) - #return (generator_output, combined_loss, adversarial_loss, mfcc_l) - return (generator_output, adversarial_loss) + # # Calculate Log STFT L1 Loss if weight is positive + # if lambda_log_stft > 0: + # log_stft_l1 = log_stft_magnitude_loss(stft_transform, hq_audio, generator_output) + + # # Calculate MFCC Loss if weight is positive + # if lambda_mfcc > 0: + # mfcc_l = gpu_mfcc_loss(mfcc_transform, hq_audio, generator_output) + # --- End of Loss Calculation Choices --- + + + # 4. Combine Losses + # Make sure calculated losses are tensors even if weights are 0 initially + # (or handle appropriately in the sum) + 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) + + # 5. Backward Pass and Optimization + 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 -- 2.43.0 From d70c86c2578bdd94772069917bed040aa8046275 Mon Sep 17 00:00:00 2001 From: NikkeDoy Date: Sat, 26 Apr 2025 17:03:28 +0300 Subject: [PATCH 14/14] :sparkles: | Implemented MFCC and STFT. --- discriminator.py | 2 +- generator.py | 2 +- training.py | 22 +++++++++------ training_utils.py | 68 +++++++++++++++++------------------------------ 4 files changed, 40 insertions(+), 54 deletions(-) diff --git a/discriminator.py b/discriminator.py index 777abf2..dfd0126 100644 --- a/discriminator.py +++ b/discriminator.py @@ -39,7 +39,7 @@ class AttentionBlock(nn.Module): return x * attention_weights class SISUDiscriminator(nn.Module): - def __init__(self, base_channels=64): + def __init__(self, base_channels=16): super(SISUDiscriminator, self).__init__() layers = base_channels self.model = nn.Sequential( diff --git a/generator.py b/generator.py index cd4d48c..a53feb7 100644 --- a/generator.py +++ b/generator.py @@ -48,7 +48,7 @@ class ResidualInResidualBlock(nn.Module): return x + residual class SISUGenerator(nn.Module): - def __init__(self, channels=64, num_rirb=8, alpha=1.0): + def __init__(self, channels=16, num_rirb=4, alpha=1.0): super(SISUGenerator, self).__init__() self.alpha = alpha diff --git a/training.py b/training.py index db7cb86..01ea749 100644 --- a/training.py +++ b/training.py @@ -34,7 +34,7 @@ parser.add_argument("--discriminator", type=str, default=None, 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", type=bool, default=False, help="Continue training using temp_generator and temp_discriminator models") +parser.add_argument("--continue_training", action="store_true", help="Continue training using temp_generator and temp_discriminator models") args = parser.parse_args() @@ -60,6 +60,10 @@ mel_transform = T.MelSpectrogram( 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 @@ -72,7 +76,7 @@ os.makedirs(audio_output_dir, exist_ok=True) # ========= SINGLE ========= -train_data_loader = DataLoader(dataset, batch_size=12, shuffle=True) +train_data_loader = DataLoader(dataset, batch_size=64, shuffle=True) # ========= MODELS ========= @@ -143,7 +147,7 @@ def start_training(): # ========= GENERATOR ========= generator.train() - generator_output, combined_loss, adversarial_loss, mel_l1_tensor = generator_train( + 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, @@ -152,11 +156,13 @@ def start_training(): criterion_d, optimizer_g, device, - mel_transform + mel_transform, + stft_transform, + mfcc_transform ) if debug: - print(combined_loss, adversarial_loss, mel_l1_tensor) + 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()) @@ -173,9 +179,9 @@ def start_training(): 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]) - if debug: - print(generator.state_dict().keys()) - print(discriminator.state_dict().keys()) + #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}) diff --git a/training_utils.py b/training_utils.py index be402d9..6f26f58 100644 --- a/training_utils.py +++ b/training_utils.py @@ -37,7 +37,6 @@ def mel_spectrogram_l2_loss(mel_transform: T.MelSpectrogram, y_true: torch.Tenso mel_spec_true = mel_spec_true[..., :min_len] mel_spec_pred = mel_spec_pred[..., :min_len] - # L2 Loss (Mean Squared Error) loss = torch.mean((mel_spec_true - mel_spec_pred)**2) return loss @@ -49,7 +48,6 @@ def log_stft_magnitude_loss(stft_transform: T.Spectrogram, y_true: torch.Tensor, stft_mag_true = stft_mag_true[..., :min_len] stft_mag_pred = stft_mag_pred[..., :min_len] - # Log Magnitude L1 Loss loss = torch.mean(torch.abs(torch.log(stft_mag_true + eps) - torch.log(stft_mag_pred + eps))) return loss @@ -61,12 +59,9 @@ def spectral_convergence_loss(stft_transform: T.Spectrogram, y_true: torch.Tenso stft_mag_true = stft_mag_true[..., :min_len] stft_mag_pred = stft_mag_pred[..., :min_len] - # Calculate Frobenius norms and the loss - # Ensure norms are calculated over frequency and time dims ([..., freq, time]) 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)) - # Average loss over the batch loss = torch.mean(norm_diff / (norm_true + eps)) return loss @@ -77,16 +72,13 @@ def discriminator_train(high_quality, low_quality, real_labels, fake_labels, dis discriminator_decision_from_real = discriminator(high_quality[0]) d_loss_real = criterion(discriminator_decision_from_real, real_labels) - # Forward pass for fake samples (from generator output) - with torch.no_grad(): # Detach generator output within no_grad context + with torch.no_grad(): generator_output = generator(low_quality[0]) - discriminator_decision_from_fake = discriminator(generator_output) # No need to detach again if inside no_grad + discriminator_decision_from_fake = discriminator(generator_output) d_loss_fake = criterion(discriminator_decision_from_fake, fake_labels.expand_as(discriminator_decision_from_fake)) - # Combine real and fake losses d_loss = (d_loss_real + d_loss_fake) / 2.0 - # Backward pass and optimization d_loss.backward() # Optional: Gradient Clipping (can be helpful) # nn.utils.clip_grad_norm_(discriminator.parameters(), max_norm=1.0) # Gradient Clipping @@ -100,65 +92,53 @@ def generator_train( real_labels, generator, discriminator, - adv_criterion, # Criterion for adversarial loss (e.g., BCEWithLogitsLoss) + adv_criterion, g_optimizer, device, - # --- Pass necessary transforms and loss weights --- - mel_transform: T.MelSpectrogram, # Example: Pass Mel transform - # stft_transform: T.Spectrogram, # Pass STFT transform if using STFT losses - # mfcc_transform: T.MFCC, # Pass MFCC transform if using MFCC loss - lambda_adv: float = 1.0, # Weight for adversarial loss - lambda_mel_l1: float = 10.0, # Example: Weight for Mel L1 loss - # lambda_log_stft: float = 0.0, # Set weights > 0 for losses you want to use - # lambda_mfcc: float = 0.0 + 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() - # 1. Generate high-quality audio from low-quality input generator_output = generator(low_quality[0]) - # 2. Calculate Adversarial Loss (Generator tries to fool discriminator) discriminator_decision = discriminator(generator_output) - # Generator wants discriminator to output "real" labels for its fakes adversarial_loss = adv_criterion(discriminator_decision, real_labels.expand_as(discriminator_decision)) - # 3. Calculate Reconstruction/Spectrogram Loss(es) - # --- Choose and calculate the losses you want to include --- mel_l1 = 0.0 - # log_stft_l1 = 0.0 - # mfcc_l = 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, hq_audio, 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, hq_audio, generator_output) - # --- End of Loss Calculation Choices --- + # Calculate MFCC Loss if weight is positive + if lambda_mfcc > 0: + mfcc_l = gpu_mfcc_loss(mfcc_transform, high_quality[0], generator_output) - - # 4. Combine Losses - # Make sure calculated losses are tensors even if weights are 0 initially - # (or handle appropriately in the sum) 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 + 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) + (lambda_mel_l1 * mel_l1_tensor) + \ + (lambda_log_stft * log_stft_l1_tensor) + \ + (lambda_mfcc * mfcc_l_tensor) - # 5. Backward Pass and Optimization 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 + return generator_output, combined_loss, adversarial_loss, mel_l1_tensor, log_stft_l1_tensor, mfcc_l_tensor -- 2.43.0