Skip to content

Commit 8e1b7a0

Browse files
spacepxlkashifsayakpaul
authored
Fix the deletion of SD3 text encoders for Dreambooth/LoRA training if the text encoders are not being trained (#8536)
* Update train_dreambooth_sd3.py to fix TE garbage collection * Update train_dreambooth_lora_sd3.py to fix TE garbage collection --------- Co-authored-by: Kashif Rasul <[email protected]> Co-authored-by: Sayak Paul <[email protected]>
1 parent 6946fac commit 8e1b7a0

File tree

2 files changed

+6
-0
lines changed

2 files changed

+6
-0
lines changed

examples/dreambooth/train_dreambooth_lora_sd3.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1316,6 +1316,9 @@ def compute_text_embeddings(prompt, text_encoders, tokenizers):
13161316
# Clear the memory here
13171317
if not train_dataset.custom_instance_prompts:
13181318
del tokenizers, text_encoders
1319+
# Explicitly delete the objects as well, otherwise only the lists are deleted and the original references remain, preventing garbage collection
1320+
del tokenizer_one, tokenizer_two, tokenizer_three
1321+
del text_encoder_one, text_encoder_two, text_encoder_three
13191322
gc.collect()
13201323
if torch.cuda.is_available():
13211324
torch.cuda.empty_cache()

examples/dreambooth/train_dreambooth_sd3.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1347,6 +1347,9 @@ def compute_text_embeddings(prompt, text_encoders, tokenizers):
13471347
# Clear the memory here
13481348
if not args.train_text_encoder and not train_dataset.custom_instance_prompts:
13491349
del tokenizers, text_encoders
1350+
# Explicitly delete the objects as well, otherwise only the lists are deleted and the original references remain, preventing garbage collection
1351+
del tokenizer_one, tokenizer_two, tokenizer_three
1352+
del text_encoder_one, text_encoder_two, text_encoder_three
13501353
gc.collect()
13511354
if torch.cuda.is_available():
13521355
torch.cuda.empty_cache()

0 commit comments

Comments
 (0)