WIP creating textual inversion training script

toolkit/config_modules.py

@@ -1,6 +1,6 @@
 import os
 import time
-from typing import List, Optional
+from typing import List, Optional, Literal
 import random
 
 
@@ -50,6 +50,13 @@ class NetworkConfig:
         self.conv_alpha: float = kwargs.get('conv_alpha', self.conv)
 
 
+class EmbeddingConfig:
+    def __init__(self, **kwargs):
+        self.trigger = kwargs.get('trigger', 'custom_embedding')
+        self.tokens = kwargs.get('tokens', 4)
+        self.init_words = kwargs.get('init_phrase', '*')
+        self.save_format = kwargs.get('save_format', 'safetensors')
+
 class TrainConfig:
     def __init__(self, **kwargs):
         self.noise_scheduler = kwargs.get('noise_scheduler', 'ddpm')
@@ -146,6 +153,20 @@ class SliderConfig:
                 self.targets.append(target)
 
 
+class DatasetConfig:
+    caption_type: Literal["txt", "caption"] = 'txt'
+
+    def __init__(self, **kwargs):
+        self.type = kwargs.get('type', 'image')  # sd, slider, reference
+        self.folder_path: str = kwargs.get('folder_path', None)
+        self.default_caption: str = kwargs.get('default_caption', None)
+        self.caption_type: str = kwargs.get('caption_type', None)
+        self.random_scale: bool = kwargs.get('random_scale', False)
+        self.random_crop: bool = kwargs.get('random_crop', False)
+        self.resolution: int = kwargs.get('resolution', 512)
+        self.scale: float = kwargs.get('scale', 1.0)
+
+
 class GenerateImageConfig:
     def __init__(
             self,