22 Commits
Author SHA1 Message Date
imi 070723c767 clarify forge -> comfy scheduling with 1+ step samplers 2026-07-28 02:36:39 +02:00
imi 635b4bbdbe last resolved prompt preview stays even when editing current prompt 2026-07-28 01:52:01 +02:00
imi 69d03a8763 readme 2026-07-25 05:10:07 +02:00
imi 8b5269bea4 more fixes, autocomplete/wildcard optimization 2026-07-25 05:01:03 +02:00
imi 9e67b317d2 Improve wildcard autocomplete performance and resolved prompt preview 2026-07-18 05:08:32 +02:00
imi 3db87dc406 Improve autocomplete, wildcard handling, and prompt previews
- Add autocomplete caching, tag prefix indexing, and stronger usage ranking
- Add bundled merged Danbooru/e621 tag data
- Estimate max wildcard and dynamic-prompt token counts
- Store original/resolved prompt metadata and restore resolved previews
- Insert short wildcard leaf names when unambiguous
- Harden prompt preview toggles and textarea rebinding
- Expand tests for autocomplete, metadata, tokenization, and wildcards
2026-07-18 00:32:40 +02:00
imi 98ad657586 Fix image path separators in README.md 2026-04-22 01:34:23 +02:00
imi 55339cede5 assets 2026-04-22 01:32:05 +02:00
imi 620f422175 changelog edit 2026-04-22 01:13:07 +02:00
imi 639c67ba1a feat: add wildcard expansion and improve autocomplete UX
- add A1111-style __wildcard__ expansion for prompt and negative nodes
- support nested wildcards and dynamic prompt syntax inside wildcard files
- add wildcard autocomplete with folder/file suggestions and content drill-down
- improve wildcard popup rendering, limits, and keep-open behavior
- normalize spacing after optional wildcard/dynamic expansions
- make autocomplete frequency ranking relevance-aware with recency weighting
- fix wildcard usage stars so selection history affects ordering
- update README, TODO, and CHANGELOG to document the new behavior
2026-04-22 00:53:38 +02:00
imi d7d12aa523 Revise CHANGELOG for upcoming release
Updated changelog to reflect unreleased features and fixes.
2026-04-03 18:11:08 +02:00
imi 84c8c03523 js hardening 2026-04-03 03:16:25 +02:00
imi 4d6257bd07 Add DeepWiki badge to README
Added a badge linking to DeepWiki for ComfyUI-A1111-cond.
2026-02-03 00:40:41 +01:00
imi c8961a3b55 workflows 2026-02-02 23:42:37 +01:00
imi 68f5fdc75f pyproject 2026-02-02 23:21:50 +01:00
imi 552946ff73 v1.0.0 2026-02-02 23:15:23 +01:00
imi c0187b4d75 Merge branch 'main' of https://github.com/Enferlain/ComfyUI-A1111-cond 2026-02-02 22:11:33 +01:00
imi 6c8c50975d readme 2026-02-02 22:02:40 +01:00
imi ab3d80d3e3 Update README with new alternation syntax examples 2026-02-02 21:45:01 +01:00
imi 1cbfe5fe06 Update images in README with new dimensions and sources 2026-02-02 21:40:58 +01:00
imi 8817cabfe7 readme 2026-02-02 21:36:20 +01:00
imi d59cdd9680 Merge pull request #5 from Enferlain/builtin-hook
removed model and step inputs from the nodes
added new alternation scheduling syntax
added token counter and sequence end markers
added tipo preview for when tipo input is received so you can see the original
added tag auto completion
2026-02-02 19:27:01 +01:00
35 changed files with 4203 additions and 988 deletions
+28
View File
@@ -0,0 +1,28 @@
name: Publish to Comfy registry
on:
workflow_dispatch:
push:
branches:
- main
- master
paths:
- "pyproject.toml"
permissions:
contents: read
issues: write
jobs:
publish-node:
name: Publish Custom Node to registry
runs-on: ubuntu-latest
if: ${{ github.repository_owner == 'enferlain' }}
steps:
- name: Check out code
uses: actions/checkout@v4
with:
submodules: true
- name: Publish Custom Node
uses: Comfy-Org/publish-node-action@v1
with:
personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}
+3 -1
View File
@@ -12,4 +12,6 @@ A1111_PARITY_FINDINGS.md
MIGRATION_GUIDE.md
MODEL_WRAPPER_ANALYSIS.md
REFACTOR_SUMMARY.md
*.zip
*.zip
data/wildcards/*
!data/wildcards/readme.md
-146
View File
@@ -1,146 +0,0 @@
# Autocomplete Features Summary
## ✅ Implemented Features
### Phase 1: Core Tag Completion (Complete)
- ✅ CSV parser for Danbooru/e621 format
- ✅ In-memory tag database (~140k tags)
- ✅ Fast prefix search with scoring
- ✅ REST API endpoint (`/a1111_prompt/autocomplete`)
- ✅ Lazy loading (loads on first search)
- ✅ Real-time popup UI
- ✅ Keyboard navigation (↑/↓/Tab/Enter/Escape)
- ✅ Mouse click selection
- ✅ Color-coded tags by type
- ✅ Post count display
- ✅ Smart tag insertion with comma handling
- ✅ Cursor position tracking
- ✅ 100ms debouncing for snappy response
### Phase 2: Alias Support (Complete)
- ✅ Parse aliases from CSV
- ✅ Search by alias names
- ✅ Display alias mapping (`sole_female → 1girl`)
- ✅ Insert canonical tag, not alias
### Phase 3: Advanced Features (Partial)
- ✅ **Frequency Sorting**: Track and prioritize frequently used tags
- localStorage persistence
- Logarithmic boost algorithm
- Visual ★ indicator with usage count
- Console utilities for management
- ✅ **Extra Files**: Load multiple tag sources
- Auto-load `extra-quality-tags.csv`
- Merge results from multiple files
- Duplicate prevention
- ⏳ **Chants**: Prompt presets (not yet implemented)
- ⏳ **Wiki Links**: Tag documentation links (not yet implemented)
- ⏳ **Configuration UI**: Settings panel (not yet implemented)
## 🎯 How It Works
### Frequency Sorting Algorithm
1. **Track Usage**: Every time you select a tag, its usage count increments
2. **Calculate Boost**: Uses logarithmic scaling: `log2(frequency + 1)`
- 1 use = 1.0 boost
- 3 uses = 2.0 boost
- 7 uses = 3.0 boost
- 15 uses = 4.0 boost
3. **Sort Results**: Frequently used tags appear first if boost difference > 0.5
4. **Visual Feedback**: ★ badge shows usage count in popup
### Multiple Tag Sources
The system can load and merge multiple CSV files:
1. Main database (e.g., `danbooru.csv`)
2. Extra files (e.g., `extra-quality-tags.csv`)
3. Custom user files in `data/tags/`
Tags are deduplicated by name, with first occurrence taking precedence.
## 🔧 Console Utilities
Open browser console and use these commands:
```javascript
// View your most used tags
window.A1111Autocomplete.getStats()
// Returns: { totalTags, totalUses, topTags: [{tag, count}, ...] }
// Reset all frequency data
window.A1111Autocomplete.resetFrequency()
// Export frequency data for backup
window.A1111Autocomplete.exportFrequency()
// Returns: JSON string of all tag frequencies
```
## 📊 Example Usage Stats
```javascript
> window.A1111Autocomplete.getStats()
{
totalTags: 45,
totalUses: 127,
topTags: [
{ tag: "1girl", count: 23 },
{ tag: "solo", count: 18 },
{ tag: "masterpiece", count: 15 },
{ tag: "best quality", count: 15 },
{ tag: "highres", count: 12 },
...
]
}
```
## 🎨 Visual Indicators
### Tag Colors (by type)
- **Light Blue** (#87ceeb) - General tags
- **Red** (#cd5c5c) - Artist tags
- **Violet** (#ee82ee) - Copyright tags
- **Green** (#90ee90) - Character tags
- **Orange** (#ffa500) - Meta tags
### Frequency Badge
- **Blue badge** (★3) - Shows usage count
- Appears next to frequently used tags
- Hover to see "Used X times"
## 🚀 Performance
- **Database Load**: ~140k tags in <1 second
- **Search Speed**: <10ms for typical queries
- **Debounce**: 100ms for responsive feel
- **Memory**: ~20MB for full Danbooru database
- **Storage**: localStorage for frequency data (~5-10KB)
## 📝 Next Steps
### Remaining Phase 3 Features
1. **Chants System**: Quick prompt presets with `@@` trigger
2. **Wiki Links**: Optional `?` button for tag documentation
3. **Configuration UI**: Settings panel for customization
### Phase 4 (Blocked)
- Wildcard completion requires parser support first
## 🐛 Known Limitations
1. **No translation support yet**: Only English tags
2. **No configuration UI**: Settings are hardcoded
3. **No chants/presets**: Manual typing only
4. **localStorage only**: No cloud sync or export UI
## 💡 Tips
1. **Use frequently**: The more you use autocomplete, the better it learns your preferences
2. **Quality tags**: Type "quality", "best", "masterpiece" to see auto-loaded quality tags
3. **Aliases work**: Try typing common aliases like "sole_female" or "1girls"
4. **Fast typing**: 100ms debounce means you can type quickly without lag
5. **Backup data**: Use `exportFrequency()` to backup your usage statistics
---
_Last updated: 2026-01-22_
-326
View File
@@ -1,326 +0,0 @@
# Tag Autocomplete Implementation Plan
This document outlines the work required to implement A1111-style tag autocomplete for the A1111 Prompt Node in ComfyUI, based on analysis of the [a1111-sd-webui-tagcomplete](https://github.com/DominikDoom/a1111-sd-webui-tagcomplete) extension.
## Executive Summary
The A1111 tag autocomplete is a mature, feature-rich extension (~68KB main JS, ~43KB Python backend). A full port would be substantial work. This plan proposes a **phased approach** starting with core functionality and progressively adding features.
---
## Phase 1: Core Tag Completion (MVP)
**Goal:** Basic tag autocomplete with Danbooru/e621 tags
**Estimated Effort:** 3-5 days
### 1.1 Backend: Tag Data Loading
**File:** `api/autocomplete.py`
- [ ] **CSV Parser**: Parse tag CSV files with format `name,type,postCount,"aliases"`
- [ ] **Tag Types** (for color coding):
- `0` = General (blue)
- `1` = Artist (red)
- `3` = Copyright (violet)
- `4` = Character (green)
- `5` = Meta (orange)
- [ ] **In-memory Tag Index**: Load tags into searchable data structure
- [ ] **Search API Endpoint**: `POST /a1111_prompt/autocomplete`
```python
# Request: {"query": "1gir", "limit": 20}
# Response: {"results": [{"name": "1girl", "type": 0, "count": 4114588, "aliases": ["1girls", "sole_female"]}]}
```
- [ ] **Lazy Loading**: Only load tags when first search is triggered (not on startup)
### 1.2 Frontend: Popup UI
**File:** `js/autocomplete.js`
- [ ] **Trigger Detection**: Detect when user is typing a tag (after comma, start of text, after `(` etc.)
- [ ] **Current Word Extraction**: Get the word being typed at cursor position
- [ ] **Popup Creation**: Create floating results div near cursor
```javascript
// Use getCaretCoordinates() to position popup at cursor
```
- [ ] **Results Display**: Show tag name, count, and category color
- [ ] **Keyboard Navigation**:
- `↑/↓` - Move selection
- `Tab/Enter` - Accept selection
- `Escape` - Close popup
- [ ] **Click Selection**: Click on result to insert
- [ ] **Debouncing**: 100-300ms delay to prevent excessive API calls
### 1.3 Tag Insertion Logic
- [ ] **Replace current word** with selected tag
- [ ] **Underscore to space**: Convert `1_girl` → `1 girl` (configurable)
- [ ] **Escape parentheses**: `name_(artist)` → `name_\(artist\)` (configurable)
- [ ] **Append comma + space** after insertion (configurable)
### 1.4 Data Files
**Directory:** `data/tags/`
- [ ] **Bundle default tag files**:
- `danbooru.csv` (~4MB, top 100k tags by post count)
- `e621.csv` (~3MB, alternative for e621-trained models)
- [ ] **Document CSV format** for users adding custom tags
---
## Phase 2: Alias & Translation Support ✅
**Goal:** Search by aliases and translations
**Status:** COMPLETED
### 2.1 Alias Search ✅
- [x] **Parse aliases from CSV** (4th column, comma-separated in quotes)
- [x] **Include aliases in search**: Match `sole_female` when searching for `1girl`
- [x] **Display alias mapping**: Show `sole_female → 1girl`
- [x] **Insert canonical tag**, not alias
### 2.2 Translation Support
- [ ] **Translation file format**: CSV with `<english_tag>,<translated_tag>`
- [ ] **Search by translation**: Type in native language, find English tag
- [ ] **Show translation in results**: `1girl (一个女孩)`
---
## Phase 3: Advanced Features
**Goal:** Quality-of-life improvements
**Estimated Effort:** 3-5 days
### 3.1 Frequency Sorting (High Priority) ✅
- [x] **Track tag usage** locally (localStorage)
- [x] **Sort frequently used tags higher** in results
- [x] **Logarithmic boost function**: Prevents over-boosting common tags
- [x] **Visual indicator**: ★ badge shows usage count
- [x] **Utility functions**: Reset, export, and view stats via console
- [x] **Persistent storage**: Survives browser restarts
### 3.2 Extra File (Custom Tags) ✅
- [x] **Support multiple tag files**: Load main + extra files
- [x] **Auto-load quality tags**: `extra-quality-tags.csv` loaded by default
- [x] **Merge results**: Combines tags from all sources
- [x] **Duplicate prevention**: Same tag won't appear twice
### 3.3 Chants (Prompt Presets)
- [ ] **JSON format**:
```json
{
"name": "HighQuality",
"terms": "Best,Quality",
"content": "(masterpiece, best quality, highres)",
"color": 5
}
```
- [ ] **Trigger character** (e.g., `@@` or custom): Show available chants
- [ ] **Insert full preset** on selection
### 3.4 Wiki Links
- [ ] **Optional `?` button** next to tags
- [ ] **Link to Danbooru/e621 wiki** for tag examples
### 3.5 Configuration UI
- [ ] **Settings panel** in ComfyUI for:
- Tag file selection
- Insert behavior (commas, spaces, underscore replacement)
- Result count
- Hotkey mapping
- Color customization
---
## Phase 4: Wildcard Completion (Requires Parser Support)
**Goal:** Complete `__wildcard__` syntax
**Estimated Effort:** 2-3 days
**Blocked by:** Wildcard support needs to be added to the parser first
### 4.1 Trigger System
- [ ] **`__` trigger**: Start showing wildcard files after `__`
- [ ] **Two-stage completion**:
1. First `__` shows available wildcard files
2. Selecting a file shows contents (possible replacements)
### 4.2 Data Collection
- [ ] **Scan `data/wildcards/`** for `.txt` files
- [ ] **Parse wildcard files**: One option per line
- [ ] **Support nested folders**: `hair/colors/light.txt` → `__hair/colors/light__`
### 4.3 Insertion
- [ ] **File selection**: Insert `__filename__` for random selection
- [ ] **Direct selection**: Insert specific option from file contents
---
## Technical Architecture
### Backend Structure
```
api/
├── __init__.py
├── tokenize.py # (existing) Token counter
└── autocomplete.py # Tag autocomplete API
├── TagDatabase # In-memory tag index
├── search_tags() # Fuzzy prefix search
├── load_csv() # CSV parser
└── register_endpoints() # FastAPI/aiohttp routes
data/
├── tags/
│ ├── danbooru.csv # Main tag database
│ ├── e621.csv # Alternative database
│ ├── extra-quality-tags.csv
│ └── translations/ # Translation files
└── wildcards/
├── README.md
└── *.txt # User wildcard files
```
### Frontend Structure
```
js/
├── a1111_prompt.js # (existing) Main extension entry
├── tokenCounter.js # (existing) Token tooltip
├── autocomplete.js # Autocomplete UI
│ ├── AutocompletePopup # Popup DOM management
│ ├── ResultsList # Results rendering
│ ├── KeyboardNav # Arrow key navigation
│ └── insertTag() # Text insertion logic
├── caretPosition.js # Cursor position tracking
└── syntaxHighlight.js # (placeholder) Syntax highlighting
```
### Data Flow
```
User Types → Debounce → Extract Word → API Call → Filter Results → Update Popup
↓ ↓
Key Press ←←←←←←←←← Navigate ←←←←←←←←←←←←←←←←←←←←←←←←←←←←←←←←← Render List
↓
Tab/Enter → Insert Selected → Close Popup → Continue Typing
```
---
## CSV Tag Format Reference
```csv
<name>,<type>,<postCount>,"<aliases>"
1girl,0,4114588,"1girls,sole_female"
solo,0,3426446,"female_solo,solo_female"
highres,5,3008413,"high_res,high_resolution,hires"
```
| Column | Description |
| --------- | ---------------------------------------- |
| name | Tag text (underscores for spaces) |
| type | Category number (see below) |
| postCount | Usage frequency (for sorting) |
| aliases | Comma-separated alternatives (in quotes) |
### Tag Types (Danbooru)
| Value | Category | Color (Dark/Light) |
| ----- | --------- | -------------------- |
| 0 | General | lightblue/dodgerblue |
| 1 | Artist | indianred/firebrick |
| 3 | Copyright | violet/darkorchid |
| 4 | Character | lightgreen/darkgreen |
| 5 | Meta | orange/darkorange |
---
## Implementation Priority
| Phase | Priority | Effort | Value | Status |
| ---------------------- | ----------- | -------- | ------------------------------- | ----------- |
| Phase 1 (Core) | 🔥 Critical | 3-5 days | High - Basic functionality | ✅ Complete |
| Phase 2 (Aliases) | ⚡ High | 2-3 days | Medium - Better discoverability | ✅ Complete |
| Phase 3 (Advanced) | 💡 Medium | 3-5 days | Medium - QoL improvements | 📋 Next |
| Phase 4 (Wildcards) | 💡 Low | 2-3 days | Low - Blocked by parser support | ⏸️ Blocked |
**Total Estimated Effort:** 10-15 days for full implementation
**Current Status:** Phase 1 & 2 complete! Autocomplete is fully functional with alias support.
> **Note:** Phase 4 (Wildcards) requires wildcard support to be added to the parser first.
> Phase 3 (Advanced Features) is now the next priority, especially frequency sorting.
---
## Key Differences from A1111
1. **No Gradio**: ComfyUI uses LiteGraph, so integration is different
2. **No shared.opts**: Settings storage needs custom solution
3. **File serving**: May need different approach for loading CSV/preview images
4. **Textareas**: ComfyUI widget textareas vs Gradio textareas
5. **Extension loading**: JS module loading differs from A1111
6. **No Extra Networks**: ComfyUI uses dedicated nodes for LoRAs/embeddings, not inline syntax
---
## Files to Port/Reference
| A1111 File | Purpose | Port Priority |
| ---------------------------- | --------------------------- | ----------------------- |
| `tagAutocomplete.js` (68KB) | Main logic | Extract core patterns |
| `_caretPosition.js` | Cursor position calculation | Port directly |
| `_utils.js` | CSV parsing, debounce | Port utilities |
| `__globals.js` | Shared state | Adapt for our structure |
| `tag_autocomplete_helper.py` | Backend API | Reference for endpoints |
| `danbooru.csv` | Tag database | Include as data file |
---
## Testing Checklist
- [ ] Basic word completion after comma
- [ ] Completion at start of prompt
- [ ] Completion inside parentheses
- [ ] Keyboard navigation (up/down/tab/enter/escape)
- [ ] Mouse click selection
- [ ] Underscore replacement
- [ ] Parenthesis escaping
- [ ] Alias search and insertion
- [ ] Wildcard completion with `__`
- [ ] Performance with large tag database (100k+ tags)
- [ ] Memory usage monitoring
- [ ] Dark/light mode colors
---
## Open Questions
1. **Settings Storage**: Where to store user preferences? (localStorage vs ComfyUI config)
2. **Tag File Distribution**: Bundle CSVs or download on first use? (Size: ~4MB each)
3. **Integration Point**: How to attach to ComfyUI textarea widgets?
---
## Next Steps
1. **Start Phase 1.1**: Create `TagDatabase` class with CSV loading
2. **Create API endpoint**: `/a1111_prompt/autocomplete`
3. **Build minimal popup**: Just show results, no styling
4. **Iterate on UX**: Add keyboard nav, styling, cursor tracking
5. **Add configuration**: User preferences for behavior
---
_Document created: 2026-01-21_
_Based on analysis of a1111-sd-webui-tagcomplete v2025.x_
+118 -1
View File
@@ -5,7 +5,84 @@ All notable changes to this project will be documented in this file.
The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
## [Unreleased] - 2026-02-02
Rules:
- Use proper sub titles "Added", "Changed", "Removed" and "Fixed"
- Keep proper track of days for where entries should go
- Be concise but mention all changes without necessarily detailing each one
## [Unreleased] - 2026-07-25
### Fixed
- Fixed the resolved prompt toggle occasionally rendering outside the node and becoming unclickable by using ComfyUI's current widget visibility and layout APIs.
- Kept the last resolved prompt preview visible while editing the prompt, replacing it only after the next execution.
- Made wildcard autocomplete deterministic by loading the complete path manifest once and filtering it locally, eliminating delayed and incomplete cached-prefix results.
- Reused the startup wildcard path catalog during prompt expansion, eliminating repeated full-tree scans for concise unique wildcard names.
## [Unreleased] - 2026-07-23
### Fixed
- Fixed resolved prompt preview state resetting after each generation; the preview now stays expanded or collapsed according to the user's last toggle.
- Hardened resolved prompt preview rendering so preview widget failures do not break autocomplete or the rest of the node frontend.
## [Unreleased] - 2026-06-15
### Added
- **Autocomplete Response Cache**: Added a client-side cache for repeated and prefix-refined autocomplete queries.
- **Autocomplete Prefix Index**: Added backend prefix indexing for tag names and aliases to speed up large tag database searches.
- **Wildcard Token Estimates**: Added deterministic longest-option wildcard and dynamic prompt expansion for live token counting.
- **Prompt PNG Metadata**: Added an `a1111_prompt_node` PNG metadata entry that records each A1111 prompt node's original and resolved prompt by node id.
### Changed
- **Autocomplete Responsiveness**: Reduced autocomplete debounce from 100 ms to 50 ms and shows cached prefix matches while fresh results load.
- **Autocomplete Ranking**: Frequent and recently selected entries now have stronger ranking influence while exact matches remain prioritized.
- **Effective Prompt Restore**: The resolved prompt preview is now stored with the workflow and shown behind a compact toggle so dragged-back images can reveal the prompt that was actually encoded without shrinking the main prompt field.
- **Wildcard Completion Text**: Nested wildcard autocomplete now inserts the leaf name when it is unambiguous, falling back to the full path only for duplicate leaf names.
- **Autocomplete Warm-Up**: The frontend now starts loading the tag database before the first typed autocomplete request, and contains searches use lazy candidate indexing to preserve substring results without repeated full-database scans.
### Fixed
- Fixed autocomplete lifecycle rebinding when ComfyUI recreates prompt textarea DOM.
- Fixed autocomplete API coercion for unusual query, limit, boolean, file, and request-body values.
- Fixed token boundary positions after A1111 syntax stripping so markers align better with original prompt text.
- Fixed token counter route registration/import resilience after adding wildcard token estimates.
- Fixed literal `__wildcard__` text being counted instead of the maximum estimated wildcard expansion in the token counter.
- Fixed wildcard-only parser imports requiring Lark when scheduler parsing is not being used.
- Fixed dragged-back generated images not showing the resolved prompt preview by writing restore data into the embedded workflow metadata during node execution.
- Fixed resolved prompt preview collapse/expand layout so the preview stays inside the node bounds after toggling.
- Fixed token sequence boundary markers disappearing for wildcard-estimated counts and later sequence cutoffs.
- Hardened resolved prompt preview toggles so restored or rebuilt widgets keep a working click handler and stray input events do not collapse an unchanged preview.
- Fixed scheduled prompt padding being written back into the global encoded prompt cache.
## [Unreleased] - 2026-04-22
### Added
- **Wildcard Expansion**: Added A1111-style `__wildcard__` expansion for prompt and negative prompt nodes, including nested wildcard resolution from `data/wildcards`.
- **Dynamic Prompt Support in Wildcards**: Added brace-style choice expansion inside wildcard content, including simple choices (`{a|b}`), ranged picks (`{1-2$$a|b|c}`), and weighted picks (`{20%a|b}`).
- **Wildcard Autocomplete**: Added wildcard-aware autocomplete when typing `__`, with support for wildcard files, inferred folders, and second-stage file-content suggestions.
- **Wildcard Documentation**: Updated README, TODO, and changelog notes to cover wildcard expansion, autocomplete flow, and current completion status.
### Changed
- **Expanded Prompt Preview**: Effective prompt output now reflects resolved wildcard expansion so executed prompt text is visible after generation.
- **Wildcard Browsing UX**: Wildcard autocomplete now keeps the popup open after selecting a wildcard file and drills into either descendant entries or file contents.
- **Wildcard Result Rendering**: Improved wildcard result styling in the autocomplete popup to make nested paths easier to scan.
- **Prompt Cleanup**: Wildcard and dynamic prompt expansion now normalizes leftover spacing from optional empty branches so final prompts do not accumulate doubled spaces or messy comma spacing.
- **Autocomplete Ranking**: Frequency sorting now uses a blended relevance model that keeps exact/prefix/contains match quality ahead of usage history, with logarithmic frequency scaling and a light recency boost.
### Fixed
- Fixed wildcard autocomplete result limits so large wildcard packs return substantially more matching entries.
- Fixed wildcard expansion to resolve a1111-style wildcard packs that combine nested wildcard files with dynamic prompt syntax.
- Fixed wildcard autocomplete usage stars being display-only; usage history now actually influences wildcard result ordering.
## [Unreleased] - 2026-04-03
### Added
- **API Normalization**: The `autocomplete` API now robustly handles the `extra_files` parameter, coercing strings to lists and providing sensible defaults.
- **Memory Management Refinements**:
@@ -13,9 +90,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
- **Internal Cleanup**: Added explicit tensor deletion in `hooks.py` to prevent accumulation.
- **Leak Prevention**: Refactored frontend extensions to properly disconnect `ResizeObserver` instances and remove global click listeners upon node removal.
### Fixed
- **Hotfixes & Stability**:
- **Backend**: Restored missing `transformers_dict` initialization in `hooks.py` to fix `AttributeError` crash.
- **Frontend**: Merged duplicate `nodeCreated` hooks in `autocomplete.js` to restore extension functionality.
### Changed
- **Performance Optimization**:
- **Shared Caching**: Implemented a structural result cache in `hooks.py` that persists across hook clones, critical for 2nd order samplers (DPM++ 2M, etc.).
- **GPU Pre-Allocation**: Moved all embeddings to the GPU during the encoding phase to eliminate device-transfer latency.
@@ -30,6 +112,41 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
- **Improved Cache Management**: Refactored cache clearing to happen before execution loops, preventing data loss during multi-prompt runs.
- **UI Robustness**: Improved autocomplete error handling to correctly reset internal selection state upon backend failure.
## [1.0.0] - 2026-02-02
### 🎉 Launch!
Stable release on the ComfyUI Registry with full parity verification and security hardening.
### Added
- **Comprehensive Test Suite**: 5 new test categories verifying the scheduler, hooks, node math, tokenizer, and API security.
- **GitHub Registry CI**: Added automated publication workflow for ComfyUI Registry.
- **Project Metadata**: Full `pyproject.toml` and `requirements.txt` integration for standardized distribution.
### Security
- **API Hardening**: Implemented a 255-character query length limit for the `autocomplete` API to prevent long-query DOS attempts.
### Changed
- **Verified A1111 Parity**:
- Verified `BREAK` isolation logic ensures zero concept bleeding between segments.
- Confirmed `_apply_direct_scaling` math exactly match A1111 emphasis (multiplication vs interpolation).
- Verified `normalization` accurately implements A1111 "Mean Rescaling" (Norm vs No Norm).
- **Hooks Scaling**: Improved hook sigma-to-step conversion for better accuracy in low-step (2-4 steps) generations.
### Refinements
- **Memory Management**:
- Bounded caching for prompt encoding to prevent memory leak on high-volume runs.
- Explicit tensor cleanup in hook executions.
- Frontend lifecycle management for `ResizeObserver` and global listeners.
- **Performance**:
- GPU pre-allocation for per-step embeddings.
- Shared result caching for cloned hooks in 2nd order samplers.
- Zero-allocation hot-path for sampling loop swapping.
## [Unreleased] - 2026-01-31
### Added
+21
View File
@@ -0,0 +1,21 @@
MIT License
Copyright (c) 2026 Enferlain
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
+111 -43
View File
@@ -14,15 +14,19 @@ pip install lark
### Positive prompt node
<img width="790" height="527" alt="image" src="https://github.com/user-attachments/assets/67c12e99-cf35-4d2e-af84-0e258fa6915d" />
![positive](assets/positive_node.png)
### Negative prompt node
<img width="363" height="241" alt="image" src="https://github.com/user-attachments/assets/ba1da788-dc47-4939-bcbe-d1699ea9f872" />
![negative](assets/negative_node.png)
### Autocomplete/Wildcards
![wildcard](assets/wildcard_example.png)
### Full setup with optional TIPO support
<img width="1755" height="813" alt="image" src="https://github.com/user-attachments/assets/d430189f-6f50-42e9-8503-afa8d6d7db8a" />
![tipo](assets/tipo_example.png)
### Core Features
@@ -30,7 +34,8 @@ pip install lark
- **Direct Scaling (Anti-Burn)**: Uses `z * weight` instead of Comfy's interpolation, avoiding artifacts at high weights
- **BREAK Support**: Fully isolated context windows - each BREAK segment is tokenized separately
- **Emphasis**: `(text:1.2)`, `(text)`, `[text]`
- **TIPO support**: TIPO prompt output can connect directly into the node, and it will show the generated prompt when the node receives it. Should use [my fork](https://github.com/Enferlain/z-tipo-extension/tree/custom) to preserve weighting emphasis in the a1111 snytax.
- **Wildcard support**: Resolves `__wildcard__` files from `data/wildcards`, including nested wildcard chains and dynamic prompt choices inside wildcard files
- **TIPO support**: TIPO prompt output can connect directly into the node, which shows the generated prompt it receives. Use [my fork](https://github.com/Enferlain/z-tipo-extension/tree/custom) to expand connected A1111 wildcards before the single TIPO generation pass and preserve weighting emphasis in A1111 syntax.
### Token Counter
@@ -63,7 +68,7 @@ The node displays a **live token count** in the header, showing tokens per 77-to
The node includes [**A1111-style tag autocomplete**](https://github.com/DominikDoom/a1111-sd-webui-tagcomplete) functionality:
- **Trigger**: Start typing any tag (2+ characters)
- **Database**: Uses Danbooru/e621 tag databases (~140k tags)
- **Database**: Uses `danbooru.csv` by default (~140k tags), with additional Danbooru/E621 databases bundled.
- **Search**: Matches tag names and aliases
- **Navigation**: Use ↑/↓ arrows, Tab/Enter to select, Escape to close
- **Color coding**: Tags are colored by type (general, artist, character, etc.)
@@ -76,17 +81,43 @@ The node includes [**A1111-style tag autocomplete**](https://github.com/DominikD
- Alias support: Type `sole_female` → suggests `1girl`
- Smart insertion: Automatically adds commas and handles spacing
- Parenthesis escaping: `name_(artist)` → `name_\(artist\)`
- Real-time search with 100ms debouncing
- Real-time search with 50ms debouncing
- Usage tracking: Tags you use often are prioritized in results
- Quality tags: Automatically includes `extra-quality-tags.csv` for common quality/style tags
- Theme-aware: Respects your ComfyUI color scheme
**Available tag databases:**
### Wildcard Support
- `danbooru.csv` - Main Danbooru database (~140k tags)
The node supports **A1111-style wildcard files** stored in `data/wildcards/`.
- **Basic syntax**: `__outfits__`
- **Nested folders**: `__characters/anime__`
- **Nested expansion**: wildcard files can reference other wildcard files
- **Dynamic prompt choices inside wildcard files**:
- `{red|blue}`
- `{1-2$$ring|necklace|bracelet}`
- `{20%hat|scarf}`
**Autocomplete integration:**
- Type `__` to browse wildcard folders and files
- Selecting a wildcard file inserts the concise `__leaf__` name when it is unique, or the full `__path/to/file__` when needed to disambiguate duplicate leaf names
- After selecting a wildcard file, the popup stays open and can show the file's contents for quick follow-up selection
- The complete wildcard path list is loaded once and filtered locally for consistent typing latency
**Notes:**
- Wildcard files are plain `.txt` files under `data/wildcards/`
- Blank lines and `#` comment lines are ignored
- Prompt spacing is normalized after expansion so optional empty branches do not leave doubled spaces behind
**Bundled tag databases:**
- `danbooru.csv` - Default Danbooru database (~140k tags)
- `e621.csv` - E621 database (furry-focused)
- `danbooru_e621_merged_*.csv` - Larger merged Danbooru/E621 databases
- `extra-quality-tags.csv` - Quality and style tags (auto-loaded)
- Custom CSV files can be added to `data/tags/`
- Custom CSV files can be added to `data/tags/`; the current autocomplete UI uses the configured backend default and does not expose a database selector
**Frequency Management:**
Open browser console and use:
@@ -97,20 +128,23 @@ Open browser console and use:
### Scheduling
| Syntax | Meaning | Requires steps? |
| --------------- | ---------------------------------------- | --------------- |
| `[cat:dog:0.5]` | Switch from "cat" to "dog" at 50% | No |
| `[cat:dog:10]` | Switch at step 10 (requires steps input) | **Yes** |
| `[glasses:0.5]` | Add "glasses" at 50% | No |
| `[glasses:10]` | Add at step 10 | **Yes** |
| `[hat::0.7]` | Remove "hat" at 70% | No |
| `[hat::15]` | Remove at step 15 | **Yes** |
| Syntax | Meaning | Auto-detects? |
| --------------- | --------------------------------- | ------------- |
| `[cat:dog:0.5]` | Switch from "cat" to "dog" at 50% | Yes |
| `[cat:dog:10]` | Switch at step 10 (literal count) | Yes |
| `[glasses:0.5]` | Add "glasses" at 50% | Yes |
| `[glasses:10]` | Add at step 10 | Yes |
| `[hat::0.7]` | Remove "hat" at 70% | Yes |
| `[hat::15]` | Remove at step 15 | Yes |
**Note**: The node automatically detects the total step count from your connected Sampler or Scheduler.
**Important**:
- **Percentage-based syntax** (with decimals like `0.5`) works automatically - no steps parameter needed
- **Step-based syntax** (integers like `10`) requires setting the `steps` parameter to match your sampler
- The node will detect which syntax you're using and show an error if steps are needed but not provided
- **Steps are automatic**: You do not need to manually set a step count; the node traverses the workflow graph to find the connected Sampler/Scheduler.
- **Scaling**:
- **Percentages** (e.g., `0.5`) scale proportionally to any sampler step count.
- **Integers** (e.g., `10`) are tied to the detected step count. If the sampler steps are changed after encoding, the node automatically scales the transition point to maintain the same relative timing.
**Nested syntax supported:**
@@ -132,9 +166,19 @@ Control when alternation starts or stops:
| Syntax | Meaning |
| -------------------- | ----------------------------------------- |
| `[as109\|fkey::0.6]` | Alternate until 60%, then nothing |
| `[as109\|fkey:0.4]` | Nothing until 40%, then start alternating |
| `[as109\|fkey::15]` | Alternate until step 15, then nothing |
| `[as109\|fkey::0.6]` | Alternate until 60%, then as109 |
| `[as109\|fkey:0.4]` | as109 until 40%, then start alternating |
| `[as109\|fkey::15]` | Alternate until step 15, then as109 |
> [!TIP]
> **Matching A1111/Forge multi-evaluation sampler behavior:** A1111/Forge can
> advance prompt scheduling once per denoiser evaluation rather than once per
> displayed sampler step. With a sampler that performs two evaluations per
> step, a 28-step alternation can therefore end around displayed step 14 and
> settle on its first option. To reproduce that effective result here, use a
> scheduled cutoff such as `[A|B::0.5]`: alternate for the first half, then
> keep `A`. A plain `[A|B]` continues across the full ComfyUI denoising
> schedule.
**Combining with scheduling to lock to a value:**
@@ -149,6 +193,24 @@ This alternates until 60%, then switches to just "as109".
- **SDXL**: Dual CLIP (clip_l + clip_g) with independent weight scaling
- **SD1.5**: Full support
### What is Normalization?
This node **always** uses A1111-style "Direct Scaling" (multiplication) instead of ComfyUI's standard interpolation. This is what provides the core "A1111 look" and prevents the default ComfyUI color burn.
The `normalization` toggle specifically controls A1111's **Mean Rescaling** (Norm vs No Norm):
**When `normalization` is enabled (True):**
- **Match A1111 "Norm"**: Rescales the final conditioning so its mathematical "energy" (mean) matches the original.
- **Effect**: Prevents saturation shifts at very high weights.
- **Commonly used with**: **SD1.5** and models sensitive to high prompt energy.
**When `normalization` is disabled (False - Default):**
- **Match A1111 "No Norm"**: Applies weights exactly as written without any rescaling.
- **Effect**: Emphasized words "punch" through more aggressively.
- **Commonly used with**: **SDXL** and modern models that handle high energy well.
---
## Usage
@@ -164,12 +226,12 @@ This pack provides **two nodes**:
#### Inputs
| Input | Type | Required | Description |
| ------------- | ------ | -------- | ---------------------------------- |
| clip | CLIP | Yes | The CLIP model |
| text | STRING | Yes | Prompt with A1111 syntax |
| normalization | BOOL | No | Enable EmphasisOriginalNoNorm |
| debug | BOOL | No | Show detailed schedule information |
| Input | Type | Required | Description |
| ------------- | ------ | -------- | ------------------------------------------------ |
| clip | CLIP | Yes | The CLIP model |
| text | STRING | Yes | Prompt with A1111 syntax |
| normalization | BOOL | No | EmphasisOriginal/EmphasisOriginalNoNorm toggle |
| debug | BOOL | No | Show detailed schedule information |
#### Outputs
@@ -179,9 +241,9 @@ This pack provides **two nodes**:
**Important**:
- **Percentage syntax** `[cat:dog:0.5]` works with `steps=0` (auto-detects from sampler)
- **Step-based syntax** `[cat:dog:10]` requires setting `steps` to match your sampler's step count
- If you use step-based syntax with `steps=0`, you'll get a clear error message
- **Auto-detection**: The node detects steps from the downstream sampler. If no sampler is connected, step-based syntax (`[a:b:10]`) will raise an error.
- **Flexibility**: Use percentage syntax (`[a:b:0.5]`) for the most portable prompts.
- **Effective prompt preview**: When TIPO or wildcards expand the text, the node shows the resolved prompt in the UI so you can see what was actually encoded. Its compact toggle preserves the expanded or collapsed state between executions.
### A1111 Style Prompt (Negative)
@@ -227,7 +289,7 @@ artist name BREAK character name, wearing a hat
# Step-based scheduling (with steps=28)
[mountains:ocean:14] at sunset
# Alternation (requires MODEL connected)
# Alternation (no MODEL connection needed)
1girl, [as109|fkey], detailed
# Add element at specific step
@@ -282,19 +344,18 @@ The node uses ComfyUI's `TransformerOptionsHook` system to swap conditioning per
The node's `INPUT_TYPES` exposes only `clip`, `text`, `normalization`, and `debug`. Step-based syntax and timing are managed as follows:
1. **Auto-detection**: The node automatically detects the total step count from the connected sampler/scheduler downstream.
2. **Step-based syntax**: Syntax like `[thing:10]` (using integers) requires a connected sampler. If no sampler is found in the graph downstream, the node will raise a clear error.
3. **Percentage-based syntax**: Syntax like `[thing:0.5]` (using decimals) works without any special requirements and scales to the sampler's step count automatically.
4. **Recommended**: Use percentage syntax for better flexibility across different step counts.
2. **Step-based syntax**: Syntax like `[thing:10]` (using integers) is relative to the detected count. The node scales these points automatically if the sampler step count changes during sampling to ensure consistent timing.
3. **Percentage-based syntax**: Syntax like `[thing:0.5]` (using decimals) is the recommended way to ensure your prompt behaves identically across different step counts.
### Known Limitations
1. **Alternation is positive-only**: Only the main node supports alternation. The negative node will use the first step's prompt if scheduling syntax is present.
2. **Step-based syntax requires matching steps**: If you write `[thing:10]` with `steps=20` but your sampler uses 30 steps, the transition will happen at step 10 (not scaled). Use percentage syntax for automatic scaling.
2. **Automatic Scaling**: Unlike A1111 (where steps are manual), this node provides **automatic scaling**. If you write `[thing:10]` for a 20-step run but later change it to 30 steps, the node will automatically move the transition to step 15 to preserve the 50% relative timing. Use percentage syntax `[thing:0.5]` for the most explicit control.
3. **Visual parity**: While the **prompt schedule** matches A1111 exactly (the same prompt text at each step), the **visual effect** may differ due to architectural differences:
3. **Visual parity**: This node follows ComfyUI's denoising progress, while A1111/Forge can advance prompt scheduling per denoiser evaluation on multi-evaluation samplers. Plain schedules can therefore differ when A1111/Forge exhausts them early. Use a scheduled cutoff such as `[a|b::0.5]` when you intentionally want that halfway behavior. The visual effect may still differ due to architectural differences:
- A1111 applies conditioning at the CFGDenoiser level (before model call)
- This node applies conditioning via model wrapper (during model call)
- This node applies conditioning through ComfyUI's hook system during sampling
Use the **scheduled alternation** syntax (`[a|b::0.6]`) if you need to control exactly when alternation stops.
@@ -302,7 +363,14 @@ Use the **scheduled alternation** syntax (`[a|b::0.6]`) if you need to control e
## Performance
- **Efficient encoding**: Unique prompts are encoded only once and cached
- `[A|B]` with 28 steps only encodes 2 prompts, not 28
- All embeddings are padded to the same length for efficient swapping
- BREAK segments are batched together for efficient processing
- **Shared Wildcard Catalog**: Wildcard paths are scanned once during startup and reused by autocomplete and runtime expansion.
- **Efficient encoding**: Unique prompts are encoded only once and cached using composite keys (CLIP + text + normalization).
- **GPU Pre-Allocation**: All embeddings are pre-moved to the GPU (`intermediate_device`) during encoding to eliminate device-transfer latency during the sampling loop.
- **Shared Hook Caching**: The step-swapping hook uses a structural result cache that persists when the hook is cloned, avoiding repeated conditioning tensor construction across sampling calls.
- **Optimized Padding**: Tensors are padded only to the lengths used within the current generation, preventing global cache bloat.
- **Direct Memory Swapping**: If prompts have identical sequence lengths, the hook uses a zero-math "fast-path" to swap embeddings instantly.
- **Batch Processing**: Multiple BREAK segments are batched into a single encoding pass.
---
[![Ask DeepWiki](https://deepwiki.com/badge.svg)](https://deepwiki.com/Enferlain/ComfyUI-A1111-cond)
+21 -2
View File
@@ -39,6 +39,15 @@
- [x] Frequency sorting with usage tracking
- [x] Theme support (respects ComfyUI color scheme)
### Wildcards (2026-04-21)
- [x] `__wildcard__` syntax expansion
- [x] Nested wildcards
- [x] Dynamic prompt syntax support inside wildcard content (`{a|b}`, `{1-2$$a|b|c}`, weighted picks)
- [x] Wildcard autocomplete when typing `__`
- [x] Wildcard folder + file suggestions with second-stage content drill-down
- [x] Wildcard prompt spacing cleanup after optional/empty expansions
---
## 🔥 High Priority
@@ -62,6 +71,12 @@
- [ ] Configuration UI for tag file selection
- [ ] Chants/prompt presets (prompt library/bookmarks)
- [ ] Wiki links for tag documentation
- [ ] Source-aware autocomplete results
- Load multiple tag sources without flattening away source identity
- Preserve per-source aliases instead of using one global alias map
- Show prepended source badge/icon in the popup (e.g. Danbooru, e621)
- Handle duplicate canonical tags by showing one result per source when needed
- [ ] Check if usage count is recorded properly
- [ ] Visual distinction between different databases (e.g., Danbooru vs e621 tags)
- Consider: badge/icon, subtle background color, or source indicator
- Useful when loading multiple databases simultaneously
@@ -69,10 +84,13 @@
### Wildcard Support
- [ ] `__wildcard__` syntax expansion
- [ ] Nested wildcards
- [x] `__wildcard__` syntax expansion
- [x] Nested wildcards
- [ ] Wildcard file browser/picker
- Current state: autocomplete can browse wildcard folders/files and drill into file contents
- [ ] Preview what wildcards will expand to
- [ ] Toggle for static wildcards (same seed same wildcard result)
- [ ] Toggle for history for what wildcards have been used so you don't get the same one too soon
---
@@ -200,3 +218,4 @@ A1111_Prompt_Node/
- [ ] Export to A1111-compatible format
- [ ] Prompt history (undo/redo)
- [ ] Multi-line prompt editor (full-screen mode)
- [ ] Wildcard/tags having fast/performant thumbnails on hover in autocomplete.js (needs autocompelte.js package split)
+638 -84
View File
@@ -7,11 +7,17 @@ Provides tag autocomplete functionality with support for:
- Post count sorting with optional frequency boosting
"""
import asyncio
import csv
import re
import threading
from pathlib import Path
from typing import Dict, List, Optional, Tuple
try:
from ..parser.wildcards import get_wildcard_file_index
except ImportError:
from parser.wildcards import get_wildcard_file_index
# Conditional server import
try:
from aiohttp import web
@@ -34,11 +40,68 @@ TAG_TYPES = {
# Get the data directory path
DATA_DIR = Path(__file__).parent.parent / "data" / "tags"
WILDCARDS_DIR = Path(__file__).parent.parent / "data" / "wildcards"
# Default tag file to use if not specified
# Default tag file to use if not specified (danbooru.csv is default)
DEFAULT_TAG_FILE = "danbooru.csv"
def _coerce_text(value) -> str:
if value is None:
return ""
if isinstance(value, str):
return value
return str(value)
def _coerce_bool(value, default: bool = True) -> bool:
if isinstance(value, bool):
return value
if isinstance(value, str):
normalized = value.strip().lower()
if normalized in {"1", "true", "yes", "on"}:
return True
if normalized in {"0", "false", "no", "off"}:
return False
return default
def _coerce_limit(value, default: int = 20, cap: int = 100) -> int:
try:
limit = int(value)
except (TypeError, ValueError):
limit = default
return max(0, min(limit, cap))
def _normalize_tag_filename(value, default: Optional[str] = None) -> Optional[str]:
filename = _coerce_text(value).strip()
if not filename:
return default
filename = filename.replace("\\", "/")
path = Path(filename)
if path.is_absolute() or ".." in path.parts:
return default
return filename
def _normalize_extra_files(value) -> List[str]:
if value is None:
value = ["extra-quality-tags.csv"]
elif isinstance(value, str):
value = [value]
elif not isinstance(value, list):
value = ["extra-quality-tags.csv"]
files = []
for item in value:
filename = _normalize_tag_filename(item)
if filename and filename not in files:
files.append(filename)
return files
class TagEntry:
"""Represents a single tag entry from the database."""
@@ -58,6 +121,29 @@ class TagEntry:
"type": self.type,
"count": self.count,
"aliases": self.aliases,
"kind": "tag",
}
class WildcardEntry:
"""Represents a single wildcard file available for completion."""
__slots__ = ("name", "kind", "search_text", "leaf_name")
def __init__(self, name: str, kind: str):
self.name = name
self.kind = kind
self.search_text = name.lower()
self.leaf_name = name.rsplit("/", 1)[-1]
def to_dict(self, completion_name: str = "") -> Dict:
completion_name = completion_name or self.leaf_name
return {
"name": self.name,
"leaf_name": self.leaf_name,
"kind": self.kind,
"completion": f"__{completion_name}__",
"meta": "Wildcard folder" if self.kind == "wildcard_folder" else "Wildcard file",
}
@@ -72,6 +158,8 @@ class TagDatabase:
def __init__(self):
self._tags: List[TagEntry] = []
self._alias_map: Dict[str, TagEntry] = {} # alias -> canonical tag
self._prefix_index: Dict[str, List[Tuple[TagEntry, Optional[str], str]]] = {}
self._contains_index: Dict[str, List[Tuple[TagEntry, Optional[str], str]]] = {}
self._loaded = False
self._current_files: List[str] = []
@@ -83,6 +171,45 @@ class TagDatabase:
def tag_count(self) -> int:
return len(self._tags)
def _add_prefix_index_entry(
self, text: str, tag: TagEntry, matched_alias: Optional[str] = None
) -> None:
normalized = text.lower().strip()
if not normalized:
return
for prefix_len in range(1, min(3, len(normalized)) + 1):
prefix = normalized[:prefix_len]
self._prefix_index.setdefault(prefix, []).append(
(tag, matched_alias, normalized)
)
def _get_contains_index_key(self, query: str) -> str:
return query[: min(3, len(query))]
def _get_contains_candidates(
self, query: str
) -> List[Tuple[TagEntry, Optional[str], str]]:
key = self._get_contains_index_key(query)
if not key:
return []
if key not in self._contains_index:
candidates: List[Tuple[TagEntry, Optional[str], str]] = []
for tag in self._tags:
if key in tag.search_text:
candidates.append((tag, None, tag.search_text))
for alias_lower, tag in self._alias_map.items():
if key in alias_lower:
original_alias = next(
(a for a in tag.aliases if a.lower() == alias_lower),
alias_lower,
)
candidates.append((tag, original_alias, alias_lower))
self._contains_index[key] = candidates
return self._contains_index[key]
def load_csv(self, filepath: Path, append: bool = False) -> int:
"""
Load tags from a CSV file.
@@ -100,7 +227,9 @@ class TagDatabase:
if not append:
self._tags = []
self._alias_map = {}
self._prefix_index = {}
self._current_files = []
self._contains_index = {}
if not filepath.exists():
print(f"[Autocomplete] Tag file not found: {filepath}")
@@ -139,10 +268,12 @@ class TagDatabase:
self._tags.append(entry)
existing_tags.add(name)
loaded_count += 1
self._add_prefix_index_entry(name, entry)
# Build alias -> canonical tag mapping
for alias in aliases:
self._alias_map[alias.lower()] = entry
self._add_prefix_index_entry(alias, entry, alias)
self._loaded = True
self._current_files.append(filepath.name)
@@ -170,7 +301,11 @@ class TagDatabase:
return total
def search(
self, query: str, limit: int = 20, search_aliases: bool = True
self,
query: str,
limit: int = 20,
search_aliases: bool = True,
contains_fallback: bool = True,
) -> List[Dict]:
"""
Search for tags matching the query.
@@ -183,11 +318,14 @@ class TagDatabase:
Returns:
List of matching tags as dictionaries
"""
if not self._loaded or not query:
query = _coerce_text(query)
limit = _coerce_limit(limit, default=20, cap=100)
if not self._loaded or not query or limit == 0:
return []
query_lower = query.lower().strip()
if not query_lower:
if not query_lower or len(query_lower) > 255:
return []
results: List[Tuple[int, TagEntry, Optional[str]]] = []
@@ -202,44 +340,282 @@ class TagDatabase:
else:
return 1 # Contains match
# Search by tag name
for tag in self._tags:
if query_lower in tag.search_text:
score = get_score(tag, tag.search_text)
results.append((score, tag, None))
seen_tags.add(tag.name)
def add_result(tag: TagEntry, matched_alias: Optional[str], matched_text: str):
if tag.name in seen_tags:
return
score = get_score(tag, matched_text)
results.append((score, tag, matched_alias))
seen_tags.add(tag.name)
# Search by aliases
if search_aliases:
for alias_lower, tag in self._alias_map.items():
def format_results() -> List[Dict]:
results.sort(key=lambda x: (-x[0], -x[1].count))
output = []
for _, tag, matched_alias in results[:limit]:
entry = tag.to_dict()
if matched_alias:
entry["matched_alias"] = matched_alias
output.append(entry)
return output
prefix_key = query_lower[: min(3, len(query_lower))]
for tag, matched_alias, matched_text in self._prefix_index.get(prefix_key, []):
if matched_alias and not search_aliases:
continue
if matched_text == query_lower or matched_text.startswith(query_lower):
add_result(tag, matched_alias, matched_text)
if len(results) >= limit:
return format_results()
if not contains_fallback:
return format_results()
if len(query_lower) >= 2:
candidates = self._get_contains_candidates(query_lower)
for tag, matched_alias, matched_text in candidates:
if matched_alias and not search_aliases:
continue
if query_lower in matched_text:
add_result(tag, matched_alias, matched_text)
else:
# Backend-only fallback for one-character API searches. The frontend
# starts at two characters, where the lazy contains index applies.
for tag in self._tags:
if tag.name in seen_tags:
continue
if query_lower in alias_lower:
score = get_score(tag, alias_lower)
# Find the original case alias
original_alias = next(
(a for a in tag.aliases if a.lower() == alias_lower),
alias_lower,
)
results.append((score, tag, original_alias))
seen_tags.add(tag.name)
if query_lower in tag.search_text:
add_result(tag, None, tag.search_text)
# Sort by: score (desc), then post count (desc)
results.sort(key=lambda x: (-x[0], -x[1].count))
if search_aliases:
for alias_lower, tag in self._alias_map.items():
if tag.name in seen_tags:
continue
if query_lower in alias_lower:
original_alias = next(
(a for a in tag.aliases if a.lower() == alias_lower),
alias_lower,
)
add_result(tag, original_alias, alias_lower)
# Format results
output = []
for _, tag, matched_alias in results[:limit]:
entry = tag.to_dict()
if matched_alias:
entry["matched_alias"] = matched_alias
output.append(entry)
return format_results()
return output
class WildcardDatabase:
"""In-memory wildcard list with simple fuzzy scoring."""
def __init__(self, wildcards_dir: Path):
self._wildcards_dir = wildcards_dir
self._entries: List[WildcardEntry] = []
self._entry_map: Dict[str, List[WildcardEntry]] = {}
self._leaf_counts: Dict[str, int] = {}
self._load_lock = threading.Lock()
self._loaded = False
@property
def is_loaded(self) -> bool:
return self._loaded
@property
def wildcard_count(self) -> int:
return len(self._entries)
def load(self) -> int:
with self._load_lock:
return self._load_unlocked()
def _load_unlocked(self) -> int:
if not self._wildcards_dir.exists():
self._entries = []
self._entry_map = {}
self._leaf_counts = {}
self._loaded = True
return 0
entries: List[WildcardEntry] = []
folder_names = set()
for path in get_wildcard_file_index(
self._wildcards_dir, refresh=True
).files:
relative = path.relative_to(self._wildcards_dir).with_suffix("")
relative_name = relative.as_posix()
entries.append(WildcardEntry(relative_name, "wildcard_file"))
parts = relative.parts[:-1]
current_parts = []
for part in parts:
current_parts.append(part)
folder_names.add("/".join(current_parts))
for folder_name in folder_names:
entries.append(WildcardEntry(folder_name, "wildcard_folder"))
entries.sort(
key=lambda entry: (
entry.search_text,
0 if entry.kind == "wildcard_folder" else 1,
len(entry.name),
)
)
entry_map: Dict[str, List[WildcardEntry]] = {}
leaf_counts: Dict[str, int] = {}
for entry in entries:
entry_map.setdefault(entry.search_text, []).append(entry)
leaf_key = entry.leaf_name.lower()
leaf_counts[leaf_key] = leaf_counts.get(leaf_key, 0) + 1
self._entries = entries
self._entry_map = entry_map
self._leaf_counts = leaf_counts
self._loaded = True
return len(entries)
def ensure_loaded(self) -> int:
if self._loaded:
return len(self._entries)
with self._load_lock:
if not self._loaded:
return self._load_unlocked()
return len(self._entries)
def _entry_to_dict(self, entry: WildcardEntry) -> Dict:
completion_name = (
entry.leaf_name
if self._leaf_counts.get(entry.leaf_name.lower(), 0) == 1
else entry.name
)
return entry.to_dict(completion_name=completion_name)
def get_manifest(self) -> List[Dict]:
"""Return every wildcard file and folder for deterministic client search."""
self.ensure_loaded()
return [self._entry_to_dict(entry) for entry in self._entries]
def _normalize_query(self, query: str) -> str:
normalized = str(query or "").strip().lower()
if normalized.startswith("__"):
normalized = normalized[2:]
if normalized.endswith("__"):
normalized = normalized[:-2]
return normalized
def _resolve_entry(self, wildcard_name: str) -> Optional[WildcardEntry]:
normalized = self._normalize_query(wildcard_name)
if not normalized:
return None
candidates = self._entry_map.get(normalized, [])
if not candidates:
return None
for kind in ("wildcard_folder", "wildcard_file"):
for entry in candidates:
if entry.kind == kind:
return entry
return candidates[0]
def search(self, query: str, limit: int = 20) -> List[Dict]:
if not self._loaded:
self.ensure_loaded()
limit = _coerce_limit(limit, default=20, cap=500)
normalized = self._normalize_query(query)
if len(normalized) > 255 or limit == 0:
return []
if not normalized:
return [self._entry_to_dict(entry) for entry in self._entries[:limit]]
scored: List[Tuple[int, WildcardEntry]] = []
for entry in self._entries:
if normalized == entry.search_text:
score = 3
elif entry.search_text.startswith(normalized):
score = 2
elif normalized in entry.search_text:
score = 1
else:
continue
scored.append((score, entry))
scored.sort(
key=lambda item: (
-item[0],
0 if item[1].kind == "wildcard_folder" else 1,
len(item[1].name),
item[1].search_text,
)
)
return [self._entry_to_dict(entry) for _, entry in scored[:limit]]
def get_contents(
self, wildcard_name: str, content_query: str = "", limit: int = 50
) -> List[Dict]:
if not self._loaded:
self.ensure_loaded()
limit = _coerce_limit(limit, default=50, cap=500)
if limit == 0:
return []
entry = self._resolve_entry(wildcard_name)
if entry is None:
return []
if entry.kind == "wildcard_folder":
prefix = f"{entry.name}/"
normalized_query = _coerce_text(content_query).strip().lower()
descendants = []
for candidate in self._entries:
if candidate.name == entry.name or not candidate.name.startswith(prefix):
continue
remainder = candidate.name[len(prefix) :]
if normalized_query and normalized_query not in remainder.lower():
continue
descendants.append(self._entry_to_dict(candidate))
if len(descendants) >= limit:
break
return descendants
wildcard_path = self._wildcards_dir.joinpath(*entry.name.split("/")).with_suffix(
".txt"
)
if not wildcard_path.exists():
return []
normalized_query = _coerce_text(content_query).strip().lower()
results = []
with open(wildcard_path, "r", encoding="utf-8", errors="replace") as handle:
for raw_line in handle:
line = raw_line.strip()
if not line or line.startswith("#"):
continue
if normalized_query and normalized_query not in line.lower():
continue
results.append(
{
"name": line,
"kind": "wildcard_content",
"completion": line,
"meta": entry.name,
}
)
if len(results) >= limit:
break
return results
# Global database instance (lazy loaded)
_database = TagDatabase()
_wildcard_database = WildcardDatabase(WILDCARDS_DIR)
_database_load_lock = threading.Lock()
_database_loading = False
_database_load_error: Optional[str] = None
_database_loading_files: List[str] = []
_wildcard_warmup_lock = threading.Lock()
_wildcard_warmup_loading = False
def get_database() -> TagDatabase:
@@ -247,6 +623,61 @@ def get_database() -> TagDatabase:
return _database
def get_wildcard_database() -> WildcardDatabase:
"""Get the global wildcard database instance."""
return _wildcard_database
def _build_tag_file_list(
tag_file: str = DEFAULT_TAG_FILE,
extra_files: Optional[List[str]] = None,
) -> List[str]:
safe_tag_file = _normalize_tag_filename(tag_file, DEFAULT_TAG_FILE)
files_to_load = [safe_tag_file]
if extra_files:
files_to_load.extend(_normalize_extra_files(extra_files))
return [filename for filename in files_to_load if filename]
def _resolve_tag_filepaths(files_to_load: List[str]) -> List[Path]:
filepaths = []
for filename in files_to_load:
# Try main data directory first
filepath = DATA_DIR / filename
# Fall back to reference directory
if not filepath.exists():
ref_path = (
Path(__file__).parent.parent
/ "autocomplete_reference"
/ "a1111-sd-webui-tagcomplete"
/ "tags"
/ filename
)
if ref_path.exists():
filepath = ref_path
if filepath.exists():
filepaths.append(filepath)
else:
print(f"[Autocomplete] Warning: Tag file not found: {filename}")
return filepaths
def get_database_status() -> Dict:
db = get_database()
return {
"loaded": db.is_loaded,
"loading": _database_loading,
"load_error": _database_load_error,
"tag_count": db.tag_count,
"current_files": db._current_files,
"loading_files": _database_loading_files,
}
def ensure_database_loaded(
tag_file: str = DEFAULT_TAG_FILE,
extra_files: Optional[List[str]] = None,
@@ -261,49 +692,107 @@ def ensure_database_loaded(
Returns:
The loaded TagDatabase instance
"""
global _database_loading, _database_load_error, _database_loading_files
db = get_database()
# Build list of files to load
files_to_load = [tag_file]
if extra_files:
files_to_load.extend(extra_files)
files_to_load = _build_tag_file_list(tag_file, extra_files)
# Check if we need to reload
needs_reload = not db.is_loaded or set(db._current_files) != set(files_to_load)
if not needs_reload:
if _database_loading and set(_database_loading_files) == set(files_to_load):
_database_loading = False
_database_loading_files = []
return db
if needs_reload:
# Find and load all tag files
filepaths = []
with _database_load_lock:
needs_reload = not db.is_loaded or set(db._current_files) != set(files_to_load)
if not needs_reload:
if _database_loading and set(_database_loading_files) == set(files_to_load):
_database_loading = False
_database_loading_files = []
return db
for filename in files_to_load:
# Try main data directory first
filepath = DATA_DIR / filename
_database_loading = True
_database_load_error = None
_database_loading_files = files_to_load[:]
# Fall back to reference directory
if not filepath.exists():
ref_path = (
Path(__file__).parent.parent
/ "autocomplete_reference"
/ "a1111-sd-webui-tagcomplete"
/ "tags"
/ filename
)
if ref_path.exists():
filepath = ref_path
if filepath.exists():
filepaths.append(filepath)
else:
print(f"[Autocomplete] Warning: Tag file not found: {filename}")
if filepaths:
db.load_multiple(filepaths)
try:
filepaths = _resolve_tag_filepaths(files_to_load)
if filepaths:
db.load_multiple(filepaths)
except Exception as exc:
_database_load_error = str(exc)
raise
finally:
_database_loading = False
_database_loading_files = []
return db
def start_database_warmup(
tag_file: str = DEFAULT_TAG_FILE,
extra_files: Optional[List[str]] = None,
) -> bool:
global _database_loading, _database_load_error, _database_loading_files
files_to_load = _build_tag_file_list(tag_file, extra_files)
db = get_database()
with _database_load_lock:
if db.is_loaded and set(db._current_files) == set(files_to_load):
return False
if _database_loading:
return False
_database_loading = True
_database_load_error = None
_database_loading_files = files_to_load[:]
def warmup_worker():
try:
ensure_database_loaded(tag_file, extra_files=extra_files)
except Exception as exc:
print(f"[Autocomplete] Warm-up failed: {exc}")
thread = threading.Thread(
target=warmup_worker,
name="A1111AutocompleteWarmup",
daemon=True,
)
thread.start()
return True
def start_wildcard_warmup() -> bool:
global _wildcard_warmup_loading
db = get_wildcard_database()
with _wildcard_warmup_lock:
if db.is_loaded or _wildcard_warmup_loading:
return False
_wildcard_warmup_loading = True
def warmup_worker():
global _wildcard_warmup_loading
try:
db.ensure_loaded()
except Exception as exc:
print(f"[Autocomplete] Wildcard warm-up failed: {exc}")
finally:
with _wildcard_warmup_lock:
_wildcard_warmup_loading = False
thread = threading.Thread(
target=warmup_worker,
name="A1111WildcardWarmup",
daemon=True,
)
thread.start()
return True
# Register API endpoints (only if server is available)
if _HAS_SERVER and PromptServer:
start_wildcard_warmup()
@PromptServer.instance.routes.post("/a1111_prompt/autocomplete")
async def autocomplete_tags(request):
@@ -338,22 +827,52 @@ if _HAS_SERVER and PromptServer:
except Exception:
return web.json_response({"error": "Invalid JSON"}, status=400)
query = data.get("query", "")
limit = min(data.get("limit", 20), 100) # Cap at 100
tag_file = data.get(
"tag_file"
) # Will use DEFAULT_TAG_FILE in ensure_database_loaded if None
if not isinstance(data, dict):
return web.json_response(
{"error": "Request body must be a JSON object"}, status=400
)
# Normalize extra_files: must be a list of strings
extra_files = data.get("extra_files")
if extra_files is None:
extra_files = ["extra-quality-tags.csv"]
elif isinstance(extra_files, str):
extra_files = [extra_files]
elif not isinstance(extra_files, list):
extra_files = ["extra-quality-tags.csv"]
query = _coerce_text(data.get("query", ""))
if len(query) > 255:
return web.json_response(
{"results": [], "tag_count": 0, "error": "Query too long"}
)
mode = str(data.get("mode", "tag")).lower()
if mode not in {"tag", "wildcard", "wildcard_contents"}:
mode = "tag"
limit_cap = 500 if mode in {"wildcard", "wildcard_contents"} else 100
limit = _coerce_limit(data.get("limit", 20), default=20, cap=limit_cap)
tag_file = _normalize_tag_filename(data.get("tag_file"), DEFAULT_TAG_FILE)
extra_files = _normalize_extra_files(data.get("extra_files"))
search_aliases = _coerce_bool(data.get("search_aliases", True), default=True)
contains_fallback = _coerce_bool(
data.get("contains_fallback", True), default=True
)
search_aliases = data.get("search_aliases", True)
if mode == "wildcard":
db = get_wildcard_database()
results = db.search(query, limit=limit)
return web.json_response(
{"results": results, "tag_count": db.wildcard_count, "mode": "wildcard"}
)
if mode == "wildcard_contents":
db = get_wildcard_database()
wildcard_name = _coerce_text(data.get("wildcard_name", ""))
content_query = _coerce_text(data.get("content_query", ""))
results = db.get_contents(
wildcard_name=wildcard_name,
content_query=content_query,
limit=limit,
)
return web.json_response(
{
"results": results,
"tag_count": len(results),
"mode": "wildcard_contents",
"wildcard_name": wildcard_name,
}
)
# Ensure database is loaded
db = ensure_database_loaded(
@@ -361,9 +880,28 @@ if _HAS_SERVER and PromptServer:
)
# Perform search
results = db.search(query, limit=limit, search_aliases=search_aliases)
results = db.search(
query,
limit=limit,
search_aliases=search_aliases,
contains_fallback=contains_fallback,
)
return web.json_response({"results": results, "tag_count": db.tag_count})
return web.json_response(
{"results": results, "tag_count": db.tag_count, "mode": "tag"}
)
@PromptServer.instance.routes.get("/a1111_prompt/autocomplete/wildcards")
async def autocomplete_wildcards(request):
db = get_wildcard_database()
results = await asyncio.to_thread(db.get_manifest)
return web.json_response(
{
"results": results,
"wildcard_count": len(results),
"mode": "wildcard_manifest",
}
)
@PromptServer.instance.routes.get("/a1111_prompt/autocomplete/status")
async def autocomplete_status(request):
@@ -378,13 +916,29 @@ if _HAS_SERVER and PromptServer:
}
"""
db = get_database()
return web.json_response(
{
"loaded": db.is_loaded,
"tag_count": db.tag_count,
"current_files": db._current_files,
}
return web.json_response(get_database_status())
@PromptServer.instance.routes.post("/a1111_prompt/autocomplete/warmup")
async def autocomplete_warmup(request):
try:
data = await request.json()
except Exception:
data = {}
if not isinstance(data, dict):
data = {}
tag_file = _normalize_tag_filename(data.get("tag_file"), DEFAULT_TAG_FILE)
extra_files = _normalize_extra_files(data.get("extra_files"))
started = start_database_warmup(
tag_file or DEFAULT_TAG_FILE, extra_files=extra_files
)
status = get_database_status()
status["started"] = started
status["wildcard_started"] = start_wildcard_warmup()
status["wildcard_loaded"] = get_wildcard_database().is_loaded
status["wildcard_count"] = get_wildcard_database().wildcard_count
return web.json_response(status)
@PromptServer.instance.routes.get("/a1111_prompt/autocomplete/files")
async def list_tag_files(request):
+239 -147
View File
@@ -7,10 +7,32 @@ Uses word-by-word tokenization with manual position tracking.
from aiohttp import web
import server
import logging
import re
try:
from ..parser.wildcards import expand_wildcards_for_token_count
except ImportError:
from parser.wildcards import expand_wildcards_for_token_count
# Lazy-loaded tokenizer instance
_tokenizer = None
logger = logging.getLogger("A1111PromptNode")
def _coerce_text(value) -> str:
if value is None:
return ""
if isinstance(value, str):
return value
return str(value)
def _find_word_position(source: str, word: str, start: int, end: int, fallback: int):
found = source.find(word, start, end)
if found == -1:
return fallback
return found
def get_tokenizer():
@@ -23,6 +45,14 @@ def get_tokenizer():
return _tokenizer
def _encode_word(hf_tokenizer, word: str):
return hf_tokenizer.encode(word, add_special_tokens=False)
def _decode_token(hf_tokenizer, token_id):
return hf_tokenizer.decode([token_id])
def strip_a1111_syntax(text: str) -> str:
r"""
Strip A1111 emphasis/scheduling syntax from text, leaving only tokenizable content.
@@ -46,6 +76,13 @@ def strip_a1111_syntax(text: str) -> str:
text = text.replace("\\[", "\x00LBRACK\x00")
text = text.replace("\\]", "\x00RBRACK\x00")
# Remove weight specifications: :1.2) at end of emphasis
# This handles (word:1.2)
text = re.sub(r":[\d.]+(?=\))", "", text)
# Remove any remaining weight specs that might be floating
text = re.sub(r":[\d.]+(?=\s|$)", "", text)
# Handle bracket expressions: [A|B|C] or [from:to:when]
# Keep the LONGEST option for worst-case token counting
def keep_longest_option(match):
@@ -55,10 +92,8 @@ def strip_a1111_syntax(text: str) -> str:
# It's alternation [A|B|C] - find longest option
options = content.split("|")
# Remove any trailing :number from last option (scheduled alternation)
# e.g., [A|B:0.5] -> options = ["A", "B:0.5"]
last = options[-1]
if ":" in last:
# Check if it ends with :number
colon_match = re.match(r"^(.+?)::?[\d.]+$", last)
if colon_match:
options[-1] = colon_match.group(1)
@@ -70,31 +105,25 @@ def strip_a1111_syntax(text: str) -> str:
# It's scheduling [from:to:when] - keep longer of from/to
parts = content.split(":")
if len(parts) >= 2:
# parts[0] = from, parts[1] = to, parts[2+] = when (numbers)
from_part = parts[0]
to_part = parts[1] if len(parts) > 1 else ""
# Return the longer of from/to
return from_part if len(from_part) >= len(to_part) else to_part
# Compare based on tokenizable length (ignoring parens) for better worst-case estimation
clean_from = from_part.replace("(", "").replace(")", "").strip()
clean_to = to_part.replace("(", "").replace(")", "").strip()
return from_part if len(clean_from) >= len(clean_to) else to_part
return content
else:
# Simple bracket emphasis [word] - just return content
return content
# Process bracket expressions - non-greedy match for innermost brackets
# Repeat to handle nested structures
prev_text = None
while prev_text != text:
prev_text = text
text = re.sub(r"\[([^\[\]]*)\]", keep_longest_option, text)
# Remove weight specifications: :1.2) at end of emphasis
# This handles (word:1.2)
text = re.sub(r":[\d.]+(?=\))", "", text)
# Remove any remaining weight specs that might be floating
text = re.sub(r":[\d.]+(?=\s|$)", "", text)
# Remove parentheses (they're just syntax markers now)
text = text.replace("(", " ")
text = text.replace(")", " ")
@@ -111,6 +140,164 @@ def strip_a1111_syntax(text: str) -> str:
return text.strip()
def build_token_info(text: str, counting_text: str, hf_tokenizer):
is_estimated = counting_text != text
# Split by BREAK first (matching parser.py behavior)
break_pattern = r"\s*\bBREAK\b\s*"
break_matches = list(re.finditer(break_pattern, counting_text))
break_segments = re.split(break_pattern, counting_text)
sequences = []
boundaries = []
tokens_detail = [] # Per-token info: [{text, id, chunk_idx}, ...]
# Track position in token-counting text and display/original text separately.
current_text_offset = 0
display_search_start = 0
current_chunk_idx = 0
for seg_idx, segment in enumerate(break_segments):
# Find where this segment starts in the counting text.
if segment:
segment_start = counting_text.find(segment, current_text_offset)
else:
segment_start = current_text_offset
leading_trim = len(segment) - len(segment.lstrip())
segment_text_start = segment_start + leading_trim
segment_text = segment.strip()
if not segment_text:
sequences.append(0)
current_chunk_idx += 1
else:
clean_segment = strip_a1111_syntax(segment_text)
if not clean_segment:
sequences.append(0)
current_chunk_idx += 1
current_text_offset = segment_start + len(segment)
continue
word_pattern = r"(\S+)"
words_with_pos = []
for match in re.finditer(word_pattern, clean_segment):
word = match.group(1)
word_start = match.start()
word_end = match.end()
words_with_pos.append((word, word_start, word_end))
chunk_size = 75
current_chunk_tokens = 0
chunk_sequences = []
counting_search_start = segment_text_start
segment_end = segment_start + len(segment)
for word, word_start_rel, _word_end_rel in words_with_pos:
word_tokens = _encode_word(hf_tokenizer, word)
word_token_count = len(word_tokens)
if is_estimated:
found = text.find(word, display_search_start)
if found != -1:
original_word_start = found
display_search_start = found + len(word)
else:
original_word_start = min(display_search_start, len(text))
else:
# Map the cleaned word back into the original segment. This keeps
# boundary markers aligned after stripping A1111 syntax.
fallback_pos = segment_text_start + word_start_rel
original_word_start = _find_word_position(
counting_text,
word,
counting_search_start,
segment_end,
fallback_pos,
)
counting_search_start = max(
counting_search_start, original_word_start + len(word)
)
if (
current_chunk_tokens + word_token_count > chunk_size
and current_chunk_tokens > 0
):
chunk_sequences.append(current_chunk_tokens)
boundaries.append(
{
"char_pos": original_word_start,
"type": "chunk",
"from_chunk": current_chunk_idx,
"to_chunk": current_chunk_idx + 1,
"estimated": is_estimated,
}
)
current_chunk_tokens = 0
current_chunk_idx += 1
for token_id in word_tokens:
tokens_detail.append(
{
"text": _decode_token(hf_tokenizer, token_id),
"id": token_id,
"chunk": current_chunk_idx,
}
)
current_chunk_tokens += word_token_count
if current_chunk_tokens > 0:
chunk_sequences.append(current_chunk_tokens)
sequences.extend(chunk_sequences if chunk_sequences else [0])
current_chunk_idx += 1
current_text_offset = segment_start + len(segment)
if seg_idx < len(break_matches):
break_char_pos = break_matches[seg_idx].start()
if is_estimated:
break_char_pos = min(display_search_start, len(text))
boundaries.append(
{
"char_pos": break_char_pos,
"type": "break",
"from_chunk": current_chunk_idx - 1,
"to_chunk": current_chunk_idx,
"estimated": is_estimated,
}
)
tokens_detail.append(
{
"text": "BREAK",
"id": None,
"chunk": current_chunk_idx - 1,
"is_break": True,
}
)
if not sequences:
sequences = [0]
return {
"sequences": sequences,
"boundaries": boundaries,
"tokens": tokens_detail,
"stats": {
"total_tokens": sum(sequences),
"chunks": len(sequences),
"words": len(text.split()),
"characters": len(text),
"estimated_from_wildcards": is_estimated,
},
"estimated_text": counting_text if is_estimated else None,
}
@server.PromptServer.instance.routes.post("/a1111_prompt/tokenize")
async def tokenize_prompt(request):
"""
@@ -120,145 +307,50 @@ async def tokenize_prompt(request):
where boundaries fall using word-by-word tokenization.
BREAK forces a new sequence (matching A1111/parser behavior).
"""
data = await request.json()
text = data.get("text", "")
try:
try:
data = await request.json()
except Exception:
return web.json_response(
{"error": "Invalid JSON", "sequences": None, "boundaries": None},
status=400,
)
if not isinstance(data, dict):
return web.json_response(
{
"error": "Request body must be a JSON object",
"sequences": None,
"boundaries": None,
},
status=400,
)
text = _coerce_text(data.get("text", ""))
tokenizer = get_tokenizer()
hf_tokenizer = tokenizer.tokenizer # Access underlying HuggingFace tokenizer
# Split by BREAK first (matching parser.py behavior)
break_pattern = r"\s*\bBREAK\b\s*"
break_matches = list(re.finditer(break_pattern, text))
break_segments = re.split(break_pattern, text)
def score_token_count(candidate: str) -> int:
clean_candidate = strip_a1111_syntax(candidate)
if not clean_candidate:
return 0
return len(hf_tokenizer.encode(clean_candidate, add_special_tokens=False))
sequences = []
boundaries = []
tokens_detail = [] # Per-token info: [{text, id, chunk_idx}, ...]
try:
counting_text = expand_wildcards_for_token_count(
text,
scorer=score_token_count,
)
except Exception as exc:
logger.warning(
"[A1111 Prompt] Wildcard token estimate failed; falling back to literal token count: %s",
exc,
)
counting_text = text
is_estimated = counting_text != text
# Track position in original text
current_text_offset = 0
current_chunk_idx = 0
for seg_idx, segment in enumerate(break_segments):
# Find where this segment starts in the original text
if segment:
segment_start = text.find(segment, current_text_offset)
else:
segment_start = current_text_offset
segment_text = segment.strip()
if not segment_text:
sequences.append(0)
current_chunk_idx += 1
else:
# FIRST: Strip A1111 syntax from entire segment
# This handles bracket expressions with spaces like [A|B C|D]
clean_segment = strip_a1111_syntax(segment_text)
if not clean_segment:
sequences.append(0)
current_chunk_idx += 1
current_text_offset = segment_start + len(segment)
continue
# Word-by-word tokenization with position tracking
# Split into words from the CLEANED segment
word_pattern = r"(\S+)"
words_with_pos = []
for match in re.finditer(word_pattern, clean_segment):
word = match.group(1)
# Position relative to segment start
word_start = match.start()
word_end = match.end()
words_with_pos.append((word, word_start, word_end))
# Tokenize each word and track cumulative token count
chunk_size = 75
current_chunk_tokens = 0
chunk_sequences = []
for word, word_start_rel, word_end_rel in words_with_pos:
# Words are already from the cleaned segment, just tokenize directly
word_tokens = hf_tokenizer.encode(word, add_special_tokens=False)
word_token_count = len(word_tokens)
# Check if adding this word would exceed chunk size
if (
current_chunk_tokens + word_token_count > chunk_size
and current_chunk_tokens > 0
):
# Save current chunk
chunk_sequences.append(current_chunk_tokens)
# Add boundary marker at the END of the previous word
# (which is the start of this word)
boundary_char = segment_start + word_start_rel
boundaries.append({"char_pos": boundary_char, "type": "chunk"})
current_chunk_tokens = 0
current_chunk_idx += 1
# Decode each token to get the text representation
for token_id in word_tokens:
token_text = hf_tokenizer.decode([token_id])
tokens_detail.append(
{
"text": token_text,
"id": token_id,
"chunk": current_chunk_idx,
}
)
current_chunk_tokens += word_token_count
# Don't forget the last chunk
if current_chunk_tokens > 0:
chunk_sequences.append(current_chunk_tokens)
sequences.extend(chunk_sequences if chunk_sequences else [0])
current_chunk_idx += 1
# Move past this segment
current_text_offset = segment_start + len(segment)
# Add BREAK boundary if not the last segment
if seg_idx < len(break_matches):
boundaries.append(
{"char_pos": break_matches[seg_idx].start(), "type": "break"}
)
# Add BREAK marker to tokens
tokens_detail.append(
{
"text": "BREAK",
"id": None,
"chunk": current_chunk_idx - 1,
"is_break": True,
}
)
if not sequences:
sequences = [0]
# Calculate stats
word_count = len(text.split())
char_count = len(text)
return web.json_response(
{
"sequences": sequences,
"boundaries": boundaries,
"tokens": tokens_detail,
"stats": {
"total_tokens": sum(sequences),
"chunks": len(sequences),
"words": word_count,
"characters": char_count,
},
}
)
return web.json_response(build_token_info(text, counting_text, hf_tokenizer))
except Exception as e:
import traceback
Binary file not shown.

After

Width:  |  Height:  |  Size: 60 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 81 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 363 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 68 KiB

+392 -91
View File
@@ -29,101 +29,270 @@ app.registerExtension({
name: "A1111PromptNode.ShowText",
async beforeRegisterNodeDef(nodeType, nodeData, app) {
if (nodeData.name === "A1111Prompt" || nodeData.name === "A1111PromptNegative") {
/**
* Populate the node with a readonly text widget showing the prompt
*/
function populate(text) {
// Find or create the display widget
let displayWidget = this.widgets?.find(
(w) => w.name === "_prompt_display"
);
// Get the actual text value (could be an array)
const textValue = Array.isArray(text) ? text[0] : text;
const currentText = this.widgets?.find((w) => w.name === "text")?.value || "";
// Only show if the effective prompt is different from the input text
// This avoids clutter when no expansion (TIPO, etc.) is happening
if (!textValue || textValue === currentText) {
if (displayWidget) {
// Hide the widget if it exists but is no longer needed
displayWidget.type = "converted-widget"; // Effectively hides it in ComfyUI
if (displayWidget.inputEl) displayWidget.inputEl.style.display = "none";
function refreshNodeLayout(node) {
requestAnimationFrame(() => {
try {
node.onResize?.(node.size);
node.setDirtyCanvas?.(true, true);
app.graph?.setDirtyCanvas?.(true, true);
} catch (error) {
console.error("[A1111 Prompt] Failed to refresh preview layout", error);
}
});
}
function setHiddenWidgetLayout(widget, hidden) {
if (hidden) {
if (!Object.hasOwn(widget, "_a1111OriginalComputeLayoutSize")) {
widget._a1111OriginalComputeLayoutSize = widget.computeLayoutSize;
}
widget.computeLayoutSize = () => ({
minHeight: 0,
minWidth: 0,
maxHeight: 0,
maxWidth: 0,
});
return;
}
if (!Object.hasOwn(widget, "_a1111OriginalComputeLayoutSize")) {
return;
}
if (widget._a1111OriginalComputeLayoutSize) {
widget.computeLayoutSize = widget._a1111OriginalComputeLayoutSize;
} else {
delete widget.computeLayoutSize;
}
delete widget._a1111OriginalComputeLayoutSize;
}
function getCurrentText(node) {
const textWidget = node.widgets?.find((w) => w.name === "text");
return textWidget?.inputEl?.value ?? textWidget?.value ?? "";
}
const PREVIEW_PROPERTY = "a1111_resolved_prompt";
const PREVIEW_EXPANDED_PROPERTY = "a1111_resolved_prompt_expanded";
const PREVIEW_SOURCE_PROPERTY = "a1111_resolved_prompt_source";
const PREVIEW_WIDGET_HEIGHT = 120;
function getPreviewText(node) {
return node.properties?.[PREVIEW_PROPERTY] || "";
}
function isPreviewExpanded(node) {
return Boolean(node.properties?.[PREVIEW_EXPANDED_PROPERTY]);
}
function setPreviewExpanded(node, expanded) {
node.properties ??= {};
node.properties[PREVIEW_EXPANDED_PROPERTY] = Boolean(expanded);
}
function setPreviewText(node, value, sourceText = "") {
node.properties ??= {};
if (value) {
node.properties[PREVIEW_PROPERTY] = value;
node.properties[PREVIEW_SOURCE_PROPERTY] = sourceText;
} else {
delete node.properties[PREVIEW_PROPERTY];
delete node.properties[PREVIEW_EXPANDED_PROPERTY];
delete node.properties[PREVIEW_SOURCE_PROPERTY];
}
}
function getSavedPreviewValue(node, info) {
const previewFromProperties =
info?.properties?.[PREVIEW_PROPERTY] ??
node.properties?.[PREVIEW_PROPERTY] ??
"";
if (previewFromProperties) return previewFromProperties;
const widgetValues = info?.widgets_values ?? node.widgets_values;
if (!Array.isArray(widgetValues)) return "";
// Compatibility with workflows saved while _prompt_display was serialized
// as an extra widget value.
return widgetValues.length > 3 ? widgetValues[3] || "" : "";
}
function getSavedPreviewSource(node, info) {
return (
info?.properties?.[PREVIEW_SOURCE_PROPERTY] ??
node.properties?.[PREVIEW_SOURCE_PROPERTY] ??
getCurrentText(node)
);
}
function setWidgetHidden(widget, hidden) {
if (!widget) return;
if (Object.hasOwn(widget, "_a1111OriginalComputeSize")) {
if (widget._a1111OriginalComputeSize) {
widget.computeSize = widget._a1111OriginalComputeSize;
} else {
delete widget.computeSize;
}
delete widget._a1111OriginalComputeSize;
}
widget.hidden = hidden;
setHiddenWidgetLayout(widget, hidden);
if (widget.inputEl) widget.inputEl.style.display = hidden ? "none" : "";
}
function setPreviewWidgetHidden(widget, hidden) {
if (!widget) return;
delete widget.computeSize;
widget.hidden = hidden;
setHiddenWidgetLayout(widget, hidden);
if (widget.inputEl) {
widget.inputEl.style.display = hidden ? "none" : "";
widget.inputEl.style.height = "100%";
widget.inputEl.style.boxSizing = "border-box";
widget.inputEl.style.border = "1px solid var(--border-color, #4b4664)";
widget.inputEl.style.borderRadius = "4px";
}
}
function ensureToggleWidget(node) {
let toggleWidget = node.widgets?.find(
(w) => w._a1111PreviewToggle || w.name === "_prompt_display_toggle"
);
if (!toggleWidget) {
toggleWidget = node.addWidget(
"button",
"_prompt_display_toggle",
"Show",
() => {}
);
toggleWidget.serialize = false;
}
toggleWidget.name = "_prompt_display_toggle";
toggleWidget._a1111PreviewToggle = true;
toggleWidget.serialize = false;
toggleWidget.type = "button";
toggleWidget.callback = () => {
setPreviewExpanded(node, !isPreviewExpanded(node));
renderPreview(node);
};
return toggleWidget;
}
function ensureDisplayWidget(node) {
let displayWidget = node.widgets?.find(
(w) => w.name === "_prompt_display"
);
if (!displayWidget) {
// Create a new readonly widget to show the prompt
const widgetResult = ComfyWidgets["STRING"](
this,
node,
"_prompt_display",
["STRING", { multiline: true }],
app
);
displayWidget = widgetResult.widget;
}
displayWidget.serialize = false;
displayWidget.type = "customtext";
displayWidget.options ??= {};
displayWidget.options.getMinHeight = () => 0;
displayWidget.options.getMaxHeight = () => PREVIEW_WIDGET_HEIGHT;
if (displayWidget.inputEl) {
displayWidget.inputEl.readOnly = true;
displayWidget.inputEl.style.opacity = "0.7";
displayWidget.inputEl.style.fontStyle = "italic";
displayWidget.inputEl.placeholder =
"(Prompt will appear after execution)";
// Mark it as a display-only widget
displayWidget.serialize = false;
}
// Ensure it's visible if we're populating it
displayWidget.type = "customtext";
if (displayWidget.inputEl) displayWidget.inputEl.style.display = "block";
displayWidget.value = textValue;
return displayWidget;
}
// Resize node to fit the new widget content
requestAnimationFrame(() => {
const sz = this.computeSize();
if (sz[0] < this.size[0]) sz[0] = this.size[0];
if (sz[1] < this.size[1]) sz[1] = this.size[1];
this.onResize?.(sz);
app.graph.setDirtyCanvas(true, false);
});
function hideDisplayWidget(displayWidget, clearValue = true) {
setPreviewWidgetHidden(displayWidget, true);
if (clearValue) displayWidget.value = "";
}
function showDisplayWidget(displayWidget, textValue) {
setPreviewWidgetHidden(displayWidget, false);
displayWidget.value = textValue;
}
function renderPreview(node) {
try {
renderPreviewInner(node);
} catch (error) {
console.error("[A1111 Prompt] Failed to render resolved prompt preview", error);
}
}
function renderPreviewInner(node) {
const previewText = getPreviewText(node);
const toggleWidget = node.widgets?.find(
(w) => w._a1111PreviewToggle || w.name === "_prompt_display_toggle"
);
const displayWidget = node.widgets?.find(
(w) => w.name === "_prompt_display"
);
if (!previewText) {
if (toggleWidget) setWidgetHidden(toggleWidget, true);
if (displayWidget) hideDisplayWidget(displayWidget, true);
refreshNodeLayout(node);
return;
}
const visibleToggle = ensureToggleWidget(node);
const isExpanded = isPreviewExpanded(node);
visibleToggle.label = isExpanded
? "Hide"
: "Show";
visibleToggle.value = visibleToggle.label;
setWidgetHidden(visibleToggle, false);
if (isExpanded) {
showDisplayWidget(ensureDisplayWidget(node), previewText);
} else if (displayWidget) {
hideDisplayWidget(displayWidget, false);
}
refreshNodeLayout(node);
}
/**
* Populate the node with a readonly text widget showing the prompt.
*/
function populate(text) {
const textValue = Array.isArray(text) ? text[0] : text;
const currentText = getCurrentText(this);
// Only show if the effective prompt is different from the input text.
// This avoids clutter when no expansion (TIPO, etc.) is happening.
if (!textValue || textValue === currentText) {
setPreviewText(this, "");
renderPreview(this);
return;
}
setPreviewText(this, textValue, currentText);
renderPreview(this);
}
// Hook into onExecuted to display the prompt after execution
const onExecuted = nodeType.prototype.onExecuted;
nodeType.prototype.onExecuted = function (message) {
onExecuted?.apply(this, arguments);
if (message?.text) {
populate.call(this, message.text);
}
populate.call(this, message?.text ?? "");
};
// Store widget values during configure for workflow reload
const VALUES = Symbol();
const configure = nodeType.prototype.configure;
nodeType.prototype.configure = function () {
this[VALUES] = arguments[0]?.widgets_values;
return configure?.apply(this, arguments);
};
// Restore display widget on workflow load
const onConfigure = nodeType.prototype.onConfigure;
nodeType.prototype.onConfigure = function () {
nodeType.prototype.onConfigure = function (info) {
onConfigure?.apply(this, arguments);
const widgets_values = this[VALUES];
// Look for stored prompt display value
if (widgets_values?.length > 1) {
requestAnimationFrame(() => {
// The display widget value might be stored after the main text widget
const displayValue = widgets_values.find(
(v) =>
typeof v === "string" &&
v.length > 0 &&
this.widgets?.find((w) => w.name === "text")?.value !== v
);
if (displayValue) {
populate.call(this, [displayValue]);
}
});
}
requestAnimationFrame(() => {
const savedPreview = getSavedPreviewValue(this, info);
if (savedPreview) setPreviewText(this, savedPreview, getSavedPreviewSource(this, info));
renderPreview(this);
});
};
}
},
@@ -288,13 +457,32 @@ app.registerExtension({
nodeType.prototype.onRemoved = function () {
if (onRemoved) onRemoved.apply(this, arguments);
// Disconnect observer to prevent leaks
if (this._boundaryTokenTimeout) {
clearTimeout(this._boundaryTokenTimeout);
this._boundaryTokenTimeout = null;
}
const textarea = this.widgets?.find((w) => w.name === "text")?.inputEl;
if (textarea && this._boundaryInputHandler) {
textarea.removeEventListener("input", this._boundaryInputHandler);
this._boundaryInputHandler = null;
}
if (textarea && this._boundaryFocusHandler) {
textarea.removeEventListener("focus", this._boundaryFocusHandler);
this._boundaryFocusHandler = null;
}
if (textarea && this._boundaryScrollHandler) {
textarea.removeEventListener("scroll", this._boundaryScrollHandler);
this._boundaryScrollHandler = null;
}
if (this._boundaryObserver) {
this._boundaryObserver.disconnect();
this._boundaryObserver = null;
}
// Cleanup overlay DOM
if (this._overlayContainer && this._overlayContainer.parentNode) {
this._overlayContainer.parentNode.removeChild(this._overlayContainer);
this._overlayContainer = null;
@@ -317,11 +505,112 @@ app.registerExtension({
// Create boundary marker overlay
let overlayContainer = null;
let mirrorDiv = null;
let activeTextarea = null;
let tokenUpdateTimeout = null;
let handleInput = null;
let handleFocus = null;
const cleanupOverlay = () => {
if (node._boundaryBoundTextarea && node._boundaryInputHandler) {
node._boundaryBoundTextarea.removeEventListener(
"input",
node._boundaryInputHandler
);
}
if (node._boundaryBoundTextarea && node._boundaryFocusHandler) {
node._boundaryBoundTextarea.removeEventListener(
"focus",
node._boundaryFocusHandler
);
}
if (activeTextarea && node._boundaryScrollHandler) {
activeTextarea.removeEventListener("scroll", node._boundaryScrollHandler);
node._boundaryScrollHandler = null;
}
node._boundaryBoundTextarea = null;
node._boundaryInputHandler = null;
node._boundaryFocusHandler = null;
if (node._boundaryObserver) {
node._boundaryObserver.disconnect();
node._boundaryObserver = null;
}
if (overlayContainer?.parentNode) {
overlayContainer.parentNode.removeChild(overlayContainer);
}
overlayContainer = null;
mirrorDiv = null;
activeTextarea = null;
node._overlayContainer = null;
};
const bindTextarea = (textarea) => {
if (!textarea || node._boundaryBoundTextarea === textarea) return;
if (node._boundaryBoundTextarea && node._boundaryInputHandler) {
node._boundaryBoundTextarea.removeEventListener(
"input",
node._boundaryInputHandler
);
}
if (node._boundaryBoundTextarea && node._boundaryFocusHandler) {
node._boundaryBoundTextarea.removeEventListener(
"focus",
node._boundaryFocusHandler
);
}
if (!handleInput) {
handleInput = () => {
updateTokenCount(textWidget.inputEl?.value || "");
requestAnimationFrame(updateBoundaryMarkers);
};
}
if (!handleFocus) {
handleFocus = () => {
ensureOverlay();
requestAnimationFrame(updateBoundaryMarkers);
};
}
textarea.addEventListener("input", handleInput);
textarea.addEventListener("focus", handleFocus);
node._boundaryInputHandler = handleInput;
node._boundaryFocusHandler = handleFocus;
node._boundaryBoundTextarea = textarea;
};
const ensureOverlay = () => {
const textarea = textWidget.inputEl;
if (!textarea) return false;
if (
textarea !== activeTextarea ||
!overlayContainer ||
!overlayContainer.isConnected ||
!mirrorDiv
) {
createOverlay();
}
bindTextarea(textarea);
return Boolean(mirrorDiv);
};
const createOverlay = () => {
if (!textWidget.inputEl) return;
const textarea = textWidget.inputEl;
cleanupOverlay();
activeTextarea = textarea;
// Create container for the overlay
overlayContainer = document.createElement("div");
@@ -354,6 +643,7 @@ app.registerExtension({
overlayContainer.appendChild(mirrorDiv);
// Insert overlay before textarea so it appears behind
if (!textarea.parentNode) return;
textarea.parentNode.style.position = "relative";
textarea.parentNode.insertBefore(overlayContainer, textarea);
node._overlayContainer = overlayContainer;
@@ -395,27 +685,26 @@ app.registerExtension({
copyStyles();
// Sync scroll position
textarea.addEventListener("scroll", () => {
const syncScroll = () => {
if (!mirrorDiv) return;
mirrorDiv.scrollTop = textarea.scrollTop;
mirrorDiv.scrollLeft = textarea.scrollLeft;
});
};
textarea.addEventListener("scroll", syncScroll);
node._boundaryScrollHandler = syncScroll;
// Use ResizeObserver to keep mirror synced with textarea size
if (node._boundaryObserver) {
node._boundaryObserver.disconnect();
}
const observer = new ResizeObserver(() => {
copyStyles();
// Force re-sync of scroll after resize as content might reflow
mirrorDiv.scrollTop = textarea.scrollTop;
mirrorDiv.scrollLeft = textarea.scrollLeft;
syncScroll();
});
observer.observe(textarea);
node._boundaryObserver = observer;
};
const updateBoundaryMarkers = () => {
if (!mirrorDiv || !textWidget.inputEl) return;
if (!ensureOverlay()) return;
const text = textWidget.inputEl.value || "";
const boundaries = node._tokenInfo?.boundaries || [];
@@ -429,23 +718,31 @@ app.registerExtension({
let html = "";
let lastPos = 0;
// Sort boundaries by position
const sortedBoundaries = [...boundaries].sort(
(a, b) => a.char_pos - b.char_pos
);
// Sort boundaries by position. Estimated wildcard boundaries can map
// multiple chunk cuts to the same source position, so equal positions
// must still render separate markers.
const sortedBoundaries = [...boundaries]
.filter((boundary) => Number.isFinite(boundary?.char_pos))
.sort((a, b) => a.char_pos - b.char_pos);
const markersAtPosition = new Map();
for (const boundary of sortedBoundaries) {
const pos = Math.min(boundary.char_pos, text.length);
if (pos <= lastPos) continue;
if (pos < lastPos) continue;
// Add text before boundary
html += escapeHtml(text.slice(lastPos, pos));
if (pos > lastPos) {
html += escapeHtml(text.slice(lastPos, pos));
}
// Add boundary marker
const markerColor =
boundary.type === "break"
? "#3498db" // Blue for BREAK
: "#e67e22"; // Orange for chunk boundary
const markerOpacity = boundary.estimated ? 0.45 : 0.7;
const markerOffset = markersAtPosition.get(pos) || 0;
markersAtPosition.set(pos, markerOffset + 1);
html += `<span style="
display: inline-block;
@@ -458,13 +755,13 @@ app.registerExtension({
"><span style="
display: block;
position: absolute;
left: -1px;
left: ${-1 + markerOffset * 3}px;
top: 0;
width: 2px;
height: 100%;
background: ${markerColor};
border-radius: 1px;
opacity: 0.7;
opacity: ${markerOpacity};
"></span></span>`;
lastPos = pos;
@@ -476,10 +773,9 @@ app.registerExtension({
mirrorDiv.innerHTML = html;
};
let timeout;
const updateTokenCount = async (text) => {
clearTimeout(timeout);
timeout = setTimeout(async () => {
clearTimeout(tokenUpdateTimeout);
tokenUpdateTimeout = setTimeout(async () => {
try {
const response = await fetch("/a1111_prompt/tokenize", {
method: "POST",
@@ -496,25 +792,30 @@ app.registerExtension({
node.setDirtyCanvas(true, false);
// Update boundary markers
requestAnimationFrame(updateBoundaryMarkers);
requestAnimationFrame(() => {
updateBoundaryMarkers();
});
} catch (e) {
// Silently ignore errors, just keep old count
}
}, 300); // 300ms debounce
node._boundaryTokenTimeout = tokenUpdateTimeout;
};
// Hook into widget callback for text changes
const origCallback = textWidget.callback;
textWidget.callback = function (value) {
origCallback?.apply(this, arguments);
updateTokenCount(value);
updateTokenCount(
typeof value === "string" ? value : textWidget.inputEl?.value || ""
);
};
// Create overlay when the textarea is available
const waitForTextarea = () => {
if (textWidget.inputEl) {
createOverlay();
updateTokenCount(textWidget.value || "");
ensureOverlay();
updateTokenCount(textWidget.inputEl.value || textWidget.value || "");
} else {
requestAnimationFrame(waitForTextarea);
}
+595 -88
View File
@@ -17,11 +17,42 @@ let currentTextarea = null;
let currentWordInfo = null;
let selectedIndex = -1;
let currentResults = [];
let debounceTimeout = null;
// Small in-memory response cache. It keeps typing responsive without making the
// backend scan the large tag database for every repeated prefix.
const AUTOCOMPLETE_CACHE_MAX = 80;
const AUTOCOMPLETE_CACHE_TTL_MS = 5 * 60 * 1000;
const autocompleteCache = new Map();
let autocompleteWarmupStarted = false;
let wildcardManifest = null;
let wildcardManifestPromise = null;
// Frequency tracking
const FREQUENCY_STORAGE_KEY = "a1111_tag_frequency";
let tagFrequency = {}; // tag_name -> usage_count
let tagFrequency = {}; // tag_name -> { count, lastUsed }
const RECENCY_WINDOW_DAYS = 30;
function normalizeUsageEntry(entry) {
if (typeof entry === "number") {
return {
count: Math.max(0, Number(entry) || 0),
lastUsed: 0,
};
}
if (entry && typeof entry === "object") {
return {
count: Math.max(0, Number(entry.count) || 0),
lastUsed: Math.max(0, Number(entry.lastUsed) || 0),
};
}
return { count: 0, lastUsed: 0 };
}
function normalizeFrequencyKey(tagName) {
return String(tagName || "").trim().toLowerCase();
}
/**
* Load tag frequency data from localStorage
@@ -30,7 +61,13 @@ function loadTagFrequency() {
try {
const stored = localStorage.getItem(FREQUENCY_STORAGE_KEY);
if (stored) {
tagFrequency = JSON.parse(stored);
const parsed = JSON.parse(stored);
tagFrequency = Object.fromEntries(
Object.entries(parsed || {}).map(([tag, usage]) => [
normalizeFrequencyKey(tag),
normalizeUsageEntry(usage),
])
);
}
} catch (e) {
console.warn("Failed to load tag frequency:", e);
@@ -54,7 +91,13 @@ function saveTagFrequency() {
* @param {string} tagName - The tag name
*/
function incrementTagFrequency(tagName) {
tagFrequency[tagName] = (tagFrequency[tagName] || 0) + 1;
const key = normalizeFrequencyKey(tagName);
if (!key) return;
const current = normalizeUsageEntry(tagFrequency[key]);
tagFrequency[key] = {
count: current.count + 1,
lastUsed: Date.now(),
};
saveTagFrequency();
}
@@ -64,20 +107,26 @@ function incrementTagFrequency(tagName) {
* @returns {number} Usage count
*/
function getTagFrequency(tagName) {
return tagFrequency[tagName] || 0;
return normalizeUsageEntry(tagFrequency[normalizeFrequencyKey(tagName)]).count;
}
/**
* Calculate frequency boost score
* Uses logarithmic scaling to prevent over-boosting
* @param {number} frequency - Usage count
* @returns {number} Boost score
*/
function calculateFrequencyBoost(frequency) {
if (frequency === 0) return 0;
// Logarithmic boost: log2(frequency + 1)
// This gives: 1 use = 1.0, 3 uses = 2.0, 7 uses = 3.0, 15 uses = 4.0
return Math.log2(frequency + 1);
function getTagLastUsed(tagName) {
return normalizeUsageEntry(tagFrequency[normalizeFrequencyKey(tagName)]).lastUsed;
}
function getUsageScore(tagName) {
const count = getTagFrequency(tagName);
if (count <= 0) return 0;
const frequencyScore = Math.log2(count + 1);
const lastUsed = getTagLastUsed(tagName);
if (!lastUsed) {
return frequencyScore;
}
const ageDays = Math.max(0, Date.now() - lastUsed) / (1000 * 60 * 60 * 24);
const recencyScore = Math.max(0, 1 - ageDays / RECENCY_WINDOW_DAYS);
return frequencyScore + recencyScore * 0.75;
}
// Load frequency data on startup
@@ -91,6 +140,255 @@ const TAG_COLORS = {
4: "#90ee90", // character - lightgreen
5: "#ffa500", // meta - orange
};
const WILDCARD_CONTENT_COLOR = "#e6edf3";
function appendWildcardPathLabel(container, wildcardPath) {
const segments = String(wildcardPath || "").split("/").filter(Boolean);
if (segments.length === 0) {
return;
}
const prefixSegments = segments.slice(0, -1);
prefixSegments.forEach((segment) => {
const segmentSpan = document.createElement("span");
segmentSpan.className = "acPathPart";
segmentSpan.textContent = `${segment}/`;
container.appendChild(segmentSpan);
});
const leafSpan = document.createElement("span");
leafSpan.textContent = segments[segments.length - 1];
leafSpan.style.fontWeight = "600";
leafSpan.style.color = "var(--fg-color, #fff)";
container.appendChild(leafSpan);
}
function normalizeAutocompleteQuery(query, kind = "") {
let normalized = String(query || "").trim().toLowerCase();
if (kind === "wildcard_file" || kind === "wildcard_folder") {
if (normalized.startsWith("__")) {
normalized = normalized.slice(2);
}
if (normalized.endsWith("__")) {
normalized = normalized.slice(0, -2);
}
}
return normalized;
}
function getSuggestionRelevance(tag, rawQuery) {
const query = normalizeAutocompleteQuery(rawQuery, tag.kind);
if (!query) return 0;
const candidates = [tag.name];
if (tag.matched_alias) {
candidates.unshift(tag.matched_alias);
}
const normalizedCandidates = candidates.map((candidate) =>
normalizeAutocompleteQuery(candidate, tag.kind)
);
if (normalizedCandidates.some((candidate) => candidate === query)) {
return 3;
}
if (normalizedCandidates.some((candidate) => candidate.startsWith(query))) {
return 2;
}
if (normalizedCandidates.some((candidate) => candidate.includes(query))) {
return 1;
}
return 0;
}
function getSortQuery(request) {
if (!request) {
return "";
}
if (request.mode === "wildcard_contents") {
return request.extra?.content_query || "";
}
return request.query || "";
}
function getRequestQueryKey(request) {
if (!request) {
return "";
}
if (request.mode === "wildcard_contents") {
return request.extra?.content_query || "";
}
return request.query || "";
}
function buildAutocompleteCacheKey(query, limit, mode, extra = {}) {
const keyPayload = {
mode,
query: normalizeAutocompleteQuery(query),
limit,
wildcard_name: extra.wildcard_name || "",
content_query: extra.content_query || "",
};
return JSON.stringify(keyPayload);
}
function setCachedSuggestions(query, limit, mode, extra, results) {
const payload = {
mode,
query: normalizeAutocompleteQuery(query),
limit,
wildcard_name: extra.wildcard_name || "",
content_query: extra.content_query || "",
};
const key = buildAutocompleteCacheKey(query, limit, mode, extra);
autocompleteCache.set(key, {
at: Date.now(),
payload,
results: Array.isArray(results) ? results : [],
});
while (autocompleteCache.size > AUTOCOMPLETE_CACHE_MAX) {
autocompleteCache.delete(autocompleteCache.keys().next().value);
}
}
function getCachedSuggestions(query, limit, mode, extra = {}) {
const key = buildAutocompleteCacheKey(query, limit, mode, extra);
const cached = autocompleteCache.get(key);
if (!cached) {
return null;
}
if (Date.now() - cached.at > AUTOCOMPLETE_CACHE_TTL_MS) {
autocompleteCache.delete(key);
return null;
}
return cached.results.slice();
}
function getCachedPrefixSuggestions(request, limit = 20) {
const rawQuery = getRequestQueryKey(request);
const query = normalizeAutocompleteQuery(rawQuery);
if (!query) {
return null;
}
for (const cached of autocompleteCache.values()) {
if (Date.now() - cached.at > AUTOCOMPLETE_CACHE_TTL_MS) {
continue;
}
if (!cached.payload || cached.payload.mode !== request.mode) {
continue;
}
if (
request.mode === "wildcard_contents" &&
cached.payload.wildcard_name !== (request.extra?.wildcard_name || "")
) {
continue;
}
if (!query.startsWith(cached.payload.query)) {
continue;
}
const matches = cached.results.filter(
(tag) => getSuggestionRelevance(tag, rawQuery) > 0
);
if (matches.length > 0) {
return matches.slice(0, limit);
}
}
return null;
}
function warmupAutocompleteDatabase() {
if (autocompleteWarmupStarted) {
return;
}
autocompleteWarmupStarted = true;
fetch("/a1111_prompt/autocomplete/warmup", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({ mode: "tag" }),
}).catch((error) => {
autocompleteWarmupStarted = false;
console.warn("[A1111 Autocomplete] Warm-up request failed:", error);
});
}
async function loadWildcardManifest() {
if (wildcardManifest) {
return wildcardManifest;
}
if (wildcardManifestPromise) {
return wildcardManifestPromise;
}
wildcardManifestPromise = fetch("/a1111_prompt/autocomplete/wildcards")
.then(async (response) => {
if (!response.ok) {
throw new Error(`Wildcard manifest request failed: ${response.status}`);
}
const data = await response.json();
const results = Array.isArray(data.results) ? data.results : [];
results.forEach((entry, index) => {
entry._searchText = String(entry.name || "").toLowerCase();
entry._manifestOrder = index;
});
wildcardManifest = results;
return wildcardManifest;
})
.catch((error) => {
wildcardManifestPromise = null;
console.warn("[A1111 Autocomplete] Wildcard manifest load failed:", error);
return null;
});
return wildcardManifestPromise;
}
function searchWildcardManifest(entries, rawQuery, limit = 200) {
const query = normalizeAutocompleteQuery(rawQuery, "wildcard_file");
if (!query) {
return entries.slice(0, limit);
}
const matches = [];
for (const entry of entries) {
const searchText = entry._searchText;
let relevance = 0;
if (searchText === query) {
relevance = 3;
} else if (searchText.startsWith(query)) {
relevance = 2;
} else if (searchText.includes(query)) {
relevance = 1;
} else {
continue;
}
matches.push({ entry, relevance });
}
matches.sort((left, right) => {
if (left.relevance !== right.relevance) {
return right.relevance - left.relevance;
}
const leftFolder = left.entry.kind === "wildcard_folder" ? 0 : 1;
const rightFolder = right.entry.kind === "wildcard_folder" ? 0 : 1;
if (leftFolder !== rightFolder) {
return leftFolder - rightFolder;
}
if (left.entry.name.length !== right.entry.name.length) {
return left.entry.name.length - right.entry.name.length;
}
return left.entry._manifestOrder - right.entry._manifestOrder;
});
return matches.slice(0, limit).map((match) => match.entry);
}
/**
* Create the autocomplete popup element
@@ -169,6 +467,15 @@ export function createAutocompletePopup() {
border-left: 3px solid #4a9eff;
padding-left: 13px !important;
}
.autocomplete-item .acPathPart:nth-child(3n+1) {
color: var(--live-translation-color-1, lightskyblue);
}
.autocomplete-item .acPathPart:nth-child(3n+2) {
color: var(--live-translation-color-2, palegoldenrod);
}
.autocomplete-item .acPathPart:nth-child(3n+3) {
color: var(--live-translation-color-3, darkseagreen);
}
`;
document.head.appendChild(style);
@@ -282,14 +589,14 @@ function getCaretCoordinates(element, position) {
* @param {HTMLElement} textarea - The textarea element
* @param {Object} wordInfo - Current word information
*/
export function showAutocompleteSuggestions(suggestions, textarea, wordInfo) {
export function showAutocompleteSuggestions(suggestions, textarea, wordInfo, request = null) {
if (!suggestions || suggestions.length === 0) {
hideAutocompletePopup();
return;
}
const popup = createAutocompletePopup();
const sortedSuggestions = sortByFrequency(suggestions);
const sortedSuggestions = sortByFrequency(suggestions, getSortQuery(request));
currentResults = sortedSuggestions;
selectedIndex = -1;
@@ -303,8 +610,11 @@ export function showAutocompleteSuggestions(suggestions, textarea, wordInfo) {
item.className = "autocomplete-item";
const nameSpan = document.createElement("span");
const isWildcardFile = tag.kind === "wildcard_file";
const isWildcardFolder = tag.kind === "wildcard_folder";
const isWildcardContent = tag.kind === "wildcard_content";
nameSpan.style.cssText = `
color: ${TAG_COLORS[tag.type] || "var(--fg-color, #ccc)"};
color: ${isWildcardContent ? WILDCARD_CONTENT_COLOR : TAG_COLORS[tag.type] || "var(--fg-color, #ccc)"};
font-weight: 500;
flex: 1;
white-space: nowrap;
@@ -312,12 +622,17 @@ export function showAutocompleteSuggestions(suggestions, textarea, wordInfo) {
text-overflow: ellipsis;
`;
let displayName = tag.name;
if (tag.matched_alias) {
let displayName = (isWildcardFile || isWildcardFolder) ? `__${tag.name}__` : tag.name;
if (isWildcardFile || isWildcardFolder) {
nameSpan.textContent = "";
appendWildcardPathLabel(nameSpan, tag.name);
} else if (tag.matched_alias) {
displayName = `${tag.matched_alias} → ${tag.name}`;
nameSpan.style.fontStyle = "italic";
nameSpan.textContent = displayName;
} else {
nameSpan.textContent = displayName;
}
nameSpan.textContent = displayName;
const frequency = getTagFrequency(tag.name);
if (frequency > 0) {
@@ -340,7 +655,11 @@ export function showAutocompleteSuggestions(suggestions, textarea, wordInfo) {
font-size: 11px;
margin-left: 8px;
`;
countSpan.textContent = formatCount(tag.count);
if (isWildcardFile || isWildcardFolder || isWildcardContent) {
countSpan.textContent = tag.meta || "Wildcard file";
} else {
countSpan.textContent = formatCount(tag.count);
}
item.appendChild(nameSpan);
item.appendChild(countSpan);
@@ -369,30 +688,48 @@ export function showAutocompleteSuggestions(suggestions, textarea, wordInfo) {
/**
* Sort suggestions by frequency and relevance
* @param {Array} suggestions - Array of tag suggestions
* @param {string} query - Active query text for relevance grouping
* @returns {Array} Sorted suggestions
*/
function sortByFrequency(suggestions) {
return suggestions.map(tag => {
function sortByFrequency(suggestions, query = "") {
return suggestions.map((tag, index) => {
const frequency = getTagFrequency(tag.name);
const frequencyBoost = calculateFrequencyBoost(frequency);
const count = Number(tag.count) || 0;
const usageScore = getUsageScore(tag.name);
const relevance = getSuggestionRelevance(tag, query);
// Combined score: base score (from backend) + frequency boost
// Backend already provides relevance score via ordering
// We add frequency boost to prioritize frequently used tags
return {
...tag,
_originalIndex: index,
_frequency: frequency,
_frequencyBoost: frequencyBoost
_count: count,
_usageScore: usageScore,
_relevance: relevance,
};
}).sort((a, b) => {
// Sort by frequency boost first (higher is better)
const freqDiff = b._frequencyBoost - a._frequencyBoost;
if (Math.abs(freqDiff) > 0.5) { // Only prioritize if significant difference
return freqDiff;
const exactDelta = Number(b._relevance === 3) - Number(a._relevance === 3);
if (exactDelta !== 0) {
return exactDelta;
}
// Then by post count (backend relevance)
return b.count - a.count;
const hasUsage = a._usageScore > 0 || b._usageScore > 0;
if (hasUsage && Math.abs(b._usageScore - a._usageScore) > 0.05) {
return b._usageScore - a._usageScore;
}
if (b._relevance !== a._relevance) {
return b._relevance - a._relevance;
}
if (!hasUsage && Math.abs(b._usageScore - a._usageScore) > 0.05) {
return b._usageScore - a._usageScore;
}
if (b._count !== a._count) {
return b._count - a._count;
}
return a._originalIndex - b._originalIndex;
});
}
@@ -452,11 +789,15 @@ function insertTag(index = selectedIndex) {
incrementTagFrequency(tag.name);
const text = textarea.value;
let tagName = tag.name;
const isWildcard = tag.kind === "wildcard_file" || tag.kind === "wildcard_folder";
const isWildcardContent = tag.kind === "wildcard_content";
let tagName = tag.completion || tag.name;
// Configurable A1111 formatting (usually on by default)
tagName = tagName.replace(/_/g, " ");
tagName = tagName.replace(/\(/g, "\\(").replace(/\)/g, "\\)");
if (!isWildcard && !isWildcardContent) {
// Configurable A1111 formatting (usually on by default)
tagName = tagName.replace(/_/g, " ");
tagName = tagName.replace(/\(/g, "\\(").replace(/\)/g, "\\)");
}
const beforeWord = text.substring(0, wordInfo.start);
const afterWord = text.substring(wordInfo.end);
@@ -465,7 +806,7 @@ function insertTag(index = selectedIndex) {
// Only add if not already present
let suffix = "";
const matchAfter = afterWord.match(/^\s*[,:]/);
if (!matchAfter) {
if (!isWildcard && !isWildcardContent && !matchAfter) {
suffix = ", ";
}
@@ -552,8 +893,23 @@ function formatCount(count) {
* @param {number} limit - Maximum number of results
* @returns {Promise<Array>} Array of tag suggestions
*/
export async function fetchTagSuggestions(query, limit = 20) {
if (!query || query.length < 2) return [];
export async function fetchTagSuggestions(query, limit = 20, mode = "tag", extra = {}) {
const isWildcard = mode === "wildcard" || mode === "wildcard_contents";
if (!query) return [];
if (!isWildcard && query.length < 2) return [];
const effectiveLimit = isWildcard ? Math.max(limit, 200) : limit;
if (mode === "wildcard") {
const entries = await loadWildcardManifest();
if (entries) {
return searchWildcardManifest(entries, query, effectiveLimit);
}
}
const cached = getCachedSuggestions(query, effectiveLimit, mode, extra);
if (cached) {
return cached;
}
try {
const response = await fetch("/a1111_prompt/autocomplete", {
@@ -561,9 +917,11 @@ export async function fetchTagSuggestions(query, limit = 20) {
headers: { "Content-Type": "application/json" },
body: JSON.stringify({
query: query,
limit: limit,
limit: effectiveLimit,
mode: mode,
search_aliases: true,
extra_files: ["extra-quality-tags.csv"] // Include quality tags
extra_files: ["extra-quality-tags.csv"], // Include quality tags
...extra,
}),
});
@@ -573,7 +931,9 @@ export async function fetchTagSuggestions(query, limit = 20) {
}
const data = await response.json();
return data.results || [];
const results = data.results || [];
setCachedSuggestions(query, effectiveLimit, mode, extra, results);
return results;
} catch (error) {
console.warn("Autocomplete fetch error:", error);
@@ -618,13 +978,25 @@ export function exportTagFrequency() {
* @returns {Object} Statistics object
*/
export function getFrequencyStats() {
const entries = Object.entries(tagFrequency);
const entries = Object.entries(tagFrequency).map(([tag, usage]) => [
tag,
normalizeUsageEntry(usage),
]);
const totalTags = entries.length;
const totalUses = entries.reduce((sum, [_, count]) => sum + count, 0);
const totalUses = entries.reduce((sum, [_, usage]) => sum + usage.count, 0);
const topTags = entries
.sort((a, b) => b[1] - a[1])
.sort((a, b) => {
if (b[1].count !== a[1].count) {
return b[1].count - a[1].count;
}
return b[1].lastUsed - a[1].lastUsed;
})
.slice(0, 10)
.map(([tag, count]) => ({ tag, count }));
.map(([tag, usage]) => ({
tag,
count: usage.count,
lastUsed: usage.lastUsed,
}));
return {
totalTags,
@@ -673,10 +1045,46 @@ export function getCurrentWord(text, cursorPos) {
};
}
function getAutocompleteRequest(word) {
const completedWildcard = word.match(/^__([^,]+?)__([^, ]*)$/);
if (completedWildcard) {
return {
mode: "wildcard_contents",
query: word,
minLength: 4,
extra: {
wildcard_name: completedWildcard[1],
content_query: completedWildcard[2] || "",
},
};
}
if (word.startsWith("__")) {
return {
mode: "wildcard",
query: word,
minLength: 2,
extra: {},
};
}
return {
mode: "tag",
query: word,
minLength: 2,
extra: {},
};
}
// Extension registration
app.registerExtension({
name: "A1111PromptNode.Autocomplete",
setup() {
warmupAutocompleteDatabase();
loadWildcardManifest();
},
async nodeCreated(node) {
if (
node.comfyClass !== "A1111Prompt" &&
@@ -686,54 +1094,103 @@ app.registerExtension({
const textWidget = node.widgets?.find((w) => w.name === "text");
if (!textWidget) return;
warmupAutocompleteDatabase();
// Wait for textarea to be available
const waitForTextarea = () => {
if (!textWidget.inputEl) {
requestAnimationFrame(waitForTextarea);
let boundTextarea = null;
let monitorFrame = null;
const bindTextarea = (textarea) => {
if (!textarea || textarea === boundTextarea) {
return;
}
const textarea = textWidget.inputEl;
node._autocompleteCleanup?.();
boundTextarea = textarea;
let suppressAutocomplete = false;
let debounceTimeout = null;
let requestVersion = 0;
const scheduleAutocomplete = (delay = 50) => {
clearTimeout(debounceTimeout);
const scheduledVersion = ++requestVersion;
debounceTimeout = setTimeout(async () => {
if (!textarea.isConnected) {
return;
}
const cursorPos = textarea.selectionStart;
const wordInfo = getCurrentWord(textarea.value, cursorPos);
const request = getAutocompleteRequest(wordInfo.word);
if (request.query.length < request.minLength) {
hideAutocompletePopup();
return;
}
const cachedSuggestions =
request.mode === "wildcard"
? null
: getCachedPrefixSuggestions(request);
if (cachedSuggestions && cachedSuggestions.length > 0) {
showAutocompleteSuggestions(
cachedSuggestions,
textarea,
wordInfo,
request
);
}
const suggestions = await fetchTagSuggestions(
request.query,
20,
request.mode,
request.extra
);
if (suggestions === null) {
return;
}
if (scheduledVersion !== requestVersion || !textarea.isConnected) {
return;
}
const latestWordInfo = getCurrentWord(
textarea.value,
textarea.selectionStart
);
if (
latestWordInfo.word !== wordInfo.word ||
latestWordInfo.start !== wordInfo.start ||
latestWordInfo.end !== wordInfo.end
) {
return;
}
if (suggestions.length > 0) {
showAutocompleteSuggestions(suggestions, textarea, latestWordInfo, request);
} else {
hideAutocompletePopup();
}
}, delay);
};
// Add input event listener for autocomplete
textarea.addEventListener("input", async (e) => {
const onInput = async (e) => {
if (suppressAutocomplete) {
suppressAutocomplete = false;
requestVersion++;
hideAutocompletePopup();
return;
}
// Clear existing debounce
if (debounceTimeout) {
clearTimeout(debounceTimeout);
}
// Debounce the search to avoid excessive API calls
debounceTimeout = setTimeout(async () => {
const cursorPos = textarea.selectionStart;
const wordInfo = getCurrentWord(textarea.value, cursorPos);
if (wordInfo.word.length >= 2) {
const suggestions = await fetchTagSuggestions(wordInfo.word);
if (suggestions === null) {
// Error already shown in popup
return;
}
if (suggestions.length > 0) {
showAutocompleteSuggestions(suggestions, textarea, wordInfo);
} else {
hideAutocompletePopup();
}
} else {
hideAutocompletePopup();
}
}, 100); // 100ms debounce - faster response
});
scheduleAutocomplete();
};
textarea.addEventListener("input", onInput);
// Add keydown event listener for navigation
textarea.addEventListener("keydown", (e) => {
const onKeydown = (e) => {
if (e.key === "Backspace") {
const start = textarea.selectionStart;
const end = textarea.selectionEnd;
@@ -757,15 +1214,17 @@ app.registerExtension({
// Event was handled by autocomplete
return;
}
});
};
textarea.addEventListener("keydown", onKeydown);
// Hide popup on blur (with delay to allow clicks)
textarea.addEventListener("blur", () => {
const onBlur = () => {
textarea._a1111_autocomplete_blur_timeout = setTimeout(() => {
hideAutocompletePopup();
textarea._a1111_autocomplete_blur_timeout = null;
}, 200);
});
};
textarea.addEventListener("blur", onBlur);
// Hide popup when clicking outside
const outsideClickListener = (e) => {
@@ -779,20 +1238,68 @@ app.registerExtension({
// Store for cleanup
node._outsideClickListener = outsideClickListener;
node._autocompleteCleanup = () => {
requestVersion++;
clearTimeout(debounceTimeout);
textarea.removeEventListener("input", onInput);
textarea.removeEventListener("keydown", onKeydown);
textarea.removeEventListener("blur", onBlur);
document.removeEventListener("click", outsideClickListener);
if (node._outsideClickListener === outsideClickListener) {
node._outsideClickListener = null;
}
if (textarea._a1111_autocomplete_blur_timeout) {
clearTimeout(textarea._a1111_autocomplete_blur_timeout);
textarea._a1111_autocomplete_blur_timeout = null;
}
if (currentTextarea === textarea) {
hideAutocompletePopup();
}
if (boundTextarea === textarea) {
boundTextarea = null;
}
};
};
requestAnimationFrame(waitForTextarea);
const monitorTextarea = () => {
if (node._autocompleteRemoved) {
return;
}
const textarea = textWidget.inputEl;
if (textarea && textarea.isConnected) {
bindTextarea(textarea);
} else if (boundTextarea) {
node._autocompleteCleanup?.();
}
monitorFrame = requestAnimationFrame(monitorTextarea);
};
monitorFrame = requestAnimationFrame(monitorTextarea);
node._autocompleteMonitorCleanup = () => {
if (monitorFrame) {
cancelAnimationFrame(monitorFrame);
monitorFrame = null;
}
};
// Add cleanup on node removal
const onRemoved = node.onRemoved;
node.onRemoved = function () {
if (onRemoved) onRemoved.apply(this, arguments);
this._autocompleteRemoved = true;
// Remove the specific document click listener for this node
if (this._outsideClickListener) {
document.removeEventListener("click", this._outsideClickListener);
this._outsideClickListener = null;
}
this._autocompleteMonitorCleanup?.();
this._autocompleteMonitorCleanup = null;
this._autocompleteCleanup?.();
this._autocompleteCleanup = null;
};
},
onRemoved() {
+8
View File
@@ -101,6 +101,14 @@ export function showTokenTooltip(node, mouseEvent, graphCanvas) {
</div>
`;
if (info.stats.estimated_from_wildcards) {
html += `
<div style="margin-bottom: 10px; padding: 8px; background: rgba(230, 126, 34, 0.15); border: 1px solid rgba(230, 126, 34, 0.45); border-radius: 4px; color: #f0b27a;">
Count uses the longest wildcard/dynamic prompt estimate.
</div>
`;
}
// Colored tokens section
html += `<div style="margin-bottom: 8px;"><strong>Tokens</strong></div>`;
html += `<div style="max-height: 150px; overflow-y: auto; margin-bottom: 10px; padding: 8px; background: var(--comfy-input-bg, #222); border-radius: 4px; line-height: 1.8;">`;
+85
View File
@@ -0,0 +1,85 @@
"""Helpers for writing A1111 prompt metadata into ComfyUI PNG info."""
from __future__ import annotations
import json
from typing import Any, Dict, Optional
METADATA_KEY = "a1111_prompt_node"
RESOLVED_PROMPT_PROPERTY = "a1111_resolved_prompt"
RESOLVED_PROMPT_EXPANDED_PROPERTY = "a1111_resolved_prompt_expanded"
RESOLVED_PROMPT_SOURCE_PROPERTY = "a1111_resolved_prompt_source"
def _coerce_workflow(extra_pnginfo: Dict[str, Any]) -> Optional[Dict[str, Any]]:
workflow = extra_pnginfo.get("workflow")
if isinstance(workflow, dict):
return workflow
if isinstance(workflow, str):
try:
parsed = json.loads(workflow)
except json.JSONDecodeError:
return None
if isinstance(parsed, dict):
extra_pnginfo["workflow"] = parsed
return parsed
return None
def _update_workflow_preview(
extra_pnginfo: Dict[str, Any],
node_id: str,
original_prompt: str,
resolved_prompt: str,
) -> None:
workflow = _coerce_workflow(extra_pnginfo)
if workflow is None:
return
nodes = workflow.get("nodes")
if not isinstance(nodes, list):
return
for node in nodes:
if not isinstance(node, dict) or str(node.get("id")) != node_id:
continue
properties = node.setdefault("properties", {})
if not isinstance(properties, dict):
properties = {}
node["properties"] = properties
if original_prompt != resolved_prompt:
properties[RESOLVED_PROMPT_PROPERTY] = resolved_prompt
properties[RESOLVED_PROMPT_SOURCE_PROPERTY] = original_prompt
properties.setdefault(RESOLVED_PROMPT_EXPANDED_PROPERTY, False)
else:
properties.pop(RESOLVED_PROMPT_PROPERTY, None)
properties.pop(RESOLVED_PROMPT_EXPANDED_PROPERTY, None)
properties.pop(RESOLVED_PROMPT_SOURCE_PROPERTY, None)
return
def record_prompt_metadata(
extra_pnginfo: Optional[Dict[str, Any]],
node_id: Optional[str],
node_type: str,
original_prompt: str,
resolved_prompt: str,
) -> None:
"""Store original/resolved prompt text in ComfyUI's extra PNG metadata."""
if extra_pnginfo is None or node_id is None:
return
key = str(node_id)
_update_workflow_preview(extra_pnginfo, key, original_prompt, resolved_prompt)
metadata = extra_pnginfo.setdefault(METADATA_KEY, {"nodes": {}})
nodes = metadata.setdefault("nodes", {})
nodes[key] = {
"node_id": key,
"type": node_type,
"original_prompt": original_prompt,
"resolved_prompt": resolved_prompt,
"changed": original_prompt != resolved_prompt,
}
+32 -4
View File
@@ -10,7 +10,8 @@ prompt and logs an info message.
"""
import logging
from ..parser import get_prompt_schedule
from ..parser import expand_wildcards, get_prompt_schedule
from ..metadata import record_prompt_metadata
from .prompt_node import A1111PromptNode
logger = logging.getLogger("A1111PromptNode")
@@ -45,6 +46,10 @@ class A1111PromptNegative:
{"default": False, "label_on": "Enable", "label_off": "Disable"},
),
},
"hidden": {
"unique_id": "UNIQUE_ID",
"extra_pnginfo": "EXTRA_PNGINFO",
},
}
RETURN_TYPES = ("CONDITIONING",)
@@ -53,19 +58,36 @@ class A1111PromptNegative:
CATEGORY = "conditioning/advanced"
OUTPUT_NODE = True
def encode(self, clip, text, normalization=False, debug=False):
def encode(
self,
clip,
text,
normalization=False,
debug=False,
unique_id=None,
extra_pnginfo=None,
):
"""
Encode negative prompt. Uses first step only if scheduling is detected.
"""
if clip is None:
raise RuntimeError("ERROR: clip input is None")
expanded_text = expand_wildcards(text)
record_prompt_metadata(
extra_pnginfo,
unique_id,
"A1111PromptNegative",
text,
expanded_text,
)
is_sdxl = hasattr(clip.cond_stage_model, "clip_l") and hasattr(
clip.cond_stage_model, "clip_g"
)
# Parse the prompt schedule (use dummy steps as it's not supported for negative)
schedule = get_prompt_schedule(text, steps=20)
schedule = get_prompt_schedule(expanded_text, steps=20)
# Check if scheduling/alternation was used
if len(schedule) > 1:
@@ -84,6 +106,12 @@ class A1111PromptNegative:
if len(text) > 100
else f"[A1111 Prompt Negative] Input text: {text}"
)
if expanded_text != text:
logger.info(
f"[A1111 Prompt Negative] Effective text: {expanded_text[:100]}..."
if len(expanded_text) > 100
else f"[A1111 Prompt Negative] Effective text: {expanded_text}"
)
logger.info(
f"[A1111 Prompt Negative] Model: {'SDXL' if is_sdxl else 'SD1.5'}"
)
@@ -102,6 +130,6 @@ class A1111PromptNegative:
)
return {
"ui": {"text": [text]},
"ui": {"text": [expanded_text]},
"result": ([[cond, {"pooled_output": pooled}]],),
}
+33 -11
View File
@@ -14,8 +14,9 @@ import logging
import re
from collections import deque
import comfy.model_management as model_management
from ..parser import get_prompt_schedule
from ..parser import expand_wildcards, get_prompt_schedule
from ..hooks import create_step_schedule_cond
from ..metadata import record_prompt_metadata
logger = logging.getLogger("A1111PromptNode")
@@ -60,6 +61,7 @@ class A1111PromptNode:
"hidden": {
"prompt": "PROMPT",
"unique_id": "UNIQUE_ID",
"extra_pnginfo": "EXTRA_PNGINFO",
},
}
@@ -70,7 +72,14 @@ class A1111PromptNode:
OUTPUT_NODE = True
def encode(
self, clip, text, normalization=False, debug=False, prompt=None, unique_id=None
self,
clip,
text,
normalization=False,
debug=False,
prompt=None,
unique_id=None,
extra_pnginfo=None,
):
"""
Main encode function - A1111 style with step-based scheduling.
@@ -85,6 +94,18 @@ class A1111PromptNode:
logger.info("[A1111 Prompt] ========== ENCODE START ==========")
logger.info(f"[A1111 Prompt] Node ID: {unique_id}")
expanded_text = expand_wildcards(text)
record_prompt_metadata(
extra_pnginfo,
unique_id,
"A1111Prompt",
text,
expanded_text,
)
if debug and expanded_text != text:
logger.info(f"[A1111 Prompt] Wildcards expanded to: {expanded_text}")
# Auto-detect steps from workflow graph
steps = None
if prompt is not None and unique_id is not None:
@@ -106,7 +127,7 @@ class A1111PromptNode:
)
# Check if prompt uses step-based syntax (integers without decimals)
uses_step_syntax = self._uses_step_based_syntax(text, debug)
uses_step_syntax = self._uses_step_based_syntax(expanded_text, debug)
if debug:
logger.info(
f"[A1111 Prompt] Syntax detection: {'step-based' if uses_step_syntax else 'percentage-based'}"
@@ -136,6 +157,12 @@ class A1111PromptNode:
if len(text) > 100
else f"[A1111 Prompt] Input text: {text}"
)
if expanded_text != text:
logger.info(
f"[A1111 Prompt] Effective text: {expanded_text[:100]}..."
if len(expanded_text) > 100
else f"[A1111 Prompt] Effective text: {expanded_text}"
)
logger.info(f"[A1111 Prompt] Model: {'SDXL' if is_sdxl else 'SD1.5'}")
logger.info(
f"[A1111 Prompt] Normalization: {'ON' if normalization else 'OFF'}"
@@ -143,7 +170,7 @@ class A1111PromptNode:
if debug:
logger.info(f"[A1111 Prompt] Parsing schedule with {parse_steps} steps...")
schedule = get_prompt_schedule(text, parse_steps)
schedule = get_prompt_schedule(expanded_text, parse_steps)
if debug:
logger.info(f"[A1111 Prompt] Schedule generated: {len(schedule)} segments")
@@ -262,7 +289,7 @@ class A1111PromptNode:
"[A1111 Prompt] ========== ENCODE COMPLETE (static) =========="
)
return {
"ui": {"text": [text]},
"ui": {"text": [expanded_text]},
"result": ([[cond, {"pooled_output": pooled}]],),
}
@@ -339,11 +366,6 @@ class A1111PromptNode:
# Update run mapping with padded version
run_embeddings[prompt_text] = (cond, pooled)
# ALSO update global cache so we don't re-pad in same run if it's used again
# (though unique_prompts already handled that, it's good for static path too)
key = self._get_cache_key(clip, prompt_text, normalization)
self._encoded_cache[key] = (cond, pooled)
padded_count += 1
if debug:
logger.info(
@@ -405,7 +427,7 @@ class A1111PromptNode:
logger.info(
"[A1111 Prompt] ========== ENCODE COMPLETE (scheduled) =========="
)
return {"ui": {"text": [text]}, "result": (conditioning,)}
return {"ui": {"text": [expanded_text]}, "result": (conditioning,)}
def _uses_step_based_syntax(self, text, debug=False):
"""
+26 -2
View File
@@ -4,12 +4,36 @@ A1111 Prompt Parser Package
Contains the Lark grammar parser and scheduling logic for A1111-style prompts.
"""
from .grammar import GRAMMAR, get_parser, reset_parser
from .scheduler import get_prompt_schedule
from .wildcards import (
expand_wildcards,
expand_wildcards_for_token_count,
get_wildcard_options,
list_available_wildcards,
)
try:
from .grammar import GRAMMAR, get_parser, reset_parser
from .scheduler import get_prompt_schedule
except ImportError as exc:
_PARSER_IMPORT_ERROR = exc
GRAMMAR = None
def get_parser():
raise _PARSER_IMPORT_ERROR
def reset_parser():
return None
def get_prompt_schedule(*args, **kwargs):
raise _PARSER_IMPORT_ERROR
__all__ = [
"GRAMMAR",
"get_parser",
"reset_parser",
"get_prompt_schedule",
"expand_wildcards",
"expand_wildcards_for_token_count",
"get_wildcard_options",
"list_available_wildcards",
]
+264
View File
@@ -0,0 +1,264 @@
"""
Classic prompt tokenization helpers.
These helpers let the node opt into two ReForge/A1111-style behaviors without
replacing the default Comfy tokenization path:
1. Classic attention parsing (`[]` de-emphasis, escaped brackets, BREAK markers)
2. Classic chunking (comma backtrack instead of pure 75-token rollover)
"""
import re
from comfy import sd1_clip as comfy_sd1_clip
_RE_ATTENTION = re.compile(
r"""
\\\(|
\\\)|
\\\[|
\\]|
\\\\|
\\|
\(|
\[|
:\s*([+-]?[.\d]+)\s*\)|
\)|
]|
[^\\()\[\]:]+|
:
""",
re.X,
)
_RE_BREAK = re.compile(r"\s*\bBREAK\b\s*", re.S)
def _parse_classic_attention(text):
res = []
round_brackets = []
square_brackets = []
round_bracket_multiplier = 1.1
square_bracket_multiplier = 1 / 1.1
def multiply_range(start_position, multiplier):
for p in range(start_position, len(res)):
res[p][1] *= multiplier
for match in _RE_ATTENTION.finditer(text):
token_text = match.group(0)
weight = match.group(1)
if token_text.startswith("\\"):
res.append([token_text[1:], 1.0])
elif token_text == "(":
round_brackets.append(len(res))
elif token_text == "[":
square_brackets.append(len(res))
elif weight is not None and round_brackets:
multiply_range(round_brackets.pop(), float(weight))
elif token_text == ")" and round_brackets:
multiply_range(round_brackets.pop(), round_bracket_multiplier)
elif token_text == "]" and square_brackets:
multiply_range(square_brackets.pop(), square_bracket_multiplier)
else:
parts = re.split(_RE_BREAK, token_text)
for i, part in enumerate(parts):
if i > 0:
res.append(["BREAK", -1.0])
if part:
res.append([part, 1.0])
for pos in round_brackets:
multiply_range(pos, round_bracket_multiplier)
for pos in square_brackets:
multiply_range(pos, square_bracket_multiplier)
if len(res) == 0:
return [("", 1.0)]
merged = []
for part_text, weight in res:
if (
merged
and merged[-1][1] == weight
and merged[-1][0] != "BREAK"
and part_text != "BREAK"
):
merged[-1][0] += part_text
else:
merged.append([part_text, weight])
return [(part_text, weight) for part_text, weight in merged if part_text or weight < 0]
def _parse_default_attention(text):
escaped = comfy_sd1_clip.escape_important(text)
parsed = comfy_sd1_clip.token_weights(escaped, 1.0)
out = []
for part_text, weight in parsed:
unescaped = comfy_sd1_clip.unescape_important(part_text)
pieces = re.split(_RE_BREAK, unescaped)
for i, piece in enumerate(pieces):
if i > 0:
out.append(("BREAK", -1.0))
if piece:
out.append((piece, weight))
if not out:
return [("", 1.0)]
return out
def _get_subtokenizers(tokenizer):
subtokenizers = {}
for name in dir(tokenizer):
if not name.startswith("clip_"):
continue
value = getattr(tokenizer, name)
suffix = name.split("clip_", 1)[1]
subtokenizers[suffix] = value
if subtokenizers:
return subtokenizers
# Fallback for tokenizers that expose a single subtokenizer directly.
if hasattr(tokenizer, "tokenize_with_weights"):
return {"l": tokenizer}
raise RuntimeError("Unsupported tokenizer layout for classic prompt tokenization.")
def _segment_to_tokens(subtokenizer, text, weight):
embedding_identifier = subtokenizer.embedding_identifier
split = re.split(r" {0}|\n{0}".format(re.escape(embedding_identifier)), text)
to_tokenize = [split[0]]
for i in range(1, len(split)):
to_tokenize.append("{}{}".format(embedding_identifier, split[i]))
flat_tokens = []
for piece in [x for x in to_tokenize if x != ""]:
if piece.startswith(embedding_identifier) and subtokenizer.embedding_directory is not None:
embedding_name = piece[len(embedding_identifier) :].strip("\n")
embed, leftover = subtokenizer._try_get_embedding(embedding_name)
if embed is not None:
if len(embed.shape) == 1:
flat_tokens.append((embed, weight))
else:
for row in range(embed.shape[0]):
flat_tokens.append((embed[row], weight))
if leftover:
piece = leftover
else:
continue
end = -1 if subtokenizer.tokenizer_adds_end_token else 999999999999
token_ids = subtokenizer.tokenizer(piece)["input_ids"][
subtokenizer.tokens_start : end
]
for token in token_ids:
flat_tokens.append((token, weight))
return flat_tokens
def _find_last_comma(tokens, comma_token):
if comma_token is None:
return -1
for index in range(len(tokens) - 1, -1, -1):
token = tokens[index][0]
if isinstance(token, int) and token == comma_token:
return index
return -1
def _pack_tokens(subtokenizer, parsed_segments, use_classic_chunking):
has_start = subtokenizer.start_token is not None
has_end = subtokenizer.end_token is not None
chunk_length = subtokenizer.max_length - int(has_start) - int(has_end)
comma_token = None
if use_classic_chunking:
comma_token = subtokenizer.tokenizer.get_vocab().get(",</w>", None)
batches = []
current = []
last_comma = -1
def flush(force=False):
nonlocal current
nonlocal last_comma
if not current and not force:
return
batch = []
if has_start:
batch.append((subtokenizer.start_token, 1.0))
batch.extend(current)
if has_end:
batch.append((subtokenizer.end_token, 1.0))
while len(batch) < subtokenizer.max_length:
batch.append((subtokenizer.pad_token, 1.0))
batches.append(batch)
current = []
last_comma = -1
for part_text, weight in parsed_segments:
if part_text == "BREAK" and weight < 0:
flush(force=False)
continue
flat_tokens = _segment_to_tokens(subtokenizer, part_text, weight)
position = 0
while position < len(flat_tokens):
if len(current) == chunk_length:
flush(force=False)
token_pair = flat_tokens[position]
token_value = token_pair[0]
if use_classic_chunking and isinstance(token_value, int):
if token_value == comma_token:
last_comma = len(current)
elif (
len(current) == chunk_length
and last_comma != -1
and len(current) - last_comma <= 20
):
break_location = last_comma + 1
reloc = current[break_location:]
current = current[:break_location]
flush(force=False)
current = reloc
last_comma = _find_last_comma(current, comma_token)
continue
if len(current) == chunk_length:
flush(force=False)
current.append(token_pair)
position += 1
flush(force=not batches)
return batches
def tokenize_prompt(clip, text, use_classic_attention=False, use_classic_chunking=False):
"""
Build token-weight batches matching Comfy's expected tokenizer output shape.
"""
parsed_segments = (
_parse_classic_attention(text)
if use_classic_attention
else _parse_default_attention(text)
)
tokens = {}
for key, subtokenizer in _get_subtokenizers(clip.tokenizer).items():
tokens[key] = _pack_tokens(
subtokenizer, parsed_segments, use_classic_chunking=use_classic_chunking
)
return tokens
+410 -41
View File
@@ -1,63 +1,432 @@
"""
Wildcard Expansion Module (Placeholder)
Wildcard expansion helpers for A1111-style prompt text.
Future implementation for A1111-style wildcard support:
- __wildcard__ syntax expansion
- Nested wildcards
- Wildcard file browser/picker
- Preview what wildcards will expand to
Supported syntax:
- ``__wildcard__`` for files like ``data/wildcards/wildcard.txt``
- Nested directories via dot notation, e.g. ``__characters.anime__``
- Nested wildcard expansion inside selected lines
- Dynamic prompt choices like ``{a|b}``, ``{1-2$$a|b|c}``, and ``{20%a|b}``
"""
from typing import Optional, List
from __future__ import annotations
import logging
import random
import re
import threading
from dataclasses import dataclass
from pathlib import Path
from typing import Callable, Dict, Iterable, List, Optional, Tuple
logger = logging.getLogger("A1111PromptNode")
DEFAULT_WILDCARDS_DIR = Path(__file__).resolve().parent.parent / "data" / "wildcards"
WILDCARD_PATTERN = re.compile(r"__([^_]+(?:_[^_]+)*)__")
DYNAMIC_PROMPT_PATTERN = re.compile(r"\{([^{}]*)\}")
MAX_WILDCARD_EXPANSION_DEPTH = 32
def expand_wildcards(text: str, wildcards_dir: Optional[str] = None) -> str:
def _get_wildcards_dir(wildcards_dir: Optional[str] = None) -> Path:
return Path(wildcards_dir) if wildcards_dir else DEFAULT_WILDCARDS_DIR
def _iter_wildcard_files(wildcards_dir: Path) -> Iterable[Path]:
if not wildcards_dir.exists():
return []
return wildcards_dir.rglob("*.txt")
def _display_name_for_file(path: Path, wildcards_dir: Path) -> str:
relative = path.relative_to(wildcards_dir).with_suffix("")
return ".".join(relative.parts)
def _normalize_wildcard_name(name: str) -> str:
return name.strip().replace("\\", ".").replace("/", ".").strip(".").lower()
@dataclass(frozen=True)
class WildcardFileIndex:
files: Tuple[Path, ...]
by_name: Dict[str, Path]
by_leaf: Dict[str, Tuple[Path, ...]]
_wildcard_index_cache: Dict[Path, WildcardFileIndex] = {}
_wildcard_index_lock = threading.Lock()
def get_wildcard_file_index(
wildcards_dir: Optional[str | Path] = None,
refresh: bool = False,
) -> WildcardFileIndex:
base_dir = _get_wildcards_dir(
str(wildcards_dir) if wildcards_dir is not None else None
).resolve()
with _wildcard_index_lock:
cached = _wildcard_index_cache.get(base_dir)
if cached is not None and not refresh:
return cached
files = tuple(_iter_wildcard_files(base_dir))
by_name: Dict[str, Path] = {}
leaf_lists: Dict[str, List[Path]] = {}
for file_path in files:
normalized_name = _normalize_wildcard_name(
_display_name_for_file(file_path, base_dir)
)
by_name.setdefault(normalized_name, file_path)
leaf_lists.setdefault(file_path.stem.lower(), []).append(file_path)
index = WildcardFileIndex(
files=files,
by_name=by_name,
by_leaf={
leaf_name: tuple(paths) for leaf_name, paths in leaf_lists.items()
},
)
_wildcard_index_cache[base_dir] = index
return index
def _resolve_wildcard_file(
wildcard_name: str, wildcards_dir: Optional[str] = None
) -> Optional[Path]:
base_dir = _get_wildcards_dir(wildcards_dir)
if not base_dir.exists():
return None
normalized_name = _normalize_wildcard_name(wildcard_name)
if not normalized_name:
return None
direct_path = base_dir.joinpath(*normalized_name.split(".")).with_suffix(".txt")
if direct_path.is_file():
return direct_path
index = get_wildcard_file_index(base_dir)
indexed_path = index.by_name.get(normalized_name)
if indexed_path is not None:
return indexed_path
leaf_matches = index.by_leaf.get(normalized_name, ())
if len(leaf_matches) == 1:
return leaf_matches[0]
if len(leaf_matches) > 1:
logger.warning(
"[A1111 Prompt] Ambiguous leaf wildcard name '%s'; use the full path to disambiguate.",
wildcard_name,
)
return None
def _read_wildcard_lines(path: Path) -> List[str]:
options: List[str] = []
with open(path, "r", encoding="utf-8", errors="replace") as handle:
for raw_line in handle:
line = raw_line.strip()
if line and not line.startswith("#"):
options.append(line)
return options
def _get_dynamic_variant_weight(variant: str) -> int:
split_variant = variant.split("%", 1)
if len(split_variant) == 2:
try:
return int(split_variant[0])
except ValueError:
return 0
return 0
def _strip_dynamic_variant_weight(variant: str) -> str:
split_variant = variant.split("%", 1)
if len(split_variant) == 2:
return split_variant[1]
return variant
def _parse_dynamic_range(range_str: Optional[str], num_variants: int) -> tuple[int, int]:
if range_str is None:
return 1, 1
parts = range_str.split("-")
if len(parts) == 1:
value = min(int(parts[0]), num_variants)
return value, value
if len(parts) == 2:
low = int(parts[0]) if parts[0] else 0
high = min(int(parts[1]), num_variants) if parts[1] else num_variants
return min(low, high), max(low, high)
raise ValueError(f"Unexpected dynamic prompt range: {range_str}")
def _expand_dynamic_prompt_match(
match: re.Match[str], chooser: random.Random | random.Random
) -> str:
combinations_str = match.group(1)
variants = [segment.strip() for segment in combinations_str.split("|")]
if not variants:
return ""
weights = [_get_dynamic_variant_weight(variant) for variant in variants]
variants = [_strip_dynamic_variant_weight(variant) for variant in variants]
splits = variants[0].split("$$", 1)
quantity_spec: Optional[str] = None
if len(splits) == 2:
quantity_spec = splits[0].strip()
variants[0] = splits[1].strip()
try:
low_range, high_range = _parse_dynamic_range(quantity_spec, len(variants))
except ValueError:
return match.group(0)
if high_range <= 0:
return ""
quantity = chooser.randint(low_range, high_range)
if quantity <= 0:
return ""
total_weight = sum(weights)
zero_weight_count = weights.count(0)
if zero_weight_count > 0 and total_weight < 100:
remaining_weight = max(0, 100 - total_weight)
fill_weight = remaining_weight / zero_weight_count if zero_weight_count else 0
weights = [fill_weight if weight == 0 else weight for weight in weights]
elif all(weight == 0 for weight in weights):
weights = [1] * len(weights)
available_variants = list(variants)
available_weights = list(weights)
picked: List[str] = []
for _ in range(min(quantity, len(available_variants))):
choice = chooser.choices(available_variants, weights=available_weights, k=1)[0]
picked.append(choice)
index = available_variants.index(choice)
available_variants.pop(index)
available_weights.pop(index)
return ", ".join(picked)
def _expand_dynamic_prompts(
text: str, chooser: random.Random | random.Random
) -> tuple[str, bool]:
changed = False
def replace_match(match: re.Match[str]) -> str:
nonlocal changed
replacement = _expand_dynamic_prompt_match(match, chooser)
if replacement != match.group(0):
changed = True
return replacement
return DYNAMIC_PROMPT_PATTERN.sub(replace_match, text), changed
def _select_highest_scoring_text(
options: List[str], scorer: Callable[[str], int]
) -> str:
if not options:
return ""
return max(options, key=lambda option: (scorer(option), len(option), option))
def _expand_dynamic_prompt_match_max(
match: re.Match[str],
scorer: Callable[[str], int],
expander: Callable[[str], str],
) -> str:
combinations_str = match.group(1)
variants = [segment.strip() for segment in combinations_str.split("|")]
if not variants:
return ""
variants = [_strip_dynamic_variant_weight(variant) for variant in variants]
splits = variants[0].split("$$", 1)
quantity_spec: Optional[str] = None
if len(splits) == 2:
quantity_spec = splits[0].strip()
variants[0] = splits[1].strip()
try:
_, high_range = _parse_dynamic_range(quantity_spec, len(variants))
except ValueError:
return match.group(0)
if high_range <= 0:
return ""
expanded_variants = [expander(variant) for variant in variants]
ranked_variants = sorted(
expanded_variants,
key=lambda variant: (scorer(variant), len(variant), variant),
reverse=True,
)
return ", ".join(ranked_variants[: min(high_range, len(ranked_variants))])
def _expand_wildcards_for_token_count(
text: str,
wildcards_dir: Optional[str],
scorer: Callable[[str], int],
depth_remaining: int,
) -> str:
if not text or depth_remaining <= 0:
return text
current = text
for _ in range(depth_remaining):
changed = False
def replace_wildcard(match: re.Match[str]) -> str:
nonlocal changed
options = get_wildcard_options(match.group(1), wildcards_dir=wildcards_dir)
if not options:
return match.group(0)
changed = True
expanded_options = [
_expand_wildcards_for_token_count(
option, wildcards_dir, scorer, depth_remaining - 1
)
for option in options
]
return _select_highest_scoring_text(expanded_options, scorer)
expanded = WILDCARD_PATTERN.sub(replace_wildcard, current)
def replace_dynamic(match: re.Match[str]) -> str:
nonlocal changed
def expand_variant(variant: str) -> str:
return _expand_wildcards_for_token_count(
variant, wildcards_dir, scorer, depth_remaining - 1
)
replacement = _expand_dynamic_prompt_match_max(
match,
scorer=scorer,
expander=expand_variant,
)
if replacement != match.group(0):
changed = True
return replacement
expanded = DYNAMIC_PROMPT_PATTERN.sub(replace_dynamic, expanded)
expanded = _normalize_prompt_spacing(expanded)
if not changed or expanded == current:
return expanded
current = expanded
return current
def _normalize_prompt_spacing(text: str) -> str:
text = re.sub(r"[ \t]*,[ \t]*", ", ", text)
text = re.sub(r"[ \t]{2,}", " ", text)
return text.strip()
def expand_wildcards(
text: str,
wildcards_dir: Optional[str] = None,
rng: Optional[random.Random] = None,
max_depth: int = MAX_WILDCARD_EXPANSION_DEPTH,
) -> str:
"""
Expand __wildcard__ syntax in prompt text.
Expand ``__wildcard__`` syntax in prompt text.
Args:
text: Prompt text containing __wildcard__ patterns
wildcards_dir: Directory containing wildcard .txt files
Returns:
Text with wildcards expanded to random selections
TODO: Implement this functionality
Missing wildcard files are left untouched so the user can spot the problem
in the prompt preview instead of silently losing text.
"""
# Placeholder - returns text unchanged
return text
if not text:
return text
chooser = rng or random
current = text
for _ in range(max_depth):
changed = False
def replace_match(match: re.Match[str]) -> str:
nonlocal changed
options = get_wildcard_options(match.group(1), wildcards_dir=wildcards_dir)
if not options:
return match.group(0)
changed = True
return chooser.choice(options)
expanded = WILDCARD_PATTERN.sub(replace_match, current)
expanded, dynamic_changed = _expand_dynamic_prompts(expanded, chooser)
expanded = _normalize_prompt_spacing(expanded)
changed = changed or dynamic_changed
if not changed or expanded == current:
return expanded
current = expanded
if WILDCARD_PATTERN.search(current):
logger.warning(
"[A1111 Prompt] Wildcard expansion depth limit reached; possible recursive loop in prompt: %s",
current[:200],
)
return current
def expand_wildcards_for_token_count(
text: str,
wildcards_dir: Optional[str] = None,
scorer: Optional[Callable[[str], int]] = None,
max_depth: int = MAX_WILDCARD_EXPANSION_DEPTH,
) -> str:
"""
Deterministically expand wildcards and dynamic prompts for worst-case token counting.
Runtime wildcard expansion remains random. This helper is only for UI estimates
and chooses the highest-scoring option at each wildcard/dynamic prompt branch.
"""
score_text = scorer or (lambda value: len(value))
return _expand_wildcards_for_token_count(
text,
wildcards_dir=wildcards_dir,
scorer=score_text,
depth_remaining=max_depth,
)
def get_wildcard_options(
wildcard_name: str, wildcards_dir: Optional[str] = None
) -> List[str]:
"""
Get all options for a given wildcard.
Get all options for a given wildcard file.
Args:
wildcard_name: Name of the wildcard (without __ delimiters)
wildcards_dir: Directory containing wildcard .txt files
Returns:
List of options from the wildcard file
TODO: Implement this functionality
Returns an empty list when the wildcard file does not exist.
"""
# Placeholder - returns empty list
return []
resolved_file = _resolve_wildcard_file(wildcard_name, wildcards_dir=wildcards_dir)
if resolved_file is None:
return []
return _read_wildcard_lines(resolved_file)
def list_available_wildcards(wildcards_dir: Optional[str] = None) -> List[str]:
"""
List all available wildcard names.
Args:
wildcards_dir: Directory containing wildcard .txt files
Returns:
List of wildcard names (without __ delimiters)
TODO: Implement this functionality
List wildcard names available under the configured wildcard directory.
"""
# Placeholder - returns empty list
return []
base_dir = _get_wildcards_dir(wildcards_dir)
if not base_dir.exists():
return []
wildcard_names = [
_display_name_for_file(file_path, base_dir)
for file_path in _iter_wildcard_files(base_dir)
]
return sorted(wildcard_names, key=lambda name: name.lower())
+30
View File
@@ -0,0 +1,30 @@
[project]
name = "a1111-prompt-node"
description = "A custom ComfyUI node that implements A1111-style prompt handling with proper isolation and emphasis math."
version = "1.0.0"
license = {file = "LICENSE"}
authors = [
{ name = "Enferlain" }
]
dependencies = ["lark", "torch", "aiohttp"]
[project.urls]
Repository = "https://github.com/Enferlain/ComfyUI-A1111-cond"
"Bug Tracker" = "https://github.com/Enferlain/ComfyUI-A1111-cond/issues"
[tool.comfy]
PublisherId = "enferlain"
DisplayName = "A1111 Prompt Node"
Icon = ""
includes = [
"nodes/",
"parser/",
"api/",
"js/",
"data/",
"__init__.py",
"hooks.py",
"LICENSE",
"README.md",
"requirements.txt"
]
+3 -1
View File
@@ -1 +1,3 @@
lark
lark
torch
aiohttp
View File
+8
View File
@@ -0,0 +1,8 @@
1girl,0,6000000,"1girls,sole_female"
1boy,0,1000000,"1boys,sole_male"
solo,0,8000000,
highres,5,5000000,"high_res"
hat,0,500000,"headwear"
artist_name,1,100,
character_name,4,1000,
first_name_last_name,0,50,
1 1girl 0 6000000 1girls,sole_female
2 1boy 0 1000000 1boys,sole_male
3 solo 0 8000000
4 highres 5 5000000 high_res
5 hat 0 500000 headwear
6 artist_name 1 100
7 character_name 4 1000
8 first_name_last_name 0 50
+290
View File
@@ -0,0 +1,290 @@
import unittest
import csv
import sys
import tempfile
from pathlib import Path
# Add the project root to sys.path
sys.path.append(str(Path(__file__).parent.parent))
from api.autocomplete import TagDatabase, WildcardDatabase
DATA_DIR = Path(__file__).parent.parent / "data" / "tags"
def load_rows_from_real_csv(filename, tag_names):
wanted = set(tag_names)
rows = {}
with (DATA_DIR / filename).open("r", encoding="utf-8", errors="replace", newline="") as handle:
reader = csv.reader(handle)
for row in reader:
if len(row) >= 3 and row[0] in wanted:
rows[row[0]] = row
if len(rows) == len(wanted):
break
missing = wanted.difference(rows)
if missing:
raise AssertionError(f"Missing matrix rows in {filename}: {sorted(missing)}")
return [rows[name] for name in tag_names]
class TestAutocomplete(unittest.TestCase):
def setUp(self):
self.db = TagDatabase()
self.mock_csv = Path(__file__).parent / "mock_tags.csv"
self.db.load_csv(self.mock_csv)
def test_prefix_search(self):
# Search for "1g" -> should find "1girl"
results = self.db.search("1g")
self.assertGreater(len(results), 0)
self.assertEqual(results[0]["name"], "1girl")
def test_alias_matching(self):
# Search for "sole_female" -> should find "1girl" (via alias)
results = self.db.search("sole_female")
self.assertGreater(len(results), 0)
self.assertEqual(results[0]["name"], "1girl")
self.assertEqual(results[0]["matched_alias"], "sole_female")
def test_sorting_by_count(self):
# "1girl" has 6M, "1boy" has 1M
# Searching for "1" should return both, girl first
results = self.db.search("1")
names = [r["name"] for r in results]
self.assertIn("1girl", names)
self.assertIn("1boy", names)
# Verify 1girl appears before 1boy due to count
girl_idx = names.index("1girl")
boy_idx = names.index("1boy")
self.assertLess(girl_idx, boy_idx)
def test_tag_types(self):
# artist_name is type 1
results = self.db.search("artist_name")
self.assertEqual(results[0]["type"], 1)
# character_name is type 4
results = self.db.search("character_name")
self.assertEqual(results[0]["type"], 4)
def test_empty_query(self):
results = self.db.search("")
self.assertEqual(len(results), 0)
def test_long_query(self):
# Even if the API handles it, the search method should be robust
long_query = "a" * 300
results = self.db.search(long_query)
self.assertEqual(len(results), 0)
def test_search_coerces_unusual_inputs(self):
self.assertEqual(self.db.search(["1girl"]), [])
self.assertEqual(self.db.search("1girl", limit="bad"), self.db.search("1girl"))
self.assertEqual(self.db.search("1girl", limit=-5), [])
def test_short_queries_keep_contains_fallback(self):
results = self.db.search("rl")
self.assertEqual(results[0]["name"], "1girl")
def test_long_queries_keep_contains_fallback(self):
results = self.db.search("ist_")
self.assertEqual(results[0]["name"], "artist_name")
def test_substring_query_finds_middle_of_tag(self):
results = self.db.search("last_name")
self.assertEqual(results[0]["name"], "first_name_last_name")
def test_contains_fallback_can_be_disabled(self):
self.assertEqual(self.db.search("ist_", contains_fallback=False), [])
class TestAutocompleteUseCaseMatrix(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.temp_dir = tempfile.TemporaryDirectory()
cls.matrix_csv = Path(cls.temp_dir.name) / "matrix_tags.csv"
rows = []
rows.extend(
load_rows_from_real_csv(
"danbooru.csv",
[
"1girl",
"looking_at_viewer",
"open_mouth",
"short_hair",
"twintails",
"school_uniform",
"depth_of_field",
"holding_hands",
],
)
)
rows.extend(
load_rows_from_real_csv(
"e621.csv",
[
"anthro",
"hi_res",
"male",
"genitals",
"clothing",
"hair",
],
)
)
with cls.matrix_csv.open("w", encoding="utf-8", newline="") as handle:
writer = csv.writer(handle)
writer.writerows(rows)
cls.db = TagDatabase()
cls.db.load_csv(cls.matrix_csv)
@classmethod
def tearDownClass(cls):
cls.temp_dir.cleanup()
def assertPopupContains(self, query, expected_name, limit=20, search_aliases=True):
results = self.db.search(query, limit=limit, search_aliases=search_aliases)
names = [result["name"] for result in results]
self.assertIn(
expected_name,
names,
msg=f"{expected_name!r} not found for query {query!r}; got {names}",
)
def test_likely_typing_prefixes_show_expected_tags(self):
cases = [
("1g", "1girl"),
("look", "looking_at_viewer"),
("open", "open_mouth"),
("short", "short_hair"),
("twin", "twintails"),
("school", "school_uniform"),
("anth", "anthro"),
("gen", "genitals"),
("cloth", "clothing"),
]
for query, expected_name in cases:
with self.subTest(query=query, expected_name=expected_name):
self.assertPopupContains(query, expected_name)
def test_likely_typing_aliases_show_canonical_tags(self):
cases = [
("sole_female", "1girl"),
("mouth_open", "open_mouth"),
("short-hair", "short_hair"),
("high_res", "hi_res"),
("1boy", "male"),
("clothes", "clothing"),
]
for query, expected_name in cases:
with self.subTest(query=query, expected_name=expected_name):
self.assertPopupContains(query, expected_name)
def test_likely_typing_middle_fragments_show_expected_tags(self):
cases = [
("at_viewer", "looking_at_viewer"),
("mouth", "open_mouth"),
("uniform", "school_uniform"),
("of_field", "depth_of_field"),
("hands", "holding_hands"),
("itals", "genitals"),
]
for query, expected_name in cases:
with self.subTest(query=query, expected_name=expected_name):
self.assertPopupContains(query, expected_name)
def test_alias_toggle_removes_alias_only_matches(self):
results = self.db.search("mouth_open", search_aliases=False)
self.assertNotIn("open_mouth", [result["name"] for result in results])
class TestWildcardAutocomplete(unittest.TestCase):
def setUp(self):
self.temp_dir = tempfile.TemporaryDirectory()
self.wildcards_dir = Path(self.temp_dir.name)
(self.wildcards_dir / "colors.txt").write_text("red\nblue\n", encoding="utf-8")
nested_dir = self.wildcards_dir / "characters"
nested_dir.mkdir(parents=True, exist_ok=True)
(nested_dir / "anime.txt").write_text("heroine\n", encoding="utf-8")
accessories_dir = nested_dir / "accessories"
accessories_dir.mkdir(parents=True, exist_ok=True)
(accessories_dir / "hats.txt").write_text("top hat\n", encoding="utf-8")
self.db = WildcardDatabase(self.wildcards_dir)
self.db.load()
def tearDown(self):
self.temp_dir.cleanup()
def test_wildcard_empty_query_returns_entries(self):
results = self.db.search("__", limit=10)
names = [result["name"] for result in results]
self.assertIn("colors", names)
self.assertIn("characters/anime", names)
def test_wildcard_prefix_search(self):
results = self.db.search("__char", limit=10)
self.assertEqual(results[0]["name"], "characters")
self.assertEqual(results[0]["completion"], "__characters__")
self.assertEqual(results[0]["kind"], "wildcard_folder")
def test_wildcard_manifest_contains_every_file_and_folder(self):
manifest = self.db.get_manifest()
names = {entry["name"] for entry in manifest}
self.assertEqual(len(manifest), self.db.wildcard_count)
self.assertIn("colors", names)
self.assertIn("characters", names)
self.assertIn("characters/anime", names)
self.assertIn("characters/accessories/hats", names)
def test_nested_wildcard_completion_uses_leaf_name(self):
results = self.db.search("__hats", limit=10)
self.assertEqual(results[0]["name"], "characters/accessories/hats")
self.assertEqual(results[0]["leaf_name"], "hats")
self.assertEqual(results[0]["completion"], "__hats__")
def test_duplicate_leaf_completion_uses_full_path(self):
(self.wildcards_dir / "props").mkdir(parents=True, exist_ok=True)
(self.wildcards_dir / "props" / "hats.txt").write_text("helmet\n", encoding="utf-8")
self.db.load()
results = self.db.search("__hats", limit=10)
completions = {result["name"]: result["completion"] for result in results}
self.assertEqual(
completions["characters/accessories/hats"],
"__characters/accessories/hats__",
)
self.assertEqual(completions["props/hats"], "__props/hats__")
def test_wildcard_contents_search(self):
results = self.db.get_contents("characters/anime", limit=10)
self.assertEqual(len(results), 1)
self.assertEqual(results[0]["name"], "heroine")
self.assertEqual(results[0]["kind"], "wildcard_content")
self.assertEqual(results[0]["meta"], "characters/anime")
def test_folder_contents_search_returns_descendants(self):
results = self.db.get_contents("characters", limit=10)
names = [result["name"] for result in results]
self.assertIn("characters/accessories", names)
self.assertIn("characters/anime", names)
folder_entry = next(
result for result in results if result["name"] == "characters/accessories"
)
self.assertEqual(folder_entry["kind"], "wildcard_folder")
def test_wildcard_limits_are_coerced(self):
self.assertEqual(self.db.search("__", limit=-1), [])
self.assertEqual(
self.db.get_contents("characters", limit="bad")[0]["name"],
"characters/accessories",
)
if __name__ == "__main__":
unittest.main()
+125
View File
@@ -0,0 +1,125 @@
import unittest
import sys
from pathlib import Path
import torch
from unittest.mock import MagicMock
# Add project root to sys.path
sys.path.append(str(Path(__file__).parent.parent))
# Mock ComfyUI dependencies
from unittest.mock import MagicMock
mock_hooks = MagicMock()
class DummyHook:
def __init__(self, *args, **kwargs):
pass
def clone(self):
return self
mock_hooks.TransformerOptionsHook = DummyHook
mock_hooks.EnumHookScope = MagicMock()
sys.modules["comfy.hooks"] = mock_hooks
sys.modules["server"] = MagicMock()
from hooks import A1111StepConditioningHook
class TestHooks(unittest.TestCase):
def test_get_step_from_sigma(self):
# mock sample_sigmas [10.0, 8.0, 5.0, 0.0]
# step 0: 10.0 -> 8.0
# step 1: 8.0 -> 5.0
# step 2: 5.0 -> 0.0
sigmas = torch.tensor([10.0, 8.0, 5.0, 0.0])
hook = A1111StepConditioningHook(step_embeddings=[])
# Exact match start
self.assertEqual(hook.get_step_from_sigma(10.0, sigmas), 0)
# Inside first step
self.assertEqual(hook.get_step_from_sigma(9.0, sigmas), 0)
# Boundary: 8.0 is the START of step 1
self.assertEqual(hook.get_step_from_sigma(8.0, sigmas), 1)
# Inside second step
self.assertEqual(hook.get_step_from_sigma(7.0, sigmas), 1)
# Exact match end of last step: 0.0 is the start of the "final" interval
# (Though samplers usually stop at 0, this returns len-2 as per hooks.py)
self.assertEqual(hook.get_step_from_sigma(0.0, sigmas), 2)
def test_step_index_scaling(self):
# sampler has 10 steps, we have 20 embeddings
sigmas = torch.linspace(10, 0, 11) # length 11 -> 10 steps
embeddings = [i for i in range(21)] # length 21 -> step 0 to 20
hook = A1111StepConditioningHook(step_embeddings=embeddings)
# Mock the wrapper logic part that calculates step_idx
def get_scaled_step(sigma_val):
raw_idx = hook.get_step_from_sigma(sigma_val, sigmas)
actual_steps = len(sigmas) - 1
emb_steps = len(embeddings) - 1
return round(raw_idx * emb_steps / actual_steps)
# Start (step 0 of 10) -> step 0 of 20
self.assertEqual(get_scaled_step(10.0), 0)
# Middle (step 5 of 10) -> step 10 of 20
self.assertEqual(get_scaled_step(5.0), 10)
# End (step 9 of 10) -> step 18 of 20 (it's N-1 for last index usually)
# actually last sigma is step 9.
self.assertEqual(get_scaled_step(sigmas[9].item()), 18)
def test_model_function_wrapper_swaps(self):
# Prepare 3 embeddings (0, 1, 2)
emb0 = torch.ones((1, 1, 4)) * 0.0
emb1 = torch.ones((1, 1, 4)) * 1.0
emb2 = torch.ones((1, 1, 4)) * 2.0
hook = A1111StepConditioningHook(
step_embeddings=[(emb0, None), (emb1, None), (emb2, None)]
)
hook.debug = True
# Mock arguments
apply_model_func = MagicMock(return_value="model_output")
input_x = torch.zeros((1, 4))
timestep = torch.tensor([10.0])
sigmas = torch.tensor([10.0, 5.0, 0.0]) # 2 steps: [10->5, 5->0]
c = {
"c_crossattn": torch.ones((2, 1, 4)) * -1.0, # original (batch=2)
"transformer_options": {"sample_sigmas": sigmas},
}
cond_or_uncond = [0, 1] # 0 = positive (swap), 1 = negative (keep)
args = {
"input": input_x,
"timestep": timestep,
"c": c,
"cond_or_uncond": cond_or_uncond,
}
# Step 0 (sigma 10.0) -> raw_idx 0
# Scaling: round(0 * 2 / 2) = 0 -> emb0
hook.model_function_wrapper(apply_model_func, args)
call_args = apply_model_func.call_args[1]
new_cond = call_args["c_crossattn"]
self.assertTrue(torch.allclose(new_cond[0], emb0[0]))
# Step 1 (sigma 5.0) -> raw_idx 1
# Scaling: round(1 * 2 / 2) = 1 -> emb1
timestep = torch.tensor([5.0])
args["timestep"] = timestep
hook.model_function_wrapper(apply_model_func, args)
call_args = apply_model_func.call_args[1]
new_cond = call_args["c_crossattn"]
self.assertTrue(torch.allclose(new_cond[0], emb1[0]))
if __name__ == "__main__":
unittest.main()
+107
View File
@@ -0,0 +1,107 @@
import sys
import unittest
from pathlib import Path
# Add project root to sys.path
sys.path.append(str(Path(__file__).parent.parent))
from metadata import (
METADATA_KEY,
RESOLVED_PROMPT_EXPANDED_PROPERTY,
RESOLVED_PROMPT_PROPERTY,
RESOLVED_PROMPT_SOURCE_PROPERTY,
record_prompt_metadata,
)
class TestPromptMetadata(unittest.TestCase):
def test_records_original_and_resolved_prompt_by_node_id(self):
extra_pnginfo = {}
record_prompt_metadata(
extra_pnginfo,
node_id="63",
node_type="A1111Prompt",
original_prompt="portrait __subject__",
resolved_prompt="portrait cat",
)
entry = extra_pnginfo[METADATA_KEY]["nodes"]["63"]
self.assertEqual(entry["node_id"], "63")
self.assertEqual(entry["type"], "A1111Prompt")
self.assertEqual(entry["original_prompt"], "portrait __subject__")
self.assertEqual(entry["resolved_prompt"], "portrait cat")
self.assertTrue(entry["changed"])
def test_ignores_missing_pnginfo_or_node_id(self):
record_prompt_metadata(None, "63", "A1111Prompt", "a", "b")
extra_pnginfo = {}
record_prompt_metadata(extra_pnginfo, None, "A1111Prompt", "a", "b")
self.assertEqual(extra_pnginfo, {})
def test_preserves_existing_prompt_entries(self):
extra_pnginfo = {}
record_prompt_metadata(extra_pnginfo, "63", "A1111Prompt", "a", "b")
record_prompt_metadata(extra_pnginfo, "64", "A1111PromptNegative", "c", "c")
nodes = extra_pnginfo[METADATA_KEY]["nodes"]
self.assertEqual(set(nodes), {"63", "64"})
self.assertFalse(nodes["64"]["changed"])
def test_writes_resolved_prompt_to_embedded_workflow_node_properties(self):
extra_pnginfo = {
"workflow": {
"nodes": [
{"id": 12, "type": "OtherNode", "properties": {}},
{"id": 63, "type": "A1111Prompt", "properties": {}},
]
}
}
record_prompt_metadata(
extra_pnginfo,
node_id="63",
node_type="A1111Prompt",
original_prompt="portrait __subject__",
resolved_prompt="portrait cat",
)
properties = extra_pnginfo["workflow"]["nodes"][1]["properties"]
self.assertEqual(properties[RESOLVED_PROMPT_PROPERTY], "portrait cat")
self.assertEqual(properties[RESOLVED_PROMPT_SOURCE_PROPERTY], "portrait __subject__")
self.assertFalse(properties[RESOLVED_PROMPT_EXPANDED_PROPERTY])
def test_removes_stale_workflow_preview_when_prompt_is_unchanged(self):
extra_pnginfo = {
"workflow": {
"nodes": [
{
"id": "63",
"type": "A1111Prompt",
"properties": {
RESOLVED_PROMPT_PROPERTY: "old resolved",
RESOLVED_PROMPT_EXPANDED_PROPERTY: True,
},
}
]
}
}
record_prompt_metadata(
extra_pnginfo,
node_id="63",
node_type="A1111Prompt",
original_prompt="same prompt",
resolved_prompt="same prompt",
)
properties = extra_pnginfo["workflow"]["nodes"][0]["properties"]
self.assertNotIn(RESOLVED_PROMPT_PROPERTY, properties)
self.assertNotIn(RESOLVED_PROMPT_EXPANDED_PROPERTY, properties)
self.assertNotIn(RESOLVED_PROMPT_SOURCE_PROPERTY, properties)
if __name__ == "__main__":
unittest.main()
+91
View File
@@ -0,0 +1,91 @@
import unittest
import torch
# We extract the pure logic from prompt_node.py to test it in isolation
# without dealing with ComfyUI's complex package/import system.
def extract_weights(token_batches):
all_weights = []
for batch in token_batches:
if not batch:
all_weights.append([])
continue
batch_weights = [
t[1] if hasattr(t, "__getitem__") and len(t) > 1 else 1.0 for t in batch
]
all_weights.append(batch_weights)
return all_weights
def apply_direct_scaling(cond, weights, normalization):
cond = cond.clone()
flat_weights = []
for chunk_weights in weights:
flat_weights.extend(chunk_weights)
seq_len = cond.shape[1]
if len(flat_weights) < seq_len:
flat_weights = flat_weights + [1.0] * (seq_len - len(flat_weights))
elif len(flat_weights) > seq_len:
flat_weights = flat_weights[:seq_len]
batch_size = cond.shape[0]
multipliers = torch.tensor(
[flat_weights] * batch_size, device=cond.device, dtype=cond.dtype
)
if normalization:
original_mean = cond.mean()
multipliers = multipliers.unsqueeze(-1).expand_as(cond)
cond = cond * multipliers
if normalization:
new_mean = cond.mean()
if new_mean.abs() > 1e-8:
scale_factor = original_mean / new_mean
cond = cond * scale_factor
return cond
class TestNodesLogic(unittest.TestCase):
def test_extract_weights(self):
tokens = [[(101, 1.0), (102, 1.2), (103, 0.8)], [(201, 1.5), (202, 1.0)]]
weights = extract_weights(tokens)
self.assertEqual(len(weights), 2)
self.assertEqual(weights[0], [1.0, 1.2, 0.8])
self.assertEqual(weights[1], [1.5, 1.0])
def test_apply_direct_scaling_basic(self):
cond = torch.ones((1, 3, 4))
weights = [[1.0, 2.0, 0.5]]
scaled = apply_direct_scaling(cond, weights, normalization=False)
self.assertTrue(torch.allclose(scaled[0, 0], torch.ones(4) * 1.0))
self.assertTrue(torch.allclose(scaled[0, 1], torch.ones(4) * 2.0))
self.assertTrue(torch.allclose(scaled[0, 2], torch.ones(4) * 0.5))
def test_apply_direct_scaling_with_break(self):
cond = torch.ones((1, 4, 4))
weights = [[1.2, 1.0], [0.8, 1.0]]
scaled = apply_direct_scaling(cond, weights, normalization=False)
expected_weights = [1.2, 1.0, 0.8, 1.0]
for i, w in enumerate(expected_weights):
self.assertTrue(torch.allclose(scaled[0, i], torch.ones(4) * w))
def test_normalization_mean_preservation(self):
cond = torch.randn((1, 10, 64))
original_mean = cond.mean().item()
weights = [[2.0] * 10]
scaled_no_norm = apply_direct_scaling(cond, weights, normalization=False)
self.assertAlmostEqual(
scaled_no_norm.mean().item(), original_mean * 2, places=4
)
scaled_norm = apply_direct_scaling(cond, weights, normalization=True)
self.assertAlmostEqual(scaled_norm.mean().item(), original_mean, places=4)
if __name__ == "__main__":
unittest.main()
+89
View File
@@ -0,0 +1,89 @@
import sys
import unittest
from pathlib import Path
from unittest.mock import MagicMock
import torch
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
mock_model_management = MagicMock()
mock_model_management.intermediate_device.return_value = "cpu"
mock_comfy = MagicMock()
mock_comfy.model_management = mock_model_management
mock_hooks = MagicMock()
class DummyHook:
def __init__(self, *args, **kwargs):
pass
def clone(self):
return self
class DummyHookGroup:
def add(self, hook):
self.hook = hook
mock_hooks.TransformerOptionsHook = DummyHook
mock_hooks.HookGroup = DummyHookGroup
mock_hooks.EnumHookScope = MagicMock()
mock_hooks.set_hooks_for_conditioning = (
lambda conditioning, hooks=None, append_hooks=True: conditioning
)
mock_comfy.hooks = mock_hooks
server = MagicMock()
server.PromptServer.instance.routes.post = lambda path: (lambda fn: fn)
server.PromptServer.instance.routes.get = lambda path: (lambda fn: fn)
sys.modules["comfy"] = mock_comfy
sys.modules["comfy.model_management"] = mock_model_management
sys.modules["comfy.hooks"] = mock_hooks
sys.modules["server"] = server
from A1111_Prompt_Node.nodes.prompt_node import A1111PromptNode
class FakeCondStageModel:
pass
class FakeClip:
def __init__(self):
self.cond_stage_model = FakeCondStageModel()
class CacheProbePromptNode(A1111PromptNode):
def _get_downstream_steps(self, prompt, start_node_id, debug=False):
return 2
def _encode_with_break_isolation(self, clip, prompt_text, normalization, is_sdxl, debug):
seq_len = 2 if prompt_text == "short" else 5
cond = torch.ones((1, seq_len, 3))
pooled = torch.ones((1, 3))
return cond, pooled
class TestPromptNodeCache(unittest.TestCase):
def test_scheduled_padding_does_not_mutate_encoded_cache(self):
node = CacheProbePromptNode()
clip = FakeClip()
node.encode(clip, "[short:long:0.5]", prompt={}, unique_id="1")
cached_shapes = {
key[2]: cond.shape[1]
for key, (cond, _pooled) in node._encoded_cache.items()
}
self.assertEqual(cached_shapes["short"], 2)
self.assertEqual(cached_shapes["long"], 5)
if __name__ == "__main__":
unittest.main()
+94
View File
@@ -0,0 +1,94 @@
import unittest
import sys
import os
from pathlib import Path
# Add the project root to sys.path so we can import modules
sys.path.append(str(Path(__file__).parent.parent))
from parser.scheduler import get_prompt_schedule
class TestScheduler(unittest.TestCase):
def test_basic_scheduling(self):
# [before:after:when]
# At 10 steps, 0.5 is step 5
schedule = get_prompt_schedule("[cat:dog:0.5]", 10)
# Expect (5, "cat"), (10, "dog")
self.assertEqual(len(schedule), 2)
self.assertEqual(schedule[0], (5, "cat"))
self.assertEqual(schedule[1], (10, "dog"))
def test_step_based_scheduling(self):
# [before:after:step]
schedule = get_prompt_schedule("[cat:dog:7]", 10)
self.assertEqual(len(schedule), 2)
self.assertEqual(schedule[0], (7, "cat"))
self.assertEqual(schedule[1], (10, "dog"))
def test_alternation(self):
# [A|B]
schedule = get_prompt_schedule("[white|black]", 4)
# Alternation changes every step
self.assertEqual(len(schedule), 4)
self.assertEqual(schedule[0][1], "white")
self.assertEqual(schedule[1][1], "black")
self.assertEqual(schedule[2][1], "white")
self.assertEqual(schedule[3][1], "black")
def test_scheduled_alternation_add(self):
# [A|B:0.5] - Yield A until 50%, then alternate
schedule = get_prompt_schedule("[A|B:0.5]", 4)
# steps=4, 0.5=2
# Step 1: A
# Step 2: A
# Step 3: A (cycle (3-1)%2=0) - wait, check implementation
# Actually _at_step calculates idx = (step - 1) % len(options)
# Step 1: yield options[0] -> A
# Step 2: yield options[0] -> A
# Step 3: (3-1)%2 = 0 -> A <-- Is this right? A1111 behavior check.
# Let's check my implementation in scheduler.py:
# idx = (step - 1) % len(options)
# So Step 3 (step > 2) yields options[0] = A
# Step 4: (4-1)%2 = 1 -> B
self.assertEqual(schedule[0][1], "A")
self.assertEqual(schedule[1][1], "A")
self.assertEqual(schedule[2][1], "A")
self.assertEqual(schedule[3][1], "B")
def test_scheduled_alternation_remove(self):
# [A|B::0.5] - Alternate until 50%, then first tag
schedule = get_prompt_schedule("[A|B::0.5]", 4)
# Step 1: A
# Step 2: B
# Step 3: A (first tag)
# Step 4: A (first tag)
self.assertEqual(schedule[0][1], "A")
self.assertEqual(schedule[1][1], "B")
self.assertEqual(schedule[2][1], "A")
self.assertEqual(schedule[3][1], "A")
def test_nested_logic(self):
# [(red:1.2):[blue:green:0.5]:0.3]
# steps=10, 0.3=3, 0.5=5
schedule = get_prompt_schedule("[(red:1.2):[blue:green:0.5]:0.3]", 10)
# Sequence expected:
# Step 1-3: (red:1.2)
# Step 4-5: blue
# Step 6-10: green
self.assertEqual(len(schedule), 3)
self.assertEqual(schedule[0], (3, "(red:1.2)"))
self.assertEqual(schedule[1], (5, "blue"))
self.assertEqual(schedule[2], (10, "green"))
def test_break_isolation(self):
# a BREAK [b:c:0.5]
schedule = get_prompt_schedule("a BREAK [b:c:0.5]", 10)
self.assertEqual(len(schedule), 2)
self.assertEqual(schedule[0], (5, "a BREAK b"))
self.assertEqual(schedule[1], (10, "a BREAK c"))
if __name__ == "__main__":
unittest.main()
+117
View File
@@ -0,0 +1,117 @@
import unittest
import sys
from pathlib import Path
from unittest.mock import MagicMock
# Add project root to sys.path
sys.path.append(str(Path(__file__).parent.parent))
# Mock ComfyUI dependencies before importing logic
sys.modules["aiohttp"] = MagicMock()
sys.modules["server"] = MagicMock()
sys.modules["comfy"] = MagicMock()
sys.modules["comfy.sd1_clip"] = MagicMock()
from api.tokenize import _find_word_position, build_token_info, strip_a1111_syntax
class FakeTokenizer:
def encode(self, word, add_special_tokens=False):
return [word]
def decode(self, token_ids):
return str(token_ids[0])
class TestTokenizer(unittest.TestCase):
def test_emphasis_stripping(self):
# (word:1.2) -> word
self.assertEqual(strip_a1111_syntax("(dog:1.3)"), "dog")
# (word) -> word
self.assertEqual(strip_a1111_syntax("(cat)"), "cat")
# multiple emphasis
self.assertEqual(
strip_a1111_syntax("a (beautiful:1.2) (cat:1.5)"), "a beautiful cat"
)
def test_bracket_emphasis(self):
# [word] -> word
self.assertEqual(strip_a1111_syntax("[bird]"), "bird")
# [word:0.5] -> word
self.assertEqual(strip_a1111_syntax("[fish:0.5]"), "fish")
def test_scheduling_stripping(self):
# [from:to:when] -> keeps longest
self.assertEqual(strip_a1111_syntax("[mountain:ocean:0.5]"), "mountain")
self.assertEqual(
strip_a1111_syntax("[forest:very_large_city:0.5]"), "very_large_city"
)
# [add:when] -> add
self.assertEqual(strip_a1111_syntax("[glasses:10]"), "glasses")
# [remove::when] -> remove
self.assertEqual(strip_a1111_syntax("[hat::15]"), "hat")
def test_alternation_stripping(self):
# [A|B|C] -> keeps longest
self.assertEqual(strip_a1111_syntax("[red|blue|green]"), "green")
self.assertEqual(strip_a1111_syntax("[a|bc|d]"), "bc")
def test_scheduled_alternation_stripping(self):
# [A|B:0.5] -> keeps longest
self.assertEqual(strip_a1111_syntax("[apple|banana:0.5]"), "banana")
self.assertEqual(strip_a1111_syntax("[apple|banana::0.5]"), "banana")
def test_escaped_stripping(self):
# \( \) -> ( )
self.assertEqual(strip_a1111_syntax(r"\(escaped\)"), "(escaped)")
self.assertEqual(strip_a1111_syntax(r"\[bracket\]"), "[bracket]")
def test_nested_stripping(self):
# [(red:1.2):[blue|green]:0.5]
# Inner [blue|green] -> green (len 5 vs blue len 4)
# Then [(red:1.2):green:0.5] -> red (len 3+tags removal) vs green (len 5)
# Result should be "green" because "red" is len 3 and "green" is len 5
self.assertEqual(strip_a1111_syntax("[(red:1.2):green:0.5]"), "green")
def test_boundary_position_uses_original_text(self):
text = "alpha (beautiful:1.2) beta"
self.assertEqual(
_find_word_position(text, "beautiful", 0, len(text), fallback=6),
text.index("beautiful"),
)
self.assertEqual(
_find_word_position(text, "missing", 0, len(text), fallback=6),
6,
)
def test_build_token_info_returns_all_chunk_boundaries(self):
text = " ".join(f"w{i}" for i in range(160))
info = build_token_info(text, text, FakeTokenizer())
self.assertEqual(info["sequences"], [75, 75, 10])
chunk_boundaries = [
boundary for boundary in info["boundaries"] if boundary["type"] == "chunk"
]
self.assertEqual(len(chunk_boundaries), 2)
self.assertEqual(chunk_boundaries[0]["to_chunk"], 1)
self.assertEqual(chunk_boundaries[1]["to_chunk"], 2)
self.assertLess(chunk_boundaries[0]["char_pos"], chunk_boundaries[1]["char_pos"])
def test_estimated_wildcard_count_keeps_approximate_boundaries(self):
text = "prefix __wildcard__ suffix"
wildcard_expansion = " ".join(f"choice{i}" for i in range(80))
counting_text = f"prefix {wildcard_expansion} suffix"
info = build_token_info(text, counting_text, FakeTokenizer())
self.assertEqual(info["sequences"], [75, 7])
self.assertTrue(info["stats"]["estimated_from_wildcards"])
self.assertEqual(len(info["boundaries"]), 1)
self.assertTrue(info["boundaries"][0]["estimated"])
self.assertGreaterEqual(info["boundaries"][0]["char_pos"], 0)
self.assertLessEqual(info["boundaries"][0]["char_pos"], len(text))
if __name__ == "__main__":
unittest.main()
+225
View File
@@ -0,0 +1,225 @@
import random
import sys
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch
# Add project root to sys.path
sys.path.append(str(Path(__file__).parent.parent))
from parser.wildcards import (
expand_wildcards,
expand_wildcards_for_token_count,
get_wildcard_options,
list_available_wildcards,
)
from parser import wildcards as wildcards_module
class TestWildcards(unittest.TestCase):
def setUp(self):
self.temp_dir = tempfile.TemporaryDirectory()
self.wildcards_dir = Path(self.temp_dir.name)
def tearDown(self):
self.temp_dir.cleanup()
def write_wildcard(self, relative_path: str, contents: str):
wildcard_file = self.wildcards_dir / relative_path
wildcard_file.parent.mkdir(parents=True, exist_ok=True)
wildcard_file.write_text(contents, encoding="utf-8")
def test_get_wildcard_options_ignores_comments_and_blank_lines(self):
self.write_wildcard(
"colors.txt",
"\n# comment\nred\n\nblue\n \n# another comment\ngreen\n",
)
self.assertEqual(
get_wildcard_options("colors", wildcards_dir=str(self.wildcards_dir)),
["red", "blue", "green"],
)
def test_list_available_wildcards_uses_dot_notation(self):
self.write_wildcard("colors.txt", "red\n")
self.write_wildcard("characters/anime.txt", "heroine\n")
self.assertEqual(
list_available_wildcards(wildcards_dir=str(self.wildcards_dir)),
["characters.anime", "colors"],
)
def test_expand_wildcards_resolves_nested_entries(self):
self.write_wildcard("colors.txt", "red\n")
self.write_wildcard("outfits.txt", "__colors__ dress\n")
expanded = expand_wildcards(
"portrait, __outfits__",
wildcards_dir=str(self.wildcards_dir),
rng=random.Random(0),
)
self.assertEqual(expanded, "portrait, red dress")
def test_expand_wildcards_supports_nested_directory_names(self):
self.write_wildcard("characters/anime.txt", "mahou shoujo\n")
expanded = expand_wildcards(
"__characters.anime__",
wildcards_dir=str(self.wildcards_dir),
rng=random.Random(0),
)
self.assertEqual(expanded, "mahou shoujo")
def test_expand_wildcards_resolves_unique_leaf_names(self):
self.write_wildcard("characters/accessories/hats.txt", "wide brim hat\n")
expanded = expand_wildcards(
"__hats__",
wildcards_dir=str(self.wildcards_dir),
rng=random.Random(0),
)
self.assertEqual(expanded, "wide brim hat")
def test_unique_leaf_resolution_scans_the_tree_only_once(self):
self.write_wildcard("nested/first.txt", "one\n")
self.write_wildcard("nested/second.txt", "two\n")
self.write_wildcard("nested/third.txt", "three\n")
with patch.object(
wildcards_module,
"_iter_wildcard_files",
wraps=wildcards_module._iter_wildcard_files,
) as iter_files:
self.assertEqual(
get_wildcard_options("first", str(self.wildcards_dir)), ["one"]
)
self.assertEqual(
get_wildcard_options("second", str(self.wildcards_dir)), ["two"]
)
self.assertEqual(
get_wildcard_options("third", str(self.wildcards_dir)), ["three"]
)
self.assertEqual(iter_files.call_count, 1)
def test_ambiguous_leaf_wildcards_are_left_unchanged(self):
self.write_wildcard("characters/hats.txt", "beret\n")
self.write_wildcard("props/hats.txt", "helmet\n")
with self.assertLogs("A1111PromptNode", level="WARNING"):
expanded = expand_wildcards(
"__hats__",
wildcards_dir=str(self.wildcards_dir),
rng=random.Random(0),
)
self.assertEqual(expanded, "__hats__")
def test_expand_wildcards_resolves_dynamic_prompt_choices(self):
expanded = expand_wildcards(
"wearing {red|blue} dress",
wildcards_dir=str(self.wildcards_dir),
rng=random.Random(0),
)
self.assertEqual(expanded, "wearing blue dress")
def test_expand_wildcards_resolves_nested_wildcards_and_dynamic_prompts(self):
self.write_wildcard(
"outfits.txt",
"__tops__{, __accessories__|}\n",
)
self.write_wildcard("tops.txt", "shirt\n")
self.write_wildcard("accessories.txt", "hat\n")
expanded = expand_wildcards(
"__outfits__",
wildcards_dir=str(self.wildcards_dir),
rng=random.Random(0),
)
self.assertEqual(expanded, "shirt, hat")
def test_expand_wildcards_supports_dynamic_prompt_ranges(self):
expanded = expand_wildcards(
"{2$$red|blue|green}",
wildcards_dir=str(self.wildcards_dir),
rng=random.Random(0),
)
self.assertEqual(expanded, "green, red")
def test_expand_wildcards_normalizes_spacing_after_empty_dynamic_choice(self):
self.write_wildcard("colors.txt", "red\n")
self.write_wildcard("shirt.txt", "{0-0$$__colors__} sleeveless turtleneck\n")
expanded = expand_wildcards(
"__shirt__",
wildcards_dir=str(self.wildcards_dir),
rng=random.Random(0),
)
self.assertEqual(expanded, "sleeveless turtleneck")
def test_missing_wildcards_are_left_unchanged(self):
self.assertEqual(
expand_wildcards(
"test __missing__",
wildcards_dir=str(self.wildcards_dir),
rng=random.Random(0),
),
"test __missing__",
)
def test_recursive_loops_do_not_hang(self):
self.write_wildcard("loop.txt", "__loop__\n")
self.assertEqual(
expand_wildcards(
"__loop__",
wildcards_dir=str(self.wildcards_dir),
rng=random.Random(0),
max_depth=4,
),
"__loop__",
)
def test_token_count_expansion_uses_highest_scoring_wildcard_option(self):
self.write_wildcard("subject.txt", "cat\nvery very long subject\n")
expanded = expand_wildcards_for_token_count(
"portrait of __subject__",
wildcards_dir=str(self.wildcards_dir),
scorer=lambda value: len(value.split()),
)
self.assertEqual(expanded, "portrait of very very long subject")
def test_token_count_expansion_uses_max_dynamic_prompt_quantity(self):
expanded = expand_wildcards_for_token_count(
"{1-2$$red|blue|extremely long green}",
wildcards_dir=str(self.wildcards_dir),
scorer=lambda value: len(value.split()),
)
self.assertEqual(expanded, "extremely long green, blue")
def test_token_count_expansion_resolves_nested_wildcards(self):
self.write_wildcard("colors.txt", "red\nultramarine blue\n")
self.write_wildcard("outfits.txt", "__colors__ robe\nhat\n")
expanded = expand_wildcards_for_token_count(
"__outfits__",
wildcards_dir=str(self.wildcards_dir),
scorer=lambda value: len(value.split()),
)
self.assertEqual(expanded, "ultramarine blue robe")
if __name__ == "__main__":
unittest.main()