inpainting_templates.py

import logging
from dataclasses import dataclass, field
from typing import Any, 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

    # Pipeline mode selection
    use_controlnet: bool = True  # False = use pure SDXL Inpainting model (more stable)
    mask_dilation: int = 0       # Pixels to expand mask for better edge blending

    # ControlNet parameters (only used when use_controlnet=True)
    controlnet_conditioning_scale: float = 0.7
    preferred_conditioning: str = "canny"  # "canny" or "depth"
    preserve_structure_in_mask: bool = False
    edge_guidance_mode: str = "boundary"

    # Generation parameters
    guidance_scale: float = 7.5
    num_inference_steps: int = 25
    strength: float = 0.99  # Use 0.99 instead of 1.0 to avoid noise issues

    # Mask parameters
    feather_radius: int = 3  # Minimal feathering, let pipeline handle blending

    # Prompt enhancement
    enhance_prompt: bool = True

    # UI metadata
    difficulty: str = "medium"
    recommended_models: List[str] = field(default_factory=lambda: ["sdxl_base"])
    example_prompts: List[str] = field(default_factory=list)
    usage_tips: List[str] = field(default_factory=list)


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

    Templates are categorized into two pipeline modes:
    - Pure Inpainting (use_controlnet=False): For replacement/removal tasks
    - ControlNet Inpainting (use_controlnet=True): For structure-preserving tasks

    Example:
        >>> manager = InpaintingTemplateManager()
        >>> template = manager.get_template("object_replacement")
        >>> if not template.use_controlnet:
        ...     # Use pure SDXL Inpainting pipeline
        ...     pass
    """

    TEMPLATES: Dict[str, InpaintingTemplate] = {
        # 1. OBJECT REPLACEMENT - Replace one object with another
        "object_replacement": InpaintingTemplate(
            key="object_replacement",
            name="Object Replacement",
            category="Replacement",
            icon="🔄",
            description="Replace objects naturally - uses dedicated inpainting model for best results",
            prompt_template="{content}, photorealistic, natural lighting, seamlessly integrated, high quality, detailed",
            negative_prompt=(
                "blurry, low quality, distorted, deformed, "
                "visible seams, harsh edges, unnatural, "
                "cartoon, anime, illustration, drawing"
            ),
            # Pipeline mode
            use_controlnet=False,  # Pure inpainting for stable results
            mask_dilation=5,       # Expand mask for seamless blending

            # Generation parameters
            guidance_scale=8.0,
            num_inference_steps=25,
            strength=0.99,

            # Mask parameters
            feather_radius=3,

            # Not used for Pure Inpainting but kept for compatibility
            controlnet_conditioning_scale=0.0,
            preferred_conditioning="canny",  # Placeholder, not used in Pure Inpainting mode
            preserve_structure_in_mask=False,
            edge_guidance_mode="none",

            enhance_prompt=True,
            difficulty="easy",
            recommended_models=["realvis_xl", "juggernaut_xl"],
            example_prompts=[
                "elegant ceramic vase with fresh roses",
                "modern minimalist desk lamp, chrome finish",
                "vintage leather-bound book with gold lettering"
            ],
            usage_tips=[
                "🎯 Purpose: Replace an object with something completely different.",
                "",
                "💡 Example Prompts:",
                "  • elegant ceramic vase with fresh roses",
                "  • modern minimalist desk lamp, chrome finish",
                "  • vintage leather-bound book with gold lettering",
                "",
                "💡 Tips:",
                "  • Draw mask slightly larger than the object",
                "  • Describe the NEW object in detail",
                "  • Include material, color, style for better results"
            ]
        ),

        # 2. OBJECT REMOVAL - Remove and fill with background (NO PROMPT NEEDED)
        "removal": InpaintingTemplate(
            key="removal",
            name="Remove Object",
            category="Removal",
            icon="🗑️",
            description="Remove unwanted objects - just draw mask, no prompt needed",
            prompt_template="seamless background, natural texture continuation, photorealistic, high quality",
            negative_prompt=(
                "object, item, thing, foreground element, new object, "
                "visible patch, inconsistent texture, "
                "blurry, low quality, artificial"
            ),
            # Pipeline mode
            use_controlnet=False,  # Pure inpainting for clean removal
            mask_dilation=8,       # Larger expansion to cover shadows/reflections

            # Generation parameters
            guidance_scale=7.0,    # Lower guidance for natural fill
            num_inference_steps=20,
            strength=0.99,

            # Mask parameters
            feather_radius=5,      # More feathering for seamless blend

            # Not used for Pure Inpainting but kept for compatibility
            controlnet_conditioning_scale=0.0,
            preferred_conditioning="canny",
            preserve_structure_in_mask=False,
            edge_guidance_mode="none",

            enhance_prompt=False,  # Do NOT enhance - keep it simple
            difficulty="easy",
            recommended_models=["realvis_xl", "juggernaut_xl"],
            example_prompts=[],    # No prompts needed for removal
            usage_tips=[
                "🎯 Purpose: Remove unwanted objects from image.",
                "",
                "📝 No prompt needed! Just:",
                "  1. Draw white mask over the object",
                "  2. Include shadows in your mask",
                "  3. Click Generate",
                "",
                "💡 Tips:",
                "  • Make mask larger than the object",
                "  • If artifacts remain, draw a bigger mask and retry"
            ]
        ),

        # CONTROLNET TEMPLATES (Structure Preserving)
        # 3. CLOTHING CHANGE - Change clothes while keeping body
        "clothing_change": InpaintingTemplate(
            key="clothing_change",
            name="Clothing Change",
            category="Replacement",
            icon="👕",
            description="Change clothing style while preserving body structure",
            prompt_template="{content}, photorealistic, realistic fabric, natural fit, high quality",
            negative_prompt=(
                "wrong proportions, distorted body, floating fabric, "
                "mismatched lighting, naked, nudity, "
                "cartoon, anime, illustration"
            ),
            # Pipeline mode
            use_controlnet=True,   # Need ControlNet to preserve body
            mask_dilation=3,       # Small expansion for clothing edges

            # ControlNet parameters
            controlnet_conditioning_scale=0.4,
            preferred_conditioning="depth",  # Depth preserves body structure
            preserve_structure_in_mask=False,
            edge_guidance_mode="soft",

            # Generation parameters
            guidance_scale=8.0,
            num_inference_steps=25,
            strength=1.0,  # Full repaint for clothing

            # Mask parameters
            feather_radius=5,

            enhance_prompt=True,
            difficulty="medium",
            recommended_models=["juggernaut_xl", "realvis_xl"],
            example_prompts=[
                "tailored charcoal suit with silk tie",
                "navy blazer with gold buttons",
                "elegant black evening dress",
                "casual white t-shirt",
                "cozy cream sweater",
                "leather motorcycle jacket",
                "formal white dress shirt",
                "vintage denim jacket",
                "red cocktail dress",
                "professional grey blazer"
            ],
            usage_tips=[
                "🎯 Purpose: Change clothing while keeping body shape.",
                "",
                "🤖 Recommended Models:",
                "  • JuggernautXL - Best for formal wear",
                "  • RealVisXL - Great for casual clothing",
                "",
                "💡 Tips:",
                "  • Mask only the clothing area",
                "  • Include fabric type: 'silk', 'cotton', 'wool'",
                "  • Body proportions are preserved automatically"
            ]
        ),

        # 4. COLOR CHANGE - Change color only, keep structure
        "change_color": InpaintingTemplate(
            key="change_color",
            name="Change Color",
            category="Color",
            icon="🎨",
            description="Change color only - strictly preserves shape and texture",
            prompt_template="{content} color, solid uniform {content}, flat color, smooth surface",
            negative_prompt=(
                "different shape, changed structure, new pattern, "
                "texture change, deformed, distorted, "
                "gradient, multiple colors, pattern"
            ),
            # Pipeline mode
            use_controlnet=True,   # Need ControlNet to preserve exact shape
            mask_dilation=0,       # No expansion - precise color change

            # ControlNet parameters
            controlnet_conditioning_scale=0.85,  # High: strict structure preservation
            preferred_conditioning="canny",       # Canny preserves edges exactly
            preserve_structure_in_mask=True,      # Keep all edges
            edge_guidance_mode="boundary",

            # Generation parameters
            guidance_scale=12.0,   # High: force the exact color
            num_inference_steps=15,
            strength=1.0,

            # Mask parameters
            feather_radius=2,      # Very small

            enhance_prompt=False,  # Use color prompt directly
            difficulty="easy",
            recommended_models=["juggernaut_xl", "realvis_xl"],
            example_prompts=[
                "vibrant red",
                "deep navy blue",
                "bright yellow",
                "emerald green",
                "soft pink",
                "pure white",
                "charcoal grey",
                "royal purple",
                "coral orange",
                "golden brown"
            ],
            usage_tips=[
                "🎯 Purpose: Change color only, shape stays exactly the same.",
                "",
                "💡 Tips:",
                "  • Enter ONLY the color name",
                "  • Use modifiers: 'bright', 'dark', 'pastel'",
                "  • Shape and texture are preserved exactly"
            ]
        ),
    }

    # Category display order
    CATEGORIES = ["Color", "Replacement", "Removal"]

    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."""
        return self.TEMPLATES

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

    def get_templates_by_category(self, category: str) -> List[InpaintingTemplate]:
        """Get all templates in a specific 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."""
        return self.CATEGORIES

    def get_template_choices_sorted(self) -> List[str]:
        """Get template choices formatted for Gradio 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."""
        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."""
        template = self.get_template(key)
        if not template:
            return {}

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

    def build_prompt(self, key: str, content: str) -> str:
        """Build complete prompt from template and user content."""
        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."""
        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."""
        template = self.get_template(key)
        if not template:
            return []
        return template.usage_tips

    def get_recommended_models(self, key: str) -> List[str]:
        """Get recommended models for a template."""
        template = self.get_template(key)
        if not template:
            return ["sdxl_base"]
        return template.recommended_models

    def get_example_prompts(self, key: str) -> List[str]:
        """Get example prompts for a template."""
        template = self.get_template(key)
        if not template:
            return []
        return template.example_prompts

    def get_primary_recommended_model(self, key: str) -> str:
        """Get the primary recommended model for a template."""
        models = self.get_recommended_models(key)
        return models[0] if models else "sdxl_base"