Compare commits
22
Commits
builtin-hook
...
main
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
070723c767 | ||
|
|
635b4bbdbe | ||
|
|
69d03a8763 | ||
|
|
8b5269bea4 | ||
|
|
9e67b317d2 | ||
|
|
3db87dc406 | ||
|
|
98ad657586 | ||
|
|
55339cede5 | ||
|
|
620f422175 | ||
|
|
639c67ba1a | ||
|
|
d7d12aa523 | ||
|
|
84c8c03523 | ||
|
|
4d6257bd07 | ||
|
|
c8961a3b55 | ||
|
|
68f5fdc75f | ||
|
|
552946ff73 | ||
|
|
c0187b4d75 | ||
|
|
6c8c50975d | ||
|
|
ab3d80d3e3 | ||
|
|
1cbfe5fe06 | ||
|
|
8817cabfe7 | ||
|
|
d59cdd9680 |
@@ -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
@@ -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
|
||||
@@ -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_
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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.
|
||||
@@ -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" />
|
||||

|
||||
|
||||
### Negative prompt node
|
||||
|
||||
<img width="363" height="241" alt="image" src="https://github.com/user-attachments/assets/ba1da788-dc47-4939-bcbe-d1699ea9f872" />
|
||||

|
||||
|
||||
### Autocomplete/Wildcards
|
||||
|
||||

|
||||
|
||||
### 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" />
|
||||

|
||||
|
||||
### 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.
|
||||
|
||||
---
|
||||
|
||||
[](https://deepwiki.com/Enferlain/ComfyUI-A1111-cond)
|
||||
|
||||
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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() {
|
||||
|
||||
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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
@@ -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())
|
||||
|
||||
@@ -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
@@ -1 +1,3 @@
|
||||
lark
|
||||
lark
|
||||
torch
|
||||
aiohttp
|
||||
@@ -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,
|
||||
|
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user