flaxdiff 0.1.21__py3-none-any.whl → 0.1.23__py3-none-any.whl

flaxdiff/data/online_loader.py

@@ -21,6 +21,7 @@ import urllib
 
 import PIL.Image
 import cv2
+import traceback 
 
 USER_AGENT = get_datasets_user_agent()
 
@@ -43,7 +44,27 @@ def fetch_single_image(image_url, timeout=None, retries=0):
     return image
 
 
-def default_image_processor(image, image_shape, interpolation=cv2.INTER_CUBIC):
+def default_image_processor(
+    image, image_shape, 
+    min_image_shape=(128, 128),
+    upscale_interpolation=cv2.INTER_CUBIC,
+    downscale_interpolation=cv2.INTER_AREA,
+):
+    image = np.array(image)
+    original_height, original_width = image.shape[:2]
+    # check if the image is too small
+    if min(original_height, original_width) < min(min_image_shape):
+        return None, original_height, original_width
+    # check if wrong aspect ratio
+    if max(original_height, original_width) / min(original_height, original_width) > 2.4:
+        return None, original_height, original_width
+    # check if the variance is too low
+    if np.std(image) < 1e-5:
+        return None, original_height, original_width
+    # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    downscale = max(original_width, original_height) > max(image_shape)
+    interpolation = downscale_interpolation if downscale else upscale_interpolation
+
     image = A.longest_max_size(image, max(
         image_shape), interpolation=interpolation)
     image = A.pad(
@@ -53,7 +74,7 @@ def default_image_processor(image, image_shape, interpolation=cv2.INTER_CUBIC):
         border_mode=cv2.BORDER_CONSTANT,
         value=[255, 255, 255],
     )
-    return image
+    return image, original_height, original_width
 
 
 def map_sample(
@@ -72,23 +93,13 @@ def map_sample(
         if image is None:
             return
 
-        image = np.array(image)
-        original_height, original_width = image.shape[:2]
-        # check if the image is too small
-        if min(original_height, original_width) < min(min_image_shape):
-            return
-        # check if wrong aspect ratio
-        if max(original_height, original_width) / min(original_height, original_width) > 2.4:
-            return
-        # check if the variance is too low
-        if np.std(image) < 1e-4:
+        image, original_height, original_width = image_processor(
+            image, image_shape, min_image_shape=min_image_shape,
+            upscale_interpolation=upscale_interpolation,
+            downscale_interpolation=downscale_interpolation,)
+        
+        if image is None:
             return
-        # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-        downscale = max(original_width, original_height) > max(image_shape)
-        interpolation = downscale_interpolation if downscale else upscale_interpolation
-
-        image = image_processor(
-            image, image_shape, interpolation=interpolation)
         
         data_queue.put({
             "url": url,
@@ -98,7 +109,8 @@ def map_sample(
             "original_width": original_width,
         })
     except Exception as e:
-        print(f"Error processing {url}", e)
+        print(f"Error maping sample {url}", e)
+        traceback.print_exc() 
         # error_queue.put_nowait({
         #     "url": url,
         #     "caption": caption,
@@ -122,7 +134,8 @@ def map_batch(
         with ThreadPoolExecutor(max_workers=num_threads) as executor:
             executor.map(map_sample_fn, batch["url"], batch['caption'])
     except Exception as e:
-        print(f"Error processing batch", e)
+        print(f"Error maping batch", e)
+        traceback.print_exc() 
         # error_queue.put_nowait({
         #     "batch": batch,
         #     "error": str(e)

flaxdiff/models/common.py

@@ -108,12 +108,13 @@ class FourierEmbedding(nn.Module):
 class TimeProjection(nn.Module):
     features:int
     activation:Callable=jax.nn.gelu
+    kernel_init:Callable=partial(kernel_init, 1.0)
 
     @nn.compact
     def __call__(self, x):
-        x = nn.DenseGeneral(self.features, kernel_init=kernel_init(1.0))(x)
+        x = nn.DenseGeneral(self.features, kernel_init=self.kernel_init())(x)
         x = self.activation(x)
-        x = nn.DenseGeneral(self.features, kernel_init=kernel_init(1.0))(x)
+        x = nn.DenseGeneral(self.features, kernel_init=self.kernel_init())(x)
         x = self.activation(x)
         return x
 
@@ -122,7 +123,7 @@ class SeparableConv(nn.Module):
     kernel_size:tuple=(3, 3)
     strides:tuple=(1, 1)
     use_bias:bool=False
-    kernel_init:Callable=kernel_init(1.0)
+    kernel_init:Callable=partial(kernel_init, 1.0)
     padding:str="SAME"
     dtype: Optional[Dtype] = None
     precision: PrecisionLike = None
@@ -132,7 +133,7 @@ class SeparableConv(nn.Module):
         in_features = x.shape[-1]
         depthwise = nn.Conv(
             features=in_features, kernel_size=self.kernel_size,
-            strides=self.strides, kernel_init=self.kernel_init,
+            strides=self.strides, kernel_init=self.kernel_init(),
             feature_group_count=in_features, use_bias=self.use_bias,
             padding=self.padding,
             dtype=self.dtype,
@@ -140,7 +141,7 @@ class SeparableConv(nn.Module):
         )(x)
         pointwise = nn.Conv(
             features=self.features, kernel_size=(1, 1),
-            strides=(1, 1), kernel_init=self.kernel_init,
+            strides=(1, 1), kernel_init=self.kernel_init(),
             use_bias=self.use_bias,
             dtype=self.dtype,
             precision=self.precision
@@ -152,7 +153,7 @@ class ConvLayer(nn.Module):
     features:int
     kernel_size:tuple=(3, 3)
     strides:tuple=(1, 1)
-    kernel_init:Callable=kernel_init(1.0)
+    kernel_init:Callable=partial(kernel_init, 1.0)
     dtype: Optional[Dtype] = None
     precision: PrecisionLike = None
 
@@ -163,7 +164,7 @@ class ConvLayer(nn.Module):
                 features=self.features,
                 kernel_size=self.kernel_size,
                 strides=self.strides,
-                kernel_init=self.kernel_init,
+                kernel_init=self.kernel_init(),
                 dtype=self.dtype,
                 precision=self.precision
             )
@@ -182,7 +183,7 @@ class ConvLayer(nn.Module):
                 features=self.features,
                 kernel_size=self.kernel_size,
                 strides=self.strides,
-                kernel_init=self.kernel_init,
+                kernel_init=self.kernel_init(),
                 dtype=self.dtype,
                 precision=self.precision
             )
@@ -191,7 +192,7 @@ class ConvLayer(nn.Module):
                 features=self.features,
                 kernel_size=self.kernel_size,
                 strides=self.strides,
-                kernel_init=self.kernel_init,
+                kernel_init=self.kernel_init(),
                 dtype=self.dtype,
                 precision=self.precision
             )
@@ -205,6 +206,7 @@ class Upsample(nn.Module):
     activation:Callable=jax.nn.swish
     dtype: Optional[Dtype] = None
     precision: PrecisionLike = None
+    kernel_init:Callable=partial(kernel_init, 1.0)
 
     @nn.compact
     def __call__(self, x, residual=None):
@@ -218,7 +220,8 @@ class Upsample(nn.Module):
             kernel_size=(3, 3),
             strides=(1, 1),
             dtype=self.dtype,
-            precision=self.precision
+            precision=self.precision,
+            kernel_init=self.kernel_init()
         )(out)
         if residual is not None:
             out = jnp.concatenate([out, residual], axis=-1)
@@ -230,6 +233,7 @@ class Downsample(nn.Module):
     activation:Callable=jax.nn.swish
     dtype: Optional[Dtype] = None
     precision: PrecisionLike = None
+    kernel_init:Callable=partial(kernel_init, 1.0)
 
     @nn.compact
     def __call__(self, x, residual=None):
@@ -239,7 +243,8 @@ class Downsample(nn.Module):
             kernel_size=(3, 3),
             strides=(2, 2),
             dtype=self.dtype,
-            precision=self.precision
+            precision=self.precision,
+            kernel_init=self.kernel_init()
         )(x)
         if residual is not None:
             if residual.shape[1] > out.shape[1]:
@@ -264,7 +269,7 @@ class ResidualBlock(nn.Module):
     direction:str=None
     res:int=2
     norm_groups:int=8
-    kernel_init:Callable=kernel_init(1.0)
+    kernel_init:Callable=partial(kernel_init, 1.0)
     dtype: Optional[Dtype] = None
     precision: PrecisionLike = None
     named_norms:bool=False
@@ -291,7 +296,7 @@ class ResidualBlock(nn.Module):
             features=self.features,
             kernel_size=self.kernel_size,
             strides=self.strides,
-            kernel_init=self.kernel_init,
+            kernel_init=self.kernel_init(),
             name="conv1",
             dtype=self.dtype,
             precision=self.precision
@@ -316,7 +321,7 @@ class ResidualBlock(nn.Module):
             features=self.features,
             kernel_size=self.kernel_size,
             strides=self.strides,
-            kernel_init=self.kernel_init,
+            kernel_init=self.kernel_init(),
             name="conv2",
             dtype=self.dtype,
             precision=self.precision
@@ -328,7 +333,7 @@ class ResidualBlock(nn.Module):
                 features=self.features,
                 kernel_size=(1, 1),
                 strides=1,
-                kernel_init=self.kernel_init,
+                kernel_init=self.kernel_init(),
                 name="residual_conv",
                 dtype=self.dtype,
                 precision=self.precision

flaxdiff/models/simple_vit.py

@@ -3,9 +3,20 @@
 import jax
 import jax.numpy as jnp
 from flax import linen as nn
-from typing import Callable, Any
+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
+import einops
+from flax.typing import Dtype, PrecisionLike
+from functools import partial
+
+def unpatchify(x, channels=3):
+    patch_size = int((x.shape[2] // channels) ** 0.5)
+    h = w = int(x.shape[1] ** .5)
+    assert h * w == x.shape[1] and patch_size ** 2 * channels == x.shape[2], f"Invalid shape: {x.shape}, should be {h*w}, {patch_size**2*channels}"
+    x = einops.rearrange(x, 'B (h w) (p1 p2 C) -> B (h p1) (w p2) C', h=h, p1=patch_size, p2=patch_size)
+    return x
 
 class PatchEmbedding(nn.Module):
     patch_size: int
@@ -37,39 +48,28 @@ class PositionalEncoding(nn.Module):
                         (1, self.max_len, self.embedding_dim))
         return x + pe[:, :x.shape[1], :]
 
-class TransformerEncoder(nn.Module):
-    num_layers: int
-    num_heads: int
-    dropout_rate: float = 0.1
-    dtype: Any = jnp.float32
-    precision: Any = jax.lax.Precision.HIGH
-    use_projection: bool = False
-
-    @nn.compact
-    def __call__(self, x, context=None):
-        for _ in range(self.num_layers):
-            x = TransformerBlock(
-                heads=self.num_heads,
-                dim_head=x.shape[-1] // self.num_heads,
-                dropout_rate=self.dropout_rate,
-                dtype=self.dtype,
-                precision=self.precision,
-                use_self_and_cross=True,
-                use_projection=self.use_projection,
-            )(x, context)
-        return x
-
-class VisionTransformer(nn.Module):
+class UViT(nn.Module):
+    output_channels:int=3
     patch_size: int = 16
-    embedding_dim: int = 768
+    emb_features:int=768,
     num_layers: int = 12
     num_heads: int = 12
-    emb_features: int = 256
     dropout_rate: float = 0.1
     dtype: Any = jnp.float32
     precision: Any = jax.lax.Precision.HIGH
     use_projection: bool = False
-
+    activation:Callable = jax.nn.swish
+    norm_groups:int=8
+    dtype: Optional[Dtype] = None
+    precision: PrecisionLike = None
+    kernel_init: Callable = partial(kernel_init, 1.0)
+
+    def setup(self):
+        if self.norm_groups > 0:
+            self.norm = partial(nn.GroupNorm, self.norm_groups)
+        else:
+            self.norm = partial(nn.RMSNorm, 1e-5)
+            
     @nn.compact
     def __call__(self, x, temb, textcontext=None):
         # Time embedding
@@ -77,44 +77,65 @@ class VisionTransformer(nn.Module):
         temb = TimeProjection(features=self.emb_features)(temb)
 
         # Patch embedding
-        x = PatchEmbedding(patch_size=self.patch_size, embedding_dim=self.embedding_dim, 
+        x = PatchEmbedding(patch_size=self.patch_size, embedding_dim=self.emb_features, 
                            dtype=self.dtype, precision=self.precision)(x)
+        num_patches = x.shape[1]
         
-        # Add positional encoding
-        x = PositionalEncoding(max_len=x.shape[1], embedding_dim=self.embedding_dim)(x)
+        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]
+        
+        # print(f'Shape of x after patch embedding: {x.shape}, numPatches: {num_patches}, temb: {temb.shape}, context_emb: {context_emb.shape}')
         
-        num_patches = x.shape[1]
-
         # Add time embedding
         temb = jnp.expand_dims(temb, axis=1)
-        x = jnp.concatenate([x, temb], axis=1)
-
-        # Transformer encoder
-        x = TransformerEncoder(
-            num_layers=self.num_layers,
-            num_heads=self.num_heads,
-            dropout_rate=self.dropout_rate,
-            dtype=self.dtype,
-            precision=self.precision,
-            use_projection=self.use_projection
-        )(x, textcontext)
-
-        x = x[:, :num_patches, :]
-
-        # Reshape to image dimensions
-        batch, _, _ = x.shape
-        height = width = int((num_patches) ** 0.5)
-        x = jnp.reshape(x, (batch, height, width, self.embedding_dim))
-
-        # Final convolution to get the desired output channels
-        x = ConvLayer(
-            conv_type="conv",
-            features=3,
-            kernel_size=(3, 3),
-            strides=(1, 1),
-            kernel_init=kernel_init(0.0),
-            dtype=self.dtype,
-            precision=self.precision
-        )(x)
-
+        x = jnp.concatenate([x, temb, context_emb], axis=1)
+        # 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)(x)
+        
+        # print(f'Shape of x after positional encoding: {x.shape}')
+        
+        skips = []
+        # In blocks
+        for i in range(self.num_layers // 2):
+            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=False, use_self_and_cross=False, force_fp32_for_softmax=True, 
+                                 only_pure_attention=False,
+                                 kernel_init=self.kernel_init())(x)
+            skips.append(x)
+            
+        # Middle block
+        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=False, use_self_and_cross=True, force_fp32_for_softmax=True, 
+                             only_pure_attention=False,
+                             kernel_init=self.kernel_init())(x)
+        
+        # # 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 = 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=False, use_self_and_cross=False, force_fp32_for_softmax=True, 
+                                 only_pure_attention=False,
+                                 kernel_init=self.kernel_init())(skip)
+        
+        # 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}')
+        x = nn.Dense(features=self.output_channels, dtype=self.dtype, precision=self.precision, kernel_init=self.kernel_init())(x)
+        
         return x

{flaxdiff-0.1.21.dist-info → flaxdiff-0.1.23.dist-info}/METADATA

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

{flaxdiff-0.1.21.dist-info → flaxdiff-0.1.23.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=w6gi1tAzWr4gtPt7onpStzOxp7Kdo_2q8Ro4Yi7OT4w,10549
+flaxdiff/data/online_loader.py,sha256=LIK_O1C3yDPvvAEOWvsJrVeBopVqjg2IOMTbiSIvH6M,11025
 flaxdiff/models/__init__.py,sha256=FAivVYXxM2JrCFIXf-C3374RB2Hth25dBrzOeNFhH1U,26
 flaxdiff/models/attention.py,sha256=ZbDGIb5Q6FRqJ6qRY660cqw4WvF9IwCnhEuYdTpLPdM,13023
-flaxdiff/models/common.py,sha256=fd-Fl0VCNEBjijHNwGBqYL5VvXe9u0347h25czNTmRw,10780
+flaxdiff/models/common.py,sha256=lzfOHB-7Bjx83ZZzywXy2mjhwP2UOMKR11vdVSKvsCo,11068
 flaxdiff/models/favor_fastattn.py,sha256=79Ew1nqarsNLPzZaBSd1ILORzJr74CupYeqGiCQK5E4,27689
 flaxdiff/models/simple_unet.py,sha256=h1o9mQlLJy7Ec8Pz_O5miRbAyUaM5UNhSs-oXzpQvZo,10763
-flaxdiff/models/simple_vit.py,sha256=xD23i1b7WEvoH4tUMsLyCe9ebDcv-PpaV0Nso38Jlb8,3887
+flaxdiff/models/simple_vit.py,sha256=tXO0WIozq2C4j2GqphSM9mMAYEwj9fKr7rfm4G6vf4A,6403
 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.21.dist-info/METADATA,sha256=k1s_EIWBL0y4oCxXxr3QIi7LQbR47_jyDFfjVbURSMY,22083
-flaxdiff-0.1.21.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
-flaxdiff-0.1.21.dist-info/top_level.txt,sha256=-2-nXnfkJgSfkki1tjm5Faw6Dso7vhtdn2szwCdX5CQ,9
-flaxdiff-0.1.21.dist-info/RECORD,,
+flaxdiff-0.1.23.dist-info/METADATA,sha256=iTgk4DY-kuALF-Y1U-a433c7pYRhSy-bXPQrtXh6d54,22083
+flaxdiff-0.1.23.dist-info/WHEEL,sha256=HiCZjzuy6Dw0hdX5R3LCFPDmFS4BWl8H-8W39XfmgX4,91
+flaxdiff-0.1.23.dist-info/top_level.txt,sha256=-2-nXnfkJgSfkki1tjm5Faw6Dso7vhtdn2szwCdX5CQ,9
+flaxdiff-0.1.23.dist-info/RECORD,,

{flaxdiff-0.1.21.dist-info → flaxdiff-0.1.23.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (72.1.0)
+Generator: setuptools (72.2.0)
 Root-Is-Purelib: true
 Tag: py3-none-any