Complete reqork of how slider training works and optimized it to hell. Can run entire algorythm in 1 batch now with less VRAM consumption than a quarter of it used to take

jobs/process/BaseSDTrainProcess.py

@@ -242,6 +242,12 @@ class BaseSDTrainProcess(BaseTrainProcess):
             unet.enable_xformers_memory_efficient_attention()
         if self.train_config.gradient_checkpointing:
             unet.enable_gradient_checkpointing()
+            # if isinstance(text_encoder, list):
+            #     for te in text_encoder:
+            #         te.enable_gradient_checkpointing()
+            # else:
+            #     text_encoder.enable_gradient_checkpointing()
+
         unet.to(self.device_torch, dtype=dtype)
         unet.requires_grad_(False)
         unet.eval()
@@ -281,6 +287,9 @@ class BaseSDTrainProcess(BaseTrainProcess):
                 default_lr=self.train_config.lr
             )
 
+            if self.train_config.gradient_checkpointing:
+                self.network.enable_gradient_checkpointing()
+
             latest_save_path = self.get_latest_save_path()
             if latest_save_path is not None:
                 self.print(f"#### IMPORTANT RESUMING FROM {latest_save_path} ####")