Added option to cache empty prompt or trigger and unload text encoders while training

toolkit/sd_device_states_presets.py

@@ -40,6 +40,7 @@ def get_train_sd_device_state_preset(
         train_adapter: bool = False,
         train_embedding: bool = False,
         train_refiner: bool = False,
+        unload_text_encoder: bool = False,
 ):
     preset = copy.deepcopy(empty_preset)
     if not cached_latents:
@@ -88,4 +89,9 @@ def get_train_sd_device_state_preset(
         preset['unet']['device'] = device
         preset['text_encoder']['device'] = device
 
+    if unload_text_encoder:
+        preset['text_encoder']['training'] = False
+        preset['text_encoder']['requires_grad'] = False
+        preset['text_encoder']['device'] = 'cpu'
+
     return preset