src/ui/generator_tab.py

"""Image generator tab for SDXL Model Merger."""

import gradio as gr

from ..config import DEFAULT_PROMPT, DEFAULT_NEGATIVE_PROMPT
from ..generator import generate_image


def create_generator_tab():
    """Create the image generation tab with all input controls."""

    with gr.Accordion("🎨 2. Generate Image", open=True, elem_classes=["feature-card"]):
        # Prompts section
        with gr.Row():
            with gr.Column(scale=1):
                prompt = gr.Textbox(
                    label="Positive Prompt",
                    value=DEFAULT_PROMPT,
                    lines=3,
                    placeholder="Describe the image you want to generate..."
                )

                cfg = gr.Slider(
                    minimum=1.0, maximum=20.0, value=7.5, step=0.5,
                    label="CFG Scale",
                    info="Higher values make outputs match prompt more strictly"
                )

                height = gr.Number(
                    value=1024, precision=0,
                    label="Height (pixels)",
                    info="Output image height"
                )

            with gr.Column(scale=1):
                negative_prompt = gr.Textbox(
                    label="Negative Prompt",
                    value=DEFAULT_NEGATIVE_PROMPT,
                    lines=3,
                    placeholder="Elements to avoid in generation..."
                )

                steps = gr.Slider(
                    minimum=1, maximum=100, value=25, step=1,
                    label="Inference Steps",
                    info="More steps = better quality but slower"
                )

                width = gr.Number(
                    value=2048, precision=0,
                    label="Width (pixels)",
                    info="Output image width"
                )

        # Tiling options
        with gr.Row():
            tile_x = gr.Checkbox(True, label="X-axis Seamless Tiling")
            tile_y = gr.Checkbox(False, label="Y-axis Seamless Tiling")

        # Generate button and outputs
        with gr.Row():
            gen_btn = gr.Button("✨ Generate Image", variant="secondary", size="lg")

        with gr.Row():
            image_output = gr.Image(
                label="Result",
                height=400,
                show_label=True
            )
            with gr.Column():
                gen_status = gr.HTML(
                    label="Generation Status",
                    value='<div class="status-success">✅ Ready to generate</div>',
                )

                gen_progress = gr.Textbox(
                    label="Generation Progress",
                    placeholder="Ready to generate...",
                    show_label=True,
                    visible=False
                )

                gr.HTML("""
                    <div style="margin-top: 16px; padding: 12px; background: #f3f4f6; border-radius: 8px;">
                        <strong>💡 Tips:</strong>
                        <ul style="margin: 8px 0; padding-left: 20px; font-size: 0.9em;">
                            <li>Use wide aspect ratios (e.g., 1024x2048) for panoramas</li>
                            <li>Enable seamless tiling for texture-like outputs</li>
                            <li>Lower CFG (3-5) for more creative results</li>
                        </ul>
                    </div>
                """)

    return (
        prompt, negative_prompt, cfg, steps, height, width,
        tile_x, tile_y, gen_btn, image_output, gen_status, gen_progress
    )


def setup_generator_events(
    prompt, negative_prompt, cfg, steps, height, width,
    tile_x, tile_y, gen_btn, image_output, gen_status, gen_progress
):
    """Setup event handlers for the generator tab."""

    def on_generate_start():
        return (
            '<div class="status-warning">⏳ Generating image...</div>',
            "Starting generation...",
            gr.update(interactive=False),
        )

    def on_generate_complete(img, msg):
        """Read both image and last progress message to determine success/failure."""
        if img is None:
            return (
                f'<div class="status-error">{msg}</div>',
                "",
                gr.update(interactive=True),
                gr.update(),
            )
        return (
            '<div class="status-success">✅ Generation complete!</div>',
            "Done",
            gr.update(interactive=True),
            gr.update(value=img),
        )

    gen_btn.click(
        fn=on_generate_start,
        inputs=[],
        outputs=[gen_status, gen_progress, gen_btn],
    ).then(
        fn=generate_image,
        inputs=[prompt, negative_prompt, cfg, steps, height, width, tile_x, tile_y],
        outputs=[image_output, gen_progress],
    ).then(
        fn=on_generate_complete,
        inputs=[image_output, gen_progress],
        outputs=[gen_status, gen_progress, gen_btn, image_output],
    )