add encode batch to lower vram usage of the clip vision encoder

This commit is contained in:
matt3o
2024-05-02 11:14:59 +02:00
parent d7ff848ac8
commit a336127797
3 changed files with 42 additions and 19 deletions
+9 -6
View File
@@ -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",),
+2
View File
@@ -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.
+31 -13
View File
@@ -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):