flaxdiff 0.1.18__py3-none-any.whl → 0.1.19__py3-none-any.whl

flaxdiff/data/online_loader.py

@@ -44,7 +44,7 @@ def fetch_single_image(image_url, timeout=None, retries=0):
     return image
 
 
-def default_image_processor(image, image_shape, interpolation=cv2.INTER_LANCZOS4):
+def default_image_processor(image, image_shape, interpolation=cv2.INTER_CUBIC):
     image = A.longest_max_size(image, max(
         image_shape), interpolation=interpolation)
     image = A.pad(
@@ -62,7 +62,7 @@ def map_sample(
     image_shape=(256, 256),
     timeout=15,
     retries=3,
-    upscale_interpolation=cv2.INTER_LANCZOS4,
+    upscale_interpolation=cv2.INTER_CUBIC,
     downscale_interpolation=cv2.INTER_AREA,
     image_processor=default_image_processor,
 ):
@@ -105,10 +105,17 @@ def map_sample(
         })
 
 
-def map_batch(batch, num_threads=256, image_shape=(256, 256), timeout=15, retries=3, image_processor=default_image_processor):
+def map_batch(
+    batch, num_threads=256, image_shape=(256, 256), 
+    timeout=15, retries=3, image_processor=default_image_processor,
+    upscale_interpolation=cv2.INTER_CUBIC,
+    downscale_interpolation=cv2.INTER_AREA,
+):
     try:
         map_sample_fn = partial(map_sample, image_shape=image_shape,
-                                timeout=timeout, retries=retries, image_processor=image_processor)
+                                timeout=timeout, retries=retries, image_processor=image_processor,
+                                upscale_interpolation=upscale_interpolation,
+                                downscale_interpolation=downscale_interpolation)
         with ThreadPoolExecutor(max_workers=num_threads) as executor:
             executor.map(map_sample_fn, batch["url"], batch['caption'])
     except Exception as e:
@@ -118,10 +125,16 @@ def map_batch(batch, num_threads=256, image_shape=(256, 256), timeout=15, retrie
         })
 
 
-def parallel_image_loader(dataset: Dataset, num_workers: int = 8, image_shape=(256, 256), 
-                          num_threads=256, timeout=15, retries=3, image_processor=default_image_processor):
+def parallel_image_loader(
+    dataset: Dataset, num_workers: int = 8, image_shape=(256, 256), 
+    num_threads=256, timeout=15, retries=3, image_processor=default_image_processor,
+    upscale_interpolation=cv2.INTER_CUBIC,
+    downscale_interpolation=cv2.INTER_AREA,
+):
     map_batch_fn = partial(map_batch, num_threads=num_threads, image_shape=image_shape,
-                           timeout=timeout, retries=retries, image_processor=image_processor)
+                           timeout=timeout, retries=retries, image_processor=image_processor,
+                           upscale_interpolation=upscale_interpolation,
+                           downscale_interpolation=downscale_interpolation)
     shard_len = len(dataset) // num_workers
     print(f"Local Shard lengths: {shard_len}")
     with multiprocessing.Pool(num_workers) as pool:
@@ -141,14 +154,21 @@ def parallel_image_loader(dataset: Dataset, num_workers: int = 8, image_shape=(2
 
 
 class ImageBatchIterator:
-    def __init__(self, dataset: Dataset, batch_size: int = 64, image_shape=(256, 256), 
-                 num_workers: int = 8, num_threads=256, timeout=15, retries=3, image_processor=default_image_processor):
+    def __init__(
+        self, dataset: Dataset, batch_size: int = 64, image_shape=(256, 256), 
+        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,
+    ):
         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, num_workers=num_workers, 
-                         timeout=timeout, retries=retries, image_processor=image_processor)
+                         timeout=timeout, retries=retries, image_processor=image_processor,
+                         upscale_interpolation=upscale_interpolation,
+                         downscale_interpolation=downscale_interpolation)
         self.thread = threading.Thread(target=loader, args=(dataset,))
         self.thread.start()
 
@@ -210,6 +230,8 @@ class OnlineStreamingDataLoader():
         timeout=15,
         retries=3,
         image_processor=default_image_processor,
+        upscale_interpolation=cv2.INTER_CUBIC,
+        downscale_interpolation=cv2.INTER_AREA,
     ):
         if isinstance(dataset, str):
             dataset_path = dataset
@@ -232,7 +254,9 @@ class OnlineStreamingDataLoader():
         print(f"Dataset length: {len(dataset)}")
         self.iterator = ImageBatchIterator(self.dataset, image_shape=image_shape,
                                            num_workers=num_workers, batch_size=batch_size, num_threads=num_threads,
-                                             timeout=timeout, retries=retries, image_processor=image_processor)
+                                             timeout=timeout, retries=retries, image_processor=image_processor,
+                                             upscale_interpolation=upscale_interpolation,
+                                             downscale_interpolation=downscale_interpolation)
         self.batch_size = batch_size
 
         # Launch a thread to load batches in the background

{flaxdiff-0.1.18.dist-info → flaxdiff-0.1.19.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flaxdiff
-Version: 0.1.18
+Version: 0.1.19
 Summary: A versatile and easy to understand Diffusion library
 Author: Ashish Kumar Singh
 Author-email: ashishkmr472@gmail.com

{flaxdiff-0.1.18.dist-info → flaxdiff-0.1.19.dist-info}/RECORD

@@ -1,7 +1,7 @@
 flaxdiff/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 flaxdiff/utils.py,sha256=B0GcHlzlVYDNEIdh2v5qmP4u0neIT-FqexNohuyuCvg,2452
 flaxdiff/data/__init__.py,sha256=PM3PkHihyohT5SHVYKc8vQ4IeVfGPpCktkSVwvqMjQ4,52
-flaxdiff/data/online_loader.py,sha256=qim6SRRGU1lRO0zQbDNjRYC7Qm6g7jtUfELEXotora0,8987
+flaxdiff/data/online_loader.py,sha256=WK4apO8Bx-RTU_z5imB53Lzq12vqGnXA9DhLq8nb0us,9991
 flaxdiff/models/__init__.py,sha256=FAivVYXxM2JrCFIXf-C3374RB2Hth25dBrzOeNFhH1U,26
 flaxdiff/models/attention.py,sha256=ZbDGIb5Q6FRqJ6qRY660cqw4WvF9IwCnhEuYdTpLPdM,13023
 flaxdiff/models/common.py,sha256=fd-Fl0VCNEBjijHNwGBqYL5VvXe9u0347h25czNTmRw,10780
@@ -34,7 +34,7 @@ flaxdiff/trainer/__init__.py,sha256=T-vUVq4zHcMK6kpCsG4Gu8vn71q6lZD-lg-Ul7yKfEk,
 flaxdiff/trainer/autoencoder_trainer.py,sha256=al7AsZ7yeDMEiDD-gbcXf0ADq_xfk1VMxvg24GfA-XQ,7008
 flaxdiff/trainer/diffusion_trainer.py,sha256=wKkg63DWZjx2MoM3VQNCDIr40rWN8fUGxH9jWWxfZao,9373
 flaxdiff/trainer/simple_trainer.py,sha256=Z77zRS5viJpd2Mpl6sonJk5WcnEWi2Cd4gl4u5tIX2M,18206
-flaxdiff-0.1.18.dist-info/METADATA,sha256=aUSr3lBb9P2mnrpmbcgQa41DT8YYM-DtVMU8NI3CZEE,22083
-flaxdiff-0.1.18.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
-flaxdiff-0.1.18.dist-info/top_level.txt,sha256=-2-nXnfkJgSfkki1tjm5Faw6Dso7vhtdn2szwCdX5CQ,9
-flaxdiff-0.1.18.dist-info/RECORD,,
+flaxdiff-0.1.19.dist-info/METADATA,sha256=NH-f1SK5obamoVRk8ZPQxvtQcz_R3mui3ToZe0Qx8Vg,22083
+flaxdiff-0.1.19.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
+flaxdiff-0.1.19.dist-info/top_level.txt,sha256=-2-nXnfkJgSfkki1tjm5Faw6Dso7vhtdn2szwCdX5CQ,9
+flaxdiff-0.1.19.dist-info/RECORD,,