remove einops exts

imagen_pytorch/imagen_video.py

@@ -11,7 +11,6 @@
 
 from einops import rearrange, repeat, pack, unpack
 from einops.layers.torch import Rearrange
-from einops_exts.torch import EinopsToAndFrom
 
 from imagen_pytorch.t5 import t5_encode_text, get_encoded_dim, DEFAULT_T5_NAME
 
@@ -1501,7 +1500,7 @@ def __init__(
         mid_dim = dims[-1]
 
         self.mid_block1 = ResnetBlock(mid_dim, mid_dim, cond_dim = cond_dim, time_cond_dim = time_cond_dim, groups = resnet_groups[-1])
-        self.mid_attn = EinopsToAndFrom('b c f h w', 'b (f h w) c', Residual(Attention(mid_dim, **attn_kwargs))) if attend_at_middle else None
+        self.mid_attn = Residual(Attention(mid_dim, **attn_kwargs)) if attend_at_middle else None
         self.mid_temporal_peg = temporal_peg(mid_dim)
         self.mid_temporal_attn = temporal_attn(mid_dim)
         self.mid_block2 = ResnetBlock(mid_dim, mid_dim, cond_dim = cond_dim, time_cond_dim = time_cond_dim, groups = resnet_groups[-1])
@@ -1872,8 +1871,14 @@ def forward(
         x = self.mid_block1(x, t, c, **conv_kwargs)
 
         if exists(self.mid_attn):
+            x = rearrange(x, 'b c f h w -> b f h w c')
+            x, ps = pack([x], 'b * c')
+
             x = self.mid_attn(x)
 
+            x, = unpack(x, ps, 'b * c')
+            x = rearrange(x, 'b f h w c -> b c f h w')
+
         if not ignore_time:
             x = self.mid_temporal_peg(x)
             x = self.mid_temporal_attn(x)

imagen_pytorch/version.py

@@ -1 +1 @@
-__version__ = '1.25.12'
+__version__ = '1.26.0'

setup.py

@@ -37,14 +37,13 @@
     'numpy',
     'packaging',
     'pillow',
-    'pydantic<=2',
+    'pydantic>=2',
     'pytorch-warmup',
     'sentencepiece',
     'torch>=1.6',
     'torchvision',
     'transformers',
-    'tqdm',
-    'einops_exts'
+    'tqdm'
   ],
   classifiers=[
     'Development Status :: 4 - Beta',