code: def __init__(
      self,
      positional_dims=2,
      num_tokens=512,
      codebook_dim=512,
      num_layers=3,
      num_resnet_blocks=0,
      hidden_dim=64,
      channels=3,
      stride=2,
      kernel_size=4,
      use_transposed_convs=True,
      encoder_norm=False,
      activation="relu",
      smooth_l1_loss=False,
      straight_through=False,
      normalization=None,  # ((0.5,) * 3, (0.5,) * 3),
      record_codes=False,
      discretization_loss_averaging_steps=100,
      lr_quantizer_args={},
  ):
      super().__init__()
      has_resblocks = num_resnet_blocks > 0
  
      self.num_tokens = num_tokens
      self.num_layers = num_layers
      self.straight_through = straight_through
      self.positional_dims = positional_dims
      self.discrete_loss = DiscretizationLoss(
          num_tokens,
          2,
          1 / (num_tokens * 2),
          discretization_loss_averaging_steps,
      )
  
      assert (
          positional_dims > 0 and positional_dims < 3
      )  # This VAE only supports 1d and 2d inputs for now.
      if positional_dims == 2:
          conv = nn.Conv2d
          conv_transpose = nn.ConvTranspose2d
      else:
          conv = nn.Conv1d
          conv_transpose = nn.ConvTranspose1d
      if not use_transposed_convs:
          conv_transpose = functools.partial(UpsampledConv, conv)
  
      if activation == "relu":
          act = nn.ReLU
      elif activation == "silu":
          act = nn.SiLU
      else:
          assert NotImplementedError()
  
      enc_layers = []
      dec_layers = []
  
      if num_layers > 0:
          enc_chans = [hidden_dim * 2**i for i in range(num_layers)]
          dec_chans = list(reversed(enc_chans))
  
          enc_chans = [channels, *enc_chans]
  
          dec_init_chan = (
              codebook_dim if not has_resblocks else dec_chans[0]
          )
          dec_chans = [dec_init_chan, *dec_chans]
  
          enc_chans_io, dec_chans_io = map(
              lambda t: list(zip(t[:-1], t[1:])), (enc_chans, dec_chans)
          )
  
          pad = (kernel_size - 1) // 2
          for (enc_in, enc_out), (dec_in, dec_out) in zip(
              enc_chans_io, dec_chans_io
          ):
              enc_layers.append(
                  nn.Sequential(
                      conv(
                          enc_in,
                          enc_out,
                          kernel_size,
                          stride=stride,
                          padding=pad,
                      ),
                      act(),
                  )
              )
              if encoder_norm:
                  enc_layers.append(nn.GroupNorm(8, enc_out))
              dec_layers.append(
                  nn.Sequential(
                      conv_transpose(
                          dec_in,
                          dec_out,
                          kernel_size,
                          stride=stride,
                          padding=pad,
                      ),
                      act(),
                  )
              )
          dec_out_chans = dec_chans[-1]
          innermost_dim = dec_chans[0]
      else:
          enc_layers.append(
              nn.Sequential(conv(channels, hidden_dim, 1), act())
          )
          dec_out_chans = hidden_dim
          innermost_dim = hidden_dim
  
      for _ in range(num_resnet_blocks):
          dec_layers.insert(0, ResBlock(innermost_dim, conv, act))
          enc_layers.append(ResBlock(innermost_dim, conv, act))
  
      if num_resnet_blocks > 0:
          dec_layers.insert(0, conv(codebook_dim, innermost_dim, 1))
  
      enc_layers.append(conv(innermost_dim, codebook_dim, 1))
      dec_layers.append(conv(dec_out_chans, channels, 1))
  
      self.encoder = nn.Sequential(*enc_layers)
      self.decoder = nn.Sequential(*dec_layers)
  
      self.loss_fn = F.smooth_l1_loss if smooth_l1_loss else F.mse_loss
      self.codebook = Quantize(
          codebook_dim, num_tokens, new_return_order=True
      )
  
      # take care of normalization within class
      self.normalization = normalization
      self.record_codes = record_codes
      if record_codes:
          self.codes = torch.zeros((1228800,), dtype=torch.long)
          self.code_ind = 0
          self.total_codes = 0
      self.internal_step = 0
  
imports: import functools
  from math import sqrt
  import torch
  import torch.distributed
  import torch.nn
  import torch.nn.functional
  import torchaudio
  from einops import rearrange
label_desc: Use image node_id 1 to create a new DiscreteVAE object from inherited base classes: nn with example: obj = DiscreteVAE(positional_dims, num_tokens, codebook_dim, num_layers, num_resnet_blocks, hidden_dim, channels, stride, kernel_size, use_transposed_convs, encoder_norm, activation, smooth_l1_loss, straight_through, normalization, record_codes, discretization_loss_averaging_steps, lr_quantizer_args)