From a336127797485aeaa255cd20ecfe64b7f3d52d75 Mon Sep 17 00:00:00 2001 From: matt3o Date: Thu, 2 May 2024 11:14:59 +0200 Subject: [PATCH] add encode batch to lower vram usage of the clip vision encoder --- IPAdapterPlus.py | 15 +++++++++------ README.md | 2 ++ utils.py | 44 +++++++++++++++++++++++++++++++------------- 3 files changed, 42 insertions(+), 19 deletions(-) diff --git a/IPAdapterPlus.py b/IPAdapterPlus.py index e36e831..cf1b9f1 100644 --- a/IPAdapterPlus.py +++ b/IPAdapterPlus.py @@ -164,7 +164,8 @@ def ipadapter_execute(model, neg_embed=None, unfold_batch=False, embeds_scaling='V only', - layer_weights=None): + layer_weights=None, + encode_batch_size=0,): device = model_management.get_torch_device() dtype = model_management.unet_dtype() if dtype not in [torch.float32, torch.float16, torch.bfloat16]: @@ -252,20 +253,20 @@ def ipadapter_execute(model, del image_iface, face if image is not None: - img_cond_embeds = encode_image_masked(clipvision, image) + img_cond_embeds = encode_image_masked(clipvision, image, batch_size=encode_batch_size) if image_composition is not None: - img_comp_cond_embeds = encode_image_masked(clipvision, image_composition) + img_comp_cond_embeds = encode_image_masked(clipvision, image_composition, batch_size=encode_batch_size) if is_plus: img_cond_embeds = img_cond_embeds.penultimate_hidden_states image_negative = image_negative if image_negative is not None else torch.zeros([1, 224, 224, 3]) - img_uncond_embeds = encode_image_masked(clipvision, image_negative).penultimate_hidden_states + img_uncond_embeds = encode_image_masked(clipvision, image_negative, batch_size=encode_batch_size).penultimate_hidden_states if image_composition is not None: img_comp_cond_embeds = img_comp_cond_embeds.penultimate_hidden_states else: img_cond_embeds = img_cond_embeds.image_embeds if not is_faceid else face_cond_embeds if image_negative is not None and not is_faceid: - img_uncond_embeds = encode_image_masked(clipvision, image_negative).image_embeds + img_uncond_embeds = encode_image_masked(clipvision, image_negative, batch_size=encode_batch_size).image_embeds else: img_uncond_embeds = torch.zeros_like(img_cond_embeds) if image_composition is not None: @@ -634,7 +635,7 @@ class IPAdapterAdvanced: FUNCTION = "apply_ipadapter" CATEGORY = "ipadapter" - def apply_ipadapter(self, model, ipadapter, start_at=0.0, end_at=1.0, weight=1.0, weight_style=1.0, weight_composition=1.0, expand_style=False, weight_type="linear", combine_embeds="concat", weight_faceidv2=None, image=None, image_style=None, image_composition=None, image_negative=None, clip_vision=None, attn_mask=None, insightface=None, embeds_scaling='V only', layer_weights=None, ipadapter_params=None): + def apply_ipadapter(self, model, ipadapter, start_at=0.0, end_at=1.0, weight=1.0, weight_style=1.0, weight_composition=1.0, expand_style=False, weight_type="linear", combine_embeds="concat", weight_faceidv2=None, image=None, image_style=None, image_composition=None, image_negative=None, clip_vision=None, attn_mask=None, insightface=None, embeds_scaling='V only', layer_weights=None, ipadapter_params=None, encode_batch_size=0): is_sdxl = isinstance(model.model, (comfy.model_base.SDXL, comfy.model_base.SDXLRefiner, comfy.model_base.SDXL_instructpix2pix)) if 'ipadapter' in ipadapter: @@ -691,6 +692,7 @@ class IPAdapterAdvanced: "embeds_scaling": embeds_scaling, "insightface": insightface if insightface is not None else ipadapter['insightface']['model'] if 'insightface' in ipadapter else None, "layer_weights": layer_weights, + "encode_batch_size": encode_batch_size, } work_model, face_image = ipadapter_execute(work_model, ipadapter_model, clip_vision, **ipa_args) @@ -714,6 +716,7 @@ class IPAdapterBatch(IPAdapterAdvanced): "start_at": ("FLOAT", { "default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001 }), "end_at": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001 }), "embeds_scaling": (['V only', 'K+V', 'K+V w/ C penalty', 'K+mean(V) w/ C penalty'], ), + "encode_batch_size": ("INT", { "default": 0, "min": 0, "max": 4096 }), }, "optional": { "image_negative": ("IMAGE",), diff --git a/README.md b/README.md index fdb6dba..db9a782 100644 --- a/README.md +++ b/README.md @@ -27,6 +27,8 @@ Please consider a [Github Sponsorship](https://github.com/sponsors/cubiq) or [Pa ## Important updates +**2024/05/02**: Add `encode_batch_size` to the Advanced batch node. This can be useful for animations with a lot of frames to reduce the VRAM usage during the image encoding. Please note that results will be slightly different based on the batch size. + **2024/04/27**: Refactored the IPAdapterWeights mostly useful for AnimateDiff animations. **2024/04/21**: Added Regional Conditioning nodes to simplify attention masking and masked text conditioning. diff --git a/utils.py b/utils.py index add3f8e..03b1350 100644 --- a/utils.py +++ b/utils.py @@ -154,21 +154,39 @@ def insightface_loader(provider): model.prepare(ctx_id=0, det_size=(640, 640)) return model -def encode_image_masked(clip_vision, image, mask=None): +def encode_image_masked(clip_vision, image, mask=None, batch_size=0): model_management.load_model_gpu(clip_vision.patcher) - image = image.to(clip_vision.load_device) - - pixel_values = clip_preprocess(image.to(clip_vision.load_device)).float() - - if mask is not None: - pixel_values = pixel_values * mask.to(clip_vision.load_device) - - out = clip_vision.model(pixel_values=pixel_values, intermediate_output=-2) - outputs = Output() - outputs["last_hidden_state"] = out[0].to(model_management.intermediate_device()) - outputs["image_embeds"] = out[2].to(model_management.intermediate_device()) - outputs["penultimate_hidden_states"] = out[1].to(model_management.intermediate_device()) + + if batch_size == 0: + batch_size = image.shape[0] + elif batch_size > image.shape[0]: + batch_size = image.shape[0] + + image_batch = torch.split(image, batch_size, dim=0) + + for img in image_batch: + img = img.to(clip_vision.load_device) + + pixel_values = clip_preprocess(img.to(clip_vision.load_device)).float() + + # TODO: support for multiple masks + if mask is not None: + pixel_values = pixel_values * mask.to(clip_vision.load_device) + + out = clip_vision.model(pixel_values=pixel_values, intermediate_output=-2) + + if not hasattr(outputs, "last_hidden_state"): + outputs["last_hidden_state"] = out[0].to(model_management.intermediate_device()) + outputs["image_embeds"] = out[2].to(model_management.intermediate_device()) + outputs["penultimate_hidden_states"] = out[1].to(model_management.intermediate_device()) + else: + outputs["last_hidden_state"] = torch.cat((outputs["last_hidden_state"], out[0].to(model_management.intermediate_device())), dim=0) + outputs["image_embeds"] = torch.cat((outputs["image_embeds"], out[2].to(model_management.intermediate_device())), dim=0) + outputs["penultimate_hidden_states"] = torch.cat((outputs["penultimate_hidden_states"], out[1].to(model_management.intermediate_device())), dim=0) + + del img, pixel_values, out + return outputs def tensor_to_size(source, dest_size):