# 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