Update iris/src/tokenizer.py

iris/src/tokenizer.py

@@ -12,15 +12,16 @@ from .models.lpips import LPIPS
 from .models.nets import Encoder, Decoder
 
 class Tokenizer(nn.Module):
-    def __init__(self, vocab_size: int, embed_dim: int, encoder: Encoder, decoder: Decoder, with_lpips: bool = True) -> None:
+    def __init__(self, config: dict, with_lpips: bool = True) -> None:
         super().__init__()
-        self.vocab_size = vocab_size
-        self.encoder = encoder
-        self.pre_quant_conv = torch.nn.Conv2d(encoder.config.z_channels, embed_dim, 1)
-        self.embedding = nn.Embedding(vocab_size, embed_dim)
-        self.post_quant_conv = torch.nn.Conv2d(embed_dim, decoder.config.z_channels, 1)
-        self.decoder = decoder
-        self.embedding.weight.data.uniform_(-1.0 / vocab_size, 1.0 / vocab_size)
+        self.vocab_size = config["vocab_size"]
+        self.embed_dim = config["embed_dim"]
+        self.encoder = Encoder(config)
+        self.pre_quant_conv = torch.nn.Conv2d(encoder.config.z_channels, self.embed_dim, 1)
+        self.embedding = nn.Embedding(self.vocab_size, self.embed_dim)
+        self.post_quant_conv = torch.nn.Conv2d(self.embed_dim, decoder.config.z_channels, 1)
+        self.decoder = Decoder(config)
+        self.embedding.weight.data.uniform_(-1.0 / self.vocab_size, 1.0 / self.vocab_size)
         self.lpips = LPIPS().eval() if with_lpips else None
 
     def __repr__(self) -> str: