Improved inverted tokens

1 files changed, 44 insertions, 23 deletions

diff --git a/data/csv.py b/data/csv.py
index d52d251..9770bec 100644
--- a/data/csv.py
+++ b/data/csv.py
@@ -178,6 +178,7 @@ class VlpnDataModule():
         shuffle: bool = False,
         interpolation: str = "bicubic",
         template_key: str = "template",
+        placeholder_tokens: list[str] = [],
         valid_set_size: Optional[int] = None,
         train_set_pad: Optional[int] = None,
         valid_set_pad: Optional[int] = None,
@@ -195,6 +196,7 @@ class VlpnDataModule():
         self.data_root = self.data_file.parent
         self.class_root = self.data_root / class_subdir
         self.class_root.mkdir(parents=True, exist_ok=True)
+        self.placeholder_tokens = placeholder_tokens
         self.num_class_images = num_class_images
         self.with_guidance = with_guidance
 
@@ -217,31 +219,50 @@ class VlpnDataModule():
         self.dtype = dtype
 
     def prepare_items(self, template, expansions, data) -> list[VlpnDataItem]:
-        image = template["image"] if "image" in template else "{}"
-        prompt = template["prompt"] if "prompt" in template else "{content}"
-        cprompt = template["cprompt"] if "cprompt" in template else "{content}"
-        nprompt = template["nprompt"] if "nprompt" in template else "{content}"
+        tpl_image = template["image"] if "image" in template else "{}"
+        tpl_prompt = template["prompt"] if "prompt" in template else "{content}"
+        tpl_cprompt = template["cprompt"] if "cprompt" in template else "{content}"
+        tpl_nprompt = template["nprompt"] if "nprompt" in template else "{content}"
+
+        items = []
+
+        for item in data:
+            image = tpl_image.format(item["image"])
+            prompt = item["prompt"] if "prompt" in item else ""
+            nprompt = item["nprompt"] if "nprompt" in item else ""
+            collection = item["collection"].split(", ") if "collection" in item else []
+
+            prompt_keywords = prompt_to_keywords(
+                tpl_prompt.format(**prepare_prompt(prompt)),
+                expansions
+            )
 
-        return [
-            VlpnDataItem(
-                self.data_root / image.format(item["image"]),
-                None,
-                prompt_to_keywords(
-                    prompt.format(**prepare_prompt(item["prompt"] if "prompt" in item else "")),
-                    expansions
-                ),
-                keywords_to_prompt(prompt_to_keywords(
-                    cprompt.format(**prepare_prompt(item["prompt"] if "prompt" in item else "")),
-                    expansions
-                )),
-                prompt_to_keywords(
-                    nprompt.format(**prepare_prompt(item["nprompt"] if "nprompt" in item else "")),
-                    expansions
-                ),
-                item["collection"].split(", ") if "collection" in item else []
+            cprompt = keywords_to_prompt(prompt_to_keywords(
+                tpl_cprompt.format(**prepare_prompt(prompt)),
+                expansions
+            ))
+
+            inverted_tokens = keywords_to_prompt([
+                f"inv_{token}"
+                for token in self.placeholder_tokens
+                if token in prompt_keywords
+            ])
+
+            nprompt_keywords = prompt_to_keywords(
+                tpl_nprompt.format(_inv=inverted_tokens, **prepare_prompt(nprompt)),
+                expansions
             )
-            for item in data
-        ]
+
+            items.append(VlpnDataItem(
+                self.data_root / image,
+                None,
+                prompt_keywords,
+                cprompt,
+                nprompt_keywords,
+                collection
+            ))
+
+        return items
 
     def filter_items(self, items: list[VlpnDataItem]) -> list[VlpnDataItem]:
         if self.filter is None: