Scope stable bucket ordering to cached DreamBooth batches

Author: azolotenkov

examples/dreambooth/train_dreambooth_lora_flux2.py

@@ -974,7 +974,9 @@ def collate_fn(examples, with_prior_preservation=False):
 
 
 class BucketBatchSampler(BatchSampler):
-    def __init__(self, dataset: DreamBoothDataset, batch_size: int, drop_last: bool = False):
+    def __init__(
+        self, dataset: DreamBoothDataset, batch_size: int, drop_last: bool = False, shuffle_batches_each_epoch: bool = True
+    ):
         if not isinstance(batch_size, int) or batch_size <= 0:
             raise ValueError("batch_size should be a positive integer value, but got batch_size={}".format(batch_size))
         if not isinstance(drop_last, bool):
@@ -983,6 +985,7 @@ def __init__(self, dataset: DreamBoothDataset, batch_size: int, drop_last: bool
         self.dataset = dataset
         self.batch_size = batch_size
         self.drop_last = drop_last
+        self.shuffle_batches_each_epoch = shuffle_batches_each_epoch
 
         # Group indices by bucket
         self.bucket_indices = [[] for _ in range(len(self.dataset.buckets))]
@@ -1004,12 +1007,14 @@ def __init__(self, dataset: DreamBoothDataset, batch_size: int, drop_last: bool
                 self.batches.append(batch)
                 self.sampler_len += 1  # Count the number of batches
 
-        # Shuffle the precomputed batches once to mix buckets while keeping
-        # the order stable across epochs for step-indexed caches.
-        random.shuffle(self.batches)
+        if not self.shuffle_batches_each_epoch:
+            # Shuffle the precomputed batches once to mix buckets while keeping
+            # the order stable across epochs for step-indexed caches.
+            random.shuffle(self.batches)
 
     def __iter__(self):
-        # Keep the precomputed batch order stable so step-indexed caches stay aligned.
+        if self.shuffle_batches_each_epoch:
+            random.shuffle(self.batches)
         for batch in self.batches:
             yield batch
 
@@ -1465,7 +1470,13 @@ def load_model_hook(models, input_dir):
         center_crop=args.center_crop,
         buckets=buckets,
     )
-    batch_sampler = BucketBatchSampler(train_dataset, batch_size=args.train_batch_size, drop_last=True)
+    has_step_indexed_caches = args.cache_latents or train_dataset.custom_instance_prompts
+    batch_sampler = BucketBatchSampler(
+        train_dataset,
+        batch_size=args.train_batch_size,
+        drop_last=True,
+        shuffle_batches_each_epoch=not has_step_indexed_caches,
+    )
     train_dataloader = torch.utils.data.DataLoader(
         train_dataset,
         batch_sampler=batch_sampler,
@@ -1582,8 +1593,7 @@ def _encode_single(prompt: str):
     # if cache_latents is set to True, we encode images to latents and store them.
     # Similar to pre-encoding in the case of a single instance prompt, if custom prompts are provided
     # we encode them in advance as well.
-    precompute_latents = args.cache_latents or train_dataset.custom_instance_prompts
-    if precompute_latents:
+    if has_step_indexed_caches:
         prompt_embeds_cache = []
         text_ids_cache = []
         latents_cache = []

examples/dreambooth/train_dreambooth_lora_flux2_img2img.py

@@ -972,7 +972,9 @@ def collate_fn(examples):
 
 
 class BucketBatchSampler(BatchSampler):
-    def __init__(self, dataset: DreamBoothDataset, batch_size: int, drop_last: bool = False):
+    def __init__(
+        self, dataset: DreamBoothDataset, batch_size: int, drop_last: bool = False, shuffle_batches_each_epoch: bool = True
+    ):
         if not isinstance(batch_size, int) or batch_size <= 0:
             raise ValueError("batch_size should be a positive integer value, but got batch_size={}".format(batch_size))
         if not isinstance(drop_last, bool):
@@ -981,6 +983,7 @@ def __init__(self, dataset: DreamBoothDataset, batch_size: int, drop_last: bool
         self.dataset = dataset
         self.batch_size = batch_size
         self.drop_last = drop_last
+        self.shuffle_batches_each_epoch = shuffle_batches_each_epoch
 
         # Group indices by bucket
         self.bucket_indices = [[] for _ in range(len(self.dataset.buckets))]
@@ -1002,12 +1005,14 @@ def __init__(self, dataset: DreamBoothDataset, batch_size: int, drop_last: bool
                 self.batches.append(batch)
                 self.sampler_len += 1  # Count the number of batches
 
-        # Shuffle the precomputed batches once to mix buckets while keeping
-        # the order stable across epochs for step-indexed caches.
-        random.shuffle(self.batches)
+        if not self.shuffle_batches_each_epoch:
+            # Shuffle the precomputed batches once to mix buckets while keeping
+            # the order stable across epochs for step-indexed caches.
+            random.shuffle(self.batches)
 
     def __iter__(self):
-        # Keep the precomputed batch order stable so step-indexed caches stay aligned.
+        if self.shuffle_batches_each_epoch:
+            random.shuffle(self.batches)
         for batch in self.batches:
             yield batch
 
@@ -1412,7 +1417,13 @@ def load_model_hook(models, input_dir):
         center_crop=args.center_crop,
         buckets=buckets,
     )
-    batch_sampler = BucketBatchSampler(train_dataset, batch_size=args.train_batch_size, drop_last=True)
+    has_step_indexed_caches = args.cache_latents or train_dataset.custom_instance_prompts
+    batch_sampler = BucketBatchSampler(
+        train_dataset,
+        batch_size=args.train_batch_size,
+        drop_last=True,
+        shuffle_batches_each_epoch=not has_step_indexed_caches,
+    )
     train_dataloader = torch.utils.data.DataLoader(
         train_dataset,
         batch_sampler=batch_sampler,
@@ -1515,8 +1526,7 @@ def _encode_single(prompt: str):
     # if cache_latents is set to True, we encode images to latents and store them.
     # Similar to pre-encoding in the case of a single instance prompt, if custom prompts are provided
     # we encode them in advance as well.
-    precompute_latents = args.cache_latents or train_dataset.custom_instance_prompts
-    if precompute_latents:
+    if has_step_indexed_caches:
         prompt_embeds_cache = []
         text_ids_cache = []
         latents_cache = []

examples/dreambooth/train_dreambooth_lora_flux2_klein.py

@@ -969,7 +969,9 @@ def collate_fn(examples, with_prior_preservation=False):
 
 
 class BucketBatchSampler(BatchSampler):
-    def __init__(self, dataset: DreamBoothDataset, batch_size: int, drop_last: bool = False):
+    def __init__(
+        self, dataset: DreamBoothDataset, batch_size: int, drop_last: bool = False, shuffle_batches_each_epoch: bool = True
+    ):
         if not isinstance(batch_size, int) or batch_size <= 0:
             raise ValueError("batch_size should be a positive integer value, but got batch_size={}".format(batch_size))
         if not isinstance(drop_last, bool):
@@ -978,6 +980,7 @@ def __init__(self, dataset: DreamBoothDataset, batch_size: int, drop_last: bool
         self.dataset = dataset
         self.batch_size = batch_size
         self.drop_last = drop_last
+        self.shuffle_batches_each_epoch = shuffle_batches_each_epoch
 
         # Group indices by bucket
         self.bucket_indices = [[] for _ in range(len(self.dataset.buckets))]
@@ -999,12 +1002,14 @@ def __init__(self, dataset: DreamBoothDataset, batch_size: int, drop_last: bool
                 self.batches.append(batch)
                 self.sampler_len += 1  # Count the number of batches
 
-        # Shuffle the precomputed batches once to mix buckets while keeping
-        # the order stable across epochs for step-indexed caches.
-        random.shuffle(self.batches)
+        if not self.shuffle_batches_each_epoch:
+            # Shuffle the precomputed batches once to mix buckets while keeping
+            # the order stable across epochs for step-indexed caches.
+            random.shuffle(self.batches)
 
     def __iter__(self):
-        # Keep the precomputed batch order stable so step-indexed caches stay aligned.
+        if self.shuffle_batches_each_epoch:
+            random.shuffle(self.batches)
         for batch in self.batches:
             yield batch
 
@@ -1458,7 +1463,13 @@ def load_model_hook(models, input_dir):
         center_crop=args.center_crop,
         buckets=buckets,
     )
-    batch_sampler = BucketBatchSampler(train_dataset, batch_size=args.train_batch_size, drop_last=True)
+    has_step_indexed_caches = args.cache_latents or train_dataset.custom_instance_prompts
+    batch_sampler = BucketBatchSampler(
+        train_dataset,
+        batch_size=args.train_batch_size,
+        drop_last=True,
+        shuffle_batches_each_epoch=not has_step_indexed_caches,
+    )
     train_dataloader = torch.utils.data.DataLoader(
         train_dataset,
         batch_sampler=batch_sampler,
@@ -1525,8 +1536,7 @@ def compute_text_embeddings(prompt, text_encoding_pipeline):
     # if cache_latents is set to True, we encode images to latents and store them.
     # Similar to pre-encoding in the case of a single instance prompt, if custom prompts are provided
     # we encode them in advance as well.
-    precompute_latents = args.cache_latents or train_dataset.custom_instance_prompts
-    if precompute_latents:
+    if has_step_indexed_caches:
         prompt_embeds_cache = []
         text_ids_cache = []
         latents_cache = []

examples/dreambooth/train_dreambooth_lora_flux2_klein_img2img.py

@@ -968,7 +968,9 @@ def collate_fn(examples):
 
 
 class BucketBatchSampler(BatchSampler):
-    def __init__(self, dataset: DreamBoothDataset, batch_size: int, drop_last: bool = False):
+    def __init__(
+        self, dataset: DreamBoothDataset, batch_size: int, drop_last: bool = False, shuffle_batches_each_epoch: bool = True
+    ):
         if not isinstance(batch_size, int) or batch_size <= 0:
             raise ValueError("batch_size should be a positive integer value, but got batch_size={}".format(batch_size))
         if not isinstance(drop_last, bool):
@@ -977,6 +979,7 @@ def __init__(self, dataset: DreamBoothDataset, batch_size: int, drop_last: bool
         self.dataset = dataset
         self.batch_size = batch_size
         self.drop_last = drop_last
+        self.shuffle_batches_each_epoch = shuffle_batches_each_epoch
 
         # Group indices by bucket
         self.bucket_indices = [[] for _ in range(len(self.dataset.buckets))]
@@ -998,12 +1001,14 @@ def __init__(self, dataset: DreamBoothDataset, batch_size: int, drop_last: bool
                 self.batches.append(batch)
                 self.sampler_len += 1  # Count the number of batches
 
-        # Shuffle the precomputed batches once to mix buckets while keeping
-        # the order stable across epochs for step-indexed caches.
-        random.shuffle(self.batches)
+        if not self.shuffle_batches_each_epoch:
+            # Shuffle the precomputed batches once to mix buckets while keeping
+            # the order stable across epochs for step-indexed caches.
+            random.shuffle(self.batches)
 
     def __iter__(self):
-        # Keep the precomputed batch order stable so step-indexed caches stay aligned.
+        if self.shuffle_batches_each_epoch:
+            random.shuffle(self.batches)
         for batch in self.batches:
             yield batch
 
@@ -1406,7 +1411,13 @@ def load_model_hook(models, input_dir):
         center_crop=args.center_crop,
         buckets=buckets,
     )
-    batch_sampler = BucketBatchSampler(train_dataset, batch_size=args.train_batch_size, drop_last=True)
+    has_step_indexed_caches = args.cache_latents or train_dataset.custom_instance_prompts
+    batch_sampler = BucketBatchSampler(
+        train_dataset,
+        batch_size=args.train_batch_size,
+        drop_last=True,
+        shuffle_batches_each_epoch=not has_step_indexed_caches,
+    )
     train_dataloader = torch.utils.data.DataLoader(
         train_dataset,
         batch_sampler=batch_sampler,
@@ -1466,8 +1477,7 @@ def compute_text_embeddings(prompt, text_encoding_pipeline):
     # if cache_latents is set to True, we encode images to latents and store them.
     # Similar to pre-encoding in the case of a single instance prompt, if custom prompts are provided
     # we encode them in advance as well.
-    precompute_latents = args.cache_latents or train_dataset.custom_instance_prompts
-    if precompute_latents:
+    if has_step_indexed_caches:
         prompt_embeds_cache = []
         text_ids_cache = []
         latents_cache = []

examples/dreambooth/train_dreambooth_lora_z_image.py

@@ -963,7 +963,9 @@ def collate_fn(examples, with_prior_preservation=False):
 
 
 class BucketBatchSampler(BatchSampler):
-    def __init__(self, dataset: DreamBoothDataset, batch_size: int, drop_last: bool = False):
+    def __init__(
+        self, dataset: DreamBoothDataset, batch_size: int, drop_last: bool = False, shuffle_batches_each_epoch: bool = True
+    ):
         if not isinstance(batch_size, int) or batch_size <= 0:
             raise ValueError("batch_size should be a positive integer value, but got batch_size={}".format(batch_size))
         if not isinstance(drop_last, bool):
@@ -972,6 +974,7 @@ def __init__(self, dataset: DreamBoothDataset, batch_size: int, drop_last: bool
         self.dataset = dataset
         self.batch_size = batch_size
         self.drop_last = drop_last
+        self.shuffle_batches_each_epoch = shuffle_batches_each_epoch
 
         # Group indices by bucket
         self.bucket_indices = [[] for _ in range(len(self.dataset.buckets))]
@@ -993,12 +996,14 @@ def __init__(self, dataset: DreamBoothDataset, batch_size: int, drop_last: bool
                 self.batches.append(batch)
                 self.sampler_len += 1  # Count the number of batches
 
-        # Shuffle the precomputed batches once to mix buckets while keeping
-        # the order stable across epochs for step-indexed caches.
-        random.shuffle(self.batches)
+        if not self.shuffle_batches_each_epoch:
+            # Shuffle the precomputed batches once to mix buckets while keeping
+            # the order stable across epochs for step-indexed caches.
+            random.shuffle(self.batches)
 
     def __iter__(self):
-        # Keep the precomputed batch order stable so step-indexed caches stay aligned.
+        if self.shuffle_batches_each_epoch:
+            random.shuffle(self.batches)
         for batch in self.batches:
             yield batch
 
@@ -1452,7 +1457,13 @@ def load_model_hook(models, input_dir):
         center_crop=args.center_crop,
         buckets=buckets,
     )
-    batch_sampler = BucketBatchSampler(train_dataset, batch_size=args.train_batch_size, drop_last=True)
+    has_step_indexed_caches = args.cache_latents or train_dataset.custom_instance_prompts
+    batch_sampler = BucketBatchSampler(
+        train_dataset,
+        batch_size=args.train_batch_size,
+        drop_last=True,
+        shuffle_batches_each_epoch=not has_step_indexed_caches,
+    )
     train_dataloader = torch.utils.data.DataLoader(
         train_dataset,
         batch_sampler=batch_sampler,
@@ -1512,8 +1523,7 @@ def compute_text_embeddings(prompt, text_encoding_pipeline):
     # if cache_latents is set to True, we encode images to latents and store them.
     # Similar to pre-encoding in the case of a single instance prompt, if custom prompts are provided
     # we encode them in advance as well.
-    precompute_latents = args.cache_latents or train_dataset.custom_instance_prompts
-    if precompute_latents:
+    if has_step_indexed_caches:
         prompt_embeds_cache = []
         latents_cache = []
         for batch in tqdm(train_dataloader, desc="Caching latents"):