2 files changed, 11 insertions, 25 deletions

diff --git a/data.py b/data/textual_inversion/csv.py
index 0d1e96e..38ffb6f 100644
--- a/data.py
+++ b/data/textual_inversion/csv.py
@@ -80,14 +80,19 @@ class CSVDataset(Dataset):
 
         self.placeholder_token = placeholder_token
 
-        self.size = size
-        self.center_crop = center_crop
         self.interpolation = {"linear": PIL.Image.LINEAR,
                               "bilinear": PIL.Image.BILINEAR,
                               "bicubic": PIL.Image.BICUBIC,
                               "lanczos": PIL.Image.LANCZOS,
                               }[interpolation]
-        self.flip = transforms.RandomHorizontalFlip(p=flip_p)
+        self.image_transforms = transforms.Compose(
+            [
+                transforms.Resize(size, interpolation=transforms.InterpolationMode.BILINEAR),
+                transforms.CenterCrop(size) if center_crop else transforms.RandomCrop(size),
+                transforms.ToTensor(),
+                transforms.Normalize([0.5], [0.5]),
+            ]
+        )
 
         self.cache = {}
 
@@ -102,9 +107,9 @@ class CSVDataset(Dataset):
 
         example = {}
         image = Image.open(image_path)
-
         if not image.mode == "RGB":
             image = image.convert("RGB")
+        image = self.image_transforms(image)
 
         text = text.format(self.placeholder_token)
 
@@ -117,24 +122,8 @@ class CSVDataset(Dataset):
             return_tensors="pt",
         ).input_ids[0]
 
-        # default to score-sde preprocessing
-        img = np.array(image).astype(np.uint8)
-
-        if self.center_crop:
-            crop = min(img.shape[0], img.shape[1])
-            h, w, = img.shape[0], img.shape[1]
-            img = img[(h - crop) // 2:(h + crop) // 2,
-                      (w - crop) // 2:(w + crop) // 2]
-
-        image = Image.fromarray(img)
-        image = image.resize((self.size, self.size),
-                             resample=self.interpolation)
-        image = self.flip(image)
-        image = np.array(image).astype(np.uint8)
-        image = (image / 127.5 - 1.0).astype(np.float32)
-
         example["key"] = "-".join([image_path, "-", str(flipped)])
-        example["pixel_values"] = torch.from_numpy(image).permute(2, 0, 1)
+        example["pixel_values"] = image
 
         self.cache[image_path] = example
         return example
diff --git a/main.py b/textual_inversion.py
index 51b64c1..aa8e744 100644
--- a/main.py
+++ b/textual_inversion.py
@@ -2,10 +2,7 @@ import argparse
 import itertools
 import math
 import os
-import random
 import datetime
-from pathlib import Path
-from typing import Optional
 
 import numpy as np
 import torch
@@ -25,7 +22,7 @@ from slugify import slugify
 import json
 import os
 
-from data import CSVDataModule
+from data.textual_inversion.csv import CSVDataModule
 
 logger = get_logger(__name__)