File size: 6,416 Bytes
8c01052
1d4eb4e
 
 
8c01052
 
 
 
 
 
 
 
 
 
b0c0c9b
e050d47
 
 
 
 
07b236d
 
 
 
 
 
e050d47
 
 
07b236d
 
e050d47
07b236d
 
 
 
 
 
 
b0c0c9b
1d4eb4e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
88b1d8e
8c01052
 
 
 
 
 
 
 
 
 
 
 
 
88b1d8e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1d4eb4e
 
 
 
 
 
 
88b1d8e
1d4eb4e
 
 
4d18b32
1d4eb4e
 
88b1d8e
 
 
 
 
 
b0c0c9b
0ea9e4a
88b1d8e
 
 
8c01052
c009d4f
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
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. "
            "Other spaces: [ImageGen1](https://huggingface.co/spaces/RioShiina/ImageGen1), "
            "[ImageGen2](https://huggingface.co/spaces/RioShiina/ImageGen2), "
            "[ImageGen3](https://huggingface.co/spaces/RioShiina/ImageGen3), "
            "[ImageGen4](https://huggingface.co/spaces/RioShiina/ImageGen4), "
            "[ImageGen5](https://huggingface.co/spaces/RioShiina/ImageGen5), "
            "[ImageGen6](https://huggingface.co/spaces/RioShiina/ImageGen6), "
            "[ImageGen7](https://huggingface.co/spaces/RioShiina/ImageGen7), "
            "[ImageGen8](https://huggingface.co/spaces/RioShiina/ImageGen8)"
        )

        # 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")
                # Main gallery component.
                gallery = gr.Gallery(label="Generated Images", show_label=False, columns=4, height="auto", type="filepath")
                ui_components["gallery"] = gallery
                # 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