missing the post vq temporal conv during inference

lucidrains · Jan 7, 2025 · 2da12f0 · 2da12f0
1 parent e5238a1
commit 2da12f0
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/genie2_pytorch/genie2.py b/genie2_pytorch/genie2.py
@@ -429,6 +429,12 @@ def generate(
         if self.latent_channel_first:
             tokens = rearrange(tokens, 'b ... d -> b d ...')
 
+        # maybe post vq temporal conv
+
+        tokens = self.post_vq_transform(tokens)
+
+        # maybe fold time into batch for certain decoders
+
         need_fold_time_into_batch = not self.is_video_enc_dec
 
         if need_fold_time_into_batch:

diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "genie2-pytorch"
-version = "0.0.20"
+version = "0.0.21"
 description = "Genie2"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }