flaxdiff 0.1.34__py3-none-any.whl → 0.1.35.1__py3-none-any.whl

flaxdiff/data/online_loader.py

@@ -86,7 +86,8 @@ def default_feature_extractor(sample):
 
 
 def map_sample(
-    sample,
+    url,
+    caption,
     image_shape=(256, 256),
     min_image_shape=(128, 128),
     timeout=15,
@@ -94,11 +95,8 @@ def map_sample(
     upscale_interpolation=cv2.INTER_CUBIC,
     downscale_interpolation=cv2.INTER_AREA,
     image_processor=default_image_processor,
-    feature_extractor=default_feature_extractor,
 ):
     try:
-        features = feature_extractor(sample)
-        url, caption = features["url"], features["caption"]
         # Assuming fetch_single_image is defined elsewhere
         image = fetch_single_image(url, timeout=timeout, retries=retries)
         if image is None:
@@ -147,8 +145,10 @@ def map_batch(
             downscale_interpolation=downscale_interpolation,
             feature_extractor=feature_extractor
         )
+        features = feature_extractor(batch)
+        url, caption = features["url"], features["caption"]
         with ThreadPoolExecutor(max_workers=num_threads) as executor:
-            executor.map(map_sample_fn, batch)
+            executor.map(map_sample_fn, url, caption)
     except Exception as e:
         print(f"Error maping batch", e)
         traceback.print_exc()
@@ -216,6 +216,7 @@ def parallel_image_loader(
             # Repeat forever
             shards = [dataset[i*shard_len:(i+1)*shard_len]
                       for i in range(num_workers)]
+            # shards = [dataset.shard(num_shards=num_workers, index=i) for i in range(num_workers)]
             print(f"mapping {len(shards)} shards")
             pool.map(map_batch_fn, shards)
             iteration += 1

{flaxdiff-0.1.34.dist-info → flaxdiff-0.1.35.1.dist-info}/METADATA

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

{flaxdiff-0.1.34.dist-info → flaxdiff-0.1.35.1.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=VpwTivqTYXM4Eu3jLWoDkYRcl7KsgJ5A6t17zDSDsDc,13226
+flaxdiff/data/online_loader.py,sha256=reMG_EYbzcNi5rTVbGnAZE7gq1iY6YxYY7K6D5JIsbw,13293
 flaxdiff/models/__init__.py,sha256=FAivVYXxM2JrCFIXf-C3374RB2Hth25dBrzOeNFhH1U,26
 flaxdiff/models/attention.py,sha256=ZbDGIb5Q6FRqJ6qRY660cqw4WvF9IwCnhEuYdTpLPdM,13023
 flaxdiff/models/common.py,sha256=hWsSs2BP2J-JN1s4qLRr-h-KYkcVyl2hOp1Wsm_L-h8,10994
@@ -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.34.dist-info/METADATA,sha256=5x3QvLbZ3AgBjghAxnWQXQD5-cJ9jZ0ifPhGxdTeuEc,22083
-flaxdiff-0.1.34.dist-info/WHEEL,sha256=cVxcB9AmuTcXqmwrtPhNK88dr7IR_b6qagTj0UvIEbY,91
-flaxdiff-0.1.34.dist-info/top_level.txt,sha256=-2-nXnfkJgSfkki1tjm5Faw6Dso7vhtdn2szwCdX5CQ,9
-flaxdiff-0.1.34.dist-info/RECORD,,
+flaxdiff-0.1.35.1.dist-info/METADATA,sha256=ktAM1HUTLykCSmVZzH3VJHj4PzeCoOEOlpslyavlOWs,22085
+flaxdiff-0.1.35.1.dist-info/WHEEL,sha256=cVxcB9AmuTcXqmwrtPhNK88dr7IR_b6qagTj0UvIEbY,91
+flaxdiff-0.1.35.1.dist-info/top_level.txt,sha256=-2-nXnfkJgSfkki1tjm5Faw6Dso7vhtdn2szwCdX5CQ,9
+flaxdiff-0.1.35.1.dist-info/RECORD,,