Using CLIP Vision Encode nodes instead of propritary image encoding nodes

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

__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",

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

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:

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)