Add Gallery tab and load images

ui/layout.py

@@ -1,48 +1,68 @@
-import os
-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 = {}
-
-    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)"
-        )
-        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())
-
-            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>")
-        
-        event_handler_function(ui_components, demo)
-
+import os
+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 = {}
+
+    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)"
+        )
+    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):
+            gallery = gr.Gallery(label="Generated Images", show_label=False).style(grid=[3], height="auto")
+            ui_components["gallery"] = gallery
+
+        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
+    def load_gallery_images():
+        # Resolve absolute path for the output directory
+        output_dir = os.path.abspath(OUTPUT_DIR)
+        if not os.path.isdir(output_dir):
+            return []
+        # Collect image files (common formats)
+        image_paths = []
+        for fname in sorted(os.listdir(output_dir)):
+            if fname.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.webp')):
+                image_paths.append(os.path.join(output_dir, fname))
+        return image_paths
+
+    demo.load(fn=load_gallery_images, outputs=ui_components["gallery"])
+
+    event_handler_function(ui_components, demo)
+
     return demo