def __init__(self, args):

self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

self.spectrogram_vq = SpectrogramVQ(args).to(self.device)

self.discriminator = Discriminator(args).to(self.device)

self.discriminator.apply(weights_init)

self.opt_vq, self.opt_disc = self.configure_optimizers(args)

self.vocoder = load_hifigan_model(args.hifigan_config_path, args.hifigan_ckpt_path, self.device)

self.prepare_training()

self.train(args)

def configure_optimizers(self, args):

lr = args.learning_rate

opt_vq = torch.optim.Adam(

list(self.spectrogram_vq.encoder.parameters()) +

list(self.spectrogram_vq.codebook.parameters()) +

list(self.spectrogram_vq.decoder.parameters()) +

list(self.spectrogram_vq.quant_conv.parameters()) +

list(self.spectrogram_vq.post_quant_conv.parameters()),

lr=lr, eps=1e-8, betas=(args.beta1, args.beta2)

)

opt_disc = torch.optim.Adam(

self.discriminator.parameters(),

lr=lr, eps=1e-8, betas=(args.beta1, args.beta2)

)

return opt_vq, opt_disc

@staticmethod

def prepare_training():

os.makedirs("results", exist_ok=True)

os.makedirs("checkpoints", exist_ok=True)

def save_checkpoint(self, epoch):

torch.save({

'spectrogram_vq_state_dict': self.spectrogram_vq.state_dict(),

'discriminator_state_dict': self.discriminator.state_dict(),

'opt_vq_state_dict': self.opt_vq.state_dict(),

'opt_disc_state_dict': self.opt_disc.state_dict(),

'epoch': epoch,

}, os.path.join("checkpoints", f"spectrogram_vq_epoch_{epoch}.pt"))

def train(self, args):

dataloader = load_data(args)

steps_per_epoch = len(dataloader)

for epoch in range(args.epochs):

pbar = tqdm(enumerate(dataloader), total=steps_per_epoch, desc=f"Epoch {epoch}")

for i, mels in pbar:

mels = mels.to(self.device)

global_step = epoch * steps_per_epoch + i

disc_factor = self.spectrogram_vq.adopt_weight(args.disc_factor, global_step, args.disc_start)

# Train Generator (VQ-VAE)

mel_hat, _, q_loss = self.spectrogram_vq(mels)

rec_loss = (args.rec_loss_factor * torch.abs(mel_hat - mels)).mean()

if disc_factor > 0:

disc_fake = self.discriminator(mel_hat)

gen_loss = -torch.mean(disc_fake)

lam = self.spectrogram_vq.calculate_lambda(rec_loss, gen_loss)

vq_loss = rec_loss + args.vq_loss_factor * q_loss + disc_factor * lam * gen_loss

else:

gen_loss = torch.tensor(0.0, device=self.device)

lam = torch.tensor(0.0, device=self.device)

vq_loss = rec_loss + args.vq_loss_factor * q_loss

self.opt_vq.zero_grad()

vq_loss.backward()

self.opt_vq.step()

# Train Discriminator

disc_real = self.discriminator(mels)

disc_fake = self.discriminator(mel_hat.detach())

d_loss_real = torch.mean(F.relu(1.0 - disc_real))

d_loss_fake = torch.mean(F.relu(1.0 + disc_fake))

gan_loss = disc_factor * 0.5 * (d_loss_real + d_loss_fake)

self.opt_disc.zero_grad()

gan_loss.backward()

self.opt_disc.step()

if i % args.log_interval == 0:

with torch.no_grad():

vis = torch.cat((mels[:4], (mel_hat[:4] + 1) * 0.5), dim=0)

save_spectrogram_to_img(vis, os.path.join("results", f"{epoch}_{i}.png"))

synthesize(self.vocoder, mel_hat[0], os.path.join("results", f"{epoch}_{i}.wav"))

pbar.set_postfix({

"VQ_Loss": f"{vq_loss.item():.5f}",

"GAN_Loss": f"{gan_loss.item():.5f}",

"rec": f"{rec_loss.item():.5f}",

"vq": f"{q_loss.item():.5f}",

"λ": f"{lam.item():.5f}",

"disc_w": f"{disc_factor:.2f}",

})

# Save model checkpoints