class_docstr: GAN-based model to add details to blurry spectrograms from TTS models like Tacotron or FastPitch. Based on StyleGAN 2 [1]
  [1] Karras et. al. - Analyzing and Improving the Image Quality of StyleGAN (https://arxiv.org/abs/1912.04958)
class_docstr_tok: ["GAN-based","model","to","add","details","to","blurry","spectrograms","from","TTS","models","like","Tacotron","or","FastPitch",".","Based","on","StyleGAN","2","[","1","]","[","1","]","Karras","et",".","al",".","-","Analyzing","and","Improving","the","Image","Quality","of","StyleGAN","(","https",":","\/\/arxiv.org\/abs\/1912.04958",")"]
code: def training_step(self, batch, batch_idx, optimizer_idx):
      input_spectrograms, target_spectrograms, lengths = batch
  
      with torch.no_grad():
          input_spectrograms = self.normalize_spectrograms(
              input_spectrograms, lengths
          )
          target_spectrograms = self.normalize_spectrograms(
              target_spectrograms, lengths
          )
  
      # train discriminator
      if optimizer_idx == 0:
          enhanced_spectrograms = self.forward(
              input_spectrograms=input_spectrograms,
              lengths=lengths,
              mixing=True,
              normalize=False,
          )
          enhanced_spectrograms = rearrange(
              enhanced_spectrograms, "b c l -> b 1 c l"
          )
          fake_logits = self.discriminator(
              enhanced_spectrograms, input_spectrograms, lengths
          )
  
          target_spectrograms_ = rearrange(
              target_spectrograms, "b c l -> b 1 c l"
          ).requires_grad_()
          real_logits = self.discriminator(
              target_spectrograms_, input_spectrograms, lengths
          )
          d_loss = self.discriminator_loss(real_logits, fake_logits)
          self.log("d_loss", d_loss, prog_bar=True)
  
          if (
              batch_idx % self._cfg.gradient_penalty_loss_every_n_steps
              == 0
          ):
              gp_loss = self.gradient_penalty_loss(
                  target_spectrograms_, real_logits
              )
              self.log("d_loss_gp", gp_loss, prog_bar=True)
              return d_loss + gp_loss
  
          return d_loss
  
      # train generator
      if optimizer_idx == 1:
          enhanced_spectrograms = self.forward(
              input_spectrograms=input_spectrograms,
              lengths=lengths,
              mixing=True,
              normalize=False,
          )
  
          input_spectrograms = rearrange(
              input_spectrograms, "b c l -> b 1 c l"
          )
          enhanced_spectrograms = rearrange(
              enhanced_spectrograms, "b c l -> b 1 c l"
          )
  
          fake_logits = self.discriminator(
              enhanced_spectrograms, input_spectrograms, lengths
          )
          g_loss = self.generator_loss(fake_logits)
          c_loss = self.consistency_loss(
              input_spectrograms, enhanced_spectrograms, lengths
          )
  
          self.log("g_loss", g_loss, prog_bar=True)
          self.log("c_loss", c_loss, prog_bar=True)
  
          with torch.no_grad():
              target_spectrograms = rearrange(
                  target_spectrograms, "b c l -> b 1 c l"
              )
              self.log_illustration(
                  target_spectrograms,
                  input_spectrograms,
                  enhanced_spectrograms,
                  lengths,
              )
          return g_loss + c_loss
  
imports: from random import random, randrange
  from typing import List, Optional
  import torch
  import torch.nn.functional
  from einops import rearrange
  from hydra.utils import instantiate
  from omegaconf import DictConfig
  from pytorch_lightning import Trainer
  from pytorch_lightning.loggers import TensorBoardLogger, WandbLogger
  from torch.utils.tensorboard.writer import SummaryWriter
  from nemo.collections.tts.losses.spectrogram_enhancer_losses import ConsistencyLoss, GeneratorLoss, GradientPenaltyLoss, HingeLoss
  from nemo.collections.tts.parts.utils.helpers import mask_sequence_tensor, to_device_recursive
  from nemo.core import Exportable, ModelPT, PretrainedModelInfo, typecheck
  from nemo.core.neural_types import LengthsType, MelSpectrogramType, NeuralType
  from nemo.core.neural_types.elements import BoolType
  from nemo.utils import logging
label_desc: Use image node_id 10 for calling the SpectrogramEnhancerModel obj's underlying member method code with example usage: obj.training_step(batch, batch_idx, optimizer_idx) and returns: d_loss, unknown, unknown