flaxdiff 0.1.38.1__py3-none-any.whl → 0.2.0__py3-none-any.whl

flaxdiff/data/sources/gcs.py

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

flaxdiff/data/sources/tfds.py

@@ -1,79 +0,0 @@
-import cv2
-import jax.numpy as jnp
-import grain.python as pygrain
-from flaxdiff.utils import AutoTextTokenizer
-from typing import Dict
-import random
-import augmax
-import jax
-
-# -----------------------------------------------------------------------------------------------#
-# Oxford flowers and other TFDS datasources -----------------------------------------------------#
-# -----------------------------------------------------------------------------------------------#
-
-PROMPT_TEMPLATES = [
-    "a photo of a {}",
-    "a photo of a {} flower",
-    "This is a photo of a {}",
-    "This is a photo of a {} flower",
-    "A photo of a {} flower",
-]
-
-def data_source_tfds(name, use_tf=True, split="all"):
-    import tensorflow_datasets as tfds
-    if use_tf:
-        def data_source(path_override):
-            return tfds.load(name, split=split, shuffle_files=True)
-    else:
-        def data_source(path_override):
-            return tfds.data_source(name, split=split, try_gcs=False)
-    return data_source
-
-def labelizer_oxford_flowers102(path):
-    with open(path, "r") as f:
-        textlabels = [i.strip() for i in f.readlines()]
-
-    def load_labels(sample):
-        raw = textlabels[int(sample['label'])]
-        # randomly select a prompt template
-        template = random.choice(PROMPT_TEMPLATES)
-        # format the template with the label
-        caption = template.format(raw)
-        # return the caption
-        return caption
-    return load_labels
-
-def tfds_augmenters(image_scale, method):
-    labelizer = labelizer_oxford_flowers102("/home/mrwhite0racle/tensorflow_datasets/oxford_flowers102/2.1.1/label.labels.txt")
-    if image_scale > 256:
-        interpolation = cv2.INTER_CUBIC
-    else:
-        interpolation = cv2.INTER_AREA
-        
-    from torchvision.transforms import v2
-
-    augments = v2.Compose([
-        v2.RandomHorizontalFlip(p=0.5),
-        v2.ColorJitter(brightness=0.2, contrast=0.05, saturation=0.2)
-    ])
-        
-    class augmenters(pygrain.MapTransform):
-        def __init__(self, *args, **kwargs):
-            super().__init__(*args, **kwargs)
-            self.tokenize = AutoTextTokenizer(tensor_type="np")
-
-        def map(self, element) -> Dict[str, jnp.array]:
-            image = element['image']
-            image = cv2.resize(image, (image_scale, image_scale),
-                            interpolation=interpolation)
-            image = augments(image)
-            # image = (image - 127.5) / 127.5
-            
-            caption = labelizer(element)
-            results = self.tokenize(caption)
-            return {
-                "image": image,
-                "input_ids": results['input_ids'][0],
-                "attention_mask": results['attention_mask'][0],
-            }
-    return augmenters

flaxdiff/trainer/video_diffusion_trainer.py

@@ -1,62 +0,0 @@
-import flax
-from flax import linen as nn
-import jax
-from typing import Callable
-from dataclasses import field
-import jax.numpy as jnp
-import optax
-import functools
-from jax.sharding import Mesh, PartitionSpec as P
-from jax.experimental.shard_map import shard_map
-from typing import Dict, Callable, Sequence, Any, Union, Tuple
-
-from ..schedulers import NoiseScheduler
-from ..predictors import DiffusionPredictionTransform, EpsilonPredictionTransform
-
-from flaxdiff.utils import RandomMarkovState
-
-from .simple_trainer import SimpleTrainer, SimpleTrainState, Metrics
-
-from flaxdiff.models.autoencoder.autoencoder import AutoEncoder
-from flax.training import dynamic_scale as dynamic_scale_lib
-
-class TrainState(SimpleTrainState):
-    rngs: jax.random.PRNGKey
-    ema_params: dict
-
-    def apply_ema(self, decay: float = 0.999):
-        new_ema_params = jax.tree_util.tree_map(
-            lambda ema, param: decay * ema + (1 - decay) * param,
-            self.ema_params,
-            self.params,
-        )
-        return self.replace(ema_params=new_ema_params)
-
-from flaxdiff.models.autoencoder.autoencoder import AutoEncoder
-from flaxdiff.trainer.diffusion_trainer import DiffusionTrainer
-
-class SimpleVideoDiffusionTrainer(DiffusionTrainer):
-    def __init__(self,
-                 model: nn.Module,
-                 input_shapes: Dict[str, Tuple[int]],
-                 optimizer: optax.GradientTransformation,
-                 noise_schedule: NoiseScheduler,
-                 rngs: jax.random.PRNGKey,
-                 unconditional_prob: float = 0.12,
-                 name: str = "SimpleVideoDiffusion",
-                 model_output_transform: DiffusionPredictionTransform = EpsilonPredictionTransform(),
-                 autoencoder: AutoEncoder = None,
-                 **kwargs
-                 ):
-        super().__init__(
-            model=model,
-            input_shapes=input_shapes,
-            optimizer=optimizer,
-            noise_schedule=noise_schedule,
-            unconditional_prob=unconditional_prob,
-            autoencoder=autoencoder,
-            model_output_transform=model_output_transform,
-            rngs=rngs,
-            name=name,
-            **kwargs
-        )

flaxdiff-0.1.38.1.dist-info/RECORD

@@ -1,50 +0,0 @@
-flaxdiff/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-flaxdiff/utils.py,sha256=b_hFXsam2NICQYCFk0EOcqtBjM-RUqnN0NKTn0lQ070,6532
-flaxdiff/data/__init__.py,sha256=PM3PkHihyohT5SHVYKc8vQ4IeVfGPpCktkSVwvqMjQ4,52
-flaxdiff/data/dataset_map.py,sha256=hcHaoR2IbNQmfyPUhYd6_8xinurxxCqawQijAsDI0Ek,3093
-flaxdiff/data/datasets.py,sha256=YUMoSvF2yAyikRvRofZVlHwfEOU3zXSSG4KkLnVfpoA,5626
-flaxdiff/data/online_loader.py,sha256=1Fi_QRixxRzbt602nORINcDeHEccvCrBpagrz4PURYg,12499
-flaxdiff/data/sources/gcs.py,sha256=11ZuQhvMyJRLg21DgVdzO5qEuae7zgzTXGNOskF-cbs,3380
-flaxdiff/data/sources/tfds.py,sha256=k6IzCXPGnoeekvB8Ul0iFBQECz1Wq6LJd2DAjA7ULHI,2761
-flaxdiff/metrics/inception.py,sha256=a5kjMCPMT9gB88c_HCKiek-2vsAyoE35K7nDt4h4pVI,31843
-flaxdiff/metrics/psnr.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-flaxdiff/metrics/ssim.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-flaxdiff/metrics/utils.py,sha256=YuuOfqvqgIjsceupwNeJ59vQ2TnGeNMIyKdkIqOmoNg,1702
-flaxdiff/models/__init__.py,sha256=FAivVYXxM2JrCFIXf-C3374RB2Hth25dBrzOeNFhH1U,26
-flaxdiff/models/attention.py,sha256=-q3xqWy4vQSLG4vXtiUN3FHVBIo7ZjpQsdLT9CkML6c,13367
-flaxdiff/models/common.py,sha256=7x9o5vY9UZvN4BNZ7LHzyuU3PNpsNym9B3m1Wfdddjo,10320
-flaxdiff/models/favor_fastattn.py,sha256=79Ew1nqarsNLPzZaBSd1ILORzJr74CupYeqGiCQK5E4,27689
-flaxdiff/models/general.py,sha256=7xMME6KVKQY8sScyHYH4f-Kek4j1pRfplKShFXwVZd4,587
-flaxdiff/models/simple_unet.py,sha256=hXdpAIA24ARfe53cufJb5Dl4CXPn1bIfWQ0iN6WPNfo,10763
-flaxdiff/models/simple_vit.py,sha256=RaWA85EmEyfquXvnutYQjImDxPEvJ6CQUoSufEjS9aU,7498
-flaxdiff/models/autoencoder/__init__.py,sha256=qY-7MldZpsfkF-_T2LqlRK7VHbqfmosz0NmvzDlBkOk,78
-flaxdiff/models/autoencoder/autoencoder.py,sha256=27_hYl0yXAdH9Mx4Xu9J79mSNo-FEKr9SxhVaS3ffn4,591
-flaxdiff/models/autoencoder/diffusers.py,sha256=DVWT4LRMvEtN36Yt0FTD0KzG8Isq_BvHkNpgDy6Gs40,3651
-flaxdiff/models/autoencoder/simple_autoenc.py,sha256=UXHPgDmwGTnv3Uts6Zj3p9R9nJXnEiEXbllgarwDfXM,805
-flaxdiff/predictors/__init__.py,sha256=A6lMnRSqDLwBdwB0a1x4xPpcXcTYCISWFMkldbiK6Vs,4554
-flaxdiff/samplers/__init__.py,sha256=EY9v1pgwEoR64Kiz9K8fAR-4_ir9c03mYeY3hrpUNhE,308
-flaxdiff/samplers/common.py,sha256=wkzalSYrnsq6oUsevEeRCVfzqwk8qfwvggAlgNTqK-o,8848
-flaxdiff/samplers/ddim.py,sha256=hTjDm0SmIj-Tkc80QRATMcN_sKVhHbqZQboRQCAn4mY,569
-flaxdiff/samplers/ddpm.py,sha256=JgkNSo7fp7Jm-8rCy4eu5m4YIzXTWzxv-iHf3EQ0z5w,2243
-flaxdiff/samplers/euler.py,sha256=QSkttB4DYnepDGwhWq3EGXYjMAqj4qLOdh7u98HttzY,2791
-flaxdiff/samplers/heun_sampler.py,sha256=EvR3hy4t_D47ZOH4luzRFqPmv2v4z78P_JhqBGEpHU8,1436
-flaxdiff/samplers/multistep_dpm.py,sha256=2M4Abb93-GUVN1f0_ZHBeA6lF0eF15Hi6QOgOu2K45s,2752
-flaxdiff/samplers/rk4_sampler.py,sha256=vcQefFhOUZdNOQGBdzNkb2NgmTC2KWd_nhUhyLtt3yI,2026
-flaxdiff/schedulers/__init__.py,sha256=EIva9gBz3DKHORuGmv1LQCKTtRqCRavFOXMNqxAR_ks,131
-flaxdiff/schedulers/common.py,sha256=PDeje2NmN7X3J5qKGauE0jYPpxjgEX44f_evJHRIG3E,4382
-flaxdiff/schedulers/continuous.py,sha256=5c_niOA20fxJ5oJDi09FfayIRogBGwtfG0XThW2IUZk,334
-flaxdiff/schedulers/cosine.py,sha256=E5pODAmINfdyC4kSYOJSPAvq3GNlKPpKEn3X82vYMz0,2055
-flaxdiff/schedulers/discrete.py,sha256=m1q3bAgeAxU3gTj5di3XFWDm4yLfMKAFJPlYdozLE2Y,3316
-flaxdiff/schedulers/exp.py,sha256=cPTnUJpYdzJRRZqMLYQz0rRUCpEmaP2tXhRumLx94jA,605
-flaxdiff/schedulers/karras.py,sha256=4GN120kGwdxxU-h2mVdhBVy9IORkUMm_vvz3XjthBcI,3355
-flaxdiff/schedulers/linear.py,sha256=pBDTXSQcOS4Z03JTh6S0f9E2qLcTQzF2E-pGoQnRoy0,572
-flaxdiff/schedulers/sqrt.py,sha256=mCd_szmOqF6vqQKiAiEOqV_3eBIPGYrW3VxK0o4rBuo,474
-flaxdiff/trainer/__init__.py,sha256=T-vUVq4zHcMK6kpCsG4Gu8vn71q6lZD-lg-Ul7yKfEk,128
-flaxdiff/trainer/autoencoder_trainer.py,sha256=hxihkRL9WCIQVGOP-pc1jjjIUaRXDLcNo3_erTKsuWM,7049
-flaxdiff/trainer/diffusion_trainer.py,sha256=is7iBV8QnL1FZNmXmmqz7K6cnNGF0LaXZlmFV1lPwHU,14185
-flaxdiff/trainer/simple_trainer.py,sha256=NzpCQZlp3pZ7jPxBnYTzPjP1Oh2neD1KHirUiUpcNH0,23222
-flaxdiff/trainer/video_diffusion_trainer.py,sha256=gMkKpnKNTo8QhTx5ptEEkc7W5-7rzXIr9queU53hXyQ,2197
-flaxdiff-0.1.38.1.dist-info/METADATA,sha256=PFvy6QShR5ldz7tyAjQTehexRUg9DK7hRwy9S_ceVuI,23985
-flaxdiff-0.1.38.1.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-flaxdiff-0.1.38.1.dist-info/top_level.txt,sha256=-2-nXnfkJgSfkki1tjm5Faw6Dso7vhtdn2szwCdX5CQ,9
-flaxdiff-0.1.38.1.dist-info/RECORD,,