diff --git a/cli/fetch_tags.py b/cli/fetch_tags.py index 5de5edfd1..1b048a581 100644 --- a/cli/fetch_tags.py +++ b/cli/fetch_tags.py @@ -71,7 +71,7 @@ MAX_RETRIES = 3 RETRY_BACKOFF = 5 # seconds, multiplied by attempt number -# ── Partial save/resume ── +# -- Partial save/resume -- def save_partial(path: str, page: int, tags: list): """Save fetch progress to a .partial file.""" @@ -105,7 +105,7 @@ def clear_partial(path: str): os.remove(path) -# ── HTTP retry helper ── +# -- HTTP retry helper -- def fetch_with_retry(session: requests.Session, url: str, params: dict | None = None, timeout: int = 30) -> requests.Response: """GET with retry and exponential backoff for transient errors.""" @@ -123,7 +123,7 @@ def fetch_with_retry(session: requests.Session, url: str, params: dict | None = raise -# ── Fetchers ── +# -- Fetchers -- def fetch_danbooru(min_count: int = 10, partial_path: str = "", **_kwargs) -> list: """Fetch tags from Danbooru API, paginated.""" diff --git a/javascript/autocomplete.js b/javascript/autocomplete.js new file mode 100644 index 000000000..ccf6a5ebc --- /dev/null +++ b/javascript/autocomplete.js @@ -0,0 +1,430 @@ +/* + * 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(); + } + }); +} diff --git a/javascript/startup.js b/javascript/startup.js index 5f6b3f926..a51a0272d 100644 --- a/javascript/startup.js +++ b/javascript/startup.js @@ -59,6 +59,7 @@ async function initStartup() { // optinally wait for modern ui if (window.waitForUiReady) await waitForUiReady(); + initAutocomplete(); monitorConnection(); removeSplash(); diff --git a/modules/api/api.py b/modules/api/api.py index 9e98be0bf..889ba6104 100644 --- a/modules/api/api.py +++ b/modules/api/api.py @@ -145,9 +145,9 @@ class Api: # hide trailing-slash duplicates from OpenAPI schema from fastapi.routing import APIRoute - paths = {r.path for r in self.app.routes if hasattr(r, 'path')} + route_paths = {r.path for r in self.app.routes if hasattr(r, 'path')} for route in self.app.routes: - if isinstance(route, APIRoute) and len(route.path) > 1 and route.path.endswith('/') and route.path[:-1] in paths: + if isinstance(route, APIRoute) and len(route.path) > 1 and route.path.endswith('/') and route.path[:-1] in route_paths: route.include_in_schema = False # upload api diff --git a/modules/api/autocomplete.py b/modules/api/autocomplete.py index a8dad5f92..cb90baee6 100644 --- a/modules/api/autocomplete.py +++ b/modules/api/autocomplete.py @@ -25,7 +25,7 @@ manifest_cache: dict = {} # {"data": [...], "fetched_at": float} def init(path: str) -> None: """Set the autocomplete directory path. Called once during API registration.""" - global autocomplete_dir # noqa: PLW0603 + global autocomplete_dir # pylint: disable=global-statement autocomplete_dir = path @@ -36,7 +36,18 @@ def get_cached(name: str) -> dict: path = os.path.join(autocomplete_dir, f"{name}.json") if not os.path.isfile(path): cache.pop(name, None) - raise HTTPException(status_code=404, detail=f"Not found: {name}") + # Auto-download from HF if available in manifest + try: + manifest = fetch_manifest_sync() + if any(e.get('name') == name for e in manifest): + log.info(f"Auto-downloading autocomplete: {name}") + download_sync(name) + else: + raise HTTPException(status_code=404, detail=f"Not found: {name}") + except HTTPException: + raise + except Exception as e: + raise HTTPException(status_code=404, detail=f"Not found: {name} ({e})") from e stat = os.stat(path) entry = cache.get(name) if entry and entry['mtime'] == stat.st_mtime: @@ -109,7 +120,7 @@ async def get_content(name: str) -> ItemAutocompleteContent: ) -# ── Remote management ── +# -- Remote management -- def fetch_manifest_sync() -> list[dict]: """Fetch manifest.json from HuggingFace, with caching.""" @@ -212,7 +223,7 @@ def download_sync(name: str) -> str: async def download(name: str): """Download a tag file from HuggingFace.""" - path = await asyncio.to_thread(download_sync, name) + await asyncio.to_thread(download_sync, name) entry = await asyncio.to_thread(get_cached, name) meta = entry['meta'] return ItemAutocomplete( diff --git a/modules/ui_definitions.py b/modules/ui_definitions.py index bc2cd3851..f563edd69 100644 --- a/modules/ui_definitions.py +++ b/modules/ui_definitions.py @@ -52,16 +52,25 @@ def get_openvino_device_list(): def list_autocomplete_names(): - """Return list of available tag autocomplete file names (JSON filenames without extension).""" + """Return list of available tag autocomplete file names from local files and HF manifest.""" from modules import shared, paths as paths_module + names = set() + # Local files autocomplete_dir = getattr(shared.opts, 'autocomplete_dir', None) or os.path.join(paths_module.models_path, 'autocomplete') - if not os.path.isdir(autocomplete_dir): - return [] - return sorted( - os.path.splitext(f)[0] - for f in os.listdir(autocomplete_dir) - if f.endswith('.json') and not f.startswith('.') and f != 'manifest.json' - ) + if os.path.isdir(autocomplete_dir): + for f in os.listdir(autocomplete_dir): + if f.endswith('.json') and not f.startswith('.') and f != 'manifest.json': + names.add(os.path.splitext(f)[0]) + # Remote manifest + try: + from modules.api.autocomplete import fetch_manifest_sync + for entry in fetch_manifest_sync(): + names.add(entry.get('name', '')) + except Exception as e: + from modules.logger import log + log.debug(f"Autocomplete manifest fetch skipped: {e}") + names.discard('') + return sorted(names) def create_settings(cmd_opts): @@ -653,7 +662,9 @@ def create_settings(cmd_opts): "wildcards_enabled": OptionInfo(True, "Enable file wildcards support"), "extra_networks_autocomplete_sep": OptionInfo("

Tag Autocomplete

", "", gr.HTML), - "autocomplete_enabled": OptionInfo([], "Enabled tag autocomplete files", gr.CheckboxGroup, lambda: {"choices": list_autocomplete_names()}), + "autocomplete_enabled": OptionInfo([], "Enabled tag autocomplete files", gr.Dropdown, lambda: {"multiselect": True, "choices": list_autocomplete_names()}), + "autocomplete_min_chars": OptionInfo(3, "Minimum characters before autocomplete triggers", gr.Slider, {"minimum": 2, "maximum": 6, "step": 1}), + "autocomplete_replace_underscores": OptionInfo(True, "Replace underscores with spaces in autocomplete results"), })) # --- Extensions ---