""" Text embedders for the variability engine. One small interface, several backends. embedder.name stable id, used as the cache key ("qwen3vl-2b", "openclip-ViT-H-14", ...) embedder.encode(texts) -> float32 array [len(texts), dim], rows L2-normalised Which space you embed in decides what "similar" means: qwen3vl Qwen3-VL-Embedding. Text embedded in a space trained jointly with images, so neighbours are close in meaning AND in what they look like. Best default. openclip CLIP text tower. Neighbours LOOK alike ("rusty" ~ "weathered copper"). comfy_clip The CLIP already loaded in your ComfyUI graph (pooled output). No extra model. sentence Any sentence-transformers model. Neighbours MEAN alike. lexical Character n-gram hashing. No model, no meaning - spelling similarity only. For tests and as a last resort so the node still runs with nothing installed. Models load lazily on first encode() and stay loaded for the session. (c) Eric Hiss. Dual license, see LICENSE. """ from __future__ import annotations import hashlib import importlib.util import os import sys from typing import Any, List, Optional import numpy as np BATCH_SIZE = 64 def _normalise(matrix: np.ndarray) -> np.ndarray: matrix = np.asarray(matrix, dtype=np.float32) norms = np.linalg.norm(matrix, axis=1, keepdims=True) norms[norms == 0] = 1.0 return matrix / norms class Embedder: name = "base" def encode(self, texts: List[str]) -> np.ndarray: raise NotImplementedError def describe(self) -> str: return self.name class LexicalEmbedder(Embedder): """Hashed character trigrams. Deterministic, dependency-free, NOT semantic.""" name = "lexical-256" def __init__(self, dim: int = 256): self.dim = dim def encode(self, texts: List[str]) -> np.ndarray: matrix = np.zeros((len(texts), self.dim), dtype=np.float32) for row, text in enumerate(texts): padded = f" {text.lower().strip()} " for index in range(len(padded) - 2): digest = hashlib.md5(padded[index:index + 3].encode("utf-8")).digest() matrix[row, int.from_bytes(digest[:4], "little") % self.dim] += 1.0 return _normalise(matrix) def describe(self) -> str: return "lexical (spelling similarity only - no model found)" class SentenceTransformerEmbedder(Embedder): def __init__(self, model_name_or_path: str, device: Optional[str] = None): self.model_name_or_path = model_name_or_path self.device = device self.name = "sentence-" + os.path.basename(os.path.normpath(model_name_or_path)).lower() self._model = None def encode(self, texts: List[str]) -> np.ndarray: if self._model is None: from sentence_transformers import SentenceTransformer self._model = SentenceTransformer(self.model_name_or_path, device=self.device) vectors = self._model.encode(texts, batch_size=BATCH_SIZE, convert_to_numpy=True, normalize_embeddings=True, show_progress_bar=False) return _normalise(vectors) class OpenClipEmbedder(Embedder): """ CLIP text tower through open_clip. `source` is a local model folder that contains open_clip_config.json (loaded with open_clip's "local-dir:" scheme), or an "architecture|pretrained-tag" pair such as "ViT-H-14|laion2b_s32b_b79k". """ def __init__(self, source: str, device: Optional[str] = None): self.source = source self.device = device label = os.path.basename(os.path.normpath(source)) if os.path.isdir(source) else source.replace("|", "-") self.name = "openclip-" + label.lower() self._model = None self._tokenizer = None def _load(self) -> None: import open_clip import torch if os.path.isdir(self.source): model_name = "local-dir:" + self.source model, _, _ = open_clip.create_model_and_transforms(model_name) self._tokenizer = open_clip.get_tokenizer(model_name) else: architecture, _, pretrained = self.source.partition("|") model, _, _ = open_clip.create_model_and_transforms(architecture, pretrained=pretrained or None) self._tokenizer = open_clip.get_tokenizer(architecture) self.device = self.device or ("cuda" if torch.cuda.is_available() else "cpu") self._model = model.to(self.device).eval() def encode(self, texts: List[str]) -> np.ndarray: import torch if self._model is None: self._load() chunks = [] with torch.no_grad(): for start in range(0, len(texts), BATCH_SIZE): tokens = self._tokenizer(texts[start:start + BATCH_SIZE]).to(self.device) chunks.append(self._model.encode_text(tokens).float().cpu().numpy()) return _normalise(np.concatenate(chunks, axis=0)) class Qwen3VLEmbedder(Embedder): """ Qwen3-VL-Embedding through the loader class in Eric's eric_semantic_search package (core/qwen3_vl_embedding.py, class Qwen3VLEmbedder). That file has no package-relative imports, so it is loaded by path; nothing in eric_semantic_search is modified or started. """ INSTRUCTION = "Represent this phrase from an image generation prompt for finding visually and semantically similar phrases." def __init__(self, model_path: str, loader_file: str, device: Optional[str] = None): self.model_path = model_path self.loader_file = loader_file self.device = device # e.g. "cuda:1" or "cpu"; None = wherever the loader puts it (cuda:0) self.name = "qwen3vl-" + os.path.basename(os.path.normpath(model_path)).lower() self._embedder: Any = None def _load(self) -> None: module_name = "vp_qwen3_vl_embedding" module = sys.modules.get(module_name) if module is None: spec = importlib.util.spec_from_file_location(module_name, self.loader_file) module = importlib.util.module_from_spec(spec) sys.modules[module_name] = module spec.loader.exec_module(module) self._embedder = module.Qwen3VLEmbedder(self.model_path) if self.device: # The loader always picks "cuda" (= cuda:0). It keeps the device in .device and sends its # inputs there, so moving the model and updating that attribute is all a relocation takes. import torch target = torch.device(self.device) self._embedder.model.to(target) self._embedder.device = target def describe(self) -> str: where = self.device or "cuda:0 (loader default)" return f"{self.name} on {where}" def encode(self, texts: List[str]) -> np.ndarray: if self._embedder is None: self._load() chunks = [] for start in range(0, len(texts), BATCH_SIZE): batch = [{"text": text, "instruction": self.INSTRUCTION} for text in texts[start:start + BATCH_SIZE]] vectors = self._embedder.process(batch, normalize=True) chunks.append(vectors.float().cpu().numpy()) return _normalise(np.concatenate(chunks, axis=0)) class ComfyClipEmbedder(Embedder): """Pooled text output of a CLIP object from the ComfyUI graph. Needs an encoder that has a pooled output.""" def __init__(self, clip: Any, label: str = "graph"): self.clip = clip self.name = "comfyclip-" + label def encode(self, texts: List[str]) -> np.ndarray: rows = [] for text in texts: tokens = self.clip.tokenize(text) _conditioning, pooled = self.clip.encode_from_tokens(tokens, return_pooled=True) if pooled is None: raise RuntimeError("this text encoder has no pooled output; connect a CLIP-L/G encoder or pick another embedder") rows.append(pooled.float().reshape(-1).cpu().numpy()) return _normalise(np.stack(rows, axis=0)) def clip_label(clip: Any) -> str: """A stable-enough id for a graph CLIP so its cache is not mixed with another encoder's.""" try: inner = getattr(clip, "cond_stage_model", None) parts = [type(inner).__name__] + sorted(name for name in dir(inner) if name.startswith("clip_")) return hashlib.md5("|".join(parts).encode("utf-8")).hexdigest()[:8] except Exception: return "graph"