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

flaxdiff/data/online_loader.py

@@ -25,7 +25,7 @@ import cv2
 USER_AGENT = get_datasets_user_agent()
 
 data_queue = Queue(16*2000)
-error_queue = Queue(16*2000)
+error_queue = Queue()
 
 
 def fetch_single_image(image_url, timeout=None, retries=0):
@@ -98,7 +98,7 @@ def map_sample(
             "original_width": original_width,
         })
     except Exception as e:
-        error_queue.put({
+        error_queue.put_nowait({
             "url": url,
             "caption": caption,
             "error": str(e)
@@ -135,6 +135,9 @@ def parallel_image_loader(dataset: Dataset, num_workers: int = 8, image_shape=(2
             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:
@@ -230,7 +233,6 @@ class OnlineStreamingDataLoader():
         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)
-        self.collate_fn = collate_fn
         self.batch_size = batch_size
 
         # Launch a thread to load batches in the background
@@ -238,7 +240,10 @@ class OnlineStreamingDataLoader():
 
         def batch_loader():
             for batch in self.iterator:
-                self.batch_queue.put(batch)
+                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()
@@ -247,8 +252,8 @@ class OnlineStreamingDataLoader():
         return self
 
     def __next__(self):
-        return self.collate_fn(self.batch_queue.get())
+        return self.batch_queue.get()
         # return self.collate_fn(next(self.iterator))
 
     def __len__(self):
-        return len(self.dataset)
+        return len(self.dataset)

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

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

{flaxdiff-0.1.17.dist-info → flaxdiff-0.1.18.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=BM4Le-4BUo8MJpRzGIA2nMHKm4-WynQ2BOdiQz0JCDs,8791
+flaxdiff/data/online_loader.py,sha256=qim6SRRGU1lRO0zQbDNjRYC7Qm6g7jtUfELEXotora0,8987
 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.17.dist-info/METADATA,sha256=2Nr_T2yg3XHFt2jBuUXo8FxLYM8si-DBLdW_PBKxzc4,22083
-flaxdiff-0.1.17.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
-flaxdiff-0.1.17.dist-info/top_level.txt,sha256=-2-nXnfkJgSfkki1tjm5Faw6Dso7vhtdn2szwCdX5CQ,9
-flaxdiff-0.1.17.dist-info/RECORD,,
+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,,