flaxdiff 0.2.7__py3-none-any.whl → 0.2.9__py3-none-any.whl

flaxdiff/data/dataloaders.py

@@ -251,6 +251,12 @@ def generate_collate_fn(media_type="image"):
     else:  # Default to image
         return image_collate
 
+class CaptionDeletionTransform(pygrain.MapTransform):
+    def map(self, element):
+        """Delete the caption from the element."""
+        if "caption" in element:
+            del element["caption"]
+        return element
 
 def get_dataset_grain(
     data_name="cc12m",
@@ -288,7 +294,6 @@ def get_dataset_grain(
     dataset = datasetMap[data_name]
     data_source = dataset["source"](dataset_source)
     augmenter = dataset["augmenter"](image_scale, method)
-
     local_batch_size = batch_size // jax.process_count()
 
     train_sampler = pygrain.IndexSampler(
@@ -310,8 +315,8 @@ def get_dataset_grain(
     def get_trainset():
         transformations = [
             augmenter(),
-            pygrain.Batch(local_batch_size, drop_remainder=True),
         ]
+        transformations.append(pygrain.Batch(local_batch_size, drop_remainder=True))
 
         loader = pygrain.DataLoader(
             data_source=data_source,
@@ -325,24 +330,23 @@ def get_dataset_grain(
         )
         return loader
     
-    # def get_valset():
-    #     transformations = [
-    #         augmenter(),
-    #         pygrain.Batch(local_batch_size, drop_remainder=True),
-    #     ]
+    def get_valset():
+        transformations = [
+            augmenter(),
+            pygrain.Batch(32, drop_remainder=True),
+        ]
 
-    #     loader = pygrain.DataLoader(
-    #         data_source=data_source,
-    #         sampler=val_sampler,
-    #         operations=transformations,
-    #         worker_count=worker_count,
-    #         read_options=pygrain.ReadOptions(
-    #             read_thread_count, read_buffer_size
-    #         ),
-    #         worker_buffer_size=worker_buffer_size,
-    #     )
-    #     return loader
-    get_valset = get_trainset  # For now, use the same function for validation
+        loader = pygrain.DataLoader(
+            data_source=data_source,
+            sampler=train_sampler,
+            operations=transformations,
+            worker_count=8,
+            read_options=pygrain.ReadOptions(
+                32, 128
+            ),
+            worker_buffer_size=32,
+        )
+        return loader
 
     return {
         "train": get_trainset,

flaxdiff/data/dataset_map.py

@@ -8,7 +8,7 @@ from .sources.videos import VideoTFDSSource, VideoLocalSource, AudioVideoAugment
 # ---------------------------------------------------------------------------------
 
 from .sources.images import data_source_tfds, tfds_augmenters, data_source_gcs
-from .sources.images import data_source_combined_gcs, gcs_augmenters
+from .sources.images import data_source_combined_gcs, gcs_augmenters, gcs_filters
 
 # Configure the following for your datasets
 datasetMap = {
@@ -23,6 +23,7 @@ datasetMap = {
     "laiona_coco": {
         "source": data_source_gcs('datasets/laion12m+mscoco'),
         "augmenter": gcs_augmenters,
+        "filter": gcs_filters,
     },
     "aesthetic_coyo": {
         "source": data_source_gcs('arrayrecords/aestheticCoyo_0.25clip_6aesthetic'),

flaxdiff/data/sources/base.py

@@ -62,6 +62,18 @@ class DataAugmenter(ABC):
         """
         pass
     
+    @abstractmethod
+    def create_filter(self, **kwargs) -> Callable[[], pygrain.FilterTransform]:
+        """Create a filter function for the data.
+        
+        Args:
+            **kwargs: Additional arguments for the filter.
+            
+        Returns:
+            A callable that returns a pygrain.FilterTransform instance.
+        """
+        pass
+    
     @staticmethod
     def create(augmenter_type: str, **kwargs) -> 'DataAugmenter':
         """Factory method to create a data augmenter of the specified type.

flaxdiff/data/sources/images.py

@@ -11,7 +11,7 @@ import struct as st
 from functools import partial
 import numpy as np
 from .base import DataSource, DataAugmenter
-
+import traceback
 
 # ----------------------------------------------------------------------------------
 # Utility functions
@@ -79,6 +79,24 @@ def labelizer_oxford_flowers102(path):
 # TFDS Image Source
 # ----------------------------------------------------------------------------------
 
+def get_oxford_valset(text_encoder):
+    # Construct a validation set by the prompts for consistency
+    val_prompts = ['water tulip', ' a water lily', ' a water lily', ' a photo of a rose', ' a photo of a rose', ' a water lily', ' a water lily', ' a photo of a marigold', ' a photo of a marigold', ' a photo of a marigold', ' a water lily', ' a photo of a sunflower', ' a photo of a lotus', ' columbine', ' columbine', ' an orchid', ' an orchid', ' an orchid', ' a water lily', ' a water lily', ' a water lily', ' columbine', ' columbine', ' a photo of a sunflower', ' a photo of a sunflower', ' a photo of a sunflower', ' a photo of a lotus', ' a photo of a lotus', ' a photo of a marigold', ' a photo of a marigold', ' a photo of a rose', ' a photo of a rose', ' a photo of a rose', ' orange dahlia', ' orange dahlia', ' a lenten rose', ' a lenten rose', ' a water lily', ' a water lily', ' a water lily', ' a water lily', ' an orchid', ' an orchid', ' an orchid', ' hard-leaved pocket orchid', ' bird of paradise', ' bird of paradise', ' a photo of a lovely rose', ' a photo of a lovely rose', ' a photo of a globe-flower', ' a photo of a globe-flower', ' a photo of a lovely rose', ' a photo of a lovely rose', ' a photo of a ruby-lipped cattleya', ' a photo of a ruby-lipped cattleya', ' a photo of a lovely rose', ' a water lily', ' a osteospermum', ' a osteospermum', ' a water lily', ' a water lily', ' a water lily', ' a red rose', ' a red rose']
+    val_prompts *= 100
+
+    def get_val_dataset(batch_size=128):
+        for i in range(0, len(val_prompts), batch_size):
+            try:
+                prompts = val_prompts[i:i + batch_size]
+                tokens = text_encoder.tokenize(prompts)
+                yield {"text": tokens}
+            except Exception as e:
+                print(f"Error in get_val_dataset: {e}")
+                traceback.print_exc()
+                continue
+
+    return get_val_dataset, len(val_prompts)
+    
 class ImageTFDSSource(DataSource):
     """Data source for TensorFlow Datasets (TFDS) image datasets."""
     
@@ -168,7 +186,11 @@ class ImageTFDSAugmenter(DataAugmenter):
                 }
         
         return TFDSTransform
-
+    
+    def create_filter(self, image_scale: int = 256):
+        class FilterTransform(pygrain.FilterTransform):
+            def map(self, element) -> bool:
+                return True
 """
 Batch structure:
 {
@@ -237,7 +259,6 @@ class CombinedImageGCSSource(DataSource):
                 records_path) if 'array_record' in i]
         return pygrain.ArrayRecordDataSource(records)
 
-
 class ImageGCSAugmenter(DataAugmenter):
     """Augmenter for GCS image datasets."""
     
@@ -297,6 +318,52 @@ class ImageGCSAugmenter(DataAugmenter):
                 }
         
         return GCSTransform
+    
+    def create_filter(self, image_scale: int = 256):
+        import torch.nn.functional as F
+        class FilterTransform(pygrain.FilterTransform):
+            """
+            Filter transform for GCS data source.
+            """
+            def __init__(self, model=None, processor=None, method=cv2.INTER_AREA):
+                super().__init__()
+                self.image_scale = image_scale
+                if model is None:
+                    from transformers import AutoProcessor, CLIPVisionModelWithProjection, FlaxCLIPModel, CLIPModel
+                    model_name = "openai/clip-vit-base-patch32"
+                    model = CLIPModel.from_pretrained(model_name)
+                    processor = AutoProcessor.from_pretrained(model_name, use_fast=False)
+                self.method = method
+                self.model = model
+                self.processor = processor
+                
+                # def _filter_(pixel_values, input_ids):
+                #     image_embeds = self.model.get_image_features(pixel_values=pixel_values)
+                #     text_embeds = self.model.get_text_features(input_ids=input_ids)
+                #     image_embeds = image_embeds / jnp.linalg.norm(image_embeds, axis=-1, keepdims=True)
+                #     text_embeds = text_embeds / jnp.linalg.norm(text_embeds, axis=-1, keepdims=True)
+                #     similarity = jnp.sum(image_embeds * text_embeds, axis=-1)
+                #     return jnp.all(similarity >= 0.25)
+                
+                # self._filter_ = _filter_
+
+            def filter(self, data: Dict[str, Any]) -> bool:
+                images = [data['image']]
+                texts = [data['caption']]
+                inputs = self.processor(text=texts, images=images, return_tensors="pt", padding=True, truncation=True)
+                # result = self._filter_(
+                #     pixel_values=inputs['pixel_values'],
+                #     input_ids=inputs['input_ids']
+                # )
+                # return result
+                
+                image_embeds = self.model.get_image_features(pixel_values=inputs['pixel_values'])
+                text_embeds = self.model.get_text_features(input_ids=inputs['input_ids'])
+                similarity = F.cosine_similarity(image_embeds, text_embeds)
+                # Filter out images with similarity less than 0.25
+                return similarity[0] >= 0.25
+                
+        return FilterTransform
 
 
 # ----------------------------------------------------------------------------------
@@ -333,3 +400,8 @@ def gcs_augmenters(image_scale, method):
     """Legacy function for GCS augmenters."""
     augmenter = ImageGCSAugmenter()
     return augmenter.create_transform(image_scale=image_scale, method=method)
+
+def gcs_filters(image_scale):
+    """Legacy function for GCS Filters."""
+    augmenter = ImageGCSAugmenter()
+    return augmenter.create_filter(image_scale=image_scale)

flaxdiff/data/sources/videos.py

@@ -216,6 +216,11 @@ class AudioVideoAugmenter(DataAugmenter):
         
         return AudioVideoTransform
 
+    
+    def create_filter(self, image_scale: int = 256):
+        class FilterTransform(pygrain.FilterTransform):
+            def map(self, element) -> bool:
+                return True
 
 # ----------------------------------------------------------------------------------
 # Helper functions for video datasets

flaxdiff/inference/utils.py

@@ -25,6 +25,9 @@ from flaxdiff.models.autoencoder.diffusers import StableDiffusionVAE
 from flaxdiff.inputs import DiffusionInputConfig, ConditionalInputConfig
 from flaxdiff.utils import defaultTextEncodeModel
 
+from flaxdiff.models.simple_vit import UViT, SimpleUDiT
+from flaxdiff.models.simple_dit import SimpleDiT
+from flaxdiff.models.simple_mmdit import SimpleMMDiT, HierarchicalMMDiT
 from orbax.checkpoint import CheckpointManager, CheckpointManagerOptions, PyTreeCheckpointer
 import os
 
@@ -116,7 +119,10 @@ def parse_config(config, overrides=None):
     MODEL_CLASSES = {
         'unet': Unet,
         'uvit': UViT,
-        'diffusers_unet_simple': FlaxUNet2DConditionModel
+        'diffusers_unet_simple': FlaxUNet2DConditionModel,
+        'simple_dit': SimpleDiT,
+        'simple_uvit': SimpleUDiT,
+        'simple_mmdit': SimpleMMDiT,
     }
     
     # Map all the leaves of the model config, converting strings to appropriate types

flaxdiff/models/common.py

@@ -335,73 +335,4 @@ class ResidualBlock(nn.Module):
 
         out = jnp.concatenate([out, extra_features], axis=-1) if extra_features is not None else out
 
-        return out
-
-# Convert Hilbert index d to 2D coordinates (x, y) for an n x n grid
-def _d2xy(n, d):
-    x = 0
-    y = 0
-    t = d
-    s = 1
-    while s < n:
-        rx = (t // 2) & 1
-        ry = (t ^ rx) & 1
-        if ry == 0:
-            if rx == 1:
-                x = n - 1 - x
-                y = n - 1 - y
-            x, y = y, x
-        x += s * rx
-        y += s * ry
-        t //= 4
-        s *= 2
-    return x, y
-
-# Hilbert index mapping for a rectangular grid of patches H_P x W_P
-
-def hilbert_indices(H_P, W_P):
-    size = max(H_P, W_P)
-    order = math.ceil(math.log2(size))
-    n = 1 << order
-    coords = []
-    for d in range(n * n):
-        x, y = _d2xy(n, d)
-        # x is column index, y is row index
-        if x < W_P and y < H_P:
-            coords.append((y, x))  # (row, col)
-            if len(coords) == H_P * W_P:
-                break
-    # Convert (row, col) to linear indices row-major
-    indices = [r * W_P + c for r, c in coords]
-    return jnp.array(indices, dtype=jnp.int32)
-
-# Inverse permutation: given idx where idx[i] = new position of element i, return inv such that inv[idx[i]] = i
-
-def inverse_permutation(idx):
-    inv = jnp.zeros_like(idx)
-    inv = inv.at[idx].set(jnp.arange(idx.shape[0], dtype=idx.dtype))
-    return inv
-
-# Patchify using Hilbert ordering: extract patches and reorder sequence
-
-def hilbert_patchify(x, patch_size):
-    B, H, W, C = x.shape
-    H_P = H // patch_size
-    W_P = W // patch_size
-    # Extract patches in row-major
-    patches = rearrange(x, 'b (h p1) (w p2) c -> b (h w) (p1 p2 c)', p1=patch_size, p2=patch_size)
-    idx = hilbert_indices(H_P, W_P)
-    return patches[:, idx, :]
-
-# Unpatchify from Hilbert ordering: reorder sequence back and reconstruct image
-
-def hilbert_unpatchify(patches, patch_size, H, W, C):
-    B, N, D = patches.shape
-    H_P = H // patch_size
-    W_P = W // patch_size
-    inv = inverse_permutation(hilbert_indices(H_P, W_P))
-    # Reorder back to row-major
-    linear = patches[:, inv, :]
-    # Reconstruct image
-    x = rearrange(linear, 'b (h w) (p1 p2 c) -> b (h p1) (w p2) c', h=H_P, w=W_P, p1=patch_size, p2=patch_size, c=C)
-    return x
+        return out