Added basic function of autocomplete

This commit is contained in:
newtextdoc1111
2025-05-01 00:03:14 +09:00
parent 41515ed2bc
commit 9b22c27b48
9 changed files with 1150 additions and 0 deletions
+4
View File
@@ -0,0 +1,4 @@
/.idea/
/venv/
/__pycache__/
**/__pycache__/*.pyc
+3
View File
@@ -0,0 +1,3 @@
# ComfyUI-Autocomplete-Plus
WIP
+8
View File
@@ -0,0 +1,8 @@
from .modules.api import *
# --- Original Node Mappings (Keep or modify as needed) ---
WEB_DIRECTORY = "./web"
NODE_CLASS_MAPPINGS = {}
__all__ = []
print("### Loading: ComfyUI-Autocomplete-Plus ###")
+97
View File
@@ -0,0 +1,97 @@
import os
import csv
import json
from aiohttp import web
import server # ComfyUIのサーバーインスタンス
# --- Helper Function for CSV Parsing ---
def parse_aliases(alias_str):
"""
Parses a comma-separated string.
Example: 'alias1,alias 2, alias3'
"""
aliases = []
if not alias_str:
return aliases
return [word.strip() for word in alias_str.split(',')]
# --- API Endpoints ---
@server.PromptServer.instance.routes.get('/autocomplete-plus/tags')
async def get_tags(request):
csv_path = os.path.join(os.path.dirname(__file__), '../', 'user', 'danbooru_tags.csv')
tags_data = []
try:
with open(csv_path, 'r', encoding='utf-8') as f:
reader = csv.reader(f)
header = next(reader) # Skip header
try:
tag_index = header.index('tag')
alias_index = header.index('alias')
count_index = header.index('count')
except ValueError:
print(f"[Autocomplete-Plus] Error: Invalid header in {csv_path}: {header}")
return web.Response(status=500, text=f"Invalid CSV header in {csv_path}")
for row in reader:
try:
tag = row[tag_index]
aliases = parse_aliases(row[alias_index])
count = int(row[count_index])
if tag: # Ensure tag is not empty
tags_data.append({'tag': tag, 'alias': aliases, 'count': count})
except (IndexError, ValueError) as e:
print(f"[Autocomplete-Plus] Warning: Skipping invalid row in {csv_path}: {row} - Error: {e}")
continue # Skip malformed rows
# Sort by count descending before sending
tags_data.sort(key=lambda x: x['count'], reverse=True)
return web.json_response(tags_data)
except FileNotFoundError:
print(f"[Autocomplete-Plus] Error: Tags file not found at {csv_path}")
return web.Response(status=404, text="Tags file not found.")
except Exception as e:
print(f"[Autocomplete-Plus] Error reading tags file: {e}")
return web.Response(status=500, text=f"Error reading tags file: {e}")
@server.PromptServer.instance.routes.get('/autocomplete-plus/cooccurrence')
async def get_cooccurrence(request):
csv_path = os.path.join(os.path.dirname(__file__), '../', 'user', 'danbooru_tags_cooccurrence.csv')
cooccurrence_data = {}
try:
with open(csv_path, 'r', encoding='utf-8') as f:
reader = csv.reader(f)
header = next(reader) # Skip header
try:
tag_a_index = header.index('tag_a')
tag_b_index = header.index('tag_b')
count_index = header.index('count')
except ValueError:
print(f"[Autocomplete-Plus] Error: Invalid header in {csv_path}: {header}")
return web.Response(status=500, text=f"Invalid CSV header in {csv_path}")
processed_count = 0
for row in reader:
try:
tag_a = row[tag_a_index]
tag_b = row[tag_b_index]
count = int(row[count_index])
if tag_a and tag_b: # Ensure tags are not empty
if tag_a not in cooccurrence_data:
cooccurrence_data[tag_a] = {}
cooccurrence_data[tag_a][tag_b] = count
processed_count += 1
except (IndexError, ValueError) as e:
print(f"[Autocomplete-Plus] Warning: Skipping invalid row in {csv_path}: {row} - Error: {e}")
continue # Skip malformed rows
print(f"[Autocomplete-Plus] Processed {processed_count} cooccurrence pairs for API.")
return web.json_response(cooccurrence_data)
except FileNotFoundError:
print(f"[Autocomplete-Plus] Error: Cooccurrence file not found at {csv_path}")
return web.Response(status=404, text="Cooccurrence file not found.")
except Exception as e:
print(f"[Autocomplete-Plus] Error reading cooccurrence file: {e}")
return web.Response(status=500, text=f"Error reading cooccurrence file: {e}")
+152
View File
@@ -0,0 +1,152 @@
# Package Modules
import os
from typing import Union, BinaryIO, Dict, List, Tuple, Optional
import time
# ComfyUI Modules
import folder_paths
from comfy.utils import ProgressBar
# Your Modules
from .modules.calculator import CalculatorModel
# Basic practice to get paths from ComfyUI
custom_nodes_script_dir = os.path.dirname(os.path.abspath(__file__))
custom_nodes_model_dir = os.path.join(folder_paths.models_dir, "my-custom-nodes")
custom_nodes_output_dir = os.path.join(folder_paths.get_output_directory(), "my-custom-nodes")
# These are example nodes that only contains basic functionalities with some comments.
# If you need detailed explanation, please refer to : https://docs.comfy.org/essentials/custom_node_walkthrough
# First Node:
class MyModelLoader:
# Define the input parameters of the node here.
@classmethod
def INPUT_TYPES(s):
my_models = ["Model A", "Model B", "Model C"]
return {
# If the key is "required", the value must be filled.
"required": {
# `my_models` is the list, so it will be shown as a dropdown menu in the node. ( So that user can select one of them. )
# You must provide the value in the tuple format. e.g. ("value",) or (3,) or ([1, 2],) etc.
"model": (my_models,),
"device": (['cuda', 'cpu', 'auto'],),
},
# If the key is "optional", the value is optional.
"optional": {
"compute_type": (['float32', 'float16'],),
}
}
# Define these constants inside the node.
# `RETURN_TYPES` is important, as it limits the parameter types that can be passed to the next node, in `INPUT_TYPES()` above.
RETURN_TYPES = ("MY_MODEL",)
RETURN_NAMES = ("my_model",)
# `FUNCTION` is the function name that will be called in the node.
FUNCTION = "load_model"
# `CATEGORY` is the category name that will be used when user searches the node.
CATEGORY = "CustomNodesTemplate"
# In the function, use same parameter names as you specified in `INPUT_TYPES()`
def load_model(self,
model: str,
device: str,
compute_type: Optional[str] = None,
) -> Tuple[CalculatorModel]:
calculator_model = CalculatorModel()
calculator_model.load_model(model, device, compute_type)
# You can use `comfy.utils.ProgressBar` to show the progress of the process.
# First, initialize the total amount of the process.
total_steps = 5
comfy_pbar = ProgressBar(total_steps)
# Then, update the progress.
for i in range(1, total_steps):
time.sleep(1)
comfy_pbar.update(i) # Alternatively, you can use `comfy_pbar.update_absolute(value)` to update the progress with absolute value.
# Return the model as a tuple.
return (calculator_model, )
# Second Node
class CalculatePlus:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"model": ("MY_MODEL", ),
},
# Specify the parameters with type and default value.
"optional": {
"a": ("INT", {"default": 5}),
"b": ("INT", {"default": 10}),
}
}
RETURN_TYPES = ("INT",)
RETURN_NAMES = ("plus_value",)
FUNCTION = "plus"
CATEGORY = "CustomNodesTemplate"
def plus(self,
model: CalculatorModel,
a: Optional[int],
b: Optional[int],
) -> Tuple[int]:
result = model.plus(a, b)
return (result, )
# Third Node
class CalculateMinus:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"model": ("MY_MODEL", ),
"a": ("INT", ),
},
"optional": {
"b": ("INT", {"default": 10}),
}
}
RETURN_TYPES = ("INT",)
RETURN_NAMES = ("minus_value",)
FUNCTION = "minus"
CATEGORY = "CustomNodesTemplate"
def minus(self,
model: CalculatorModel,
a: Optional[int],
b: Optional[int],
) -> Tuple[int]:
result = model.minus(a, b)
return (result, )
# Output Node
class ExampleOutputNode:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"value": ("INT", ),
},
}
# If the node is output node, set this to True.
OUTPUT_NODE = True
RETURN_TYPES = ("INT",)
RETURN_NAMES = ("int",)
FUNCTION = "result"
CATEGORY = "CustomNodesTemplate"
def result(self,
value: int,) -> Tuple[int]:
return (value, )
+15
View File
@@ -0,0 +1,15 @@
[project]
name = "comfyui-autocomplete-plus"
description = "Custom node to add autocomplete functionality [ComfyUI-Autocomplete-Plus](https://github.com/newtextdoc1111/ComfyUI-Autocomplete-Plus)."
version = "1.0.0"
license = {file = "LICENSE"}
dependencies = ["",]
[project.urls]
Repository = "https://github.com/newtextdoc1111/ComfyUI-Autocomplete-Plus"
# Used by Comfy Registry https://comfyregistry.org
[tool.comfy]
PublisherId = "newtextdoc1111"
DisplayName = "ComfyUI-Autocomplete-Plus"
Icon = ""
+49
View File
@@ -0,0 +1,49 @@
/* Basic styles for autocomplete - will be refined later */
#autocomplete-plus-list {
background-color: var(--comfy-input-bg);
border: 1px solid var(--border-color);
border-radius: 4px;
box-shadow: 0 2px 5px rgba(0, 0, 0, 0.2);
min-width: 200px; /* Increased min width */
overflow-y: auto;
position: absolute; /* Ensure position is absolute */
z-index: 10000; /* High z-index */
color: var(--input-text);
font-size: 0.9em; /* Slightly smaller font */
}
.autocomplete-plus-item {
padding: 5px 10px; /* Adjusted padding */
cursor: pointer;
white-space: nowrap; /* Prevent wrapping */
overflow: hidden;
text-overflow: ellipsis; /* Show ellipsis if too long */
border-bottom: 1px solid var(--border-color); /* Separator line */
}
.autocomplete-plus-item:last-child {
border-bottom: none; /* No border for the last item */
}
/* Alternating row colors */
.autocomplete-plus-item:nth-child(even) {
background-color: var(--comfy-menu-bg); /* Slightly different background for even rows */
}
.autocomplete-plus-item:nth-child(odd) {
background-color: var(--comfy-input-bg);
}
/* Hover effect */
.autocomplete-plus-item:hover {
background-color: var(--comfy-hover-bg); /* Use a hover variable if available, or adjust */
color: var(--desc-text); /* Adjust text color on hover if needed */
}
/* Selected item style */
.autocomplete-plus-item.selected {
background-color: var(--comfy-menu-bg-selected, var(--comfy-button-bg)); /* Use selection color, fallback to button */
color: var(--comfy-menu-text-selected, var(--button-text)); /* Use selected text color, fallback */
font-weight: bold;
}
+379
View File
@@ -0,0 +1,379 @@
// --- Autocomplete UI Class ---
export class AutocompleteUI {
constructor() {
this.element = document.createElement('div');
this.element.id = 'autocomplete-plus-list';
this.element.style.display = 'none'; // Initially hidden
this.element.style.position = 'absolute'; // Position near the input
this.element.style.zIndex = '10000'; // Ensure it's on top
this.activeInput = null;
this.selectedIndex = -1;
this.candidates = [];
document.body.appendChild(this.element);
// Add event listener for clicks on items
this.element.addEventListener('mousedown', (e) => {
if (e.target.tagName === 'DIV') {
const index = parseInt(e.target.dataset.index, 10);
if (!isNaN(index)) {
this.#selectItem(index);
e.preventDefault(); // Prevent focus loss from input
e.stopPropagation();
}
}
});
}
show(textareaElement, candidates) {
this.#update(candidates);
if (candidates.length > 0) {
this.#show(textareaElement);
} else {
this.#hide();
}
}
hide() {
this.#hide();
}
isVisible() {
return this.element.style.display !== 'none';
}
/** Moves the selection up or down */
navigate(direction) {
if (this.candidates.length === 0) return;
this.selectedIndex += direction;
if (this.selectedIndex < 0) {
this.selectedIndex = this.candidates.length - 1; // Wrap around to bottom
} else if (this.selectedIndex >= this.candidates.length) {
this.selectedIndex = 0; // Wrap around to top
}
this.#highlightItem(this.selectedIndex);
}
/** Selects the currently highlighted item */
getSelectedTag() {
if (this.selectedIndex >= 0 && this.selectedIndex < this.candidates.length) {
return this.candidates[this.selectedIndex].tag;
}
return null; // No valid selection
}
/**
* Shows the autocomplete list near the target input element.
* Position calculation logic inspired by:
* https://github.com/pythongosssss/ComfyUI-Custom-Scripts/blob/main/web/js/common/autocomplete.js
* License: MIT License (assumed based on repository root LICENSE file)
* Considers ComfyUI canvas scale.
* Adjusts max-height and enables scrolling if the list exceeds viewport bounds.
* @param {HTMLTextAreaElement} inputElement The textarea being typed into.
*/
#show(inputElement) {
this.activeInput = inputElement;
this.selectedIndex = 0;
// Calculate caret position using the helper function (returns viewport-relative coordinates)
this.#updateCoordinate();
this.element.style.overflowY = 'auto';
this.element.style.display = 'block'; // Make it visible
// Highlight the first item
this.#highlightItem(this.selectedIndex);
}
#updateCoordinate() {
const { top: caretTop, left: caretLeft, lineHeight: caretLineHeight } = this.#getCaretCoordinates(this.activeInput);
// Reset scroll position and max-height before calculating position
this.element.scrollTop = 0;
this.element.style.maxHeight = ''; // Reset max-height for accurate measurement
// Get ComfyUI canvas scale if available, otherwise default to 1
const scale = window.app?.canvas?.ds?.scale ?? 1.0;
// Initial desired position: below the current text line where the caret is.
let topPosition = caretTop + (caretLineHeight * scale);
let leftPosition = caretLeft;
// Make the list visible *before* getting its dimensions to ensure they are accurate
// Use visibility instead of display to measure without affecting layout yet
this.element.style.visibility = 'hidden';
this.element.style.display = 'block';
const listRect = this.element.getBoundingClientRect(); // Dimensions without max-height constraint
const naturalHeight = listRect.height;
this.element.style.display = 'none'; // Hide again until final position is set
this.element.style.visibility = 'visible';
const viewportWidth = window.innerWidth;
const viewportHeight = window.innerHeight;
const margin = 5; // Small margin from viewport edges
// --- Horizontal Collision Detection and Adjustment ---
if (leftPosition + listRect.width > viewportWidth - margin) {
leftPosition = viewportWidth - listRect.width - margin;
}
if (leftPosition < margin) {
leftPosition = margin;
}
// --- Vertical Collision Detection and Adjustment ---
const availableSpaceBelow = viewportHeight - topPosition - margin;
const availableSpaceAbove = caretTop - margin; // Space above the caret line
// Reset max-height before deciding
this.element.style.maxHeight = '';
if (naturalHeight <= availableSpaceBelow) {
// Fits perfectly below the caret
// topPosition remains as calculated initially
} else {
// Doesn't fit below, check if it fits perfectly above
const topAboveCaret = caretTop - naturalHeight - margin; // Position list bottom just above caret line
if (naturalHeight <= availableSpaceAbove) {
// Fits perfectly above
topPosition = topAboveCaret;
} else {
// Doesn't fit perfectly either below or above, needs scrolling.
// Choose the position (above or below) that offers more space.
if (availableSpaceBelow >= availableSpaceAbove) {
// Scroll below: topPosition remains as initially calculated
this.element.style.maxHeight = `${availableSpaceBelow}px`;
} else {
// Scroll above: Position near the top edge and set max-height
topPosition = margin;
this.element.style.maxHeight = `${availableSpaceAbove}px`;
}
}
}
// Final check to prevent going off the top edge
if (topPosition < margin) {
topPosition = margin;
// If pushed down, recalculate max-height if it was set based on top alignment
if (this.element.style.maxHeight && availableSpaceBelow < availableSpaceAbove) {
// Recalculate max-height based on space from the top margin
this.element.style.maxHeight = `${viewportHeight - margin - margin}px`;
}
}
// Apply the calculated position and display the element
this.element.style.left = `${leftPosition}px`;
this.element.style.top = `${topPosition}px`;
}
/**
* Hides the autocomplete list.
*/
#hide() {
this.element.style.display = 'none';
this.selectedIndex = -1;
this.activeInput = null;
this.candidates = [];
}
/**
* Updates the list with new candidates.
* @param {Array<{tag: string, count: number, alias?: string}>} candidates List of candidate tags.
*/
#update(candidates) {
this.candidates = candidates;
this.element.innerHTML = ''; // Clear previous items
if (candidates.length === 0) {
this.#hide();
return;
}
candidates.forEach((candidate, index) => {
const item = document.createElement('div');
item.classList.add('autocomplete-plus-item');
item.dataset.index = index;
// Display tag and count (alias if matched by alias)
item.textContent = `${candidate.tag} (${candidate.count})${candidate.alias ? ` [${candidate.alias}]` : ''}`;
item.style.padding = '4px 8px';
item.style.cursor = 'pointer';
this.element.appendChild(item);
});
}
/** Highlights the item at the given index */
#highlightItem() {
const items = this.element.children;
for (let i = 0; i < items.length; i++) {
if (i === this.selectedIndex) {
items[i].classList.add('selected'); // Use CSS class for selection
items[i].scrollIntoView({ block: 'nearest' });
} else {
items[i].classList.remove('selected');
}
}
}
/**
* Handles the selection of an item (e.g., inserts into input).
* @param {number} index The index of the selected candidate.
*/
#selectItem(index) {
if (!this.activeInput || index < 0 || index >= this.candidates.length) {
this.#hide();
return;
}
this.#hide();
}
/**
* Gets the pixel coordinates of the caret in the input element.
* Uses a temporary div to calculate the position accurately.
* Based on https://github.com/component/textarea-caret-position
* @param {HTMLTextAreaElement} element The textarea element.
* @returns {{ top: number, left: number, lineHeight: number }}
*/
#getCaretCoordinates(element) {
const properties = [
'direction', // RTL support
'boxSizing',
'width', // on Chrome and IE, exclude the scrollbar, so the mirror div wraps exactly as the textarea does
'height',
'overflowX',
'overflowY', // copy the scrollbar for IE
'borderTopWidth',
'borderRightWidth',
'borderBottomWidth',
'borderLeftWidth',
'borderStyle',
'paddingTop',
'paddingRight',
'paddingBottom',
'paddingLeft',
// https://developer.mozilla.org/en-US/docs/Web/CSS/font
'fontStyle',
'fontVariant',
'fontWeight',
'fontStretch',
'fontSize',
'fontSizeAdjust',
'lineHeight',
'fontFamily',
'textAlign',
'textTransform',
'textIndent',
'textDecoration', // might not make a difference, but better be safe
'letterSpacing',
'wordSpacing',
'tabSize',
'MozTabSize' // Firefox
];
const isBrowser = typeof window !== 'undefined';
const isFirefox = isBrowser && window.mozInnerScreenX != null;
// The mirror div will replicate the textarea's style
const div = document.createElement('div');
div.id = 'input-textarea-caret-position-mirror-div';
document.body.appendChild(div);
const style = div.style;
const computed = window.getComputedStyle(element);
const isInput = element.nodeName === 'INPUT';
// Default textarea styles
style.whiteSpace = 'pre-wrap';
if (!isInput) style.wordWrap = 'break-word'; // only for textarea-s
// Position off-screen
style.position = 'absolute'; // required to return coordinates properly
style.visibility = 'hidden'; // not 'display: none' because we want rendering
// Transfer the element's properties to the div
properties.forEach(prop => {
if (isInput && prop === "lineHeight") {
// Special case for <input>s because text is rendered centered and line height may be != height
if (computed.boxSizing === "border-box") {
var height = parseInt(computed.height);
var outerHeight =
parseInt(computed.paddingTop) +
parseInt(computed.paddingBottom) +
parseInt(computed.borderTopWidth) +
parseInt(computed.borderBottomWidth);
var targetHeight = outerHeight + parseInt(computed.lineHeight);
if (height > targetHeight) {
style.lineHeight = height - outerHeight + "px";
} else if (height === targetHeight) {
style.lineHeight = computed.lineHeight;
} else {
style.lineHeight = 0;
}
} else {
style.lineHeight = computed.height;
}
} else {
style[prop] = computed[prop];
}
});
// Calculate lineHeight more robustly
let computedLineHeight = computed.lineHeight;
let numericLineHeight;
if (computedLineHeight === 'normal') {
// Calculate fallback based on font size
const fontSize = parseFloat(computed.fontSize);
numericLineHeight = Math.round(fontSize * 1.2); // Common approximation
} else {
numericLineHeight = parseFloat(computedLineHeight); // Use parseFloat for pixel values like "16px"
}
// Ensure we have a valid number, fallback if somehow still NaN
if (isNaN(numericLineHeight)) {
const fontSize = parseFloat(computed.fontSize);
numericLineHeight = Math.round(fontSize * 1.2) || 16; // Final fallback
}
if (isFirefox) {
// Firefox lies about the overflow property for textareas: https://bugzilla.mozilla.org/show_bug.cgi?id=984275
if (element.scrollHeight > parseInt(computed.height)) style.overflowY = 'scroll';
} else {
style.overflow = 'hidden'; // for Chrome to not render a scrollbar; IE keeps overflowY = 'scroll'
}
div.textContent = element.value.substring(0, element.selectionStart);
// The second special handling for input type=text doesn't need to be copied:
// If isInput then usage is https://github.com/component/textarea-caret-position#usage-input-typetext
const span = document.createElement('span');
// Wrapping must be replicated *exactly*, including whitespace spaces and carriage returns
span.textContent = element.value.substring(element.selectionStart) || '.'; // || '.' because a completely empty faux span doesn't render at all
div.appendChild(span);
const coordinates = {
top: span.offsetTop + (parseInt(computed['borderTopWidth']) || 0),
left: span.offsetLeft + (parseInt(computed['borderLeftWidth']) || 0),
lineHeight: numericLineHeight // Use the calculated numeric lineHeight
};
// Calculate the bounding rect of the input element relative to the viewport
const rect = element.getBoundingClientRect();
// Adjust the coordinates to be relative to the viewport
coordinates.top = rect.top + element.scrollTop + coordinates.top;
coordinates.left = rect.left + element.scrollLeft + coordinates.left;
document.body.removeChild(div);
return coordinates;
}
}
+443
View File
@@ -0,0 +1,443 @@
import { app } from "../../../scripts/app.js";
import { AutocompleteUI } from "./autocomplete.js";
(function() {
// Function to load a CSS file
function loadCSS(href) {
const link = document.createElement('link');
link.rel = 'stylesheet';
link.type = 'text/css';
link.href = href;
// Ensure the CSS is loaded before other scripts might rely on its styles
// by adding it to the head.
document.head.appendChild(link);
console.log(`Loaded CSS: ${href}`); // Optional: Log loading
}
// Construct the path to the CSS file
// Assumes your custom node directory is 'ComfyUI-Autocomplete-Plus'
const cssPath = `/extensions/ComfyUI-Autocomplete-Plus/css/autocomplete.css`;
// Load the CSS file
loadCSS(cssPath);
// --- ここから下に既存の main.js のコードを追加 ---
// 例:
// import { app } from "/scripts/app.js";
// app.registerExtension({...});
})(); // IIFE to avoid polluting the global scope
// Data storage
let tagMap = new Map();
let aliasMap = new Map();
let sortedTags = []; // Now populated directly from sorted API response
let cooccurrenceMap = new Map(); // Structure remains Map<string, Map<string, number>>
let tagsLoaded = false;
let cooccurrenceLoaded = false;
// --- Data Loading Functions ---
/**
* Loads and processes tag data from the Python API endpoint.
*/
async function loadTags() {
const url = '/autocomplete-plus/tags';
try {
const response = await fetch(url);
if (!response.ok) {
throw new Error(`HTTP error! status: ${response.status}`);
}
const data = await response.json(); // Data is already sorted by count
sortedTags = data; // Assign directly as it's pre-sorted
data.forEach(tagData => {
tagMap.set(tagData.tag, tagData);
if (tagData.alias && Array.isArray(tagData.alias)) {
tagData.alias.forEach(alias => {
if (!aliasMap.has(alias)) {
aliasMap.set(alias, tagData.tag); // Map alias back to the main tag
}
});
}
});
tagsLoaded = true;
console.log(`[Autocomplete-Plus] Processed ${sortedTags.length} tags from API.`);
} catch (error) {
console.error(`[Autocomplete-Plus] Failed to fetch or process tags from ${url}:`, error);
tagsLoaded = false;
}
}
/**
* Loads and processes cooccurrence data from the Python API endpoint.
*/
async function loadCooccurrence() {
const url = '/autocomplete-plus/cooccurrence';
try {
const response = await fetch(url);
if (!response.ok) {
throw new Error(`HTTP error! status: ${response.status}`);
}
const data = await response.json(); // Data is in { tag_a: { tag_b: count, ... }, ... } format
let primaryTagCount = 0;
let pairCount = 0;
for (const tagA in data) {
if (Object.hasOwnProperty.call(data, tagA)) {
const innerMap = new Map();
const pairs = data[tagA];
for (const tagB in pairs) {
if (Object.hasOwnProperty.call(pairs, tagB)) {
innerMap.set(tagB, pairs[tagB]);
pairCount++;
}
}
if (innerMap.size > 0) {
cooccurrenceMap.set(tagA, innerMap);
primaryTagCount++;
}
}
}
cooccurrenceLoaded = true;
console.log(`[Autocomplete-Plus] Processed ${pairCount} cooccurrence pairs for ${primaryTagCount} primary tags from API.`);
} catch (error) {
console.error(`[Autocomplete-Plus] Failed to fetch or process cooccurrence data from ${url}:`, error);
cooccurrenceLoaded = false;
}
}
// --- End Data Loading Functions ---
// --- Helper Functions ---
/**
* Converts Hiragana to Katakana.
* @param {string} str Input string.
* @returns {string} Katakana string.
*/
function hiraToKata(str) {
return str.replace(/[\u3041-\u3096]/g, function(match) {
const chr = match.charCodeAt(0) + 0x60;
return String.fromCharCode(chr);
});
}
/**
* Converts Katakana to Hiragana.
* @param {string} str Input string.
* @returns {string} Hiragana string.
*/
function kataToHira(str) {
return str.replace(/[\u30a1-\u30f6]/g, function(match) {
const chr = match.charCodeAt(0) - 0x60;
return String.fromCharCode(chr);
});
}
// --- End Helper Functions ---
// --- Autocomplete Logic ---
const MAX_CANDIDATES = 50; // Max number of suggestions to show
let autocompleteUI = null; // Singleton UI instance
/**
* Finds tag completion candidates based on the input query.
* Handles Hiragana/Katakana conversion for matching.
* @param {string} query The partial tag input.
* @returns {Array<{tag: string, count: number, alias?: string}>}
*/
function findCompletionCandidates(query) {
const startTime = performance.now(); // 検索開始時間を記録
if (!query || !tagsLoaded) {
return [];
}
const lowerQuery = query.toLowerCase();
const candidates = [];
const addedTags = new Set(); // Keep track of added tags to avoid duplicates
// Generate Hiragana/Katakana variations if applicable
const queryVariations = new Set([lowerQuery]);
const kataQuery = hiraToKata(lowerQuery);
if (kataQuery !== lowerQuery) {
queryVariations.add(kataQuery);
}
const hiraQuery = kataToHira(lowerQuery);
if (hiraQuery !== lowerQuery) {
queryVariations.add(hiraQuery);
}
// Search in sortedTags (already sorted by count)
for (const tagData of sortedTags) {
let matched = false;
let matchedAlias = null;
// Check primary tag against all variations
for (const variation of queryVariations) {
// Ensure tagData.tag is treated as lowercase for comparison
if (tagData.tag.toLowerCase().includes(variation)) {
matched = true;
break;
}
}
// If primary tag didn't match, check aliases against all variations
if (!matched && tagData.alias && Array.isArray(tagData.alias) && tagData.alias.length > 0) {
for (const alias of tagData.alias) {
const lowerAlias = alias.toLowerCase();
for (const variation of queryVariations) {
if (lowerAlias.includes(variation)) {
matched = true;
matchedAlias = alias; // Store the alias that matched
break;
}
}
if (matched) break; // Stop checking aliases for this tag if one matched
}
}
// Add candidate if matched and not already added
if (matched && !addedTags.has(tagData.tag)) {
candidates.push({
tag: tagData.tag,
count: tagData.count,
...(matchedAlias && { alias: matchedAlias }) // Add alias property only if matched via alias
});
addedTags.add(tagData.tag);
if (candidates.length >= MAX_CANDIDATES) {
// 早期リターンする場合もログを出力
const endTime = performance.now();
const duration = endTime - startTime;
console.log(`[Autocomplete-Plus] Search for "${query}" took ${duration.toFixed(2)}ms. Found ${candidates.length} candidates (max reached).`);
return candidates; // Early exit
}
}
}
// 検索終了時間を記録し、コンソールに出力
const endTime = performance.now();
const duration = endTime - startTime;
console.log(`[Autocomplete-Plus] Search for "${query}" took ${duration.toFixed(2)}ms. Found ${candidates.length} candidates.`);
return candidates;
}
/**
* Extracts the current tag being typed before the cursor.
* Assumes tags are separated by commas.
* @param {HTMLTextAreaElement} inputElement
* @returns {string} The partial tag or empty string.
*/
function getCurrentPartialTag(inputElement) {
const text = inputElement.value;
const cursorPos = inputElement.selectionStart;
// Find the last comma before the cursor
const lastComma = text.lastIndexOf(',', cursorPos - 1);
const start = lastComma === -1 ? 0 : lastComma + 1;
// Extract the text between the last comma (or start) and the cursor
const partial = text.substring(start, cursorPos).trimStart();
return partial;
}
/**
* Inserts the selected tag into the textarea, replacing the partial tag.
* @param {HTMLTextAreaElement} inputElement
* @param {string} tagToInsert
*/
function insertTag(inputElement, tagToInsert) {
const text = inputElement.value;
const cursorPos = inputElement.selectionStart;
const lastComma = text.lastIndexOf(',', cursorPos - 1);
const start = lastComma === -1 ? 0 : lastComma + 1;
const actualTag = tagToInsert.replace("_", " ");
// Find the end of the word/tag at the cursor (if any)
// This basic version assumes we replace up to the next comma or end of string
let end = text.indexOf(',', cursorPos);
if (end === -1) {
end = text.length;
}
// More precise: find word boundary if not comma
const nextSpace = text.indexOf(' ', cursorPos);
if (nextSpace !== -1 && nextSpace < end) {
// This logic might need refinement depending on desired behavior
// For now, replace up to the original cursor position if inserting mid-word
end = cursorPos;
}
// Ensure space after comma if inserting after a comma
const prefix = text.substring(0, start).trimEnd();
const needsSpaceBefore = start > 0 && prefix[prefix.length - 1] === ',';
const textBefore = text.substring(0, start) + (needsSpaceBefore ? ' ' : '');
const textAfter = text.substring(cursorPos); // Use cursorPos to replace only typed part
// Add comma and space if needed after insertion
const suffix = ', '; // Standard separator
inputElement.value = textBefore + actualTag + suffix + textAfter;
// Set cursor position after the inserted tag and the following comma+space
const newCursorPos = textBefore.length + actualTag.length + suffix.length;
inputElement.selectionStart = inputElement.selectionEnd = newCursorPos;
// Trigger input event for ComfyUI to recognize change
inputElement.dispatchEvent(new Event('input', { bubbles: true }));
}
// --- Event Handlers ---
function handleInput(event) {
if (!autocompleteUI) return;
const textareaElement = event.target;
const partialTag = getCurrentPartialTag(textareaElement);
if (partialTag.length > 0) {
const candidates = findCompletionCandidates(partialTag);
autocompleteUI.show(textareaElement, candidates);
} else {
autocompleteUI.hide();
}
}
function handleFocus(event) {
// Potentially show suggestions immediately on focus?
// For now, only show on input
if (!autocompleteUI) {
autocompleteUI = new AutocompleteUI();
}
// Maybe check if there's already text and show suggestions?
// handleInput(event); // Trigger check immediately
}
function handleBlur(event) {
// Need a slight delay because clicking the autocomplete list causes blur
setTimeout(() => {
if (autocompleteUI && !autocompleteUI.element.contains(document.activeElement)) {
autocompleteUI.hide();
}
}, 150);
}
function handleKeyDown(event) {
if (!autocompleteUI || !autocompleteUI.isVisible()) return;
const textareaElement = event.target;
switch (event.key) {
case 'ArrowDown':
event.preventDefault();
autocompleteUI.navigate(1);
break;
case 'ArrowUp':
event.preventDefault();
autocompleteUI.navigate(-1);
break;
case 'Enter':
case 'Tab':
if (autocompleteUI.getSelectedTag() !== null) {
event.preventDefault();
insertTag(textareaElement, autocompleteUI.getSelectedTag());
} else {
// Allow default Tab/Enter if no item is selected
autocompleteUI.hide();
}
break;
case 'Escape':
event.preventDefault();
autocompleteUI.hide();
break;
}
}
// --- Initialization ---
function initializeAutocomplete() {
if (!tagsLoaded) {
console.warn("[Autocomplete-Plus] Tags not loaded, cannot initialize autocomplete.");
return;
}
console.log("[Autocomplete-Plus] Initializing autocomplete features...");
autocompleteUI = new AutocompleteUI();
// Find relevant textareas (e.g., prompt inputs)
// This selector might need adjustment based on ComfyUI's structure
const targetSelectors = [
'.comfy-multiline-input',
// Add other selectors if needed
];
// Use MutationObserver to detect dynamically added textareas
const observer = new MutationObserver((mutations) => {
mutations.forEach((mutation) => {
mutation.addedNodes.forEach((node) => {
if (node.nodeType === Node.ELEMENT_NODE) {
targetSelectors.forEach(selector => {
// Check if the added node itself matches or contains matching elements
if (node.matches(selector)) {
attachListeners(node);
} else {
node.querySelectorAll(selector).forEach(attachListeners);
}
});
}
});
});
});
// Function to attach listeners
function attachListeners(element) {
if (element.dataset.autocompleteAttached) return; // Prevent double attachment
console.log("[Autocomplete-Plus] Attaching listeners to:", element);
element.addEventListener('input', handleInput);
element.addEventListener('focus', handleFocus);
element.addEventListener('blur', handleBlur);
element.addEventListener('keydown', handleKeyDown);
element.dataset.autocompleteAttached = 'true';
}
// Initial scan for existing elements
targetSelectors.forEach(selector => {
document.querySelectorAll(selector).forEach(attachListeners);
});
// Start observing the document body for changes
observer.observe(document.body, { childList: true, subtree: true });
console.log("[Autocomplete-Plus] Autocomplete initialized and observer started.");
}
app.registerExtension({
name: "comfyui-autocomplete-plus",
async setup() {
console.log("[Autocomplete-Plus] Starting setup...");
// Load data asynchronously from API endpoints
Promise.all([loadTags(), loadCooccurrence()])
.then(() => {
if (tagsLoaded) { // Only need tags for autocomplete
console.log("[Autocomplete-Plus] Tag data loaded successfully via API.");
initializeAutocomplete(); // Initialize after tags are loaded
} else {
console.error("[Autocomplete-Plus] Failed to load tag data via API. Autocomplete disabled.");
}
if (!cooccurrenceLoaded) {
console.warn("[Autocomplete-Plus] Failed to load cooccurrence data via API. Similar tags disabled.");
}
})
.catch(error => {
console.error("[Autocomplete-Plus] Error during API data loading:", error);
});
console.log("[Autocomplete-Plus] Setup complete (API data loading initiated)!");
},
});