From c90099f06e0b461660b326fb6d86b69d86e78289 Mon Sep 17 00:00:00 2001 From: Volpeon Date: Mon, 3 Oct 2022 14:47:01 +0200 Subject: Added negative prompt support for training scripts --- data/textual_inversion/csv.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) (limited to 'data/textual_inversion/csv.py') diff --git a/data/textual_inversion/csv.py b/data/textual_inversion/csv.py index 3ac57df..64f0c28 100644 --- a/data/textual_inversion/csv.py +++ b/data/textual_inversion/csv.py @@ -43,9 +43,10 @@ class CSVDataModule(pl.LightningDataModule): def prepare_data(self): 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] - self.data_full = [(img, cap) for img, cap, skip in zip(image_paths, captions, skips) if skip != "x"] + prompts = metadata['prompt'].values + nprompts = metadata['nprompt'].values if 'nprompt' in metadata else [""] * len(image_paths) + skips = metadata['skip'].values if 'skip' in metadata else [""] * len(image_paths) + self.data_full = [(i, p, n) for i, p, n, s in zip(image_paths, prompts, nprompts, skips) if s != "x"] def setup(self, stage=None): valid_set_size = int(len(self.data_full) * 0.2) @@ -109,7 +110,7 @@ class CSVDataset(Dataset): return math.ceil(self._length / self.batch_size) * self.batch_size def get_example(self, i): - image_path, text = self.data[i % self.num_instance_images] + image_path, prompt, nprompt = self.data[i % self.num_instance_images] if image_path in self.cache: return self.cache[image_path] @@ -120,12 +121,13 @@ class CSVDataset(Dataset): if not instance_image.mode == "RGB": instance_image = instance_image.convert("RGB") - text = text.format(self.placeholder_token) + prompt = prompt.format(self.placeholder_token) - example["prompts"] = text + example["prompts"] = prompt + example["nprompts"] = nprompt example["pixel_values"] = instance_image example["input_ids"] = self.tokenizer( - text, + prompt, padding="max_length", truncation=True, max_length=self.tokenizer.model_max_length, @@ -140,6 +142,7 @@ class CSVDataset(Dataset): unprocessed_example = self.get_example(i) example["prompts"] = unprocessed_example["prompts"] + example["nprompts"] = unprocessed_example["nprompts"] example["input_ids"] = unprocessed_example["input_ids"] example["pixel_values"] = self.image_transforms(unprocessed_example["pixel_values"]) -- cgit v1.2.3-70-g09d2