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

{flaxdiff-0.1.36.dist-info → flaxdiff-0.1.36.1.dist-info}/METADATA

@@ -1,21 +1,24 @@
 Metadata-Version: 2.4
 Name: flaxdiff
-Version: 0.1.36
+Version: 0.1.36.1
 Summary: A versatile and easy to understand Diffusion library
-Author: Ashish Kumar Singh
-Author-email: ashishkmr472@gmail.com
+Author-email: Ashish Kumar Singh <ashishkmr472@gmail.com>
+License-Expression: MIT
 Description-Content-Type: text/markdown
 Requires-Dist: flax>=0.8.4
-Requires-Dist: optax>=0.2.2
 Requires-Dist: jax>=0.4.28
+Requires-Dist: optax>=0.2.2
 Requires-Dist: orbax
+Requires-Dist: numpy
 Requires-Dist: clu
-Dynamic: author
-Dynamic: author-email
-Dynamic: description
-Dynamic: description-content-type
-Dynamic: requires-dist
-Dynamic: summary
+Requires-Dist: einops
+Requires-Dist: tqdm
+Requires-Dist: grain
+Requires-Dist: termcolor
+Requires-Dist: augmax
+Requires-Dist: albumentations
+Requires-Dist: rich
+Requires-Dist: python-dotenv
 
 # ![](images/logo.jpeg "FlaxDiff")
 

flaxdiff-0.1.36.1.dist-info/RECORD

@@ -0,0 +1,6 @@
+flaxdiff/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+flaxdiff/utils.py,sha256=b_hFXsam2NICQYCFk0EOcqtBjM-RUqnN0NKTn0lQ070,6532
+flaxdiff-0.1.36.1.dist-info/METADATA,sha256=Fl9tlGh_BgRnT-f8k4cEYnFj7G03VecUNOX_1zbJrmE,22310
+flaxdiff-0.1.36.1.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+flaxdiff-0.1.36.1.dist-info/top_level.txt,sha256=-2-nXnfkJgSfkki1tjm5Faw6Dso7vhtdn2szwCdX5CQ,9
+flaxdiff-0.1.36.1.dist-info/RECORD,,

flaxdiff/data/__init__.py

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

flaxdiff/data/dataset_map.py

@@ -1,71 +0,0 @@
-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

@@ -1,169 +0,0 @@
-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

@@ -1,363 +0,0 @@
-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/models/__init__.py

@@ -1 +0,0 @@
-from .simple_unet import *