tabs/images/ui.py

import random

import gradio as gr

from config import Config
from .events import *


def image_tab():
    with gr.Row():
        with gr.Column():
            with gr.Group():
                model = gr.Dropdown(label='Model', choices=[model['repo_id'] for model in Config.IMAGES_MODELS], value=Config.IMAGES_MODELS[0]['repo_id'], interactive=True)
                prompt = gr.Textbox(lines=5, label='Prompt', placeholder='Enter your prompt here...', value='A beautiful sunset over the mountains.')
                negative_prompt = gr.Textbox(lines=2, label='Negative Prompt', placeholder='Enter your negative prompt here...', visible=False)
                fast_generation = gr.Checkbox(label='Fast Generation (Hyper-SD 🧪)', value=False)
                
            
            with gr.Accordion('Loras', open=True):
                for m in Config.IMAGES_MODELS:
                    if m['repo_id'] == model.value:
                        lora_gallery_values = []
                        if m['loader'] == 'flux':
                            lora_gallery_values = [(lora['image'], lora['title']) for lora in Config.IMAGES_LORAS_FLUX]
                        elif m['loader'] == 'sdxl':
                            lora_gallery_values = [(lora['image'], lora['title']) for lora in Config.IMAGES_LORAS_SDXL]
                
                lora_gallery = gr.Gallery(
                    label='Loras',
                    value=lora_gallery_values,
                    allow_preview=False,
                    interactive=True,
                    rows=2,
                    columns=3,
                )
                
                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, visible=False) as embeddings_accordion:
                with gr.Group():
                    with gr.Row():
                        with gr.Group():
                            custom_embedding = gr.Textbox(label="Custom Embedding", info="Enter a Huggingface repo path")
                            add_embedding = gr.Button(value="Add Embedding")
                            custom_embedding_info = gr.HTML(visible=False)
                    with gr.Row():
                        enabled_embeddings = gr.State(value=[])
                        enabled_embeddings_list = gr.Checkboxgroup(label="Enabled Embeddings", choices=[], visible=False)
            
            
            with gr.Accordion('Image Options', open=False):
            
                with gr.Tabs():
                    image_options = [
                        ('img2img', 'Image to Image', 'image', True),
                        ('inpaint', 'Inpainting', 'imageeditor', True),
                        ('canny', 'Edge Detection', 'image', True),
                        ('pose', 'Pose Detection', 'image', True),
                        ('depth', 'Depth Estimation', 'image', True),
                        ('scribble', 'Scribble', 'imageeditor', False),
                    ]
                    for image_option, label, type, visible in image_options:
                        with gr.Tab(label=image_option) as globals()[f"{image_option}_tab"]:
                            if type == 'image':
                                globals()[f"{image_option}_image"] = gr.Image(label=label, visible=visible, interactive=True, type='pil')
                            elif type == 'imageeditor':
                                globals()[f"{image_option}_image"] = gr.ImageEditor(label=label, visible=visible, interactive=True,
                                                                                    brush=gr.Brush(colors=["#FFFFFF"], color_mode="fixed") if image_option == 'inpaint' else gr.Brush(),
                                                                                    type='pil', image_mode='RGB', layers=False)
                            
                            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 Image', type='pil', interactive=False, value=[], allow_preview=True)
                generate = gr.Button(value="Generate", variant="primary")
            
            
            with gr.Accordion('Advance Settings', open=True):
                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", "fm_euler"
                    ],
                    value="fm_euler",
                    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, 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
    model.change(update_model_options, [model], [negative_prompt, lora_gallery, embeddings_accordion, scribble_tab, scheduler, image_clip_skip, image_guidance_scale]) # type: ignore
    fast_generation.change(update_fast_generation, [model, fast_generation], [image_num_inference_steps, image_guidance_scale]) # type: ignore
    
    # Loras
    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, [model, 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, enabled_embeddings], [custom_embedding, custom_embedding_info, enabled_embeddings])
    enabled_embeddings.change(update_enabled_embeddings_list, enabled_embeddings, [enabled_embeddings_list]) # type: ignore
    enabled_embeddings_list.change(update_enabled_embeddings, enabled_embeddings_list, [enabled_embeddings]) # type: ignore
    
    # Generate Image
    generate.click(
        generate_image, # type: ignore
        [
            model, prompt, negative_prompt, fast_generation, enabled_loras, enabled_embeddings,
            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]
    )