flaxdiff 0.1.15__py3-none-any.whl → 0.1.16__py3-none-any.whl

flaxdiff/data/online_loader.py

@@ -27,7 +27,6 @@ USER_AGENT = get_datasets_user_agent()
 data_queue = Queue(16*2000)
 error_queue = Queue(16*2000)
 
-
 def fetch_single_image(image_url, timeout=None, retries=0):
     for _ in range(retries + 1):
         try:
@@ -46,11 +45,13 @@ def fetch_single_image(image_url, timeout=None, retries=0):
 def map_sample(
     url, caption, 
     image_shape=(256, 256),
+    timeout=15,
+    retries=3,
     upscale_interpolation=cv2.INTER_LANCZOS4,
     downscale_interpolation=cv2.INTER_AREA,
 ):
     try:
-        image = fetch_single_image(url, timeout=15, retries=3)  # Assuming fetch_single_image is defined elsewhere
+        image = fetch_single_image(url, timeout=timeout, retries=retries)  # Assuming fetch_single_image is defined elsewhere
         if image is None:
             return
         
@@ -84,15 +85,24 @@ def map_sample(
             "original_width": original_width,
         })
     except Exception as e:
+        print(f"Error in map_sample: {str(e)}")
         error_queue.put({
             "url": url,
             "caption": caption,
             "error": str(e)
         })
-        
-def map_batch(batch, num_threads=256, image_shape=(256, 256), timeout=None, retries=0):
-    with ThreadPoolExecutor(max_workers=num_threads) as executor:
-        executor.map(map_sample, batch["url"], batch['caption'], image_shape=image_shape, timeout=timeout, retries=retries)
+
+def map_batch(batch, num_threads=256, image_shape=(256, 256), timeout=15, retries=3):
+    try:
+        map_sample_fn = partial(map_sample, image_shape=image_shape, timeout=timeout, retries=retries)
+        with ThreadPoolExecutor(max_workers=num_threads) as executor:
+            executor.map(map_sample_fn, batch["url"], batch['caption'])
+    except Exception as e:
+        print(f"Error in map_batch: {str(e)}")
+        error_queue.put({
+            "batch": batch,
+            "error": str(e)
+        })
     
 def parallel_image_loader(dataset: Dataset, num_workers: int = 8, image_shape=(256, 256), num_threads=256):
     map_batch_fn = partial(map_batch, num_threads=num_threads, image_shape=image_shape)
@@ -102,8 +112,10 @@ def parallel_image_loader(dataset: Dataset, num_workers: int = 8, image_shape=(2
         iteration = 0
         while True:
             # Repeat forever
-            dataset = dataset.shuffle(seed=iteration)
+            print(f"Shuffling dataset with seed {iteration}")
+            # dataset = dataset.shuffle(seed=iteration)
             shards = [dataset[i*shard_len:(i+1)*shard_len] for i in range(num_workers)]
+            print(f"mapping {len(shards)} shards")
             pool.map(map_batch_fn, shards)
             iteration += 1
             
@@ -205,4 +217,3 @@ class OnlineStreamingDataLoader():
         
     def __len__(self):
         return len(self.dataset)
-    

{flaxdiff-0.1.15.dist-info → flaxdiff-0.1.16.dist-info}/METADATA

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

{flaxdiff-0.1.15.dist-info → flaxdiff-0.1.16.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=Xd4xAerzOU65tvb0cuM8S3Hnm8xv1dJ62xbxAxZkeBw,7307
+flaxdiff/data/online_loader.py,sha256=nrtZU4srZHsg3iN0sG91y_6nY7QtYXRPLk5rGn_BTIU,7728
 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.15.dist-info/METADATA,sha256=45ifUa-j2rPqP3s734HRy2rIM1StRvSZ4lIJI0R3sTw,22083
-flaxdiff-0.1.15.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
-flaxdiff-0.1.15.dist-info/top_level.txt,sha256=-2-nXnfkJgSfkki1tjm5Faw6Dso7vhtdn2szwCdX5CQ,9
-flaxdiff-0.1.15.dist-info/RECORD,,
+flaxdiff-0.1.16.dist-info/METADATA,sha256=BM2RLOiCDqRSWO_owxvWmL_PS3aFtNokPbf-qAuyK4o,22083
+flaxdiff-0.1.16.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
+flaxdiff-0.1.16.dist-info/top_level.txt,sha256=-2-nXnfkJgSfkki1tjm5Faw6Dso7vhtdn2szwCdX5CQ,9
+flaxdiff-0.1.16.dist-info/RECORD,,