Major Features: - New modular prompt generator architecture (core/prompt_generator/) - 15+ taggers/analyzers: WD14, PixAI, JoyTag, NudeNet, Places365, etc. - Decision engine for intelligent prompt construction - Template-based prompt building with dynamic sections Standard Mode Enhancements: - Tag injection for unmapped high-confidence tags - Scene-specific prompts (20 scene groups from Places365) - Animal-specific prompts (20+ animal types) - Clothing-specific prompts (30+ clothing types) - Tag-driven detail enhancements (40+ tag mappings) - Semantic inclusion analysis using sentence-transformers Detailed Mode (API with reasoning): - 3-pass LLM generation with full enhancements - Falls back to Standard mode for local models - Comprehensive metadata with pass-by-pass details Removed: - Extreme mode (consolidated into Detailed mode) 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
907 lines
27 KiB
Markdown
907 lines
27 KiB
Markdown
# SID_PromptGenerator Implementation Plan (4.2.1)
|
|
|
|
## Overview
|
|
|
|
This plan implements the `SID_PromptGenerator` node as designed in `4.2.1-design.md`.
|
|
|
|
**Target Branch**: `dev/4.2.1` (to be created from main)
|
|
**Reference Branch**: `dev/4.3.0` (tagger/analyzer code)
|
|
|
|
---
|
|
|
|
## Phase Summary
|
|
|
|
| Phase | Description | Dependencies | Complexity |
|
|
|-------|-------------|--------------|------------|
|
|
| 1 | Core Infrastructure | None | Low |
|
|
| 2 | Config & Templates | Phase 1 | Medium |
|
|
| 3 | Decision Engine | Phase 1, 2 | Medium |
|
|
| 4 | Florence Mode | Phase 1 | Low |
|
|
| 5 | Quick Mode | Phase 1 | Low |
|
|
| 6 | Tagger Integration | Phase 1, 3 | High |
|
|
| 7 | Standard Mode | Phase 1-3, 6 | High |
|
|
| 8 | Detailed Mode | Phase 7 | Medium |
|
|
| 9 | Extreme Mode | Phase 7 | Medium |
|
|
| 10 | Node Integration | All | Low |
|
|
| 11 | Testing | All | Medium |
|
|
|
|
---
|
|
|
|
## Phase 1: Core Infrastructure
|
|
|
|
**Goal**: Set up base classes, types, and project structure.
|
|
|
|
### Tasks
|
|
|
|
- [ ] 1.1 Create directory structure
|
|
```
|
|
core/prompt_generator/
|
|
├── __init__.py
|
|
├── generator.py
|
|
├── types.py
|
|
├── modes/
|
|
│ └── __init__.py
|
|
├── templates/
|
|
│ └── __init__.py
|
|
├── decisions/
|
|
│ └── __init__.py
|
|
└── config/
|
|
```
|
|
|
|
- [ ] 1.2 Create `types.py` with core types
|
|
```python
|
|
from enum import Enum
|
|
from dataclasses import dataclass
|
|
from typing import Dict, List, Optional, Any
|
|
|
|
class AnalysisMode(Enum):
|
|
QUICK = "Quick"
|
|
STANDARD = "Standard"
|
|
DETAILED = "Detailed"
|
|
EXTREME = "Extreme"
|
|
|
|
class SubjectType(Enum):
|
|
WOMAN = "woman"
|
|
MAN = "man"
|
|
COUPLE = "couple"
|
|
GROUP = "group"
|
|
LANDSCAPE = "landscape"
|
|
ANIMAL = "animal"
|
|
VEHICLE = "vehicle"
|
|
PRODUCT = "product"
|
|
UNKNOWN = "unknown"
|
|
|
|
class NSFWLevel(Enum):
|
|
SAFE = "safe"
|
|
SUGGESTIVE = "suggestive"
|
|
EXPLICIT = "explicit"
|
|
|
|
class ShotType(Enum):
|
|
ECU = "extreme close-up"
|
|
CU = "close-up"
|
|
MCU = "medium close-up"
|
|
MS = "medium shot"
|
|
MFS = "medium full shot"
|
|
FS = "full shot"
|
|
MLS = "medium long shot"
|
|
LS = "long shot"
|
|
ELS = "extreme long shot"
|
|
UNKNOWN = "unknown"
|
|
|
|
@dataclass
|
|
class GeneratorResult:
|
|
prompt: str
|
|
metadata: Dict[str, Any]
|
|
|
|
@dataclass
|
|
class Decisions:
|
|
subject_type: SubjectType
|
|
gender: Optional[str]
|
|
nsfw_level: NSFWLevel
|
|
shot_type: ShotType
|
|
scene_type: Optional[str]
|
|
template_sections: List[str]
|
|
```
|
|
|
|
- [ ] 1.3 Create base mode interface in `modes/base.py`
|
|
```python
|
|
from abc import ABC, abstractmethod
|
|
from ..types import GeneratorResult
|
|
|
|
class BaseMode(ABC):
|
|
@abstractmethod
|
|
def execute(self, image, **kwargs) -> GeneratorResult:
|
|
pass
|
|
|
|
@property
|
|
@abstractmethod
|
|
def name(self) -> str:
|
|
pass
|
|
```
|
|
|
|
- [ ] 1.4 Create `generator.py` skeleton
|
|
```python
|
|
class SID_PromptGenerator:
|
|
CATEGORY = "SID Nodes"
|
|
RETURN_TYPES = ("STRING", "STRING")
|
|
RETURN_NAMES = ("prompt", "prompt_metadata")
|
|
FUNCTION = "generate"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
},
|
|
"optional": {
|
|
"llm_model": ("LLM_MODEL",),
|
|
"user_prompt": ("STRING",),
|
|
"analysis_mode": (["Quick", "Standard", "Detailed", "Extreme"],),
|
|
},
|
|
}
|
|
|
|
def generate(self, image, llm_model=None, user_prompt=None, analysis_mode="Standard"):
|
|
# Route to appropriate mode
|
|
pass
|
|
```
|
|
|
|
---
|
|
|
|
## Phase 2: Config & Templates
|
|
|
|
**Goal**: Create YAML config files and template loading system.
|
|
|
|
### Tasks
|
|
|
|
- [ ] 2.1 Create `config/prompt_templates.yaml`
|
|
- Woman template with all sections
|
|
- Man template
|
|
- Landscape template
|
|
- Animal template
|
|
- Vehicle template
|
|
- Product template
|
|
- Generic fallback template
|
|
|
|
- [ ] 2.2 Create `config/tag_decisions.yaml`
|
|
- Subject detection rules
|
|
- Gender detection rules
|
|
- NSFW detection rules
|
|
- Shot type mapping
|
|
- Section trigger rules
|
|
|
|
- [ ] 2.3 Create `templates/loader.py`
|
|
```python
|
|
import yaml
|
|
from pathlib import Path
|
|
|
|
class TemplateLoader:
|
|
def __init__(self, config_dir: Path = None):
|
|
self.config_dir = config_dir or Path(__file__).parent.parent / "config"
|
|
self._templates = None
|
|
self._decisions = None
|
|
|
|
def load_templates(self) -> dict:
|
|
if self._templates is None:
|
|
with open(self.config_dir / "prompt_templates.yaml") as f:
|
|
self._templates = yaml.safe_load(f)
|
|
return self._templates
|
|
|
|
def load_decisions(self) -> dict:
|
|
if self._decisions is None:
|
|
with open(self.config_dir / "tag_decisions.yaml") as f:
|
|
self._decisions = yaml.safe_load(f)
|
|
return self._decisions
|
|
```
|
|
|
|
- [ ] 2.4 Create `templates/selector.py`
|
|
- Select template based on subject type
|
|
- Handle fallback to generic template
|
|
|
|
- [ ] 2.5 Create `templates/builder.py`
|
|
- Build prompt from template sections
|
|
- Handle section conditions
|
|
- Assemble final prompt string
|
|
|
|
---
|
|
|
|
## Phase 3: Decision Engine
|
|
|
|
**Goal**: Create tag-based decision logic.
|
|
|
|
### Tasks
|
|
|
|
- [ ] 3.1 Create `decisions/subject.py`
|
|
```python
|
|
from ..types import SubjectType
|
|
|
|
class SubjectDetector:
|
|
SUBJECT_TAGS = {
|
|
SubjectType.WOMAN: ["1girl", "woman", "female", "lady", "female_focus"],
|
|
SubjectType.MAN: ["1boy", "man", "male", "gentleman", "male_focus"],
|
|
SubjectType.LANDSCAPE: ["no_humans", "scenery", "landscape"],
|
|
SubjectType.ANIMAL: [], # Uses wildlife_detected flag
|
|
SubjectType.VEHICLE: ["car", "motorcycle", "vehicle", "train", "airplane"],
|
|
}
|
|
|
|
def detect(self, tags: dict, analyzers: dict) -> SubjectType:
|
|
# Check tags against SUBJECT_TAGS
|
|
# Check wildlife_detected in analyzers
|
|
# Return highest confidence match
|
|
pass
|
|
```
|
|
|
|
- [ ] 3.2 Create `decisions/gender.py`
|
|
- Detect gender from tags (1girl, 1boy, etc.)
|
|
- Cross-reference with NudeNet output
|
|
- Handle ambiguous cases
|
|
|
|
- [ ] 3.3 Create `decisions/nsfw.py`
|
|
```python
|
|
from ..types import NSFWLevel
|
|
|
|
class NSFWDetector:
|
|
SUGGESTIVE_PARTS = ["cleavage", "covered_breast", "bare_shoulders"]
|
|
EXPLICIT_PARTS = ["exposed_breast", "exposed_genitalia"]
|
|
|
|
def detect(self, nudenet_output: dict) -> NSFWLevel:
|
|
parts = nudenet_output.get("parts", [])
|
|
if any(p in self.EXPLICIT_PARTS for p in parts):
|
|
return NSFWLevel.EXPLICIT
|
|
if any(p in self.SUGGESTIVE_PARTS for p in parts):
|
|
return NSFWLevel.SUGGESTIVE
|
|
return NSFWLevel.SAFE
|
|
```
|
|
|
|
- [ ] 3.4 Create `decisions/shot_type.py`
|
|
- Map shot type analyzer output to ShotType enum
|
|
- Handle confidence thresholds
|
|
|
|
- [ ] 3.5 Create `decisions/sections.py`
|
|
```python
|
|
class SectionTriggerEngine:
|
|
"""Determines which template sections to activate based on decisions."""
|
|
|
|
def get_active_sections(self, decisions: Decisions, tags: dict) -> List[str]:
|
|
sections = ["appearance", "expression", "lighting", "composition"] # Always active
|
|
|
|
# Body details: enabled for MS+ shots or NSFW
|
|
if decisions.shot_type in [ShotType.MS, ShotType.MFS, ShotType.FS, ...]:
|
|
sections.append("body")
|
|
if decisions.nsfw_level in [NSFWLevel.SUGGESTIVE, NSFWLevel.EXPLICIT]:
|
|
sections.append("body")
|
|
|
|
# Pose: enabled for MCU+ shots
|
|
if decisions.shot_type in [ShotType.MCU, ShotType.MS, ...]:
|
|
sections.append("pose")
|
|
|
|
# Environment: enabled unless simple background
|
|
if not self._is_simple_background(tags):
|
|
sections.append("environment")
|
|
|
|
# Clothing: enabled if clothing tags present
|
|
if self._has_clothing_tags(tags):
|
|
sections.append("clothing")
|
|
|
|
return sections
|
|
```
|
|
|
|
- [ ] 3.6 Create `decisions/engine.py` (orchestrator)
|
|
```python
|
|
class DecisionEngine:
|
|
def __init__(self):
|
|
self.subject_detector = SubjectDetector()
|
|
self.gender_detector = GenderDetector()
|
|
self.nsfw_detector = NSFWDetector()
|
|
self.shot_detector = ShotTypeDetector()
|
|
self.section_engine = SectionTriggerEngine()
|
|
|
|
def make_decisions(self, tags: dict, analyzers: dict) -> Decisions:
|
|
subject = self.subject_detector.detect(tags, analyzers)
|
|
gender = self.gender_detector.detect(tags, analyzers)
|
|
nsfw = self.nsfw_detector.detect(analyzers.get("nudenet", {}))
|
|
shot = self.shot_detector.detect(analyzers.get("shot_type", {}))
|
|
|
|
decisions = Decisions(
|
|
subject_type=subject,
|
|
gender=gender,
|
|
nsfw_level=nsfw,
|
|
shot_type=shot,
|
|
scene_type=analyzers.get("places365", {}).get("top_category"),
|
|
template_sections=[] # Filled next
|
|
)
|
|
|
|
decisions.template_sections = self.section_engine.get_active_sections(
|
|
decisions, tags
|
|
)
|
|
|
|
return decisions
|
|
```
|
|
|
|
---
|
|
|
|
## Phase 4: Florence Mode
|
|
|
|
**Goal**: Implement fallback mode when no LLM is provided.
|
|
|
|
### Tasks
|
|
|
|
- [ ] 4.1 Create `modes/florence.py`
|
|
```python
|
|
from ..types import GeneratorResult
|
|
from .base import BaseMode
|
|
|
|
class FlorenceMode(BaseMode):
|
|
MODEL_ID = "MiaoshouAI/Florence-2-large-PromptGen-v2.0"
|
|
|
|
@property
|
|
def name(self) -> str:
|
|
return "florence"
|
|
|
|
def execute(self, image, **kwargs) -> GeneratorResult:
|
|
# Load Florence model
|
|
# Run MORE_DETAILED_CAPTION
|
|
# Return result
|
|
pass
|
|
|
|
def _load_model(self):
|
|
# Lazy load Florence model
|
|
pass
|
|
|
|
def _run_caption(self, image) -> str:
|
|
# Execute Florence captioning
|
|
pass
|
|
```
|
|
|
|
- [ ] 4.2 Integrate Florence model loading from existing codebase
|
|
- Reference: `core/compose/tokenizer/caption_extractor.py` from 4.3.0
|
|
- Reuse existing Florence utilities if available
|
|
|
|
- [ ] 4.3 Add caching for Florence model to avoid reloading
|
|
|
|
---
|
|
|
|
## Phase 5: Quick Mode
|
|
|
|
**Goal**: Implement single-shot LLM mode (no taggers).
|
|
|
|
### Tasks
|
|
|
|
- [ ] 5.1 Create `modes/quick.py`
|
|
```python
|
|
class QuickMode(BaseMode):
|
|
PROMPT_TEMPLATE = """
|
|
Describe this image in detail for AI image generation.
|
|
Include: subject, appearance, clothing, pose, environment,
|
|
lighting, composition. Write as flowing natural sentences.
|
|
Do not include meta-commentary about the image itself.
|
|
"""
|
|
|
|
@property
|
|
def name(self) -> str:
|
|
return "quick"
|
|
|
|
def execute(self, image, llm_model, **kwargs) -> GeneratorResult:
|
|
# Call LLM with single prompt
|
|
response = self._call_llm(image, llm_model, self.PROMPT_TEMPLATE)
|
|
|
|
return GeneratorResult(
|
|
prompt=response,
|
|
metadata={
|
|
"mode": "quick",
|
|
"source": "llm_single_shot",
|
|
"llm_provider": llm_model.provider,
|
|
"llm_model": llm_model.model_name,
|
|
"taggers_executed": False,
|
|
"llm_executed": True,
|
|
"llm_passes": 1,
|
|
}
|
|
)
|
|
|
|
def _call_llm(self, image, llm_model, prompt: str) -> str:
|
|
# Use existing LLM infrastructure
|
|
pass
|
|
```
|
|
|
|
- [ ] 5.2 Integrate with existing LLM infrastructure
|
|
- Reference: `core/llm/` or existing LLM calling code
|
|
- Support both local (Qwen) and API (Claude, GPT-4) models
|
|
|
|
---
|
|
|
|
## Phase 6: Tagger Integration
|
|
|
|
**Goal**: Port 22 taggers/analyzers from 4.3.0.
|
|
|
|
### Tasks
|
|
|
|
- [ ] 6.1 Create `taggers/__init__.py` with unified interface
|
|
```python
|
|
class TaggerRunner:
|
|
def __init__(self):
|
|
self.taggers = self._load_taggers()
|
|
|
|
def run_all(self, image) -> dict:
|
|
results = {}
|
|
for name, tagger in self.taggers.items():
|
|
try:
|
|
results[name] = tagger.run(image)
|
|
except Exception as e:
|
|
results[name] = {"error": str(e)}
|
|
return results
|
|
```
|
|
|
|
- [ ] 6.2 Port tagger implementations from 4.3.0
|
|
- Reference: `core/compose/tokenizer/tagger_extractor.py`
|
|
- Taggers to port:
|
|
- WD14
|
|
- PixAI
|
|
- JoyTag
|
|
- NudeNet
|
|
- Pose
|
|
|
|
- [ ] 6.3 Port analyzer implementations from 4.3.0
|
|
- Reference: `core/compose/tokenizer/analyzer_extractor.py`
|
|
- Analyzers to port:
|
|
- Photography
|
|
- IQA (Image Quality Assessment)
|
|
- Composition
|
|
- Shot Type
|
|
- Places365
|
|
- CLIP Scene
|
|
- Saliency
|
|
- Lighting
|
|
- Fashion Color
|
|
- DeepFashion
|
|
- Wildlife
|
|
|
|
- [ ] 6.4 Create tagger config loader
|
|
- Reference: `core/config/prompt_generator.yaml` from 4.3.0
|
|
- Load confidence thresholds
|
|
- Load enabled/disabled flags
|
|
|
|
- [ ] 6.5 Add parallel execution for taggers
|
|
```python
|
|
import concurrent.futures
|
|
|
|
def run_taggers_parallel(self, image, max_workers=4) -> dict:
|
|
results = {}
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
|
|
futures = {
|
|
executor.submit(tagger.run, image): name
|
|
for name, tagger in self.taggers.items()
|
|
}
|
|
for future in concurrent.futures.as_completed(futures):
|
|
name = futures[future]
|
|
try:
|
|
results[name] = future.result()
|
|
except Exception as e:
|
|
results[name] = {"error": str(e)}
|
|
return results
|
|
```
|
|
|
|
---
|
|
|
|
## Phase 7: Standard Mode
|
|
|
|
**Goal**: Implement tag-driven template-based LLM mode.
|
|
|
|
### Tasks
|
|
|
|
- [ ] 7.1 Create `modes/standard.py`
|
|
```python
|
|
class StandardMode(BaseMode):
|
|
def __init__(self):
|
|
self.tagger_runner = TaggerRunner()
|
|
self.decision_engine = DecisionEngine()
|
|
self.template_loader = TemplateLoader()
|
|
self.template_builder = TemplateBuilder()
|
|
|
|
@property
|
|
def name(self) -> str:
|
|
return "standard"
|
|
|
|
def execute(self, image, llm_model, **kwargs) -> GeneratorResult:
|
|
# Step 1: Run all taggers
|
|
tagger_results = self.tagger_runner.run_all(image)
|
|
|
|
# Step 2: Make decisions based on tags
|
|
decisions = self.decision_engine.make_decisions(
|
|
tags=self._extract_tags(tagger_results),
|
|
analyzers=self._extract_analyzers(tagger_results)
|
|
)
|
|
|
|
# Step 3: Build template prompt from decisions
|
|
template_prompt = self.template_builder.build(
|
|
decisions=decisions,
|
|
templates=self.template_loader.load_templates()
|
|
)
|
|
|
|
# Step 4: Call LLM with template prompt
|
|
llm_response = self._call_llm(image, llm_model, template_prompt)
|
|
|
|
return GeneratorResult(
|
|
prompt=llm_response,
|
|
metadata={
|
|
"mode": "standard",
|
|
"source": "llm_template",
|
|
"decisions": asdict(decisions),
|
|
"taggers": tagger_results,
|
|
"template": {
|
|
"selected": decisions.subject_type.value,
|
|
"sections_used": decisions.template_sections,
|
|
},
|
|
# ... other metadata
|
|
}
|
|
)
|
|
```
|
|
|
|
- [ ] 7.2 Implement template builder
|
|
```python
|
|
class TemplateBuilder:
|
|
def build(self, decisions: Decisions, templates: dict) -> str:
|
|
# Get template for subject type
|
|
template = templates["templates"].get(
|
|
decisions.subject_type.value,
|
|
templates["templates"]["generic"]
|
|
)
|
|
|
|
# Build prompt from active sections
|
|
sections = []
|
|
sections.append(template["base_prompt"])
|
|
|
|
for section_name in decisions.template_sections:
|
|
if section_name in template["sections"]:
|
|
sections.append(template["sections"][section_name]["prompt"])
|
|
|
|
return "\n".join(sections)
|
|
```
|
|
|
|
- [ ] 7.3 Test Standard mode end-to-end
|
|
|
|
---
|
|
|
|
## Phase 8: Detailed Mode
|
|
|
|
**Goal**: Implement multi-pass LLM mode (2-3 passes).
|
|
|
|
### Tasks
|
|
|
|
- [ ] 8.1 Create `modes/detailed.py`
|
|
```python
|
|
class DetailedMode(BaseMode):
|
|
PASS_CONFIG = [
|
|
{
|
|
"name": "pass_1",
|
|
"sections": ["subject", "appearance"],
|
|
"prompt": "Focus on the subject's identity and physical appearance..."
|
|
},
|
|
{
|
|
"name": "pass_2",
|
|
"sections": ["clothing", "pose", "environment"],
|
|
"prompt": "Focus on clothing, pose, and environment..."
|
|
},
|
|
{
|
|
"name": "pass_3",
|
|
"sections": ["lighting", "composition", "assembly"],
|
|
"prompt": "Focus on lighting, composition, and assemble final description..."
|
|
},
|
|
]
|
|
|
|
def execute(self, image, llm_model, **kwargs) -> GeneratorResult:
|
|
# Check if LLM supports reasoning
|
|
if not self._supports_reasoning(llm_model):
|
|
# Fall back to Standard mode
|
|
return StandardMode().execute(image, llm_model, **kwargs)
|
|
|
|
# Run taggers
|
|
tagger_results = self.tagger_runner.run_all(image)
|
|
decisions = self.decision_engine.make_decisions(...)
|
|
|
|
# Multi-pass execution
|
|
pass_outputs = {}
|
|
context = ""
|
|
|
|
for pass_config in self.PASS_CONFIG:
|
|
prompt = self._build_pass_prompt(pass_config, decisions, context)
|
|
response = self._call_llm(image, llm_model, prompt)
|
|
pass_outputs[pass_config["name"]] = response
|
|
context += f"\n{response}"
|
|
|
|
# Final assembly
|
|
final_prompt = self._assemble_passes(pass_outputs)
|
|
|
|
return GeneratorResult(
|
|
prompt=final_prompt,
|
|
metadata={
|
|
"mode": "detailed",
|
|
"llm_passes": len(self.PASS_CONFIG),
|
|
"pass_outputs": pass_outputs,
|
|
# ...
|
|
}
|
|
)
|
|
```
|
|
|
|
- [ ] 8.2 Implement `_supports_reasoning()` check
|
|
- Check `llm_model.supports_reasoning` flag
|
|
- Or infer from model name (claude, gpt-4, gemini)
|
|
|
|
- [ ] 8.3 Implement pass assembly logic
|
|
|
|
---
|
|
|
|
## Phase 9: Extreme Mode
|
|
|
|
**Goal**: Implement component-by-component LLM mode (10+ passes).
|
|
|
|
### Tasks
|
|
|
|
- [ ] 9.1 Create `modes/extreme.py`
|
|
```python
|
|
class ExtremeMode(BaseMode):
|
|
COMPONENT_PASSES = [
|
|
{"name": "subject", "prompt": "Identify the main subject..."},
|
|
{"name": "ethnicity", "prompt": "Analyze ethnicity/race..."},
|
|
{"name": "hair", "prompt": "Describe hair in detail..."},
|
|
{"name": "face", "prompt": "Describe facial features..."},
|
|
{"name": "eyes", "prompt": "Describe eyes in detail..."},
|
|
{"name": "body", "prompt": "Describe body/figure..."},
|
|
{"name": "clothing", "prompt": "Describe clothing..."},
|
|
{"name": "accessories", "prompt": "Describe accessories..."},
|
|
{"name": "pose", "prompt": "Describe pose..."},
|
|
{"name": "expression", "prompt": "Describe expression..."},
|
|
{"name": "environment", "prompt": "Describe environment..."},
|
|
{"name": "lighting", "prompt": "Describe lighting..."},
|
|
{"name": "composition", "prompt": "Describe composition..."},
|
|
{"name": "assembly", "prompt": "Assemble final description..."},
|
|
]
|
|
|
|
def execute(self, image, llm_model, **kwargs) -> GeneratorResult:
|
|
if not self._supports_reasoning(llm_model):
|
|
return StandardMode().execute(image, llm_model, **kwargs)
|
|
|
|
# Run taggers for decisions
|
|
tagger_results = self.tagger_runner.run_all(image)
|
|
decisions = self.decision_engine.make_decisions(...)
|
|
|
|
# Filter passes based on decisions
|
|
active_passes = self._filter_passes(decisions)
|
|
|
|
# Execute each component pass
|
|
component_outputs = {}
|
|
for pass_config in active_passes:
|
|
response = self._call_llm(image, llm_model, pass_config["prompt"])
|
|
component_outputs[pass_config["name"]] = response
|
|
|
|
# Final assembly
|
|
final_prompt = self._assemble_components(component_outputs)
|
|
|
|
return GeneratorResult(
|
|
prompt=final_prompt,
|
|
metadata={
|
|
"mode": "extreme",
|
|
"llm_passes": len(active_passes),
|
|
"component_outputs": component_outputs,
|
|
# ...
|
|
}
|
|
)
|
|
```
|
|
|
|
- [ ] 9.2 Implement pass filtering based on decisions
|
|
- Skip body pass for ECU/CU shots
|
|
- Skip environment for simple backgrounds
|
|
- Add/remove passes based on subject type
|
|
|
|
- [ ] 9.3 Implement component assembly logic
|
|
|
|
---
|
|
|
|
## Phase 10: Node Integration
|
|
|
|
**Goal**: Wire everything together in the ComfyUI node.
|
|
|
|
### Tasks
|
|
|
|
- [ ] 10.1 Complete `generator.py` with mode routing
|
|
```python
|
|
class SID_PromptGenerator:
|
|
def __init__(self):
|
|
self.modes = {
|
|
"override": OverrideMode(),
|
|
"florence": FlorenceMode(),
|
|
"quick": QuickMode(),
|
|
"standard": StandardMode(),
|
|
"detailed": DetailedMode(),
|
|
"extreme": ExtremeMode(),
|
|
}
|
|
|
|
def generate(self, image, llm_model=None, user_prompt=None, analysis_mode="Standard"):
|
|
import json
|
|
import time
|
|
|
|
start_time = time.time()
|
|
|
|
# Override mode
|
|
if user_prompt is not None and user_prompt.strip():
|
|
result = self.modes["override"].execute(user_prompt=user_prompt)
|
|
return self._format_output(result, start_time)
|
|
|
|
# Florence mode
|
|
if llm_model is None:
|
|
result = self.modes["florence"].execute(image=image)
|
|
return self._format_output(result, start_time)
|
|
|
|
# LLM modes
|
|
mode_key = analysis_mode.lower()
|
|
result = self.modes[mode_key].execute(
|
|
image=image,
|
|
llm_model=llm_model
|
|
)
|
|
return self._format_output(result, start_time)
|
|
|
|
def _format_output(self, result: GeneratorResult, start_time: float):
|
|
result.metadata["timing"] = {
|
|
"total_ms": int((time.time() - start_time) * 1000)
|
|
}
|
|
return (result.prompt, json.dumps(result.metadata, indent=2))
|
|
```
|
|
|
|
- [ ] 10.2 Register node in `__init__.py`
|
|
```python
|
|
from .generator import SID_PromptGenerator
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"SID_PromptGenerator": SID_PromptGenerator,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"SID_PromptGenerator": "SID Prompt Generator",
|
|
}
|
|
```
|
|
|
|
- [ ] 10.3 Add node to main `nodes/__init__.py` exports
|
|
|
|
- [ ] 10.4 Create `modes/override.py`
|
|
```python
|
|
class OverrideMode(BaseMode):
|
|
@property
|
|
def name(self) -> str:
|
|
return "override"
|
|
|
|
def execute(self, user_prompt, **kwargs) -> GeneratorResult:
|
|
return GeneratorResult(
|
|
prompt=user_prompt,
|
|
metadata={
|
|
"mode": "override",
|
|
"source": "user_prompt",
|
|
"processing": {
|
|
"taggers_executed": False,
|
|
"florence_executed": False,
|
|
"llm_executed": False,
|
|
}
|
|
}
|
|
)
|
|
```
|
|
|
|
---
|
|
|
|
## Phase 11: Testing
|
|
|
|
**Goal**: Comprehensive testing of all components.
|
|
|
|
### Tasks
|
|
|
|
- [ ] 11.1 Create test fixtures
|
|
- Sample images for each subject type
|
|
- Mock tagger outputs
|
|
- Mock LLM responses
|
|
|
|
- [ ] 11.2 Unit tests for decision engine
|
|
```python
|
|
def test_subject_detection_woman():
|
|
detector = SubjectDetector()
|
|
tags = {"wd14": {"1girl": 0.98, "long_hair": 0.9}}
|
|
assert detector.detect(tags, {}) == SubjectType.WOMAN
|
|
|
|
def test_nsfw_detection_suggestive():
|
|
detector = NSFWDetector()
|
|
nudenet = {"parts": ["cleavage"]}
|
|
assert detector.detect(nudenet) == NSFWLevel.SUGGESTIVE
|
|
```
|
|
|
|
- [ ] 11.3 Unit tests for template builder
|
|
- Test template selection
|
|
- Test section activation
|
|
- Test prompt assembly
|
|
|
|
- [ ] 11.4 Integration tests for each mode
|
|
- Override mode: pass-through works
|
|
- Florence mode: generates description
|
|
- Quick mode: LLM called once
|
|
- Standard mode: taggers + LLM
|
|
- Detailed mode: multiple passes
|
|
- Extreme mode: component passes
|
|
|
|
- [ ] 11.5 End-to-end tests
|
|
- Full pipeline with real images
|
|
- Verify metadata completeness
|
|
- Verify prompt quality
|
|
|
|
- [ ] 11.6 Performance benchmarks
|
|
- Measure tagger execution time
|
|
- Measure LLM call latency
|
|
- Optimize bottlenecks
|
|
|
|
---
|
|
|
|
## File Checklist
|
|
|
|
When complete, the following files should exist:
|
|
|
|
```
|
|
core/prompt_generator/
|
|
├── __init__.py # Phase 10
|
|
├── generator.py # Phase 1, 10
|
|
├── types.py # Phase 1
|
|
├── modes/
|
|
│ ├── __init__.py # Phase 1
|
|
│ ├── base.py # Phase 1
|
|
│ ├── override.py # Phase 10
|
|
│ ├── florence.py # Phase 4
|
|
│ ├── quick.py # Phase 5
|
|
│ ├── standard.py # Phase 7
|
|
│ ├── detailed.py # Phase 8
|
|
│ └── extreme.py # Phase 9
|
|
├── templates/
|
|
│ ├── __init__.py # Phase 2
|
|
│ ├── loader.py # Phase 2
|
|
│ ├── selector.py # Phase 2
|
|
│ └── builder.py # Phase 2, 7
|
|
├── decisions/
|
|
│ ├── __init__.py # Phase 3
|
|
│ ├── engine.py # Phase 3
|
|
│ ├── subject.py # Phase 3
|
|
│ ├── gender.py # Phase 3
|
|
│ ├── nsfw.py # Phase 3
|
|
│ ├── shot_type.py # Phase 3
|
|
│ └── sections.py # Phase 3
|
|
├── taggers/
|
|
│ ├── __init__.py # Phase 6
|
|
│ ├── runner.py # Phase 6
|
|
│ └── [individual tagger files] # Phase 6
|
|
├── config/
|
|
│ ├── prompt_templates.yaml # Phase 2
|
|
│ └── tag_decisions.yaml # Phase 2
|
|
└── tests/
|
|
├── __init__.py # Phase 11
|
|
├── test_decisions.py # Phase 11
|
|
├── test_templates.py # Phase 11
|
|
├── test_modes.py # Phase 11
|
|
└── fixtures/ # Phase 11
|
|
```
|
|
|
|
---
|
|
|
|
## Recommended Implementation Order
|
|
|
|
1. **Phase 1** - Core infrastructure (foundation)
|
|
2. **Phase 2** - Config & Templates (needed by all modes)
|
|
3. **Phase 4** - Florence mode (simplest, no dependencies)
|
|
4. **Phase 5** - Quick mode (simple LLM integration)
|
|
5. **Phase 3** - Decision engine (needed for Standard+)
|
|
6. **Phase 6** - Tagger integration (needed for Standard+)
|
|
7. **Phase 7** - Standard mode (core functionality)
|
|
8. **Phase 8** - Detailed mode (builds on Standard)
|
|
9. **Phase 9** - Extreme mode (builds on Standard)
|
|
10. **Phase 10** - Node integration (wire everything)
|
|
11. **Phase 11** - Testing (throughout + final)
|
|
|
|
---
|
|
|
|
## Notes
|
|
|
|
- Start with Phases 1, 4, 5 for a minimal working prototype
|
|
- Phase 6 (Tagger Integration) is the most complex - copy carefully from 4.3.0
|
|
- Phase 7 (Standard Mode) is the core value - invest most effort here
|
|
- Phases 8, 9 can be simplified initially with fewer passes
|
|
- Test incrementally after each phase
|