Added base for ultimate slider. WIP

toolkit/config_modules.py

@@ -4,7 +4,6 @@ from typing import List, Optional
 import random
 
 
-
 class SaveConfig:
     def __init__(self, **kwargs):
         self.save_every: int = kwargs.get('save_every', 1000)
@@ -87,6 +86,24 @@ class ModelConfig:
             raise ValueError('name_or_path must be specified')
 
 
+class ReferenceDatasetConfig:
+    def __init__(self, **kwargs):
+        # can pass with a side by side pait or a folder with pos and neg folder
+        self.pair_folder: str = kwargs.get('pair_folder', None)
+        self.pos_folder: str = kwargs.get('pos_folder', None)
+        self.neg_folder: str = kwargs.get('neg_folder', None)
+
+        self.network_weight: float = float(kwargs.get('network_weight', 1.0))
+        self.pos_weight: float = float(kwargs.get('pos_weight', self.network_weight))
+        self.neg_weight: float = float(kwargs.get('neg_weight', self.network_weight))
+        # make sure they are all absolute values no negatives
+        self.pos_weight = abs(self.pos_weight)
+        self.neg_weight = abs(self.neg_weight)
+
+        self.target_class: str = kwargs.get('target_class', '')
+        self.size: int = kwargs.get('size', 512)
+
+
 class SliderTargetConfig:
     def __init__(self, **kwargs):
         self.target_class: str = kwargs.get('target_class', '')

toolkit/data_loader.py

@@ -163,7 +163,7 @@ class PairedImageDataset(Dataset):
             self.pos_file_list = [os.path.join(self.pos_folder, file) for file in os.listdir(self.pos_folder) if
                                   file.lower().endswith(supported_exts)]
             self.neg_file_list = [os.path.join(self.neg_folder, file) for file in os.listdir(self.neg_folder) if
-                                    file.lower().endswith(supported_exts)]
+                                  file.lower().endswith(supported_exts)]
 
             matched_files = []
             for pos_file in self.pos_file_list:
@@ -177,7 +177,6 @@ class PairedImageDataset(Dataset):
             # remove duplicates
             matched_files = [t for t in (set(tuple(i) for i in matched_files))]
 
-
             self.file_list = matched_files
             print(f"  -  Found {len(self.file_list)} matching pairs")
         else:
@@ -190,6 +189,15 @@ class PairedImageDataset(Dataset):
             transforms.Normalize([0.5], [0.5]),  # normalize to [-1, 1]
         ])
 
+    def get_all_prompts(self):
+        prompts = []
+        for index in range(len(self.file_list)):
+            prompts.append(self.get_prompt_item(index))
+
+        # remove duplicates
+        prompts = list(set(prompts))
+        return prompts
+
     def __len__(self):
         return len(self.file_list)
 
@@ -202,19 +210,9 @@ class PairedImageDataset(Dataset):
         else:
             return default
 
-    def __getitem__(self, index):
+    def get_prompt_item(self, index):
         img_path_or_tuple = self.file_list[index]
         if isinstance(img_path_or_tuple, tuple):
-            # load both images
-            img_path = img_path_or_tuple[0]
-            img1 = exif_transpose(Image.open(img_path)).convert('RGB')
-            img_path = img_path_or_tuple[1]
-            img2 = exif_transpose(Image.open(img_path)).convert('RGB')
-            # combine them side by side
-            img = Image.new('RGB', (img1.width + img2.width, max(img1.height, img2.height)))
-            img.paste(img1, (0, 0))
-            img.paste(img2, (img1.width, 0))
-
             # check if either has a prompt file
             path_no_ext = os.path.splitext(img_path_or_tuple[0])[0]
             prompt_path = path_no_ext + '.txt'
@@ -223,7 +221,6 @@ class PairedImageDataset(Dataset):
                 prompt_path = path_no_ext + '.txt'
         else:
             img_path = img_path_or_tuple
-            img = exif_transpose(Image.open(img_path)).convert('RGB')
             # see if prompt file exists
             path_no_ext = os.path.splitext(img_path)[0]
             prompt_path = path_no_ext + '.txt'
@@ -242,6 +239,25 @@ class PairedImageDataset(Dataset):
                 prompt = ', '.join(prompt_split)
         else:
             prompt = self.default_prompt
+        return prompt
+
+    def __getitem__(self, index):
+        img_path_or_tuple = self.file_list[index]
+        if isinstance(img_path_or_tuple, tuple):
+            # load both images
+            img_path = img_path_or_tuple[0]
+            img1 = exif_transpose(Image.open(img_path)).convert('RGB')
+            img_path = img_path_or_tuple[1]
+            img2 = exif_transpose(Image.open(img_path)).convert('RGB')
+            # combine them side by side
+            img = Image.new('RGB', (img1.width + img2.width, max(img1.height, img2.height)))
+            img.paste(img1, (0, 0))
+            img.paste(img2, (img1.width, 0))
+        else:
+            img_path = img_path_or_tuple
+            img = exif_transpose(Image.open(img_path)).convert('RGB')
+
+        prompt = self.get_prompt_item(index)
 
         height = self.size
         # determine width to keep aspect ratio
@@ -252,4 +268,3 @@ class PairedImageDataset(Dataset):
         img = self.transform(img)
 
         return img, prompt, (self.neg_weight, self.pos_weight)
-