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

flaxdiff/data/online_loader.py

@@ -25,7 +25,6 @@ import cv2
 USER_AGENT = get_datasets_user_agent()
 
 data_queue = Queue(16*2000)
-error_queue = Queue()
 
 
 def fetch_single_image(image_url, timeout=None, retries=0):
@@ -60,6 +59,7 @@ def default_image_processor(image, image_shape, interpolation=cv2.INTER_CUBIC):
 def map_sample(
     url, caption,
     image_shape=(256, 256),
+    min_image_shape=(128, 128),
     timeout=15,
     retries=3,
     upscale_interpolation=cv2.INTER_CUBIC,
@@ -75,21 +75,21 @@ def map_sample(
         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(image_shape):
+        if min(original_height, original_width) < min(min_image_shape):
             return
         # check if wrong aspect ratio
-        if max(original_height, original_width) / min(original_height, original_width) > 2:
+        if max(original_height, original_width) / min(original_height, original_width) > 2.4:
             return
         # check if the variance is too low
         if np.std(image) < 1e-4:
             return
-        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        # 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 = image_processor(
             image, image_shape, interpolation=interpolation)
-
+        
         data_queue.put({
             "url": url,
             "caption": caption,
@@ -98,40 +98,47 @@ def map_sample(
             "original_width": original_width,
         })
     except Exception as e:
-        error_queue.put_nowait({
-            "url": url,
-            "caption": caption,
-            "error": str(e)
-        })
+        print(f"Error processing {url}", e)
+        # error_queue.put_nowait({
+        #     "url": url,
+        #     "caption": caption,
+        #     "error": str(e)
+        # })
+        pass
 
 
 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,
 ):
     try:
-        map_sample_fn = partial(map_sample, image_shape=image_shape,
+        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:
             executor.map(map_sample_fn, batch["url"], batch['caption'])
     except Exception as e:
-        error_queue.put({
-            "batch": batch,
-            "error": str(e)
-        })
+        print(f"Error processing batch", e)
+        # 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,
 ):
-    map_batch_fn = partial(map_batch, num_threads=num_threads, image_shape=image_shape,
+    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)
@@ -149,13 +156,14 @@ def parallel_image_loader(
             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()
+            # 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,
@@ -165,7 +173,9 @@ class ImageBatchIterator:
         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, 
+                         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)
@@ -215,6 +225,7 @@ class OnlineStreamingDataLoader():
         dataset,
         batch_size=64,
         image_shape=(256, 256),
+        min_image_shape=(128, 128),
         num_workers=16,
         num_threads=512,
         default_split="all",
@@ -253,8 +264,9 @@ class OnlineStreamingDataLoader():
             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,
+                                            timeout=timeout, retries=retries, image_processor=image_processor,
                                              upscale_interpolation=upscale_interpolation,
                                              downscale_interpolation=downscale_interpolation)
         self.batch_size = batch_size

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

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

{flaxdiff-0.1.19.dist-info → flaxdiff-0.1.21.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=WK4apO8Bx-RTU_z5imB53Lzq12vqGnXA9DhLq8nb0us,9991
+flaxdiff/data/online_loader.py,sha256=w6gi1tAzWr4gtPt7onpStzOxp7Kdo_2q8Ro4Yi7OT4w,10549
 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.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,,
+flaxdiff-0.1.21.dist-info/METADATA,sha256=k1s_EIWBL0y4oCxXxr3QIi7LQbR47_jyDFfjVbURSMY,22083
+flaxdiff-0.1.21.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
+flaxdiff-0.1.21.dist-info/top_level.txt,sha256=-2-nXnfkJgSfkki1tjm5Faw6Dso7vhtdn2szwCdX5CQ,9
+flaxdiff-0.1.21.dist-info/RECORD,,