flaxdiff 0.1.36__tar.gz → 0.1.36.1__tar.gz

{flaxdiff-0.1.36 → flaxdiff-0.1.36.1}/PKG-INFO

@@ -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 → flaxdiff-0.1.36.1}/flaxdiff.egg-info/PKG-INFO

@@ -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/flaxdiff.egg-info/SOURCES.txt

@@ -0,0 +1,9 @@
+README.md
+pyproject.toml
+flaxdiff/__init__.py
+flaxdiff/utils.py
+flaxdiff.egg-info/PKG-INFO
+flaxdiff.egg-info/SOURCES.txt
+flaxdiff.egg-info/dependency_links.txt
+flaxdiff.egg-info/requires.txt
+flaxdiff.egg-info/top_level.txt

flaxdiff-0.1.36.1/flaxdiff.egg-info/requires.txt

@@ -0,0 +1,14 @@
+flax>=0.8.4
+jax>=0.4.28
+optax>=0.2.2
+orbax
+numpy
+clu
+einops
+tqdm
+grain
+termcolor
+augmax
+albumentations
+rich
+python-dotenv

flaxdiff-0.1.36.1/pyproject.toml

@@ -0,0 +1,32 @@
+[build-system]
+requires = ["setuptools", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "flaxdiff"
+version = "0.1.36.1"
+description = "A versatile and easy to understand Diffusion library"
+readme = "README.md"
+authors = [
+  { name="Ashish Kumar Singh", email="ashishkmr472@gmail.com" }
+]
+dependencies = [
+    "flax>=0.8.4",
+    "jax>=0.4.28",
+    "optax>=0.2.2",
+    "orbax",
+    "numpy",
+    "clu",
+    "einops",
+    "tqdm",
+    "grain",
+    "termcolor",
+    "augmax",
+    "albumentations",
+    "rich",
+    "python-dotenv",
+]
+license = "MIT"
+
+[tool.setuptools]
+packages = ["flaxdiff"]

flaxdiff-0.1.36/flaxdiff/data/__init__.py

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

flaxdiff-0.1.36/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-0.1.36/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,
-    }