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

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

    # 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
        "change_color": InpaintingTemplate(
            key="change_color",
            name="Change Color",
            category="Color",
            icon="🎨",
            description="Change color ONLY - works on any object (clothes, walls, furniture) based on mask",
            prompt_template="{content}",  # Minimal: only color word
            negative_prompt=(
                "original color, keeping same color, unchanged color, "
                "black clothing, dark colors, partial color change, "
                "texture details, fabric patterns, material description, object names, "
                "realistic lighting effects, complex shadows, photorealistic rendering, "
                "style description, clothing type, object description"
            ),
            controlnet_conditioning_scale=0.12,   # Low: allows color freedom while maintaining shape
            feather_radius=6,                     # Low: reduces blending with original color
            guidance_scale=18.0,                  # Very high: forces color transformation
            num_inference_steps=10,               # Optimized for speed
            preferred_conditioning="canny",       # Edge-based: preserves shape without locking color
            enhance_prompt=False,                 # Disabled: avoid interference from object descriptions
            difficulty="easy",
            usage_tips=[
                "🎯 Purpose: Pure color transformation (any object type)",
                "✅ Use SIMPLE color words ONLY: 'red', 'blue', 'white', 'bright yellow'",
                "❌ DON'T describe objects: avoid 'red shirt', 'blue wall', 'white car'",
                "✅ Mask defines what to change - prompt defines what color",
                "✅ Works on: clothes, walls, furniture, vehicles, any object",
                "📝 Example prompts:",
                "  • 'red'",
                "  • 'bright blue'",
                "  • 'pure white'",
                "  • 'deep black'",
                "💡 For extreme changes (black→red on anime images), this template is optimized"
            ]
        ),

        # 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
            preferred_conditioning="depth",       # Depth: preserves fabric folds and body structure
            enhance_prompt=True,                  # Enabled: enriches clothing details
            difficulty="easy",
            usage_tips=[
                "🎯 Purpose: Change clothing style, material, or design",
                "✅ Describe clothing type: 'white polo shirt', 'formal black suit', 'denim jacket'",
                "✅ Can include color: 'red polo shirt with collar'",
                "✅ Can describe details: 'casual shirt with buttons and pockets'",
                "✅ Preserves body structure and natural fit",
                "📝 Example prompts:",
                "  • 'white polo shirt with clean collar'",
                "  • 'formal business suit with tie'",
                "  • 'casual denim jacket with pockets'",
                "  • 'red t-shirt with simple design'",
                "💡 Use after 'Change Color' if you want to refine style details"
            ]
        ),

        # 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
            preferred_conditioning="canny",       # Edge-based: preserves scene perspective
            enhance_prompt=True,                  # Enabled: enriches object details
            difficulty="medium",
            usage_tips=[
                "🎯 Purpose: Replace objects with new ones",
                "⚠️ IMPORTANT: One object type at a time only!",
                "⚠️ All masked areas become the SAME object",
                "✅ Describe target object: 'ceramic vase with flowers', 'wooden chair'",
                "✅ Can add details: 'modern glass table with metal legs'",
                "📝 Example workflow:",
                "  • Step 1: Mask position A → prompt 'red vase' → Generate",
                "  • Step 2: Mask position B → prompt 'wooden chair' → Generate",
                "📝 Example prompts:",
                "  • 'white ceramic vase with fresh flowers'",
                "  • 'modern wooden chair with cushion'",
                "  • 'elegant glass bottle with cork'",
                "💡 For multiple object types, do it in batches (one type per generation)"
            ]
        ),

        # 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
            preferred_conditioning="depth",       # Depth: preserves spatial perspective
            enhance_prompt=False,                 # Disabled: avoid generating new objects
            difficulty="medium",
            usage_tips=[
                "🎯 Purpose: Remove objects and fill with background",
                "✅ Describe the BACKGROUND material: 'wooden floor', 'white wall', 'grass field'",
                "❌ DON'T describe the object to remove",
                "❌ Leave prompt EMPTY if you want to match surrounding automatically",
                "✅ Mask should completely cover the object to remove",
                "📝 Example prompts:",
                "  • 'wooden floor texture'",
                "  • 'white painted wall'",
                "  • 'green grass field'",
                "  • 'concrete pavement'",
                "  • '' (empty - auto match surrounding)",
                "💡 Works best with simple, uniform 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,
            "preferred_conditioning": template.preferred_conditioning
        }

    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)