diff --git a/README.md b/README.md index ed7db3a..991374e 100644 --- a/README.md +++ b/README.md @@ -22,7 +22,7 @@ This repository provides several custom nodes for ComfyUI that enable efficient ## After creating a database, you can search for similar images using a workflow like the one shown above: ![Using a ClipVision-Database](examples/UsingClipVision-Database.png) -- Load an image into the Get Image Embeddings node, connected to a "Load CLIP Vision" node and select a Clip Vision model. → This node generates the embedding data of your input image. +- Load an image into a "CLIP Vision Encode" node, connected to a "Load CLIP Vision" node and select a ClipVision model. → This node generates the image embedding data of your input image. - Pass this embedding data to the "Image Searcher" node, which calculates image similarities using a cosine similarity algorithm. - "Result Browser": Browse the search results - image_index: Index of the image, depending on match direction. @@ -32,12 +32,25 @@ This repository provides several custom nodes for ComfyUI that enable efficient - image_count: the amount of images to output as a image batch - match: see match of "Result Browser" -Additionally you can chain up as many databases has you want by using the optional "img_db" socket. +Additionally you can chain up multiple databases by using the optional "img_db" socket. + → Please note, all databases need to be generated with the same ClipVision model -## You can directly compare the similarity of two images by connecting the outputs of two "Get Image Embeddings" nodes to a "Compare Embeddings" node. +## You can directly compare the similarity of two images by connecting the outputs of two "CLIP Vision Encode" nodes to a "Compare Embeddings" node. ![Simple image comparison](examples/SimpleImageComparison.png) This node calculates the similarity as a factor between 0.0 and 1.0, which can be displayed using Crystools’ “Show any value to console/display” node or any similar output node. +## Using a "CLIP Text Encode" in combination with the experimental node "Condition 2 Embeddings" you can search through the database by text description. +![Simple image comparison](examples/GeneralWorkflow_experimental.png) + + → Please note, not all CLIP text encoder models are compatible with all CLIP vision encoders. Please pay attention to the CLIP vision model with which you created the databse! + I have very good results using the "clip_g" CLIP text encoder model in combination with "CLIP-ViT-bigG-14-laion2B-39B-b160k" CLIP vision model. + + You can check text and vision model compatibility using the "Embeddings Shape Info" node. Models with the same "shape" can be combined. + +## Using a "CLIP Text Encode" and other experimental nodes enables you to search for images in combination with text. +![Simple image comparison](examples/GeneralWorkflow_experimental.png) + → Please note, this is a very experimental workflow! Results may vary due to the algorithms of the "Calculate Embeddings" node. + ## Manual install ```bash # install latest version of ClipVision_Tools diff --git a/__init__.py b/__init__.py index 95ca02f..d223a74 100644 --- a/__init__.py +++ b/__init__.py @@ -1,4 +1,5 @@ -from .clipvision_embeds import GetImageEmbeds +from .clipvision_embeds import EmbedsInfo +from .clipvision_embeds import Cond2Embeds from .clipvision_embeds import CalcEmbeds from .clipvision_embeds import ScaleEmbeds from .clipvision_embeds import CompareEmbeds @@ -17,7 +18,8 @@ import os folder_paths.folder_names_and_paths["EmbDBs"] = ([os.path.join(folder_paths.models_dir, "EmbDBs")], [".json"]) NODE_CLASS_MAPPINGS = { - "GetImageEmbeds": GetImageEmbeds, + "EmbedsInfo": EmbedsInfo, + "Cond2Embeds": Cond2Embeds, "ScaleEmbeds": ScaleEmbeds, "CalcEmbeds": CalcEmbeds, "CompareEmbeds": CompareEmbeds, @@ -31,7 +33,8 @@ NODE_CLASS_MAPPINGS = { # A dictionary that contains the friendly/humanly readable titles for the nodes NODE_DISPLAY_NAME_MAPPINGS = { - "GetImageEmbeds": "Get Image Embeddings", + "EmbedsInfo": "Embeddings Shape Info (experimental)", + "Cond2Embeds": "Condition 2 Embeddings (experimental)", "ScaleEmbeds": "Scale Embeddings (experimental)", "CalcEmbeds": "Calculate Embeddings (experimental)", "CompareEmbeds": "Compare Embeddings", diff --git a/clipvision_embeds.py b/clipvision_embeds.py index a408eeb..c5edca7 100644 --- a/clipvision_embeds.py +++ b/clipvision_embeds.py @@ -1,9 +1,10 @@ from typing import Any import numpy as np -from .utils import (get_image_clip_embeddings) - -class GetImageEmbeds: +import torch +from comfy.clip_vision import Output +from comfy.comfy_types import IO +class EmbedsInfo: def __init__(self): pass @@ -11,18 +12,41 @@ class GetImageEmbeds: def INPUT_TYPES(s): return { "required": { - "clip_vision": ("CLIP_VISION",), - "image": ("IMAGE",), - } + "clip_vision_output": ("CLIP_VISION_OUTPUT",), + }, } - RETURN_TYPES = ("img_emb", ) - RETURN_NAMES = ("IMG_EMB", ) - FUNCTION = "GetImgEmbeds" - CATEGORY = "ClipVisionTools" + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("CLIP TEXT Info",) + FUNCTION = "cond_to_embeds" + CATEGORY = "ClipVisionTools/experimental" - def GetImgEmbeds(self, clip_vision, image): - return get_image_clip_embeddings(clip_vision, image), + def cond_to_embeds(self, clip_vision_output): + info = clip_vision_output["image_embeds"].shape + return info, + +class Cond2Embeds: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "conditioning": (IO.CONDITIONING,), + }, + } + + RETURN_TYPES = ("CLIP_VISION_OUTPUT", ) + RETURN_NAMES = ("CLIP_VISION_OUTPUT", ) + FUNCTION = "cond_to_embeds" + CATEGORY = "ClipVisionTools/experimental" + + def cond_to_embeds(self, conditioning): + tmp = Output() + cond = conditioning[0][1].get("pooled_output", None) + tmp["image_embeds"] = cond + return tmp, class ScaleEmbeds: def __init__(self): @@ -32,21 +56,25 @@ class ScaleEmbeds: def INPUT_TYPES(s): return { "required": { - "img_emb1": ("img_emb",), + "clip_vision_output": ("CLIP_VISION_OUTPUT",), "scale": ("FLOAT", { "default": 1, "step": 0.01 }), }, } - RETURN_TYPES = ("img_emb", ) - RETURN_NAMES = ("IMG_EMBEDS", ) + RETURN_TYPES = ("CLIP_VISION_OUTPUT", ) + RETURN_NAMES = ("clip_vision_output", ) FUNCTION = "Scale_Embeds" - CATEGORY = "ClipVisionTools" + CATEGORY = "ClipVisionTools/experimental" - def Scale_Embeds(self, img_emb1, scale): - image_embeds = (np.array(img_emb1) * scale) - return image_embeds, + def Scale_Embeds(self, clip_vision_output, scale): + tmp = Output() + cv1 = clip_vision_output["image_embeds"].numpy().flatten().tolist() + oshape = clip_vision_output["image_embeds"].shape + image_embeds = np.array(cv1) * scale + tmp["image_embeds"] = torch.tensor(np.array(image_embeds).reshape(oshape)) + return tmp, #These algorithms are experimental. Some of them seem to work, some of them seem to have unexpected results. class CalcEmbeds: @@ -58,29 +86,36 @@ class CalcEmbeds: return { "required": { "calculation": (["normalize", "add", "subtract", "most common", "remove image2 from image1", "average of both images", "or", "multiply"], {"default": "subtract"}), - "img_emb1": ("img_emb",), + "clip_vision_output1": ("CLIP_VISION_OUTPUT",), }, "optional": { - "img_emb2": ("img_emb",), + "clip_vision_output2": ("CLIP_VISION_OUTPUT",), } } - RETURN_TYPES = ("img_emb", ) - RETURN_NAMES = ("IMG_EMBEDS", ) + RETURN_TYPES = ("CLIP_VISION_OUTPUT", ) + RETURN_NAMES = ("clip_vision_output", ) FUNCTION = "Calc_Embeds" - CATEGORY = "ClipVisionTools" + CATEGORY = "ClipVisionTools/experimental" + + def Calc_Embeds(self, calculation, clip_vision_output1, clip_vision_output2=None): + tmp = Output() + cv1 = clip_vision_output1["image_embeds"].numpy().flatten().tolist() + oshape = clip_vision_output1["image_embeds"].shape + + if clip_vision_output2 is not None: + cv2 = clip_vision_output2["image_embeds"].numpy().flatten().tolist() - def Calc_Embeds(self, calculation, img_emb1, img_emb2=None): if calculation == "normalize": - if img_emb2 is not None: - a_min, a_max = np.min(np.array(img_emb1)), np.max(np.array(img_emb1)) - b_min, b_max = np.min(np.array(img_emb2)), np.max(np.array(img_emb2)) + if cv2 is not None: + a_min, a_max = np.min(np.array(cv1)), np.max(np.array(cv1)) + b_min, b_max = np.min(np.array(cv2)), np.max(np.array(cv2)) if np.isclose(a_max, a_min): - image_embeds = np.full_like(np.array(img_emb1), (b_min + b_max) / 2) + image_embeds = np.full_like(np.array(cv1), (b_min + b_max) / 2) else: - image_embeds = (np.array(img_emb1) - a_min) / (a_max - a_min) * (b_max - b_min) + b_min + image_embeds = (np.array(cv1) - a_min) / (a_max - a_min) * (b_max - b_min) + b_min else: - x = np.array(img_emb1, dtype=float) + x = np.array(cv1, dtype=float) min_x = np.min(x) max_x = np.max(x) if max_x == min_x: @@ -92,28 +127,30 @@ class CalcEmbeds: image_embeds = norm * 2 - 1 if calculation == "add": - image_embeds = (np.array(img_emb1) + np.array(img_emb2)) + image_embeds = (np.array(cv1) + np.array(cv2)) if calculation == "subtract": - image_embeds = (np.array(img_emb1) - np.array(img_emb2)) + image_embeds = (np.array(cv1) - np.array(cv2)) if calculation == "most common": - cos_sim = np.dot(np.array(img_emb1), np.array(img_emb2)) / (np.linalg.norm(np.array(img_emb1)) * np.linalg.norm(np.array(img_emb2))) - image_embeds = cos_sim * (np.array(img_emb1) + np.array(img_emb2)) / 2 + cos_sim = np.dot(np.array(cv1), np.array(cv2)) / (np.linalg.norm(np.array(cv1)) * np.linalg.norm(np.array(cv2))) + image_embeds = cos_sim * (np.array(cv1) + np.array(cv2)) / 2 if calculation == "remove image2 from image1": - image_embeds = (np.array(img_emb1) - np.array(img_emb2)) + image_embeds = (np.array(cv1) - np.array(cv2)) image_embeds = image_embeds / np.linalg.norm(image_embeds) if calculation == "average of both images": - image_embeds = (np.array(img_emb1) + np.array(img_emb2)) / 2 + image_embeds = (np.array(cv1) + np.array(cv2)) / 2 if calculation == "or": - image_embeds = np.minimum(np.array(img_emb1), np.array(img_emb2)) + image_embeds = np.minimum(np.array(cv1), np.array(cv2)) if calculation == "multiply": - image_embeds = (np.array(img_emb1) * np.array(img_emb2)) - return image_embeds, + image_embeds = (np.array(cv1) * np.array(cv2)) + + tmp["image_embeds"] = torch.tensor(np.array(image_embeds).reshape(oshape)) + return tmp, class CompareEmbeds: def __init__(self): @@ -123,8 +160,8 @@ class CompareEmbeds: def INPUT_TYPES(s): return { "required": { - "img_emb1": ("img_emb",), - "img_emb2": ("img_emb",), + "clip_vision_output1": ("CLIP_VISION_OUTPUT",), + "clip_vision_output2": ("CLIP_VISION_OUTPUT",), } } @@ -133,7 +170,9 @@ class CompareEmbeds: FUNCTION = "Compare_Embeds" CATEGORY = "ClipVisionTools" - def Compare_Embeds(self, img_emb1, img_emb2): + def Compare_Embeds(self, clip_vision_output1, clip_vision_output2): + img_emb1 = clip_vision_output1["image_embeds"].numpy().flatten().tolist() + img_emb2 = clip_vision_output2["image_embeds"].numpy().flatten().tolist() image_norm = np.linalg.norm(img_emb1) feature_norm = np.linalg.norm(img_emb2) dot_products = np.dot(img_emb1, img_emb2) # Matrix-Vector-Product diff --git a/examples/CreatingClipVision-Database.png b/examples/CreatingClipVision-Database.png index 8076c87..a11b091 100644 Binary files a/examples/CreatingClipVision-Database.png and b/examples/CreatingClipVision-Database.png differ diff --git a/examples/GeneralWorkflow_experimental.png b/examples/GeneralWorkflow_experimental.png new file mode 100644 index 0000000..0ea310d Binary files /dev/null and b/examples/GeneralWorkflow_experimental.png differ diff --git a/examples/GeneralWorkflow_experimental2.png b/examples/GeneralWorkflow_experimental2.png new file mode 100644 index 0000000..dd97df1 Binary files /dev/null and b/examples/GeneralWorkflow_experimental2.png differ diff --git a/examples/Overview.png b/examples/Overview.png index 5b5cf19..b444c21 100644 Binary files a/examples/Overview.png and b/examples/Overview.png differ diff --git a/examples/SimpleImageComparison.png b/examples/SimpleImageComparison.png index e7b26c1..a0bb242 100644 Binary files a/examples/SimpleImageComparison.png and b/examples/SimpleImageComparison.png differ diff --git a/examples/UsingClipVision-Database.png b/examples/UsingClipVision-Database.png index 60e49f7..6d1f3c0 100644 Binary files a/examples/UsingClipVision-Database.png and b/examples/UsingClipVision-Database.png differ diff --git a/image_search.py b/image_search.py index 8b4577d..394ae6b 100644 --- a/image_search.py +++ b/image_search.py @@ -1,7 +1,8 @@ from typing import Any import numpy as np from .clipvision_db import LoadDB -from .utils import (get_image, repairImage) +from .utils import (repairImage, image_to_tensor) +from PIL import Image, ImageFile, UnidentifiedImageError from nodes import (ImageBatch) class results_class: distances: Any = None @@ -23,7 +24,7 @@ class ImageSearcher: return { "required": { "img_db": ("LoadDB",), - "img_emb": ("img_emb",), + "clip_vision_output": ("CLIP_VISION_OUTPUT",), }, } @@ -32,9 +33,10 @@ class ImageSearcher: FUNCTION = "get_similar_images" CATEGORY = "ClipVisionTools" - def get_similar_images(self, img_db: LoadDB, img_emb): + def get_similar_images(self, img_db: LoadDB, clip_vision_output): distances = [] - image_embeds = img_emb + + image_embeds = clip_vision_output["image_embeds"].numpy().flatten().tolist() clip_features = img_db.LOADED_DB self.LOADED_clip_features = clip_features @@ -98,6 +100,7 @@ class ResultBrowser: CATEGORY = "ClipVisionTools" def Result_Browser(self, results, image_index, match): + ImageFile.LOAD_TRUNCATED_IMAGES = True distances1 = results.distances clip_features = results.clip_features @@ -111,7 +114,7 @@ class ResultBrowser: idx, score = distances_sorted[image_index] file_name, _ = clip_features[idx] - t_image = repairImage(get_image(file_name)) + t_image = repairImage(image_to_tensor(Image.open(file_name))) return t_image, file_name, score class ResultBrowserAdvanced: @@ -138,6 +141,8 @@ class ResultBrowserAdvanced: CATEGORY = "ClipVisionTools" def AdvResult_Browser(self, results, offset_index, image_count, match, batch_frame_image=None ): + ImageFile.LOAD_TRUNCATED_IMAGES = True + scores = [] filenames = [] distances1 = results.distances @@ -159,7 +164,8 @@ class ResultBrowserAdvanced: for i in range(start_idx, end_idx): idx, score = distances_sorted[i] file_name, _ = clip_features[idx] - t_image = repairImage(get_image(file_name)) + t_image = repairImage(image_to_tensor(Image.open(file_name).convert("RGB"))) + if batch_frame_image is None: batch_frame_image = t_image else: diff --git a/utils.py b/utils.py index d8d7c22..034da78 100644 --- a/utils.py +++ b/utils.py @@ -1,27 +1,28 @@ -import json +import orjson import os from pathlib import Path import numpy as np - +from nodes import CLIPTextEncode import comfy.model_management as model_management import torch -from comfy.clip_vision import Output, clip_preprocess +from comfy.clip_vision import Output, clip_preprocess, ClipVisionModel from PIL import Image, ImageFile, UnidentifiedImageError from torchvision.transforms.functional import pil_to_tensor from tqdm import tqdm -def generate_clip_features_json(clip_vision, path_to_images_folder: Path, +def generate_clip_features_json(clip_vision: ClipVisionModel, path_to_images_folder: Path, output_json_path: Path): clip_features = [] + errors = [] + ImageFile.LOAD_TRUNCATED_IMAGES = True image_path_list = list(path_to_images_folder.glob('**/*.*')) imagetypes = {".png", ".jpg", ".jpeg", ".gif", ".bmp", ".tiff", ".tif", ".webp"} - model_management.load_model_gpu(clip_vision.patcher) - # print("--------- Generating embeddings for images ---------") + image_tqdm = tqdm(image_path_list) - errors = "" + for image_path in image_tqdm: if os.path.splitext(str(image_path.name).lower())[1] in imagetypes: image_tqdm.set_description("Processing " + str(image_path.name)) @@ -52,34 +53,20 @@ def generate_clip_features_json(clip_vision, path_to_images_folder: Path, if error != "": errors = errors + error + "\r\n" - with open(output_json_path, 'w') as f: - json.dump(clip_features, f) + with open(output_json_path, "wb") as f: # Binary-Mode! + f.write(orjson.dumps(clip_features)) + return errors +def get_image_clip_embeddings(clip_vision:ClipVisionModel, image: Image): + with torch.no_grad(): # Disable gradient calculation during inference + cv_image = clip_vision.encode_image(image_to_tensor(image), crop=False) + return cv_image["image_embeds"].numpy().flatten().tolist() -def get_image_clip_embeddings(clip_vision, image: Image): - model_management.load_model_gpu(clip_vision.patcher) - if type(image) is not torch.Tensor: - image = image_to_tensor(image) - image = image.to(clip_vision.load_device) - - pixel_values = clip_preprocess(image, crop=False) - 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()) - - return outputs["image_embeds"].numpy().flatten().tolist() - -def get_image(file_name: str): - image = image_to_tensor(Image.open(file_name)) - return image +#def get_image(file_name: str): +# image = image_to_tensor(Image.open(file_name)) +# return image def image_to_tensor(image): tensor = torch.clamp(pil_to_tensor(image).float() / 255., 0, 1)