feat: Added notes for working with qdrant vector database

.gitignore

@@ -1 +1,2 @@
 **/__pycache__
+.mypy_cache

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.

__init__.py

@@ -1,12 +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,
@@ -17,6 +17,7 @@ NODE_CLASS_MAPPINGS = {
     **vec_NCM,
     **control_NCM,
     **graphics_NCM,
+    **rag_NCM,
 }
 
 

pyproject.toml

@@ -3,16 +3,20 @@ name = "comfymath"
 version = "0.1.0"
 description = "Math nodes for ComfyUI"
 authors = ["Evan Spearman <evan@spearman.mb.ca>"]
-license = { text = "Apache License 2.0" }
+license = "Apache-2.0"
 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"]

requirements.txt

@@ -1 +1,4 @@
 numpy
+qdrant-client
+tqdm
+click

src/comfymath/bool.py

@@ -1,6 +1,6 @@
 from typing import Any, Callable, Mapping
 
-DEFAULT_BOOL = ("BOOLEAN", {"default": False})
+DEFAULT_BOOL = ("BOOL", {"default": False})
 
 
 BOOL_UNARY_OPERATIONS: Mapping[str, Callable[[bool], bool]] = {
@@ -26,7 +26,7 @@ class BoolUnaryOperation:
             "required": {"op": (list(BOOL_UNARY_OPERATIONS.keys()),), "a": DEFAULT_BOOL}
         }
 
-    RETURN_TYPES = ("BOOLEAN",)
+    RETURN_TYPES = ("BOOL",)
     FUNCTION = "op"
     CATEGORY = "math/bool"
 
@@ -45,7 +45,7 @@ class BoolBinaryOperation:
             }
         }
 
-    RETURN_TYPES = ("BOOLEAN",)
+    RETURN_TYPES = ("BOOL",)
     FUNCTION = "op"
     CATEGORY = "math/bool"
 

src/comfymath/convert.py

@@ -1,13 +1,13 @@
 from typing import Any, Mapping
 
-from .vec import VEC2_ZERO, VEC3_ZERO, VEC4_ZERO
-from .types import Number, Vec2, Vec3, Vec4
+from .vec import Vec2, VEC2_ZERO, Vec3, VEC3_ZERO, Vec4, VEC4_ZERO
+from .number import number
 
 
 class BoolToInt:
     @classmethod
     def INPUT_TYPES(cls) -> Mapping[str, Any]:
-        return {"required": {"a": ("BOOLEAN", {"default": False})}}
+        return {"required": {"a": ("BOOL", {"default": False})}}
 
     RETURN_TYPES = ("INT",)
     FUNCTION = "op"
@@ -22,7 +22,7 @@ class IntToBool:
     def INPUT_TYPES(cls) -> Mapping[str, Any]:
         return {"required": {"a": ("INT", {"default": 0})}}
 
-    RETURN_TYPES = ("BOOLEAN",)
+    RETURN_TYPES = ("BOOL",)
     FUNCTION = "op"
     CATEGORY = "math/conversion"
 
@@ -33,7 +33,7 @@ class IntToBool:
 class FloatToInt:
     @classmethod
     def INPUT_TYPES(cls) -> Mapping[str, Any]:
-        return {"required": {"a": ("FLOAT", {"default": 0.0, "round": False})}}
+        return {"required": {"a": ("FLOAT", {"default": 0.0})}}
 
     RETURN_TYPES = ("INT",)
     FUNCTION = "op"
@@ -65,7 +65,7 @@ class IntToNumber:
     FUNCTION = "op"
     CATEGORY = "math/conversion"
 
-    def op(self, a: int) -> tuple[Number]:
+    def op(self, a: int) -> tuple[number]:
         return (a,)
 
 
@@ -78,20 +78,20 @@ class NumberToInt:
     FUNCTION = "op"
     CATEGORY = "math/conversion"
 
-    def op(self, a: Number) -> tuple[int]:
+    def op(self, a: number) -> tuple[int]:
         return (int(a),)
 
 
 class FloatToNumber:
     @classmethod
     def INPUT_TYPES(cls) -> Mapping[str, Any]:
-        return {"required": {"a": ("FLOAT", {"default": 0.0, "round": False})}}
+        return {"required": {"a": ("FLOAT", {"default": 0.0})}}
 
     RETURN_TYPES = ("NUMBER",)
     FUNCTION = "op"
     CATEGORY = "math/conversion"
 
-    def op(self, a: float) -> tuple[Number]:
+    def op(self, a: float) -> tuple[number]:
         return (a,)
 
 
@@ -104,7 +104,7 @@ class NumberToFloat:
     FUNCTION = "op"
     CATEGORY = "math/conversion"
 
-    def op(self, a: Number) -> tuple[float]:
+    def op(self, a: number) -> tuple[float]:
         return (float(a),)
 
 
@@ -113,8 +113,8 @@ class ComposeVec2:
     def INPUT_TYPES(cls) -> Mapping[str, Any]:
         return {
             "required": {
-                "x": ("FLOAT", {"default": 0.0, "round": False}),
-                "y": ("FLOAT", {"default": 0.0, "round": False}),
+                "x": ("FLOAT", {"default": 0.0}),
+                "y": ("FLOAT", {"default": 0.0}),
             }
         }
 
@@ -131,7 +131,7 @@ class FillVec2:
     def INPUT_TYPES(cls) -> Mapping[str, Any]:
         return {
             "required": {
-                "a": ("FLOAT", {"default": 0.0, "round": False}),
+                "a": ("FLOAT", {"default": 0.0}),
             }
         }
 

src/comfymath/float.py

@@ -2,7 +2,7 @@ import math
 
 from typing import Any, Callable, Mapping
 
-DEFAULT_FLOAT = ("FLOAT", {"default": 0.0, "step": 0.001, "round": False})
+DEFAULT_FLOAT = ("FLOAT", {"default": 0.0})
 
 FLOAT_UNARY_OPERATIONS: Mapping[str, Callable[[float], float]] = {
     "Neg": lambda a: -a,
@@ -105,7 +105,7 @@ class FloatUnaryCondition:
             }
         }
 
-    RETURN_TYPES = ("BOOLEAN",)
+    RETURN_TYPES = ("BOOL",)
     FUNCTION = "op"
     CATEGORY = "math/float"
 
@@ -143,7 +143,7 @@ class FloatBinaryCondition:
             }
         }
 
-    RETURN_TYPES = ("BOOLEAN",)
+    RETURN_TYPES = ("BOOL",)
     FUNCTION = "op"
     CATEGORY = "math/float"
 

src/comfymath/graphics.py

@@ -1,5 +1,6 @@
-from abc import ABC, abstractmethod
-from typing import Any, Mapping, Sequence, Tuple
+from typing import Any, Mapping
+
+from torch import Tensor
 
 
 SDXL_SUPPORTED_RESOLUTIONS = [
@@ -14,60 +15,15 @@ SDXL_SUPPORTED_RESOLUTIONS = [
     (640, 1536, 0.4166666666666667),
 ]
 
-SDXL_EXTENDED_RESOLUTIONS = [
-    (512, 2048, 0.25),
-    (512, 1984, 0.26),
-    (512, 1920, 0.27),
-    (512, 1856, 0.28),
-    (576, 1792, 0.32),
-    (576, 1728, 0.33),
-    (576, 1664, 0.35),
-    (640, 1600, 0.4),
-    (640, 1536, 0.42),
-    (704, 1472, 0.48),
-    (704, 1408, 0.5),
-    (704, 1344, 0.52),
-    (768, 1344, 0.57),
-    (768, 1280, 0.6),
-    (832, 1216, 0.68),
-    (832, 1152, 0.72),
-    (896, 1152, 0.78),
-    (896, 1088, 0.82),
-    (960, 1088, 0.88),
-    (960, 1024, 0.94),
-    (1024, 1024, 1.0),
-    (1024, 960, 1.8),
-    (1088, 960, 1.14),
-    (1088, 896, 1.22),
-    (1152, 896, 1.30),
-    (1152, 832, 1.39),
-    (1216, 832, 1.47),
-    (1280, 768, 1.68),
-    (1344, 768, 1.76),
-    (1408, 704, 2.0),
-    (1472, 704, 2.10),
-    (1536, 640, 2.4),
-    (1600, 640, 2.5),
-    (1664, 576, 2.90),
-    (1728, 576, 3.0),
-    (1792, 576, 3.12),
-    (1856, 512, 3.63),
-    (1920, 512, 3.76),
-    (1984, 512, 3.89),
-    (2048, 512, 4.0),
-]
-
-
-class Resolution(ABC):
-    @classmethod
-    @abstractmethod
-    def resolutions(cls) -> Sequence[Tuple[int, int, float]]: ...
 
+class SDXLResolution:
     @classmethod
     def INPUT_TYPES(cls) -> Mapping[str, Any]:
         return {
             "required": {
-                "resolution": ([f"{res[0]}x{res[1]}" for res in cls.resolutions()],)
+                "resolution": (
+                    [f"{res[0]}x{res[1]}" for res in SDXL_SUPPORTED_RESOLUTIONS],
+                )
             }
         }
 
@@ -81,11 +37,7 @@ class Resolution(ABC):
         return (int(width), int(height))
 
 
-class NearestResolution(ABC):
-    @classmethod
-    @abstractmethod
-    def resolutions(cls) -> Sequence[Tuple[int, int, float]]: ...
-
+class NearestSDXLResolution:
     @classmethod
     def INPUT_TYPES(cls) -> Mapping[str, Any]:
         return {"required": {"image": ("IMAGE",)}}
@@ -95,14 +47,14 @@ class NearestResolution(ABC):
     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}")
         image_ratio = image_width / image_height
         differences = [
             (abs(image_ratio - resolution[2]), resolution)
-            for resolution in self.resolutions()
+            for resolution in SDXL_SUPPORTED_RESOLUTIONS
         ]
         smallest = None
         for difference in differences:
@@ -117,37 +69,11 @@ class NearestResolution(ABC):
         else:
             width = 1024
             height = 1024
-        print(f"Selected resolution: {width}x{height}")
+        print(f"Selected SDXL resolution: {width}x{height}")
         return (width, height)
 
 
-class SDXLResolution(Resolution):
-    @classmethod
-    def resolutions(cls):
-        return SDXL_SUPPORTED_RESOLUTIONS
-
-
-class SDXLExtendedResolution(Resolution):
-    @classmethod
-    def resolutions(cls):
-        return SDXL_EXTENDED_RESOLUTIONS
-
-
-class NearestSDXLResolution(NearestResolution):
-    @classmethod
-    def resolutions(cls):
-        return SDXL_SUPPORTED_RESOLUTIONS
-
-
-class NearestSDXLExtendedResolution(NearestResolution):
-    @classmethod
-    def resolutions(cls):
-        return SDXL_EXTENDED_RESOLUTIONS
-
-
 NODE_CLASS_MAPPINGS = {
     "CM_SDXLResolution": SDXLResolution,
     "CM_NearestSDXLResolution": NearestSDXLResolution,
-    "CM_SDXLExtendedResolution": SDXLExtendedResolution,
-    "CM_NearestSDXLExtendedResolution": NearestSDXLExtendedResolution,
 }

src/comfymath/number.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import Any, Callable, Mapping
+from typing import Any, Callable, Mapping, TypeAlias
 
 from .float import (
     FLOAT_UNARY_OPERATIONS,
@@ -7,10 +7,11 @@ from .float import (
     FLOAT_BINARY_OPERATIONS,
     FLOAT_BINARY_CONDITIONS,
 )
-from .types import Number
 
 DEFAULT_NUMBER = ("NUMBER", {"default": 0.0})
 
+number: TypeAlias = int | float
+
 
 class NumberUnaryOperation:
     @classmethod
@@ -26,7 +27,7 @@ class NumberUnaryOperation:
     FUNCTION = "op"
     CATEGORY = "math/number"
 
-    def op(self, op: str, a: Number) -> tuple[float]:
+    def op(self, op: str, a: number) -> tuple[float]:
         return (FLOAT_UNARY_OPERATIONS[op](float(a)),)
 
 
@@ -42,9 +43,9 @@ class NumberUnaryCondition:
 
     RETURN_TYPES = ("BOOL",)
     FUNCTION = "op"
-    CATEGORY = "math/Number"
+    CATEGORY = "math/number"
 
-    def op(self, op: str, a: Number) -> tuple[bool]:
+    def op(self, op: str, a: number) -> tuple[bool]:
         return (FLOAT_UNARY_CONDITIONS[op](float(a)),)
 
 
@@ -63,7 +64,7 @@ class NumberBinaryOperation:
     FUNCTION = "op"
     CATEGORY = "math/number"
 
-    def op(self, op: str, a: Number, b: Number) -> tuple[float]:
+    def op(self, op: str, a: number, b: number) -> tuple[float]:
         return (FLOAT_BINARY_OPERATIONS[op](float(a), float(b)),)
 
 
@@ -82,7 +83,7 @@ class NumberBinaryCondition:
     FUNCTION = "op"
     CATEGORY = "math/float"
 
-    def op(self, op: str, a: Number, b: Number) -> tuple[bool]:
+    def op(self, op: str, a: number, b: number) -> tuple[bool]:
         return (FLOAT_BINARY_CONDITIONS[op](float(a), float(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,
+}

src/comfymath/types.py

@@ -1,16 +0,0 @@
-import sys
-
-if sys.version_info[1] < 10:
-    from typing import Tuple, Union
-
-    Number = Union[int, float]
-    Vec2 = Tuple[float, float]
-    Vec3 = Tuple[float, float, float]
-    Vec4 = Tuple[float, float, float, float]
-else:
-    from typing import TypeAlias
-
-    Number: TypeAlias = int | float
-    Vec2: TypeAlias = tuple[float, float]
-    Vec3: TypeAlias = tuple[float, float, float]
-    Vec4: TypeAlias = tuple[float, float, float, float]

src/comfymath/vec.py

@@ -1,15 +1,16 @@
 import numpy
 
-from typing import Any, Callable, Mapping
-
-from .types import Vec2, Vec3, Vec4
+from typing import Any, Callable, Mapping, TypeAlias
 
+Vec2: TypeAlias = tuple[float, float]
 VEC2_ZERO = (0.0, 0.0)
 DEFAULT_VEC2 = ("VEC2", {"default": VEC2_ZERO})
 
+Vec3: TypeAlias = tuple[float, float, float]
 VEC3_ZERO = (0.0, 0.0, 0.0)
 DEFAULT_VEC3 = ("VEC3", {"default": VEC3_ZERO})
 
+Vec4: TypeAlias = tuple[float, float, float, float]
 VEC4_ZERO = (0.0, 0.0, 0.0, 0.0)
 DEFAULT_VEC4 = ("VEC4", {"default": VEC4_ZERO})