flaxdiff 0.1.27__py3-none-any.whl → 0.1.28__py3-none-any.whl

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,6 +81,8 @@ 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, 
@@ -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.27.dist-info → flaxdiff-0.1.28.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flaxdiff
-Version: 0.1.27
+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.27.dist-info → flaxdiff-0.1.28.dist-info}/RECORD

@@ -7,7 +7,7 @@ flaxdiff/models/attention.py,sha256=ZbDGIb5Q6FRqJ6qRY660cqw4WvF9IwCnhEuYdTpLPdM,
 flaxdiff/models/common.py,sha256=fw_gP7PZayO6RVe6xSf-7FtVq-S0pp5U6NgHg4PlKO8,10990
 flaxdiff/models/favor_fastattn.py,sha256=79Ew1nqarsNLPzZaBSd1ILORzJr74CupYeqGiCQK5E4,27689
 flaxdiff/models/simple_unet.py,sha256=h1o9mQlLJy7Ec8Pz_O5miRbAyUaM5UNhSs-oXzpQvZo,10763
-flaxdiff/models/simple_vit.py,sha256=-xGeiRztVssisf0CRd9CvlBQNrIUXaRSQbNckYvkuac,6972
+flaxdiff/models/simple_vit.py,sha256=W2LxTKWA0wJHbPLf4hd2eUO4-ZV5u0Y-M168QulGwTg,7786
 flaxdiff/models/autoencoder/__init__.py,sha256=qY-7MldZpsfkF-_T2LqlRK7VHbqfmosz0NmvzDlBkOk,78
 flaxdiff/models/autoencoder/autoencoder.py,sha256=27_hYl0yXAdH9Mx4Xu9J79mSNo-FEKr9SxhVaS3ffn4,591
 flaxdiff/models/autoencoder/diffusers.py,sha256=JHeFLCxiHhu-QHwhKiCuKsQJn4AZumquiuxgZkiYGQ0,3643
@@ -34,7 +34,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=wKkg63DWZjx2MoM3VQNCDIr40rWN8fUGxH9jWWxfZao,9373
 flaxdiff/trainer/simple_trainer.py,sha256=Z77zRS5viJpd2Mpl6sonJk5WcnEWi2Cd4gl4u5tIX2M,18206
-flaxdiff-0.1.27.dist-info/METADATA,sha256=-344uFzDA8b17cd1LV5RpDDx4bGZ6i8kdNFJ439FD9g,22083
-flaxdiff-0.1.27.dist-info/WHEEL,sha256=cVxcB9AmuTcXqmwrtPhNK88dr7IR_b6qagTj0UvIEbY,91
-flaxdiff-0.1.27.dist-info/top_level.txt,sha256=-2-nXnfkJgSfkki1tjm5Faw6Dso7vhtdn2szwCdX5CQ,9
-flaxdiff-0.1.27.dist-info/RECORD,,
+flaxdiff-0.1.28.dist-info/METADATA,sha256=AeUPnS3eT-lJSMSM9p1J_HBkJz9f5QFey5r0wdZddH8,22083
+flaxdiff-0.1.28.dist-info/WHEEL,sha256=cVxcB9AmuTcXqmwrtPhNK88dr7IR_b6qagTj0UvIEbY,91
+flaxdiff-0.1.28.dist-info/top_level.txt,sha256=-2-nXnfkJgSfkki1tjm5Faw6Dso7vhtdn2szwCdX5CQ,9
+flaxdiff-0.1.28.dist-info/RECORD,,