flaxdiff 0.1.6__py3-none-any.whl → 0.1.8__py3-none-any.whl

flaxdiff/models/attention.py

@@ -162,65 +162,6 @@ class NormalAttention(nn.Module):
         proj = proj.reshape(orig_x_shape)
         return proj
 
-class BasicTransformerBlock(nn.Module):
-    # Has self and cross attention
-    query_dim: int
-    heads: int = 4
-    dim_head: int = 64
-    dtype: Optional[Dtype] = None
-    precision: PrecisionLike = None
-    use_bias: bool = True
-    kernel_init: Callable = lambda : kernel_init(1.0)
-    use_flash_attention:bool = False
-    use_cross_only:bool = False
-    
-    def setup(self):
-        if self.use_flash_attention:
-            attenBlock = EfficientAttention
-        else:
-            attenBlock = NormalAttention
-            
-        self.attention1 = attenBlock(
-         query_dim=self.query_dim,
-            heads=self.heads,
-            dim_head=self.dim_head,
-            name=f'Attention1',
-            precision=self.precision,
-            use_bias=self.use_bias,
-            dtype=self.dtype,
-            kernel_init=self.kernel_init
-        )
-        self.attention2 = attenBlock(
-            query_dim=self.query_dim,
-            heads=self.heads,
-            dim_head=self.dim_head,
-            name=f'Attention2',
-            precision=self.precision,
-            use_bias=self.use_bias,
-            dtype=self.dtype,
-            kernel_init=self.kernel_init
-        )
-        
-        self.ff = FlaxFeedForward(dim=self.query_dim)
-        self.norm1 = nn.RMSNorm(epsilon=1e-5, dtype=self.dtype)
-        self.norm2 = nn.RMSNorm(epsilon=1e-5, dtype=self.dtype)
-        self.norm3 = nn.RMSNorm(epsilon=1e-5, dtype=self.dtype)
-        
-    @nn.compact
-    def __call__(self, hidden_states, context=None):
-        # self attention
-        if not self.use_cross_only:
-            print("Using self attention")
-            hidden_states = hidden_states + self.attention1(self.norm1(hidden_states))
-
-        # cross attention
-        hidden_states = hidden_states + self.attention2(self.norm2(hidden_states), context)
-
-        # feed forward
-        hidden_states = hidden_states + self.ff(self.norm3(hidden_states))
-        
-        return hidden_states
-
 class FlaxGEGLU(nn.Module):
     r"""
     Flax implementation of a Linear layer followed by the variant of the gated linear unit activation function from
@@ -330,7 +271,7 @@ class BasicTransformerBlock(nn.Module):
     @nn.compact
     def __call__(self, hidden_states, context=None):
         if self.only_pure_attention:
-            return self.attention2(self.norm2(hidden_states), context)
+            return self.attention2(hidden_states, context)
         
         # self attention
         if not self.use_cross_only:

flaxdiff/models/common.py

@@ -270,8 +270,8 @@ class ResidualBlock(nn.Module):
     @nn.compact
     def __call__(self, x:jax.Array, temb:jax.Array, textemb:jax.Array=None, extra_features:jax.Array=None):
         residual = x
-        # out = nn.GroupNorm(self.norm_groups)(x)
-        out = nn.RMSNorm()(x)
+        out = nn.GroupNorm(self.norm_groups)(x)
+        # out = nn.RMSNorm()(x)
         out = self.activation(out)
 
         out = ConvLayer(
@@ -295,8 +295,8 @@ class ResidualBlock(nn.Module):
         # out = out * (1 + scale) + shift
         out = out + temb
 
-        # out = nn.GroupNorm(self.norm_groups)(out)
-        out = nn.RMSNorm()(out)
+        out = nn.GroupNorm(self.norm_groups)(out)
+        # out = nn.RMSNorm()(out)
         out = self.activation(out)
 
         out = ConvLayer(

flaxdiff/models/simple_unet.py

@@ -4,7 +4,7 @@ from flax import linen as nn
 from flax.typing import Dtype, PrecisionLike
 from typing import Dict, Callable, Sequence, Any, Union, Optional
 import einops
-from .common import kernel_init, ConvLayer, Downsample, Upsample, FourierEmbedding, TimeProjection
+from .common import kernel_init, ConvLayer, Downsample, Upsample, FourierEmbedding, TimeProjection, ResidualBlock
 from .attention import TransformerBlock
 
 class Unet(nn.Module):
@@ -65,7 +65,7 @@ class Unet(nn.Module):
                 if attention_config is not None and j == self.num_res_blocks - 1:   # Apply attention only on the last block
                     x = TransformerBlock(heads=attention_config['heads'], dtype=attention_config.get('dtype', jnp.float32),
                                         dim_head=dim_in // attention_config['heads'],
-                                        use_flash_attention=attention_config.get("flash_attention", True),
+                                        use_flash_attention=attention_config.get("flash_attention", False),
                                         use_projection=attention_config.get("use_projection", False),
                                         use_self_and_cross=attention_config.get("use_self_and_cross", True),
                                         precision=attention_config.get("precision", self.precision),
@@ -103,7 +103,7 @@ class Unet(nn.Module):
             if middle_attention is not None and j == self.num_middle_res_blocks - 1:   # Apply attention only on the last block
                 x = TransformerBlock(heads=middle_attention['heads'], dtype=middle_attention.get('dtype', jnp.float32), 
                                     dim_head=middle_dim_out // middle_attention['heads'],
-                                    use_flash_attention=middle_attention.get("flash_attention", True),
+                                    use_flash_attention=middle_attention.get("flash_attention", False),
                                     use_linear_attention=False,
                                     use_projection=middle_attention.get("use_projection", False),
                                     use_self_and_cross=False,
@@ -146,7 +146,7 @@ class Unet(nn.Module):
                 if attention_config is not None and j == self.num_res_blocks - 1:   # Apply attention only on the last block
                     x = TransformerBlock(heads=attention_config['heads'], dtype=attention_config.get('dtype', jnp.float32), 
                                         dim_head=dim_out // attention_config['heads'],
-                                        use_flash_attention=attention_config.get("flash_attention", True),
+                                        use_flash_attention=attention_config.get("flash_attention", False),
                                         use_projection=attention_config.get("use_projection", False),
                                         use_self_and_cross=attention_config.get("use_self_and_cross", True),
                                         precision=attention_config.get("precision", self.precision),

{flaxdiff-0.1.6.dist-info → flaxdiff-0.1.8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flaxdiff
-Version: 0.1.6
+Version: 0.1.8
 Summary: A versatile and easy to understand Diffusion library
 Author: Ashish Kumar Singh
 Author-email: ashishkmr472@gmail.com

{flaxdiff-0.1.6.dist-info → flaxdiff-0.1.8.dist-info}/RECORD

@@ -1,10 +1,10 @@
 flaxdiff/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 flaxdiff/utils.py,sha256=B0GcHlzlVYDNEIdh2v5qmP4u0neIT-FqexNohuyuCvg,2452
 flaxdiff/models/__init__.py,sha256=FAivVYXxM2JrCFIXf-C3374RB2Hth25dBrzOeNFhH1U,26
-flaxdiff/models/attention.py,sha256=OhpKQXdxWbf8K2_yotLfS0DYdHb-zNpL2p8--ql_FAg,14503
-flaxdiff/models/common.py,sha256=RYNxX9K19hvwSWaB9Wtv7MIZLhcacdugDgD9uZDh8XM,10358
+flaxdiff/models/attention.py,sha256=pDGXG2DT7znvHJWyx7_vTUx235s_D9cubwmA6FDq4qE,12526
+flaxdiff/models/common.py,sha256=lBY2ffKikNeSFlt2umsCTUUe43UgonTVMyQPgzWoAM8,10358
 flaxdiff/models/favor_fastattn.py,sha256=79Ew1nqarsNLPzZaBSd1ILORzJr74CupYeqGiCQK5E4,27689
-flaxdiff/models/simple_unet.py,sha256=hAcz074E9NVdUtECPMi1c1Kw-52Dc6l_ME-5FqIg-n8,9255
+flaxdiff/models/simple_unet.py,sha256=lakCwUkCODEiuS5T6j45Z_sHamcQa9ZWk77NowgjZyc,9273
 flaxdiff/models/simple_vit.py,sha256=vTu2CQRoSOxetBHTrnCWddm-vxrZDkMe8EpdNxtpJMk,4015
 flaxdiff/models/autoencoder/__init__.py,sha256=qY-7MldZpsfkF-_T2LqlRK7VHbqfmosz0NmvzDlBkOk,78
 flaxdiff/models/autoencoder/autoencoder.py,sha256=27_hYl0yXAdH9Mx4Xu9J79mSNo-FEKr9SxhVaS3ffn4,591
@@ -32,7 +32,7 @@ flaxdiff/trainer/__init__.py,sha256=T-vUVq4zHcMK6kpCsG4Gu8vn71q6lZD-lg-Ul7yKfEk,
 flaxdiff/trainer/autoencoder_trainer.py,sha256=al7AsZ7yeDMEiDD-gbcXf0ADq_xfk1VMxvg24GfA-XQ,7008
 flaxdiff/trainer/diffusion_trainer.py,sha256=h5YxIMjBI553xDNeapzLDGF0_4y0MfGRMuHume5sPtM,7785
 flaxdiff/trainer/simple_trainer.py,sha256=f4g2KGuGM__d9v_4Ip3ng8wQubmenWZUW60VEu2ANOg,16774
-flaxdiff-0.1.6.dist-info/METADATA,sha256=sWY_oQgQhhuyW89KyRwIBrpVHBPJjRMmsk5twfgIBlo,20090
-flaxdiff-0.1.6.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
-flaxdiff-0.1.6.dist-info/top_level.txt,sha256=-2-nXnfkJgSfkki1tjm5Faw6Dso7vhtdn2szwCdX5CQ,9
-flaxdiff-0.1.6.dist-info/RECORD,,
+flaxdiff-0.1.8.dist-info/METADATA,sha256=RVH7dPknslUCneKMAY_ira3uzPWfr1whKhsoZXEKiqU,20090
+flaxdiff-0.1.8.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
+flaxdiff-0.1.8.dist-info/top_level.txt,sha256=-2-nXnfkJgSfkki1tjm5Faw6Dso7vhtdn2szwCdX5CQ,9
+flaxdiff-0.1.8.dist-info/RECORD,,