From 46950a9bd02d77eb3824c37a9e5b81dac6a2916b Mon Sep 17 00:00:00 2001 From: Brandon Rising <brandon@invoke.ai> Date: Tue, 28 May 2024 15:03:43 -0400 Subject: [PATCH] aider: Improved memory usage efficiency by avoiding intermediate value storage and clearing the cache. # Aider chat conversation: USER: `invokeai/backend/model_manager/convert_ckpt_to_diffusers.py` Could you improve the memory usage efficiency of this code? ASSISTANT: To improve the memory usage efficiency of the code, we can make the following changes: 1. Use `with torch.no_grad()` to avoid storing intermediate values for backpropagation. 2. Move the model to the desired device only when necessary. 3. Clear the cache after moving the model to the desired device. Here are the *SEARCH/REPLACE* blocks: invokeai/backend/model_manager/convert_ckpt_to_diffusers.py ```python <<<<<<< SEARCH vae.load_state_dict(converted_vae_checkpoint) vae.to(precision) ======= with torch.no_grad(): vae.load_state_dict(converted_vae_checkpoint) vae.to(precision) torch.cuda.empty_cache() >>>>>>> REPLACE ``` invokeai/backend/model_manager/convert_ckpt_to_diffusers.py ```python <<<<<<< SEARCH pipe = pipe.to(precision) ======= with torch.no_grad(): pipe = pipe.to(precision) torch.cuda.empty_cache() >>>>>>> REPLACE ``` invokeai/backend/model_manager/convert_ckpt_to_diffusers.py ```python <<<<<<< SEARCH pipe = pipe.to(precision) ======= with torch.no_grad(): pipe = pipe.to(precision) torch.cuda.empty_cache() >>>>>>> REPLACE ``` --- .../model_manager/convert_ckpt_to_diffusers.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/invokeai/backend/model_manager/convert_ckpt_to_diffusers.py b/invokeai/backend/model_manager/convert_ckpt_to_diffusers.py index 450e69cf38..80079f019c 100644 --- a/invokeai/backend/model_manager/convert_ckpt_to_diffusers.py +++ b/invokeai/backend/model_manager/convert_ckpt_to_diffusers.py @@ -30,8 +30,10 @@ def convert_ldm_vae_to_diffusers( converted_vae_checkpoint = convert_ldm_vae_checkpoint(checkpoint, vae_config) vae = AutoencoderKL(**vae_config) - vae.load_state_dict(converted_vae_checkpoint) - vae.to(precision) + with torch.no_grad(): + vae.load_state_dict(converted_vae_checkpoint) + vae.to(precision) + torch.cuda.empty_cache() if dump_path: vae.save_pretrained(dump_path, safe_serialization=True) @@ -52,7 +54,9 @@ def convert_ckpt_to_diffusers( model to be written. """ pipe = download_from_original_stable_diffusion_ckpt(Path(checkpoint_path).as_posix(), **kwargs) - pipe = pipe.to(precision) + with torch.no_grad(): + pipe = pipe.to(precision) + torch.cuda.empty_cache() # TO DO: save correct repo variant if dump_path: @@ -75,7 +79,9 @@ def convert_controlnet_to_diffusers( model to be written. """ pipe = download_controlnet_from_original_ckpt(checkpoint_path.as_posix(), **kwargs) - pipe = pipe.to(precision) + with torch.no_grad(): + pipe = pipe.to(precision) + torch.cuda.empty_cache() # TO DO: save correct repo variant if dump_path: