feedforward should be agnostic to images or video

lucidrains · Jun 19, 2023 · 57ed18b · 57ed18b
1 parent 875bd25
commit 57ed18b
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 2 deletions.
diff --git a/make_a_video_pytorch/make_a_video.py b/make_a_video_pytorch/make_a_video.py
@@ -84,12 +84,24 @@ def __init__(self, dim, mult = 4):
         )
 
     def forward(self, x, enable_time = True):
+
+        is_video = x.ndim == 5
+        enable_time &= is_video
+
+        if not is_video:
+            x = rearrange(x, 'b c h w -> b c 1 h w')
+
         x = self.proj_in(x)
 
         if enable_time:
             x = shift_token(x)
 
-        return self.proj_out(x)
+        out = self.proj_out(x)
+
+        if not is_video:
+            out = rearrange(out, 'b c 1 h w -> b c h w')
+
+        return out
 
 # best relative positional encoding
 

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'make-a-video-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.2.0',
+  version = '0.2.1',
   license='MIT',
   description = 'Make-A-Video - Pytorch',
   author = 'Phil Wang',