flaxdiff 0.1.36.1__py3-none-any.whl → 0.1.36.3__py3-none-any.whl

flaxdiff/data/__init__.py

@@ -0,0 +1 @@
+from .online_loader import OnlineStreamingDataLoader

flaxdiff/data/dataset_map.py

@@ -0,0 +1,71 @@
+from .sources.tfds import data_source_tfds, tfds_augmenters
+from .sources.gcs import data_source_gcs, data_source_combined_gcs, gcs_augmenters
+
+# Configure the following for your datasets
+datasetMap = {
+    "oxford_flowers102": {
+        "source": data_source_tfds("oxford_flowers102", use_tf=False),
+        "augmenter": tfds_augmenters,
+    },
+    "cc12m": {
+        "source": data_source_gcs('arrayrecord2/cc12m'),
+        "augmenter": gcs_augmenters,
+    },
+    "laiona_coco": {
+        "source": data_source_gcs('arrayrecord2/laion-aesthetics-12m+mscoco-2017'),
+        "augmenter": gcs_augmenters,
+    },
+    "aesthetic_coyo": {
+        "source": data_source_gcs('arrayrecords/aestheticCoyo_0.25clip_6aesthetic'),
+        "augmenter": gcs_augmenters,
+    },
+    "combined_aesthetic": {
+        "source": data_source_combined_gcs([
+                'arrayrecord2/laion-aesthetics-12m+mscoco-2017',
+                'arrayrecords/aestheticCoyo_0.25clip_6aesthetic',
+                'arrayrecord2/cc12m',
+                'arrayrecords/aestheticCoyo_0.25clip_6aesthetic',
+            ]),
+        "augmenter": gcs_augmenters,
+    },
+    "laiona_coco_coyo": {
+        "source": data_source_combined_gcs([
+                'arrayrecords/aestheticCoyo_0.25clip_6aesthetic',
+                'arrayrecord2/laion-aesthetics-12m+mscoco-2017',
+                'arrayrecords/aestheticCoyo_0.25clip_6aesthetic',
+            ]),
+        "augmenter": gcs_augmenters,
+    },
+    "combined_30m": {
+        "source": data_source_combined_gcs([
+                'arrayrecord2/laion-aesthetics-12m+mscoco-2017',
+                'arrayrecord2/cc12m',
+                'arrayrecord2/aestheticCoyo_0.26_clip_5.5aesthetic_256plus',
+                "arrayrecord2/playground+leonardo_x4+cc3m.parquet",
+            ]),
+        "augmenter": gcs_augmenters,
+    }
+}
+
+onlineDatasetMap = {
+    "combined_online": {
+        "source": [
+            # "gs://flaxdiff-datasets-regional/datasets/laion-aesthetics-12m+mscoco-2017.parquet"
+            # "ChristophSchuhmann/MS_COCO_2017_URL_TEXT",
+            # "dclure/laion-aesthetics-12m-umap",
+            "gs://flaxdiff-datasets-regional/datasets/laion-aesthetics-12m+mscoco-2017",
+            "gs://flaxdiff-datasets-regional/datasets/coyo700m-aesthetic-5.4_25M",
+            "gs://flaxdiff-datasets-regional/datasets/leonardo-liked-1.8m",
+            "gs://flaxdiff-datasets-regional/datasets/leonardo-liked-1.8m",
+            "gs://flaxdiff-datasets-regional/datasets/leonardo-liked-1.8m",
+            "gs://flaxdiff-datasets-regional/datasets/cc12m",
+            "gs://flaxdiff-datasets-regional/datasets/playground-liked",
+            "gs://flaxdiff-datasets-regional/datasets/leonardo-liked-1.8m",
+            "gs://flaxdiff-datasets-regional/datasets/leonardo-liked-1.8m",
+            "gs://flaxdiff-datasets-regional/datasets/cc3m",
+            "gs://flaxdiff-datasets-regional/datasets/cc3m",
+            "gs://flaxdiff-datasets-regional/datasets/laion2B-en-aesthetic-4.2_37M",
+            # "gs://flaxdiff-datasets-regional/datasets/laiion400m-185M"
+        ]
+    }
+}

flaxdiff/data/datasets.py

@@ -0,0 +1,169 @@
+import jax.numpy as jnp
+import grain.python as pygrain
+from typing import Dict
+import numpy as np
+import jax
+from flaxdiff.utils import convert_to_global_tree, AutoTextTokenizer
+from .dataset_map import datasetMap, onlineDatasetMap
+import traceback
+from .online_loader import OnlineStreamingDataLoader
+import queue
+from jax.sharding import Mesh
+import threading
+
+def batch_mesh_map(mesh):
+    class augmenters(pygrain.MapTransform):
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, **kwargs)
+
+        def map(self, batch) -> Dict[str, jnp.array]:
+            return convert_to_global_tree(mesh, batch)
+    return augmenters
+
+def get_dataset_grain(
+    data_name="cc12m",
+    batch_size=64,
+    image_scale=256,
+    count=None,
+    num_epochs=None,
+    method=jax.image.ResizeMethod.LANCZOS3,
+    worker_count=32,
+    read_thread_count=64,
+    read_buffer_size=50,
+    worker_buffer_size=20,
+    seed=0,
+    dataset_source="/mnt/gcs_mount/arrayrecord2/cc12m/",
+):
+    dataset = datasetMap[data_name]
+    data_source = dataset["source"](dataset_source)
+    augmenter = dataset["augmenter"](image_scale, method)
+
+    local_batch_size = batch_size // jax.process_count()
+
+    sampler = pygrain.IndexSampler(
+        num_records=len(data_source) if count is None else count,
+        shuffle=True,
+        seed=seed,
+        num_epochs=num_epochs,
+        shard_options=pygrain.ShardByJaxProcess(),
+    )
+
+    def get_trainset():
+        transformations = [
+            augmenter(),
+            pygrain.Batch(local_batch_size, drop_remainder=True),
+        ]
+        
+        # if mesh != None:
+        #     transformations += [batch_mesh_map(mesh)]
+
+        loader = pygrain.DataLoader(
+            data_source=data_source,
+            sampler=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
+    
+
+    return {
+        "train": get_trainset,
+        "train_len": len(data_source),
+        "local_batch_size": local_batch_size,
+        "global_batch_size": batch_size,
+        # "null_labels": null_labels,
+        # "null_labels_full": null_labels_full,
+        # "model": model,
+        # "tokenizer": tokenizer,
+    }
+
+def generate_collate_fn():
+    auto_tokenize = AutoTextTokenizer(tensor_type="np")
+    def default_collate(batch):
+        try:
+            # urls = [sample["url"] for sample in batch]
+            captions = [sample["caption"] for sample in batch]
+            results = auto_tokenize(captions)
+            images = np.stack([sample["image"] for sample in batch], axis=0)
+            return {
+                "image": images,
+                "input_ids": results['input_ids'],
+                "attention_mask": results['attention_mask'],
+            }
+        except Exception as e:
+            print("Error in collate function", e, [sample["image"].shape for sample in batch])
+            traceback.print_exc()
+            
+    return default_collate
+    
+def get_dataset_online(
+        data_name="combined_online",
+        batch_size=64,
+        image_scale=256,
+        count=None,
+        num_epochs=None,
+        method=jax.image.ResizeMethod.LANCZOS3,
+        worker_count=32,
+        read_thread_count=64,
+        read_buffer_size=50,
+        worker_buffer_size=20,
+        seed=0,
+        dataset_source="/mnt/gcs_mount/arrayrecord2/cc12m/",
+    ):
+    local_batch_size = batch_size // jax.process_count()
+    
+    sources = onlineDatasetMap[data_name]["source"]
+    dataloader = OnlineStreamingDataLoader(
+            sources, 
+            batch_size=local_batch_size,
+            num_workers=worker_count,
+            num_threads=read_thread_count,
+            image_shape=(image_scale, image_scale),
+            global_process_count=jax.process_count(),
+            global_process_index=jax.process_index(),
+            prefetch=worker_buffer_size,
+            collate_fn=generate_collate_fn(),
+            default_split="train",
+        )
+    
+    def get_trainset(mesh: Mesh = None):
+        if mesh != None:
+            class dataLoaderWithMesh:
+                def __init__(self, dataloader, mesh):
+                    self.dataloader = dataloader
+                    self.mesh = mesh
+                    self.tmp_queue = queue.Queue(worker_buffer_size)
+                    def batch_loader():
+                        for batch in self.dataloader:
+                            try:
+                                self.tmp_queue.put(convert_to_global_tree(mesh, batch))
+                            except Exception as e:
+                                print("Error processing batch", e)
+                    self.loader_thread = threading.Thread(target=batch_loader)
+                    self.loader_thread.start()
+                    
+                def __iter__(self):
+                    return self
+                
+                def __next__(self):
+                    return self.tmp_queue.get()
+                
+            dataloader_with_mesh = dataLoaderWithMesh(dataloader, mesh)
+                    
+            return dataloader_with_mesh
+        return dataloader
+    
+    return {
+        "train": get_trainset,
+        "train_len": len(dataloader) * jax.process_count(),
+        "local_batch_size": local_batch_size,
+        "global_batch_size": batch_size,
+        # "null_labels": null_labels,
+        # "null_labels_full": null_labels_full,
+        # "model": model,
+        # "tokenizer": tokenizer,
+    }

flaxdiff/data/online_loader.py

@@ -0,0 +1,363 @@
+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, load_from_disk
+from datasets.utils.file_utils import get_datasets_user_agent
+from concurrent.futures import ThreadPoolExecutor
+import io
+import urllib
+
+import PIL.Image
+import cv2
+import traceback
+
+USER_AGENT = get_datasets_user_agent()
+
+data_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 default_image_processor(
+    image, image_shape, 
+    min_image_shape=(128, 128),
+    upscale_interpolation=cv2.INTER_CUBIC,
+    downscale_interpolation=cv2.INTER_AREA,
+):
+    try:
+        image = np.array(image)
+        if len(image.shape) != 3 or image.shape[2] != 3:
+            return None, 0, 0
+        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(
+            image,
+            min_height=image_shape[0],
+            min_width=image_shape[1],
+            border_mode=cv2.BORDER_CONSTANT,
+            value=[255, 255, 255],
+        )
+        return image, original_height, original_width
+    except Exception as e:
+        # print("Error processing image", e, image_shape, interpolation)
+        # traceback.print_exc()
+        return None, 0, 0
+
+
+def map_sample(
+    url,
+    caption,
+    image_shape=(256, 256),
+    min_image_shape=(128, 128),
+    timeout=15,
+    retries=3,
+    upscale_interpolation=cv2.INTER_CUBIC,
+    downscale_interpolation=cv2.INTER_AREA,
+    image_processor=default_image_processor,
+):
+    try:
+        # Assuming fetch_single_image is defined elsewhere
+        image = fetch_single_image(url, timeout=timeout, retries=retries)
+        if image is None:
+            return
+
+        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
+
+        data_queue.put({
+            "url": url,
+            "caption": caption,
+            "image": image,
+            "original_height": original_height,
+            "original_width": original_width,
+        })
+    except Exception as e:
+        # print(f"Error maping sample {url}", e)
+        # traceback.print_exc()
+        # error_queue.put_nowait({
+        #     "url": url,
+        #     "caption": caption,
+        #     "error": str(e)
+        # })
+        pass
+
+def default_feature_extractor(sample):
+    url = None
+    if "url" in sample:
+        url = sample["url"]
+    elif "URL" in sample:
+        url = sample["URL"]
+    elif "image_url" in sample:
+        url = sample["image_url"]
+    else:
+        print("No url found in sample, skipping", sample.keys())
+    
+    caption = None
+    if "caption" in sample:
+        caption = sample["caption"]
+    elif "CAPTION" in sample:
+        caption = sample["CAPTION"]
+    elif "txt" in sample:
+        caption = sample["txt"]
+    elif "TEXT" in sample:
+        caption = sample["TEXT"]
+    elif "text" in sample:
+        caption = sample["text"]
+    else:
+        print("No caption found in sample, skipping", sample.keys())
+        
+    return {
+        "url": url,
+        "caption": caption,
+    }
+    
+def map_batch(
+    batch, num_threads=256, image_shape=(256, 256),
+    min_image_shape=(128, 128),
+    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,
+            timeout=timeout, retries=retries, image_processor=image_processor,
+            upscale_interpolation=upscale_interpolation,
+            downscale_interpolation=downscale_interpolation
+        )
+        with ThreadPoolExecutor(max_workers=num_threads) as executor:
+            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()
+        # error_queue.put_nowait({
+        #     "batch": batch,
+        #     "error": str(e)
+        # })
+        pass
+
+
+def parallel_image_loader(
+    dataset: Dataset, num_workers: int = 8, image_shape=(256, 256),
+    min_image_shape=(128, 128),
+    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,
+        feature_extractor=feature_extractor
+    )
+    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
+            shards = [dataset[i*shard_len:(i+1)*shard_len]
+                      for i in range(num_workers)]
+            print(f"mapping {len(shards)} shards")
+            pool.map(map_batch_fn, shards)
+            iteration += 1
+            print(f"Shuffling dataset with seed {iteration}")
+            dataset = dataset.shuffle(seed=iteration)
+            # Clear the error queue
+            # while not error_queue.empty():
+            #     error_queue.get_nowait()
+
+
+class ImageBatchIterator:
+    def __init__(
+        self, dataset: Dataset, batch_size: int = 64, image_shape=(256, 256),
+        min_image_shape=(128, 128),
+        num_workers: int = 8, 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,
+    ):
+        self.dataset = dataset
+        self.num_workers = num_workers
+        self.batch_size = batch_size
+        loader = partial(
+            parallel_image_loader,
+            num_threads=num_threads,
+            image_shape=image_shape,
+            min_image_shape=min_image_shape,
+            num_workers=num_workers,
+            timeout=timeout, retries=retries,
+            image_processor=image_processor,
+            upscale_interpolation=upscale_interpolation,
+            downscale_interpolation=downscale_interpolation,
+            feature_extractor=feature_extractor
+        )
+        self.thread = threading.Thread(target=loader, args=(dataset,))
+        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,
+        image_shape=(256, 256),
+        min_image_shape=(128, 128),
+        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,
+        timeout=15,
+        retries=3,
+        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
+            print("Loading dataset from path")
+            if "gs://" in dataset:
+                dataset = load_from_disk(dataset_path)
+            else:
+                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_from_disk(dataset_path) if "gs://" in dataset_path else load_dataset(
+                    dataset_path, split=default_split) for dataset_path in dataset]
+            print("Concatenating multiple datasets")
+            dataset = concatenate_datasets(dataset)
+            dataset = dataset.shuffle(seed=0)
+        # dataset = dataset.map(pre_map_maker(pre_map_def), batched=True, batch_size=10000000)
+        self.dataset = dataset.shard(
+            num_shards=global_process_count, index=global_process_index)
+        print(f"Dataset length: {len(dataset)}")
+        self.iterator = ImageBatchIterator(self.dataset, image_shape=image_shape,
+                                           min_image_shape=min_image_shape,
+                                           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,
+                                             feature_extractor=feature_extractor)
+        self.batch_size = batch_size
+
+        # Launch a thread to load batches in the background
+        self.batch_queue = queue.Queue(prefetch)
+
+        def batch_loader():
+            for batch in self.iterator:
+                try:
+                    self.batch_queue.put(collate_fn(batch))
+                except Exception as e:
+                    print("Error collating batch", e)
+
+        self.loader_thread = threading.Thread(target=batch_loader)
+        self.loader_thread.start()
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        return self.batch_queue.get()
+        # return self.collate_fn(next(self.iterator))
+
+    def __len__(self):
+        return len(self.dataset)

flaxdiff/data/sources/gcs.py

@@ -0,0 +1,81 @@
+import cv2
+import jax.numpy as jnp
+import grain.python as pygrain
+from flaxdiff.utils import AutoTextTokenizer
+from typing import Dict
+import os
+import struct as st
+from functools import partial
+import numpy as np
+
+# -----------------------------------------------------------------------------------------------#
+# CC12m and other GCS data sources --------------------------------------------------------------#
+# -----------------------------------------------------------------------------------------------#
+
+def data_source_gcs(source='arrayrecord/laion-aesthetics-12m+mscoco-2017'):
+    def data_source(base="/home/mrwhite0racle/gcs_mount"):
+        records_path = os.path.join(base, source)
+        records = [os.path.join(records_path, i) for i in os.listdir(
+            records_path) if 'array_record' in i]
+        ds = pygrain.ArrayRecordDataSource(records)
+        return ds
+    return data_source
+
+def data_source_combined_gcs(
+    sources=[]):
+    def data_source(base="/home/mrwhite0racle/gcs_mount"):
+        records_paths = [os.path.join(base, source) for source in sources]
+        records = []
+        for records_path in records_paths:
+            records += [os.path.join(records_path, i) for i in os.listdir(
+                records_path) if 'array_record' in i]
+        ds = pygrain.ArrayRecordDataSource(records)
+        return ds
+    return data_source
+
+def unpack_dict_of_byte_arrays(packed_data):
+    unpacked_dict = {}
+    offset = 0
+    while offset < len(packed_data):
+        # Unpack the key length
+        key_length = st.unpack_from('I', packed_data, offset)[0]
+        offset += st.calcsize('I')
+        # Unpack the key bytes and convert to string
+        key = packed_data[offset:offset+key_length].decode('utf-8')
+        offset += key_length
+        # Unpack the byte array length
+        byte_array_length = st.unpack_from('I', packed_data, offset)[0]
+        offset += st.calcsize('I')
+        # Unpack the byte array
+        byte_array = packed_data[offset:offset+byte_array_length]
+        offset += byte_array_length
+        unpacked_dict[key] = byte_array
+    return unpacked_dict
+
+def image_augmenter(image, image_scale, method=cv2.INTER_AREA):
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    image = cv2.resize(image, (image_scale, image_scale),
+                            interpolation=cv2.INTER_AREA)
+    return image
+
+def gcs_augmenters(image_scale, method):
+    labelizer = lambda sample : sample['txt']
+    class augmenters(pygrain.MapTransform):
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, **kwargs)
+            self.auto_tokenize = AutoTextTokenizer(tensor_type="np")
+            self.image_augmenter = partial(image_augmenter, image_scale=image_scale, method=method)
+
+        def map(self, element) -> Dict[str, jnp.array]:
+            element = unpack_dict_of_byte_arrays(element)
+            image = np.asarray(bytearray(element['jpg']), dtype="uint8")
+            image = cv2.imdecode(image, cv2.IMREAD_UNCHANGED)
+            image = self.image_augmenter(image)
+            caption = labelizer(element).decode('utf-8')
+            results = self.auto_tokenize(caption)
+            return {
+                "image": image,
+                "input_ids": results['input_ids'][0],
+                "attention_mask": results['attention_mask'][0],
+            }
+    return augmenters