at least two attention heads in karras unet

Author: lucidrains

denoising_diffusion_pytorch/karras_unet.py

@@ -240,7 +240,7 @@ def __init__(
         if has_attn:
             self.attn = Attention(
                 dim = dim_out,
-                heads = ceil(dim_out / attn_dim_head),
+                heads = max(ceil(dim_out / attn_dim_head), 2),
                 dim_head = attn_dim_head,
                 mp_add_t = attn_res_mp_add_t,
                 flash = attn_flash
@@ -322,7 +322,7 @@ def __init__(
         if has_attn:
             self.attn = Attention(
                 dim = dim_out,
-                heads = ceil(dim_out / attn_dim_head),
+                heads = max(ceil(dim_out / attn_dim_head), 2),
                 dim_head = attn_dim_head,
                 mp_add_t = attn_res_mp_add_t,
                 flash = attn_flash

denoising_diffusion_pytorch/version.py

@@ -1 +1 @@
-__version__ = '1.10.2'
+__version__ = '1.10.3'