ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-01-26 16:39:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

Reordered logs

Browse Source
This commit is contained in:
Jaret Burkett
2025-10-15 09:15:06 -06:00
parent 645c54d617
commit 91b87e06a1
1 changed files with 10 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
extensions_built_in/sd_trainer/SDTrainer.py
View File

@@ -304,19 +304,11 @@ class SDTrainer(BaseSDTrainProcess):
# handle unload text encoder
if self.train_config.unload_text_encoder or self.is_caching_text_embeddings:
print_acc("Caching embeddings and unloading text encoder")
with torch.no_grad():
if self.train_config.train_text_encoder:
raise ValueError("Cannot unload text encoder if training text encoder")
# cache embeddings
print_acc("\n***** UNLOADING TEXT ENCODER *****")
if self.is_caching_text_embeddings:
print_acc("Embeddings cached to disk. We dont need the text encoder anymore")
else:
print_acc("This will train only with a blank prompt or trigger word, if set")
print_acc("If this is not what you want, remove the unload_text_encoder flag")
print_acc("***********************************")
print_acc("")
self.sd.text_encoder_to(self.device_torch)
encode_kwargs = {}
if self.sd.encode_control_in_text_embeddings:
@@ -332,6 +324,15 @@ class SDTrainer(BaseSDTrainProcess):
self.diff_output_preservation_embeds = self.sd.encode_prompt(self.train_config.diff_output_preservation_class)
self.cache_sample_prompts()
print_acc("\n***** UNLOADING TEXT ENCODER *****")
if self.is_caching_text_embeddings:
print_acc("Embeddings cached to disk. We dont need the text encoder anymore")
else:
print_acc("This will train only with a blank prompt or trigger word, if set")
print_acc("If this is not what you want, remove the unload_text_encoder flag")
print_acc("***********************************")
print_acc("")
# unload the text encoder
if self.is_caching_text_embeddings: