class_docstr: Improved version over VectorQuantizer, can be used as a drop-in replacement. Mostly
  avoids costly matrix multiplications and allows for post-hoc remapping of indices.
class_docstr_tok: ["Improved","version","over","VectorQuantizer",",","can","be","used","as","a","drop-in","replacement",".","Mostly","avoids","costly","matrix","multiplications","and","allows","for","post-hoc","remapping","of","indices","."]
code: def __init__(
      self,
      n_e,
      e_dim,
      beta,
      remap=None,
      unknown_index="random",
      sane_index_shape=False,
      legacy=True,
  ):
      super().__init__()
      self.n_e = n_e
      self.e_dim = e_dim
      self.beta = beta
      self.legacy = legacy
  
      self.embedding = nn.Embedding(self.n_e, self.e_dim)
      self.embedding.weight.data.uniform_(
          -1.0 / self.n_e, 1.0 / self.n_e
      )
  
      self.remap = remap
      if self.remap is not None:
          self.register_buffer(
              "used", torch.tensor(np.load(self.remap))
          )
          self.re_embed = self.used.shape[0]
          self.unknown_index = (
              unknown_index  # "random" or "extra" or integer
          )
          if self.unknown_index == "extra":
              self.unknown_index = self.re_embed
              self.re_embed = self.re_embed + 1
          print(
              f"Remapping {self.n_e} indices to {self.re_embed} indices. "
              f"Using {self.unknown_index} for unknown indices."
          )
      else:
          self.re_embed = n_e
  
      self.sane_index_shape = sane_index_shape
  
file_path: stable-diffusion-webui/extensions-builtin/LDSR/vqvae_quantize.py
imports: import torch
  import torch.nn
  import numpy
  from einops import rearrange
label_desc: Use image node_id 1 to create a new VectorQuantizer2 object from inherited base classes: nn with example: obj = VectorQuantizer2(n_e, e_dim, beta, remap, unknown_index, sane_index_shape, legacy)