ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-05-01 11:41:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Jaret Burkett
2024-09-21 20:54:09 -06:00
parent f85ad452c6
commit 2776221497
5 changed files with 74 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
toolkit/sd_device_states_presets.py
View File

@@ -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