able to customize attention heads and dimension per head

Author: lucidrains

denoising_diffusion_pytorch/classifier_free_guidance.py

@@ -263,6 +263,8 @@ def __init__(
         learned_sinusoidal_cond = False,
         random_fourier_features = False,
         learned_sinusoidal_dim = 16,
+        attn_dim_head = 32,
+        attn_heads = 4
     ):
         super().__init__()
 
@@ -334,7 +336,7 @@ def __init__(
 
         mid_dim = dims[-1]
         self.mid_block1 = block_klass(mid_dim, mid_dim, time_emb_dim = time_dim, classes_emb_dim = classes_dim)
-        self.mid_attn = Residual(PreNorm(mid_dim, Attention(mid_dim)))
+        self.mid_attn = Residual(PreNorm(mid_dim, Attention(mid_dim, dim_head = attn_dim_head, heads = attn_heads)))
         self.mid_block2 = block_klass(mid_dim, mid_dim, time_emb_dim = time_dim, classes_emb_dim = classes_dim)
 
         for ind, (dim_in, dim_out) in enumerate(reversed(in_out)):

denoising_diffusion_pytorch/denoising_diffusion_pytorch.py

@@ -271,6 +271,8 @@ def __init__(
         learned_sinusoidal_cond = False,
         random_fourier_features = False,
         learned_sinusoidal_dim = 16,
+        attn_dim_head = 32,
+        attn_heads = 4,
         full_attn = (False, False, False, True),
         flash_attn = False
     ):
@@ -331,7 +333,7 @@ def __init__(
             self.downs.append(nn.ModuleList([
                 block_klass(dim_in, dim_in, time_emb_dim = time_dim),
                 block_klass(dim_in, dim_in, time_emb_dim = time_dim),
-                attn_klass(dim_in),
+                attn_klass(dim_in, dim_head = attn_dim_head, heads = attn_heads),
                 Downsample(dim_in, dim_out) if not is_last else nn.Conv2d(dim_in, dim_out, 3, padding = 1)
             ]))
 
@@ -348,7 +350,7 @@ def __init__(
             self.ups.append(nn.ModuleList([
                 block_klass(dim_out + dim_in, dim_out, time_emb_dim = time_dim),
                 block_klass(dim_out + dim_in, dim_out, time_emb_dim = time_dim),
-                attn_klass(dim_out),
+                attn_klass(dim_out, dim_head = attn_dim_head, heads = attn_heads),
                 Upsample(dim_out, dim_in) if not is_last else  nn.Conv2d(dim_out, dim_in, 3, padding = 1)
             ]))
 

denoising_diffusion_pytorch/denoising_diffusion_pytorch_1d.py

@@ -265,7 +265,9 @@ def __init__(
         learned_variance = False,
         learned_sinusoidal_cond = False,
         random_fourier_features = False,
-        learned_sinusoidal_dim = 16
+        learned_sinusoidal_dim = 16,
+        attn_dim_head = 32,
+        attn_heads = 4
     ):
         super().__init__()
 
@@ -321,7 +323,7 @@ def __init__(
 
         mid_dim = dims[-1]
         self.mid_block1 = block_klass(mid_dim, mid_dim, time_emb_dim = time_dim)
-        self.mid_attn = Residual(PreNorm(mid_dim, Attention(mid_dim)))
+        self.mid_attn = Residual(PreNorm(mid_dim, Attention(mid_dim, dim_head = attn_dim_head, heads = attn_heads)))
         self.mid_block2 = block_klass(mid_dim, mid_dim, time_emb_dim = time_dim)
 
         for ind, (dim_in, dim_out) in enumerate(reversed(in_out)):

denoising_diffusion_pytorch/version.py

@@ -1 +1 @@
-__version__ = '1.8.5'
+__version__ = '1.8.6'