├── .gitmodules ├── LICENSE ├── README.md ├── __init__.py ├── example └── illustrate.gif ├── requirements.txt └── script └── index.js /.gitmodules: -------------------------------------------------------------------------------- 1 | [submodule "danbooru_tag_db"] 2 | path = danbooru_tag_db 3 | url = https://github.com/waterminer/danbooru_tag_db/ 4 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2024 WaterMiner 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Feature 2 | 3 | ![illustrate](https://github.com/waterminer/ComfyUI-tagcomplete/blob/master/example/illustrate.gif) 4 | 5 | # Install 6 | 7 | 1. clone this project to `ComfyUI/custom_nodes` 8 | 9 | ```bash 10 | ~/ComfyUI/custom_nodes$ git clone https://github.com/waterminer/ComfyUI-tagcomplete 11 | ``` 12 | 13 | 2. init submodule 14 | 15 | ```bash 16 | ~/ComfyUI/custom_nodes/ComfyUI-tagcomplete$ git submodule init 17 | ``` 18 | 19 | 3. install requirements 20 | 21 | ```bash 22 | ~/ComfyUI/custom_nodes/ComfyUI-tagcomplete$ pip install -r ./requirements.txt 23 | ``` 24 | -------------------------------------------------------------------------------- /__init__.py: -------------------------------------------------------------------------------- 1 | import sys,os,aiosqlite,json 2 | 3 | from server import PromptServer 4 | from aiohttp import web 5 | 6 | ROOT_PATH = os.path.dirname(__file__) 7 | sys.path.append(ROOT_PATH) 8 | 9 | WEB_DIRECTORY = "./script" 10 | 11 | __all__ =["WEB_DIRECTORY"] 12 | 13 | routes = PromptServer.instance.routes 14 | @routes.get("/water_miner/database") 15 | async def select_database(request): 16 | input = request.query.get('name') 17 | SQL=f''' 18 | SELECT 19 | t.id, t.name, t.category, t.post_count, 20 | 21 | CASE 22 | WHEN t.name = '{input}' THEN 1000 * t.post_count 23 | WHEN t.name LIKE '{input}%' AND t.category = 0 THEN 100 * t.post_count 24 | WHEN t.name LIKE '%{input}%' AND t.category = 0 THEN 50 * t.post_count 25 | WHEN t.name LIKE '{input}%' AND t.category IN (1, 3, 4) THEN 10 * t.post_count 26 | WHEN t.name LIKE '%{input}%' AND t.category IN (1, 3, 4) THEN 1 * t.post_count 27 | WHEN t.name LIKE '{input}%' AND t.category = 5 THEN 1 * t.post_count 28 | WHEN t.name LIKE '%{input}%' AND t.category = 5 THEN 1 * t.post_count 29 | WHEN alias.name = '{input}' THEN 100 * t.post_count 30 | WHEN alias.name LIKE '{input}%' THEN 10 * t.post_count 31 | WHEN alias.name LIKE '%{input}%' THEN 1 * t.post_count 32 | END AS weight, 33 | 34 | CASE 35 | WHEN t.name NOT LIKE '%{input}%' AND alias.name IS NOT NULL THEN alias.name 36 | ELSE NULL 37 | END AS alias 38 | 39 | FROM 40 | Tag t 41 | LEFT JOIN 42 | TagAlias alias ON t.id = alias.consequent_id 43 | AND (alias.name = '{input}' 44 | OR alias.name LIKE '{input}%' 45 | OR alias.name LIKE '%{input}%') 46 | 47 | WHERE 48 | t.name LIKE '%{input}%' 49 | OR alias.consequent_id IS NOT NULL 50 | 51 | ORDER BY 52 | weight DESC 53 | LIMIT 10; 54 | 55 | ''' 56 | async with aiosqlite.connect(f"{ROOT_PATH}/danbooru_tag_db/database.db") as db: 57 | db.row_factory = aiosqlite.Row 58 | cur = await db.execute(SQL) 59 | rows = await cur.fetchall() 60 | result = [dict(row) for row in rows] 61 | body = json.dumps(result) 62 | return web.Response(body=body,status=200) 63 | 64 | NODE_CLASS_MAPPINGS = {} 65 | -------------------------------------------------------------------------------- /example/illustrate.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/waterminer/ComfyUI-tagcomplete/15d50d92c907385172cfe1090c4fe7ae9f4fce07/example/illustrate.gif -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | aiosqlite -------------------------------------------------------------------------------- /script/index.js: -------------------------------------------------------------------------------- 1 | import { app } from "../../scripts/app.js"; 2 | import { api } from "../../scripts/api.js"; 3 | 4 | function extensionAddMultilineWidget(node, name, options, app) { 5 | const emptyTooltipList = document.createElement('div'); 6 | Object.assign(emptyTooltipList, { length: 0 }); 7 | const container = document.createElement('div'); 8 | const tooltip = document.createElement('div'); 9 | const inputBox = document.createElement('textarea'); 10 | let selectTag = ""; 11 | inputBox.className = 'comfy-multiline-input'; 12 | inputBox.value = options.defaultVal; 13 | inputBox.placeholder = options.placeholder || name; 14 | let tooltipListLen = 0 15 | if (app.vueAppReady) { 16 | api.fetchApi("/settings/Comfy.TextareaWidget.Spellcheck").then(res => { 17 | return res.json(); 18 | }).then(json => { 19 | inputBox.spellcheck = json; 20 | }); 21 | } 22 | Object.assign(inputBox.style, { 23 | width: "100%", 24 | height: "100%" 25 | }); 26 | Object.assign(tooltip.style, { 27 | position: "absolute", 28 | display: "none", 29 | color: "white", 30 | border: "1px solid #ccc", 31 | padding: "0px", 32 | zIndex: "1000", 33 | borderRadius: "4px" 34 | }); 35 | 36 | function updateSelectWold(cursorPosition) { 37 | const text = inputBox.value; 38 | const front = text.slice(0, cursorPosition); 39 | selectTag = front.split(",").at(-1).trim().replace(/\s/g, '_'); 40 | if (tooltipListLen !== 0) 41 | tooltip.style.display = 'block'; 42 | else 43 | tooltip.style.display = 'none'; 44 | } 45 | function insertText(tagData) { 46 | let front = inputBox.value.slice(0, inputBox.selectionStart); 47 | let frontArr = front.split(','); 48 | frontArr.pop(); 49 | front = frontArr.join(); 50 | if (frontArr.length !== 0) { 51 | front = front + ','; 52 | } 53 | const behind = inputBox.value.slice(inputBox.selectionStart) 54 | inputBox.value = `${front}${tagData.name.replace(/_/g, ' ')},${behind}` 55 | tooltip.style.display = 'none' 56 | tooltip.replaceChildren(emptyTooltipList) 57 | } 58 | function createTooltipList(json) { 59 | const tooltipList = document.createElement('ul'); 60 | const first_chose = json.at(0); 61 | inputBox.addEventListener('keydown', (event) => { 62 | if (event.key === 'Tab') { 63 | event.preventDefault(); 64 | insertText(first_chose) 65 | } 66 | }, { once: true }) 67 | Object.assign(tooltipList, { 68 | length: 0 69 | }) 70 | Object.assign(tooltipList.style, { 71 | listStyleType: "none", 72 | padding: "0px", 73 | margin: "0px" 74 | }) 75 | for (const tagData of json) { 76 | const li = document.createElement('li') 77 | const element = document.createElement('button'); 78 | tooltipList.length += 1 79 | let textColor = "#ca0000" 80 | switch (tagData.category) { 81 | case 0: 82 | textColor = '#009be6' 83 | break; 84 | case 1: 85 | textColor = '#ff8a8b' 86 | break; 87 | case 3: 88 | textColor = '#a800aa' 89 | break; 90 | case 4: 91 | textColor = '#00ab2c' 92 | break; 93 | case 5: 94 | textColor = '#fd9200' 95 | break; 96 | default: 97 | break; 98 | } 99 | Object.assign(element.style, { 100 | width: "100%", 101 | color: textColor, 102 | textAlign: 'left' 103 | }) 104 | if(tagData.alias){ 105 | element.innerText = `${tagData.alias}->${tagData.name}`; 106 | }else{ 107 | element.innerText = `${tagData.name}`; 108 | } 109 | element.addEventListener('click', (event) => { 110 | insertText(tagData) 111 | }) 112 | li.append(element); 113 | tooltipList.append(li); 114 | 115 | } 116 | return tooltipList 117 | } 118 | 119 | async function updateTooltip(tag) { 120 | const params = new URLSearchParams({ name: tag }) 121 | const result = await api.fetchApi( 122 | `/water_miner/database?${params.toString()}`, 123 | { method: "GET" } 124 | ) 125 | if (result.status !== 200) 126 | throw new Error(`Fetch Failed,code:${result.status}`); 127 | const json = await result.json(); 128 | const TooltipList = createTooltipList(json) 129 | tooltipListLen = TooltipList.length; 130 | if (tooltipListLen === 0) { 131 | tooltip.style.display = 'none'; 132 | } 133 | tooltip.replaceChildren(TooltipList); 134 | } 135 | 136 | function getTextWidth(text, font) { 137 | const canvas = document.createElement('canvas'); 138 | const context = canvas.getContext('2d'); 139 | context.font = font; 140 | return context.measureText(text).width; 141 | } 142 | 143 | inputBox.addEventListener('focus', () => { 144 | if (tooltipListLen === 0) 145 | tooltip.style.display = 'none'; 146 | else 147 | tooltip.style.display = 'block'; 148 | }); 149 | 150 | inputBox.addEventListener('blur', () => { 151 | setTimeout(() => { 152 | if (!tooltip.contains(document.activeElement)) { 153 | tooltip.style.display = 'none'; 154 | } 155 | }, 100); 156 | }); 157 | 158 | inputBox.addEventListener('input', async () => { 159 | const cursorPosition = inputBox.selectionStart; 160 | const textBeforeCursor = inputBox.value.slice(0, cursorPosition); 161 | const textWidth = getTextWidth(textBeforeCursor, window.getComputedStyle(inputBox).font); 162 | updateSelectWold(cursorPosition); 163 | tooltip.style.left = `${inputBox.offsetLeft + textWidth}px`; 164 | tooltip.style.top = `${inputBox.offsetTop + 14}px`; 165 | await updateTooltip(selectTag); 166 | }); 167 | 168 | container.append(inputBox, tooltip) 169 | 170 | const widget = node.addDOMWidget(name, 'extensioncustomtext', container, { 171 | getValue() { 172 | return inputBox.value; 173 | }, 174 | setValue(v) { 175 | inputBox.value = v; 176 | } 177 | }) 178 | 179 | widget.inputEl = container; 180 | inputBox.addEventListener('input', () => { 181 | widget.callback?.(widget.value); 182 | }) 183 | return { minWidth: 400, minHeight: 200, widget } 184 | } 185 | 186 | function hijackString(app) { 187 | const o_WidgetSTRING = app.widgets.STRING; 188 | app.widgets.STRING = function (node, inputName, inputData, app) { 189 | let res = undefined 190 | const multiline = !!inputData[1].multiline; 191 | const defaultVal = inputData[1].default || ''; 192 | if (multiline) { 193 | res = extensionAddMultilineWidget(node, inputName, { defaultVal, ...inputData[1] }, app); 194 | if (inputData[1].dynamicPrompts != undefined) { 195 | res.widget.dynamicPrompts = inputData[1].dynamicPrompts 196 | } 197 | } else { 198 | res = o_WidgetSTRING(node, inputName, inputData, app); 199 | } 200 | return res; 201 | } 202 | } 203 | 204 | const extension = { 205 | name: "WaterMiner.TagComplete", 206 | async beforeRegisterNodeDef(nodeType, nodeData, app) { 207 | if (!g_execute) { 208 | hijackString(app) 209 | g_execute = true 210 | } 211 | } 212 | } 213 | let g_execute = false 214 | app.registerExtension(extension); --------------------------------------------------------------------------------