ui/layout.py

import os
import shutil
import zipfile
import tempfile
import gradio as gr
from core.settings import *

from .shared import txt2img_ui, img2img_ui, inpaint_ui, outpaint_ui, hires_fix_ui

MAX_DYNAMIC_CONTROLS = 10

def build_ui(event_handler_function):
    ui_components = {}

    # Helper function to load thumbnails from the output folder.
    # It returns a **deduplicated** list of absolute image file paths.
    # Using a set ensures that even if the same filename appears multiple
    # times (e.g., due to accidental duplicate saves), the Gallery tab will
    # only display each image once.
    def load_gallery_images():
        """Return a deduplicated list of image paths.
        If multiple files share the same stem (e.g., ``img.png`` and ``img.webp``),
        only the first encountered file (based on sorted order) is kept. This prevents
        the Gallery tab from displaying duplicate visual entries that arise from
        different extensions of the same generated image.
        """
        output_dir = os.path.abspath(OUTPUT_DIR)
        if not os.path.isdir(output_dir):
            return []
        seen_bases = set()
        image_paths = []
        for fname in sorted(os.listdir(output_dir)):
            low = fname.lower()
            if low.endswith((".png", ".jpg", ".jpeg", ".gif", ".webp")):
                base = os.path.splitext(fname)[0]
                if base not in seen_bases:
                    seen_bases.add(base)
                    image_paths.append(os.path.join(output_dir, fname))
        return image_paths

    # Clear all images from the gallery folder and return an empty list for the UI.
    def clear_gallery_images():
        """Remove every file in ``OUTPUT_DIR`` and return an empty list for the gallery.
        This provides a quick way for users to reset the gallery view.
        """
        output_dir = os.path.abspath(OUTPUT_DIR)
        if os.path.isdir(output_dir):
            for f in os.listdir(output_dir):
                try:
                    os.remove(os.path.join(output_dir, f))
                except Exception:
                    pass
        return []

    # Create a zip archive of all images in the gallery for downloading.
    def download_gallery_zip():
        """Package all files in ``OUTPUT_DIR`` into a temporary zip file.
        Returns the file path which Gradio will serve as a downloadable file.
        """
        output_dir = os.path.abspath(OUTPUT_DIR)
        tmp_fd, tmp_path = tempfile.mkstemp(suffix=".zip")
        os.close(tmp_fd)
        with zipfile.ZipFile(tmp_path, "w", zipfile.ZIP_DEFLATED) as zf:
            if os.path.isdir(output_dir):
                for fname in sorted(os.listdir(output_dir)):
                    file_path = os.path.join(output_dir, fname)
                    if os.path.isfile(file_path):
                        zf.write(file_path, arcname=fname)
        return tmp_path

    # All UI components, including tabs, must be created inside the same Gradio Blocks context.
    with gr.Blocks() as demo:
        gr.Markdown("# ImageGen")
        gr.Markdown(
            "This demo is a streamlined version of the [Comfy web UI](https://github.com/RioShiina47/comfy-webui)'s [ImageGen](https://huggingface.co/spaces/RioShiina/ImageGen) functionality. "
        )

        # Tabs container – now correctly nested within the Blocks context.
        with gr.Tabs(elem_id="tabs_container") as tabs:
            with gr.TabItem("Txt2Img", id=0):
                ui_components.update(txt2img_ui.create_ui())

            with gr.TabItem("Img2Img", id=1):
                ui_components.update(img2img_ui.create_ui())

            with gr.TabItem("Inpaint", id=2):
                ui_components.update(inpaint_ui.create_ui())

            with gr.TabItem("Outpaint", id=3):
                ui_components.update(outpaint_ui.create_ui())

            with gr.TabItem("Hires. Fix", id=4):
                ui_components.update(hires_fix_ui.create_ui())

            # New Gallery tab to display generated images
            with gr.TabItem("Gallery", id=5):
                # Row of control buttons for the gallery.
                with gr.Row():
                    refresh_btn = gr.Button("Refresh Gallery", variant="secondary")
                    clear_btn = gr.Button("Clear Gallery", variant="secondary")
                    download_btn = gr.Button("Download All", variant="secondary")
                # Gallery and metadata side by side.
                with gr.Row():
                    # Left side: Gallery (larger column)
                    with gr.Column(scale=3):
                        gallery = gr.Gallery(label="Generated Images", show_label=False, columns=4, height=622, type="filepath")
                        ui_components["gallery"] = gallery
                    # Right side: Metadata textbox + drag‑and‑drop (smaller column)
                    with gr.Column(scale=1):
                        metadata_box = gr.Textbox(label="Generation info", lines=15, interactive=False)
                        ui_components["gallery_metadata"] = metadata_box
                        info_upload = gr.File(label="Drag image here for info", file_count="single", type="filepath", interactive=True)
                        ui_components["info_upload"] = info_upload
                # Hidden file component for zip download.
                download_file = gr.File(label="Download", visible=False)
                # Bind buttons.
                refresh_btn.click(fn=load_gallery_images, outputs=gallery)
                clear_btn.click(fn=clear_gallery_images, outputs=gallery)
                download_btn.click(fn=download_gallery_zip, outputs=download_file)

            ui_components["tabs"] = tabs
            ui_components["image_gen_tabs"] = tabs

        gr.Markdown("<div style='text-align: center; margin-top: 20px;'>Made by RioShiina with ❤️<br><a href='https://github.com/RioShiina47' target='_blank'>GitHub</a> | <a href='https://huggingface.co/RioShiina' target='_blank'>Hugging Face</a> | <a href='https://civitai.com/user/RioShiina' target='_blank'>Civitai</a></div>")

        # Load gallery images on startup using the helper defined earlier.
        demo.load(fn=load_gallery_images, inputs=[], outputs=ui_components["gallery"])

        # Connect event handlers (e.g., button clicks) with the UI components.
        event_handler_function(ui_components, demo)

    return demo