Hude rework to move the batch to a DTO to make it far more modular to the future ui

extensions_built_in/sd_trainer/SDTrainer.py

@@ -35,6 +35,7 @@ class SDTrainer(BaseSDTrainProcess):
     def hook_train_loop(self, batch):
         dtype = get_torch_dtype(self.train_config.dtype)
         noisy_latents, noise, timesteps, conditioned_prompts, imgs = self.process_general_training_batch(batch)
+        network_weight_list = batch.get_network_weight_list()
 
         self.optimizer.zero_grad()
         flush()
@@ -53,6 +54,9 @@ class SDTrainer(BaseSDTrainProcess):
         else:
             network = BlankNetwork()
 
+        # set the weights
+        network.multiplier = network_weight_list
+
         # activate network if it exits
         with network:
             with torch.set_grad_enabled(grad_on_text_encoder):
@@ -114,5 +118,7 @@ class SDTrainer(BaseSDTrainProcess):
         loss_dict = OrderedDict(
             {'loss': loss.item()}
         )
+        # reset network multiplier
+        network.multiplier = 1.0
 
         return loss_dict

extensions_built_in/sd_trainer/config/train.example.yaml

@@ -19,7 +19,7 @@ config:
         max_step_saves_to_keep: 5 # only affects step counts
       datasets:
         - folder_path: "/path/to/dataset"
-          caption_type: "txt"
+          caption_ext: "txt"
           default_caption: "[trigger]"
           buckets: true
           resolution: 512