import free_memory

Signed-off-by: jiqing-feng <[email protected]>

Author: jiqing-feng

examples/dreambooth/train_dreambooth_lora.py

@@ -54,7 +54,11 @@
 )
 from diffusers.loaders import StableDiffusionLoraLoaderMixin
 from diffusers.optimization import get_scheduler
-from diffusers.training_utils import _set_state_dict_into_text_encoder, cast_training_params
+from diffusers.training_utils import (
+    _set_state_dict_into_text_encoder,
+    cast_training_params,
+    free_memory,
+)
 from diffusers.utils import (
     check_min_version,
     convert_state_dict_to_diffusers,
@@ -75,13 +79,6 @@
 logger = get_logger(__name__)
 
 
-def free_memory():
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-    if hasattr(torch, "xpu") and torch.xpu.is_available():
-        torch.xpu.empty_cache()
-
-
 def save_model_card(
     repo_id: str,
     images=None,

src/diffusers/training_utils.py

@@ -299,6 +299,8 @@ def free_memory():
         torch.mps.empty_cache()
     elif is_torch_npu_available():
         torch_npu.npu.empty_cache()
+    elif hasattr(torch, "xpu") and torch.xpu.is_available():
+        torch.xpu.empty_cache()
 
 
 # Adapted from torch-ema https://github.com/fadel/pytorch_ema/blob/master/torch_ema/ema.py#L14