flaxdiff 0.1.30__py3-none-any.whl → 0.1.32__py3-none-any.whl

flaxdiff/data/online_loader.py

@@ -117,7 +117,12 @@ def map_sample(
         #     "error": str(e)
         # })
         pass
-
+    
+def default_feature_extractor(sample):
+    return {
+        "url": sample["url"],
+        "caption": sample["caption"],
+    }
 
 def map_batch(
     batch, num_threads=256, image_shape=(256, 256), 
@@ -125,6 +130,7 @@ def map_batch(
     timeout=15, retries=3, image_processor=default_image_processor,
     upscale_interpolation=cv2.INTER_CUBIC,
     downscale_interpolation=cv2.INTER_AREA,
+    feature_extractor=default_feature_extractor,
 ):
     try:
         map_sample_fn = partial(map_sample, image_shape=image_shape, min_image_shape=min_image_shape,
@@ -132,7 +138,9 @@ def map_batch(
                                 upscale_interpolation=upscale_interpolation,
                                 downscale_interpolation=downscale_interpolation)
         with ThreadPoolExecutor(max_workers=num_threads) as executor:
-            executor.map(map_sample_fn, batch["url"], batch['caption'])
+            features = feature_extractor(batch)
+            url, caption = features["url"], features["caption"]
+            executor.map(map_sample_fn, url, caption)
     except Exception as e:
         print(f"Error maping batch", e)
         traceback.print_exc() 
@@ -149,12 +157,14 @@ def parallel_image_loader(
     num_threads=256, timeout=15, retries=3, image_processor=default_image_processor,
     upscale_interpolation=cv2.INTER_CUBIC,
     downscale_interpolation=cv2.INTER_AREA,
+    feature_extractor=default_feature_extractor,
 ):
     map_batch_fn = partial(map_batch, num_threads=num_threads, image_shape=image_shape, 
                            min_image_shape=min_image_shape,
                            timeout=timeout, retries=retries, image_processor=image_processor,
                            upscale_interpolation=upscale_interpolation,
-                           downscale_interpolation=downscale_interpolation)
+                           downscale_interpolation=downscale_interpolation,
+                           feature_extractor=feature_extractor)
     shard_len = len(dataset) // num_workers
     print(f"Local Shard lengths: {shard_len}")
     with multiprocessing.Pool(num_workers) as pool:
@@ -181,6 +191,7 @@ class ImageBatchIterator:
         image_processor=default_image_processor,
         upscale_interpolation=cv2.INTER_CUBIC,
         downscale_interpolation=cv2.INTER_AREA,
+        feature_extractor=default_feature_extractor,
     ):
         self.dataset = dataset
         self.num_workers = num_workers
@@ -191,7 +202,8 @@ class ImageBatchIterator:
                          num_workers=num_workers, 
                          timeout=timeout, retries=retries, image_processor=image_processor,
                          upscale_interpolation=upscale_interpolation,
-                         downscale_interpolation=downscale_interpolation)
+                         downscale_interpolation=downscale_interpolation,
+                         feature_extractor=feature_extractor)
         self.thread = threading.Thread(target=loader, args=(dataset,))
         self.thread.start()
 
@@ -256,6 +268,7 @@ class OnlineStreamingDataLoader():
         image_processor=default_image_processor,
         upscale_interpolation=cv2.INTER_CUBIC,
         downscale_interpolation=cv2.INTER_AREA,
+        feature_extractor=default_feature_extractor,
     ):
         if isinstance(dataset, str):
             dataset_path = dataset
@@ -281,7 +294,8 @@ class OnlineStreamingDataLoader():
                                            num_workers=num_workers, batch_size=batch_size, num_threads=num_threads,
                                             timeout=timeout, retries=retries, image_processor=image_processor,
                                              upscale_interpolation=upscale_interpolation,
-                                             downscale_interpolation=downscale_interpolation)
+                                             downscale_interpolation=downscale_interpolation,
+                                             feature_extractor=feature_extractor)
         self.batch_size = batch_size
 
         # Launch a thread to load batches in the background

flaxdiff/models/simple_vit.py

@@ -23,7 +23,7 @@ class PatchEmbedding(nn.Module):
     embedding_dim: int
     dtype: Any = jnp.float32
     precision: Any = jax.lax.Precision.HIGH
-    kernel_init: Callable = kernel_init(1.0)
+    kernel_init: Callable = partial(kernel_init, 1.0)
 
     @nn.compact
     def __call__(self, x):
@@ -34,7 +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,
+                    kernel_init=self.kernel_init(),
                     precision=self.precision)(x)
         x = jnp.reshape(x, (batch, -1, self.embedding_dim))
         return x
@@ -67,7 +67,7 @@ class UViT(nn.Module):
     norm_groups:int=8
     dtype: Optional[Dtype] = None
     precision: PrecisionLike = None
-    kernel_init: Callable = partial(kernel_init)
+    kernel_init: Callable = partial(kernel_init, scale=1.0)
     add_residualblock_output: bool = False
 
     def setup(self):
@@ -86,10 +86,10 @@ class UViT(nn.Module):
 
         # Patch embedding
         x = PatchEmbedding(patch_size=self.patch_size, embedding_dim=self.emb_features, 
-                           dtype=self.dtype, precision=self.precision, kernel_init=self.kernel_init(1.0))(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(1.0), 
+        context_emb = nn.DenseGeneral(features=self.emb_features, kernel_init=self.kernel_init(), 
                                dtype=self.dtype, precision=self.precision)(textcontext)
         num_text_tokens = textcontext.shape[1]
         
@@ -112,7 +112,7 @@ class UViT(nn.Module):
                                  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(1.0))(x)
+                                 kernel_init=self.kernel_init())(x)
             skips.append(x)
             
         # Middle block
@@ -120,24 +120,24 @@ class UViT(nn.Module):
                              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(1.0))(x)
+                             kernel_init=self.kernel_init())(x)
         
         # # Out blocks
         for i in range(self.num_layers // 2):
             x = jnp.concatenate([x, skips.pop()], axis=-1)
-            x = nn.DenseGeneral(features=self.emb_features, kernel_init=self.kernel_init(1.0), 
+            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(1.0))(x)
+                                 kernel_init=self.kernel_init())(x)
         
         # print(f'Shape of x after transformer blocks: {x.shape}')
         x = self.norm()(x)
         
         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(1.0))(x)
+        x = nn.Dense(features=patch_dim, dtype=self.dtype, precision=self.precision, kernel_init=self.kernel_init())(x)
         x = x[:, 1 + num_text_tokens:, :]
         x = unpatchify(x, channels=self.output_channels)
         
@@ -151,7 +151,7 @@ class UViT(nn.Module):
                 kernel_size=(3, 3),
                 strides=(1, 1),
                 # activation=jax.nn.mish
-                kernel_init=self.kernel_init(0.0),
+                kernel_init=self.kernel_init(scale=0.0),
                 dtype=self.dtype,
                 precision=self.precision
             )(x)
@@ -165,7 +165,7 @@ class UViT(nn.Module):
             kernel_size=(3, 3),
             strides=(1, 1),
             # activation=jax.nn.mish
-            kernel_init=self.kernel_init(0.0),
+            kernel_init=self.kernel_init(scale=0.0),
             dtype=self.dtype,
             precision=self.precision
         )(x)

{flaxdiff-0.1.30.dist-info → flaxdiff-0.1.32.dist-info}/METADATA

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

{flaxdiff-0.1.30.dist-info → flaxdiff-0.1.32.dist-info}/RECORD

@@ -1,13 +1,13 @@
 flaxdiff/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 flaxdiff/utils.py,sha256=B0GcHlzlVYDNEIdh2v5qmP4u0neIT-FqexNohuyuCvg,2452
 flaxdiff/data/__init__.py,sha256=PM3PkHihyohT5SHVYKc8vQ4IeVfGPpCktkSVwvqMjQ4,52
-flaxdiff/data/online_loader.py,sha256=LIK_O1C3yDPvvAEOWvsJrVeBopVqjg2IOMTbiSIvH6M,11025
+flaxdiff/data/online_loader.py,sha256=Yi1thc2izJoXeB9UPTAkfWhWp8m4UYWi6cDwW6I6zUc,11661
 flaxdiff/models/__init__.py,sha256=FAivVYXxM2JrCFIXf-C3374RB2Hth25dBrzOeNFhH1U,26
 flaxdiff/models/attention.py,sha256=ZbDGIb5Q6FRqJ6qRY660cqw4WvF9IwCnhEuYdTpLPdM,13023
 flaxdiff/models/common.py,sha256=hWsSs2BP2J-JN1s4qLRr-h-KYkcVyl2hOp1Wsm_L-h8,10994
 flaxdiff/models/favor_fastattn.py,sha256=79Ew1nqarsNLPzZaBSd1ILORzJr74CupYeqGiCQK5E4,27689
 flaxdiff/models/simple_unet.py,sha256=h1o9mQlLJy7Ec8Pz_O5miRbAyUaM5UNhSs-oXzpQvZo,10763
-flaxdiff/models/simple_vit.py,sha256=Nnrlo5T9IUu3lu6y-SIWIgfISc07uOztBB4kyfBrQVY,7443
+flaxdiff/models/simple_vit.py,sha256=NHt6v-teGjiI65fk1l1WN3WqfeqTE7xY9VYqBiYUDgI,7454
 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.30.dist-info/METADATA,sha256=lzEiqudjsqRLsDrI1icVnN3NM8hHrAqWloafwhxbhBE,22083
-flaxdiff-0.1.30.dist-info/WHEEL,sha256=cVxcB9AmuTcXqmwrtPhNK88dr7IR_b6qagTj0UvIEbY,91
-flaxdiff-0.1.30.dist-info/top_level.txt,sha256=-2-nXnfkJgSfkki1tjm5Faw6Dso7vhtdn2szwCdX5CQ,9
-flaxdiff-0.1.30.dist-info/RECORD,,
+flaxdiff-0.1.32.dist-info/METADATA,sha256=yqODOUEAlcE_60tOOfSP7XoH0VTBm3dY9J1ali36VFY,22083
+flaxdiff-0.1.32.dist-info/WHEEL,sha256=cVxcB9AmuTcXqmwrtPhNK88dr7IR_b6qagTj0UvIEbY,91
+flaxdiff-0.1.32.dist-info/top_level.txt,sha256=-2-nXnfkJgSfkki1tjm5Faw6Dso7vhtdn2szwCdX5CQ,9
+flaxdiff-0.1.32.dist-info/RECORD,,