fixed val pipeline for oxford

Author: AshishKumar4

flaxdiff/data/dataloaders.py

@@ -333,7 +333,7 @@ def get_trainset():
     def get_valset():
         transformations = [
             augmenter(),
-            pygrain.Batch(32, drop_remainder=True),
+            pygrain.Batch(64, drop_remainder=True),
         ]
 
         loader = pygrain.DataLoader(
@@ -342,7 +342,7 @@ def get_valset():
             operations=transformations,
             worker_count=8,
             read_options=pygrain.ReadOptions(
-                32, 100
+                32, 128
             ),
             worker_buffer_size=32,
         )

training.py

@@ -577,19 +577,7 @@ def main(args):
     print(f"Training on {CONFIG['dataset']['name']} dataset with {batches} samples")
 
     # Hardcoding these cuz don't have much time for project submission
-    if dataset_name == 'oxford_flowers102':
-        # Construct a validation set by the prompts
-        val_prompts = ['water tulip', ' a water lily', ' a water lily', ' a photo of a rose', ' a photo of a rose', ' a water lily', ' a water lily', ' a photo of a marigold', ' a photo of a marigold', ' a photo of a marigold', ' a water lily', ' a photo of a sunflower', ' a photo of a lotus', ' columbine', ' columbine', ' an orchid', ' an orchid', ' an orchid', ' a water lily', ' a water lily', ' a water lily', ' columbine', ' columbine', ' a photo of a sunflower', ' a photo of a sunflower', ' a photo of a sunflower', ' a photo of a lotus', ' a photo of a lotus', ' a photo of a marigold', ' a photo of a marigold', ' a photo of a rose', ' a photo of a rose', ' a photo of a rose', ' orange dahlia', ' orange dahlia', ' a lenten rose', ' a lenten rose', ' a water lily', ' a water lily', ' a water lily', ' a water lily', ' an orchid', ' an orchid', ' an orchid', ' hard-leaved pocket orchid', ' bird of paradise', ' bird of paradise', ' a photo of a lovely rose', ' a photo of a lovely rose', ' a photo of a globe-flower', ' a photo of a globe-flower', ' a photo of a lovely rose', ' a photo of a lovely rose', ' a photo of a ruby-lipped cattleya', ' a photo of a ruby-lipped cattleya', ' a photo of a lovely rose', ' a water lily', ' a osteospermum', ' a osteospermum', ' a water lily', ' a water lily', ' a water lily', ' a red rose', ' a red rose']
-        val_prompts *= 100
-        def get_val_dataset(batch_size=64):
-            for i in range(0, len(val_prompts), batch_size):
-                prompts = val_prompts[i:i + batch_size]
-                tokens = text_encoder.tokenize(prompts)
-                yield {"text": tokens}
-                
-        data['val'] = get_val_dataset
-        data['val_len'] = len(val_prompts)
-    elif dataset_name == 'laiona_coco':
+    if dataset_name == 'laiona_coco':
         import pickle
         val_set = pickle.load(open("/home/mrwhite0racle/gcs_mount/datasets/laion12m+mscoco_filtered-new/validation_set_small.pkl", "rb"))
         def get_val_dataset():