nodes/cubiq-ComfyUI_IPAdapter_plus
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files
utils.py
+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):