1 files changed, 37 insertions, 46 deletions

diff --git a/data/csv.py b/data/csv.py
index 316c099..4c91ded 100644
--- a/data/csv.py
+++ b/data/csv.py
@@ -1,11 +1,14 @@
 import math
 import pandas as pd
+import torch
 from pathlib import Path
 import pytorch_lightning as pl
 from PIL import Image
 from torch.utils.data import Dataset, DataLoader, random_split
 from torchvision import transforms
-from typing import NamedTuple, List
+from typing import NamedTuple, List, Optional
+
+from models.clip.prompt import PromptProcessor
 
 
 class CSVDataItem(NamedTuple):
@@ -18,19 +21,19 @@ class CSVDataItem(NamedTuple):
 class CSVDataModule(pl.LightningDataModule):
     def __init__(
             self,
-            batch_size,
-            data_file,
-            tokenizer,
-            instance_identifier,
-            class_identifier=None,
-            class_subdir="cls",
-            num_class_images=100,
-            size=512,
-            repeats=100,
-            interpolation="bicubic",
-            center_crop=False,
-            valid_set_size=None,
-            generator=None,
+            batch_size: int,
+            data_file: str,
+            prompt_processor: PromptProcessor,
+            instance_identifier: str,
+            class_identifier: Optional[str] = None,
+            class_subdir: str = "cls",
+            num_class_images: int = 100,
+            size: int = 512,
+            repeats: int = 1,
+            interpolation: str = "bicubic",
+            center_crop: bool = False,
+            valid_set_size: Optional[int] = None,
+            generator: Optional[torch.Generator] = None,
             collate_fn=None
     ):
         super().__init__()
@@ -45,7 +48,7 @@ class CSVDataModule(pl.LightningDataModule):
         self.class_root.mkdir(parents=True, exist_ok=True)
         self.num_class_images = num_class_images
 
-        self.tokenizer = tokenizer
+        self.prompt_processor = prompt_processor
         self.instance_identifier = instance_identifier
         self.class_identifier = class_identifier
         self.size = size
@@ -65,7 +68,7 @@ class CSVDataModule(pl.LightningDataModule):
                 self.data_root.joinpath(item.image),
                 self.class_root.joinpath(f"{Path(item.image).stem}_{i}{Path(item.image).suffix}"),
                 item.prompt,
-                item.nprompt if "nprompt" in item else ""
+                item.nprompt
             )
             for item in data
             for i in range(image_multiplier)
@@ -88,12 +91,12 @@ class CSVDataModule(pl.LightningDataModule):
         self.data_val = self.prepare_subdata(data_val)
 
     def setup(self, stage=None):
-        train_dataset = CSVDataset(self.data_train, self.tokenizer, batch_size=self.batch_size,
+        train_dataset = CSVDataset(self.data_train, self.prompt_processor, batch_size=self.batch_size,
                                    instance_identifier=self.instance_identifier, class_identifier=self.class_identifier,
                                    num_class_images=self.num_class_images,
                                    size=self.size, interpolation=self.interpolation,
                                    center_crop=self.center_crop, repeats=self.repeats)
-        val_dataset = CSVDataset(self.data_val, self.tokenizer, batch_size=self.batch_size,
+        val_dataset = CSVDataset(self.data_val, self.prompt_processor, batch_size=self.batch_size,
                                  instance_identifier=self.instance_identifier,
                                  size=self.size, interpolation=self.interpolation,
                                  center_crop=self.center_crop, repeats=self.repeats)
@@ -113,19 +116,19 @@ class CSVDataset(Dataset):
     def __init__(
         self,
         data: List[CSVDataItem],
-        tokenizer,
-        instance_identifier,
-        batch_size=1,
-        class_identifier=None,
-        num_class_images=0,
-        size=512,
-        repeats=1,
-        interpolation="bicubic",
-        center_crop=False,
+        prompt_processor: PromptProcessor,
+        instance_identifier: str,
+        batch_size: int = 1,
+        class_identifier: Optional[str] = None,
+        num_class_images: int = 0,
+        size: int = 512,
+        repeats: int = 1,
+        interpolation: str = "bicubic",
+        center_crop: bool = False,
     ):
 
         self.data = data
-        self.tokenizer = tokenizer
+        self.prompt_processor = prompt_processor
         self.batch_size = batch_size
         self.instance_identifier = instance_identifier
         self.class_identifier = class_identifier
@@ -163,12 +166,6 @@ class CSVDataset(Dataset):
 
         example = {}
 
-        if isinstance(item.prompt, str):
-            item.prompt = [item.prompt]
-
-        if isinstance(item.nprompt, str):
-            item.nprompt = [item.nprompt]
-
         example["prompts"] = item.prompt
         example["nprompts"] = item.nprompt
 
@@ -181,12 +178,9 @@ class CSVDataset(Dataset):
             self.image_cache[item.instance_image_path] = instance_image
 
         example["instance_images"] = instance_image
-        example["instance_prompt_ids"] = self.tokenizer(
-            item.prompt.format(self.instance_identifier),
-            padding="max_length",
-            truncation=True,
-            max_length=self.tokenizer.model_max_length,
-        ).input_ids
+        example["instance_prompt_ids"] = self.prompt_processor.get_input_ids(
+            item.prompt.format(self.instance_identifier)
+        )
 
         if self.num_class_images != 0:
             class_image = Image.open(item.class_image_path)
@@ -194,12 +188,9 @@ class CSVDataset(Dataset):
                 class_image = class_image.convert("RGB")
 
             example["class_images"] = class_image
-            example["class_prompt_ids"] = self.tokenizer(
-                item.prompt.format(self.class_identifier),
-                padding="max_length",
-                truncation=True,
-                max_length=self.tokenizer.model_max_length,
-            ).input_ids
+            example["class_prompt_ids"] = self.prompt_processor.get_input_ids(
+                item.nprompt.format(self.class_identifier)
+            )
 
         self.cache[item.instance_image_path] = example
         return example