cont: remove all training related code

Author: cnoco1at3

configs/open_model.yaml

@@ -17,7 +17,6 @@ shape_model:
   encoder_with_cls_token: true
   num_encoder_latents: 512 
   num_decoder_latents: 0 
-  encoder_dynamic_latents: true
   embed_dim: 32
   width: 768
   num_heads: 12
@@ -31,9 +30,5 @@ shape_model:
   num_decoder_layers: 24
   dropout: 0.0
   num_codes: 16384
-  vq_beta: 0.8
-  sync_nu: 0.1
-  replace_freq: 0
-  skip_p: 0.5
 
 text_model_pretrained_model_name_or_path: "openai/clip-vit-large-patch14"

src/cube/model/autoencoder/one_d_autoencoder.py

@@ -15,7 +15,6 @@
     generate_dense_grid_points,
     marching_cubes_with_warp,
 )
-from cube.model.autoencoder.sinkhorn_vq import SinkhornVectorQuantizer
 from cube.model.autoencoder.spherical_vq import SphericalVectorQuantizer
 from cube.model.transformers.attention import (
     EncoderCrossAttentionLayer,
@@ -265,6 +264,7 @@ class Config:
         num_decoder_layers: int = 23
 
         encoder_with_cls_token: bool = True
+        num_codes: int = 16384
         dropout: float = 0.0
 
     def __init__(self, cfg: Config) -> None:
@@ -289,13 +289,11 @@ def __init__(self, cfg: Config) -> None:
             eps=self.cfg.eps,
         )
 
-        block = SinkhornVectorQuantizer(
+        block = SphericalVectorQuantizer(
             self.cfg.embed_dim,
             self.cfg.num_codes,
             self.cfg.width,
-            self.cfg.vq_beta,
-            self.cfg.sync_nu,
-            self.cfg.skip_p,
+            codebook_regularization="kl",
         )
         self.bottleneck = OneDBottleNeck(block=block)
 

src/cube/model/autoencoder/sinkhorn_vq.py

@@ -13,8 +13,6 @@ def __init__(
         embed_dim: int,
         num_codes: int,
         width: Optional[int] = None,
-        beta: float = 0.8,
-        sync_nu: float = 0.0,
         codebook_regularization: Literal["batch_norm", "kl"] = "batch_norm",
     ):
         """
@@ -23,14 +21,12 @@ def __init__(
             embed_dim (int): The dimensionality of the embeddings.
             num_codes (int): The number of codes in the codebook.
             width (Optional[int], optional): The width of the input. Defaults to None.
-            sync_nu (float, optional): Synchronization parameter. Defaults to 0.0.
         Raises:
             ValueError: If beta is not in the range [0, 1].
         """
         super().__init__()
 
         self.num_codes = num_codes
-        self.sync_nu = sync_nu
 
         self.codebook = nn.Embedding(num_codes, embed_dim)
         self.codebook.weight.data.uniform_(-1.0 / num_codes, 1.0 / num_codes)