Make the tokenizer a pytorch module

Ease model weight loading.

Author: LTMeyer

aion/codecs/tokenizers/base.py

@@ -6,7 +6,7 @@
 from aion.codecs.quantizers import Quantizer
 
 
-class Codec(ABC):
+class Codec(ABC, torch.nn.Module):
     """Abstract definition of a Codec.
 
     A codec embeds a specific type of data into a sequence of either
@@ -53,6 +53,13 @@ def decode(
         """Encodes a given batch of samples into latent space."""
         return self._decode(z)
 
+    def forward(
+        self,
+        x: Float[torch.Tensor, " b c *input_shape"],
+        channel_mask: Bool[torch.Tensor, " b c"],
+    ) -> Float[torch.Tensor, " b c1 *code_shape"]:
+        return self.encode(x, channel_mask)
+
 
 class QuantizedCodec(Codec):
     def __init__(self, quantizer: Quantizer):

aion/codecs/tokenizers/image.py

@@ -114,7 +114,7 @@ def __init__(
             mult_factor: Multiplication factor.
         """
         # Get MagViT architecture
-        self.model = MagVitAE(
+        model = MagVitAE(
             n_bands=multisurvey_projection_dims,
             hidden_dims=hidden_dims,
             n_compressions=n_compressions,
@@ -123,11 +123,12 @@ def __init__(
         super().__init__(
             n_bands,
             quantizer,
-            self.model.encode,
-            self.model.decode,
+            model.encode,
+            model.decode,
             hidden_dims,
             embedding_dim,
             multisurvey_projection_dims,
             range_compression_factor,
             mult_factor,
         )
+        self.model = model