ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-04-29 18:51:37 +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
jobs/process/BaseSDTrainProcess.py
View File

@@ -174,6 +174,7 @@ class BaseSDTrainProcess(BaseTrainProcess):
train_adapter=is_training_adapter,
train_embedding=self.embed_config is not None,
train_refiner=self.train_config.train_refiner,
unload_text_encoder=self.train_config.unload_text_encoder
)
# fine_tuning here is for training actual SD network, not LoRA, embeddings, etc. it is (Dreambooth, etc)
@@ -1750,6 +1751,11 @@ class BaseSDTrainProcess(BaseTrainProcess):
if self.train_config.free_u:
self.sd.pipeline.disable_freeu()
self.sample(self.step_num)
if self.train_config.unload_text_encoder:
# make sure the text encoder is unloaded
self.sd.text_encoder_to('cpu')
flush()
self.ensure_params_requires_grad()
self.progress_bar.unpause()