tabs/image_tab.py

# 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, loras


def image_tab():
    with gr.Tab(label="🖼️ Image"):
        with gr.Tabs():
            flux_tab()
            sdxl_tab()


def flux_tab():
    with gr.Tab("Flux"):
        with gr.Row():
            with gr.Column():
                with gr.Group() as image_options:
                    model = gr.Dropdown(label="Models", choices=flux_models, value=flux_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
                    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): # only support max 6 loras due to inference time
                                        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): # Embeddings
                    gr.Label("To be implemented")
                
                
                with gr.Accordion("Image Options"): # Image Options
                    with gr.Tabs():
                        image_options = {
                            "img2img": "Upload Image",
                            "inpaint": "Upload Image",
                            "canny": "Upload Image",
                            "pose": "Upload Image",
                            "depth": "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",
                                    )
                                
                                # Image Strength (Co-relates to controlnet strength, strength for img2img n inpaint)
                                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 = [
                                "fm_euler"
                            ],
                            value="fm_euler",
                            interactive=True
                        )

                    with gr.Row():
                        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),
                                    ("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),
                                    ("Guidance Scale", "image_guidance_scale", 0, 20, 0.5, 3.5, 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, [custom_lora, selected_lora])
    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]
        )
    

    # 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, # type: ignore
            img2img_strength, inpaint_strength, canny_strength, pose_strength, depth_strength, # type: ignore
            resize_mode,
            scheduler, image_height, image_width, image_num_images_per_prompt, # type: ignore
            image_num_inference_steps, image_guidance_scale, image_seed, # type: ignore
            refiner, vae
        ],
        [output_images]
    )


def sdxl_tab():
    with gr.Tab("SDXL"):
        gr.Label("To be implemented")