2 files changed, 25 insertions, 21 deletions

diff --git a/data/dreambooth/csv.py b/data/dreambooth/csv.py
index 04df4c6..e70c068 100644
--- a/data/dreambooth/csv.py
+++ b/data/dreambooth/csv.py
@@ -11,7 +11,7 @@ from torchvision import transforms
 class CSVDataModule(pl.LightningDataModule):
     def __init__(self,
                  batch_size,
-                 data_root,
+                 data_file,
                  tokenizer,
                  instance_prompt,
                  class_data_root=None,
@@ -24,7 +24,12 @@ class CSVDataModule(pl.LightningDataModule):
                  collate_fn=None):
         super().__init__()
 
-        self.data_root = data_root
+        self.data_file = Path(data_file)
+
+        if not self.data_file.is_file():
+            raise ValueError("data_file must be a file")
+
+        self.data_root = self.data_file.parent
         self.tokenizer = tokenizer
         self.instance_prompt = instance_prompt
         self.class_data_root = class_data_root
@@ -38,7 +43,7 @@ class CSVDataModule(pl.LightningDataModule):
         self.batch_size = batch_size
 
     def prepare_data(self):
-        metadata = pd.read_csv(f'{self.data_root}/list.csv')
+        metadata = pd.read_csv(self.data_file)
         image_paths = [os.path.join(self.data_root, f_path) for f_path in metadata['image'].values]
         captions = [caption for caption in metadata['caption'].values]
         skips = [skip for skip in metadata['skip'].values]
@@ -50,14 +55,13 @@ class CSVDataModule(pl.LightningDataModule):
         self.data_train, self.data_val = random_split(self.data_full, [train_set_size, valid_set_size])
 
         train_dataset = CSVDataset(self.data_train, self.tokenizer, instance_prompt=self.instance_prompt,
-                                   class_data_root=self.class_data_root,
-                                   class_prompt=self.class_prompt, size=self.size, repeats=self.repeats,
-                                   interpolation=self.interpolation, identifier=self.identifier,
-                                   center_crop=self.center_crop)
+                                   class_data_root=self.class_data_root, class_prompt=self.class_prompt,
+                                   size=self.size, interpolation=self.interpolation, identifier=self.identifier,
+                                   center_crop=self.center_crop, repeats=self.repeats)
         val_dataset = CSVDataset(self.data_val, self.tokenizer, instance_prompt=self.instance_prompt,
-                                 class_data_root=self.class_data_root,
-                                 class_prompt=self.class_prompt, size=self.size, interpolation=self.interpolation,
-                                 identifier=self.identifier, center_crop=self.center_crop)
+                                 class_data_root=self.class_data_root, class_prompt=self.class_prompt,
+                                 size=self.size, interpolation=self.interpolation, identifier=self.identifier,
+                                 center_crop=self.center_crop)
         self.train_dataloader_ = DataLoader(train_dataset, batch_size=self.batch_size,
                                             shuffle=True, collate_fn=self.collate_fn)
         self.val_dataloader_ = DataLoader(val_dataset, batch_size=self.batch_size, collate_fn=self.collate_fn)
diff --git a/dreambooth.py b/dreambooth.py
index 89ed96a..45a0497 100644
--- a/dreambooth.py
+++ b/dreambooth.py
@@ -46,7 +46,7 @@ def parse_args():
         help="Pretrained tokenizer name or path if not the same as model_name",
     )
     parser.add_argument(
-        "--train_data_dir",
+        "--train_data_file",
         type=str,
         default=None,
         help="A folder containing the training data."
@@ -269,8 +269,8 @@ def parse_args():
     if env_local_rank != -1 and env_local_rank != args.local_rank:
         args.local_rank = env_local_rank
 
-    if args.train_data_dir is None:
-        raise ValueError("You must specify --train_data_dir")
+    if args.train_data_file is None:
+        raise ValueError("You must specify --train_data_file")
 
     if args.pretrained_model_name_or_path is None:
         raise ValueError("You must specify --pretrained_model_name_or_path")
@@ -587,7 +587,7 @@ def main():
         return batch
 
     datamodule = CSVDataModule(
-        data_root=args.train_data_dir,
+        data_file=args.train_data_file,
         batch_size=args.train_batch_size,
         tokenizer=tokenizer,
         instance_prompt=args.instance_prompt,
@@ -680,12 +680,12 @@ def main():
         0,
         args.resolution, args.resolution, 7.5, 0.0, args.sample_steps)
 
-    progress_bar = tqdm(range(args.max_train_steps), disable=not accelerator.is_local_main_process)
-    progress_bar.set_description("Global steps")
-
     local_progress_bar = tqdm(range(num_update_steps_per_epoch), disable=not accelerator.is_local_main_process)
     local_progress_bar.set_description("Steps")
 
+    progress_bar = tqdm(range(args.max_train_steps), disable=not accelerator.is_local_main_process)
+    progress_bar.set_description("Global steps")
+
     try:
         for epoch in range(args.num_train_epochs):
             local_progress_bar.reset()
@@ -733,14 +733,14 @@ def main():
 
                 # Checks if the accelerator has performed an optimization step behind the scenes
                 if accelerator.sync_gradients:
-                    progress_bar.update(1)
                     local_progress_bar.update(1)
+                    progress_bar.update(1)
 
                     global_step += 1
 
                     if global_step % args.checkpoint_frequency == 0 and global_step > 0 and accelerator.is_main_process:
-                        progress_bar.clear()
                         local_progress_bar.clear()
+                        progress_bar.clear()
 
                         checkpointer.save_samples(
                             "training",
@@ -782,8 +782,8 @@ def main():
                     val_loss += loss
 
                 if accelerator.sync_gradients:
-                    progress_bar.update(1)
                     local_progress_bar.update(1)
+                    progress_bar.update(1)
 
                 logs = {"mode": "validation", "loss": loss}
                 local_progress_bar.set_postfix(**logs)
@@ -792,8 +792,8 @@ def main():
 
             accelerator.log({"train/loss": train_loss, "val/loss": val_loss}, step=global_step)
 
-            progress_bar.clear()
             local_progress_bar.clear()
+            progress_bar.clear()
 
             if min_val_loss > val_loss:
                 accelerator.print(f"Validation loss reached new minimum: {min_val_loss:.2e} -> {val_loss:.2e}")