ui/imagegen_ui.py

import gradio as gr
from core.settings import MODEL_MAP_CHECKPOINT, MODEL_DEFAULTS_CONFIG
from comfy_integration.nodes import SAMPLER_CHOICES, SCHEDULER_CHOICES
from .shared.ui_components import (
    RESOLUTION_MAP,
    create_lora_settings_ui,
    create_controlnet_ui, create_ipadapter_ui, create_embedding_ui,
    create_conditioning_ui, create_vae_override_ui,
    create_model_architecture_filter_ui, create_category_filter_ui,
    create_sd3_ipadapter_ui, create_flux1_ipadapter_ui, create_style_ui,
    create_reference_latent_ui, create_hidream_o1_reference_ui, create_sensenova_reference_ui, create_joyai_reference_ui, create_reference_image_ui, create_krea2_identity_edit_ui, create_krea2_reference_edit_ui, create_qwen_image_edit_ui,
    create_pid_ui, create_qwen_image_2_1_prompt_enhancer_ui, create_ming_image_prompt_enhancer_ui
)

default_vals = MODEL_DEFAULTS_CONFIG.get('Default', {})
DEFAULT_STEPS = default_vals.get('steps', 20)
DEFAULT_CFG = default_vals.get('cfg', 5.0)
DEFAULT_SAMPLER = default_vals.get('sampler_name', 'euler')
DEFAULT_SCHEDULER = default_vals.get('scheduler', 'simple')
DEFAULT_POS_PROMPT = default_vals.get('positive_prompt', '')
DEFAULT_NEG_PROMPT = default_vals.get('negative_prompt', '')

TYPE_CHOICES = ["Txt2Img", "Img2Img", "Inpaint", "Outpaint", "Hires. Fix"]

def create_ui():
    """Creates the unified UI components for ImageGen."""
    prefix = "imagegen"
    components = {}
    
    with gr.Column():
        components[f'type_{prefix}'] = gr.Radio(
            choices=TYPE_CHOICES,
            value="Txt2Img",
            label="Task",
            interactive=True
        )

        with gr.Row() as arch_row:
            components.update(create_model_architecture_filter_ui(prefix))
        components[f'arch_row_{prefix}'] = arch_row
        
        with gr.Row() as model_and_run_row:
            components.update(create_category_filter_ui(prefix))
            components[f'base_model_{prefix}'] = gr.Dropdown(
                label="Base Model", 
                choices=list(MODEL_MAP_CHECKPOINT.keys()), 
                value=list(MODEL_MAP_CHECKPOINT.keys())[0], 
                scale=3,
                allow_custom_value=True
            )
            with gr.Column(scale=1): 
                components[f'run_{prefix}'] = gr.Button("Run", variant="primary")
        components[f'model_and_run_row_{prefix}'] = [arch_row, model_and_run_row]

        with gr.Row() as inputs_prompts_row:
            with gr.Column(scale=1, visible=False) as image_input_col:
                components[f'input_image_{prefix}'] = gr.Image(
                    type="pil", 
                    label="Input Image (No Scaling)", 
                    height=255, 
                    visible=False
                )
                with gr.Column(visible=False) as inpaint_box:
                    components[f'view_mode_{prefix}'] = gr.Radio(
                        ["Normal View", "Fullscreen View"], 
                        label="Editor View", 
                        value="Normal View", 
                        interactive=True
                    )
                    components[f'input_image_dict_{prefix}'] = gr.ImageEditor(
                        type="pil", 
                        label="Image & Mask (No Scaling)", 
                        height=272
                    )
                components[f'inpaint_box_{prefix}'] = inpaint_box
            components[f'image_input_col_{prefix}'] = image_input_col

            with gr.Column(scale=2) as prompts_col:
                components[f'prompt_{prefix}'] = gr.Text(label="Prompt", lines=3, value=DEFAULT_POS_PROMPT)
                components[f'neg_prompt_{prefix}'] = gr.Text(label="Negative prompt", lines=3, value=DEFAULT_NEG_PROMPT)
            components[f'prompts_col_{prefix}'] = prompts_col
        components[f'inputs_prompts_row_{prefix}'] = inputs_prompts_row

        with gr.Row() as params_and_gallery_row:
            with gr.Column(scale=1) as params_col:
                # Txt2Img parameters
                with gr.Row(visible=True) as aspect_ratio_row:
                    components[f'aspect_ratio_{prefix}'] = gr.Dropdown(
                        label="Aspect Ratio",
                        choices=list(RESOLUTION_MAP.get('sdxl', {}).keys()),
                        value="1:1 (Square)",
                        interactive=True,
                        allow_custom_value=True
                    )
                components[f'aspect_ratio_row_{prefix}'] = aspect_ratio_row

                with gr.Row(visible=True) as width_height_row:
                    components[f'width_{prefix}'] = gr.Number(label="Width", value=1024, interactive=True)
                    components[f'height_{prefix}'] = gr.Number(label="Height", value=1024, interactive=True)
                components[f'width_height_row_{prefix}'] = width_height_row

                # Img2Img / Inpaint / Hires. Fix Denoise parameter & Inpaint Grow Mask By parameter
                with gr.Row(visible=False) as denoise_row:
                    components[f'denoise_{prefix}'] = gr.Slider(
                        label="Denoise Strength", 
                        minimum=0.0, 
                        maximum=1.0, 
                        step=0.01, 
                        value=0.7
                    )
                    components[f'grow_mask_by_{prefix}'] = gr.Slider(
                        label="Grow Mask By", 
                        minimum=0, 
                        maximum=64, 
                        step=1, 
                        value=6,
                        visible=False
                    )
                components[f'denoise_row_{prefix}'] = denoise_row

                # Outpaint Pad & Feathering parameters
                with gr.Column(visible=False) as outpaint_pads_col:
                    with gr.Row():
                        components[f'left_{prefix}'] = gr.Slider(label="Pad Left", minimum=0, maximum=512, step=64, value=64)
                        components[f'right_{prefix}'] = gr.Slider(label="Pad Right", minimum=0, maximum=512, step=64, value=64)
                    with gr.Row():
                        components[f'top_{prefix}'] = gr.Slider(label="Pad Top", minimum=0, maximum=512, step=64, value=64)
                        components[f'bottom_{prefix}'] = gr.Slider(label="Pad Bottom", minimum=0, maximum=512, step=64, value=64)
                    components[f'feathering_{prefix}'] = gr.Slider(label="Feathering / Grow Mask", minimum=0, maximum=100, step=1, value=10)
                components[f'outpaint_pads_col_{prefix}'] = outpaint_pads_col

                # Hires. Fix Upscaler parameters
                with gr.Row(visible=False) as hires_upscaler_row:
                    components[f'hires_upscaler_{prefix}'] = gr.Dropdown(
                        label="Upscaler", 
                        choices=["nearest-exact", "bilinear", "area", "bicubic", "bislerp"], 
                        value="nearest-exact"
                    )
                    components[f'hires_scale_by_{prefix}'] = gr.Slider(
                        label="Upscale by", 
                        minimum=1.0, 
                        maximum=4.0, 
                        step=0.1, 
                        value=1.5
                    )
                components[f'hires_upscaler_row_{prefix}'] = hires_upscaler_row

                # Common parameters
                with gr.Row():
                    components[f'sampler_{prefix}'] = gr.Dropdown(
                        label="Sampler", 
                        choices=SAMPLER_CHOICES, 
                        value=DEFAULT_SAMPLER if DEFAULT_SAMPLER in SAMPLER_CHOICES else (SAMPLER_CHOICES[0] if SAMPLER_CHOICES else 'euler')
                    )
                    components[f'scheduler_{prefix}'] = gr.Dropdown(
                        label="Scheduler", 
                        choices=SCHEDULER_CHOICES, 
                        value=DEFAULT_SCHEDULER if DEFAULT_SCHEDULER in SCHEDULER_CHOICES else (SCHEDULER_CHOICES[0] if SCHEDULER_CHOICES else 'simple')
                    )
                with gr.Row():
                    components[f'steps_{prefix}'] = gr.Slider(label="Steps", minimum=1, maximum=100, step=1, value=DEFAULT_STEPS)
                    components[f'cfg_{prefix}'] = gr.Slider(label="CFG Scale", minimum=1.0, maximum=20.0, step=0.1, value=DEFAULT_CFG)
                with gr.Row():
                    components[f'seed_{prefix}'] = gr.Number(label="Seed (-1 for random)", value=-1, precision=0)
                    components[f'batch_size_{prefix}'] = gr.Slider(label="Batch Size", minimum=1, maximum=16, step=1, value=1)
                with gr.Row():
                    components[f'zero_gpu_{prefix}'] = gr.Number(label="ZeroGPU Duration (s)", value=60, placeholder="Default: 60s", info="Optional: Set how long to reserve the GPU. Longer jobs may need more time.")

            with gr.Column(scale=1):
                components[f'result_{prefix}'] = gr.Gallery(label="Result", show_label=False, columns=2, object_fit="contain", height=627)
        components[f'params_and_gallery_row_{prefix}'] = params_and_gallery_row

        with gr.Column() as accordion_wrapper:
            components.update(create_lora_settings_ui(prefix))
            components.update(create_controlnet_ui(prefix))
            components.update(create_ipadapter_ui(prefix))
            components.update(create_flux1_ipadapter_ui(prefix))
            components.update(create_sd3_ipadapter_ui(prefix))
            components.update(create_embedding_ui(prefix))
            components.update(create_style_ui(prefix))
            components.update(create_conditioning_ui(prefix))
            components.update(create_reference_latent_ui(prefix))
            components.update(create_hidream_o1_reference_ui(prefix))
            components.update(create_sensenova_reference_ui(prefix))
            components.update(create_joyai_reference_ui(prefix))
            components.update(create_krea2_identity_edit_ui(prefix))
            components.update(create_krea2_reference_edit_ui(prefix))
            components.update(create_qwen_image_edit_ui(prefix))
            components.update(create_reference_image_ui(prefix))
            components.update(create_qwen_image_2_1_prompt_enhancer_ui(prefix))
            components.update(create_ming_image_prompt_enhancer_ui(prefix))
            components.update(create_vae_override_ui(prefix))
            components.update(create_pid_ui(prefix))
        components[f'accordion_wrapper_{prefix}'] = accordion_wrapper
        
    return components