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

flaxdiff/data/__init__.py

@@ -0,0 +1 @@
+from .online_loader import OnlineStreamingDataLoader

flaxdiff/data/online_loader.py

@@ -79,7 +79,9 @@ def map_sample(
         data_queue.put({
             "url": url,
             "caption": caption,
-            "image": image
+            "image": image,
+            "original_height": original_height,
+            "original_width": original_width,
         })
     except Exception as e:
         error_queue.put({
@@ -88,12 +90,12 @@ def map_sample(
             "error": str(e)
         })
         
-def map_batch(batch, num_threads=256, timeout=None, retries=0):
+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'])
+        executor.map(map_sample, batch["url"], batch['caption'], image_shape=image_shape, timeout=timeout, retries=retries)
     
-def parallel_image_loader(dataset: Dataset, num_workers: int = 8, num_threads=256):
-    map_batch_fn = partial(map_batch, num_threads=num_threads)
+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)
     shard_len = len(dataset) // num_workers
     print(f"Local Shard lengths: {shard_len}")
     with multiprocessing.Pool(num_workers) as pool:
@@ -106,12 +108,12 @@ def parallel_image_loader(dataset: Dataset, num_workers: int = 8, num_threads=25
             iteration += 1
             
 class ImageBatchIterator:
-    def __init__(self, dataset: Dataset, batch_size: int = 64, num_workers: int = 8, num_threads=256):
+    def __init__(self, dataset: Dataset, batch_size: int = 64, image_shape=(256, 256), num_workers: int = 8, num_threads=256):
         self.dataset = dataset
         self.num_workers = num_workers
         self.batch_size = batch_size
-        loader = partial(parallel_image_loader, num_threads=num_threads)
-        self.thread = threading.Thread(target=loader, args=(dataset, num_workers))
+        loader = partial(parallel_image_loader, num_threads=num_threads, image_shape=image_shape, num_workers=num_workers)
+        self.thread = threading.Thread(target=loader, args=(dataset,))
         self.thread.start()
         
     def __iter__(self):
@@ -129,7 +131,7 @@ class ImageBatchIterator:
         
     def __len__(self):
         return len(self.dataset) // self.batch_size
-    
+
 def default_collate(batch):
     urls = [sample["url"] for sample in batch]
     captions = [sample["caption"] for sample in batch]
@@ -153,6 +155,7 @@ class OnlineStreamingDataLoader():
         self, 
         dataset, 
         batch_size=64, 
+        image_shape=(256, 256),
         num_workers=16, 
         num_threads=512,
         default_split="all",
@@ -174,14 +177,14 @@ class OnlineStreamingDataLoader():
             if isinstance(dataset[0], str):
                 print("Loading multiple datasets from paths")
                 dataset = [load_dataset(dataset_path, split=default_split) for dataset_path in dataset]
-            else:
-                print("Concatenating multiple datasets")
-                dataset = concatenate_datasets(dataset)
-        dataset = dataset.map(pre_map_maker(pre_map_def))
+            print("Concatenating multiple datasets")
+            dataset = concatenate_datasets(dataset)
+        dataset = dataset.map(pre_map_maker(pre_map_def), batched=True, batch_size=10000000)
         self.dataset = dataset.shard(num_shards=global_process_count, index=global_process_index)
         print(f"Dataset length: {len(dataset)}")
-        self.iterator = ImageBatchIterator(self.dataset, num_workers=num_workers, batch_size=batch_size, num_threads=num_threads)
+        self.iterator = ImageBatchIterator(self.dataset, image_shape=image_shape, num_workers=num_workers, batch_size=batch_size, num_threads=num_threads)
         self.collate_fn = collate_fn
+        self.batch_size = batch_size
         
         # Launch a thread to load batches in the background
         self.batch_queue = queue.Queue(prefetch)
@@ -201,5 +204,5 @@ class OnlineStreamingDataLoader():
         # return self.collate_fn(next(self.iterator))
         
     def __len__(self):
-        return len(self.dataset) // self.batch_size
+        return len(self.dataset)
     

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

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

{flaxdiff-0.1.13.dist-info → flaxdiff-0.1.15.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=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-flaxdiff/data/online_loader.py,sha256=_q0YgXGif1zLCNOLYJB0w3QFqE_p1zOjE9-hAApZFA4,6938
+flaxdiff/data/__init__.py,sha256=PM3PkHihyohT5SHVYKc8vQ4IeVfGPpCktkSVwvqMjQ4,52
+flaxdiff/data/online_loader.py,sha256=Xd4xAerzOU65tvb0cuM8S3Hnm8xv1dJ62xbxAxZkeBw,7307
 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.13.dist-info/METADATA,sha256=sx4MOQ9jPKy1yl_sUjkJbeu6vUJJG9HX6x6W8ILmW5I,22083
-flaxdiff-0.1.13.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
-flaxdiff-0.1.13.dist-info/top_level.txt,sha256=-2-nXnfkJgSfkki1tjm5Faw6Dso7vhtdn2szwCdX5CQ,9
-flaxdiff-0.1.13.dist-info/RECORD,,
+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,,