Fix the deletion of SD3 text encoders for Dreambooth/LoRA training if…

… the text encoders are not being trained (huggingface#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]>
bghira · Jun 16, 2024 · 8e1b7a0 · 8e1b7a0
1 parent 6946fac
commit 8e1b7a0
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/examples/dreambooth/train_dreambooth_lora_sd3.py b/examples/dreambooth/train_dreambooth_lora_sd3.py
@@ -1316,6 +1316,9 @@ def compute_text_embeddings(prompt, text_encoders, tokenizers):
     # Clear the memory here
     if not train_dataset.custom_instance_prompts:
         del tokenizers, text_encoders
+        # Explicitly delete the objects as well, otherwise only the lists are deleted and the original references remain, preventing garbage collection
+        del tokenizer_one, tokenizer_two, tokenizer_three
+        del text_encoder_one, text_encoder_two, text_encoder_three
         gc.collect()
         if torch.cuda.is_available():
             torch.cuda.empty_cache()

diff --git a/examples/dreambooth/train_dreambooth_sd3.py b/examples/dreambooth/train_dreambooth_sd3.py
@@ -1347,6 +1347,9 @@ def compute_text_embeddings(prompt, text_encoders, tokenizers):
     # Clear the memory here
     if not args.train_text_encoder and not train_dataset.custom_instance_prompts:
         del tokenizers, text_encoders
+        # Explicitly delete the objects as well, otherwise only the lists are deleted and the original references remain, preventing garbage collection
+        del tokenizer_one, tokenizer_two, tokenizer_three
+        del text_encoder_one, text_encoder_two, text_encoder_three
         gc.collect()
         if torch.cuda.is_available():
             torch.cuda.empty_cache()