code: def __init__(
      self,
      img_size,
      nf,
      ch_mult,
      quantizer="nearest",
      res_blocks=2,
      attn_resolutions=None,
      codebook_size=1024,
      emb_dim=256,
      beta=0.25,
      gumbel_straight_through=False,
      gumbel_kl_weight=1e-8,
      model_path=None,
  ):
      super().__init__()
      logger = get_root_logger()
      self.in_channels = 3
      self.nf = nf
      self.n_blocks = res_blocks
      self.codebook_size = codebook_size
      self.embed_dim = emb_dim
      self.ch_mult = ch_mult
      self.resolution = img_size
      self.attn_resolutions = attn_resolutions or [16]
      self.quantizer_type = quantizer
      self.encoder = Encoder(
          self.in_channels,
          self.nf,
          self.embed_dim,
          self.ch_mult,
          self.n_blocks,
          self.resolution,
          self.attn_resolutions,
      )
      if self.quantizer_type == "nearest":
          self.beta = beta  # 0.25
          self.quantize = VectorQuantizer(
              self.codebook_size, self.embed_dim, self.beta
          )
      elif self.quantizer_type == "gumbel":
          self.gumbel_num_hiddens = emb_dim
          self.straight_through = gumbel_straight_through
          self.kl_weight = gumbel_kl_weight
          self.quantize = GumbelQuantizer(
              self.codebook_size,
              self.embed_dim,
              self.gumbel_num_hiddens,
              self.straight_through,
              self.kl_weight,
          )
      self.generator = Generator(
          self.nf,
          self.embed_dim,
          self.ch_mult,
          self.n_blocks,
          self.resolution,
          self.attn_resolutions,
      )
  
      if model_path is not None:
          chkpt = torch.load(model_path, map_location="cpu")
          if "params_ema" in chkpt:
              self.load_state_dict(
                  torch.load(model_path, map_location="cpu")[
                      "params_ema"
                  ]
              )
              logger.info(
                  f"vqgan is loaded from: {model_path} [params_ema]"
              )
          elif "params" in chkpt:
              self.load_state_dict(
                  torch.load(model_path, map_location="cpu")["params"]
              )
              logger.info(
                  f"vqgan is loaded from: {model_path} [params]"
              )
          else:
              raise ValueError("Wrong params!")
  
imports: import torch
  import torch.nn
  import torch.nn.functional
  from basicsr.utils import get_root_logger
  from basicsr.utils.registry import ARCH_REGISTRY
label_desc: Use image node_id 1 to create a new VQAutoEncoder object from inherited base classes: nn with example: obj = VQAutoEncoder(img_size, nf, ch_mult, quantizer, res_blocks, attn_resolutions, codebook_size, emb_dim, beta, gumbel_straight_through, gumbel_kl_weight, model_path)