Refactor image_tab.py and flux_tab.py to update LoRA gallery functionality

tabs/image_tab.py

@@ -1,12 +1,7 @@
 # tabs/image_tab.py
 
 import gradio as gr
-from modules.events.flux_events import *
-from modules.events.sdxl_events import *
 from modules.helpers.common_helpers import *
-from modules.helpers.flux_helpers import *
-from modules.helpers.sdxl_helpers import *
-from config import flux_models, sdxl_models, flux_loras
 
 
 def image_tab():
@@ -18,6 +13,17 @@ def image_tab():
 
 
 def flux_tab():
+    from modules.events.flux_events import (
+        update_fast_generation,
+        selected_lora_from_gallery,
+        update_selected_lora,
+        add_to_enabled_loras,
+        update_lora_sliders,
+        remove_from_enabled_loras,
+        generate_image
+    )
+    from config import flux_models, flux_loras
+    
     loras = flux_loras
     with gr.Row():
         with gr.Column():
@@ -122,8 +128,8 @@ def flux_tab():
                     for column in range(2):
                         with gr.Column():
                             options = [
-                                ("Height", "image_height", 64, 1024, 64, 1024, True),
-                                ("Width", "image_width", 64, 1024, 64, 1024, True),
+                                ("Height", "image_height", 64, 2048, 64, 1024, True),
+                                ("Width", "image_width", 64, 2048, 64, 1024, True),
                                 ("Num Images Per Prompt", "image_num_images_per_prompt", 1, 4, 1, 1, True),
                                 ("Num Inference Steps", "image_num_inference_steps", 1, 100, 1, 20, True),
                                 ("Clip Skip", "image_clip_skip", 0, 2, 1, 2, False),
@@ -180,4 +186,197 @@ def flux_tab():
 
 
 def sdxl_tab():
-    gr.Label("To be implemented")
+    from modules.events.sdxl_events import (
+        update_fast_generation,
+        selected_lora_from_gallery,
+        update_selected_lora,
+        add_to_enabled_loras,
+        update_lora_sliders,
+        remove_from_enabled_loras,
+        add_to_embeddings,
+        update_custom_embedding,
+        remove_from_embeddings,
+        generate_image
+    )
+    from config import sdxl_models, sdxl_loras
+    
+    loras = sdxl_loras
+    with gr.Row():
+        with gr.Column():
+            with gr.Group() as image_options:
+                model = gr.Dropdown(label="Models", choices=sdxl_models, value=sdxl_models[0], interactive=True)
+                prompt = gr.Textbox(lines=5, label="Prompt")
+                fast_generation = gr.Checkbox(label="Fast Generation (Hyper-SD) 🧪")
+            
+            
+            with gr.Accordion("Loras", open=True):
+                lora_gallery = gr.Gallery(
+                    label="Gallery",
+                    value=[(lora['image'], lora['title']) for lora in loras],
+                    allow_preview=False,
+                    columns=3,
+                    rows=3,
+                    type="pil"
+                )
+                
+                with gr.Group():
+                    with gr.Column():
+                        with gr.Row():
+                            custom_lora = gr.Textbox(label="Custom Lora", info="Enter a Huggingface repo path")
+                            selected_lora = gr.Textbox(label="Selected Lora", info="Choose from the gallery or enter a custom LoRA")
+                        
+                        custom_lora_info = gr.HTML(visible=False)
+                        add_lora = gr.Button(value="Add LoRA")
+                        
+                        enabled_loras = gr.State(value=[])
+                        with gr.Group():
+                            with gr.Row():
+                                for i in range(6):
+                                    with gr.Column():
+                                        with gr.Column(scale=2):
+                                            globals()[f"lora_slider_{i}"] = gr.Slider(label=f"LoRA {i+1}", minimum=0, maximum=1, step=0.01, value=0.8, visible=False, interactive=True)
+                                        with gr.Column():
+                                            globals()[f"lora_remove_{i}"] = gr.Button(value="Remove LoRA", visible=False)
+                                            
+            with gr.Accordion("Embeddings", open=False):
+                custom_embedding = gr.Textbox(label="Custom Embedding")
+                custom_embedding_info = gr.HTML(visible=False)
+                add_embedding = gr.Button(value="Add Embedding")
+                embeddings = gr.State(value=[])
+                with gr.Group():
+                    with gr.Row():
+                        for i in range(6):
+                            with gr.Column():
+                                with gr.Column(scale=2):
+                                    globals()[f"embedding_list_{i}"] = gr.Label(label=f"Embedding {i+1}", visible=False)
+                                with gr.Column():
+                                    globals()[f"embedding_remove_{i}"] = gr.Button(value="Remove Embedding", visible=False)
+            
+            with gr.Accordion("Image Options", open=False):
+                with gr.Tabs():
+                    image_options = {
+                        "img2img": "Upload Image",
+                        "inpaint": "Upload Image",
+                        "canny": "Upload Image",
+                        "pose": "Upload Image",
+                        "depth": "Upload Image",
+                        "scribble": "Upload Image",
+                    }
+                    
+                    for image_option, label in image_options.items():
+                        with gr.Tab(image_option):
+                            if not image_option in ['inpaint', 'scribble']:
+                                globals()[f"{image_option}_image"] = gr.Image(label=label, type="pil")
+                            elif image_option in ['inpaint', 'scribble']:
+                                globals()[f"{image_option}_image"] = gr.ImageEditor(
+                                    label=label,
+                                    image_mode='RGB',
+                                    layers=False,
+                                    brush=gr.Brush(colors=["#FFFFFF"], color_mode="fixed") if image_option == 'inpaint' else gr.Brush(),
+                                    interactive=True,
+                                    type="pil",
+                                )
+                            
+                            globals()[f"{image_option}_strength"] = gr.Slider(label="Strength", minimum=0, maximum=1, step=0.01, value=1.0, interactive=True)
+                    
+                    resize_mode = gr.Radio(
+                        label="Resize Mode",
+                        choices=["crop and resize", "resize only", "resize and fill"],
+                        value="resize and fill",
+                        interactive=True
+                    )
+        
+        with gr.Column():
+            with gr.Group():
+                output_images = gr.Gallery(
+                        label="Output Images",
+                        value=[],
+                        allow_preview=True,
+                        type="pil",
+                        interactive=False,
+                    )
+                generate_images = gr.Button(value="Generate Images", variant="primary")
+            
+            with gr.Accordion("Advance Settings", open=True):
+                with gr.Row():
+                    scheduler = gr.Dropdown(
+                        label="Scheduler",
+                        choices = [
+                            "dpmpp_2m", "dpmpp_2m_k", "dpmpp_2m_sde", "dpmpp_2m_sde_k",
+                            "dpmpp_sde", "dpmpp_sde_k", "dpm2", "dpm2_k", "dpm2_a",
+                            "dpm2_a_k", "euler", "euler_a", "heun", "lms", "lms_k",
+                            "deis", "unipc"
+                        ]
+                        value="dpmpp_2m_sde_k",
+                        interactive=True
+                    )
+
+                with gr.Row():
+                    for column in range(2):
+                        with gr.Column():
+                            options = [
+                                ("Height", "image_height", 64, 2048, 64, 1024, True),
+                                ("Width", "image_width", 64, 2048, 64, 1024, True),
+                                ("Num Images Per Prompt", "image_num_images_per_prompt", 1, 4, 1, 1, True),
+                                ("Num Inference Steps", "image_num_inference_steps", 1, 100, 1, 20, True),
+                                ("Clip Skip", "image_clip_skip", 0, 2, 1, 2, True),
+                                ("Guidance Scale", "image_guidance_scale", 0, 20, 0.5, 7.0, True),
+                                ("Seed", "image_seed", 0, 100000, 1, random.randint(0, 100000), True),
+                            ]
+                            for label, var_name, min_val, max_val, step, value, visible in options[column::2]:
+                                globals()[var_name] = gr.Slider(label=label, minimum=min_val, maximum=max_val, step=step, value=value, visible=visible, interactive=True)
+                
+                with gr.Row():
+                    refiner = gr.Checkbox(
+                        label="Refiner 🧪",
+                        value=False,
+                    )
+                    vae = gr.Checkbox(
+                        label="VAE",
+                        value=True,
+                    )
+    
+    # Events
+    # Base Options
+    fast_generation.change(update_fast_generation, [fast_generation], [image_guidance_scale, image_num_inference_steps]) # Fast Generation # type: ignore
+    
+    
+    # Lora Gallery
+    lora_gallery.select(selected_lora_from_gallery, None, selected_lora)
+    custom_lora.change(update_selected_lora, custom_lora, [selected_lora, custom_lora_info])
+    add_lora.click(add_to_enabled_loras, [selected_lora, enabled_loras], [selected_lora, custom_lora_info, enabled_loras])
+    enabled_loras.change(update_lora_sliders, enabled_loras, [lora_slider_0, lora_slider_1, lora_slider_2, lora_slider_3, lora_slider_4, lora_slider_5, lora_remove_0, lora_remove_1, lora_remove_2, lora_remove_3, lora_remove_4, lora_remove_5]) # type: ignore
+    
+    for i in range(6):
+        globals()[f"lora_remove_{i}"].click(
+            lambda enabled_loras, index=i: remove_from_enabled_loras(enabled_loras, index),
+            [enabled_loras],
+            [enabled_loras]
+        )
+
+    
+    # Embeddings
+    custom_embedding.change(update_custom_embedding, custom_embedding, [custom_embedding_info])
+    add_embedding.click(add_to_embeddings, [custom_embedding, embeddings], [custom_embedding, custom_embedding_info, embeddings])
+    for i in range(6):
+        globals()[f"embedding_remove_{i}"].click(
+            lambda embeddings, index=i: remove_from_embeddings(embeddings, index),
+            [embeddings],
+            [embeddings]
+        )
+    
+    # Generate Image
+    generate_images.click(
+        generate_image, # type: ignore
+        [
+            model, prompt, fast_generation, enabled_loras,
+            lora_slider_0, lora_slider_1, lora_slider_2, lora_slider_3, lora_slider_4, lora_slider_5, # type: ignore
+            img2img_image, inpaint_image, canny_image, pose_image, depth_image, scribble_image, # type: ignore
+            img2img_strength, inpaint_strength, canny_strength, pose_strength, depth_strength, scribble_strength, # type: ignore
+            resize_mode,
+            scheduler, image_height, image_width, image_num_images_per_prompt, # type: ignore
+            image_num_inference_steps, image_clip_skip, image_guidance_scale, image_seed, # type: ignore
+            refiner, vae
+        ],
+        [output_images]
+    )