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

flaxdiff/models/simple_vit.py

@@ -136,32 +136,25 @@ class UViT(nn.Module):
         # print(f'Shape of x after transformer blocks: {x.shape}')
         x = self.norm()(x)
         
-        # print(f'Shape of x after norm: {x.shape}')
-        
         patch_dim = self.patch_size ** 2 * self.output_channels
         x = nn.Dense(features=patch_dim, dtype=self.dtype, precision=self.precision, kernel_init=self.kernel_init())(x)
-        # print(f'Shape of x after patch dense layer: {x.shape}, patch_dim: {patch_dim}')
         x = x[:, 1 + num_text_tokens:, :]
         x = unpatchify(x, channels=self.output_channels)
-        # print(f'Shape of x after final dense layer: {x.shape}')
         
         if self.add_residualblock_output:
             # Concatenate the original image
             x = jnp.concatenate([original_img, x], axis=-1)
             
-            x = ResidualBlock(
+            x = ConvLayer(
                 "conv",
-                name="final_residual",
                 features=64,
-                kernel_init=self.kernel_init(1.0),
-                kernel_size=(3,3),
+                kernel_size=(3, 3),
                 strides=(1, 1),
-                activation=self.activation,
-                norm_groups=self.norm_groups,
+                # activation=jax.nn.mish
+                kernel_init=self.kernel_init(0.0),
                 dtype=self.dtype,
-                precision=self.precision,
-                named_norms=False
-            )(x, temb)
+                precision=self.precision
+            )(x)
 
             x = self.norm()(x)
             x = self.activation(x)

{flaxdiff-0.1.28.dist-info → flaxdiff-0.1.29.dist-info}/METADATA

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

{flaxdiff-0.1.28.dist-info → flaxdiff-0.1.29.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=W2LxTKWA0wJHbPLf4hd2eUO4-ZV5u0Y-M168QulGwTg,7786
+flaxdiff/models/simple_vit.py,sha256=atjeXc22w8WYub_6d0JAFFgvQ4TP1wt4N1ubIzZlQZ0,7436
 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.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,,
+flaxdiff-0.1.29.dist-info/METADATA,sha256=PcevgEjt61-62ccMC_CI4EvHYUX-tdrpEBptKXkTudA,22083
+flaxdiff-0.1.29.dist-info/WHEEL,sha256=cVxcB9AmuTcXqmwrtPhNK88dr7IR_b6qagTj0UvIEbY,91
+flaxdiff-0.1.29.dist-info/top_level.txt,sha256=-2-nXnfkJgSfkki1tjm5Faw6Dso7vhtdn2szwCdX5CQ,9
+flaxdiff-0.1.29.dist-info/RECORD,,