import torch import torch.nn as nn import torch.nn.functional as F import torchaudio # Added for audio loading from torch.nn import Conv1d, ConvTranspose1d, AvgPool1d, Conv2d from torch.nn.utils import weight_norm, remove_weight_norm, spectral_norm from torch.utils.data import DataLoader, Dataset from torch.optim import AdamW from tqdm import tqdm import numpy as np import itertools import glob from torch.cuda.amp import autocast, GradScaler # Constants and Utilities LRELU_SLOPE = 0.1 SAMPLE_RATE = 16000 # Added sample rate constant def init_weights(m, mean=0.0, std=0.01): classname = m.__class__.__name__ if classname.find("Conv") != -1: m.weight.data.normal_(mean, std) def get_padding(kernel_size, dilation=1): return int((kernel_size*dilation - dilation)/2) # [Keep all model components unchanged] # Model Components class ResBlock1(nn.Module): def __init__(self, h, channels, kernel_size=3, dilation=(1, 3, 5)): super().__init__() self.h = h self.convs1 = nn.ModuleList([ weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0], padding=get_padding(kernel_size, dilation[0]))), weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1], padding=get_padding(kernel_size, dilation[1]))), weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2], padding=get_padding(kernel_size, dilation[2]))) ]) self.convs1.apply(init_weights) self.convs2 = nn.ModuleList([ weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1))), weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1))), weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1))) ]) self.convs2.apply(init_weights) def forward(self, x): for c1, c2 in zip(self.convs1, self.convs2): xt = F.leaky_relu(x, LRELU_SLOPE) xt = c1(xt) xt = F.leaky_relu(xt, LRELU_SLOPE) xt = c2(xt) x = xt + x return x def remove_weight_norm(self): for l in self.convs1: remove_weight_norm(l) for l in self.convs2: remove_weight_norm(l) class EmbeddingGenerator(nn.Module): def __init__(self, h): super().__init__() self.h = h self.num_kernels = len(h.resblock_kernel_sizes) self.num_upsamples = len(h.upsample_rates) self.conv_pre = weight_norm(Conv1d(h.embedding_dim, h.upsample_initial_channel, 7, 1, padding=3)) self.ups = nn.ModuleList() for i, (u, k) in enumerate(zip(h.upsample_rates, h.upsample_kernel_sizes)): self.ups.append(weight_norm( ConvTranspose1d(h.upsample_initial_channel//(2**i), h.upsample_initial_channel//(2**(i+1)), k, u, padding=(k-u)//2))) self.resblocks = nn.ModuleList() for i in range(len(self.ups)): ch = h.upsample_initial_channel//(2**(i+1)) for j, (k, d) in enumerate(zip(h.resblock_kernel_sizes, h.resblock_dilation_sizes)): self.resblocks.append(ResBlock1(h, ch, k, d)) self.conv_post = weight_norm(Conv1d(ch, 1, 7, 1, padding=3)) self.ups.apply(init_weights) self.conv_post.apply(init_weights) def forward(self, x): # Remove the transpose operation here x = self.conv_pre(x) # Input shape should be (B, embedding_dim, T) for i in range(self.num_upsamples): x = F.leaky_relu(x, LRELU_SLOPE) x = torch.utils.checkpoint.checkpoint(self.ups[i], x) # Checkpointing xs = None for j in range(self.num_kernels): res_out = torch.utils.checkpoint.checkpoint(self.resblocks[i*self.num_kernels+j], x) xs = res_out if xs is None else xs + res_out x = xs / self.num_kernels x = F.leaky_relu(x) x = self.conv_post(x) return torch.tanh(x) def remove_weight_norm(self): for l in self.ups: remove_weight_norm(l) for l in self.resblocks: l.remove_weight_norm() remove_weight_norm(self.conv_pre) remove_weight_norm(self.conv_post) class DiscriminatorP(nn.Module): def __init__(self, period, kernel_size=5, stride=3, use_spectral_norm=False): super().__init__() self.period = period norm_f = spectral_norm if use_spectral_norm else weight_norm self.convs = nn.ModuleList([ norm_f(Conv2d(1, 32, (kernel_size, 1), (stride, 1), padding=(get_padding(5, 1), 0))), norm_f(Conv2d(32, 128, (kernel_size, 1), (stride, 1), padding=(get_padding(5, 1), 0))), norm_f(Conv2d(128, 512, (kernel_size, 1), (stride, 1), padding=(get_padding(5, 1), 0))), norm_f(Conv2d(512, 1024, (kernel_size, 1), (stride, 1), padding=(get_padding(5, 1), 0))), norm_f(Conv2d(1024, 1024, (kernel_size, 1), 1, padding=(2, 0))), ]) self.conv_post = norm_f(Conv2d(1024, 1, (3, 1), 1, padding=(1, 0))) def forward(self, x): fmap = [] b, c, t = x.shape # Correct padding calculation required_length = ((t + self.period - 1) // self.period) * self.period n_pad = required_length - t if n_pad != 0: x = F.pad(x, (0, n_pad), "reflect") x = x.view(b, c, required_length // self.period, self.period) for conv in self.convs: x = conv(x) x = F.leaky_relu(x, LRELU_SLOPE) fmap.append(x) x = self.conv_post(x) fmap.append(x) return torch.flatten(x, 1, -1), fmap class MultiPeriodDiscriminator(nn.Module): def __init__(self): super().__init__() self.discriminators = nn.ModuleList([ DiscriminatorP(2), DiscriminatorP(3), DiscriminatorP(5), DiscriminatorP(7), DiscriminatorP(11), ]) def forward(self, y, y_hat): y_d_rs, y_d_gs, fmap_rs, fmap_gs = [], [], [], [] for d in self.discriminators: y_d_r, fmap_r = d(y) y_d_g, fmap_g = d(y_hat) y_d_rs.append(y_d_r) y_d_gs.append(y_d_g) fmap_rs.append(fmap_r) fmap_gs.append(fmap_g) return y_d_rs, y_d_gs, fmap_rs, fmap_gs class DiscriminatorS(nn.Module): def __init__(self, use_spectral_norm=False): super().__init__() norm_f = spectral_norm if use_spectral_norm else weight_norm self.convs = nn.ModuleList([ norm_f(Conv1d(1, 128, 15, 1, padding=7)), norm_f(Conv1d(128, 128, 41, 2, groups=4, padding=20)), norm_f(Conv1d(128, 256, 41, 2, groups=16, padding=20)), norm_f(Conv1d(256, 512, 41, 4, groups=16, padding=20)), norm_f(Conv1d(512, 1024, 41, 4, groups=16, padding=20)), norm_f(Conv1d(1024, 1024, 41, 1, groups=16, padding=20)), norm_f(Conv1d(1024, 1024, 5, 1, padding=2)), ]) self.conv_post = norm_f(Conv1d(1024, 1, 3, 1, padding=1)) def forward(self, x): fmap = [] for conv in self.convs: x = conv(x) x = F.leaky_relu(x, LRELU_SLOPE) fmap.append(x) x = self.conv_post(x) fmap.append(x) return torch.flatten(x, 1, -1), fmap class MultiScaleDiscriminator(nn.Module): def __init__(self): super().__init__() self.discriminators = nn.ModuleList([ DiscriminatorS(use_spectral_norm=True), DiscriminatorS(), DiscriminatorS(), ]) self.meanpools = nn.ModuleList([ AvgPool1d(4, 2, padding=2), AvgPool1d(4, 2, padding=2) ]) def forward(self, y, y_hat): y_d_rs, y_d_gs, fmap_rs, fmap_gs = [], [], [], [] for i, d in enumerate(self.discriminators): if i != 0: y = self.meanpools[i-1](y) y_hat = self.meanpools[i-1](y_hat) y_d_r, fmap_r = d(y) y_d_g, fmap_g = d(y_hat) y_d_rs.append(y_d_r) y_d_gs.append(y_d_g) fmap_rs.append(fmap_r) fmap_gs.append(fmap_g) return y_d_rs, y_d_gs, fmap_rs, fmap_gs # Loss Functions def feature_loss(fmap_r, fmap_g): loss = 0 for dr, dg in zip(fmap_r, fmap_g): for rl, gl in zip(dr, dg): loss += torch.mean(torch.abs(rl - gl)) return loss * 2 def discriminator_loss(disc_real_outputs, disc_generated_outputs): loss = 0 for dr, dg in zip(disc_real_outputs, disc_generated_outputs): loss += (torch.mean((1-dr)**2) + torch.mean(dg**2)) return loss def generator_loss(disc_outputs): loss = 0 for dg in disc_outputs: loss += torch.mean((1-dg)**2) return loss # Training Infrastructure class Hparams: def __init__(self): self.embedding_dim = 1024 self.upsample_rates = [10, 8, 4, 1] self.upsample_kernel_sizes = [20, 16, 8, 4] self.upsample_initial_channel = 256 self.resblock_kernel_sizes = [3, 7] self.resblock_dilation_sizes = [[1, 3, 5], [1, 3, 5], [1, 3, 5]] # Modified Dataset Class class EmbeddingAudioDataset(Dataset): def __init__(self, embedding_files, audio_files, max_length=16000*3): self.embedding_files = embedding_files self.audio_files = audio_files self.max_length = max_length self.resampler = torchaudio.transforms.Resample(orig_freq=SAMPLE_RATE, new_freq=SAMPLE_RATE) def __len__(self): return len(self.embedding_files) def __getitem__(self, idx): # Load embedding and convert to tensor immediately embedding = np.load(self.embedding_files[idx], allow_pickle=True) if embedding.ndim == 1: embedding = embedding.reshape(1, -1) embedding = torch.from_numpy(embedding).float() # Convert to tensor here # Load and process audio waveform, orig_sr = torchaudio.load(self.audio_files[idx]) if waveform.shape[0] > 1: waveform = torch.mean(waveform, dim=0, keepdim=True) if orig_sr != SAMPLE_RATE: waveform = self.resampler(waveform) # Audio truncation/padding if waveform.shape[1] > self.max_length: start = torch.randint(0, waveform.shape[1] - self.max_length, (1,)) waveform = waveform[:, start:start+self.max_length] else: waveform = F.pad(waveform, (0, self.max_length - waveform.shape[1])) # Embedding truncation/padding emb_len = self.max_length // 320 if embedding.shape[0] > emb_len: embedding = embedding[:emb_len] else: # Pad time dimension (axis 0) with zeros embedding = F.pad(embedding, (0, 0, 0, emb_len - embedding.shape[0])) return embedding, waveform.squeeze().float() # # Load embedding # embedding = np.load(self.embedding_files[idx], allow_pickle=True) # Fix 1: Allow pickles # embedding.shape # # Ensure 2D shape: (time_steps, embedding_dim) # if embedding.ndim == 1: # embedding = embedding.reshape(1, -1) # Add time dimension if missing # # Load and process audio # waveform, orig_sr = torchaudio.load(self.audio_files[idx]) # Fix 2: Proper audio loading # if waveform.shape[0] > 1: # Convert to mono # waveform = torch.mean(waveform, dim=0, keepdim=True) # if orig_sr != SAMPLE_RATE: # waveform = self.resampler(waveform) # waveform = waveform / torch.max(torch.abs(waveform)) # Normalize # # Truncate/pad audio to fixed length # if waveform.shape[1] > self.max_length: # start = torch.randint(0, waveform.shape[1] - self.max_length, (1,)) # waveform = waveform[:, start:start+self.max_length] # else: # waveform = F.pad(waveform, (0, self.max_length - waveform.shape[1])) # # Truncate/pad embeddings correspondingly # emb_len = self.max_length // 320 # 320 samples per embedding step # if embedding.shape[0] > emb_len: # embedding = embedding[:emb_len] # else: # embedding = F.pad(embedding, (0, 0, 0, emb_len - embedding.shape[0])) # return embedding.float(), waveform.squeeze().float() # return torch.from_numpy(embedding).float(), waveform.squeeze().float() # [Keep rest of the code unchanged until main execution] if __name__ == "__main__": h = Hparams() # In main execution device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Initialize models generator = EmbeddingGenerator(h).to(device) mpd = MultiPeriodDiscriminator().to(device) msd = MultiScaleDiscriminator().to(device) # Optimizers optim_g = AdamW(generator.parameters(), lr=0.0002, betas=(0.8, 0.99)) optim_d = AdamW(itertools.chain(mpd.parameters(), msd.parameters()), lr=0.0002, betas=(0.8, 0.99)) scaler_g = torch.amp.GradScaler() scaler_d = torch.amp.GradScaler() # Dataset embedding_files = sorted(glob.glob('/home/vikrant/Conversational-AI-Model/embedding_vocoder/embeddings/*.npy')) audio_files = sorted(glob.glob('/home/vikrant/Conversational-AI-Model/embedding_vocoder/non_empty_wavs/*.wav')) assert len(embedding_files) == len(audio_files), "Mismatched files" loader = DataLoader( EmbeddingAudioDataset(embedding_files, audio_files), batch_size=8, shuffle=True,persistent_workers=False, num_workers=4, pin_memory=False ) # # Training loop modifications # for epoch in range(num_epochs): # generator.train() # mpd.train() # msd.train() # for batch_idx, (embeddings, audio) in enumerate(tqdm(loader)): # embeddings, audio = embeddings.to(device), audio.to(device) # # Add channel dimension to audio # audio = audio.unsqueeze(1) # Shape: (B, 1, T) # # Convert embeddings to (B, D, T) # if embeddings.ndim == 2: # Single sample case # embeddings = embeddings.unsqueeze(0).permute(0, 2, 1) # else: # Batch case # embeddings = embeddings.permute(0, 2, 1) # # --- Mixed precision forward pass --- # with autocast(): # # Generate audio # fake_audio = generator(embeddings) # # Pad/Crop to match target length # target_len = audio.size(2) # fake_audio = fake_audio[:, :, :target_len] if fake_audio.size(2) > target_len \ # else F.pad(fake_audio, (0, target_len - fake_audio.size(2))) # # --------------------------- # # Discriminator training # # --------------------------- # optim_d.zero_grad() # # Mixed precision for discriminator forward # with autocast(): # # Get discriminator outputs # y_d_rs_mpd, y_d_gs_mpd, _, _ = mpd(audio, fake_audio.detach()) # y_d_rs_msd, y_d_gs_msd, _, _ = msd(audio, fake_audio.detach()) # # Calculate discriminator loss # loss_disc = discriminator_loss(y_d_rs_mpd + y_d_rs_msd, # y_d_gs_mpd + y_d_gs_msd) # # Scaled backward # scaler_d.scale(loss_disc).backward() # scaler_d.step(optim_d) # scaler_d.update() # # --------------------------- # # Generator training # # --------------------------- # optim_g.zero_grad() # # Mixed precision for generator forward # with autocast(): # # Get discriminator outputs and feature maps # y_d_gs_mpd, _, fmap_rs_mpd, fmap_gs_mpd = mpd(audio, fake_audio) # y_d_gs_msd, _, fmap_rs_msd, fmap_gs_msd = msd(audio, fake_audio) # # Calculate losses # loss_gen = generator_loss(y_d_gs_mpd + y_d_gs_msd) # loss_feat = feature_loss(fmap_rs_mpd + fmap_rs_msd, # fmap_gs_mpd + fmap_gs_msd) # total_loss = loss_gen + loss_feat # # Scaled backward # scaler_g.scale(total_loss).backward() # scaler_g.step(optim_g) # scaler_g.update() # if batch_idx % 100 == 0: # print(f"Epoch {epoch} Batch {batch_idx} | G Loss: {loss_gen.item():.3f} | D Loss: {loss_disc.item():.3f}") # # # Training loop # # num_epochs = 1 # # for epoch in range(num_epochs): # # generator.train() # # mpd.train() # # msd.train() # # for batch_idx, (embeddings, audio) in enumerate(tqdm(loader)): # # embeddings, audio = embeddings.to(device), audio.to(device) # # # Add channel dimension to audio # # audio = audio.unsqueeze(1) # Shape: (B, 1, T) # # # Convert embeddings to (B, D, T) # # if embeddings.ndim == 2: # Single sample case # # embeddings = embeddings.unsqueeze(0).permute(0, 2, 1) # # else: # Batch case # # embeddings = embeddings.permute(0, 2, 1) # # # Generate audio # # fake_audio = generator(embeddings) # # # # Calculate required length based on embeddings # # # expected_samples = embeddings.shape[2] * 320 # 320 samples per embedding step # # # fake_audio = fake_audio[:, :, :expected_samples] # # target_len = audio.size(2) # audio is now (B, 1, T) # # fake_audio = fake_audio[:, :, :target_len] if fake_audio.size(2) > target_len \ # # else F.pad(fake_audio, (0, target_len - fake_audio.size(2))) # # # Discriminator training # # optim_d.zero_grad() # # # _, disc_real_mpd, _, disc_real_msd = mpd(audio, fake_audio.detach()) # # # _, disc_fake_mpd, _, disc_fake_msd = msd(audio, fake_audio.detach()) # # # Get discriminator outputs # # y_d_rs_mpd, y_d_gs_mpd, _, _ = mpd(audio, fake_audio.detach()) # First 2 returns are real/generated outputs # # y_d_rs_msd, y_d_gs_msd, _, _ = msd(audio, fake_audio.detach()) # First 2 returns are real/generated outputs # # # loss_disc = discriminator_loss(disc_real_mpd + disc_real_msd, disc_fake_mpd + disc_fake_msd) # # loss_disc = discriminator_loss(y_d_rs_mpd + y_d_rs_msd, # All real outputs # # y_d_gs_mpd + y_d_gs_msd) # All generated outputs # # loss_disc.backward() # # optim_d.step() # # # Generator training # # optim_g.zero_grad() # # # disc_out_mpd, disc_out_msd, fmap_mpd, fmap_msd = mpd(audio, fake_audio) # # # Get discriminator outputs and feature maps # # y_d_gs_mpd, _, fmap_rs_mpd, fmap_gs_mpd = mpd(audio, fake_audio) # # y_d_gs_msd, _, fmap_rs_msd, fmap_gs_msd = msd(audio, fake_audio) # # # loss_gen = generator_loss(disc_out_mpd + disc_out_msd) # # # loss_feat = feature_loss(fmap_mpd + fmap_msd) # # # (loss_gen + loss_feat).backward() # # # optim_g.step() # # # Calculate losses # # loss_gen = generator_loss(y_d_gs_mpd + y_d_gs_msd) # # loss_feat = feature_loss(fmap_rs_mpd + fmap_rs_msd, # # fmap_gs_mpd + fmap_gs_msd) # # (loss_gen + loss_feat).backward() # # optim_g.step() # # if batch_idx % 100 == 0: # # print(f"Epoch {epoch} Batch {batch_idx} | G Loss: {loss_gen.item():.3f} | D Loss: {loss_disc.item():.3f}") # # Save checkpoint # torch.save({ # 'generator': generator.state_dict(), # 'mpd': mpd.state_dict(), # 'msd': msd.state_dict(), # 'optim_g': optim_g.state_dict(), # 'optim_d': optim_d.state_dict(), # 'epoch': epoch, # }, f"checkpoint_epoch_{epoch}.pt") # # Finalize # generator.eval() # generator.remove_weight_norm() # print("Training completed") # Training parameters num_epochs = 100 start_epoch = 1 # Load checkpoint if resuming checkpoint_path = None # Set to path if resuming if checkpoint_path: checkpoint = torch.load(checkpoint_path) generator.load_state_dict(checkpoint['generator']) mpd.load_state_dict(checkpoint['mpd']) msd.load_state_dict(checkpoint['msd']) optim_g.load_state_dict(checkpoint['optim_g']) optim_d.load_state_dict(checkpoint['optim_d']) scaler_g.load_state_dict(checkpoint['scaler_g']) scaler_d.load_state_dict(checkpoint['scaler_d']) start_epoch = checkpoint['epoch'] + 1 for epoch in range(start_epoch, num_epochs): generator.train() mpd.train() msd.train() for batch_idx, (embeddings, audio) in enumerate(tqdm(loader)): embeddings, audio = embeddings.to(device), audio.to(device) audio = audio.unsqueeze(1) # Convert embeddings if embeddings.ndim == 2: embeddings = embeddings.unsqueeze(0).permute(0, 2, 1) else: embeddings = embeddings.permute(0, 2, 1) # --- Mixed precision forward --- with autocast(): fake_audio = generator(embeddings) target_len = audio.size(2) fake_audio = fake_audio[:, :, :target_len] if fake_audio.size(2) > target_len \ else F.pad(fake_audio, (0, target_len - fake_audio.size(2))) # --- Discriminator update --- optim_d.zero_grad() with autocast(): y_d_rs_mpd, y_d_gs_mpd, _, _ = mpd(audio, fake_audio.detach()) y_d_rs_msd, y_d_gs_msd, _, _ = msd(audio, fake_audio.detach()) loss_disc = discriminator_loss(y_d_rs_mpd + y_d_rs_msd, y_d_gs_mpd + y_d_gs_msd) scaler_d.scale(loss_disc).backward() scaler_d.step(optim_d) scaler_d.update() # --- Generator update --- optim_g.zero_grad() with autocast(): y_d_gs_mpd, _, fmap_rs_mpd, fmap_gs_mpd = mpd(audio, fake_audio) y_d_gs_msd, _, fmap_rs_msd, fmap_gs_msd = msd(audio, fake_audio) loss_gen = generator_loss(y_d_gs_mpd + y_d_gs_msd) loss_feat = feature_loss(fmap_rs_mpd + fmap_rs_msd, fmap_gs_mpd + fmap_gs_msd) total_loss = loss_gen + loss_feat scaler_g.scale(total_loss).backward() scaler_g.step(optim_g) scaler_g.update() if batch_idx % 100 == 0: print(f"Epoch {epoch} Batch {batch_idx} | G Loss: {loss_gen.item():.3f} | D Loss: {loss_disc.item():.3f}") print(f"Memory allocated: {torch.cuda.memory_allocated()/1e9:.2f}GB") print(f"Memory reserved: {torch.cuda.memory_reserved()/1e9:.2f}GB") # Memory cleanup torch.cuda.empty_cache() # Save checkpoint torch.save({ 'generator': generator.state_dict(), 'mpd': mpd.state_dict(), 'msd': msd.state_dict(), 'optim_g': optim_g.state_dict(), 'optim_d': optim_d.state_dict(), 'scaler_g': scaler_g.state_dict(), 'scaler_d': scaler_d.state_dict(), 'epoch': epoch, 'hparams': h.__dict__, }, f"checkpoint_epoch_{epoch}.pt") # Finalize generator.eval().cpu() generator.remove_weight_norm() torch.save(generator.state_dict(), "final_generator.pth") print("Training completed")