Format DreamBooth bucket sampler updates

Author: azolotenkov

examples/dreambooth/train_dreambooth_lora_flux2.py

@@ -975,7 +975,11 @@ def collate_fn(examples, with_prior_preservation=False):
 
 class BucketBatchSampler(BatchSampler):
     def __init__(
-        self, dataset: DreamBoothDataset, batch_size: int, drop_last: bool = False, shuffle_batches_each_epoch: bool = True
+        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))

examples/dreambooth/train_dreambooth_lora_flux2_img2img.py

@@ -973,7 +973,11 @@ def collate_fn(examples):
 
 class BucketBatchSampler(BatchSampler):
     def __init__(
-        self, dataset: DreamBoothDataset, batch_size: int, drop_last: bool = False, shuffle_batches_each_epoch: bool = True
+        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))

examples/dreambooth/train_dreambooth_lora_flux2_klein.py

@@ -970,7 +970,11 @@ def collate_fn(examples, with_prior_preservation=False):
 
 class BucketBatchSampler(BatchSampler):
     def __init__(
-        self, dataset: DreamBoothDataset, batch_size: int, drop_last: bool = False, shuffle_batches_each_epoch: bool = True
+        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))

examples/dreambooth/train_dreambooth_lora_flux2_klein_img2img.py

@@ -969,7 +969,11 @@ def collate_fn(examples):
 
 class BucketBatchSampler(BatchSampler):
     def __init__(
-        self, dataset: DreamBoothDataset, batch_size: int, drop_last: bool = False, shuffle_batches_each_epoch: bool = True
+        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))

examples/dreambooth/train_dreambooth_lora_z_image.py

@@ -964,7 +964,11 @@ def collate_fn(examples, with_prior_preservation=False):
 
 class BucketBatchSampler(BatchSampler):
     def __init__(
-        self, dataset: DreamBoothDataset, batch_size: int, drop_last: bool = False, shuffle_batches_each_epoch: bool = True
+        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))