inpainting_templates.py

import logging
from dataclasses import dataclass, field
from typing import Dict, List, Optional

logger = logging.getLogger(__name__)


@dataclass
class InpaintingTemplate:
    """Data class representing an inpainting template."""

    key: str
    name: str
    category: str
    icon: str
    description: str

    # Prompt templates
    prompt_template: str
    negative_prompt: str

    # Recommended parameters
    controlnet_conditioning_scale: float = 0.7
    feather_radius: int = 8
    guidance_scale: float = 7.5
    num_inference_steps: int = 25

    # Inpainting strength (0.0-1.0)
    # 1.0 = fully repaint masked area, 0.0 = keep original
    strength: float = 1.0

    # Conditioning type preference
    preferred_conditioning: str = "canny"  # "canny" or "depth"

    # Structure preservation in masked area
    # True = keep edges in mask (for color change), False = clear edges (for replacement/removal)
    preserve_structure_in_mask: bool = False

    # Prompt enhancement control
    enhance_prompt: bool = True  # Whether to use OpenCLIP prompt enhancement

    # Difficulty level for UI display
    difficulty: str = "medium"  # "easy", "medium", "advanced"

    # Tips for users
    usage_tips: List[str] = field(default_factory=list)


class InpaintingTemplateManager:
    """
    Manages inpainting templates for various use cases.

    Provides categorized presets optimized for different inpainting scenarios
    including object replacement, removal, style transfer, and enhancement.

    Attributes:
        TEMPLATES: Dictionary of all available templates
        CATEGORIES: List of category names in display order

    Example:
        >>> manager = InpaintingTemplateManager()
        >>> template = manager.get_template("object_replacement")
        >>> print(template.prompt_template)
    """

    TEMPLATES: Dict[str, InpaintingTemplate] = {
        # ========================================
        # 4 CORE TEMPLATES - Optimized for Speed & Quality
        # ========================================

        # 1. CHANGE COLOR - Pure color transformation (preserves structure)
        "change_color": InpaintingTemplate(
            key="change_color",
            name="Change Color",
            category="Color",
            icon="🎨",
            description="Change color ONLY - preserves original structure, only changes the fill color",
            prompt_template="{content} colored, solid {content} color, flat color, same structure",
            negative_prompt=(
                "original color, keeping same color, unchanged color, "
                "black, dark, keeping black, maintaining black color, "
                "black clothing, dark colors, dark fabric, black fabric, "
                "patterns, floral, stripes, plaid, checkered, decorative patterns, "
                "texture change, material change, fabric change, "
                "new design, different style, style change, "
                "complex patterns, printed patterns, embroidery"
            ),
            controlnet_conditioning_scale=0.5,    # Medium: balance structure preservation and color freedom
            feather_radius=4,                     # Low: clean color boundaries
            guidance_scale=13.0,                  # High: follow color prompt strongly
            num_inference_steps=10,               # Optimized for speed
            strength=1.0,                         # Full repaint for color change
            preferred_conditioning="canny",       # Edge-based: preserves shape
            preserve_structure_in_mask=True,      # KEY: keep edges to preserve clothing structure
            enhance_prompt=False,                 # Disabled: use color prompt directly
            difficulty="easy",
            usage_tips=[
                "🎯 Purpose: Change the color while keeping the original shape and structure.",
                "",
                "📝 Example Prompts:",
                "  • 'vibrant red' - for a bold, saturated red color",
                "  • 'soft pastel pink' - for a gentle, light pink tone",
                "  • 'deep navy blue' - for a rich, dark blue shade",
                "  • 'bright yellow' - for an eye-catching yellow",
                "  • 'pure white' - for a clean, solid white",
                "",
                "💡 Tips:",
                "  • Describe ONLY the color, not the object itself",
                "  • Paint the entire area you want to recolor",
                "  • Add modifiers like 'bright', 'dark', 'pastel' for better results"
            ]
        ),

        # 2. CLOTHING CHANGE - Style and garment transformation
        "clothing_change": InpaintingTemplate(
            key="clothing_change",
            name="Clothing Change",
            category="Replacement",
            icon="👕",
            description="Change clothing style, material, or design - can include color change",
            prompt_template="{content} with proper fit and natural appearance",
            negative_prompt=(
                "wrong body proportions, floating fabric, unrealistic wrinkles, "
                "mismatched lighting, visible edges, original clothing style, "
                "keeping same color, original color, faded colors, unchanged appearance, partial change, "
                "black clothing, dark original color, distorted body, naked, nudity"
            ),
            controlnet_conditioning_scale=0.30,   # Medium: preserves body structure, allows clothing change
            feather_radius=14,                    # Medium: natural blending with body
            guidance_scale=11.5,                  # Medium-high: accurate clothing generation
            num_inference_steps=10,               # Optimized for speed
            strength=1.0,                         # Full repaint: completely replace clothing
            preferred_conditioning="depth",       # Depth: preserves fabric folds and body structure
            enhance_prompt=True,                  # Enabled: enriches clothing details
            difficulty="easy",
            usage_tips=[
                "🎯 Purpose: Replace clothing with a completely different style or design.",
                "",
                "📝 Example Prompts:",
                "  • 'white polo shirt with collar' - casual business look",
                "  • 'formal black suit with tie' - professional attire",
                "  • 'cozy knit sweater in cream color' - warm casual style",
                "  • 'vintage denim jacket with patches' - retro fashion",
                "  • 'elegant silk blouse in emerald green' - sophisticated look",
                "",
                "💡 Tips:",
                "  • Include both clothing type AND color for best results",
                "  • Add details like 'with buttons', 'v-neck', 'long sleeves' for specificity",
                "  • Body structure is preserved automatically"
            ]
        ),

        # 3. OBJECT REPLACEMENT - Replace one object with another
        "object_replacement": InpaintingTemplate(
            key="object_replacement",
            name="Object Replacement",
            category="Replacement",
            icon="🔄",
            description="Replace objects (one type at a time) - all masked areas become the SAME object",
            prompt_template="{content} in natural lighting, fitting the scene",
            negative_prompt=(
                "inconsistent lighting, wrong perspective, mismatched colors, "
                "visible seams, floating objects, unrealistic placement, original object, "
                "poorly integrated, disconnected from scene, keeping original, remnants of original, "
                "multiple different objects, mixed objects, various items"
            ),
            controlnet_conditioning_scale=0.25,   # Low-medium: allows complete object replacement
            feather_radius=10,                    # Medium: natural scene integration
            guidance_scale=13.0,                  # Medium-high: accurate object generation
            num_inference_steps=10,               # Optimized for speed
            strength=1.0,                         # Full repaint: completely replace object
            preferred_conditioning="canny",       # Edge-based: preserves scene perspective
            enhance_prompt=True,                  # Enabled: enriches object details
            difficulty="medium",
            usage_tips=[
                "🎯 Purpose: Replace an object in the scene with something completely different.",
                "",
                "📝 Example Prompts:",
                "  • 'ceramic vase with colorful flowers' - decorative item",
                "  • 'modern laptop on a stand' - tech gadget",
                "  • 'stack of vintage books' - classic decoration",
                "  • 'potted succulent plant' - natural element",
                "  • 'antique brass lamp' - lighting fixture",
                "",
                "💡 Tips:",
                "  • Replace ONE object type at a time for best results",
                "  • Describe what you want, not what you're removing",
                "  • Include material and style details for realism"
            ]
        ),

        # 4. REMOVAL - Remove objects and fill with background
        "removal": InpaintingTemplate(
            key="removal",
            name="Remove Object",
            category="Removal",
            icon="🗑️",
            description="Remove objects and naturally fill with background - describe the background material",
            prompt_template="continue the background with {content}, seamless blending, natural continuation",
            negative_prompt=(
                "new object appearing, adding items, inserting objects, "
                "foreground elements, visible object, thing, item, "
                "unnatural filling, visible patches, inconsistent texture, "
                "mismatched pattern, color discontinuity, artificial blending"
            ),
            controlnet_conditioning_scale=0.20,   # Low: allows creative background filling
            feather_radius=12,                    # Medium: smooth background blending
            guidance_scale=12.0,                  # Medium: balanced control and naturalness
            num_inference_steps=10,               # Optimized for speed
            strength=1.0,                         # Full repaint: completely remove and fill
            preferred_conditioning="depth",       # Depth: preserves spatial perspective
            enhance_prompt=False,                 # Disabled: avoid generating new objects
            difficulty="medium",
            usage_tips=[
                "🎯 Purpose: Remove unwanted objects and fill the area with background.",
                "",
                "📝 Example Prompts:",
                "  • 'smooth wooden floor texture' - for indoor floors",
                "  • 'clean white painted wall' - for wall backgrounds",
                "  • 'natural green grass' - for outdoor lawn areas",
                "  • 'beige carpet texture' - for carpeted floors",
                "  • 'clear blue sky with clouds' - for sky backgrounds",
                "",
                "💡 Tips:",
                "  • Describe the BACKGROUND, not the object being removed",
                "  • Leave empty to auto-match surrounding textures",
                "  • Works best with uniform, simple backgrounds"
            ]
        ),
    }


    # Category display order
    CATEGORIES = ["Color", "Replacement", "Removal"]  # 4 core templates only

    def __init__(self):
        """Initialize the InpaintingTemplateManager."""
        logger.info(f"InpaintingTemplateManager initialized with {len(self.TEMPLATES)} templates")

    def get_all_templates(self) -> Dict[str, InpaintingTemplate]:
        """
        Get all available templates.

        Returns
        -------
        dict
            Dictionary of all templates keyed by template key
        """
        return self.TEMPLATES

    def get_template(self, key: str) -> Optional[InpaintingTemplate]:
        """
        Get a specific template by key.

        Parameters
        ----------
        key : str
            Template identifier

        Returns
        -------
        InpaintingTemplate or None
            Template if found, None otherwise
        """
        return self.TEMPLATES.get(key)

    def get_templates_by_category(self, category: str) -> List[InpaintingTemplate]:
        """
        Get all templates in a specific category.

        Parameters
        ----------
        category : str
            Category name

        Returns
        -------
        list
            List of templates in the category
        """
        return [t for t in self.TEMPLATES.values() if t.category == category]

    def get_categories(self) -> List[str]:
        """
        Get list of all categories in display order.

        Returns
        -------
        list
            Category names
        """
        return self.CATEGORIES

    def get_template_choices_sorted(self) -> List[str]:
        """
        Get template choices formatted for Gradio dropdown.

        Returns list of display strings sorted by category then A-Z.
        Format: "icon Name"

        Returns
        -------
        list
            Formatted display strings for dropdown
        """
        display_list = []

        for category in self.CATEGORIES:
            templates = self.get_templates_by_category(category)
            for template in sorted(templates, key=lambda t: t.name):
                display_name = f"{template.icon} {template.name}"
                display_list.append(display_name)

        return display_list

    def get_template_key_from_display(self, display_name: str) -> Optional[str]:
        """
        Get template key from display name.

        Parameters
        ----------
        display_name : str
            Display string like "🔄 Object Replacement"

        Returns
        -------
        str or None
            Template key if found
        """
        if not display_name:
            return None

        for key, template in self.TEMPLATES.items():
            if f"{template.icon} {template.name}" == display_name:
                return key
        return None

    def get_parameters_for_template(self, key: str) -> Dict[str, any]:
        """
        Get recommended parameters for a template.

        Parameters
        ----------
        key : str
            Template key

        Returns
        -------
        dict
            Dictionary of parameter names and values
        """
        template = self.get_template(key)
        if not template:
            return {}

        return {
            "controlnet_conditioning_scale": template.controlnet_conditioning_scale,
            "feather_radius": template.feather_radius,
            "guidance_scale": template.guidance_scale,
            "num_inference_steps": template.num_inference_steps,
            "strength": template.strength,
            "preferred_conditioning": template.preferred_conditioning,
            "preserve_structure_in_mask": template.preserve_structure_in_mask,
            "enhance_prompt": template.enhance_prompt
        }

    def build_prompt(self, key: str, content: str) -> str:
        """
        Build complete prompt from template and user content.

        Parameters
        ----------
        key : str
            Template key
        content : str
            User-provided content description

        Returns
        -------
        str
            Formatted prompt with content inserted
        """
        template = self.get_template(key)
        if not template:
            return content

        return template.prompt_template.format(content=content)

    def get_negative_prompt(self, key: str) -> str:
        """
        Get negative prompt for a template.

        Parameters
        ----------
        key : str
            Template key

        Returns
        -------
        str
            Negative prompt string
        """
        template = self.get_template(key)
        if not template:
            return ""
        return template.negative_prompt

    def get_usage_tips(self, key: str) -> List[str]:
        """
        Get usage tips for a template.

        Parameters
        ----------
        key : str
            Template key

        Returns
        -------
        list
            List of tip strings
        """
        template = self.get_template(key)
        if not template:
            return []
        return template.usage_tips

    def build_gallery_html(self) -> str:
        """
        Build HTML for template gallery display.

        Returns
        -------
        str
            HTML string for Gradio display
        """
        html_parts = ['<div class="inpainting-gallery">']

        for category in self.CATEGORIES:
            templates = self.get_templates_by_category(category)
            if not templates:
                continue

            html_parts.append(f'''
            <div class="inpainting-category">
                <h4 class="inpainting-category-title">{category}</h4>
                <div class="inpainting-grid">
            ''')

            for template in sorted(templates, key=lambda t: t.name):
                html_parts.append(f'''
                <div class="inpainting-card" data-template="{template.key}">
                    <span class="inpainting-icon">{template.icon}</span>
                    <span class="inpainting-name">{template.name}</span>
                    <span class="inpainting-desc">{template.description[:50]}...</span>
                </div>
                ''')

            html_parts.append('</div></div>')

        html_parts.append('</div>')
        return ''.join(html_parts)