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 log_illustration(
      self,
      target_spectrograms,
      input_spectrograms,
      enhanced_spectrograms,
      lengths,
  ):
      if self.global_rank != 0:
          return
  
      if not self.loggers:
          return
  
      step = self.trainer.global_step // 2  # because of G/D training
      if step % self.trainer.log_every_n_steps != 0:
          return
  
      idx = 0
      length = int(lengths.flatten()[idx].item())
      tensor = torch.stack(
          [
              enhanced_spectrograms - input_spectrograms,
              input_spectrograms,
              enhanced_spectrograms,
              target_spectrograms,
          ],
          dim=0,
      ).cpu()[:, idx, :, :, :length]
  
      assert (
          TORCHVISION_AVAILABLE
      ), "Torchvision imports failed but they are required."
      grid = torchvision.utils.make_grid(tensor, nrow=1).clamp(0.0, 1.0)
  
      for logger in self.loggers:
          if isinstance(logger, TensorBoardLogger):
              writer: SummaryWriter = logger.experiment
              writer.add_image("spectrograms", grid, global_step=step)
              writer.flush()
          elif isinstance(logger, WandbLogger):
              logger.log_image(
                  "spectrograms",
                  [grid],
                  caption=["residual, input, output, ground truth"],
                  step=step,
              )
          else:
              logging.warning(
                  "Unsupported logger type: %s", str(type(logger))
              )
  
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 15 for calling the SpectrogramEnhancerModel obj's underlying member method code with example usage: obj.log_illustration(target_spectrograms, input_spectrograms, enhanced_spectrograms, lengths) without return types