flaxdiff 0.1.26__tar.gz → 0.1.28__tar.gz

{flaxdiff-0.1.26 → flaxdiff-0.1.28}/PKG-INFO

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

{flaxdiff-0.1.26 → flaxdiff-0.1.28}/flaxdiff/models/simple_vit.py

@@ -6,7 +6,7 @@ from flax import linen as nn
 from typing import Callable, Any, Optional, Tuple
 from .simple_unet import FourierEmbedding, TimeProjection, ConvLayer, kernel_init
 from .attention import TransformerBlock
-from flaxdiff.models.simple_unet import FourierEmbedding, TimeProjection, ConvLayer, kernel_init
+from flaxdiff.models.simple_unet import FourierEmbedding, TimeProjection, ConvLayer, kernel_init, ResidualBlock
 import einops
 from flax.typing import Dtype, PrecisionLike
 from functools import partial
@@ -68,6 +68,7 @@ class UViT(nn.Module):
     dtype: Optional[Dtype] = None
     precision: PrecisionLike = None
     kernel_init: Callable = partial(kernel_init, 1.0)
+    add_residualblock_output: bool = False
 
     def setup(self):
         if self.norm_groups > 0:
@@ -80,10 +81,12 @@ class UViT(nn.Module):
         # Time embedding
         temb = FourierEmbedding(features=self.emb_features)(temb)
         temb = TimeProjection(features=self.emb_features)(temb)
+        
+        original_img = x
 
         # 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(), 
@@ -98,7 +101,7 @@ class UViT(nn.Module):
         # print(f'Shape of x after time embedding: {x.shape}')
         
         # Add positional encoding
-        x = PositionalEncoding(max_len=x.shape[1], embedding_dim=self.emb_features, kernel_init=self.kernel_init)(x)
+        x = PositionalEncoding(max_len=x.shape[1], embedding_dim=self.emb_features)(x)
         
         # print(f'Shape of x after positional encoding: {x.shape}')
         
@@ -141,14 +144,36 @@ 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.Conv(
-            features=self.output_channels, 
-            kernel_size=(3, 3), 
+        
+        if self.add_residualblock_output:
+            # Concatenate the original image
+            x = jnp.concatenate([original_img, x], axis=-1)
+            
+            x = ResidualBlock(
+                "conv",
+                name="final_residual",
+                features=64,
+                kernel_init=self.kernel_init(1.0),
+                kernel_size=(3,3),
+                strides=(1, 1),
+                activation=self.activation,
+                norm_groups=self.norm_groups,
+                dtype=self.dtype,
+                precision=self.precision,
+                named_norms=False
+            )(x, temb)
+
+            x = self.norm()(x)
+            x = self.activation(x)
+
+        x = ConvLayer(
+            "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),
+            # activation=jax.nn.mish
+            kernel_init=self.kernel_init(0.0),
+            dtype=self.dtype,
+            precision=self.precision
         )(x)
-        
         return x

{flaxdiff-0.1.26 → flaxdiff-0.1.28}/flaxdiff.egg-info/PKG-INFO

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

{flaxdiff-0.1.26 → flaxdiff-0.1.28}/setup.py

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