free memory.

Author: sayakpaul

examples/flux-control/train_control_flux.py

@@ -1166,6 +1166,11 @@ def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
             flux_transformer.to(torch.float32)
         flux_transformer.save_pretrained(args.output_dir)
 
+        del flux_transformer
+        del text_encoding_pipeline
+        del vae
+        free_memory()
+
         # Run a final round of validation.
         image_logs = None
         if args.validation_prompt is not None:

examples/flux-control/train_control_lora_flux.py

@@ -1319,6 +1319,11 @@ def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
             transformer_lora_layers=transformer_lora_layers,
         )
 
+        del flux_transformer
+        del text_encoding_pipeline
+        del vae
+        free_memory()
+
         # Run a final round of validation.
         image_logs = None
         if args.validation_prompt is not None: