Better dataset prompt handling

3 files changed, 18 insertions, 11 deletions

diff --git a/data/csv.py b/data/csv.py
index edce2b1..265293b 100644
--- a/data/csv.py
+++ b/data/csv.py
@@ -15,10 +15,11 @@ def prepare_prompt(prompt: Union[str, Dict[str, str]]):
     return {"content": prompt} if isinstance(prompt, str) else prompt
 
 
-def keywords_to_prompt(prompt: list[str], dropout: float = 0) -> str:
+def keywords_to_prompt(prompt: list[str], dropout: float = 0, shuffle: bool = False) -> str:
     if dropout != 0:
         prompt = [keyword for keyword in prompt if np.random.random() > dropout]
-    np.random.shuffle(prompt)
+    if shuffle:
+        np.random.shuffle(prompt)
     return ", ".join(prompt)
 
 
@@ -38,8 +39,8 @@ class CSVDataItem(NamedTuple):
     instance_image_path: Path
     class_image_path: Path
     prompt: list[str]
-    cprompt: str
-    nprompt: str
+    cprompt: list[str]
+    nprompt: list[str]
 
 
 class CSVDataModule():
@@ -104,8 +105,14 @@ class CSVDataModule():
                     prompt.format(**prepare_prompt(item["prompt"] if "prompt" in item else "")),
                     expansions
                 ),
-                cprompt.format(**prepare_prompt(item["cprompt"] if "cprompt" in item else "")),
-                nprompt.format(**prepare_prompt(item["nprompt"] if "nprompt" in item else "")),
+                prompt_to_keywords(
+                    cprompt.format(**prepare_prompt(item["cprompt"] if "cprompt" in item else "")),
+                    expansions
+                ),
+                prompt_to_keywords(
+                    prompt.format(**prepare_prompt(item["nprompt"] if "nprompt" in item else "")),
+                    expansions
+                ),
             )
             for item in data
         ]
@@ -253,9 +260,9 @@ class CSVDataset(Dataset):
 
         example = {}
 
-        example["prompts"] = keywords_to_prompt(unprocessed_example["prompts"], self.dropout)
-        example["cprompts"] = unprocessed_example["cprompts"]
-        example["nprompts"] = unprocessed_example["nprompts"]
+        example["prompts"] = keywords_to_prompt(unprocessed_example["prompts"], self.dropout, True)
+        example["cprompts"] = keywords_to_prompt(unprocessed_example["cprompts"])
+        example["nprompts"] = keywords_to_prompt(unprocessed_example["nprompts"])
 
         example["instance_images"] = self.image_transforms(unprocessed_example["instance_images"])
         example["instance_prompt_ids"] = self.prompt_processor.get_input_ids(example["prompts"])
diff --git a/train_dreambooth.py b/train_dreambooth.py
index 2f913e7..1a79b2b 100644
--- a/train_dreambooth.py
+++ b/train_dreambooth.py
@@ -687,7 +687,7 @@ def main():
             ).to(accelerator.device)
             pipeline.set_progress_bar_config(dynamic_ncols=True)
 
-            with torch.autocast("cuda"), torch.inference_mode():
+            with torch.inference_mode():
                 for batch in batched_data:
                     image_name = [item.class_image_path for item in batch]
                     prompt = [item.cprompt for item in batch]
diff --git a/train_ti.py b/train_ti.py
index e272b5d..cc208f0 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -621,7 +621,7 @@ def main():
             ).to(accelerator.device)
             pipeline.set_progress_bar_config(dynamic_ncols=True)
 
-            with torch.autocast("cuda"), torch.inference_mode():
+            with torch.inference_mode():
                 for batch in batched_data:
                     image_name = [item.class_image_path for item in batch]
                     prompt = [item.cprompt for item in batch]