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

data/sources/tfds.py

@@ -0,0 +1,67 @@
+import cv2
+import jax.numpy as jnp
+import grain.python as pygrain
+from flaxdiff.utils import AutoTextTokenizer
+from typing import Dict
+import random
+
+# -----------------------------------------------------------------------------------------------#
+# 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
+    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 = (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-0.1.36.1.dist-info → flaxdiff-0.1.36.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flaxdiff
-Version: 0.1.36.1
+Version: 0.1.36.2
 Summary: A versatile and easy to understand Diffusion library
 Author-email: Ashish Kumar Singh <ashishkmr472@gmail.com>
 License-Expression: MIT

flaxdiff-0.1.36.2.dist-info/RECORD

@@ -0,0 +1,47 @@
+__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+utils.py,sha256=b_hFXsam2NICQYCFk0EOcqtBjM-RUqnN0NKTn0lQ070,6532
+data/__init__.py,sha256=PM3PkHihyohT5SHVYKc8vQ4IeVfGPpCktkSVwvqMjQ4,52
+data/dataset_map.py,sha256=hcHaoR2IbNQmfyPUhYd6_8xinurxxCqawQijAsDI0Ek,3093
+data/datasets.py,sha256=YUMoSvF2yAyikRvRofZVlHwfEOU3zXSSG4KkLnVfpoA,5626
+data/online_loader.py,sha256=1Fi_QRixxRzbt602nORINcDeHEccvCrBpagrz4PURYg,12499
+data/sources/gcs.py,sha256=11ZuQhvMyJRLg21DgVdzO5qEuae7zgzTXGNOskF-cbs,3380
+data/sources/tfds.py,sha256=WA3h9lyR4yotCNEmJON2noIN-2HNcqhf6zigx1XXsMI,2481
+metrics/inception.py,sha256=a5kjMCPMT9gB88c_HCKiek-2vsAyoE35K7nDt4h4pVI,31843
+metrics/utils.py,sha256=YuuOfqvqgIjsceupwNeJ59vQ2TnGeNMIyKdkIqOmoNg,1702
+models/__init__.py,sha256=FAivVYXxM2JrCFIXf-C3374RB2Hth25dBrzOeNFhH1U,26
+models/attention.py,sha256=JvrP7-09MV6IfRLRBhqjPmNUU-lkEMk9TOnJSBKcar8,13289
+models/common.py,sha256=hWsSs2BP2J-JN1s4qLRr-h-KYkcVyl2hOp1Wsm_L-h8,10994
+models/favor_fastattn.py,sha256=79Ew1nqarsNLPzZaBSd1ILORzJr74CupYeqGiCQK5E4,27689
+models/simple_unet.py,sha256=L5m2j5580QP7pJ5VIme7U5xYA22PZiGP7qdvcKUnB38,11463
+models/simple_vit.py,sha256=UCDDr0XVnpf6tbJWKFtEt3_nAqMqOoakXf5amyVWZNo,7929
+models/autoencoder/__init__.py,sha256=qY-7MldZpsfkF-_T2LqlRK7VHbqfmosz0NmvzDlBkOk,78
+models/autoencoder/autoencoder.py,sha256=27_hYl0yXAdH9Mx4Xu9J79mSNo-FEKr9SxhVaS3ffn4,591
+models/autoencoder/diffusers.py,sha256=JHeFLCxiHhu-QHwhKiCuKsQJn4AZumquiuxgZkiYGQ0,3643
+models/autoencoder/simple_autoenc.py,sha256=UXHPgDmwGTnv3Uts6Zj3p9R9nJXnEiEXbllgarwDfXM,805
+predictors/__init__.py,sha256=SKkYYRF9Wfgk2zhtZw4vCXOdOeRlrm2Mk6cvuaEvAzc,4403
+samplers/__init__.py,sha256=_S-9TwDeshrI0VmapV-J2hqjTByOa0-oOeUs_IdovjU,285
+samplers/common.py,sha256=ZA08VyovxegpRx4wOQq9LSwZi0gSCz2lrbS5oVYOEYg,8488
+samplers/ddim.py,sha256=pB8Kod8ZLJ3GXev4uM3cOj1Uy6ibR0jsaZa-VE0fyJM,552
+samplers/ddpm.py,sha256=u1OchQu0XPhc_6w9JXoaFp2wo4y-zXyQNtGAIJwxNLg,2209
+samplers/euler.py,sha256=Htb-IJeu7jSgY6mvgYr9yl9pUnos49vijlVk5IQsRps,2740
+samplers/heun_sampler.py,sha256=UyI-hSlyWvt-7VEUJj27zjgyzKkGVl8fDUHV-YpSOCc,1421
+samplers/multistep_dpm.py,sha256=3Wu3MrMLYaBb1ObraTbWrJmtEtU0adl1dDbz5fPJ4Gs,2735
+samplers/rk4_sampler.py,sha256=1j1pES_Q2QiaURvEWeedbbT1LHmkc3jsu0GgH83qBL0,1926
+schedulers/__init__.py,sha256=3id390WEfdf-MN-oLSPAhlRFIXrFWr6ioAHPAwURJyE,375
+schedulers/common.py,sha256=b-W4iI-aqScpVE8VZbBpiYvAVI6rqDkUP-C_hEVBwCI,4151
+schedulers/continuous.py,sha256=5c_niOA20fxJ5oJDi09FfayIRogBGwtfG0XThW2IUZk,334
+schedulers/cosine.py,sha256=9ban0dFHLMm35wQvaBT4nCQwPGmzNsXwQ1xI0oppmJI,2005
+schedulers/discrete.py,sha256=O54wH2HVu3olJA71NxgAXFW9cr6B6Gl-DR_uZeytpds,3319
+schedulers/exp.py,sha256=cPTnUJpYdzJRRZqMLYQz0rRUCpEmaP2tXhRumLx94jA,605
+schedulers/karras.py,sha256=4GN120kGwdxxU-h2mVdhBVy9IORkUMm_vvz3XjthBcI,3355
+schedulers/linear.py,sha256=6003F5ISq1Wc0h6UAzY95MJgsDIKGMhBzbiVALpea0k,581
+schedulers/sqrt.py,sha256=1F84ZgQPuoNMhe6yxGTR2G0h7dPOZtm4UDQOakbSsEU,445
+trainer/__init__.py,sha256=T-vUVq4zHcMK6kpCsG4Gu8vn71q6lZD-lg-Ul7yKfEk,128
+trainer/autoencoder_trainer.py,sha256=hxihkRL9WCIQVGOP-pc1jjjIUaRXDLcNo3_erTKsuWM,7049
+trainer/diffusion_trainer.py,sha256=ajOWBgFFwXP_VQScUjcuPoaB4Gk02aF0Ls5LNlA8wqA,12691
+trainer/simple_trainer.py,sha256=lmRo8N0bMupIyS3ejPvPtxoskY_3GLC8iyJE6u4TIWc,21990
+trainer/video_diffusion_trainer.py,sha256=gMkKpnKNTo8QhTx5ptEEkc7W5-7rzXIr9queU53hXyQ,2197
+flaxdiff-0.1.36.2.dist-info/METADATA,sha256=4c3Kl8wDAN6pzg_nLfQ3VSbN6e9xVv3flLCsh1ub-4U,22310
+flaxdiff-0.1.36.2.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+flaxdiff-0.1.36.2.dist-info/top_level.txt,sha256=FMc3ZGZsKPjeaG9fWYdPl5HYvG4JBNSxqXJj2WwJz6s,74
+flaxdiff-0.1.36.2.dist-info/RECORD,,

flaxdiff-0.1.36.2.dist-info/top_level.txt

@@ -0,0 +1,9 @@
+__init__
+data
+metrics
+models
+predictors
+samplers
+schedulers
+trainer
+utils