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

flaxdiff/utils.py

@@ -2,7 +2,12 @@ import jax
 import jax.numpy as jnp
 import flax.struct as struct
 import flax.linen as nn
-from typing import Any
+from typing import Any, Callable
+from dataclasses import dataclass
+from functools import partial
+import numpy as np
+from jax.sharding import Mesh, PartitionSpec as P
+from abc import ABC, abstractmethod
 
 class MarkovState(struct.PyTreeNode):
     pass
@@ -17,6 +22,30 @@ class RandomMarkovState(MarkovState):
 def clip_images(images, clip_min=-1, clip_max=1):
     return jnp.clip(images, clip_min, clip_max)
 
+def _build_global_shape_and_sharding(
+    local_shape: tuple[int, ...], global_mesh: Mesh
+) -> tuple[tuple[int, ...], jax.sharding.NamedSharding]:
+  sharding = jax.sharding.NamedSharding(global_mesh, P(global_mesh.axis_names))
+  global_shape = (jax.process_count() * local_shape[0],) + local_shape[1:]
+  return global_shape, sharding
+
+
+def form_global_array(path, array: np.ndarray, global_mesh: Mesh) -> jax.Array:
+  """Put local sharded array into local devices"""
+  global_shape, sharding = _build_global_shape_and_sharding(np.shape(array), global_mesh)
+  try:
+    local_device_arrays = np.split(array, len(global_mesh.local_devices), axis=0)
+  except ValueError as array_split_error:
+    raise ValueError(
+        f"Unable to put to devices shape {array.shape} with "
+        f"local device count {len(global_mesh.local_devices)} "
+    ) from array_split_error
+  local_device_buffers = jax.device_put(local_device_arrays, global_mesh.local_devices)
+  return jax.make_array_from_single_device_arrays(global_shape, sharding, local_device_buffers)
+
+def convert_to_global_tree(global_mesh, pytree):
+    return jax.tree_util.tree_map_with_path(partial(form_global_array, global_mesh=global_mesh), pytree)
+
 class RMSNorm(nn.Module):
     """
     From "Root Mean Square Layer Normalization" by https://arxiv.org/abs/1910.07467
@@ -86,4 +115,78 @@ class RMSNorm(nn.Module):
             ).reshape(feature_shape)
             mul *= scale
         y = mul * x
-        return jnp.asarray(y, dtype)
+        return jnp.asarray(y, dtype)
+
+@dataclass
+class ConditioningEncoder(ABC):
+    model: nn.Module
+    tokenizer: Callable
+
+    def __call__(self, data):
+        tokens = self.tokenize(data)
+        outputs = self.encode_from_tokens(tokens)
+        return outputs
+        
+    def encode_from_tokens(self, tokens):
+        outputs = self.model(input_ids=tokens['input_ids'],
+                        attention_mask=tokens['attention_mask'])
+        last_hidden_state = outputs.last_hidden_state
+        return last_hidden_state
+    
+    def tokenize(self, data):
+        tokens = self.tokenizer(data, padding="max_length",
+                        max_length=self.tokenizer.model_max_length, truncation=True, return_tensors="np")
+        return tokens
+    
+@dataclass
+class TextEncoder(ConditioningEncoder):
+    # def __call__(self, data):
+    #     tokens = self.tokenize(data)
+    #     outputs = self.encode_from_tokens(tokens)
+    #     return outputs
+        
+    # def encode_from_tokens(self, tokens):
+    #     outputs = self.model(input_ids=tokens['input_ids'],
+    #                     attention_mask=tokens['attention_mask'])
+    #     last_hidden_state = outputs.last_hidden_state
+    #     # pooler_output = outputs.pooler_output  # pooled (EOS token) states
+    #     # embed_pooled = pooler_output  # .astype(jnp.float16)
+    #     embed_labels_full = last_hidden_state  # .astype(jnp.float16)
+
+    #     return embed_labels_full
+    pass
+
+class AutoTextTokenizer:
+    def __init__(self, tensor_type="pt", modelname="openai/clip-vit-large-patch14"):
+        from transformers import AutoTokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(modelname)
+        self.tensor_type = tensor_type
+
+    def __call__(self, inputs):
+        # print(caption)
+        tokens = self.tokenizer(inputs, padding="max_length", max_length=self.tokenizer.model_max_length,
+                                truncation=True, return_tensors=self.tensor_type)
+        # print(tokens.keys())
+        return {
+            "input_ids": tokens["input_ids"],
+            "attention_mask": tokens["attention_mask"],
+            "caption": inputs,
+        }
+
+    def __repr__(self):
+        return self.__class__.__name__ + '()'
+
+def defaultTextEncodeModel(backend="jax"):
+    from transformers import (
+        CLIPTextModel,
+        FlaxCLIPTextModel,
+        AutoTokenizer,
+    )
+    modelname = "openai/clip-vit-large-patch14"
+    if backend == "jax":
+        model = FlaxCLIPTextModel.from_pretrained(
+            modelname, dtype=jnp.bfloat16)
+    else:
+        model = CLIPTextModel.from_pretrained(modelname)
+    tokenizer = AutoTokenizer.from_pretrained(modelname, dtype=jnp.float16)
+    return TextEncoder(model, tokenizer)

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

@@ -1,15 +1,24 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.4
 Name: flaxdiff
-Version: 0.1.35.6
+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: flax>=0.8.4
+Requires-Dist: jax>=0.4.28
+Requires-Dist: optax>=0.2.2
 Requires-Dist: orbax
+Requires-Dist: numpy
 Requires-Dist: clu
+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-0.1.35.6.dist-info → flaxdiff-0.1.36.1.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.1.0)
+Generator: setuptools (78.1.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

flaxdiff/data/__init__.py

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

flaxdiff/data/online_loader.py

@@ -1,336 +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,
-):
-    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(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
-
-
-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):
-    return {
-        "url": sample["url"],
-        "caption": sample["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 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.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 *