diff --git a/AudioUtils.py b/AudioUtils.py index 04f75db..f4866dd 100644 --- a/AudioUtils.py +++ b/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 diff --git a/data.py b/data.py index bc7574f..88364b6 100644 --- a/data.py +++ b/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 - 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] + splitted_high_quality_audio = AudioUtils.split_audio(high_quality_audio, 128) + splitted_high_quality_audio[-1] = AudioUtils.pad_tensor(splitted_high_quality_audio[-1], 128) + splitted_high_quality_audio = [tensor.to(self.device) for tensor in splitted_high_quality_audio] - # 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] + splitted_low_quality_audio = AudioUtils.split_audio(low_quality_audio, 128) + splitted_low_quality_audio[-1] = AudioUtils.pad_tensor(splitted_low_quality_audio[-1], 128) + splitted_low_quality_audio = [tensor.to(self.device) for tensor in splitted_low_quality_audio] - 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) + return (splitted_high_quality_audio, original_sample_rate), (splitted_low_quality_audio, mangled_sample_rate) diff --git a/training.py b/training.py index 01ea749..f6ab2f4 100644 --- a/training.py +++ b/training.py @@ -43,11 +43,11 @@ print(f"Using device: {device}") # Parameters sample_rate = 44100 -n_fft = 2048 -hop_length = 256 +n_fft = 128 +hop_length = 128 win_length = n_fft -n_mels = 128 -n_mfcc = 20 # If using MFCC +n_mels = 40 +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,61 +115,69 @@ 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) - ai_enhanced_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) 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_clip, low_quality_clip in train_data_loader: - high_quality_sample = (high_quality_clip[0], high_quality_clip[1]) - low_quality_sample = (low_quality_clip[0], low_quality_clip[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}"): + ## Data structure: + # [[float..., float..., float...], sample_rate] # ========= 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) - # ========= DISCRIMINATOR ========= - discriminator.train() - d_loss = discriminator_train( - high_quality_sample, - low_quality_sample, - real_labels, - fake_labels, - discriminator, - generator, - criterion_d, - optimizer_d - ) + high_quality_audio = high_quality_data + low_quality_audio = low_quality_data - # ========= GENERATOR ========= - 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, - generator, - discriminator, - criterion_d, - optimizer_g, - device, - mel_transform, - stft_transform, - mfcc_transform - ) + ai_enhanced_outputs = [] - 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()) + 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( + 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, mel_l1_tensor, log_stft_l1_tensor, mfcc_l_tensor = generator_train( + low_quality_sample, + high_quality_sample, + real_labels, + generator, + discriminator, + criterion_d, + optimizer_g, + device, + mel_transform, + stft_transform, + mfcc_transform + ) + + 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]) - low_quality_audio = (low_quality_clip[0][0], low_quality_clip[1][0]) - ai_enhanced_audio = (generator_output[0], high_quality_clip[1][0]) + high_quality_audio = (torch.cat(high_quality_data[0]), high_quality_data[1]) + low_quality_audio = (torch.cat(low_quality_data[0]), low_quality_data[1]) + ai_enhanced_audio = (torch.cat(ai_enhanced_outputs), high_quality_data[1]) new_epoch = generator_epoch+epoch diff --git a/training_utils.py b/training_utils.py index 6f26f58..c7d43e5 100644 --- a/training_utils.py +++ b/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