flaxdiff 0.1.25__tar.gz → 0.1.27__tar.gz

{flaxdiff-0.1.25 → flaxdiff-0.1.27}/PKG-INFO

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

{flaxdiff-0.1.25 → flaxdiff-0.1.27}/flaxdiff/models/simple_vit.py

@@ -23,6 +23,7 @@ class PatchEmbedding(nn.Module):
     embedding_dim: int
     dtype: Any = jnp.float32
     precision: Any = jax.lax.Precision.HIGH
+    kernel_init: Callable = partial(kernel_init, 1.0)
 
     @nn.compact
     def __call__(self, x):
@@ -33,6 +34,7 @@ class PatchEmbedding(nn.Module):
                     kernel_size=(self.patch_size, self.patch_size), 
                     strides=(self.patch_size, self.patch_size),
                     dtype=self.dtype,
+                    kernel_init=self.kernel_init(),
                     precision=self.precision)(x)
         x = jnp.reshape(x, (batch, -1, self.embedding_dim))
         return x
@@ -81,7 +83,7 @@ class UViT(nn.Module):
 
         # Patch embedding
         x = PatchEmbedding(patch_size=self.patch_size, embedding_dim=self.emb_features, 
-                           dtype=self.dtype, precision=self.precision)(x)
+                           dtype=self.dtype, precision=self.precision, kernel_init=self.kernel_init)(x)
         num_patches = x.shape[1]
         
         context_emb = nn.DenseGeneral(features=self.emb_features, kernel_init=self.kernel_init(), 
@@ -119,14 +121,14 @@ class UViT(nn.Module):
         
         # # Out blocks
         for i in range(self.num_layers // 2):
-            skip = jnp.concatenate([x, skips.pop()], axis=-1)
-            skip = nn.DenseGeneral(features=self.emb_features, kernel_init=self.kernel_init(), 
-                                   dtype=self.dtype, precision=self.precision)(skip)
+            x = jnp.concatenate([x, skips.pop()], axis=-1)
+            x = nn.DenseGeneral(features=self.emb_features, kernel_init=self.kernel_init(), 
+                                   dtype=self.dtype, precision=self.precision)(x)
             x = TransformerBlock(heads=self.num_heads, dim_head=self.emb_features // self.num_heads, 
                                  dtype=self.dtype, precision=self.precision, use_projection=self.use_projection, 
                                  use_flash_attention=self.use_flash_attention, use_self_and_cross=self.use_self_and_cross, force_fp32_for_softmax=self.force_fp32_for_softmax, 
                                  only_pure_attention=False,
-                                 kernel_init=self.kernel_init())(skip)
+                                 kernel_init=self.kernel_init())(x)
         
         # print(f'Shape of x after transformer blocks: {x.shape}')
         x = self.norm()(x)
@@ -139,6 +141,14 @@ class UViT(nn.Module):
         x = x[:, 1 + num_text_tokens:, :]
         x = unpatchify(x, channels=self.output_channels)
         # print(f'Shape of x after final dense layer: {x.shape}')
-        x = nn.Dense(features=self.output_channels, dtype=self.dtype, precision=self.precision, kernel_init=self.kernel_init())(x)
+        x = nn.Conv(
+            features=self.output_channels, 
+            kernel_size=(3, 3), 
+            strides=(1, 1),
+            padding='SAME', 
+            dtype=self.dtype, 
+            precision=self.precision, 
+            kernel_init=kernel_init(0.0),
+        )(x)
         
         return x

{flaxdiff-0.1.25 → flaxdiff-0.1.27}/flaxdiff.egg-info/PKG-INFO

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

{flaxdiff-0.1.25 → flaxdiff-0.1.27}/setup.py

@@ -11,7 +11,7 @@ required_packages=[
 setup(
     name='flaxdiff',
     packages=find_packages(),
-    version='0.1.25',
+    version='0.1.27',
     description='A versatile and easy to understand Diffusion library',
     long_description=open('README.md').read(),
     long_description_content_type='text/markdown',