reformat

kappamodules/attention/linear_attention.py

@@ -93,4 +93,4 @@ def to_channel_last(self, x):
 
     def to_channel_first(self, x, og_shape):
         _, _, h, w, d = og_shape
-        return einops.rearrange(x, "bs (h w d) dim -> bs dim h w d", h=h, w=w, d=d)
+        return einops.rearrange(x, "bs (h w d) dim -> bs dim h w d", h=h, w=w, d=d)

kappamodules/attention/perceiver_attention.py

@@ -9,6 +9,7 @@
     init_truncnormal_zero_bias,
 )
 
+
 class PerceiverAttention(nn.Module):
     def __init__(self, dim, num_heads=8, bias=True, concat_query_to_kv=False, init_weights="truncnormal"):
         super().__init__()
@@ -65,4 +66,4 @@ def forward(self, q, kv, attn_mask=None):
 
 
 class PerceiverAttention1d(PerceiverAttention):
-    pass
+    pass

kappamodules/init/functional.py

@@ -42,6 +42,7 @@ def init_norm_as_noaffine(m):
         if m.weight is not None:
             nn.init.constant_(m.weight, 1.)
 
+
 # LEGACY remove
 def init_norms_as_noaffine(m):
     if isinstance(m, ALL_NORMS):
@@ -74,6 +75,7 @@ def init_norm_as_identity(m):
         if m.weight is not None:
             nn.init.constant_(m.weight, 0.)
 
+
 # LEGACY remove
 def init_norms_as_identity(m):
     if isinstance(m, ALL_NORMS):

kappamodules/layers/__init__.py

@@ -3,7 +3,7 @@
 from .identity import Identity
 from .normalize import Normalize
 from .paramless_batchnorm import ParamlessBatchNorm1d
-from .residual import Residual
-from .weight_norm_linear import WeightNormLinear
 from .regular_grid_sincos_embed import RegularGridSincosEmbed
+from .residual import Residual
 from .rms_norm import RMSNorm
+from .weight_norm_linear import WeightNormLinear

kappamodules/modulation/__init__.py

@@ -1,3 +1,3 @@
 from .dit import Dit
 from .film import Film
-from .timestep_embed import TimestepEmbed
+from .timestep_embed import TimestepEmbed

kappamodules/modulation/film.py

@@ -1,10 +1,8 @@
-import math
-
 from torch import nn
 
+from kappamodules.init import init_xavier_uniform_merged_linear
 from kappamodules.utils.shapes import to_ndim
 
-from kappamodules.init import init_xavier_uniform_merged_linear
 
 class Film(nn.Module):
     def __init__(self, dim_cond, dim_out, init_weights="xavier_uniform"):

kappamodules/transformer/__init__.py

@@ -1,2 +1,2 @@
 from .perceiver_block import PerceiverBlock
-from .postnorm_block import PostnormBlock
+from .postnorm_block import PostnormBlock

kappamodules/unet/__init__.py

@@ -1 +1 @@
-from .unet_denoising_diffusion import UnetDenoisingDiffusion
+from .unet_denoising_diffusion import UnetDenoisingDiffusion

kappamodules/vit/vit_batchnorm.py

@@ -1,4 +1,3 @@
-import einops
 from torch import nn
 
 

kappamodules/vit/vit_pos_embed.py

@@ -34,6 +34,7 @@ def forward(self, x):
             embed = self.embed
         return x + embed
 
+
 # LEGACY remove
 class VitPosEmbedNd(VitPosEmbed):
     pass

original_modules/mae_pos_embed.py

@@ -1,5 +1,6 @@
 import numpy as np
 
+
 def get_2d_sincos_pos_embed(embed_dim, grid_size):
     if isinstance(grid_size, int):
         grid_size = (grid_size, grid_size)

original_modules/original_perceiver_attention.py

@@ -1,6 +1,7 @@
+import einops
 import torch
 from torch import nn
-import einops
+
 
 class OriginalPerceiverAttention(nn.Module):
     """
@@ -13,7 +14,7 @@ class OriginalPerceiverAttention(nn.Module):
 
     def __init__(self, dim, dim_head=64, heads=8):
         super().__init__()
-        self.scale = dim_head**-0.5
+        self.scale = dim_head ** -0.5
         self.heads = heads
         inner_dim = dim_head * heads
 
@@ -45,4 +46,4 @@ def forward(self, x, latents):
 
         out = torch.einsum("... i j, ... j d -> ... i d", attn, v)
         out = einops.rearrange(out, "b h tn d -> b tn (h d)", h=h)
-        return self.to_out(out)
+        return self.to_out(out)

tests_unit/attention/test_perceiver_attention.py

@@ -1,8 +1,10 @@
 import unittest
 
+import torch
+
 from kappamodules.attention import PerceiverAttention1d
 from original_modules.original_perceiver_attention import OriginalPerceiverAttention
-import torch
+
 
 class TestPerceiverAttention(unittest.TestCase):
     def test_shape(self):
@@ -51,4 +53,3 @@ def test_equal_to_original(self):
         self.assertEqual(q.shape, y_og.shape)
         self.assertEqual(y_kc.shape, y_og.shape)
         self.assertTrue(torch.allclose(y_kc, y_og))
-

tests_unit/functional/test_pos_embed.py

@@ -1,9 +1,12 @@
+import unittest
+
 import einops
 import torch
-import unittest
+
 from kappamodules.functional.pos_embed import get_sincos_pos_embed_from_seqlens
 from original_modules.mae_pos_embed import get_2d_sincos_pos_embed
 
+
 class TestPosEmbed(unittest.TestCase):
     def test_shapes(self):
         self.assertEqual((10, 5), get_sincos_pos_embed_from_seqlens(seqlens=(10,), dim=5).shape)

tests_unit/layers/test_continuous_sincos_embed.py

@@ -2,8 +2,10 @@
 
 import einops
 import torch
-from kappamodules.layers import ContinuousSincosEmbed
+
 from kappamodules.functional.pos_embed import get_sincos_pos_embed_from_seqlens
+from kappamodules.layers import ContinuousSincosEmbed
+
 
 class TestContinuousSincosEmbed(unittest.TestCase):
     def test_shape(self):

tests_unit/transformer/test_perceiver_block.py

@@ -1,7 +1,10 @@
 import unittest
+
 import torch
+
 from kappamodules.transformer import PerceiverBlock
 
+
 class TestPerceiverBlock(unittest.TestCase):
     def test_shape(self):
         dim = 8

tests_unit/unet/test_unet_denoising_diffusion.py

@@ -1,8 +1,10 @@
-import torch
 import unittest
 
+import torch
+
 from kappamodules.unet import UnetDenoisingDiffusion
 
+
 class TestUnetDenoisingDiffusion(unittest.TestCase):
     def test_1d_uncond(self):
         torch.manual_seed(9823)
@@ -32,4 +34,4 @@ def test_1d_cond(self):
         cond = torch.randn(2, 4)
         y = model(x, cond=cond)
         self.assertEqual(x.shape, y.shape)
-        self.assertTrue(torch.isclose(y.mean(), torch.tensor(0.54554283618927)))
+        self.assertTrue(torch.isclose(y.mean(), torch.tensor(0.54554283618927)))

tests_unit/vit/test_dit_block.py

@@ -5,6 +5,7 @@
 from kappamodules.vit import DitBlock
 from original_modules.original_dit_block import OriginalDitBlock
 
+
 class TestDitBlock(unittest.TestCase):
     def test_equal_to_original(self):
         dim = 12

tests_unit/vit/test_vit_batchnorm.py

@@ -1,8 +1,10 @@
 import unittest
 
-from kappamodules.vit import VitBatchNorm
 import torch
 
+from kappamodules.vit import VitBatchNorm
+
+
 class TestVitBatchNorm(unittest.TestCase):
     def test_0d(self):
         dim = 4
@@ -38,4 +40,4 @@ def test_3d(self):
         y = bn(x)
         self.assertEqual(x.shape, y.shape)
         self.assertTrue(torch.allclose(y.mean(dim=[0, 1, 2, 3]), torch.zeros(size=(dim,)), atol=1e-6))
-        self.assertTrue(torch.allclose(y.std(dim=[0, 1, 2, 3]), torch.ones(size=(dim,)), atol=1e-1))
+        self.assertTrue(torch.allclose(y.std(dim=[0, 1, 2, 3]), torch.ones(size=(dim,)), atol=1e-1))

tests_unit/vit/test_vit_block.py

@@ -4,10 +4,11 @@
 
 from kappamodules.vit import VitBlock
 
+
 class TestVitBlock(unittest.TestCase):
     def test(self):
         dim = 4
         block = VitBlock(dim=dim, num_heads=2)
         x = torch.randn(2, 6, dim, generator=torch.Generator().manual_seed(9834))
         y = block(x)
-        self.assertEqual(x.shape, y.shape)
+        self.assertEqual(x.shape, y.shape)

tests_unit/vit/test_vit_pos_embed.py

@@ -51,7 +51,6 @@ def test_3d_fixed(self):
     def test_3d_learnable(self):
         self._test3d(is_learnable=True)
 
-
     def test_interpolate_2d(self):
         seqlens = (8, 12)
         dim = 64
@@ -63,4 +62,3 @@ def test_interpolate_2d(self):
         x_half = torch.zeros(2, *[seqlen // 2 for seqlen in seqlens], dim)
         y_half = pos_embed(x_half)
         self.assertEqual(x_half.shape, y_half.shape)
-