ui_manager.py

import logging
import time
import traceback
from pathlib import Path
from typing import Optional, Tuple, Dict, Any, List

import cv2
import gradio as gr
import numpy as np
from PIL import Image

from css_styles import CSSStyles
from scene_templates import SceneTemplateManager
from inpainting_templates import InpaintingTemplateManager
from scene_weaver_core import SceneWeaverCore
from gpu_handlers import GPUHandlers

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(name)s] %(levelname)s: %(message)s',
    datefmt='%H:%M:%S'
)


class UIManager:
    """
    Gradio UI Manager with support for background generation and inpainting.

    Provides a professional interface with mode switching, template selection,
    and advanced parameter controls. GPU operations are delegated to GPUHandlers.

    Attributes:
        gpu_handlers: GPUHandlers instance for GPU operations
        template_manager: Scene template manager
        inpainting_template_manager: Inpainting template manager
    """

    def __init__(self):
        self.sceneweaver = SceneWeaverCore()
        self.gpu_handlers = GPUHandlers(
            core=self.sceneweaver,
            inpainting_template_manager=InpaintingTemplateManager()
        )
        self.template_manager = SceneTemplateManager()
        self.inpainting_template_manager = InpaintingTemplateManager()
        self.generation_history = []
        self.inpainting_history = []
        self._preview_sensitivity = 0.5
        self._current_mode = "background"  # "background" or "inpainting"

    def apply_template(self, display_name: str, current_negative: str) -> Tuple[str, str, float]:
        """
        Apply a scene template to the prompt fields.

        Args:
            display_name: The display name from dropdown (e.g., "🏢 Modern Office")
            current_negative: Current negative prompt value

        Returns:
            Tuple of (prompt, negative_prompt, guidance_scale)
        """
        if not display_name:
            return "", current_negative, 7.5

        # Convert display name to template key
        template_key = self.template_manager.get_template_key_from_display(display_name)
        if not template_key:
            return "", current_negative, 7.5

        template = self.template_manager.get_template(template_key)
        if template:
            prompt = template.prompt
            negative = self.template_manager.get_negative_prompt_for_template(
                template_key, current_negative
            )
            guidance = template.guidance_scale
            return prompt, negative, guidance

        return "", current_negative, 7.5

    def quick_preview(
        self,
        uploaded_image: Optional[Image.Image],
        sensitivity: float = 0.5
    ) -> Optional[Image.Image]:
        """
        Generate quick foreground preview using lightweight traditional methods.

        Args:
            uploaded_image: Uploaded PIL Image
            sensitivity: Detection sensitivity (0.0 - 1.0)

        Returns:
            Preview image with colored overlay or None
        """
        if uploaded_image is None:
            return None

        try:
            logger.info(f"Generating quick preview (sensitivity={sensitivity:.2f})")

            img_array = np.array(uploaded_image.convert('RGB'))
            height, width = img_array.shape[:2]

            max_preview_size = 512
            if max(width, height) > max_preview_size:
                scale = max_preview_size / max(width, height)
                new_w = int(width * scale)
                new_h = int(height * scale)
                img_array = cv2.resize(img_array, (new_w, new_h), interpolation=cv2.INTER_AREA)
                height, width = new_h, new_w

            gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
            blurred = cv2.GaussianBlur(gray, (5, 5), 0)

            low_threshold = int(30 + (1 - sensitivity) * 50)
            high_threshold = int(100 + (1 - sensitivity) * 100)
            edges = cv2.Canny(blurred, low_threshold, high_threshold)

            kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7))
            dilated = cv2.dilate(edges, kernel, iterations=2)

            contours, _ = cv2.findContours(dilated, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

            mask = np.zeros((height, width), dtype=np.uint8)

            if contours:
                sorted_contours = sorted(contours, key=cv2.contourArea, reverse=True)
                min_area = (width * height) * 0.01 * (1 - sensitivity)
                for contour in sorted_contours:
                    if cv2.contourArea(contour) > min_area:
                        cv2.fillPoly(mask, [contour], 255)

            kernel_close = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (11, 11))
            mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel_close)

            overlay = img_array.copy().astype(np.float32)

            fg_mask = mask > 127
            overlay[fg_mask] = overlay[fg_mask] * 0.5 + np.array([0, 255, 0]) * 0.5

            bg_mask = mask <= 127
            overlay[bg_mask] = overlay[bg_mask] * 0.5 + np.array([255, 0, 0]) * 0.5

            overlay = np.clip(overlay, 0, 255).astype(np.uint8)

            original_size = uploaded_image.size
            preview_image = Image.fromarray(overlay)
            if preview_image.size != original_size:
                preview_image = preview_image.resize(original_size, Image.LANCZOS)

            logger.info("Quick preview generated successfully")
            return preview_image

        except Exception as e:
            logger.error(f"Quick preview failed: {e}")
            return None

    def _save_result(self, combined_image: Image.Image, prompt: str):
        """Save result with memory-conscious history management"""
        if not combined_image:
            return

        output_dir = Path("outputs")
        output_dir.mkdir(exist_ok=True)

        combined_image.save(output_dir / "latest_combined.png")

        self.generation_history.append({
            "prompt": prompt,
            "timestamp": time.time()
        })

        max_history = self.sceneweaver.max_history
        if len(self.generation_history) > max_history:
            self.generation_history = self.generation_history[-max_history:]

    def generate_handler(
        self,
        uploaded_image: Optional[Image.Image],
        prompt: str,
        combination_mode: str,
        focus_mode: str,
        negative_prompt: str,
        steps: int,
        guidance: float,
        progress=gr.Progress()
    ):
        """
        Generation handler - delegates GPU work to GPUHandlers.

        Parameters
        ----------
        uploaded_image : PIL.Image
            Input image
        prompt : str
            Background description
        combination_mode : str
            Composition mode
        focus_mode : str
            Focus mode
        negative_prompt : str
            Negative prompt
        steps : int
            Inference steps
        guidance : float
            Guidance scale
        progress : gr.Progress
            Progress callback

        Returns
        -------
        tuple
            (combined, generated, original, status, download_btn_update)
        """
        if uploaded_image is None:
            return None, None, None, "Please upload an image to get started!", gr.update(visible=False)

        if not prompt.strip():
            return None, None, None, "Please describe the background scene you'd like!", gr.update(visible=False)

        try:
            # Delegate to GPUHandlers
            # Note: progress_callback removed for ZeroGPU compatibility (can't pickle local functions)
            result = self.gpu_handlers.background_generate(
                image=uploaded_image,
                prompt=prompt,
                negative_prompt=negative_prompt,
                composition_mode=combination_mode,
                focus_mode=focus_mode,
                num_steps=int(steps),
                guidance_scale=float(guidance),
                progress_callback=None  # ZeroGPU doesn't support local function callbacks
            )

            if result["success"]:
                combined = result["combined_image"]
                generated = result["generated_scene"]
                original = result["original_image"]

                self._save_result(combined, prompt)

                status_msg = "Image created successfully!"

                return combined, generated, original, status_msg, gr.update(visible=True)
            else:
                error_msg = result.get("error", "Something went wrong")
                return None, None, None, f"Error: {error_msg}", gr.update(visible=False)

        except Exception as e:
            error_traceback = traceback.format_exc()
            logger.error(f"Generation handler error: {str(e)}")
            logger.error(f"Traceback:\n{error_traceback}")
            return None, None, None, f"Error: {str(e)}", gr.update(visible=False)

    def create_interface(self):
        """Create professional user interface"""

        self._css = CSSStyles.get_main_css()

        # Check Gradio version for API compatibility
        self._gradio_version = gr.__version__
        self._gradio_major = int(self._gradio_version.split('.')[0])

        # Compatible with Gradio 4.44.0+
        # Use minimal constructor arguments for maximum compatibility
        with gr.Blocks() as interface:

            # Inject CSS (compatible with all Gradio versions)
            gr.HTML(f"<style>{self._css}</style>")

            # Header
            gr.HTML("""
            <div class="main-header">
                <h1 class="main-title">
                    <span class="title-emoji">🎨</span>
                    SceneWeaver
                </h1>
                <p class="main-subtitle">AI-powered background generation and inpainting with professional edge processing</p>
            </div>
            """)

            # Main Tabs for Mode Selection
            with gr.Tabs(elem_id="main-mode-tabs") as main_tabs:

                # Background Generation Tab
                with gr.Tab("Background Generation", elem_id="bg-gen-tab"):

                    with gr.Row():
                        # Left Column - Input controls
                        with gr.Column(scale=1, min_width=350, elem_classes=["feature-card"]):
                            gr.HTML("""
                            <div class="card-content">
                                <h3 class="card-title">
                                    <span class="section-emoji">📸</span>
                                    Upload & Generate
                                </h3>
                            </div>
                            """)

                            uploaded_image = gr.Image(
                                label="Upload Your Image",
                                type="pil",
                                height=280,
                                elem_classes=["input-field"]
                            )

                            # Scene Template Selector (without Accordion to fix dropdown positioning in Gradio 5.x)
                            template_dropdown = gr.Dropdown(
                                label="Scene Templates",
                                choices=[""] + self.template_manager.get_template_choices_sorted(),
                                value="",
                                info="24 curated scenes sorted A-Z (optional)",
                                elem_classes=["template-dropdown"]
                            )

                            prompt_input = gr.Textbox(
                                label="Background Scene Description",
                                placeholder="Select a template above or describe your own scene...",
                                lines=3,
                                elem_classes=["input-field"]
                            )

                            combination_mode = gr.Dropdown(
                                label="Composition Mode",
                                choices=["center", "left_half", "right_half", "full"],
                                value="center",
                                info="center=Smart Center | left_half=Left Half | right_half=Right Half | full=Full Image",
                                elem_classes=["input-field"]
                            )

                            focus_mode = gr.Dropdown(
                                label="Focus Mode",
                                choices=["person", "scene"],
                                value="person",
                                info="person=Tight Crop | scene=Include Surrounding Objects",
                                elem_classes=["input-field"]
                            )

                            with gr.Accordion("Advanced Options", open=False):
                                negative_prompt = gr.Textbox(
                                    label="Negative Prompt",
                                    value="blurry, low quality, distorted, people, characters",
                                    lines=2,
                                    elem_classes=["input-field"]
                                )

                                steps_slider = gr.Slider(
                                    label="Quality Steps",
                                    minimum=15,
                                    maximum=50,
                                    value=25,
                                    step=5,
                                    elem_classes=["input-field"]
                                )

                                guidance_slider = gr.Slider(
                                    label="Guidance Scale",
                                    minimum=5.0,
                                    maximum=15.0,
                                    value=7.5,
                                    step=0.5,
                                    elem_classes=["input-field"]
                                )

                            generate_btn = gr.Button(
                                "Generate Background",
                                variant="primary",
                                size="lg",
                                elem_classes=["primary-button"]
                            )

                        # Right Column - Results display
                        with gr.Column(scale=2, elem_classes=["feature-card"], elem_id="results-gallery-centered"):
                            gr.HTML("""
                            <div class="card-content">
                                <h3 class="card-title">
                                    <span class="section-emoji">🎭</span>
                                    Results Gallery
                                </h3>
                            </div>
                            """)

                            # Loading notice
                            gr.HTML("""
                            <div class="loading-notice">
                                <span class="loading-notice-icon">⏱️</span>
                                <span class="loading-notice-text">
                                    <strong>First-time users:</strong> Initial model loading takes 1-2 minutes.
                                    Subsequent generations are much faster (~30s).
                                </span>
                            </div>
                            """)

                            # Quick start guide
                            gr.HTML("""
                            <details class="user-guidance-panel">
                                <summary class="guidance-summary">
                                    <span class="emoji-enhanced">💡</span>
                                    Quick Start Guide
                                </summary>
                                <div class="guidance-content">
                                    <p><strong>Step 1:</strong> Upload any image with a clear subject</p>
                                    <p><strong>Step 2:</strong> Describe or Choose your desired background scene</p>
                                    <p><strong>Step 3:</strong> Choose composition mode (center works best)</p>
                                    <p><strong>Step 4:</strong> Click Generate and wait for the magic!</p>
                                    <p><strong>Tip:</strong> For dark clothing, ensure good lighting in original photo.</p>
                                </div>
                            </details>
                            """)

                            with gr.Tabs():
                                with gr.TabItem("Final Result"):
                                    combined_output = gr.Image(
                                        label="Your Generated Image",
                                        elem_classes=["result-gallery"],
                                        show_label=False
                                    )
                                with gr.TabItem("Background"):
                                    generated_output = gr.Image(
                                        label="Generated Background",
                                        elem_classes=["result-gallery"],
                                        show_label=False
                                    )
                                with gr.TabItem("Original"):
                                    original_output = gr.Image(
                                        label="Processed Original",
                                        elem_classes=["result-gallery"],
                                        show_label=False
                                    )

                            status_output = gr.Textbox(
                                label="Status",
                                value="Ready to create! Upload an image and describe your vision.",
                                interactive=False,
                                elem_classes=["status-panel", "status-ready"]
                            )

                            with gr.Row():
                                download_btn = gr.DownloadButton(
                                    "Download Result",
                                    value=None,
                                    visible=False,
                                    elem_classes=["secondary-button"]
                                )
                                clear_btn = gr.Button(
                                    "Clear All",
                                    elem_classes=["secondary-button"]
                                )
                                memory_btn = gr.Button(
                                    "Clean Memory",
                                    elem_classes=["secondary-button"]
                                )

                    # Event handlers for Background Generation Tab
                    # Template selection handler
                    template_dropdown.change(
                        fn=self.apply_template,
                        inputs=[template_dropdown, negative_prompt],
                        outputs=[prompt_input, negative_prompt, guidance_slider]
                    )

                    generate_btn.click(
                        fn=self.generate_handler,
                        inputs=[
                            uploaded_image,
                            prompt_input,
                            combination_mode,
                            focus_mode,
                            negative_prompt,
                            steps_slider,
                            guidance_slider
                        ],
                        outputs=[
                            combined_output,
                            generated_output,
                            original_output,
                            status_output,
                            download_btn
                        ]
                    )

                    clear_btn.click(
                        fn=lambda: (None, None, None, "Ready to create!", gr.update(visible=False)),
                        outputs=[combined_output, generated_output, original_output, status_output, download_btn]
                    )

                    memory_btn.click(
                        fn=lambda: self.sceneweaver._ultra_memory_cleanup() or "Memory cleaned!",
                        outputs=[status_output]
                    )

                    combined_output.change(
                        fn=lambda img: gr.update(value="outputs/latest_combined.png", visible=True) if (img is not None) else gr.update(visible=False),
                        inputs=[combined_output],
                        outputs=[download_btn]
                    )

                # End of Background Generation Tab

                # Inpainting Tab
                self.create_inpainting_tab()

            # Footer with tech credits (outside tabs)
            gr.HTML("""
            <div class="app-footer">
                <div class="footer-powered">
                    <p class="footer-powered-title">Powered By</p>
                    <div class="footer-tech-grid">
                        <span class="footer-tech-item">Stable Diffusion XL</span>
                        <span class="footer-tech-item">OpenCLIP</span>
                        <span class="footer-tech-item">BiRefNet</span>
                        <span class="footer-tech-item">rembg</span>
                        <span class="footer-tech-item">PyTorch</span>
                        <span class="footer-tech-item">Gradio</span>
                    </div>
                </div>
                <div class="footer-divider"></div>
                <p class="footer-copyright">
                    SceneWeaver &copy; 2025 &nbsp;|&nbsp;
                    Built with <a href="https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0" target="_blank">SDXL</a>
                    and <a href="https://github.com/mlfoundations/open_clip" target="_blank">OpenCLIP</a>
                </p>
            </div>
            """)

        return interface

    def launch(self, share: bool = True, debug: bool = False):
        """Launch the UI interface"""
        interface = self.create_interface()

        # Launch kwargs compatible with Gradio 4.44.0+
        # Keep minimal for maximum compatibility
        launch_kwargs = {
            "share": share,
            "debug": debug,
            "show_error": True,
            "quiet": False
        }

        return interface.launch(**launch_kwargs)

    # INPAINTING UI METHODS
    def apply_inpainting_template(
        self,
        display_name: str,
        current_prompt: str
    ) -> Tuple[str, float, int, str, Any, Any, Any]:
        """
        Apply an inpainting template to the UI fields.

        Parameters
        ----------
        display_name : str
            Template display name from dropdown
        current_prompt : str
            Current prompt content

        Returns
        -------
        tuple
            (prompt, conditioning_scale, feather_radius, conditioning_type,
             controlnet_settings_visibility, mode_info_html, model_selection_visibility)
        """
        # Default returns for no template selected
        default_return = (
            current_prompt,
            0.7,
            8,
            "canny",
            gr.update(visible=True),   # Show ControlNet settings by default
            "",                         # No mode info
            gr.update(visible=True)    # Show model selection by default
        )

        if not display_name:
            return default_return

        template_key = self.inpainting_template_manager.get_template_key_from_display(display_name)
        if not template_key:
            return default_return

        template = self.inpainting_template_manager.get_template(template_key)
        if template:
            params = self.inpainting_template_manager.get_parameters_for_template(template_key)
            use_controlnet = params.get('use_controlnet', True)

            # Determine visibility and info based on mode
            if use_controlnet:
                controlnet_visibility = gr.update(visible=True)
                model_visibility = gr.update(visible=True)
                mode_info = """
                <div style="background: linear-gradient(135deg, #e8f5e9 0%, #c8e6c9 100%);
                            border-left: 4px solid #4CAF50;
                            padding: 10px 14px;
                            border-radius: 8px;
                            margin: 8px 0;">
                    <p style="margin: 0; color: #2e7d32; font-size: 13px;">
                        🎛️ <strong>ControlNet Mode</strong> - Structure will be preserved using edge/depth guidance.
                        You can adjust ControlNet settings and select model below.
                    </p>
                </div>
                """
            else:
                # Pure Inpainting mode - hide both ControlNet and Model Selection
                controlnet_visibility = gr.update(visible=False)
                model_visibility = gr.update(visible=False)
                mode_info = """
                <div style="background: linear-gradient(135deg, #fff3e0 0%, #ffe0b2 100%);
                            border-left: 4px solid #ff9800;
                            padding: 10px 14px;
                            border-radius: 8px;
                            margin: 8px 0;">
                    <p style="margin: 0; color: #e65100; font-size: 13px;">
                        🚀 <strong>Pure Inpainting Mode</strong> - Using dedicated SDXL Inpainting model.<br>
                        Model and ControlNet settings are automatically configured for best results.
                    </p>
                </div>
                """

            return (
                current_prompt,
                params.get('controlnet_conditioning_scale', 0.7),
                params.get('feather_radius', 8),
                params.get('preferred_conditioning', 'canny'),
                controlnet_visibility,
                mode_info,
                model_visibility
            )

        return default_return

    def extract_mask_from_editor(self, editor_output: Dict[str, Any]) -> Optional[Image.Image]:
        """
        Extract mask from Gradio ImageEditor output.

        Handles different Gradio versions' output formats.

        Parameters
        ----------
        editor_output : dict
            Output from gr.ImageEditor component

        Returns
        -------
        PIL.Image or None
            Extracted mask as grayscale image
        """
        if editor_output is None:
            return None

        try:
            # Gradio 5.x format
            if isinstance(editor_output, dict):
                # Check for 'layers' key (Gradio 5.x ImageEditor)
                if 'layers' in editor_output and editor_output['layers']:
                    # Get the first layer as mask
                    layer = editor_output['layers'][0]
                    if isinstance(layer, np.ndarray):
                        mask_array = layer
                    elif isinstance(layer, Image.Image):
                        mask_array = np.array(layer)
                    else:
                        return None

                # Check for 'composite' key
                elif 'composite' in editor_output:
                    composite = editor_output['composite']
                    if isinstance(composite, np.ndarray):
                        mask_array = composite
                    elif isinstance(composite, Image.Image):
                        mask_array = np.array(composite)
                    else:
                        return None
                else:
                    return None

            elif isinstance(editor_output, np.ndarray):
                mask_array = editor_output
            elif isinstance(editor_output, Image.Image):
                mask_array = np.array(editor_output)
            else:
                logger.warning(f"Unexpected editor output type: {type(editor_output)}")
                return None

            # Convert to grayscale mask
            if len(mask_array.shape) == 3:
                if mask_array.shape[2] == 4:
                    # RGBA format - extract white brush strokes from RGB channels
                    # White brush strokes have high RGB values AND high alpha
                    rgb_part = mask_array[:, :, :3]
                    alpha_part = mask_array[:, :, 3]

                    # Convert RGB to grayscale to detect white areas
                    gray = cv2.cvtColor(rgb_part, cv2.COLOR_RGB2GRAY)

                    # Combine: white areas (high gray value) with opacity (high alpha)
                    # This captures white brush strokes
                    mask_gray = np.minimum(gray, alpha_part)
                elif mask_array.shape[2] == 3:
                    # RGB - convert to grayscale (white areas become white in mask)
                    mask_gray = cv2.cvtColor(mask_array, cv2.COLOR_RGB2GRAY)
                else:
                    mask_gray = mask_array[:, :, 0]
            else:
                # Already grayscale
                mask_gray = mask_array

            return Image.fromarray(mask_gray.astype(np.uint8), mode='L')

        except Exception as e:
            logger.error(f"Failed to extract mask from editor: {e}")
            return None

    def inpainting_handler(
        self,
        image: Optional[Image.Image],
        mask_editor: Dict[str, Any],
        prompt: str,
        template_dropdown: str,
        model_choice: str,
        conditioning_type: str,
        conditioning_scale: float,
        feather_radius: int,
        guidance_scale: float,
        num_steps: int,
        seed: int,
        progress: gr.Progress = gr.Progress()
    ) -> Tuple[Optional[Image.Image], Optional[Image.Image], str, int]:
        """
        Handle inpainting generation request - delegates GPU work to GPUHandlers.

        Parameters
        ----------
        image : PIL.Image
            Original image to inpaint
        mask_editor : dict
            Mask editor output
        prompt : str
            Text description of desired content
        template_dropdown : str
            Selected template (optional)
        model_choice : str
            Model key to use (juggernaut_xl, realvis_xl, sdxl_base, animagine_xl)
        conditioning_type : str
            ControlNet conditioning type
        conditioning_scale : float
            ControlNet influence strength
        feather_radius : int
            Mask feathering radius
        guidance_scale : float
            Guidance scale for generation
        num_steps : int
            Number of inference steps
        seed : int
            Random seed (-1 for random)
        progress : gr.Progress
            Progress callback

        Returns
        -------
        tuple
            (result_image, control_image, status_message, used_seed)
        """
        if image is None:
            return None, None, "⚠️ Please upload an image first", -1

        # Extract mask
        mask = self.extract_mask_from_editor(mask_editor)
        if mask is None:
            return None, None, "⚠️ Please draw a mask on the image", -1

        # Validate mask
        mask_array = np.array(mask)
        coverage = np.count_nonzero(mask_array > 127) / mask_array.size
        if coverage < 0.01:
            return None, None, "⚠️ Mask too small - please select a larger area", -1
        if coverage > 0.95:
            return None, None, "⚠️ Mask too large - consider using background generation instead", -1

        try:
            # Get template key if selected
            template_key = None
            if template_dropdown:
                template_key = self.inpainting_template_manager.get_template_key_from_display(
                    template_dropdown
                )

            # Delegate to GPUHandlers
            # Note: progress_callback removed for ZeroGPU compatibility (can't pickle local functions)
            result_image, control_image, status, used_seed = self.gpu_handlers.inpainting_generate(
                image=image,
                mask=mask,
                prompt=prompt,
                template_key=template_key,
                model_key=model_choice,
                conditioning_type=conditioning_type,
                conditioning_scale=conditioning_scale,
                feather_radius=feather_radius,
                guidance_scale=guidance_scale,
                num_steps=num_steps,
                seed=int(seed) if seed is not None else -1,
                progress_callback=None  # ZeroGPU doesn't support local function callbacks
            )

            # Store in history if successful
            if result_image is not None:
                self.inpainting_history.append({
                    'result': result_image,
                    'prompt': prompt,
                    'seed': used_seed,
                    'time': time.time()
                })
                if len(self.inpainting_history) > 3:
                    self.inpainting_history.pop(0)

            return result_image, control_image, status, used_seed

        except Exception as e:
            logger.error(f"Inpainting handler error: {e}")
            logger.error(traceback.format_exc())
            return None, None, f"❌ Error: {str(e)}", -1

    def create_inpainting_tab(self) -> gr.Tab:
        """
        Create the inpainting tab UI.

        Returns
        -------
        gr.Tab
            Configured inpainting tab component
        """
        with gr.Tab("Inpainting", elem_id="inpainting-tab") as tab:
            gr.HTML("""
            <div class="inpainting-header">
                <h3 style="display: flex; align-items: center; gap: 10px; margin-bottom: 8px;">
                    ControlNet Inpainting
                    <span style="background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
                                 color: white;
                                 padding: 3px 10px;
                                 border-radius: 12px;
                                 font-size: 0.65em;
                                 font-weight: 700;
                                 letter-spacing: 0.5px;
                                 box-shadow: 0 2px 4px rgba(102, 126, 234, 0.3);">
                        BETA
                    </span>
                </h3>
                <p style="color: #666; margin-bottom: 12px;">Draw a mask to select the area you want to regenerate</p>
            </div>
            """)

            # Model Selection Guide
            gr.HTML("""
            <div style="background: linear-gradient(135deg, #f5f7fa 0%, #e4e8ec 100%);
                        padding: 16px;
                        border-radius: 12px;
                        margin: 12px 0;
                        border: 1px solid #ddd;">
                <h4 style="margin: 0 0 12px 0; color: #333; font-size: 16px;">
                    📸 Model Selection Guide
                </h4>
                <div style="display: grid; grid-template-columns: 1fr 1fr; gap: 12px;">
                    <div style="background: white; padding: 12px; border-radius: 8px; border-left: 4px solid #4CAF50;">
                        <p style="margin: 0 0 8px 0; font-weight: bold; color: #4CAF50;">
                            🖼️ Photo Mode (Real Photos)
                        </p>
                        <p style="margin: 0; font-size: 13px; color: #555;">
                            <strong>Best for:</strong> Photographs, portraits, product shots, nature photos
                        </p>
                        <p style="margin: 8px 0 0 0; font-size: 12px; color: #777;">
                            • <strong>JuggernautXL</strong> - Best for portraits and people<br>
                            • <strong>RealVisXL</strong> - Best for scenes and objects
                        </p>
                    </div>
                    <div style="background: white; padding: 12px; border-radius: 8px; border-left: 4px solid #9C27B0;">
                        <p style="margin: 0 0 8px 0; font-weight: bold; color: #9C27B0;">
                            🎨 Anime Mode (Illustrations)
                        </p>
                        <p style="margin: 0; font-size: 13px; color: #555;">
                            <strong>Best for:</strong> Anime, manga, illustrations, digital art, cartoons
                        </p>
                        <p style="margin: 8px 0 0 0; font-size: 12px; color: #777;">
                            • <strong>Animagine XL</strong> - Best for anime/manga style<br>
                            • <strong>SDXL Base</strong> - Versatile for general art
                        </p>
                    </div>
                </div>
            </div>
            """)

            with gr.Row():
                # Left column - Input
                with gr.Column(scale=1):
                    # Image upload
                    inpaint_image = gr.Image(
                        label="Upload Image",
                        type="pil",
                        height=300
                    )

                    # Mask editor
                    mask_editor = gr.ImageEditor(
                        label="Draw Mask (white = area to inpaint)",
                        type="pil",
                        height=300,
                        brush=gr.Brush(colors=["#FFFFFF"], default_size=20),
                        eraser=gr.Eraser(default_size=20),
                        layers=True,
                        sources=["upload"],
                        image_mode="RGBA"
                    )

                    # Template selection
                    with gr.Accordion("Inpainting Templates", open=False):
                        inpaint_template = gr.Dropdown(
                            choices=[""] + self.inpainting_template_manager.get_template_choices_sorted(),
                            value="",
                            label="Select Template",
                            elem_classes=["template-dropdown"]
                        )
                        template_tips = gr.Markdown("")

                    # Mode info (dynamically updated based on template)
                    mode_info_html = gr.HTML("")

                    # Prompt
                    inpaint_prompt = gr.Textbox(
                        label="Prompt",
                        placeholder="Describe what you want to generate in the masked area...",
                        lines=2
                    )

                # Right column - Settings and Output
                with gr.Column(scale=1):
                    # Model Selection (hidden for Pure Inpainting templates)
                    # Using gr.Column for better Gradio 4.x compatibility
                    with gr.Column(visible=True) as model_selection_group:
                        with gr.Accordion("Model Selection", open=True):
                            model_choice = gr.Dropdown(
                                choices=[
                                    ("🖼️ JuggernautXL v9 - Best for portraits & real photos", "juggernaut_xl"),
                                    ("🖼️ RealVisXL v4 - Best for realistic scenes", "realvis_xl"),
                                    ("🎨 SDXL Base - Versatile for general art", "sdxl_base"),
                                    ("🎨 Animagine XL 3.1 - Best for anime/manga", "animagine_xl"),
                                ],
                                value="juggernaut_xl",
                                label="Select Model",
                                info="Choose based on your image type (photo vs illustration)"
                            )

                    # ControlNet Settings (hidden for Pure Inpainting templates)
                    # Using gr.Column for better Gradio 4.x compatibility
                    with gr.Column(visible=True) as controlnet_settings_group:
                        with gr.Accordion("ControlNet Settings", open=True):
                            conditioning_type = gr.Radio(
                                choices=["canny", "depth"],
                                value="canny",
                                label="ControlNet Mode",
                                info="Canny: preserves edges | Depth: preserves 3D structure"
                            )

                            conditioning_scale = gr.Slider(
                                minimum=0.05,
                                maximum=1.0,
                                value=0.7,
                                step=0.05,
                                label="ControlNet Strength",
                                info="Higher = more structure preservation"
                            )

                    # General Settings (always visible)
                    with gr.Accordion("General Settings", open=True):
                        feather_radius = gr.Slider(
                            minimum=0,
                            maximum=20,
                            value=8,
                            step=1,
                            label="Feather Radius (px)",
                            info="Edge blending softness"
                        )

                    with gr.Accordion("Advanced Settings", open=False):
                        inpaint_guidance = gr.Slider(
                            minimum=5.0,
                            maximum=15.0,
                            value=7.5,
                            step=0.5,
                            label="Guidance Scale"
                        )

                        inpaint_steps = gr.Slider(
                            minimum=15,
                            maximum=50,
                            value=25,
                            step=5,
                            label="Inference Steps"
                        )

                        # Seed control for reproducibility
                        seed_input = gr.Number(
                            label="Seed",
                            value=-1,
                            precision=0,
                            info="-1 = random seed, or enter a specific number to reproduce results"
                        )

                    # Generate button
                    inpaint_btn = gr.Button(
                        "Generate Inpainting",
                        variant="primary",
                        elem_classes=["primary-button"]
                    )

                    # Processing time reminder
                    gr.Markdown(
                        """
                        <div style="background: linear-gradient(135deg, #fff8e1 0%, #ffecb3 100%);
                                    border-left: 4px solid #ffa000;
                                    padding: 12px 16px;
                                    border-radius: 8px;
                                    margin: 12px 0;">
                            <p style="margin: 0; color: #5d4037; font-size: 14px;">
                                ⏳ <strong>Please be patient!</strong><br>
                                • <strong>First run:</strong> 5-7 minutes (model initialization)<br>
                                • <strong>Subsequent runs:</strong> 2-3 minutes (model cached)
                            </p>
                        </div>
                        <div style="background: linear-gradient(135deg, #e3f2fd 0%, #bbdefb 100%);
                                    border-left: 4px solid #1976d2;
                                    padding: 12px 16px;
                                    border-radius: 8px;
                                    margin: 12px 0;">
                            <p style="margin: 0; color: #0d47a1; font-size: 14px;">
                                🔄 <strong>Want to make more changes?</strong> After each generation, please
                                <strong>re-upload your image</strong> and draw a new mask if you want to apply additional edits.
                            </p>
                        </div>
                        <div style="background: linear-gradient(135deg, #e8f5e9 0%, #c8e6c9 100%);
                                    border-left: 4px solid #4CAF50;
                                    padding: 12px 16px;
                                    border-radius: 8px;
                                    margin: 12px 0;">
                            <p style="margin: 0; color: #2e7d32; font-size: 14px;">
                                💡 <strong>Good to know:</strong><br>
                                • <strong>No progress bar:</strong> Due to GPU limitations, progress won't update during generation. Please wait for the result.<br>
                                • <strong>Seed for reproducibility:</strong> Each generation uses a random seed. To recreate the same result, copy the "Used Seed" value and enter it in Advanced Settings → Seed.<br>
                                • <strong>Pure Inpainting templates:</strong> "Object Replacement" and "Remove Object" use a dedicated inpainting model for best results. Model selection is ignored for these templates.
                            </p>
                        </div>
                        """
                    )

                    # Status and Seed display
                    inpaint_status = gr.Textbox(
                        label="Status",
                        value="Ready for inpainting",
                        interactive=False
                    )

                    # Display used seed for reproducibility
                    with gr.Row():
                        used_seed_display = gr.Number(
                            label="Used Seed (copy this to reproduce)",
                            value=-1,
                            precision=0,
                            interactive=False
                        )
                        copy_seed_btn = gr.Button(
                            "📋 Use This Seed",
                            size="sm",
                            scale=0
                        )

            # Output row
            with gr.Row():
                with gr.Column(scale=1):
                    inpaint_result = gr.Image(
                        label="Result",
                        type="pil",
                        height=400
                    )

                with gr.Column(scale=1):
                    # Control image (structure guidance visualization)
                    inpaint_control = gr.Image(
                        label="Control Image (Structure Guidance)",
                        type="pil",
                        height=400
                    )

            # Event handlers
            inpaint_template.change(
                fn=self.apply_inpainting_template,
                inputs=[inpaint_template, inpaint_prompt],
                outputs=[
                    inpaint_prompt,
                    conditioning_scale,
                    feather_radius,
                    conditioning_type,
                    controlnet_settings_group,
                    mode_info_html,
                    model_selection_group
                ]
            )

            inpaint_template.change(
                fn=lambda x: self._get_template_tips(x),
                inputs=[inpaint_template],
                outputs=[template_tips]
            )

            # Copy uploaded image to mask editor (as background)
            def set_mask_editor_background(image):
                """Set uploaded image as mask editor background."""
                if image is None:
                    return None
                # Return dict format for ImageEditor with background
                return {"background": image, "layers": [], "composite": None}

            inpaint_image.change(
                fn=set_mask_editor_background,
                inputs=[inpaint_image],
                outputs=[mask_editor]
            )

            inpaint_btn.click(
                fn=self.inpainting_handler,
                inputs=[
                    inpaint_image,
                    mask_editor,
                    inpaint_prompt,
                    inpaint_template,
                    model_choice,
                    conditioning_type,
                    conditioning_scale,
                    feather_radius,
                    inpaint_guidance,
                    inpaint_steps,
                    seed_input
                ],
                outputs=[
                    inpaint_result,
                    inpaint_control,
                    inpaint_status,
                    used_seed_display
                ]
            )

            # Copy seed button - copies used seed to input
            copy_seed_btn.click(
                fn=lambda x: x,
                inputs=[used_seed_display],
                outputs=[seed_input]
            )

        return tab

    def _get_template_tips(self, display_name: str) -> str:
        """Get usage tips for selected template."""
        if not display_name:
            return ""

        template_key = self.inpainting_template_manager.get_template_key_from_display(display_name)
        if not template_key:
            return ""

        tips = self.inpainting_template_manager.get_usage_tips(template_key)
        if tips:
            return "**Tips:**\n" + "\n".join(f"- {tip}" for tip in tips)
        return ""