flaxdiff 0.1.11__tar.gz → 0.1.13__tar.gz

{flaxdiff-0.1.11 → flaxdiff-0.1.13}/PKG-INFO

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

flaxdiff-0.1.13/flaxdiff/data/online_loader.py

@@ -0,0 +1,205 @@
+import multiprocessing
+import threading
+from multiprocessing import Queue
+# from arrayqueues.shared_arrays import ArrayQueue
+# from faster_fifo import Queue
+import time
+import albumentations as A
+import queue
+import cv2
+from functools import partial
+from typing import Any, Dict, List, Tuple
+
+import numpy as np
+from functools import partial
+
+from datasets import load_dataset, concatenate_datasets, Dataset
+from datasets.utils.file_utils import get_datasets_user_agent
+from concurrent.futures import ThreadPoolExecutor
+import io
+import urllib
+
+import PIL.Image
+import cv2
+
+USER_AGENT = get_datasets_user_agent()
+
+data_queue = Queue(16*2000)
+error_queue = Queue(16*2000)
+
+
+def fetch_single_image(image_url, timeout=None, retries=0):
+    for _ in range(retries + 1):
+        try:
+            request = urllib.request.Request(
+                image_url,
+                data=None,
+                headers={"user-agent": USER_AGENT},
+            )
+            with urllib.request.urlopen(request, timeout=timeout) as req:
+                image = PIL.Image.open(io.BytesIO(req.read()))
+            break
+        except Exception:
+            image = None
+    return image
+
+def map_sample(
+    url, caption, 
+    image_shape=(256, 256),
+    upscale_interpolation=cv2.INTER_LANCZOS4,
+    downscale_interpolation=cv2.INTER_AREA,
+):
+    try:
+        image = fetch_single_image(url, timeout=15, retries=3)  # Assuming fetch_single_image is defined elsewhere
+        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(image_shape):
+            return
+        # check if wrong aspect ratio
+        if max(original_height, original_width) / min(original_height, original_width) > 2:
+            return
+        # check if the variance is too low
+        if np.std(image) < 1e-4:
+            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 = A.longest_max_size(image, max(image_shape), interpolation=interpolation)
+        image = A.pad(
+            image,
+            min_height=image_shape[0],
+            min_width=image_shape[1],
+            border_mode=cv2.BORDER_CONSTANT,
+            value=[255, 255, 255],
+        )
+        data_queue.put({
+            "url": url,
+            "caption": caption,
+            "image": image
+        })
+    except Exception as e:
+        error_queue.put({
+            "url": url,
+            "caption": caption,
+            "error": str(e)
+        })
+        
+def map_batch(batch, num_threads=256, timeout=None, retries=0):
+    with ThreadPoolExecutor(max_workers=num_threads) as executor:
+        executor.map(map_sample, batch["url"], batch['caption'])
+    
+def parallel_image_loader(dataset: Dataset, num_workers: int = 8, num_threads=256):
+    map_batch_fn = partial(map_batch, num_threads=num_threads)
+    shard_len = len(dataset) // num_workers
+    print(f"Local Shard lengths: {shard_len}")
+    with multiprocessing.Pool(num_workers) as pool:
+        iteration = 0
+        while True:
+            # Repeat forever
+            dataset = dataset.shuffle(seed=iteration)
+            shards = [dataset[i*shard_len:(i+1)*shard_len] for i in range(num_workers)]
+            pool.map(map_batch_fn, shards)
+            iteration += 1
+            
+class ImageBatchIterator:
+    def __init__(self, dataset: Dataset, batch_size: int = 64, num_workers: int = 8, num_threads=256):
+        self.dataset = dataset
+        self.num_workers = num_workers
+        self.batch_size = batch_size
+        loader = partial(parallel_image_loader, num_threads=num_threads)
+        self.thread = threading.Thread(target=loader, args=(dataset, num_workers))
+        self.thread.start()
+        
+    def __iter__(self):
+        return self
+    
+    def __next__(self):
+        def fetcher(_):
+            return data_queue.get()
+        with ThreadPoolExecutor(max_workers=self.batch_size) as executor:
+            batch = list(executor.map(fetcher, range(self.batch_size)))
+        return batch
+    
+    def __del__(self):
+        self.thread.join()
+        
+    def __len__(self):
+        return len(self.dataset) // self.batch_size
+    
+def default_collate(batch):
+    urls = [sample["url"] for sample in batch]
+    captions = [sample["caption"] for sample in batch]
+    images = np.stack([sample["image"] for sample in batch], axis=0)
+    return {
+        "url": urls,
+        "caption": captions,
+        "image": images,
+    }
+    
+def dataMapper(map: Dict[str, Any]):
+    def _map(sample) -> Dict[str, Any]:
+        return {
+            "url": sample[map["url"]],
+            "caption": sample[map["caption"]],
+        }
+    return _map
+
+class OnlineStreamingDataLoader():
+    def __init__(
+        self, 
+        dataset, 
+        batch_size=64, 
+        num_workers=16, 
+        num_threads=512,
+        default_split="all",
+        pre_map_maker=dataMapper, 
+        pre_map_def={
+            "url": "URL",
+            "caption": "TEXT",
+        },
+        global_process_count=1,
+        global_process_index=0,
+        prefetch=1000,
+        collate_fn=default_collate,
+    ):
+        if isinstance(dataset, str):
+            dataset_path = dataset
+            print("Loading dataset from path")
+            dataset = load_dataset(dataset_path, split=default_split)
+        elif isinstance(dataset, list):
+            if isinstance(dataset[0], str):
+                print("Loading multiple datasets from paths")
+                dataset = [load_dataset(dataset_path, split=default_split) for dataset_path in dataset]
+            else:
+                print("Concatenating multiple datasets")
+                dataset = concatenate_datasets(dataset)
+        dataset = dataset.map(pre_map_maker(pre_map_def))
+        self.dataset = dataset.shard(num_shards=global_process_count, index=global_process_index)
+        print(f"Dataset length: {len(dataset)}")
+        self.iterator = ImageBatchIterator(self.dataset, num_workers=num_workers, batch_size=batch_size, num_threads=num_threads)
+        self.collate_fn = collate_fn
+        
+        # Launch a thread to load batches in the background
+        self.batch_queue = queue.Queue(prefetch)
+        
+        def batch_loader():
+            for batch in self.iterator:
+                self.batch_queue.put(batch)
+        
+        self.loader_thread = threading.Thread(target=batch_loader)
+        self.loader_thread.start()
+        
+    def __iter__(self):
+        return self
+    
+    def __next__(self):
+        return self.collate_fn(self.batch_queue.get())
+        # return self.collate_fn(next(self.iterator))
+        
+    def __len__(self):
+        return len(self.dataset) // self.batch_size
+    

{flaxdiff-0.1.11 → flaxdiff-0.1.13}/flaxdiff/models/attention.py

@@ -22,7 +22,7 @@ class EfficientAttention(nn.Module):
     dtype: Optional[Dtype] = None
     precision: PrecisionLike = None
     use_bias: bool = True
-    kernel_init: Callable = lambda : kernel_init(1.0)
+    kernel_init: Callable = kernel_init(1.0)
     force_fp32_for_softmax: bool = True
 
     def setup(self):
@@ -33,7 +33,7 @@ class EfficientAttention(nn.Module):
             self.heads * self.dim_head,
             precision=self.precision, 
             use_bias=self.use_bias, 
-            kernel_init=self.kernel_init(), 
+            kernel_init=self.kernel_init, 
             dtype=self.dtype
         )
         self.query = dense(name="to_q")
@@ -41,7 +41,7 @@ class EfficientAttention(nn.Module):
         self.value = dense(name="to_v")
         
         self.proj_attn = nn.DenseGeneral(self.query_dim, use_bias=False, precision=self.precision, 
-                                     kernel_init=self.kernel_init(), dtype=self.dtype, name="to_out_0")
+                                     kernel_init=self.kernel_init, dtype=self.dtype, name="to_out_0")
         # self.attnfn = make_fast_generalized_attention(qkv_dim=inner_dim, lax_scan_unroll=16)
     
     def _reshape_tensor_to_head_dim(self, tensor):
@@ -114,7 +114,7 @@ class NormalAttention(nn.Module):
     dtype: Optional[Dtype] = None
     precision: PrecisionLike = None
     use_bias: bool = True
-    kernel_init: Callable = lambda : kernel_init(1.0)
+    kernel_init: Callable = kernel_init(1.0)
     force_fp32_for_softmax: bool = True
 
     def setup(self):
@@ -125,7 +125,7 @@ class NormalAttention(nn.Module):
             axis=-1, 
             precision=self.precision, 
             use_bias=self.use_bias, 
-            kernel_init=self.kernel_init(), 
+            kernel_init=self.kernel_init, 
             dtype=self.dtype
         )
         self.query = dense(name="to_q")
@@ -139,7 +139,7 @@ class NormalAttention(nn.Module):
             use_bias=self.use_bias, 
             dtype=self.dtype, 
             name="to_out_0",
-            kernel_init=self.kernel_init()
+            kernel_init=self.kernel_init
             # kernel_init=jax.nn.initializers.xavier_uniform()
         )
 
@@ -235,7 +235,7 @@ class BasicTransformerBlock(nn.Module):
     dtype: Optional[Dtype] = None
     precision: PrecisionLike = None
     use_bias: bool = True
-    kernel_init: Callable = lambda : kernel_init(1.0)
+    kernel_init: Callable = kernel_init(1.0)
     use_flash_attention:bool = False
     use_cross_only:bool = False
     only_pure_attention:bool = False
@@ -302,7 +302,7 @@ class TransformerBlock(nn.Module):
     use_self_and_cross:bool = True
     only_pure_attention:bool = False
     force_fp32_for_softmax: bool = True
-    kernel_init: Callable = lambda : kernel_init(1.0)
+    kernel_init: Callable = kernel_init(1.0)
 
     @nn.compact
     def __call__(self, x, context=None):
@@ -313,12 +313,12 @@ class TransformerBlock(nn.Module):
             if self.use_linear_attention:
                 projected_x = nn.Dense(features=inner_dim, 
                                        use_bias=False, precision=self.precision, 
-                                       kernel_init=self.kernel_init(),
+                                       kernel_init=self.kernel_init,
                                        dtype=self.dtype, name=f'project_in')(normed_x)
             else:
                 projected_x = nn.Conv(
                     features=inner_dim, kernel_size=(1, 1),
-                    kernel_init=self.kernel_init(),
+                    kernel_init=self.kernel_init,
                     strides=(1, 1), padding='VALID', use_bias=False, dtype=self.dtype,
                     precision=self.precision, name=f'project_in_conv',
                 )(normed_x)
@@ -347,12 +347,12 @@ class TransformerBlock(nn.Module):
             if self.use_linear_attention:
                 projected_x = nn.Dense(features=C, precision=self.precision, 
                                        dtype=self.dtype, use_bias=False, 
-                                       kernel_init=self.kernel_init(),
+                                       kernel_init=self.kernel_init,
                                        name=f'project_out')(projected_x)
             else:
                 projected_x = nn.Conv(
                     features=C, kernel_size=(1, 1),
-                    kernel_init=self.kernel_init(),
+                    kernel_init=self.kernel_init,
                     strides=(1, 1), padding='VALID', use_bias=False, dtype=self.dtype,
                     precision=self.precision, name=f'project_out_conv',
                 )(projected_x)

{flaxdiff-0.1.11 → flaxdiff-0.1.13}/flaxdiff/models/autoencoder/diffusers.py

@@ -11,15 +11,15 @@ All credits for the model go to the developers of Stable Diffusion VAE and all c
 """
 
 class StableDiffusionVAE(AutoEncoder):
-    def __init__(self, modelname = "CompVis/stable-diffusion-v1-4"):
+    def __init__(self, modelname = "CompVis/stable-diffusion-v1-4", revision="bf16", dtype=jnp.bfloat16):
         
         from diffusers.models.vae_flax import FlaxEncoder, FlaxDecoder
         from diffusers import FlaxStableDiffusionPipeline
         
         pipeline, params = FlaxStableDiffusionPipeline.from_pretrained(
             modelname,
-            revision="bf16",
-            dtype=jnp.bfloat16,
+            revision=revision,
+            dtype=dtype,
         )
         
         vae = pipeline.vae

{flaxdiff-0.1.11 → flaxdiff-0.1.13}/flaxdiff.egg-info/PKG-INFO

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

{flaxdiff-0.1.11 → flaxdiff-0.1.13}/flaxdiff.egg-info/SOURCES.txt

@@ -7,6 +7,8 @@ flaxdiff.egg-info/SOURCES.txt
 flaxdiff.egg-info/dependency_links.txt
 flaxdiff.egg-info/requires.txt
 flaxdiff.egg-info/top_level.txt
+flaxdiff/data/__init__.py
+flaxdiff/data/online_loader.py
 flaxdiff/models/__init__.py
 flaxdiff/models/attention.py
 flaxdiff/models/common.py

{flaxdiff-0.1.11 → flaxdiff-0.1.13}/setup.py

@@ -11,7 +11,7 @@ required_packages=[
 setup(
     name='flaxdiff',
     packages=find_packages(),
-    version='0.1.11',
+    version='0.1.13',
     description='A versatile and easy to understand Diffusion library',
     long_description=open('README.md').read(),
     long_description_content_type='text/markdown',