Merge pull request 'Merge new-arch, because it has proven to give the best results' (#1) from new-arch into main

Reviewed-on: #1

AudioUtils.py

@@ -16,37 +16,3 @@ def stretch_tensor(tensor, target_length):
     tensor = F.interpolate(tensor, scale_factor=scale_factor, mode='linear', align_corners=False)
 
     return tensor
-
-def pad_tensor(audio_tensor: torch.Tensor, target_length: int = 128):
-    current_length = audio_tensor.shape[-1]
-
-    if current_length < target_length:
-        padding_needed = target_length - current_length
-
-        padding_tuple = (0, padding_needed)
-        padded_audio_tensor = F.pad(audio_tensor, padding_tuple, mode='constant', value=0)
-    else:
-        padded_audio_tensor = audio_tensor
-
-    return padded_audio_tensor
-
-def split_audio(audio_tensor: torch.Tensor, chunk_size: int = 128) -> list[torch.Tensor]:
-    if not isinstance(chunk_size, int) or chunk_size <= 0:
-        raise ValueError("chunk_size must be a positive integer.")
-
-    # Handle scalar tensor edge case if necessary
-    if audio_tensor.dim() == 0:
-        return [audio_tensor] if audio_tensor.numel() > 0 else []
-
-    # Identify the dimension to split (usually the last one, representing time/samples)
-    split_dim = -1
-    num_samples = audio_tensor.shape[split_dim]
-
-    if num_samples == 0:
-        return [] # Return empty list if the dimension to split is empty
-
-    # Use torch.split to divide the tensor into chunks
-    # It handles the last chunk being potentially smaller automatically.
-    chunks = list(torch.split(audio_tensor, chunk_size, dim=split_dim))
-
-    return chunks

data.py

@@ -5,42 +5,49 @@ import torchaudio
 import os
 import random
 import torchaudio.transforms as T
-import tqdm
 import AudioUtils
 
 class AudioDataset(Dataset):
     audio_sample_rates = [11025]
+    MAX_LENGTH = 44100  # Define your desired maximum length here
 
-    def __init__(self, input_dir, device, clip_length = 256):
+    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
-        input_files = [os.path.join(root, f) for root, _, files in os.walk(input_dir) for f in files if f.endswith('.wav') or f.endswith('.mp3') or f.endswith('.flac')]
-
-        data = []
-        for audio_clip in tqdm.tqdm(input_files, desc=f"Processing {len(input_files)} audio file(s)"):
-            audio, original_sample_rate = torchaudio.load(audio_clip, normalize=True)
-            audio = AudioUtils.stereo_tensor_to_mono(audio)
-
-            # Generate low-quality audio with random downsampling
-            mangled_sample_rate = random.choice(self.audio_sample_rates)
-            resample_transform_low = torchaudio.transforms.Resample(original_sample_rate, mangled_sample_rate)
-            resample_transform_high = torchaudio.transforms.Resample(mangled_sample_rate, original_sample_rate)
-
-            low_audio = resample_transform_low(audio)
-            low_audio = resample_transform_high(low_audio)
-
-            splitted_high_quality_audio = AudioUtils.split_audio(audio, clip_length)
-            splitted_high_quality_audio[-1] = AudioUtils.pad_tensor(splitted_high_quality_audio[-1], clip_length)
-
-            splitted_low_quality_audio = AudioUtils.split_audio(low_audio, clip_length)
-            splitted_low_quality_audio[-1] = AudioUtils.pad_tensor(splitted_low_quality_audio[-1], clip_length)
-
-            for high_quality_sample, low_quality_sample in zip(splitted_high_quality_audio, splitted_low_quality_audio):
-                data.append(((high_quality_sample, low_quality_sample), (original_sample_rate, mangled_sample_rate)))
-
-        self.audio_data = data
 
     def __len__(self):
-        return len(self.audio_data)
+        return len(self.input_files)
 
     def __getitem__(self, idx):
-        return self.audio_data[idx]
+        # Load high-quality audio
+        high_quality_audio, original_sample_rate = torchaudio.load(self.input_files[idx], normalize=True)
+
+        # Generate low-quality audio with random downsampling
+        mangled_sample_rate = random.choice(self.audio_sample_rates)
+        resample_transform_low = torchaudio.transforms.Resample(original_sample_rate, mangled_sample_rate)
+        low_quality_audio = resample_transform_low(high_quality_audio)
+
+        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)
+        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)

file_utils.py

@@ -2,22 +2,20 @@ import json
 
 filepath = "my_data.json"
 
-def write_data(filepath, data, debug=False):
+def write_data(filepath, data):
   try:
     with open(filepath, 'w') as f:
       json.dump(data, f, indent=4)  # Use indent for pretty formatting
-    if debug:
+    print(f"Data written to '{filepath}'")
-        print(f"Data written to '{filepath}'")
   except Exception as e:
     print(f"Error writing to file: {e}")
 
 
-def read_data(filepath, debug=False):
+def read_data(filepath):
   try:
     with open(filepath, 'r') as f:
       data = json.load(f)
-    if debug:
+    print(f"Data read from '{filepath}'")
-        print(f"Data read from '{filepath}'")
     return data
   except FileNotFoundError:
     print(f"File not found: {filepath}")

training.py

@@ -43,11 +43,11 @@ print(f"Using device: {device}")
 
 # Parameters
 sample_rate = 44100
-n_fft = 128
+n_fft = 2048
-hop_length = 128
+hop_length = 256
 win_length = n_fft
-n_mels = 40
+n_mels = 128
-n_mfcc = 13 # If using MFCC
+n_mfcc = 20 # If using MFCC
 
 mfcc_transform = T.MFCC(
     sample_rate,
@@ -76,7 +76,7 @@ os.makedirs(audio_output_dir, exist_ok=True)
 
 # ========= SINGLE =========
 
-train_data_loader = DataLoader(dataset, batch_size=1024, shuffle=True)
+train_data_loader = DataLoader(dataset, batch_size=64, shuffle=True)
 
 
 # ========= MODELS =========
@@ -115,35 +115,28 @@ scheduler_d = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer_d, mode='min'
 def start_training():
     generator_epochs = 5000
     for generator_epoch in range(generator_epochs):
-        high_quality_audio = ([torch.empty((1))], 1)
+        low_quality_audio = (torch.empty((1)), 1)
-        low_quality_audio = ([torch.empty((1))], 1)
+        high_quality_audio = (torch.empty((1)), 1)
-        ai_enhanced_audio = ([torch.empty((1))], 1)
+        ai_enhanced_audio = (torch.empty((1)), 1)
 
         times_correct = 0
 
         # ========= TRAINING =========
-        for training_data 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 tqdm.tqdm(train_data_loader, desc=f"Training epoch {generator_epoch+1}/{generator_epochs}, Current epoch {epoch+1}"):
-            ## Data structure:
+        # for high_quality_clip, low_quality_clip in train_data_loader:
-            # [[[float..., float..., float...], [float..., float..., float...]], [original_sample_rate, mangled_sample_rate]]
+            high_quality_sample = (high_quality_clip[0], high_quality_clip[1])
+            low_quality_sample = (low_quality_clip[0], low_quality_clip[1])
 
             # ========= LABELS =========
-            good_quality_data = training_data[0][0].to(device)
+            batch_size = high_quality_clip[0].size(0)
-            bad_quality_data = training_data[0][1].to(device)
-            original_sample_rate = training_data[1][0]
-            mangled_sample_rate = training_data[1][1]
-
-            batch_size = good_quality_data.size(0)
             real_labels = torch.ones(batch_size, 1).to(device)
             fake_labels = torch.zeros(batch_size, 1).to(device)
 
-            high_quality_audio = (good_quality_data, original_sample_rate)
-            low_quality_audio = (bad_quality_data, mangled_sample_rate)
-
             # ========= DISCRIMINATOR =========
             discriminator.train()
             d_loss = discriminator_train(
-                good_quality_data,
+                high_quality_sample,
-                bad_quality_data,
+                low_quality_sample,
                 real_labels,
                 fake_labels,
                 discriminator,
@@ -155,8 +148,8 @@ def start_training():
             # ========= GENERATOR =========
             generator.train()
             generator_output, combined_loss, adversarial_loss, mel_l1_tensor, log_stft_l1_tensor, mfcc_l_tensor = generator_train(
-                bad_quality_data,
+                low_quality_sample,
-                good_quality_data,
+                high_quality_sample,
                 real_labels,
                 generator,
                 discriminator,
@@ -174,17 +167,17 @@ def start_training():
             scheduler_g.step(adversarial_loss.detach())
 
             # ========= SAVE LATEST AUDIO =========
-            high_quality_audio = (good_quality_data, original_sample_rate)
+            high_quality_audio = (high_quality_clip[0][0], high_quality_clip[1][0])
-            low_quality_audio = (bad_quality_data, original_sample_rate)
+            low_quality_audio = (low_quality_clip[0][0], low_quality_clip[1][0])
-            ai_enhanced_audio = (generator_output, original_sample_rate)
+            ai_enhanced_audio = (generator_output[0], high_quality_clip[1][0])
 
         new_epoch = generator_epoch+epoch
 
-        # if generator_epoch % 25 == 0:
+        if generator_epoch % 25 == 0:
-        #     print(f"Saved epoch {new_epoch}!")
+            print(f"Saved epoch {new_epoch}!")
-        #     torchaudio.save(f"{audio_output_dir}/epoch-{new_epoch}-audio-orig.wav", high_quality_audio[0][-1].cpu().detach(), high_quality_audio[1][-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-crap.wav", low_quality_audio[0][-1].cpu().detach(), high_quality_audio[1][-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-ai.wav", ai_enhanced_audio[0][-1].cpu().detach(), high_quality_audio[1][-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())

training_utils.py

@@ -20,10 +20,12 @@ def mel_spectrogram_l1_loss(mel_transform: T.MelSpectrogram, y_true: torch.Tenso
     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
 
@@ -67,11 +69,11 @@ def discriminator_train(high_quality, low_quality, real_labels, fake_labels, dis
     optimizer.zero_grad()
 
     # Forward pass for real samples
-    discriminator_decision_from_real = discriminator(high_quality)
+    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)
+        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))
 
@@ -103,7 +105,7 @@ def generator_train(
 ):
     g_optimizer.zero_grad()
 
-    generator_output = generator(low_quality)
+    generator_output = generator(low_quality[0])
 
     discriminator_decision = discriminator(generator_output)
     adversarial_loss = adv_criterion(discriminator_decision, real_labels.expand_as(discriminator_decision))
@@ -114,15 +116,15 @@ def generator_train(
 
     # Calculate Mel L1 Loss if weight is positive
     if lambda_mel_l1 > 0:
-        mel_l1 = mel_spectrogram_l1_loss(mel_transform, high_quality, generator_output)
+        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, generator_output)
+        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, generator_output)
+        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