Added support for new concept slider training script to CLI and UI

extensions_built_in/sd_trainer/__init__.py

@@ -16,8 +16,10 @@ class SDTrainerExtension(Extension):
     def get_process(cls):
         # import your process class here so it is only loaded when needed and return it
         from .SDTrainer import SDTrainer
+
         return SDTrainer
 
+
 # This is for generic training (LoRA, Dreambooth, FineTuning)
 class UITrainerExtension(Extension):
     # uid must be unique, it is how the extension is identified
@@ -32,9 +34,28 @@ class UITrainerExtension(Extension):
     def get_process(cls):
         # import your process class here so it is only loaded when needed and return it
         from .UITrainer import UITrainer
+
         return UITrainer
 
 
+# This is a universal trainer that can be from ui or api
+class DiffusionTrainerExtension(Extension):
+    # uid must be unique, it is how the extension is identified
+    uid = "diffusion_trainer"
+
+    # name is the name of the extension for printing
+    name = "Diffusion Trainer"
+
+    # This is where your process class is loaded
+    # keep your imports in here so they don't slow down the rest of the program
+    @classmethod
+    def get_process(cls):
+        # import your process class here so it is only loaded when needed and return it
+        from .DiffusionTrainer import DiffusionTrainer
+
+        return DiffusionTrainer
+
+
 # for backwards compatability
 class TextualInversionTrainer(SDTrainerExtension):
     uid = "textual_inversion_trainer"
@@ -42,5 +63,8 @@ class TextualInversionTrainer(SDTrainerExtension):
 
 AI_TOOLKIT_EXTENSIONS = [
     # you can put a list of extensions here
-    SDTrainerExtension, TextualInversionTrainer, UITrainerExtension
+    SDTrainerExtension,
+    TextualInversionTrainer,
+    UITrainerExtension,
+    DiffusionTrainerExtension,
 ]