diff --git a/.gitignore b/.gitignore index eeb8a6e..7fbf810 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,2 @@ **/__pycache__ +.mypy_cache diff --git a/README.md b/README.md index 54ec13d..c2cc4db 100644 --- a/README.md +++ b/README.md @@ -9,11 +9,377 @@ Provides nodes for: * Integer Arithmetic * Floating Point Arithmetic and Functions * Vec2, Vec3, and Vec4 Arithmetic and Functions +* Vector Database (Qdrant) ## Installation -From the `custom_nodes` directory in your ComfyUI installation, run: +### Easy + +ComfyMath can be installed using [ComfyUI-Manager](https://github.com/ltdrdata/ComfyUI-Manager). + +### Manual + +From the `custom_nodes` directory in your ComfyUI installation, with your venv activated, run: ```sh git clone https://github.com/evanspearman/ComfyMath.git +cd ComfyMath +pip install -r requirements.txt ``` + +## Nodes + +### Boolean Logic + +#### BooleanUnaryOperation + +Perform an operation on a single boolean value. Supported operations are: + +* Not + +#### BooleanBinaryOperation + +Perform an operation on two boolean values. Supported operations are: + +* Nor +* Xor +* Nand +* And +* Xnor +* Or +* Eq +* Neq + +### Integer Arithmetic and Logic + +#### IntUnaryOperation + +Perform an operation on a sigle integer value. Supported operations are: + +* Abs +* Neg +* Inc +* Dec +* Sqr +* Cube +* Not +* Factorial + +#### IntUnaryCondition + +Perform a comparison on a single integer value. Supported conditions are: + +* IsZero +* IsNonZero +* IsPositive +* IsNegative +* IsEven +* IsOdd + + +#### IntBinaryOperation + +Perform an operation on two integer values. Supported operations are: + +* Add +* Sub +* Mul +* Div +* Mod +* Pow +* And +* Nand +* Or +* Nor +* Xor +* Xnor +* Shl +* Shr +* Max +* Min + +#### IntBinaryCondition + +Perform a comparison on two integer values. Supported conditions are: + +* Eq +* Neq +* Gt +* Lt +* Geq +* Leq + + +### Floating Point Math + +#### FloatUnaryOperation + +Perform an operation on a single floating point value. Supported operations are: + +* Neg +* Inc +* Dec +* Abs +* Sqr +* Cube +* Sqrt +* Exp +* Ln +* Log10 +* Log2 +* Sin +* Cos +* Tan +* Asin +* Acos +* Atan +* Sinh +* Cosh +* Tanh +* Asinh +* Acosh +* Atanh +* Round +* Floor +* Ceil +* Trunc +* Erf +* Erfc +* Gamma +* Radians +* Degrees + +#### FloatUnaryCondition + +Perform a comparison on a single floating point value. Supported conditions are: + +* IsZero +* IsPositive +* IsNegative +* IsNonZero +* IsPositiveInfinity +* IsNegativeInfinity +* IsNaN +* IsFinite +* IsInfinite +* IsEven +* IsOdd + +#### FloatBinaryOperation + +Perform an operation on two floating point values. Supported operations are: + +* Add +* Sub +* Mul +* Div +* Mod +* Pow +* FloorDiv +* Max +* Min +* Log +* Atan2 + +#### FloatBinaryCondition + +Perform a comparison on two floating point values. Supported conditions are: + +* Eq +* Neq +* Gt +* Gte +* Lt +* Lte + +### Number Math + +`NUMBER` is a type found in some custom nodes that can be either an `int` or a `float`. + +#### NumberUnaryOperation + +Same operations as FloatUnaryOperation + +#### NumberUnaryCondition + +Same conditions as FloatUnaryCondition + +#### NumberBinaryOperation + +Same conditions as FloatBinaryOperation + +#### NumberBinaryCondition + +Same conditions as FloatBinaryCondition + +### Vector Math + +Nodes for performing vector math operations in Euclidean 2-space, 3-space, and 4-space. Under the hood, the `VEC2`, `VEC3`, and `VEC4` types are implemented as `tuple[float, float]`, `tuple[float, float, float]`, and `tuple[float, float, float, float]` respectively. Each size of vector has it's own set of nodes, but the nodes have all the same operations and conditions. The actual processing is performed using numpy. + +#### VecNUnaryOperation + +Perform an operation on a single vector. Available operations are: + +* Neg +* Normalize + +#### VecNUnaryCondition + +Perform a comparison on a single vector. Available conditions are: + +* IsZero +* IsNotZero +* IsNormalized +* IsNotNormalized + +#### VecNToScalarUnaryOperation + +Perform an operation on a single vector that results in a scalar. Available operations are: + +* Norm + +#### VecNBinaryOperation + +Perform an operation on two vectors. Available operations are: + +* Add +* Sub +* Cross + +#### VecNBinaryCondition + +Perform a comparision on two vectors. Available conditions are: + +* Eq +* Neq + +#### VecNToScalarBinaryOperation + +Perform an operation on two vectors that results in a scalar. Available operations are: + +* Dot +* Distance + +#### VecNScalarOperation + +Perform an operation on a vector and a scalar. Available operations are: + +* Mul +* Div + +### Type Conversion + +Nodes to convert between different types. + +#### BoolToInt + +`True` is converted to `1` and `False` is converted to `0` + +#### IntToBool + +`0` is converted to `False` and Non-zero is converted to `True` + +#### FloatToInt + +#### IntToFloat + +#### IntToNumber + +#### NumberToInt + +#### FloatToNumber + +#### NumberToFloat + +#### ComposeVecN + +Build a vector by composing floating point values + +#### FillVecN + +Build a vector by repeating a single floating point value + +#### BreakoutVecN + +Retrieve the floating point values that make up a vector + +### Graphics + +#### SDXLResolution + +Allows for selecting one of the officially supported resolutions of SDXL-based models and outputs the width and height. + +#### NearestSDXLResolution + +Given an `IMAGE` find the SDXL resolution that has the closest aspect ratio. This is useful for Image to Image or ControlNet workflows where you want the image to have as close as possible an aspect ratio to the original image. + +### Data Structures + +Utilities for working with different data structures such as lists and dictionaries. + +#### AddStringToDict + +If no value is provided for `input_dict`, create a new dictionary and set the value of key `key` to `value. + +If `input_dict` is provided, add or override the value of key `key` to `value`. + +#### RetrieveStringFromDict + +Give a key and existing dictionary with string keys, retrieve the value of key `key`. If the value is not a `STRING` or there is no such key, an error will occur. + +#### StringAtIndex + +Given a list of `STRING`, retrieve the value at the given index. + +#### FloatAtIndex + +Given a list of `FLOAT`, retrieve the value at the given index. + +### Image Files + +Nodes for dealing with image files on disk. + +#### LoadImageFromPath + +Given a `STRING` representing the path to an image file on the local file system, load the image as an `IMAGE`. + +#### GetImageSequence + +Given a `STRING` representing the path to an image file on the local file system, return a list of paths (including the given path) where the last block of digits in the file name contains the _only_ differences compared to the given file. For instance, `image_1000.png` and `image_0210.png` would be considered part of the same sequence. `image_100_1000.png` and `image_200_1001.png` would not. `image_1000.png` and `render_1001.png` would also not. Note that this node does not try to load the images into memory or validate that they are image files. + +### Vector Databases + +Nodes for interacting with the vector database Qdrant. Vector databases allow indexing data on vectors, including large vectors such as embeddings. They are useful for searching for data that is "near" a given vector. This is useful for looking for semantically similar text, or images that are visually similar or contain similar subjects. + +#### QdrantConnectionFromFile + +Either opens or creates a connection to a Qdrant database at a given path. + +#### QdrantCollection + +Either creates or references an existing collection in the connected Qdrant database. The `vector_size` and `distance` inputs are optional and are only used when creating a new collection. Otherwise they are ignored. All vectors inserted into the collection must be of size `vector_size`. For CLIP, this size should be 1280 because a CLIP embedding has 1280 dimensions. `distance` sets the type of distance calculation that is used for searching the collection to determine how far apart a given vector is from the vector used as the search query. `Cosine` refers to consine similarity, `Euclid` refers to the euclidean distance, and `Dot` is the dot product. + +#### CLIPVisionOutputToQdrantVector + +This node is used to convert a `CLIP_VISION_OUTPUT` value into a format that can be inserted into a Qdrant collection. + + +#### ConditioningToQdrantVector + +This node is used to convert a `CONDITIONING` value into a format that can be inserted into a Qdrant collection. + +#### QdrantInsertVector + +Insert a given vector into a given Qdrant collection. Optionally include a `DICT` as a payload. The id (a `STRING` representing a UUID) is output, which can be used to retrieve the inserted entry from the collection. The payload is other data that is associated with the vector. For instance, the original text that the vector is an embedding of, or the path to an image that the vector is an embedding of. This node functions as an output node so no other output node needs to exist in the workflow in order to execute it. + +#### QdrantRetrievePayloadById + +Given an id, get the payload from the given qdrant collection that is associated with that id. This is useful after performing a search as the QdrantSearch node does not retrieve the payloads of similar vectors it found. Only the associated ids. + +#### QdrantSearch + +Given a vector search for nearby vectors in the given collection. The limit value is the maximum number of entries to output. This node does not output the vectors or payloads of the nearby vectors, rather, it outputs a list of ids associated with those vectors in the collection and a list of scores that represent how close the vector associated with the id at the same index in the lists is to the original vector. Use the QdrantRetrievePayloadById node to retrieve the payload associated with the output ids. + +#### QdrantInsertImageSequence + +Given an image sequence (list of paths to image files) and a CLIP Vision Model, encode the images using the CLIP Vision model and insert the embedding into the given collection. Each embedding will be inserted with a payload of `{"path": "/path/to/image"}`. This node functions as an output node so no other output node needs to exist in the workflow in order to execute it. diff --git a/__init__.py b/__init__.py index 0888cb3..12d5b47 100644 --- a/__init__.py +++ b/__init__.py @@ -1,14 +1,12 @@ -from .src.comfymath.convert import NODE_CLASS_MAPPINGS as convert_NCM from .src.comfymath.bool import NODE_CLASS_MAPPINGS as bool_NCM -from .src.comfymath.int import NODE_CLASS_MAPPINGS as int_NCM -from .src.comfymath.float import NODE_CLASS_MAPPINGS as float_NCM -from .src.comfymath.number import NODE_CLASS_MAPPINGS as number_NCM -from .src.comfymath.vec import NODE_CLASS_MAPPINGS as vec_NCM from .src.comfymath.control import NODE_CLASS_MAPPINGS as control_NCM +from .src.comfymath.convert import NODE_CLASS_MAPPINGS as convert_NCM +from .src.comfymath.float import NODE_CLASS_MAPPINGS as float_NCM from .src.comfymath.graphics import NODE_CLASS_MAPPINGS as graphics_NCM - - - +from .src.comfymath.int import NODE_CLASS_MAPPINGS as int_NCM +from .src.comfymath.number import NODE_CLASS_MAPPINGS as number_NCM +from .src.comfymath.rag import NODE_CLASS_MAPPINGS as rag_NCM +from .src.comfymath.vec import NODE_CLASS_MAPPINGS as vec_NCM NODE_CLASS_MAPPINGS = { **convert_NCM, @@ -19,6 +17,7 @@ NODE_CLASS_MAPPINGS = { **vec_NCM, **control_NCM, **graphics_NCM, + **rag_NCM, } diff --git a/pyproject.toml b/pyproject.toml index 41689f3..62f49e5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,10 +9,14 @@ readme = "README.md" [tool.poetry.dependencies] python = "^3.10" numpy = "^1.25.1" +qdrant-client = "^1.4.0" +tqdm = "^4.66.1" +click = "^8.1.7" [tool.poetry.group.dev.dependencies] mypy = "^1.4.1" black = "^23.7.0" +isort = "^5.12.0" [build-system] requires = ["poetry-core"] diff --git a/requirements.txt b/requirements.txt index 24ce15a..4157c03 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1 +1,4 @@ numpy +qdrant-client +tqdm +click diff --git a/src/comfymath/graphics.py b/src/comfymath/graphics.py index 6a8d4c1..a17315e 100644 --- a/src/comfymath/graphics.py +++ b/src/comfymath/graphics.py @@ -1,5 +1,7 @@ from typing import Any, Mapping +from torch import Tensor + SDXL_SUPPORTED_RESOLUTIONS = [ (1024, 1024, 1.0), @@ -45,7 +47,7 @@ class NearestSDXLResolution: FUNCTION = "op" CATEGORY = "math/graphics" - def op(self, image) -> tuple[int, int]: + def op(self, image: Tensor) -> tuple[int, int]: image_width = image.size()[2] image_height = image.size()[1] print(f"Input image resolution: {image_width}x{image_height}") diff --git a/src/comfymath/rag.py b/src/comfymath/rag.py new file mode 100644 index 0000000..b1c106b --- /dev/null +++ b/src/comfymath/rag.py @@ -0,0 +1,475 @@ +import re + +from typing import Any, Optional, Protocol, TypeAlias +from pathlib import Path + +import numpy +import torch + +from qdrant_client import QdrantClient +from qdrant_client.models import Distance, VectorParams, CollectionStatus, PointStruct +from uuid import uuid4 +from torch import Tensor +from tqdm import tqdm +from PIL import Image, ImageOps + +_qdrant_connections: dict[str, QdrantClient] = {} + + +def _get_qdrant_connection(connection_id: str) -> QdrantClient: + if connection_id not in _qdrant_connections: + raise ValueError(f"Connection {connection_id} not found") + return _qdrant_connections[connection_id] + + +def _get_unused_qdrant_id(qdrant_connection: QdrantClient, collection: str) -> str: + while ( + len( + qdrant_connection.retrieve( + collection_name=collection, ids=[id := uuid4().hex] + ) + ) + > 0 + ): + pass + return id + + +def _get_sequence_from_file_path(original_path: Path) -> list[Path]: + match = SEQUENCE_REGEX.match(original_path.name) + if match is None: + return [original_path] + results = {int(match.group(2)): original_path} + directory = original_path.parent + for child_path in directory.iterdir(): + if child_path.is_file(): + child_match = SEQUENCE_REGEX.match(child_path.name) + if child_match is not None: + if child_match.group(1) == match.group(1) and child_match.group( + 3 + ) == match.group(3): + results[int(child_match.group(2))] = child_path + sorted_keys = sorted(results.keys()) + return [results[key] for key in sorted_keys] + + +def _load_image_to_tensor(path: Path) -> Tensor: + image = ImageOps.exif_transpose(Image.open(path)).convert("RGB") + image_array = numpy.array(image).astype(numpy.float32) / 255.0 + return torch.from_numpy(image_array)[None,] + + +SEQUENCE_REGEX = re.compile(r"^(.*?)(\d+)(\D*)$") + + +_Conditioning: TypeAlias = list[list[dict[str, Tensor]]] + + +class _CLIPVisionOutput(Protocol): + image_embeds: Tensor + + +class _CLIPVisionModel(Protocol): + def encode_image(self, Tensor) -> _CLIPVisionOutput: + pass + + +class AddStringToDict: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "key": ("STRING", {"default": "", "multiline": False}), + "value": ("STRING", {"default": "", "multiline": False}), + }, + "optional": {"input_dict": ("DICT",)}, + } + + RETURN_TYPES = ("DICT",) + FUNCTION = "op" + CATEGORY = "math/data_structures" + + def op( + self, key: str, value: str, input_dict: Optional[dict[str, Any]] = None + ) -> tuple[dict]: + output_dict = input_dict.copy() if input_dict else {} + output_dict[key] = value + return (output_dict,) + + +class RetrieveStringFromDict: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "key": ("STRING", {"default": "", "multiline": False}), + "input_dict": ("DICT",), + } + } + + RETURN_TYPES = ("STRING",) + FUNCTION = "op" + CATEGORY = "math/data_structures" + + def op(self, key: str, input_dict: dict[str, Any]) -> tuple[str]: + value = input_dict[key] + if not isinstance(value, str): + raise RuntimeError(f"Value at key {key} is not a string.") + return (value,) + + +class StringAtIndex: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "string_list": ("STRING", {"forceInput": True}), + "index": ("INT", {"default": 0, "min": 0}), + }, + } + + RETURN_TYPES = ("STRING",) + INPUT_IS_LIST = True + FUNCTION = "op" + CATEGORY = "math/data_structures" + + def op(self, string_list: list[str], index: list[int]) -> tuple[str]: + return (string_list[index[0]],) + + +class FloatAtIndex: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "float_list": ("FLOAT", {"forceInput": True}), + "index": ("INT", {"default": 0, "min": 0}), + } + } + + RETURN_TYPES = ("FLOAT",) + INPUT_IS_LIST = True + FUNCTION = "op" + CATEGORY = "math/data_structures" + + def op(self, float_list: list[float], index: list[int]) -> tuple[float]: + return (float_list[index[0]],) + + +class LoadImageFromPath: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image_path": ("STRING", {"default": "", "multiline": False}), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "op" + CATEGORY = "math/image_files" + + def op(self, image_path: str) -> tuple[Tensor]: + return (_load_image_to_tensor(Path(image_path)),) + + +class GetImageSequence: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image_sequence_path": ("STRING", {"default": "", "multiline": False}), + } + } + + RETURN_TYPES = ("STRING", "INT") + RETURN_NAMES = ("image_paths", "sequence_length") + OUTPUT_IS_LIST = (True, False) + FUNCTION = "op" + CATEGORY = "math/image_files" + + def op(self, image_sequence_path: str) -> tuple[list[str], int]: + file_sequence = _get_sequence_from_file_path(Path(image_sequence_path)) + return ( + [str(path) for path in file_sequence], + len(file_sequence), + ) + + +class QdrantConnectionFromFile: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "path": ("STRING", {"default": "", "multiline": False}), + }, + } + + RETURN_TYPES = ("QDRANT_CONNECTION",) + FUNCTION = "op" + CATEGORY = "math/vectordb/qdrant" + + def op(self, path: str) -> tuple[str]: + client = QdrantClient(path=path) + while (id := uuid4().hex) in _qdrant_connections: + pass + _qdrant_connections[id] = client + return (id,) + + +class QdrantCollection: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "connection": ("QDRANT_CONNECTION",), + "collection_name": ("STRING", {"default": "", "multiline": False}), + }, + "optional": { + "vector_size": ("INT", {"default": 1280, "min": 0}), + "distance": (["Cosine", "Euclid", "Dot"],), + }, + } + + RETURN_TYPES = ("QDRANT_COLLECTION",) + FUNCTION = "op" + CATEGORY = "math/vectordb/qdrant" + + def op( + self, + connection: str, + collection_name: str, + vector_size: Optional[int], + distance: Optional[str], + ) -> tuple[dict[str, str]]: + qdrant_connection = _get_qdrant_connection(connection) + try: + connection_info = qdrant_connection.get_collection( + collection_name=collection_name + ) + if connection_info.status != CollectionStatus.GREEN: + raise RuntimeError(f"Collection {collection_name} is not available.") + except ValueError: + real_vector_size = vector_size if vector_size is not None else 512 + real_distance = Distance.COSINE if distance is None else distance + match real_distance: + case "Cosine": + real_distance_value = Distance.COSINE + case "Euclid": + real_distance_value = Distance.EUCLID + case "Dot": + real_distance_value = Distance.DOT + case _: + raise ValueError(f"Unknown distance {distance}.") + qdrant_connection.recreate_collection( + collection_name=collection_name, + vectors_config=VectorParams( + size=real_vector_size, distance=real_distance_value + ), + ) + return ({"connection": connection, "name": collection_name},) + + +class CLIPVisionOutputToQdrantVector: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "clip_vision_output": ("CLIP_VISION_OUTPUT",), + }, + } + + RETURN_TYPES = ("QDRANT_VECTOR",) + FUNCTION = "op" + CATEGORY = "math/vectordb/qdrant" + + def op( + self, + clip_vision_output: _CLIPVisionOutput, + ) -> tuple[list[float]]: + return (clip_vision_output.image_embeds.flatten().tolist(),) + + +class ConditioningToQdrantVector: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "conditioning": ("CONDITIONING",), + }, + } + + RETURN_TYPES = ("QDRANT_VECTOR",) + FUNCTION = "op" + CATEGORY = "math/vectordb/qdrant" + + def op( + self, + conditioning: _Conditioning, + ) -> tuple[list[float]]: + return (conditioning[0][1]["pooled_output"].flatten().tolist(),) + + +class QdrantInsertVector: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "collection": ("QDRANT_COLLECTION",), + "vector": ("QDRANT_VECTOR",), + }, + "optional": {"payload": ("DICT",)}, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("id",) + OUTPUT_NODE = True + FUNCTION = "op" + CATEGORY = "math/vectordb/qdrant" + + def op( + self, + collection: dict[str, str], + vector: list[float], + payload: Optional[dict] = None, + ) -> tuple[str]: + qdrant_connection = _get_qdrant_connection(collection["connection"]) + collection_name = collection["name"] + id = _get_unused_qdrant_id(qdrant_connection, collection_name) + + qdrant_connection.upsert( + collection_name=collection_name, + points=[ + PointStruct( + id=id, + vector=vector, + payload=payload, + ) + ], + ) + return (id,) + + +class QdrantRetrievePayloadById: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "collection": ("QDRANT_COLLECTION",), + "id": ("STRING", {"default": "", "multiline": False}), + } + } + + RETURN_TYPES = ("DICT",) + FUNCTION = "op" + CATEGORY = "math/vectordb/qdrant" + + def op(self, collection: dict[str, str], id: str) -> tuple[dict[str, Any]]: + qdrant_connection = _get_qdrant_connection(collection["connection"]) + collection_name = collection["name"] + result = qdrant_connection.retrieve( + collection_name=collection_name, ids=[id], with_payload=True + ) + if result[0].payload is None: + raise RuntimeError(f"No such id {id} in collection {collection_name}") + return (result[0].payload,) + + +class QdrantSearch: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "collection": ("QDRANT_COLLECTION",), + "query_vector": ("QDRANT_VECTOR",), + "limit": ("INT", {"default": 1, "min": 1}), + }, + } + + RETURN_TYPES = ("STRING", "FLOAT") + RETURN_NAMES = ("ids", "scores") + OUTPUT_IS_LIST = (True, True) + FUNCTION = "op" + CATEGORY = "math/vectordb/qdrant" + + def op( + self, collection: dict[str, str], query_vector: list[float], limit: int + ) -> tuple[list[str], list[float]]: + qdrant_connection = _get_qdrant_connection(collection["connection"]) + collection_name = collection["name"] + results = qdrant_connection.search( + collection_name=collection_name, query_vector=query_vector, limit=limit + ) + return ( + [str(result.id) for result in results], + [result.score for result in results], + ) + + +class QdrantInsertImageSequence: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "collection": ("QDRANT_COLLECTION",), + "clip_vision": ("CLIP_VISION",), + "image_paths": ("STRING", {"forceInput": True}), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("ids",) + OUTPUT_IS_LIST = (True,) + INPUT_IS_LIST = True + OUTPUT_NODE = True + FUNCTION = "op" + CATEGORY = "math/vectordb/qdrant" + + def op( + self, + collection: list[dict[str, str]], + clip_vision: list[_CLIPVisionModel], + image_paths: list[str], + ) -> tuple[list[str]]: + qdrant_connection = _get_qdrant_connection(collection[0]["connection"]) + collection_name = collection[0]["name"] + clip_vision_model = clip_vision[0] + + ids = [] + for path in tqdm(image_paths, desc="Encoding and upserting images to Qdrant DB"): + image = _load_image_to_tensor(Path(path)) + image_embedding = ( + clip_vision_model.encode_image(image).image_embeds.flatten().tolist() + ) + ids.append(id := _get_unused_qdrant_id(qdrant_connection, collection_name)) + qdrant_connection.upsert( + collection_name=collection_name, + points=[ + PointStruct( + id=id, + vector=image_embedding, + payload={"path": path}, + ) + ], + ) + + return (ids,) + + +NODE_CLASS_MAPPINGS = { + "CM_AddStringToDict": AddStringToDict, + "CM_RetrieveStringFromDict": RetrieveStringFromDict, + "CM_StringAtIndex": StringAtIndex, + "CM_FloatAtIndex": FloatAtIndex, + "CM_LoadImageFromPath": LoadImageFromPath, + "CM_GetImageSequence": GetImageSequence, + "CM_QdrantConnectionFromFile": QdrantConnectionFromFile, + "CM_QdrantCollection": QdrantCollection, + "CM_CLIPVisionOutputToQdrantVector": CLIPVisionOutputToQdrantVector, + "CM_ConditioningToQdrantVector": ConditioningToQdrantVector, + "CM_QdrantInsertVector": QdrantInsertVector, + "CM_QdrantRetrievePayloadById": QdrantRetrievePayloadById, + "CM_QdrantSearch": QdrantSearch, + "CM_QdrantInsertImageSequence": QdrantInsertImageSequence, +}