From 178c5673184df23356d668a7c16eea04bcdba04b Mon Sep 17 00:00:00 2001
From: Yasser Baalla <70327506+yasser-baalla@users.noreply.github.com>
Date: Wed, 19 Mar 2025 13:56:48 +0100
Subject: [PATCH] Add SemanticImageFetch class for retrieving closest images

---
 __init__.py | 36 ++++++++++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)
 create mode 100644 __init__.py

diff --git a/__init__.py b/__init__.py
new file mode 100644
index 0000000..1997168
--- /dev/null
+++ b/__init__.py
@@ -0,0 +1,36 @@
+import torch
+
+
+class SemanticImageFetch:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE", {"tooltip": "The list of images to fetch the semantic map from."}), 
+                "conditionning": ("CLIP", {"tooltip": "The CLIP model used for encoding the text."}),
+                "number_of_candidates": ("INT", {"default": 1, "min": 1, "max": 10, "tooltip": "Number of closest images to retrieve."}),
+            }
+        }
+    RETURN_TYPES = ("IMAGE", "CONDITIONING",)
+
+    FUNCTION = "fetch"
+
+    CATEGORY = "conditioning"
+    DESCRIPTION = "Select the closest images to an input prompt"
+    
+    def fetch(self, image, conditioning, prompt, number_of_candidates):
+    # Encode the prompt using the CLIP model
+        with torch.no_grad():
+            text_embeddings = conditioning.encode(prompt)
+
+        # Encode the images using the CLIP model
+        image_embeddings = self.encode_images(image, conditioning)
+
+        # Compute similarity between text and image embeddings
+        similarities = self.compute_similarity(text_embeddings, image_embeddings)
+
+        # Select the top-k closest images
+        top_k_indices = torch.topk(similarities, k=number_of_candidates, dim=0).indices
+        closest_images = image[top_k_indices]
+
+        return (closest_images, conditioning)
\ No newline at end of file