Added inbuild plugins and made one for image referenced. WIP

jobs/process/BaseSDTrainProcess.py

@@ -1,6 +1,9 @@
 import glob
 from collections import OrderedDict
 import os
+from typing import Union
+
+from torch.utils.data import DataLoader
 
 from toolkit.lora_special import LoRASpecialNetwork
 from toolkit.optimizer import get_optimizer
@@ -54,6 +57,7 @@ class BaseSDTrainProcess(BaseTrainProcess):
         self.logging_config = LogingConfig(**self.get_conf('logging', {}))
         self.optimizer = None
         self.lr_scheduler = None
+        self.data_loader: Union[DataLoader, None] = None
 
         self.sd = StableDiffusion(
             device=self.device,
@@ -193,7 +197,7 @@ class BaseSDTrainProcess(BaseTrainProcess):
     def hook_before_train_loop(self):
         pass
 
-    def hook_train_loop(self):
+    def hook_train_loop(self, batch=None):
         # return loss
         return 0.0
 
@@ -358,12 +362,29 @@ class BaseSDTrainProcess(BaseTrainProcess):
             iterable=range(0, self.train_config.steps),
         )
 
+        if self.data_loader is not None:
+            dataloader = self.data_loader
+            dataloader_iterator = iter(dataloader)
+        else:
+            dataloader = None
+            dataloader_iterator = None
+
         # self.step_num = 0
         for step in range(self.step_num, self.train_config.steps):
-            # todo handle dataloader here maybe, not sure
+            if dataloader is not None:
+                try:
+                    batch = next(dataloader_iterator)
+                except StopIteration:
+                    # hit the end of an epoch, reset
+                    # todo, should we do something else here? like blow up balloons?
+                    dataloader_iterator = iter(dataloader)
+                    batch = next(dataloader_iterator)
+            else:
+                batch = None
 
             ### HOOK ###
-            loss_dict = self.hook_train_loop()
+            loss_dict = self.hook_train_loop(batch)
+            flush()
 
             if self.train_config.optimizer.lower().startswith('dadaptation') or \
                     self.train_config.optimizer.lower().startswith('prodigy'):

jobs/process/BaseTrainProcess.py

@@ -29,11 +29,11 @@ class BaseTrainProcess(BaseProcess):
         super().__init__(process_id, job, config)
         self.progress_bar = None
         self.writer = None
-        self.training_folder = self.get_conf('training_folder', self.job.training_folder)
-        self.save_root = os.path.join(self.training_folder, self.job.name)
+        self.training_folder = self.get_conf('training_folder', self.job.training_folder if hasattr(self.job, 'training_folder') else None)
+        self.save_root = os.path.join(self.training_folder, self.name)
         self.step = 0
         self.first_step = 0
-        self.log_dir = self.get_conf('log_dir', self.job.log_dir)
+        self.log_dir = self.get_conf('log_dir', self.job.log_dir if hasattr(self.job, 'log_dir') else None)
         self.setup_tensorboard()
         self.save_training_config()
 
@@ -62,7 +62,7 @@ class BaseTrainProcess(BaseProcess):
 
     def save_training_config(self):
         timestamp = datetime.now().strftime('%Y%m%d-%H%M%S')
-        os.makedirs(self.training_folder, exist_ok=True)
-        save_dif = os.path.join(self.training_folder, f'process_config_{timestamp}.yaml')
+        os.makedirs(self.save_root, exist_ok=True)
+        save_dif = os.path.join(self.save_root, f'process_config_{timestamp}.yaml')
         with open(save_dif, 'w') as f:
             yaml.dump(self.raw_process_config, f)

jobs/process/TrainLoRAHack.py

@@ -68,7 +68,7 @@ class TrainLoRAHack(BaseSDTrainProcess):
 
         return loss_dict
 
-    def hook_train_loop(self):
+    def hook_train_loop(self, batch):
         if self.hack_config.type == 'suppression':
             return self.supress_loop()
         else:

jobs/process/TrainSDRescaleProcess.py

@@ -210,7 +210,7 @@ class TrainSDRescaleProcess(BaseSDTrainProcess):
         flush()
         # end hook_before_train_loop
 
-    def hook_train_loop(self):
+    def hook_train_loop(self, batch):
         dtype = get_torch_dtype(self.train_config.dtype)
 
         loss_function = torch.nn.MSELoss()

jobs/process/TrainSliderProcess.py

@@ -173,7 +173,7 @@ class TrainSliderProcess(BaseSDTrainProcess):
         flush()
         # end hook_before_train_loop
 
-    def hook_train_loop(self):
+    def hook_train_loop(self, batch):
         dtype = get_torch_dtype(self.train_config.dtype)
 
         # get a random pair

jobs/process/TrainSliderProcessOld.py

@@ -221,7 +221,7 @@ class TrainSliderProcessOld(BaseSDTrainProcess):
         flush()
         # end hook_before_train_loop
 
-    def hook_train_loop(self):
+    def hook_train_loop(self, batch):
         dtype = get_torch_dtype(self.train_config.dtype)
 
         # get a random pair

jobs/process/__init__.py

@@ -13,3 +13,4 @@ from .ModRescaleLoraProcess import ModRescaleLoraProcess
 from .GenerateProcess import GenerateProcess
 from .BaseExtensionProcess import BaseExtensionProcess
 from .TrainESRGANProcess import TrainESRGANProcess
+from .BaseSDTrainProcess import BaseSDTrainProcess