feat: add KikoEmbeddingAutocomplete with settings registry system
- Implement centralized settings registry for all KikoTools
- Create KikoEmbeddingAutocomplete node with backend API
- Add frontend JavaScript autocomplete widget with ComfyUI integration
- Support for embeddings and LoRAs with smart filtering
- Configurable settings in ComfyUI UI with 🫶 branding
- Include keyboard navigation and real-time suggestions
This commit is contained in:
Vito Sansevero
+32
@@ -15,6 +15,38 @@ except ImportError:
|
||||
# Tell ComfyUI where to find our JavaScript extensions
|
||||
WEB_DIRECTORY = "./web"
|
||||
|
||||
def setup_autocomplete_api():
|
||||
"""Setup API routes for embedding autocomplete."""
|
||||
try:
|
||||
from aiohttp import web
|
||||
from server import PromptServer
|
||||
from kikotools.tools.embedding_autocomplete.node import KikoEmbeddingAutocompleteAPI
|
||||
|
||||
@PromptServer.instance.routes.get("/kikotools/autocomplete/suggestions")
|
||||
async def get_suggestions(request):
|
||||
"""API endpoint for getting autocomplete suggestions."""
|
||||
prefix = request.query.get("prefix", "")
|
||||
max_results = int(request.query.get("max", 20))
|
||||
include_embeddings = request.query.get("embeddings", "true").lower() == "true"
|
||||
include_loras = request.query.get("loras", "true").lower() == "true"
|
||||
case_sensitive = request.query.get("case_sensitive", "false").lower() == "true"
|
||||
|
||||
suggestions = KikoEmbeddingAutocompleteAPI.get_suggestions(
|
||||
prefix=prefix,
|
||||
max_results=max_results,
|
||||
include_embeddings=include_embeddings,
|
||||
include_loras=include_loras,
|
||||
case_sensitive=case_sensitive
|
||||
)
|
||||
|
||||
return web.json_response(suggestions)
|
||||
|
||||
except ImportError:
|
||||
pass # Server not available when not in ComfyUI
|
||||
|
||||
# Setup API if available
|
||||
setup_autocomplete_api()
|
||||
|
||||
|
||||
def get_version():
|
||||
"""Parse version from pyproject.toml"""
|
||||
|
||||
@@ -23,6 +23,7 @@ from .tools.xyz_helpers import (
|
||||
PlotParametersNode,
|
||||
LoRAFolderBatchNode,
|
||||
)
|
||||
from .tools.embedding_autocomplete import KikoEmbeddingAutocomplete
|
||||
|
||||
# ComfyUI node registration mappings
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
@@ -45,6 +46,7 @@ NODE_CLASS_MAPPINGS = {
|
||||
"FluxSamplerParams": FluxSamplerParamsNode,
|
||||
"PlotParameters+": PlotParametersNode,
|
||||
"LoRAFolderBatch": LoRAFolderBatchNode,
|
||||
"KikoEmbeddingAutocomplete": KikoEmbeddingAutocomplete,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
@@ -67,6 +69,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"FluxSamplerParams": "Flux Sampler Parameters",
|
||||
"PlotParameters+": "Plot Parameters",
|
||||
"LoRAFolderBatch": "LoRA Folder Batch",
|
||||
"KikoEmbeddingAutocomplete": "🫶 Embedding Autocomplete",
|
||||
}
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
|
||||
@@ -0,0 +1,96 @@
|
||||
"""Tool registry for KikoTools.
|
||||
|
||||
This module provides the central registration system for all KikoTools nodes.
|
||||
"""
|
||||
|
||||
import importlib
|
||||
import os
|
||||
from typing import Dict, List, Any, Optional
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
class ToolRegistry:
|
||||
"""Central registry for all KikoTools."""
|
||||
|
||||
def __init__(self):
|
||||
self.tools: Dict[str, Any] = {}
|
||||
self.node_classes: Dict[str, Any] = {}
|
||||
|
||||
def register_tool(self, tool_name: str, node_class: Any) -> None:
|
||||
"""Register a tool and its node class.
|
||||
|
||||
Args:
|
||||
tool_name: Name of the tool
|
||||
node_class: The ComfyUI node class
|
||||
"""
|
||||
self.tools[tool_name] = node_class
|
||||
|
||||
# Also register by class name for ComfyUI
|
||||
class_name = node_class.__name__
|
||||
self.node_classes[class_name] = node_class
|
||||
|
||||
def discover_tools(self) -> None:
|
||||
"""Automatically discover and load all tools in the tools directory."""
|
||||
tools_dir = Path(__file__).parent.parent / "tools"
|
||||
|
||||
if not tools_dir.exists():
|
||||
return
|
||||
|
||||
for tool_dir in tools_dir.iterdir():
|
||||
if tool_dir.is_dir() and not tool_dir.name.startswith("_"):
|
||||
self._load_tool(tool_dir.name)
|
||||
|
||||
def _load_tool(self, tool_name: str) -> None:
|
||||
"""Load a single tool module.
|
||||
|
||||
Args:
|
||||
tool_name: Name of the tool directory
|
||||
"""
|
||||
try:
|
||||
# Try to import the tool's node module
|
||||
module = importlib.import_module(f"kikotools.tools.{tool_name}.node")
|
||||
|
||||
# Look for node classes (classes with ComfyUI node attributes)
|
||||
for attr_name in dir(module):
|
||||
attr = getattr(module, attr_name)
|
||||
if (
|
||||
isinstance(attr, type)
|
||||
and hasattr(attr, "INPUT_TYPES")
|
||||
and hasattr(attr, "FUNCTION")
|
||||
):
|
||||
self.register_tool(tool_name, attr)
|
||||
|
||||
# If the tool has settings, register them
|
||||
if hasattr(attr, "SETTINGS"):
|
||||
from .settings import settings_registry
|
||||
|
||||
settings_registry.register_tool_settings(
|
||||
tool_name,
|
||||
getattr(
|
||||
attr,
|
||||
"DISPLAY_NAME",
|
||||
tool_name.replace("_", " ").title(),
|
||||
),
|
||||
attr.SETTINGS,
|
||||
)
|
||||
|
||||
except ImportError as e:
|
||||
# Tool might not have a node.py file yet
|
||||
pass
|
||||
|
||||
def get_node_class_mappings(self) -> Dict[str, Any]:
|
||||
"""Get node class mappings for ComfyUI registration."""
|
||||
return self.node_classes.copy()
|
||||
|
||||
def get_node_display_name_mappings(self) -> Dict[str, str]:
|
||||
"""Get display name mappings for ComfyUI."""
|
||||
mappings = {}
|
||||
for class_name, node_class in self.node_classes.items():
|
||||
if hasattr(node_class, "DISPLAY_NAME"):
|
||||
mappings[class_name] = node_class.DISPLAY_NAME
|
||||
else:
|
||||
# Generate a display name from class name
|
||||
mappings[class_name] = class_name.replace("Kiko", "").replace(
|
||||
"Node", ""
|
||||
)
|
||||
return mappings
|
||||
@@ -0,0 +1,201 @@
|
||||
"""Settings registry for KikoTools.
|
||||
|
||||
This module provides a centralized settings management system for all KikoTools.
|
||||
Tools can register their settings, which are then exposed in ComfyUI's settings UI.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
from typing import Dict, Any, List, Optional, Union
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
|
||||
@dataclass
|
||||
class SettingDefinition:
|
||||
"""Definition of a single setting."""
|
||||
|
||||
id: str
|
||||
name: str
|
||||
type: str # "boolean", "combo", "number", "string", "custom"
|
||||
default: Any
|
||||
description: Optional[str] = None
|
||||
options: Optional[Union[List[Any], Dict[str, Any]]] = None
|
||||
min_value: Optional[float] = None
|
||||
max_value: Optional[float] = None
|
||||
step: Optional[float] = None
|
||||
on_change: Optional[str] = None # JavaScript callback as string
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolSettings:
|
||||
"""Settings collection for a single tool."""
|
||||
|
||||
tool_name: str
|
||||
display_name: str
|
||||
settings: List[SettingDefinition] = field(default_factory=list)
|
||||
|
||||
|
||||
class SettingsRegistry:
|
||||
"""Central registry for all KikoTools settings."""
|
||||
|
||||
def __init__(self):
|
||||
self.tools: Dict[str, ToolSettings] = {}
|
||||
self.settings_by_id: Dict[str, SettingDefinition] = {}
|
||||
|
||||
def register_tool_settings(
|
||||
self, tool_name: str, display_name: str, settings: Dict[str, Dict[str, Any]]
|
||||
) -> None:
|
||||
"""Register settings for a tool.
|
||||
|
||||
Args:
|
||||
tool_name: Internal tool identifier (e.g., "embedding_autocomplete")
|
||||
display_name: Display name for the tool (e.g., "Embedding Autocomplete")
|
||||
settings: Dictionary of setting configurations
|
||||
{
|
||||
"enabled": {
|
||||
"type": "boolean",
|
||||
"default": True,
|
||||
"description": "Enable embedding autocomplete"
|
||||
},
|
||||
"max_suggestions": {
|
||||
"type": "combo",
|
||||
"default": 20,
|
||||
"options": [10, 20, 50],
|
||||
"description": "Maximum number of suggestions"
|
||||
}
|
||||
}
|
||||
"""
|
||||
tool_settings = ToolSettings(tool_name, display_name)
|
||||
|
||||
for setting_key, config in settings.items():
|
||||
# Generate fully qualified setting ID
|
||||
setting_id = f"kikotools.{tool_name}.{setting_key}"
|
||||
|
||||
# Create display name with branding
|
||||
setting_name = f"🫶 {display_name}: {setting_key.replace('_', ' ').title()}"
|
||||
|
||||
setting_def = SettingDefinition(
|
||||
id=setting_id,
|
||||
name=setting_name,
|
||||
type=config.get("type", "string"),
|
||||
default=config.get("default"),
|
||||
description=config.get("description"),
|
||||
options=config.get("options"),
|
||||
min_value=config.get("min"),
|
||||
max_value=config.get("max"),
|
||||
step=config.get("step"),
|
||||
on_change=config.get("on_change"),
|
||||
)
|
||||
|
||||
tool_settings.settings.append(setting_def)
|
||||
self.settings_by_id[setting_id] = setting_def
|
||||
|
||||
self.tools[tool_name] = tool_settings
|
||||
|
||||
def get_setting(self, setting_id: str) -> Optional[SettingDefinition]:
|
||||
"""Get a setting definition by ID."""
|
||||
return self.settings_by_id.get(setting_id)
|
||||
|
||||
def get_tool_settings(self, tool_name: str) -> Optional[ToolSettings]:
|
||||
"""Get all settings for a tool."""
|
||||
return self.tools.get(tool_name)
|
||||
|
||||
def generate_frontend_registration(self) -> str:
|
||||
"""Generate JavaScript code for frontend settings registration."""
|
||||
js_lines = [
|
||||
"// Auto-generated KikoTools settings registration",
|
||||
"// This file is automatically generated by the settings registry",
|
||||
"",
|
||||
"import { app } from '../../scripts/app.js';",
|
||||
"",
|
||||
"app.registerExtension({",
|
||||
" name: 'kikotools.settings',",
|
||||
" async init() {",
|
||||
" // Register all KikoTools settings",
|
||||
]
|
||||
|
||||
for tool_name, tool_settings in self.tools.items():
|
||||
js_lines.append(f" // {tool_settings.display_name} settings")
|
||||
|
||||
for setting in tool_settings.settings:
|
||||
js_lines.append(f" app.ui.settings.addSetting({{")
|
||||
js_lines.append(f' id: "{setting.id}",')
|
||||
js_lines.append(f' name: "{setting.name}",')
|
||||
js_lines.append(
|
||||
f" defaultValue: {self._js_value(setting.default)},"
|
||||
)
|
||||
js_lines.append(f' type: "{setting.type}",')
|
||||
|
||||
if setting.description:
|
||||
js_lines.append(f' tooltip: "{setting.description}",')
|
||||
|
||||
if setting.type == "combo" and setting.options:
|
||||
js_lines.append(f" options: (value) => {{")
|
||||
js_lines.append(
|
||||
f" const options = {json.dumps(setting.options)};"
|
||||
)
|
||||
js_lines.append(f" return options.map(opt => ({{")
|
||||
js_lines.append(f" value: opt,")
|
||||
js_lines.append(f" text: String(opt),")
|
||||
js_lines.append(f" selected: opt === value")
|
||||
js_lines.append(f" }}));")
|
||||
js_lines.append(f" }},")
|
||||
|
||||
if setting.type == "number":
|
||||
if setting.min_value is not None:
|
||||
js_lines.append(f" min: {setting.min_value},")
|
||||
if setting.max_value is not None:
|
||||
js_lines.append(f" max: {setting.max_value},")
|
||||
if setting.step is not None:
|
||||
js_lines.append(f" step: {setting.step},")
|
||||
|
||||
if setting.on_change:
|
||||
js_lines.append(f" onChange(value) {{")
|
||||
js_lines.append(f" {setting.on_change}")
|
||||
js_lines.append(f" }}")
|
||||
|
||||
js_lines.append(f" }});")
|
||||
js_lines.append("")
|
||||
|
||||
js_lines.extend([" }", "});", ""])
|
||||
|
||||
return "\n".join(js_lines)
|
||||
|
||||
def _js_value(self, value: Any) -> str:
|
||||
"""Convert Python value to JavaScript literal."""
|
||||
if isinstance(value, bool):
|
||||
return "true" if value else "false"
|
||||
elif isinstance(value, str):
|
||||
return f'"{value}"'
|
||||
elif value is None:
|
||||
return "null"
|
||||
else:
|
||||
return str(value)
|
||||
|
||||
def save_frontend_settings(
|
||||
self, output_path: str = "web/js/kikoSettings.js"
|
||||
) -> None:
|
||||
"""Save the generated frontend settings to a file."""
|
||||
js_content = self.generate_frontend_registration()
|
||||
|
||||
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
||||
with open(output_path, "w") as f:
|
||||
f.write(js_content)
|
||||
|
||||
def get_all_settings(self) -> Dict[str, Any]:
|
||||
"""Get all registered settings as a dictionary."""
|
||||
result = {}
|
||||
for tool_name, tool_settings in self.tools.items():
|
||||
result[tool_name] = {
|
||||
"display_name": tool_settings.display_name,
|
||||
"settings": {
|
||||
setting.id.split(".")[-1]: {
|
||||
"type": setting.type,
|
||||
"default": setting.default,
|
||||
"description": setting.description,
|
||||
"options": setting.options,
|
||||
}
|
||||
for setting in tool_settings.settings
|
||||
},
|
||||
}
|
||||
return result
|
||||
@@ -0,0 +1,5 @@
|
||||
"""Embedding Autocomplete tool for KikoTools."""
|
||||
|
||||
from .node import KikoEmbeddingAutocomplete
|
||||
|
||||
__all__ = ["KikoEmbeddingAutocomplete"]
|
||||
@@ -0,0 +1,254 @@
|
||||
"""KikoEmbeddingAutocomplete node for ComfyUI.
|
||||
|
||||
Provides autocomplete functionality for embeddings and LoRAs in text inputs.
|
||||
"""
|
||||
|
||||
import os
|
||||
import json
|
||||
from typing import Dict, List, Any, Optional
|
||||
import folder_paths
|
||||
|
||||
|
||||
class KikoEmbeddingAutocomplete:
|
||||
"""Node that provides embedding autocomplete functionality."""
|
||||
|
||||
DISPLAY_NAME = "Embedding Autocomplete"
|
||||
CATEGORY = "ComfyAssets"
|
||||
|
||||
# Settings definition for the settings registry
|
||||
SETTINGS = {
|
||||
"enabled": {
|
||||
"type": "boolean",
|
||||
"default": True,
|
||||
"description": "Enable embedding autocomplete",
|
||||
},
|
||||
"trigger_chars": {
|
||||
"type": "combo",
|
||||
"default": 2,
|
||||
"options": [1, 2, 3, 4],
|
||||
"description": "Number of characters before showing suggestions",
|
||||
},
|
||||
"max_suggestions": {
|
||||
"type": "combo",
|
||||
"default": 20,
|
||||
"options": [10, 20, 30, 50],
|
||||
"description": "Maximum number of suggestions to show",
|
||||
},
|
||||
"show_embeddings": {
|
||||
"type": "boolean",
|
||||
"default": True,
|
||||
"description": "Show embeddings in suggestions",
|
||||
},
|
||||
"show_loras": {
|
||||
"type": "boolean",
|
||||
"default": True,
|
||||
"description": "Show LoRAs in suggestions",
|
||||
},
|
||||
"case_sensitive": {
|
||||
"type": "boolean",
|
||||
"default": False,
|
||||
"description": "Case sensitive matching",
|
||||
},
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
"""Define input types for the node."""
|
||||
return {
|
||||
"required": {},
|
||||
"optional": {
|
||||
"refresh": ("BOOLEAN", {"default": False}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("DICT",)
|
||||
RETURN_NAMES = ("autocomplete_data",)
|
||||
FUNCTION = "get_autocomplete_data"
|
||||
|
||||
def __init__(self):
|
||||
self.embeddings_cache = None
|
||||
self.loras_cache = None
|
||||
|
||||
def get_autocomplete_data(self, refresh=False):
|
||||
"""Get autocomplete data for embeddings and LoRAs.
|
||||
|
||||
This node doesn't process data in the traditional sense - it provides
|
||||
autocomplete data to the frontend JavaScript component.
|
||||
"""
|
||||
if refresh or self.embeddings_cache is None:
|
||||
self.refresh_cache()
|
||||
|
||||
return (
|
||||
{
|
||||
"embeddings": self.embeddings_cache,
|
||||
"loras": self.loras_cache,
|
||||
"timestamp": os.path.getmtime(folder_paths.base_path),
|
||||
},
|
||||
)
|
||||
|
||||
def refresh_cache(self):
|
||||
"""Refresh the cache of embeddings and LoRAs."""
|
||||
self.embeddings_cache = self.get_embeddings()
|
||||
self.loras_cache = self.get_loras()
|
||||
|
||||
def get_embeddings(self) -> List[Dict[str, Any]]:
|
||||
"""Get list of available embeddings."""
|
||||
embeddings = []
|
||||
|
||||
# Get embedding files from ComfyUI's folder system
|
||||
try:
|
||||
embedding_files = folder_paths.get_filename_list("embeddings")
|
||||
for file in embedding_files:
|
||||
name = os.path.splitext(file)[0]
|
||||
embeddings.append(
|
||||
{
|
||||
"name": name,
|
||||
"file": file,
|
||||
"type": "embedding",
|
||||
"display": f"embedding:{name}",
|
||||
"value": f"embedding:{name}",
|
||||
}
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"Error loading embeddings: {e}")
|
||||
|
||||
return embeddings
|
||||
|
||||
def get_loras(self) -> List[Dict[str, Any]]:
|
||||
"""Get list of available LoRAs."""
|
||||
loras = []
|
||||
|
||||
# Get LoRA files from ComfyUI's folder system
|
||||
try:
|
||||
lora_files = folder_paths.get_filename_list("loras")
|
||||
for file in lora_files:
|
||||
name = os.path.splitext(file)[0]
|
||||
loras.append(
|
||||
{
|
||||
"name": name,
|
||||
"file": file,
|
||||
"type": "lora",
|
||||
"display": f"<lora:{name}:1.0>",
|
||||
"value": f"<lora:{name}:1.0>",
|
||||
}
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"Error loading LoRAs: {e}")
|
||||
|
||||
return loras
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, **kwargs):
|
||||
"""Check if the node needs to be re-executed."""
|
||||
# Always re-execute if refresh is True
|
||||
if kwargs.get("refresh", False):
|
||||
return float("NaN")
|
||||
|
||||
# Check if embeddings/loras folders have changed
|
||||
try:
|
||||
embeddings_path = folder_paths.get_folder_paths("embeddings")[0]
|
||||
loras_path = folder_paths.get_folder_paths("loras")[0]
|
||||
|
||||
# Return combined modification time
|
||||
return os.path.getmtime(embeddings_path) + os.path.getmtime(loras_path)
|
||||
except:
|
||||
return 0
|
||||
|
||||
|
||||
class KikoEmbeddingAutocompleteAPI:
|
||||
"""API endpoints for embedding autocomplete."""
|
||||
|
||||
@staticmethod
|
||||
def get_suggestions(
|
||||
prefix: str,
|
||||
max_results: int = 20,
|
||||
include_embeddings: bool = True,
|
||||
include_loras: bool = True,
|
||||
case_sensitive: bool = False,
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Get autocomplete suggestions for a given prefix.
|
||||
|
||||
Args:
|
||||
prefix: The text prefix to match
|
||||
max_results: Maximum number of results to return
|
||||
include_embeddings: Include embeddings in results
|
||||
include_loras: Include LoRAs in results
|
||||
case_sensitive: Use case-sensitive matching
|
||||
|
||||
Returns:
|
||||
List of suggestion dictionaries
|
||||
"""
|
||||
suggestions = []
|
||||
|
||||
# Normalize prefix for matching
|
||||
match_prefix = prefix if case_sensitive else prefix.lower()
|
||||
|
||||
# Get embeddings
|
||||
if include_embeddings:
|
||||
try:
|
||||
embedding_files = folder_paths.get_filename_list("embeddings")
|
||||
for file in embedding_files:
|
||||
name = os.path.splitext(file)[0]
|
||||
match_name = name if case_sensitive else name.lower()
|
||||
|
||||
# Check for match
|
||||
if match_name.startswith(match_prefix):
|
||||
suggestions.append(
|
||||
{
|
||||
"name": name,
|
||||
"type": "embedding",
|
||||
"display": f"embedding:{name}",
|
||||
"value": f"embedding:{name}",
|
||||
"priority": 1 if match_name == match_prefix else 0,
|
||||
}
|
||||
)
|
||||
elif match_prefix in match_name:
|
||||
suggestions.append(
|
||||
{
|
||||
"name": name,
|
||||
"type": "embedding",
|
||||
"display": f"embedding:{name}",
|
||||
"value": f"embedding:{name}",
|
||||
"priority": -1,
|
||||
}
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"Error loading embeddings: {e}")
|
||||
|
||||
# Get LoRAs
|
||||
if include_loras:
|
||||
try:
|
||||
lora_files = folder_paths.get_filename_list("loras")
|
||||
for file in lora_files:
|
||||
name = os.path.splitext(file)[0]
|
||||
match_name = name if case_sensitive else name.lower()
|
||||
|
||||
# Check for match
|
||||
if match_name.startswith(match_prefix):
|
||||
suggestions.append(
|
||||
{
|
||||
"name": name,
|
||||
"type": "lora",
|
||||
"display": f"<lora:{name}:1.0>",
|
||||
"value": f"<lora:{name}:1.0>",
|
||||
"priority": 1 if match_name == match_prefix else 0,
|
||||
}
|
||||
)
|
||||
elif match_prefix in match_name:
|
||||
suggestions.append(
|
||||
{
|
||||
"name": name,
|
||||
"type": "lora",
|
||||
"display": f"<lora:{name}:1.0>",
|
||||
"value": f"<lora:{name}:1.0>",
|
||||
"priority": -1,
|
||||
}
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"Error loading LoRAs: {e}")
|
||||
|
||||
# Sort by priority and name
|
||||
suggestions.sort(key=lambda x: (-x["priority"], x["name"]))
|
||||
|
||||
# Limit results
|
||||
return suggestions[:max_results]
|
||||
@@ -0,0 +1,87 @@
|
||||
"""Basic tests for KikoEmbeddingAutocomplete."""
|
||||
|
||||
import sys
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
# Mock ComfyUI's folder_paths module
|
||||
sys.modules['folder_paths'] = MagicMock()
|
||||
sys.modules['folder_paths'].get_filename_list = MagicMock(return_value=[])
|
||||
sys.modules['folder_paths'].get_folder_paths = MagicMock(return_value=["/mock/path"])
|
||||
sys.modules['folder_paths'].base_path = "/mock/base"
|
||||
|
||||
|
||||
def test_import():
|
||||
"""Test that the module can be imported."""
|
||||
from kikotools.tools.embedding_autocomplete import KikoEmbeddingAutocomplete
|
||||
assert KikoEmbeddingAutocomplete is not None
|
||||
assert KikoEmbeddingAutocomplete.DISPLAY_NAME == "Embedding Autocomplete"
|
||||
assert KikoEmbeddingAutocomplete.CATEGORY == "ComfyAssets"
|
||||
|
||||
|
||||
def test_settings_defined():
|
||||
"""Test that settings are properly defined."""
|
||||
from kikotools.tools.embedding_autocomplete import KikoEmbeddingAutocomplete
|
||||
|
||||
settings = KikoEmbeddingAutocomplete.SETTINGS
|
||||
assert "enabled" in settings
|
||||
assert "trigger_chars" in settings
|
||||
assert "max_suggestions" in settings
|
||||
assert "show_embeddings" in settings
|
||||
assert "show_loras" in settings
|
||||
assert "case_sensitive" in settings
|
||||
|
||||
# Check settings structure
|
||||
assert settings["enabled"]["type"] == "boolean"
|
||||
assert settings["enabled"]["default"] is True
|
||||
assert settings["trigger_chars"]["type"] == "combo"
|
||||
assert settings["trigger_chars"]["options"] == [1, 2, 3, 4]
|
||||
|
||||
|
||||
def test_input_types():
|
||||
"""Test INPUT_TYPES class method."""
|
||||
from kikotools.tools.embedding_autocomplete import KikoEmbeddingAutocomplete
|
||||
|
||||
input_types = KikoEmbeddingAutocomplete.INPUT_TYPES()
|
||||
assert "required" in input_types
|
||||
assert "optional" in input_types
|
||||
assert "refresh" in input_types["optional"]
|
||||
|
||||
|
||||
def test_api_suggestions():
|
||||
"""Test the API suggestions method."""
|
||||
from kikotools.tools.embedding_autocomplete.node import KikoEmbeddingAutocompleteAPI
|
||||
|
||||
# Mock folder_paths to return some test files
|
||||
sys.modules['folder_paths'].get_filename_list = MagicMock(
|
||||
side_effect=lambda x: ["test1.pt", "test2.safetensors"] if x == "embeddings" else ["lora1.pt", "lora2.safetensors"]
|
||||
)
|
||||
|
||||
# Test with embeddings
|
||||
suggestions = KikoEmbeddingAutocompleteAPI.get_suggestions(
|
||||
prefix="test",
|
||||
include_embeddings=True,
|
||||
include_loras=False
|
||||
)
|
||||
|
||||
assert len(suggestions) == 2
|
||||
assert suggestions[0]["type"] == "embedding"
|
||||
assert suggestions[0]["name"] == "test1"
|
||||
|
||||
# Test with LoRAs
|
||||
suggestions = KikoEmbeddingAutocompleteAPI.get_suggestions(
|
||||
prefix="lora",
|
||||
include_embeddings=False,
|
||||
include_loras=True
|
||||
)
|
||||
|
||||
assert len(suggestions) == 2
|
||||
assert suggestions[0]["type"] == "lora"
|
||||
assert "<lora:" in suggestions[0]["value"]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_import()
|
||||
test_settings_defined()
|
||||
test_input_types()
|
||||
test_api_suggestions()
|
||||
print("All tests passed!")
|
||||
@@ -0,0 +1,416 @@
|
||||
/**
|
||||
* KikoTools Embedding Autocomplete
|
||||
* Provides autocomplete functionality for embeddings and LoRAs in ComfyUI text inputs
|
||||
*/
|
||||
|
||||
import { app } from "../../scripts/app.js";
|
||||
import { ComfyWidgets } from "../../scripts/widgets.js";
|
||||
|
||||
// Autocomplete class
|
||||
class KikoEmbeddingAutocomplete {
|
||||
constructor() {
|
||||
this.suggestions = [];
|
||||
this.currentIndex = -1;
|
||||
this.dropdown = null;
|
||||
this.activeWidget = null;
|
||||
this.activeNode = null;
|
||||
|
||||
// Settings
|
||||
this.settings = {
|
||||
enabled: true,
|
||||
triggerChars: 2,
|
||||
maxSuggestions: 20,
|
||||
showEmbeddings: true,
|
||||
showLoras: true,
|
||||
caseSensitive: false
|
||||
};
|
||||
|
||||
// Load settings from ComfyUI
|
||||
this.loadSettings();
|
||||
}
|
||||
|
||||
loadSettings() {
|
||||
this.settings.enabled = app.ui.settings.getSettingValue("kikotools.embedding_autocomplete.enabled", true);
|
||||
this.settings.triggerChars = app.ui.settings.getSettingValue("kikotools.embedding_autocomplete.trigger_chars", 2);
|
||||
this.settings.maxSuggestions = app.ui.settings.getSettingValue("kikotools.embedding_autocomplete.max_suggestions", 20);
|
||||
this.settings.showEmbeddings = app.ui.settings.getSettingValue("kikotools.embedding_autocomplete.show_embeddings", true);
|
||||
this.settings.showLoras = app.ui.settings.getSettingValue("kikotools.embedding_autocomplete.show_loras", true);
|
||||
this.settings.caseSensitive = app.ui.settings.getSettingValue("kikotools.embedding_autocomplete.case_sensitive", false);
|
||||
}
|
||||
|
||||
async fetchSuggestions(prefix) {
|
||||
if (!this.settings.enabled || prefix.length < this.settings.triggerChars) {
|
||||
return [];
|
||||
}
|
||||
|
||||
try {
|
||||
const params = new URLSearchParams({
|
||||
prefix: prefix,
|
||||
max: this.settings.maxSuggestions,
|
||||
embeddings: this.settings.showEmbeddings,
|
||||
loras: this.settings.showLoras,
|
||||
case_sensitive: this.settings.caseSensitive
|
||||
});
|
||||
|
||||
const response = await fetch(`/kikotools/autocomplete/suggestions?${params}`);
|
||||
if (response.ok) {
|
||||
return await response.json();
|
||||
}
|
||||
} catch (error) {
|
||||
console.error("[KikoAutocomplete] Error fetching suggestions:", error);
|
||||
}
|
||||
|
||||
return [];
|
||||
}
|
||||
|
||||
extractCurrentWord(text, cursorPos) {
|
||||
// Extract word at cursor position
|
||||
const beforeCursor = text.substring(0, cursorPos);
|
||||
|
||||
// Look for special triggers
|
||||
const embeddingMatch = beforeCursor.match(/embedding:([^,\s]*)$/);
|
||||
if (embeddingMatch) {
|
||||
return {
|
||||
word: embeddingMatch[1],
|
||||
start: beforeCursor.length - embeddingMatch[1].length,
|
||||
type: "embedding"
|
||||
};
|
||||
}
|
||||
|
||||
const loraMatch = beforeCursor.match(/<lora:([^:>]*)$/);
|
||||
if (loraMatch) {
|
||||
return {
|
||||
word: loraMatch[1],
|
||||
start: beforeCursor.length - loraMatch[1].length,
|
||||
type: "lora"
|
||||
};
|
||||
}
|
||||
|
||||
// General word extraction
|
||||
const wordMatch = beforeCursor.match(/([^,\s]+)$/);
|
||||
if (wordMatch) {
|
||||
return {
|
||||
word: wordMatch[1],
|
||||
start: beforeCursor.length - wordMatch[1].length,
|
||||
type: "general"
|
||||
};
|
||||
}
|
||||
|
||||
return null;
|
||||
}
|
||||
|
||||
createDropdown() {
|
||||
if (this.dropdown) {
|
||||
this.removeDropdown();
|
||||
}
|
||||
|
||||
this.dropdown = document.createElement("div");
|
||||
this.dropdown.className = "kiko-autocomplete-dropdown";
|
||||
this.dropdown.style.cssText = `
|
||||
position: absolute;
|
||||
background: var(--comfy-menu-bg, #353535);
|
||||
border: 1px solid var(--border-color, #4e4e4e);
|
||||
border-radius: 4px;
|
||||
max-height: 200px;
|
||||
overflow-y: auto;
|
||||
z-index: 10000;
|
||||
box-shadow: 0 2px 5px rgba(0,0,0,0.3);
|
||||
min-width: 200px;
|
||||
`;
|
||||
|
||||
document.body.appendChild(this.dropdown);
|
||||
}
|
||||
|
||||
removeDropdown() {
|
||||
if (this.dropdown) {
|
||||
this.dropdown.remove();
|
||||
this.dropdown = null;
|
||||
this.currentIndex = -1;
|
||||
}
|
||||
}
|
||||
|
||||
positionDropdown(textarea) {
|
||||
if (!this.dropdown) return;
|
||||
|
||||
const rect = textarea.getBoundingClientRect();
|
||||
|
||||
// Simple positioning below the textarea
|
||||
// In a more advanced version, we'd calculate caret position
|
||||
this.dropdown.style.left = rect.left + "px";
|
||||
this.dropdown.style.top = (rect.bottom + 2) + "px";
|
||||
this.dropdown.style.width = rect.width + "px";
|
||||
}
|
||||
|
||||
showSuggestions(suggestions, textarea, wordInfo) {
|
||||
if (!suggestions || suggestions.length === 0) {
|
||||
this.removeDropdown();
|
||||
return;
|
||||
}
|
||||
|
||||
this.suggestions = suggestions;
|
||||
this.createDropdown();
|
||||
this.positionDropdown(textarea);
|
||||
|
||||
// Populate dropdown
|
||||
this.dropdown.innerHTML = "";
|
||||
suggestions.forEach((suggestion, index) => {
|
||||
const item = document.createElement("div");
|
||||
item.className = "kiko-autocomplete-item";
|
||||
item.style.cssText = `
|
||||
padding: 5px 10px;
|
||||
cursor: pointer;
|
||||
color: var(--fg-color, #fff);
|
||||
display: flex;
|
||||
justify-content: space-between;
|
||||
align-items: center;
|
||||
`;
|
||||
|
||||
// Highlight matching part
|
||||
const displayText = suggestion.display || suggestion.value;
|
||||
const matchIndex = displayText.toLowerCase().indexOf(wordInfo.word.toLowerCase());
|
||||
|
||||
if (matchIndex >= 0 && wordInfo.word) {
|
||||
const before = displayText.substring(0, matchIndex);
|
||||
const match = displayText.substring(matchIndex, matchIndex + wordInfo.word.length);
|
||||
const after = displayText.substring(matchIndex + wordInfo.word.length);
|
||||
|
||||
item.innerHTML = `
|
||||
<span>${before}<strong style="color: var(--primary-color, #ff6b6b);">${match}</strong>${after}</span>
|
||||
<span style="opacity: 0.6; font-size: 0.9em;">${suggestion.type}</span>
|
||||
`;
|
||||
} else {
|
||||
item.innerHTML = `
|
||||
<span>${displayText}</span>
|
||||
<span style="opacity: 0.6; font-size: 0.9em;">${suggestion.type}</span>
|
||||
`;
|
||||
}
|
||||
|
||||
// Hover effect
|
||||
item.addEventListener("mouseenter", () => {
|
||||
this.selectIndex(index);
|
||||
});
|
||||
|
||||
// Click to select
|
||||
item.addEventListener("click", () => {
|
||||
this.applySuggestion(textarea, wordInfo, suggestion);
|
||||
});
|
||||
|
||||
this.dropdown.appendChild(item);
|
||||
});
|
||||
|
||||
// Select first item by default
|
||||
if (suggestions.length > 0) {
|
||||
this.selectIndex(0);
|
||||
}
|
||||
}
|
||||
|
||||
selectIndex(index) {
|
||||
// Remove previous selection
|
||||
const items = this.dropdown?.querySelectorAll(".kiko-autocomplete-item");
|
||||
if (!items) return;
|
||||
|
||||
items.forEach((item, i) => {
|
||||
if (i === index) {
|
||||
item.style.background = "var(--comfy-input-bg, #222)";
|
||||
this.currentIndex = i;
|
||||
} else {
|
||||
item.style.background = "";
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
applySuggestion(textarea, wordInfo, suggestion) {
|
||||
const text = textarea.value;
|
||||
const beforeWord = text.substring(0, wordInfo.start);
|
||||
const afterWord = text.substring(wordInfo.start + wordInfo.word.length);
|
||||
|
||||
// Determine what to insert based on context
|
||||
let insertion = suggestion.value;
|
||||
|
||||
// Handle different insertion contexts
|
||||
if (wordInfo.type === "embedding" && !suggestion.value.startsWith("embedding:")) {
|
||||
insertion = "embedding:" + suggestion.value;
|
||||
} else if (wordInfo.type === "lora" && !suggestion.value.startsWith("<lora:")) {
|
||||
insertion = `<lora:${suggestion.value}:1.0>`;
|
||||
}
|
||||
|
||||
// Apply the suggestion
|
||||
textarea.value = beforeWord + insertion + afterWord;
|
||||
|
||||
// Position cursor after insertion
|
||||
const newPos = beforeWord.length + insertion.length;
|
||||
textarea.setSelectionRange(newPos, newPos);
|
||||
|
||||
// Trigger input event to update the widget
|
||||
textarea.dispatchEvent(new Event("input", { bubbles: true }));
|
||||
|
||||
// Remove dropdown
|
||||
this.removeDropdown();
|
||||
}
|
||||
|
||||
attachToWidget(widget, node) {
|
||||
const textarea = widget.inputEl;
|
||||
if (!textarea || textarea.kikoAutocompleteAttached) return;
|
||||
|
||||
textarea.kikoAutocompleteAttached = true;
|
||||
|
||||
// Input handler with debouncing
|
||||
let debounceTimer;
|
||||
textarea.addEventListener("input", async (e) => {
|
||||
clearTimeout(debounceTimer);
|
||||
debounceTimer = setTimeout(async () => {
|
||||
const cursorPos = textarea.selectionStart;
|
||||
const wordInfo = this.extractCurrentWord(textarea.value, cursorPos);
|
||||
|
||||
if (wordInfo && wordInfo.word.length >= this.settings.triggerChars) {
|
||||
const suggestions = await this.fetchSuggestions(wordInfo.word);
|
||||
this.showSuggestions(suggestions, textarea, wordInfo);
|
||||
this.activeWidget = widget;
|
||||
this.activeNode = node;
|
||||
} else {
|
||||
this.removeDropdown();
|
||||
}
|
||||
}, 200);
|
||||
});
|
||||
|
||||
// Keyboard navigation
|
||||
textarea.addEventListener("keydown", (e) => {
|
||||
if (!this.dropdown || this.suggestions.length === 0) return;
|
||||
|
||||
switch(e.key) {
|
||||
case "ArrowDown":
|
||||
e.preventDefault();
|
||||
this.selectIndex((this.currentIndex + 1) % this.suggestions.length);
|
||||
break;
|
||||
|
||||
case "ArrowUp":
|
||||
e.preventDefault();
|
||||
this.selectIndex(this.currentIndex <= 0 ? this.suggestions.length - 1 : this.currentIndex - 1);
|
||||
break;
|
||||
|
||||
case "Enter":
|
||||
case "Tab":
|
||||
if (this.currentIndex >= 0) {
|
||||
e.preventDefault();
|
||||
const cursorPos = textarea.selectionStart;
|
||||
const wordInfo = this.extractCurrentWord(textarea.value, cursorPos);
|
||||
if (wordInfo) {
|
||||
this.applySuggestion(textarea, wordInfo, this.suggestions[this.currentIndex]);
|
||||
}
|
||||
}
|
||||
break;
|
||||
|
||||
case "Escape":
|
||||
e.preventDefault();
|
||||
this.removeDropdown();
|
||||
break;
|
||||
}
|
||||
});
|
||||
|
||||
// Hide on blur
|
||||
textarea.addEventListener("blur", () => {
|
||||
// Delay to allow click events to fire
|
||||
setTimeout(() => this.removeDropdown(), 200);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// Initialize and register
|
||||
app.registerExtension({
|
||||
name: "kikotools.embeddingAutocomplete",
|
||||
|
||||
async init() {
|
||||
// Add settings
|
||||
app.ui.settings.addSetting({
|
||||
id: "kikotools.embedding_autocomplete.enabled",
|
||||
name: "🫶 Embedding Autocomplete: Enabled",
|
||||
defaultValue: true,
|
||||
type: "boolean",
|
||||
onChange(value) {
|
||||
if (window.kikoAutocomplete) {
|
||||
window.kikoAutocomplete.settings.enabled = value;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
app.ui.settings.addSetting({
|
||||
id: "kikotools.embedding_autocomplete.trigger_chars",
|
||||
name: "🫶 Embedding Autocomplete: Trigger After Characters",
|
||||
defaultValue: 2,
|
||||
type: "combo",
|
||||
options: (value) => [
|
||||
{ value: 1, text: "1", selected: value === 1 },
|
||||
{ value: 2, text: "2", selected: value === 2 },
|
||||
{ value: 3, text: "3", selected: value === 3 },
|
||||
{ value: 4, text: "4", selected: value === 4 }
|
||||
],
|
||||
onChange(value) {
|
||||
if (window.kikoAutocomplete) {
|
||||
window.kikoAutocomplete.settings.triggerChars = value;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
app.ui.settings.addSetting({
|
||||
id: "kikotools.embedding_autocomplete.max_suggestions",
|
||||
name: "🫶 Embedding Autocomplete: Maximum Suggestions",
|
||||
defaultValue: 20,
|
||||
type: "combo",
|
||||
options: (value) => [
|
||||
{ value: 10, text: "10", selected: value === 10 },
|
||||
{ value: 20, text: "20", selected: value === 20 },
|
||||
{ value: 30, text: "30", selected: value === 30 },
|
||||
{ value: 50, text: "50", selected: value === 50 }
|
||||
],
|
||||
onChange(value) {
|
||||
if (window.kikoAutocomplete) {
|
||||
window.kikoAutocomplete.settings.maxSuggestions = value;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
app.ui.settings.addSetting({
|
||||
id: "kikotools.embedding_autocomplete.show_embeddings",
|
||||
name: "🫶 Embedding Autocomplete: Show Embeddings",
|
||||
defaultValue: true,
|
||||
type: "boolean",
|
||||
onChange(value) {
|
||||
if (window.kikoAutocomplete) {
|
||||
window.kikoAutocomplete.settings.showEmbeddings = value;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
app.ui.settings.addSetting({
|
||||
id: "kikotools.embedding_autocomplete.show_loras",
|
||||
name: "🫶 Embedding Autocomplete: Show LoRAs",
|
||||
defaultValue: true,
|
||||
type: "boolean",
|
||||
onChange(value) {
|
||||
if (window.kikoAutocomplete) {
|
||||
window.kikoAutocomplete.settings.showLoras = value;
|
||||
}
|
||||
}
|
||||
});
|
||||
},
|
||||
|
||||
async setup() {
|
||||
// Create autocomplete instance
|
||||
window.kikoAutocomplete = new KikoEmbeddingAutocomplete();
|
||||
|
||||
// Override STRING widget creation to add autocomplete
|
||||
const originalStringWidget = ComfyWidgets.STRING;
|
||||
ComfyWidgets.STRING = function(node, inputName, inputData, app) {
|
||||
const widget = originalStringWidget.apply(this, arguments);
|
||||
|
||||
// Only attach to multiline text widgets (prompts)
|
||||
if (inputData[1]?.multiline && widget.inputEl) {
|
||||
setTimeout(() => {
|
||||
window.kikoAutocomplete.attachToWidget(widget, node);
|
||||
}, 100);
|
||||
}
|
||||
|
||||
return widget;
|
||||
};
|
||||
}
|
||||
});
|
||||
Reference in New Issue
Block a user