slahiri-ComfyUI-AI-Photography-Toolkit/core/compose/assembler/base.py

"""Base classes for prompt assembly (Phase 4).

The assembler takes classified tokens from Phase 3 and combines them
into natural language prompts using templates and grammar rules.
"""

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from enum import Enum
from typing import Dict, List, Optional, Any

from ..classifier.base import ClassifiedImage, TokenClassification
from ..classifier.categories import (
    CanonicalCategory,
    SubjectDetailType,
    CATEGORY_ORDER,
    SUBJECT_DETAIL_ORDER,
    CATEGORY_DISPLAY_NAMES,
)


class PromptStyle(Enum):
    """Output prompt style."""
    NATURAL = "natural"  # Full sentences, flowing prose (legacy)
    TAGS = "tags"  # Comma-separated tags (booru/SDXL style)
    HYBRID = "hybrid"  # Sentence intro + tag details (legacy)
    STRUCTURED = "structured"  # Category-segregated output with labels
    FULL = "full"  # Florence description + structured tags (preserves all info)
    ZIMAGE = "zimage"  # Z-Image optimized: Florence prose + organized attributes


@dataclass
class AssemblerConfig:
    """Configuration for prompt assembly."""
    style: PromptStyle = PromptStyle.NATURAL
    max_tokens_per_category: int = 50  # Default limit
    include_quality_boosters: bool = True
    include_technical: bool = True
    min_confidence: float = 0.3
    deduplicate: bool = True
    use_pronouns: bool = True  # Use pronouns after first subject mention

    # Category-specific token limits (overrides max_tokens_per_category)
    # Lower for dense categories, higher for scene-critical categories
    category_limits: Dict[str, int] = field(default_factory=lambda: {
        "quality_boosters": 10,
        "style_medium": 8,
        "subject": 5,
        "subject_details": 25,  # CRITICAL: Cap this to prevent overwhelming
        "action_pose": 10,
        "environment": 15,      # Important for scene recreation
        "lighting": 12,         # Important for mood
        "composition": 12,
        "technical": 5,
    })

    def get_limit_for_category(self, category: CanonicalCategory) -> int:
        """Get the token limit for a specific category.

        If max_tokens_per_category is set higher than the default (50),
        use it as the limit for all categories (user wants more detail).
        Otherwise, use the category-specific limits (user wants defaults or less).
        """
        default_limit = 50  # Default value
        category_limit = self.category_limits.get(category.value, default_limit)

        # If user explicitly set a higher limit, use it
        if self.max_tokens_per_category > default_limit:
            return self.max_tokens_per_category
        # If user set a lower limit, cap all categories to it
        elif self.max_tokens_per_category < default_limit:
            return min(category_limit, self.max_tokens_per_category)
        # Default: use category-specific limits
        return category_limit


@dataclass
class CategorySection:
    """A section of the assembled prompt for one category."""
    category: CanonicalCategory
    tokens: List[str]  # The token texts included
    text: str  # The assembled text for this section
    confidence: float  # Average confidence of tokens


@dataclass
class AssembledPrompt:
    """Result from prompt assembly."""
    prompt: str
    style: PromptStyle
    sections: List[CategorySection] = field(default_factory=list)
    word_count: int = 0
    token_count: int = 0
    categories_used: List[str] = field(default_factory=list)
    metadata: Dict[str, Any] = field(default_factory=dict)

    def __post_init__(self):
        self.word_count = len(self.prompt.split())
        self.token_count = sum(len(s.tokens) for s in self.sections)
        self.categories_used = [s.category.value for s in self.sections if s.tokens]

    def get_section(self, category: CanonicalCategory) -> Optional[CategorySection]:
        """Get section for a specific category."""
        for section in self.sections:
            if section.category == category:
                return section
        return None


class BaseAssembler(ABC):
    """Base class for prompt assemblers."""

    def __init__(self, config: Optional[AssemblerConfig] = None):
        self.config = config or AssemblerConfig()

    @abstractmethod
    def assemble(self, classified: ClassifiedImage) -> AssembledPrompt:
        """Assemble classified tokens into a prompt.

        Args:
            classified: ClassifiedImage from the classifier

        Returns:
            AssembledPrompt with the generated prompt
        """
        pass

    def _get_tokens_for_category(
        self,
        classified: ClassifiedImage,
        category: CanonicalCategory,
        max_tokens: Optional[int] = None
    ) -> List[TokenClassification]:
        """Get tokens for a category, sorted by confidence.

        Args:
            classified: ClassifiedImage to get tokens from
            category: Category to get tokens for
            max_tokens: Maximum tokens to return (None = use category-specific limit)

        Returns:
            List of TokenClassification sorted by confidence
        """
        tokens = classified.get_category(category)
        if not tokens:
            return []

        # Filter by minimum confidence
        tokens = [t for t in tokens if t.token.confidence >= self.config.min_confidence]

        # For subject_details, prioritize SHORT tags over LONG sentences
        # Short tags (< 30 chars) are more useful for image generation
        if category == CanonicalCategory.SUBJECT_DETAILS:
            # Split into short tags and long sentences
            short_tags = [t for t in tokens if len(t.token.text) < 30]
            long_sentences = [t for t in tokens if len(t.token.text) >= 30]

            # Sort each group by confidence
            short_tags.sort(key=lambda t: t.token.confidence, reverse=True)
            long_sentences.sort(key=lambda t: t.token.confidence, reverse=True)

            # Combine: short tags first (up to 80% of limit), then long sentences
            limit = max_tokens if max_tokens else self.config.get_limit_for_category(category)
            short_limit = int(limit * 0.8)
            # Calculate long_limit to ensure total doesn't exceed limit
            actual_short = min(len(short_tags), short_limit)
            long_limit = limit - actual_short

            tokens = short_tags[:actual_short] + long_sentences[:long_limit]
            # Final safety check - truncate to exact limit
            tokens = tokens[:limit]
        else:
            # Sort by confidence (highest first = most important)
            tokens.sort(key=lambda t: t.token.confidence, reverse=True)

        # Use category-specific limit, falling back to provided max_tokens or default
        if max_tokens is not None:
            limit = max_tokens
        else:
            limit = self.config.get_limit_for_category(category)

        return tokens[:limit]

    def _deduplicate_tokens(
        self,
        tokens: List[TokenClassification]
    ) -> List[TokenClassification]:
        """Remove duplicate token texts."""
        seen = set()
        unique = []
        for t in tokens:
            text_lower = t.token.text.lower().strip()
            if text_lower not in seen:
                seen.add(text_lower)
                unique.append(t)
        return unique

    def _join_natural(self, items: List[str]) -> str:
        """Join items in natural language style."""
        if not items:
            return ""
        if len(items) == 1:
            return items[0]
        if len(items) == 2:
            return f"{items[0]} and {items[1]}"
        return ', '.join(items[:-1]) + f", and {items[-1]}"

    def _capitalize_first(self, text: str) -> str:
        """Capitalize first letter only."""
        if not text:
            return text
        return text[0].upper() + text[1:]