class_docstr: VQ-VAE model with XCiT attention.
  
  This model follows the overall architecture of the autoencoder from the
  VQ-GAN paper:
  
  Taming Transformers for High-Resolution Image Synthesis
  Patrick Esser, Robin Rombach, Bjorn Ommer
  
  However, there is one key modification: the use of XCiT attention as
  opposed to standard attention. This allows the resulting autoencoders to
  easily handle very high-resolution images that are used in compression.
  
  XCiT attention was introduced in the following paper:
  
  XCiT: Cross-Covariance Image Transformers
  A El-Nouby, H Touvron, M Caron, P Bojanowski, M Douze, A Joulin, I Laptev,
  N Neverova, G Synnaeve, J Verbeek, H Jegou
  
  This class does not implement the full VQ-VAE, just the encoder/decoder.
  The user is expected to pass in their own bottleneck and quantization
  operations via the `bottleneck_op` functional parameter.
  
  Args:
      ch: Number of top-level convolution channels.
      out_ch: Number of output channels.
      ch_mult: A multiplication factor for scaling the number of channels.
      num_res_blocks: Number of residual blocks for convolution layers.
      attn_resolution: The resolutions at which to apply attention.
      dropout: Dropout probability.
      resamp_with_conv: Whether to resample in upsampling layers with
          convolutions.
      in_channels: Number of input channels.
      resolution: Overall image resolution. Used for estimating size of
          latent on model instantiation.
      embed_dim: Size of embedding dimension.
      freeze_encoder: Whether to freeze the encoder.
      freeze_bottleneck: Whether to freeze the bottleneck (codebook).
      norm_type: Type of normalization.
class_docstr_tok: ["VQ-VAE","model","with","XCiT","attention",".","This","model","follows","the","overall","architecture","of","the","autoencoder","from","the","VQ-GAN","paper",":","Taming","Transformers","for","High-Resolution","Image","Synthesis","Patrick","Esser",",","Robin","Rombach",",","Bjorn","Ommer","However",",","there","is","one","key","modification",":","the","use","of","XCiT","attention","as","opposed","to","standard","attention",".","This","allows","the","resulting","autoencoders","to","easily","handle","very","high-resolution","images","that","are","used","in","compression",".","XCiT","attention","was","introduced","in","the","following","paper",":","XCiT",":","Cross-Covariance","Image","Transformers","A","El-Nouby",",","H","Touvron",",","M","Caron",",","P","Bojanowski",",","M","Douze",",","A","Joulin",",","I","Laptev",",","N","Neverova",",","G","Synnaeve",",","J","Verbeek",",","H","Jegou","This","class","does","not","implement","the","full","VQ-VAE",",","just","the","encoder\/decoder",".","The","user","is","expected","to","pass","in","their","own","bottleneck","and","quantization","operations","via","the","`","bottleneck_op","`","functional","parameter",".","Args",":","ch",":","Number","of","top-level","convolution","channels",".","out_ch",":","Number","of","output","channels",".","ch_mult",":","A","multiplication","factor","for","scaling","the","number","of","channels",".","num_res_blocks",":","Number","of","residual","blocks","for","convolution","layers",".","attn_resolution",":","The","resolutions","at","which","to","apply","attention",".","dropout",":","Dropout","probability",".","resamp_with_conv",":","Whether","to","resample","in","upsampling","layers","with","convolutions",".","in_channels",":","Number","of","input","channels",".","resolution",":","Overall","image","resolution",".","Used","for","estimating","size","of","latent","on","model","instantiation",".","embed_dim",":","Size","of","embedding","dimension",".","freeze_encoder",":","Whether","to","freeze","the","encoder",".","freeze_bottleneck",":","Whether","to","freeze","the","bottleneck","(","codebook",")",".","norm_type",":","Type","of","normalization","."]
code: def __init__(
      self,
      ch: int = 128,
      out_ch: int = 3,
      ch_mult: Optional[Sequence] = None,
      num_res_blocks: int = 2,
      attn_resolutions: Optional[Sequence] = None,
      dropout: float = 0.0,
      resamp_with_conv: bool = True,
      in_channels: int = 3,
      resolution: int = 256,
      embed_dim: int = 256,
      bottleneck_op: Optional[nn.Module] = None,
      freeze_encoder: bool = False,
      freeze_bottleneck: bool = False,
      norm_type: str = "group",
  ):
      super().__init__()
      self.logger = logging.getLogger(self.__class__.__name__)
      if ch_mult is None:
          ch_mult = (1, 1, 2, 2, 4)
      if attn_resolutions is None:
          attn_resolutions = [16]
  
      self.freeze_encoder = freeze_encoder
      self.freeze_bottleneck = freeze_bottleneck
      self.factor = 2 ** (len(ch_mult) - 1)
  
      self.encoder = _VqVaeXCiTConvEncoder(
          ch=ch,
          ch_mult=ch_mult,
          num_res_blocks=num_res_blocks,
          attn_resolutions=attn_resolutions,
          dropout=dropout,
          resamp_with_conv=resamp_with_conv,
          in_channels=in_channels,
          resolution=resolution,
          embed_dim=embed_dim,
          norm_type=norm_type,
      )
      self.decoder = _VqvaeXCiTConvDecoder(
          ch=ch,
          out_ch=out_ch,
          ch_mult=ch_mult,
          num_res_blocks=num_res_blocks,
          attn_resolutions=attn_resolutions,
          dropout=dropout,
          resamp_with_conv=resamp_with_conv,
          in_channels=in_channels,
          resolution=resolution,
          embed_dim=embed_dim,
          norm_type=norm_type,
      )
      self.bottleneck_op = bottleneck_op
  
      if self.freeze_encoder:
          self.logger.info("Freezing encoder!")
          self.encoder.eval()
          for param in self.encoder.parameters():
              param.requires_grad_(False)
  
      if self.freeze_bottleneck:
          self.logger.info("Freezing bottleneck!")
          if self.bottleneck_op is not None:
              self.bottleneck_op.eval()
              for param in self.bottleneck_op.parameters():
                  param.requires_grad_(False)
  
imports: import logging
  from typing import Optional, Sequence
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from torch import Tensor
  from neuralcompression.layers import XCA, ChannelNorm2D
  from .._outputs._vqvae_autoencoder_output import VqVaeAutoencoderOutput
label_desc: Use image node_id 1 to create a new VqVaeXCiTAutoencoder object from inherited base classes: nn with example: obj = VqVaeXCiTAutoencoder(ch, out_ch, ch_mult, num_res_blocks, attn_resolutions, dropout, resamp_with_conv, in_channels, resolution, embed_dim, bottleneck_op, freeze_encoder, freeze_bottleneck, norm_type)