:albemic: | Real-time testing...

AudioUtils.py

@@ -16,3 +16,37 @@ 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

@@ -21,33 +21,25 @@ class AudioDataset(Dataset):
     def __getitem__(self, idx):
         # Load high-quality audio
         high_quality_audio, original_sample_rate = torchaudio.load(self.input_files[idx], normalize=True)
+        # Change to mono
+        high_quality_audio = AudioUtils.stereo_tensor_to_mono(high_quality_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)
-        low_quality_audio = resample_transform_low(high_quality_audio)
 
+        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_quality_audio = resample_transform_low(high_quality_audio)
         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
+        splitted_high_quality_audio = AudioUtils.split_audio(high_quality_audio, 128)
-        if high_quality_audio.shape[1] < self.MAX_LENGTH:
+        splitted_high_quality_audio[-1] = AudioUtils.pad_tensor(splitted_high_quality_audio[-1], 128)
-            padding = self.MAX_LENGTH - high_quality_audio.shape[1]
+        splitted_high_quality_audio = [tensor.to(self.device) for tensor in splitted_high_quality_audio]
-            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
+        splitted_low_quality_audio = AudioUtils.split_audio(low_quality_audio, 128)
-        if low_quality_audio.shape[1] < self.MAX_LENGTH:
+        splitted_low_quality_audio[-1] = AudioUtils.pad_tensor(splitted_low_quality_audio[-1], 128)
-            padding = self.MAX_LENGTH - low_quality_audio.shape[1]
+        splitted_low_quality_audio = [tensor.to(self.device) for tensor in splitted_low_quality_audio]
-            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)
+        return (splitted_high_quality_audio, original_sample_rate), (splitted_low_quality_audio, mangled_sample_rate)
-        low_quality_audio = low_quality_audio.to(self.device)
-
-        return (high_quality_audio, original_sample_rate), (low_quality_audio, mangled_sample_rate)

training.py

@@ -43,11 +43,11 @@ print(f"Using device: {device}")
 
 # Parameters
 sample_rate = 44100
-n_fft = 2048
+n_fft = 128
-hop_length = 256
+hop_length = 128
 win_length = n_fft
-n_mels = 128
+n_mels = 40
-n_mfcc = 20 # If using MFCC
+n_mfcc = 13 # 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=64, shuffle=True)
+train_data_loader = DataLoader(dataset, batch_size=1, shuffle=True)
 
 
 # ========= MODELS =========
@@ -115,23 +115,29 @@ scheduler_d = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer_d, mode='min'
 def start_training():
     generator_epochs = 5000
     for generator_epoch in range(generator_epochs):
-        low_quality_audio = (torch.empty((1)), 1)
+        high_quality_audio = ([torch.empty((1))], 1)
-        high_quality_audio = (torch.empty((1)), 1)
+        low_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 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_data, low_quality_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 train_data_loader:
+            ## Data structure:
-            high_quality_sample = (high_quality_clip[0], high_quality_clip[1])
+            # [[float..., float..., float...], sample_rate]
-            low_quality_sample = (low_quality_clip[0], low_quality_clip[1])
 
             # ========= LABELS =========
-            batch_size = high_quality_clip[0].size(0)
+
+            batch_size = high_quality_data[0][0].size(0)
             real_labels = torch.ones(batch_size, 1).to(device)
             fake_labels = torch.zeros(batch_size, 1).to(device)
 
+            high_quality_audio = high_quality_data
+            low_quality_audio = low_quality_data
+
+            ai_enhanced_outputs = []
+
+            for high_quality_sample, low_quality_sample in tqdm.tqdm(zip(high_quality_data[0], low_quality_data[0]), desc=f"Processing audio clip.. Length: {len(high_quality_data[0])}"):
                 # ========= DISCRIMINATOR =========
                 discriminator.train()
                 d_loss = discriminator_train(
@@ -161,15 +167,17 @@ def start_training():
                     mfcc_transform
                 )
 
+                ai_enhanced_outputs.append(generator_output)
+
                 if debug:
                     print(f"D_LOSS: {d_loss.item():.4f}, COMBINED_LOSS: {combined_loss.item():.4f}, ADVERSARIAL_LOSS: {adversarial_loss.item():.4f}, MEL_L1_LOSS: {mel_l1_tensor.item():.4f}, LOG_STFT_L1_LOSS: {log_stft_l1_tensor.item():.4f}, MFCC_LOSS: {mfcc_l_tensor.item():.4f}")
                 scheduler_d.step(d_loss.detach())
                 scheduler_g.step(adversarial_loss.detach())
 
             # ========= SAVE LATEST AUDIO =========
-            high_quality_audio = (high_quality_clip[0][0], high_quality_clip[1][0])
+            high_quality_audio = (torch.cat(high_quality_data[0]), high_quality_data[1])
-            low_quality_audio = (low_quality_clip[0][0], low_quality_clip[1][0])
+            low_quality_audio = (torch.cat(low_quality_data[0]), low_quality_data[1])
-            ai_enhanced_audio = (generator_output[0], high_quality_clip[1][0])
+            ai_enhanced_audio = (torch.cat(ai_enhanced_outputs), high_quality_data[1])
 
         new_epoch = generator_epoch+epoch
 

training_utils.py

@@ -20,12 +20,10 @@ 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
 
@@ -69,11 +67,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[0])
+    discriminator_decision_from_real = discriminator(high_quality)
     d_loss_real = criterion(discriminator_decision_from_real, real_labels)
 
     with torch.no_grad():
-        generator_output = generator(low_quality[0])
+        generator_output = generator(low_quality)
     discriminator_decision_from_fake = discriminator(generator_output)
     d_loss_fake = criterion(discriminator_decision_from_fake, fake_labels.expand_as(discriminator_decision_from_fake))
 
@@ -105,7 +103,7 @@ def generator_train(
 ):
     g_optimizer.zero_grad()
 
-    generator_output = generator(low_quality[0])
+    generator_output = generator(low_quality)
 
     discriminator_decision = discriminator(generator_output)
     adversarial_loss = adv_criterion(discriminator_decision, real_labels.expand_as(discriminator_decision))
@@ -116,15 +114,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[0], generator_output)
+        mel_l1 = mel_spectrogram_l1_loss(mel_transform, high_quality, 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)
+        log_stft_l1 = log_stft_magnitude_loss(stft_transform, high_quality, generator_output)
 
     # Calculate MFCC Loss if weight is positive
     if lambda_mfcc > 0:
-        mfcc_l = gpu_mfcc_loss(mfcc_transform, high_quality[0], generator_output)
+        mfcc_l = gpu_mfcc_loss(mfcc_transform, high_quality, 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