mirror of
https://github.com/vladmandic/automatic
synced 2026-08-29 08:31:00 +02:00
CalamitousFelicitousness
431 lines
14 KiB
JavaScript
431 lines
14 KiB
JavaScript
/*
|
|
* Tag autocomplete for SD.Next prompt textareas.
|
|
*
|
|
* Ported from Enso's CodeMirror-based autocomplete (autocomplete.ts).
|
|
* Uses binary search on sorted tag arrays for O(log n) prefix lookup,
|
|
* with substring fallback for 4+ char queries.
|
|
*/
|
|
|
|
// -- Category colors (unified 14-category scheme) --
|
|
|
|
const CATEGORY_COLORS = {
|
|
0: '#0075f8', // general
|
|
1: '#cc0000', // artist
|
|
2: '#ff4500', // studio
|
|
3: '#9900ff', // copyright
|
|
4: '#00ab2c', // character
|
|
5: '#ed5d1f', // species
|
|
6: '#8a66ff', // genre
|
|
7: '#00cccc', // medium
|
|
8: '#6b7280', // meta
|
|
9: '#228b22', // lore
|
|
10: '#e67e22', // lens
|
|
11: '#f1c40f', // lighting
|
|
12: '#1abc9c', // composition
|
|
13: '#e84393', // color
|
|
};
|
|
|
|
// -- Utilities (ported from Enso) --
|
|
|
|
/** Binary search for the first tag where tag.name >= prefix. */
|
|
function lowerBound(tags, prefix) {
|
|
let lo = 0;
|
|
let hi = tags.length;
|
|
while (lo < hi) {
|
|
const mid = (lo + hi) >>> 1;
|
|
if (tags[mid].name < prefix) lo = mid + 1;
|
|
else hi = mid;
|
|
}
|
|
return lo;
|
|
}
|
|
|
|
/** Format post count as abbreviated string. */
|
|
function formatCount(count) {
|
|
if (count >= 1_000_000) return `${(count / 1_000_000).toFixed(1)}M`;
|
|
if (count >= 1_000) return `${Math.round(count / 1_000)}k`;
|
|
return String(count);
|
|
}
|
|
|
|
// -- TagIndex --
|
|
|
|
class TagIndex {
|
|
constructor(data) {
|
|
this.categories = data.categories || {};
|
|
// Build sorted array of {name, category, count} from raw [name, catId, count] tuples
|
|
this.tags = data.tags.map(([name, category, count]) => ({
|
|
name: name.toLowerCase(),
|
|
display: name,
|
|
category,
|
|
count,
|
|
}));
|
|
this.tags.sort((a, b) => a.name.localeCompare(b.name));
|
|
}
|
|
|
|
/** Prefix search with binary search. Returns matches sorted by count descending. */
|
|
search(prefix, limit = 20) {
|
|
const query = prefix.toLowerCase().replace(/ /g, '_');
|
|
if (!query) return [];
|
|
const start = lowerBound(this.tags, query);
|
|
const matches = [];
|
|
for (let i = start; i < this.tags.length && matches.length < limit * 5; i++) {
|
|
if (!this.tags[i].name.startsWith(query)) break;
|
|
matches.push(this.tags[i]);
|
|
}
|
|
// Substring fallback for 4+ chars if prefix found nothing
|
|
if (matches.length === 0 && query.length >= 4) {
|
|
for (let i = 0; i < this.tags.length && matches.length < limit * 5; i++) {
|
|
if (this.tags[i].name.includes(query)) matches.push(this.tags[i]);
|
|
}
|
|
}
|
|
matches.sort((a, b) => b.count - a.count);
|
|
return matches.slice(0, limit);
|
|
}
|
|
}
|
|
|
|
// -- Engine --
|
|
|
|
const engine = {
|
|
indices: new Map(), // name -> TagIndex
|
|
categoryColors: { ...CATEGORY_COLORS },
|
|
|
|
async loadEnabled() {
|
|
const enabled = window.opts?.autocomplete_enabled || [];
|
|
if (!enabled.length) {
|
|
this.indices.clear();
|
|
return;
|
|
}
|
|
const toLoad = enabled.filter((n) => !this.indices.has(n));
|
|
const toRemove = [...this.indices.keys()].filter((n) => !enabled.includes(n));
|
|
toRemove.forEach((n) => this.indices.delete(n));
|
|
await Promise.all(toLoad.map(async (name) => {
|
|
try {
|
|
const resp = await fetch(`${window.api}/autocomplete/${name}`, { credentials: 'include' });
|
|
if (!resp.ok) throw new Error(`${resp.status}`);
|
|
const data = await resp.json();
|
|
this.indices.set(name, new TagIndex(data));
|
|
// Extract category colors from first loaded file
|
|
if (data.categories) {
|
|
Object.entries(data.categories).forEach(([id, cat]) => {
|
|
if (cat.color) this.categoryColors[id] = cat.color;
|
|
});
|
|
}
|
|
log('autocomplete', `loaded ${name}: ${data.tags?.length || 0} tags`);
|
|
} catch (e) {
|
|
log('autocomplete', `failed to load ${name}: ${e}`);
|
|
}
|
|
}));
|
|
},
|
|
|
|
searchAll(prefix, limit = 20) {
|
|
if (this.indices.size === 0) return [];
|
|
const all = [];
|
|
this.indices.forEach((index) => {
|
|
all.push(...index.search(prefix, limit));
|
|
});
|
|
// Deduplicate by name, keeping highest count
|
|
const seen = new Map();
|
|
all.forEach((tag) => {
|
|
const existing = seen.get(tag.name);
|
|
if (!existing || tag.count > existing.count) seen.set(tag.name, tag);
|
|
});
|
|
const results = [...seen.values()];
|
|
results.sort((a, b) => b.count - a.count);
|
|
return results.slice(0, limit);
|
|
},
|
|
};
|
|
|
|
// -- Textarea integration --
|
|
|
|
/** Extract the current word being typed at the cursor position. */
|
|
function getCurrentWord(textarea) {
|
|
const { value, selectionStart } = textarea;
|
|
if (selectionStart !== textarea.selectionEnd) return null; // has selection
|
|
// Scan backward from cursor to find word start
|
|
let start = selectionStart;
|
|
while (start > 0) {
|
|
const ch = value[start - 1];
|
|
if (ch === ',' || ch === '\n') break;
|
|
start--;
|
|
}
|
|
// Skip leading whitespace
|
|
while (start < selectionStart && value[start] === ' ') start++;
|
|
const word = value.slice(start, selectionStart);
|
|
if (!word) return null;
|
|
// Skip if inside angle brackets (LoRA/embedding syntax)
|
|
const before = value.slice(0, selectionStart);
|
|
const lastOpen = before.lastIndexOf('<');
|
|
const lastClose = before.lastIndexOf('>');
|
|
if (lastOpen > lastClose) return null;
|
|
// Skip if inside wildcard syntax
|
|
const wcBefore = before.slice(start);
|
|
if (wcBefore.startsWith('__') && !wcBefore.endsWith('__')) return null;
|
|
return { word, start, end: selectionStart };
|
|
}
|
|
|
|
/** Insert a tag at the current word position, replacing the typed prefix. */
|
|
function insertTag(textarea, tagName) {
|
|
const info = getCurrentWord(textarea);
|
|
if (!info) return;
|
|
const { value } = textarea;
|
|
const before = value.slice(0, info.start);
|
|
const after = value.slice(info.end);
|
|
// Build insertion: tag + separator
|
|
const needsCommaBefore = before.length > 0 && !before.trimEnd().endsWith(',') && before.trimEnd().length > 0;
|
|
const prefix = needsCommaBefore ? ', ' : '';
|
|
let suffix = ', ';
|
|
if (after.length > 0 && after.trimStart().startsWith(',')) suffix = ' ';
|
|
const insertion = `${prefix}${tagName}${suffix}`;
|
|
textarea.value = before.trimEnd() + (before.trimEnd().length > 0 ? ' ' : '') + insertion + after.trimStart();
|
|
// Position cursor after the inserted tag + separator
|
|
const cursorPos = before.trimEnd().length + (before.trimEnd().length > 0 ? 1 : 0) + insertion.length;
|
|
textarea.selectionStart = cursorPos;
|
|
textarea.selectionEnd = cursorPos;
|
|
// Sync with Gradio
|
|
if (typeof updateInput === 'function') updateInput(textarea);
|
|
}
|
|
|
|
// -- Dropdown --
|
|
|
|
const dropdown = {
|
|
el: null,
|
|
listEl: null,
|
|
selectedIndex: -1,
|
|
results: [],
|
|
textarea: null,
|
|
visible: false,
|
|
|
|
init() {
|
|
this.el = document.createElement('div');
|
|
this.el.className = 'autocompleteResults';
|
|
this.el.style.display = 'none';
|
|
this.listEl = document.createElement('ul');
|
|
this.listEl.className = 'autocompleteResultsList';
|
|
this.el.appendChild(this.listEl);
|
|
document.body.appendChild(this.el);
|
|
this.el.addEventListener('mousedown', (e) => e.preventDefault()); // prevent blur on click
|
|
this.el.addEventListener('click', (e) => {
|
|
const li = e.target.closest('li');
|
|
if (!li) return;
|
|
const idx = [...this.listEl.children].indexOf(li);
|
|
if (idx >= 0 && idx < this.results.length) {
|
|
this.selectedIndex = idx;
|
|
this.accept();
|
|
}
|
|
});
|
|
},
|
|
|
|
show(results, textarea) {
|
|
if (results.length === 0) { this.hide(); return; }
|
|
this.results = results;
|
|
this.textarea = textarea;
|
|
this.selectedIndex = -1;
|
|
this.render();
|
|
this.position();
|
|
this.el.style.display = '';
|
|
this.visible = true;
|
|
},
|
|
|
|
hide() {
|
|
this.el.style.display = 'none';
|
|
this.visible = false;
|
|
this.results = [];
|
|
this.selectedIndex = -1;
|
|
},
|
|
|
|
render() {
|
|
const replaceUnderscores = window.opts?.autocomplete_replace_underscores ?? true;
|
|
this.listEl.innerHTML = '';
|
|
this.results.forEach((tag, i) => {
|
|
const li = document.createElement('li');
|
|
if (i === this.selectedIndex) li.classList.add('selected');
|
|
const dot = document.createElement('span');
|
|
dot.className = 'autocomplete-category';
|
|
dot.style.color = engine.categoryColors[tag.category] || '#888';
|
|
dot.textContent = '\u25CF';
|
|
const name = document.createElement('span');
|
|
name.className = 'autocomplete-tag';
|
|
name.textContent = replaceUnderscores ? tag.display.replace(/_/g, ' ') : tag.display;
|
|
const count = document.createElement('span');
|
|
count.className = 'autocomplete-count';
|
|
count.textContent = tag.count > 0 ? formatCount(tag.count) : '';
|
|
li.append(dot, name, count);
|
|
li.addEventListener('mouseenter', () => {
|
|
this.selectedIndex = i;
|
|
this.updateSelection();
|
|
});
|
|
this.listEl.appendChild(li);
|
|
});
|
|
},
|
|
|
|
position() {
|
|
if (!this.textarea) return;
|
|
const rect = this.textarea.getBoundingClientRect();
|
|
const spaceBelow = window.innerHeight - rect.bottom;
|
|
const dropHeight = Math.min(this.el.scrollHeight, 300);
|
|
if (spaceBelow >= dropHeight || spaceBelow >= rect.top) {
|
|
this.el.style.top = `${rect.bottom + 2}px`;
|
|
} else {
|
|
this.el.style.top = `${rect.top - dropHeight - 2}px`;
|
|
}
|
|
this.el.style.left = `${rect.left}px`;
|
|
this.el.style.width = `${rect.width}px`;
|
|
},
|
|
|
|
updateSelection() {
|
|
[...this.listEl.children].forEach((li, i) => {
|
|
li.classList.toggle('selected', i === this.selectedIndex);
|
|
});
|
|
const selected = this.listEl.children[this.selectedIndex];
|
|
if (selected) selected.scrollIntoView({ block: 'nearest' });
|
|
},
|
|
|
|
navigate(dir) {
|
|
if (this.results.length === 0) return;
|
|
if (this.selectedIndex === -1) {
|
|
this.selectedIndex = dir > 0 ? 0 : this.results.length - 1;
|
|
} else {
|
|
this.selectedIndex = (this.selectedIndex + dir + this.results.length) % this.results.length;
|
|
}
|
|
this.updateSelection();
|
|
},
|
|
|
|
accept() {
|
|
if (this.selectedIndex < 0 || this.selectedIndex >= this.results.length) {
|
|
// Tab with no selection: select first
|
|
if (this.results.length > 0) {
|
|
this.selectedIndex = 0;
|
|
this.updateSelection();
|
|
}
|
|
return;
|
|
}
|
|
const tag = this.results[this.selectedIndex];
|
|
if (this.textarea) insertTag(this.textarea, tag.display);
|
|
this.hide();
|
|
},
|
|
};
|
|
|
|
// -- Event handlers --
|
|
|
|
let debounceTimer = null;
|
|
|
|
function onInput(textarea) {
|
|
const minChars = window.opts?.autocomplete_min_chars ?? 3;
|
|
const info = getCurrentWord(textarea);
|
|
if (!info || info.word.length < minChars) {
|
|
dropdown.hide();
|
|
return;
|
|
}
|
|
clearTimeout(debounceTimer);
|
|
debounceTimer = setTimeout(() => {
|
|
const results = engine.searchAll(info.word);
|
|
dropdown.show(results, textarea);
|
|
}, 150);
|
|
}
|
|
|
|
function onKeyDown(e) {
|
|
if (!dropdown.visible) return;
|
|
switch (e.key) {
|
|
case 'ArrowDown':
|
|
e.preventDefault();
|
|
e.stopPropagation();
|
|
dropdown.navigate(1);
|
|
break;
|
|
case 'ArrowUp':
|
|
e.preventDefault();
|
|
e.stopPropagation();
|
|
dropdown.navigate(-1);
|
|
break;
|
|
case 'Enter':
|
|
if (dropdown.selectedIndex >= 0) {
|
|
e.preventDefault();
|
|
e.stopPropagation();
|
|
dropdown.accept();
|
|
}
|
|
break;
|
|
case 'Tab':
|
|
e.preventDefault();
|
|
e.stopPropagation();
|
|
dropdown.accept();
|
|
break;
|
|
case 'Escape':
|
|
e.preventDefault();
|
|
e.stopPropagation();
|
|
dropdown.hide();
|
|
break;
|
|
default:
|
|
break;
|
|
}
|
|
}
|
|
|
|
/** Attach autocomplete to a single textarea. */
|
|
function attachAutocomplete(textarea) {
|
|
textarea.addEventListener('input', () => onInput(textarea));
|
|
textarea.addEventListener('keydown', onKeyDown);
|
|
textarea.addEventListener('focusout', () => {
|
|
setTimeout(() => dropdown.hide(), 200);
|
|
});
|
|
}
|
|
|
|
// -- Prompt textarea IDs --
|
|
|
|
const PROMPT_IDS = [
|
|
'txt2img_prompt', 'txt2img_neg_prompt',
|
|
'img2img_prompt', 'img2img_neg_prompt',
|
|
'control_prompt', 'control_neg_prompt',
|
|
'video_prompt', 'video_neg_prompt',
|
|
];
|
|
|
|
// -- Initialization --
|
|
|
|
async function initAutocomplete() {
|
|
const enabled = window.opts?.autocomplete_enabled || [];
|
|
if (!enabled.length) {
|
|
log('autocomplete', 'no dictionaries enabled');
|
|
return;
|
|
}
|
|
log('autocomplete', `init: ${enabled.join(', ')}`);
|
|
// Inject styles (CSS files in javascript/ are not auto-loaded)
|
|
const style = document.createElement('style');
|
|
style.textContent = [
|
|
'.autocompleteResults { position: fixed; z-index: 9999; max-height: 300px; overflow-y: auto;',
|
|
' background: var(--sd-main-background-color, var(--background-fill-primary, #1f2937));',
|
|
' border: 1px solid var(--sd-input-border-color, var(--border-color-primary, #374151));',
|
|
' border-radius: var(--sd-border-radius, 6px); box-shadow: 0 4px 12px rgba(0,0,0,0.3);',
|
|
' font-size: 13px; scrollbar-width: thin; }',
|
|
'.autocompleteResultsList { list-style: none; margin: 0; padding: 4px 0; }',
|
|
'.autocompleteResultsList > li { display: flex; align-items: center; padding: 4px 10px; cursor: pointer; gap: 8px; line-height: 1.4; }',
|
|
'.autocompleteResultsList > li:hover { background: var(--sd-panel-background-color, var(--input-background-fill-focus, #374151)); }',
|
|
'.autocompleteResultsList > li.selected { background: var(--sd-main-accent-color, var(--button-primary-background-fill, #4b5563)); }',
|
|
'.autocomplete-category { font-size: 10px; flex-shrink: 0; width: 10px; text-align: center; }',
|
|
'.autocomplete-tag { flex: 1; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; }',
|
|
'.autocomplete-count { font-size: 0.8em; opacity: 0.5; flex-shrink: 0; font-variant-numeric: tabular-nums; }',
|
|
].join('\n');
|
|
document.head.appendChild(style);
|
|
dropdown.init();
|
|
await engine.loadEnabled();
|
|
if (engine.indices.size === 0) {
|
|
log('autocomplete', 'no dictionaries loaded');
|
|
return;
|
|
}
|
|
// Attach to all prompt textareas
|
|
let attached = 0;
|
|
PROMPT_IDS.forEach((id) => {
|
|
const textarea = gradioApp().querySelector(`#${id} > label > textarea`);
|
|
if (textarea) {
|
|
attachAutocomplete(textarea);
|
|
attached++;
|
|
}
|
|
});
|
|
log('autocomplete', `attached to ${attached} textareas, ${engine.indices.size} dictionaries`);
|
|
// Reload when settings change
|
|
onOptionsChanged(async () => {
|
|
const newEnabled = window.opts?.autocomplete_enabled || [];
|
|
const currentKeys = [...engine.indices.keys()].sort().join(',');
|
|
const newKeys = [...newEnabled].sort().join(',');
|
|
if (currentKeys !== newKeys) {
|
|
log('autocomplete', `reloading: ${newEnabled.join(', ')}`);
|
|
await engine.loadEnabled();
|
|
}
|
|
});
|
|
}
|