🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
541 lines
17 KiB
Markdown
541 lines
17 KiB
Markdown
# Image Tokenization & Composition - Implementation Plan
|
|
|
|
## Related Documents
|
|
|
|
- **Specification**: [image_tokenization_spec.md](./image_tokenization_spec.md) - Full technical specification
|
|
- **Project Root**: `E:\ComfyUI\custom_nodes\ComfyUI-AI-Photography-Toolkit`
|
|
- **Target Module**: `core/compose/` (replaces current StandardGenerator)
|
|
|
|
---
|
|
|
|
## Project Overview
|
|
|
|
**Goal**: Replace the current broken `StandardGenerator` with a robust tokenization → classification → assembly pipeline that produces high-quality natural language prompts from image analysis metadata.
|
|
|
|
**Two Modes**:
|
|
- `Standard` → Rule-based assembly (fast, free, deterministic)
|
|
- `Enhance with AI` → LLM-based assembly (higher quality, costs money)
|
|
|
|
---
|
|
|
|
## Implementation Phases
|
|
|
|
### Phase 1: Core Infrastructure [Priority: HIGH]
|
|
|
|
**Files to Create:**
|
|
|
|
```
|
|
core/compose/
|
|
├── tokenizer/
|
|
│ ├── __init__.py
|
|
│ ├── base.py # ImageToken dataclass, TokenType enum
|
|
│ ├── tagger_extractor.py # Extract from wd14, pixai, joytag, nudenet
|
|
│ ├── analyzer_extractor.py # Extract from photography, iqa, composition, saliency
|
|
│ ├── caption_extractor.py # Extract from florence captions (sentence/clause parsing)
|
|
│ └── normalizer.py # Deduplication, confidence aggregation
|
|
├── classifier/
|
|
│ ├── __init__.py
|
|
│ ├── base.py # TokenClassification dataclass
|
|
│ ├── deterministic.py # Layer 1-2: structured keys, source routing
|
|
│ ├── dictionary.py # Layer 3: keyword dictionaries
|
|
│ ├── embeddings.py # Layer 4: semantic similarity (optional)
|
|
│ └── categories.py # Category definitions and dictionaries
|
|
├── assembler/
|
|
│ ├── __init__.py
|
|
│ ├── base.py # AssemblyResult dataclass
|
|
│ ├── rule_based.py # Phase 5: Template-based assembly
|
|
│ ├── llm_based.py # Phase 6: LLM-based assembly
|
|
│ ├── templates.py # Template definitions
|
|
│ └── phrase_mappings.py # Token → natural phrase dictionaries
|
|
├── validator/
|
|
│ ├── __init__.py
|
|
│ ├── content_preservation.py
|
|
│ ├── tag_alignment.py
|
|
│ ├── coherence.py
|
|
│ └── report.py
|
|
└── pipeline.py # Main orchestrator (replaces current pipeline.py)
|
|
```
|
|
|
|
**Tasks:**
|
|
|
|
1. [x] Create `ImageToken` dataclass with fields: text, confidence, source, token_type, metadata
|
|
2. [x] Create `TokenType` enum: TAG, PHRASE, KEY_VALUE, SENTENCE
|
|
3. [x] Create `CanonicalCategory` enum with all 9 categories
|
|
4. [x] Create `TokenClassification` dataclass with primary/secondary categories
|
|
|
|
**Commit:** `c7859cf` - feat(compose): implement Phase 1 - Core Infrastructure
|
|
|
|
---
|
|
|
|
### Phase 2: Tokenization [Priority: HIGH]
|
|
|
|
**Spec Reference**: Phase 1 in image_tokenization_spec.md
|
|
|
|
**Tasks:**
|
|
|
|
1. [x] **Tagger Extractor** (`tagger_extractor.py`)
|
|
- Extract from wd14, pixai, joytag dictionaries
|
|
- Handle both formats: `{tag: conf}` and `{tags: [...], attributes: {...}}`
|
|
- Set token_type = TAG
|
|
|
|
2. [x] **Analyzer Extractor** (`analyzer_extractor.py`)
|
|
- Extract from photography, iqa, composition, saliency
|
|
- Preserve source info for Layer 2 classification
|
|
- Set token_type = TAG or KEY_VALUE
|
|
|
|
3. [x] **Caption Extractor** (`caption_extractor.py`)
|
|
- Sentence segmentation (split by `.!?`)
|
|
- Clause segmentation (split by `,`, `and`, `with`, `while`)
|
|
- Pattern matching for category-specific phrases
|
|
- META pattern filtering (remove "the image shows", etc.)
|
|
- Set token_type = PHRASE or SENTENCE
|
|
|
|
4. [x] **Normalizer** (`normalizer.py`)
|
|
- Lowercase normalization
|
|
- Whitespace cleanup
|
|
- Duplicate detection (exact + near-duplicate)
|
|
- Confidence aggregation for duplicates
|
|
- Source reliability weighting
|
|
|
|
**Test Cases:**
|
|
- [x] Test tagger extraction with sample wd14 output
|
|
- [x] Test caption extraction with sample florence_description
|
|
- [x] Test META filtering removes philosophical content
|
|
- [x] Test deduplication merges "brown hair" from multiple sources
|
|
|
|
**Commit:** `bdb58fd` - feat(compose): implement Phase 2 - Tokenization
|
|
|
|
---
|
|
|
|
### Phase 3: Classification [Priority: HIGH]
|
|
|
|
**Spec Reference**: Phase 3 in image_tokenization_spec.md
|
|
|
|
**Tasks:**
|
|
|
|
1. [x] **Category Definitions** (`categories.py`)
|
|
- Define all 9 canonical categories
|
|
- Define subcategories for Subject Details (hair, eyes, face, body, clothing, skin)
|
|
- Create keyword dictionaries with exact_matches, prefix_patterns, contains_patterns
|
|
|
|
2. [x] **Deterministic Classifier** (`deterministic.py`)
|
|
- Layer 1: florence_analyze key mapping (camera_angle → Composition, etc.)
|
|
- Layer 2: Source-based routing (photography → Lighting, iqa → Quality)
|
|
|
|
3. [x] **Dictionary Classifier** (`dictionary.py`)
|
|
- Layer 3: Keyword lookup with exact, prefix, suffix, contains matching
|
|
- Handle multi-category tokens (primary + secondary)
|
|
|
|
4. [ ] **Embedding Classifier** (`embeddings.py`) [OPTIONAL - defer if time-constrained]
|
|
- Layer 4: Semantic similarity using sentence-transformers
|
|
- Category centroid matching
|
|
- Cache embeddings for performance
|
|
|
|
5. [x] **Conflict Resolution**
|
|
- Handle mutually exclusive tokens (standing vs sitting)
|
|
- Keep highest confidence for conflicts
|
|
|
|
**Test Cases:**
|
|
- [x] Test "brown hair" → [Subject Details/hair]
|
|
- [x] Test "standing" → [Action/Pose]
|
|
- [x] Test "soft light" → [Lighting]
|
|
- [x] Test florence_analyze key "camera_angle: front" → [Composition]
|
|
- [x] Test multi-category: "looking at viewer" → [Action/Pose] primary, [Subject Details/eyes] secondary
|
|
|
|
**Commit:** `574b976` - feat(compose): implement Phase 3 - Classification
|
|
|
|
---
|
|
|
|
### Phase 4: Rule-Based Assembly [Priority: HIGH]
|
|
|
|
**Spec Reference**: Phase 5 in image_tokenization_spec.md
|
|
|
|
**Tasks:**
|
|
|
|
1. [ ] **Phrase Mappings** (`phrase_mappings.py`)
|
|
- Token → natural phrase dictionaries
|
|
- `"1girl"` → `"a woman"`
|
|
- `"looking at viewer"` → `"looking directly at the camera"`
|
|
- `"cowboy shot"` → `"framed from mid-thigh up"`
|
|
|
|
2. [ ] **Templates** (`templates.py`)
|
|
- Subject templates with variations
|
|
- Hair/eyes/clothing templates
|
|
- Action/environment/lighting templates
|
|
- Master template for full assembly
|
|
|
|
3. [ ] **Rule-Based Assembler** (`rule_based.py`)
|
|
- Implement all category generators:
|
|
- `generate_subject()`
|
|
- `generate_subject_details()` (hair, eyes, face, body, clothing)
|
|
- `generate_action()`
|
|
- `generate_environment()`
|
|
- `generate_lighting()`
|
|
- `generate_style()`
|
|
- `generate_composition()`
|
|
- Adjective ordering rules (length, texture, color)
|
|
- Pronoun resolution (She/Her based on gender detection)
|
|
- Sentence fusion for same-subject attributes
|
|
- Connector variation (and, with, featuring)
|
|
|
|
4. [ ] **Empty Category Handling**
|
|
- Skip empty categories gracefully
|
|
- Fallback phrases for critical categories
|
|
|
|
5. [ ] **Confidence-Based Filtering**
|
|
- Only include tokens above threshold
|
|
- Prioritize high-confidence tokens
|
|
|
|
**Test Cases:**
|
|
- [ ] Test hair tokens → "long, wavy brown hair" (correct adjective order)
|
|
- [ ] Test clothing grouping → "white Calvin Klein crop top with matching briefs"
|
|
- [ ] Test full assembly produces valid 3-4 paragraph output
|
|
- [ ] Test empty lighting category → graceful omission
|
|
|
|
---
|
|
|
|
### Phase 5: LLM-Based Assembly [Priority: MEDIUM]
|
|
|
|
**Spec Reference**: Phase 6 in image_tokenization_spec.md
|
|
|
|
**Tasks:**
|
|
|
|
1. [ ] **LLM Assembler** (`llm_based.py`)
|
|
- Single comprehensive prompt approach (primary)
|
|
- Section-by-section approach (optional)
|
|
- Template + LLM hybrid approach (optional)
|
|
|
|
2. [ ] **Prompt Templates**
|
|
- System prompt with rules and structure
|
|
- User prompt with JSON token input
|
|
- Negative instructions (don't invent, don't use raw tags)
|
|
|
|
3. [ ] **Provider Integration**
|
|
- Reuse existing providers from `core/compose/llm.py`:
|
|
- Local (Qwen)
|
|
- Anthropic
|
|
- OpenAI
|
|
- Gemini
|
|
|
|
4. [ ] **Output Validation**
|
|
- Schema validation (length, no raw tags, proper sentences)
|
|
- Semantic validation (key elements present)
|
|
- Retry on failure
|
|
|
|
5. [ ] **Caching**
|
|
- Token hash → description cache
|
|
- Section-level caching for hybrid approach
|
|
|
|
6. [ ] **Fallback Chain**
|
|
- Retry with backoff
|
|
- Fall back to rule-based on failure
|
|
|
|
**Test Cases:**
|
|
- [ ] Test LLM output doesn't contain raw tags like "1girl"
|
|
- [ ] Test output contains key elements from input tokens
|
|
- [ ] Test fallback to rule-based when LLM fails
|
|
|
|
---
|
|
|
|
### Phase 6: Validation Module [Priority: LOW]
|
|
|
|
**Spec Reference**: Validation section in image_tokenization_spec.md
|
|
|
|
**Tasks:**
|
|
|
|
1. [ ] **Content Preservation Checker**
|
|
- Word-level accounting
|
|
- Coverage percentage calculation
|
|
- Unaccounted words report
|
|
|
|
2. [ ] **Tag Alignment Checker**
|
|
- Compare extractions to original tags
|
|
- Match type detection (exact, semantic, partial)
|
|
- Category accuracy scoring
|
|
|
|
3. [ ] **Coherence Analyzer** [OPTIONAL - requires embeddings]
|
|
- Intra-category similarity
|
|
- Inter-category separation
|
|
- Contamination detection
|
|
|
|
4. [ ] **Report Generator**
|
|
- Aggregate all metrics
|
|
- Generate actionable items
|
|
- Overall quality score
|
|
|
|
---
|
|
|
|
### Phase 7: Integration [Priority: HIGH]
|
|
|
|
**Tasks:**
|
|
|
|
1. [ ] **Update Pipeline** (`pipeline.py`)
|
|
- Replace current `StandardGenerator` with new tokenizer → classifier → assembler pipeline
|
|
- Keep interface compatible with existing `ComposePipeline`
|
|
|
|
2. [ ] **Update Node** (`nodes/prompt_compose.py`)
|
|
- No changes needed if pipeline interface unchanged
|
|
- Verify mode switching works (Standard vs Enhance with AI)
|
|
|
|
3. [ ] **Configuration**
|
|
- Move dictionaries to YAML config files if too large for code
|
|
- Add verbosity control settings
|
|
|
|
4. [ ] **Cleanup**
|
|
- Remove old `standard.py` once new pipeline verified
|
|
- Update imports in `__init__.py`
|
|
|
|
---
|
|
|
|
### Phase 8: Token Grounding Visualization [Priority: LOW]
|
|
|
|
**Spec Reference**: Phase 7 in image_tokenization_spec.md
|
|
|
|
**Purpose**: Visual debugging and QA - show WHERE in the image each token spatially corresponds.
|
|
|
|
**Files to Create:**
|
|
|
|
```
|
|
core/compose/grounding/
|
|
├── __init__.py
|
|
├── router.py # Select grounding method by category
|
|
├── human_parsing.py # SCHP/Graphonomy integration
|
|
├── grounding_dino.py # Grounding DINO + SAM
|
|
├── clip_attention.py # CLIP-based heatmaps (fallback)
|
|
└── visualizer.py # Overlay composition, export
|
|
```
|
|
|
|
**Tasks:**
|
|
|
|
1. [ ] **Category Router** (`router.py`)
|
|
- Map token categories to appropriate grounding method
|
|
- Subject Details → human_parsing
|
|
- Environment → grounding_dino
|
|
- Abstract (lighting, style) → clip_attention or skip
|
|
|
|
2. [ ] **Human Parsing Integration** (`human_parsing.py`)
|
|
- SCHP or Graphonomy model loading
|
|
- Token → segment class mapping dictionary
|
|
- Mask extraction for hair, clothing, body parts
|
|
|
|
3. [ ] **Grounding DINO + SAM** (`grounding_dino.py`) [OPTIONAL]
|
|
- Token → natural language prompt conversion
|
|
- Bounding box detection
|
|
- SAM mask refinement
|
|
|
|
4. [ ] **CLIP Attention Fallback** (`clip_attention.py`) [OPTIONAL]
|
|
- GradCAM-based attention heatmaps
|
|
- For abstract concepts that can't be localized
|
|
|
|
5. [ ] **Visualizer** (`visualizer.py`)
|
|
- Single token heatmap overlay
|
|
- Multi-token composite (color-coded by category)
|
|
- Static image export (PNG)
|
|
- JSON annotation export
|
|
- Interactive HTML viewer [OPTIONAL]
|
|
|
|
6. [ ] **Integration with Validation**
|
|
- Spatial coherence checks
|
|
- Visual QA report generation
|
|
|
|
**Dependencies:**
|
|
- SCHP/Graphonomy (~150MB)
|
|
- face-parsing.PyTorch (~100MB) [optional]
|
|
- Grounding DINO (~700MB) [optional]
|
|
- SAM ViT-B (~400MB) [optional]
|
|
|
|
**Test Cases:**
|
|
- [ ] Human parsing correctly segments hair, clothing, body
|
|
- [ ] Token → segment mapping returns correct regions
|
|
- [ ] Visualization overlays are correctly positioned
|
|
- [ ] JSON export contains valid bounding boxes and masks
|
|
|
|
---
|
|
|
|
## File Dependency Order
|
|
|
|
Implement in this order to minimize blocking:
|
|
|
|
```
|
|
1. core/compose/tokenizer/base.py # No dependencies
|
|
2. core/compose/classifier/categories.py # No dependencies
|
|
3. core/compose/tokenizer/tagger_extractor.py
|
|
4. core/compose/tokenizer/analyzer_extractor.py
|
|
5. core/compose/tokenizer/caption_extractor.py
|
|
6. core/compose/tokenizer/normalizer.py
|
|
7. core/compose/classifier/deterministic.py
|
|
8. core/compose/classifier/dictionary.py
|
|
9. core/compose/assembler/phrase_mappings.py
|
|
10. core/compose/assembler/templates.py
|
|
11. core/compose/assembler/rule_based.py
|
|
12. core/compose/assembler/llm_based.py # Can reuse existing llm.py logic
|
|
13. core/compose/pipeline.py # Orchestrates everything
|
|
14. core/compose/validator/* # Optional, implement last
|
|
15. core/compose/grounding/* # Optional, Phase 8 - implement after validation
|
|
```
|
|
|
|
---
|
|
|
|
## Testing Strategy
|
|
|
|
### Unit Tests
|
|
|
|
Location: `tests/compose/`
|
|
|
|
```
|
|
tests/compose/
|
|
├── test_tokenizer.py
|
|
├── test_classifier.py
|
|
├── test_assembler.py
|
|
└── test_pipeline.py
|
|
```
|
|
|
|
### Integration Test
|
|
|
|
Create sample metadata file with real output from Analysis node, run through full pipeline, verify output quality.
|
|
|
|
### Manual Testing
|
|
|
|
1. Run in ComfyUI with real images
|
|
2. Compare output to current broken StandardGenerator
|
|
3. Verify "Enhance with AI" mode works with local Qwen model
|
|
|
|
---
|
|
|
|
## Dependencies
|
|
|
|
### Required (Already Installed)
|
|
- Python 3.10+
|
|
- PyYAML (for config files)
|
|
- re (regex, stdlib)
|
|
|
|
### Optional (For Embedding Classifier)
|
|
- sentence-transformers (`all-MiniLM-L6-v2`)
|
|
|
|
### Optional (For LLM Assembly)
|
|
- anthropic
|
|
- openai
|
|
- google-generativeai
|
|
- transformers (for local models)
|
|
|
|
---
|
|
|
|
## Success Criteria
|
|
|
|
### Minimum Viable Product (MVP)
|
|
|
|
- [ ] Tokenization extracts from all metadata sources
|
|
- [ ] Classification assigns tokens to correct categories
|
|
- [ ] Rule-based assembly produces grammatically correct output
|
|
- [ ] No more "her eyes are brown eyes" errors
|
|
- [ ] No more "in the middle of the image" meta-commentary
|
|
- [ ] Output follows canonical schema order
|
|
|
|
### Full Implementation
|
|
|
|
- [ ] All MVP criteria
|
|
- [ ] LLM assembly works as fallback/enhancement
|
|
- [ ] Validation module can score output quality
|
|
- [ ] Caching reduces redundant computation
|
|
- [ ] Performance: <100ms for rule-based, <3s for LLM
|
|
|
|
---
|
|
|
|
## Risk Mitigation
|
|
|
|
| Risk | Mitigation |
|
|
|------|------------|
|
|
| Caption parsing too fragile | Start with simple patterns, iterate based on failures |
|
|
| Classification accuracy low | Extensive keyword dictionaries, fall back to uncategorized |
|
|
| LLM output inconsistent | Strict validation, fallback to rule-based |
|
|
| Performance too slow | Cache aggressively, defer embedding classifier |
|
|
| Breaking existing workflows | Keep pipeline interface identical |
|
|
|
|
---
|
|
|
|
## Estimated Effort
|
|
|
|
| Phase | Effort | Priority |
|
|
|-------|--------|----------|
|
|
| Phase 1: Infrastructure | 2-3 hours | HIGH |
|
|
| Phase 2: Tokenization | 4-6 hours | HIGH |
|
|
| Phase 3: Classification | 4-6 hours | HIGH |
|
|
| Phase 4: Rule-Based Assembly | 6-8 hours | HIGH |
|
|
| Phase 5: LLM Assembly | 3-4 hours | MEDIUM |
|
|
| Phase 6: Validation | 4-6 hours | LOW |
|
|
| Phase 7: Integration | 2-3 hours | HIGH |
|
|
| Phase 8: Token Grounding Visualization | 6-10 hours | LOW |
|
|
| **Total (MVP: Phases 1-4,7)** | **18-26 hours** | |
|
|
| **Total (Full)** | **31-46 hours** | |
|
|
|
|
---
|
|
|
|
## Quick Start Commands
|
|
|
|
```bash
|
|
# Navigate to project
|
|
cd E:\ComfyUI\custom_nodes\ComfyUI-AI-Photography-Toolkit
|
|
|
|
# Create new module structure
|
|
mkdir -p core/compose/tokenizer
|
|
mkdir -p core/compose/classifier
|
|
mkdir -p core/compose/assembler
|
|
mkdir -p core/compose/validator
|
|
|
|
# Create __init__.py files
|
|
touch core/compose/tokenizer/__init__.py
|
|
touch core/compose/classifier/__init__.py
|
|
touch core/compose/assembler/__init__.py
|
|
touch core/compose/validator/__init__.py
|
|
|
|
# Run tests (once created)
|
|
python -m pytest tests/compose/ -v
|
|
```
|
|
|
|
---
|
|
|
|
## Context Recovery Notes
|
|
|
|
**If context is lost, read these files in order:**
|
|
|
|
1. `docs/image_tokenization_spec.md` - Full technical specification
|
|
2. `docs/implementation_plan.md` - This file, implementation roadmap
|
|
3. `CLAUDE.md` - Project overview and architecture
|
|
4. `core/compose/` - Current implementation (to be replaced)
|
|
|
|
**Current State Summary:**
|
|
|
|
The existing `StandardGenerator` in `core/compose/standard.py` produces broken output like:
|
|
- "her eyes are brown eyes" (grammar error)
|
|
- "standing in the middle of the image" (meta-commentary not filtered)
|
|
- Missing clothing details
|
|
- Redundant content
|
|
|
|
The new implementation will fix these issues with a proper tokenization → classification → assembly pipeline.
|
|
|
|
**Key Design Decisions:**
|
|
|
|
1. **Tokenization**: Extract ALL content from metadata into unified ImageToken format
|
|
2. **Classification**: 5-layer cascade (deterministic → dictionary → embeddings)
|
|
3. **Assembly (Rule-Based)**: Templates + phrase mappings + grammar rules
|
|
4. **Assembly (LLM)**: Fallback/enhancement using existing LLM providers
|
|
5. **Validation**: Optional quality scoring module
|
|
6. **Token Grounding**: Optional visual debugging showing spatial token locations
|
|
|
|
---
|
|
|
|
## Next Steps
|
|
|
|
When resuming work:
|
|
|
|
1. Read this plan and the spec document
|
|
2. Check which tasks are marked complete
|
|
3. Continue from the next uncompleted task
|
|
4. Update checkboxes as tasks complete
|
|
5. Test incrementally - don't wait until end
|
|
|
|
---
|
|
|
|
*Last Updated: December 2024*
|
|
*Version: 1.0*
|