From 2e37308cee70a1e83f59495e5c23a971608b86e6 Mon Sep 17 00:00:00 2001 From: nullquant Date: Sun, 12 May 2024 18:21:51 +0300 Subject: [PATCH] memory management --- brushnet_nodes.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/brushnet_nodes.py b/brushnet_nodes.py index 6fe79b9..5b417fa 100644 --- a/brushnet_nodes.py +++ b/brushnet_nodes.py @@ -39,6 +39,10 @@ ComfySDLayers = [comfy.ops.disable_weight_init.Conv2d, comfy.ldm.modules.diffusionmodules.openaimodel.ResBlock ] +ModelsToUnload = [comfy.sd1_clip.SD1ClipModel, + comfy.ldm.models.autoencoder.AutoencoderKL + ] + class BrushNetLoader: @classmethod @@ -282,6 +286,15 @@ class PowerPaint: prompt_embeds_pp = (prompt_embedsA * fitting + (1.0 - fitting) * prompt_embedsB).to(dtype=torch_dtype).to(powerpaint['brushnet'].device) negative_prompt_embeds_pp = (negative_prompt_embedsA * fitting + (1.0 - fitting) * negative_prompt_embedsB).to(dtype=torch_dtype).to(powerpaint['brushnet'].device) + # unload vae and CLIPs + del vae + del clip + for loaded_model in comfy.model_management.current_loaded_models: + if type(loaded_model.model.model) in ModelsToUnload: + comfy.model_management.current_loaded_models.remove(loaded_model) + loaded_model.model_unload() + del loaded_model + # apply patch to code if comfy.samplers.sample.__doc__ is None or 'BrushNet' not in comfy.samplers.sample.__doc__: @@ -364,6 +377,14 @@ class BrushNet: conditioning_latents[0] = conditioning_latents[0].to(dtype=torch_dtype).to(brushnet['brushnet'].device) conditioning_latents[1] = conditioning_latents[1].to(dtype=torch_dtype).to(brushnet['brushnet'].device) + # unload vae + del vae + for loaded_model in comfy.model_management.current_loaded_models: + if type(loaded_model.model.model) in ModelsToUnload: + comfy.model_management.current_loaded_models.remove(loaded_model) + loaded_model.model_unload() + del loaded_model + # prepare embeddings prompt_embeds = positive[0][0].to(dtype=torch_dtype).to(brushnet['brushnet'].device)