3 files changed, 47 insertions, 10 deletions

diff --git a/data/csv.py b/data/csv.py
index 9125212..9c3c3f8 100644
--- a/data/csv.py
+++ b/data/csv.py
@@ -54,8 +54,10 @@ class CSVDataModule(pl.LightningDataModule):
             dropout: float = 0,
             interpolation: str = "bicubic",
             center_crop: bool = False,
+            template_key: str = "template",
             valid_set_size: Optional[int] = None,
             generator: Optional[torch.Generator] = None,
+            keyword_filter: list[str] = [],
             collate_fn=None,
             num_workers: int = 0
     ):
@@ -78,38 +80,60 @@ class CSVDataModule(pl.LightningDataModule):
         self.repeats = repeats
         self.dropout = dropout
         self.center_crop = center_crop
+        self.template_key = template_key
         self.interpolation = interpolation
         self.valid_set_size = valid_set_size
         self.generator = generator
+        self.keyword_filter = keyword_filter
         self.collate_fn = collate_fn
         self.num_workers = num_workers
         self.batch_size = batch_size
 
-    def prepare_subdata(self, template, data, num_class_images=1):
+    def prepare_items(self, template, data) -> list[CSVDataItem]:
         image = template["image"] if "image" in template else "{}"
         prompt = template["prompt"] if "prompt" in template else "{content}"
         nprompt = template["nprompt"] if "nprompt" in template else "{content}"
 
-        image_multiplier = max(math.ceil(num_class_images / len(data)), 1)
-
         return [
             CSVDataItem(
                 self.data_root.joinpath(image.format(item["image"])),
-                self.class_root.joinpath(f"{Path(item['image']).stem}_{i}{Path(item['image']).suffix}"),
+                None,
                 prompt.format(**prepare_prompt(item["prompt"] if "prompt" in item else "")),
                 nprompt.format(**prepare_prompt(item["nprompt"] if "nprompt" in item else ""))
             )
             for item in data
+        ]
+
+    def filter_items(self, items: list[CSVDataItem]) -> list[CSVDataItem]:
+        if len(self.keyword_filter) == 0:
+            return items
+
+        return [item for item in items if any(keyword in item.prompt for keyword in self.keyword_filter)]
+
+    def pad_items(self, items: list[CSVDataItem], num_class_images: int = 1) -> list[CSVDataItem]:
+        image_multiplier = max(math.ceil(num_class_images / len(items)), 1)
+
+        return [
+            CSVDataItem(
+                item.instance_image_path,
+                self.class_root.joinpath(f"{item.instance_image_path.stem}_{i}{item.instance_image_path.suffix}"),
+                item.prompt,
+                item.nprompt
+            )
+            for item in items
             for i in range(image_multiplier)
         ]
 
     def prepare_data(self):
         with open(self.data_file, 'rt') as f:
             metadata = json.load(f)
-        template = metadata["template"] if "template" in metadata else {}
+        template = metadata[self.template_key] if self.template_key in metadata else {}
         items = metadata["items"] if "items" in metadata else []
 
         items = [item for item in items if not "skip" in item or item["skip"] != True]
+        items = self.prepare_items(template, items)
+        items = self.filter_items(items)
+
         num_images = len(items)
 
         valid_set_size = int(num_images * 0.1)
@@ -120,8 +144,8 @@ class CSVDataModule(pl.LightningDataModule):
 
         data_train, data_val = random_split(items, [train_set_size, valid_set_size], self.generator)
 
-        self.data_train = self.prepare_subdata(template, data_train, self.num_class_images)
-        self.data_val = self.prepare_subdata(template, data_val)
+        self.data_train = self.pad_items(data_train, self.num_class_images)
+        self.data_val = self.pad_items(data_val)
 
     def setup(self, stage=None):
         train_dataset = CSVDataset(self.data_train, self.prompt_processor, batch_size=self.batch_size,
diff --git a/dreambooth.py b/dreambooth.py
index 31416e9..5521b21 100644
--- a/dreambooth.py
+++ b/dreambooth.py
@@ -57,6 +57,11 @@ def parse_args():
         help="A folder containing the training data."
     )
     parser.add_argument(
+        "--train_data_template",
+        type=str,
+        default="template",
+    )
+    parser.add_argument(
         "--instance_identifier",
         type=str,
         default=None,
@@ -768,6 +773,7 @@ def main():
         repeats=args.repeats,
         dropout=args.tag_dropout,
         center_crop=args.center_crop,
+        template_key=args.train_data_template,
         valid_set_size=args.valid_set_size,
         num_workers=args.dataloader_num_workers,
         collate_fn=collate_fn
diff --git a/textual_inversion.py b/textual_inversion.py
index 19b8993..fd4a313 100644
--- a/textual_inversion.py
+++ b/textual_inversion.py
@@ -58,6 +58,11 @@ def parse_args():
         help="A CSV file containing the training data."
     )
     parser.add_argument(
+        "--train_data_template",
+        type=str,
+        default="template",
+    )
+    parser.add_argument(
         "--instance_identifier",
         type=str,
         default=None,
@@ -121,7 +126,7 @@ def parse_args():
     parser.add_argument(
         "--tag_dropout",
         type=float,
-        default=0.1,
+        default=0,
         help="Tag dropout probability.",
     )
     parser.add_argument(
@@ -170,7 +175,7 @@ def parse_args():
     parser.add_argument(
         "--lr_scheduler",
         type=str,
-        default="constant_with_warmup",
+        default="one_cycle",
         help=(
             'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
             ' "constant", "constant_with_warmup", "one_cycle"]'
@@ -670,8 +675,10 @@ def main():
         repeats=args.repeats,
         dropout=args.tag_dropout,
         center_crop=args.center_crop,
+        template_key=args.train_data_template,
         valid_set_size=args.valid_set_size,
         num_workers=args.dataloader_num_workers,
+        keyword_filter=args.placeholder_token,
         collate_fn=collate_fn
     )
 
@@ -740,7 +747,7 @@ def main():
             num_warmup_steps=warmup_steps,
             num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
             num_cycles=args.lr_cycles or math.ceil(math.sqrt(
-                ((args.max_train_steps - args.lr_warmup_steps) / num_update_steps_per_epoch))),
+                ((args.max_train_steps - warmup_steps) / num_update_steps_per_epoch))),
         )
     else:
         lr_scheduler = get_scheduler(