ui/shared/ui_components.py

import gradio as gr
from comfy_integration.nodes import SAMPLER_CHOICES, SCHEDULER_CHOICES
from core.settings import (
    MAX_LORAS, LORA_SOURCE_CHOICES, MAX_EMBEDDINGS, MAX_CONDITIONINGS,
    MAX_CONTROLNETS, RESOLUTION_MAP
)
import yaml
import os
from functools import lru_cache

def create_base_parameter_ui(prefix, defaults=None):
    if defaults is None:
        defaults = {}
    
    components = {}
    
    with gr.Row():
        components[f'aspect_ratio_{prefix}'] = gr.Dropdown(
            label="Aspect Ratio",
            choices=list(RESOLUTION_MAP['sdxl'].keys()),
            value="1:1 (Square)",
            interactive=True
        )
    with gr.Row():
        components[f'width_{prefix}'] = gr.Number(label="Width", value=defaults.get('w', 1024), interactive=True)
        components[f'height_{prefix}'] = gr.Number(label="Height", value=defaults.get('h', 1024), interactive=True)
    with gr.Row():
        components[f'sampler_{prefix}'] = gr.Dropdown(label="Sampler", choices=SAMPLER_CHOICES, value="res_multistep")
        components[f'scheduler_{prefix}'] = gr.Dropdown(label="Scheduler", choices=SCHEDULER_CHOICES, value="simple")
    with gr.Row():
        components[f'steps_{prefix}'] = gr.Slider(label="Steps", minimum=1, maximum=100, step=1, value=9)
        components[f'cfg_{prefix}'] = gr.Slider(label="CFG Scale", minimum=1.0, maximum=20.0, step=0.1, value=1.0)
    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=None, placeholder="Default: 60s, Max: 120s", info="Optional: Set how long to reserve the GPU. Longer jobs may need more time.")
    
    components[f'clip_skip_{prefix}'] = gr.State(value=1)
            
    return components


def create_api_key_ui(prefix: str):
    components = {}
    with gr.Accordion("API Key Settings", open=False) as api_key_accordion:
        components[f'api_key_accordion_{prefix}'] = api_key_accordion
        gr.Markdown("💡 **Tip:** Enter API key (optional). An API key is required for resources that need a login to download. The key will be used for all Civitai downloads on this tab. You can also manually upload the corresponding files to avoid API Key leakage caused by potential vulnerabilities.")
        with gr.Row():
             components[f'civitai_api_key_{prefix}'] = gr.Textbox(
                 label="Civitai API Key", 
                 type="password", 
                 placeholder="Enter your Civitai API key here (optional)"
            )
    return components


def create_lora_settings_ui(prefix: str):
    components = {}
    
    lora_rows, lora_sources, lora_ids, lora_scales, lora_uploads = [], [], [], [], []

    with gr.Accordion("LoRA Settings", open=False) as lora_accordion:
        components[f'lora_accordion_{prefix}'] = lora_accordion
        gr.Markdown("💡 **Tip:** When downloading from Civitai, please use the **Version ID**, not the Model ID. You can find the Version ID in the URL (e.g., `civitai.com/models/123?modelVersionId=456`) or under the model's download button.")
        components[f'lora_count_state_{prefix}'] = gr.State(1)

        for i in range(MAX_LORAS):
            with gr.Row(visible=i==0) as row:
                source = gr.Dropdown(label=f"LoRA Source {i+1}", choices=LORA_SOURCE_CHOICES, value=LORA_SOURCE_CHOICES[0], scale=1)
                lora_id = gr.Textbox(label=f"Civitai Version ID / File", placeholder="Civitai Version ID or Filename", scale=2, type="text")
                scale = gr.Slider(label=f"Scale", minimum=-2.0, maximum=2.0, step=0.05, value=0.8, scale=1)
                upload = gr.UploadButton(label="Upload", file_types=[".safetensors"], scale=1)
            
            lora_rows.append(row)
            lora_sources.append(source)
            lora_ids.append(lora_id)
            lora_scales.append(scale)
            lora_uploads.append(upload)
            
        with gr.Row():
            components[f'add_lora_button_{prefix}'] = gr.Button("Add LoRA", variant="secondary")
            components[f'delete_lora_button_{prefix}'] = gr.Button("Remove LoRA", variant="secondary", visible=False)

    components[f'lora_rows_{prefix}'] = lora_rows
    components[f'lora_sources_{prefix}'] = lora_sources
    components[f'lora_ids_{prefix}'] = lora_ids
    components[f'lora_scales_{prefix}'] = lora_scales
    components[f'lora_uploads_{prefix}'] = lora_uploads
    
    all_lora_components_flat = []
    for i in range(MAX_LORAS):
        all_lora_components_flat.extend([lora_sources[i], lora_ids[i], lora_scales[i], lora_uploads[i]])
    components[f'all_lora_components_flat_{prefix}'] = all_lora_components_flat
    
    return components

def create_controlnet_ui(prefix: str, max_units=MAX_CONTROLNETS):
    components = {}
    key = lambda name: f"{name}_{prefix}"

    with gr.Accordion("ControlNet Settings", open=False) as accordion:
        components[key('controlnet_accordion')] = accordion
        
        cn_rows, images, series, types, strengths, filepaths = [], [], [], [], [], []
        components.update({
            key('controlnet_rows'): cn_rows,
            key('controlnet_images'): images,
            key('controlnet_series'): series,
            key('controlnet_types'): types,
            key('controlnet_strengths'): strengths,
            key('controlnet_filepaths'): filepaths
        })
        
        for i in range(max_units):
            with gr.Row(visible=(i < 1)) as row:
                with gr.Column(scale=1):
                    images.append(gr.Image(label=f"Control Image {i+1}", type="pil", sources=["upload"], height=256))
                with gr.Column(scale=2):
                    types.append(gr.Dropdown(label="Type", choices=[], interactive=True))
                    series.append(gr.Dropdown(label="Series", choices=[], interactive=True))
                    strengths.append(gr.Slider(label="Strength", minimum=0.0, maximum=2.0, step=0.05, value=1.0, interactive=True))
                    filepaths.append(gr.State(None))
                cn_rows.append(row)

        with gr.Row():
            components[key('add_controlnet_button')] = gr.Button("✚ Add ControlNet")
            components[key('delete_controlnet_button')] = gr.Button("➖ Delete ControlNet", visible=False)
        components[key('controlnet_count_state')] = gr.State(1)
        
        all_cn_components_flat = []
        for i in range(max_units):
            all_cn_components_flat.extend([
                images[i], types[i], series[i], strengths[i], filepaths[i]
            ])
        components[key('all_controlnet_components_flat')] = all_cn_components_flat

    return components

def create_diffsynth_controlnet_ui(prefix: str, max_units=MAX_CONTROLNETS):
    components = {}
    key = lambda name: f"{name}_{prefix}"

    with gr.Accordion("DiffSynth ControlNet Settings", open=False) as accordion:
        components[key('diffsynth_controlnet_accordion')] = accordion
        
        cn_rows, images, series, types, strengths, filepaths = [], [], [], [], [], []
        components.update({
            key('diffsynth_controlnet_rows'): cn_rows,
            key('diffsynth_controlnet_images'): images,
            key('diffsynth_controlnet_series'): series,
            key('diffsynth_controlnet_types'): types,
            key('diffsynth_controlnet_strengths'): strengths,
            key('diffsynth_controlnet_filepaths'): filepaths
        })
        
        for i in range(max_units):
            with gr.Row(visible=(i < 1)) as row:
                with gr.Column(scale=1):
                    images.append(gr.Image(label=f"Control Image {i+1}", type="pil", sources=["upload"], height=256))
                with gr.Column(scale=2):
                    types.append(gr.Dropdown(label="Type", choices=[], interactive=True))
                    series.append(gr.Dropdown(label="Series", choices=[], interactive=True))
                    strengths.append(gr.Slider(label="Strength", minimum=0.0, maximum=2.0, step=0.05, value=1.0, interactive=True))
                    filepaths.append(gr.State(None))
                cn_rows.append(row)

        with gr.Row():
            components[key('add_diffsynth_controlnet_button')] = gr.Button("✚ Add DiffSynth ControlNet")
            components[key('delete_diffsynth_controlnet_button')] = gr.Button("➖ Delete DiffSynth ControlNet", visible=False)
        components[key('diffsynth_controlnet_count_state')] = gr.State(1)
        
        all_cn_components_flat = []
        for i in range(max_units):
            all_cn_components_flat.extend([
                images[i], types[i], series[i], strengths[i], filepaths[i]
            ])
        components[key('all_diffsynth_controlnet_components_flat')] = all_cn_components_flat

    return components

def create_embedding_ui(prefix: str):
    components = {}
    key = lambda name: f"{name}_{prefix}"

    with gr.Accordion("Embedding Settings", open=False, visible=True) as accordion:
        components[key('embedding_accordion')] = accordion
        gr.Markdown("💡 **Tip:** Embeddings are automatically added to your prompt using `embedding:filename` syntax. When downloading from Civitai, please use the **Version ID**, not the Model ID. You can find the Version ID in the URL (e.g., `civitai.com/models/123?modelVersionId=456`) or under the model's download button. For instance, using the Version ID `456` from the example above would automatically append `embedding:civitai_456` to your positive prompt.")
        
        embedding_rows, sources, ids, files, upload_buttons = [], [], [], [], []
        components.update({
            key('embedding_rows'): embedding_rows, 
            key('embeddings_sources'): sources, 
            key('embeddings_ids'): ids, 
            key('embeddings_files'): files,
            key('embeddings_uploads'): upload_buttons
        })
        
        for i in range(MAX_EMBEDDINGS):
            with gr.Row(visible=(i < 1)) as row:
                sources.append(gr.Dropdown(label=f"Embedding Source {i+1}", choices=LORA_SOURCE_CHOICES, value="Civitai", scale=1, interactive=True))
                ids.append(gr.Textbox(label="Civitai Version ID / File", placeholder="Civitai Version ID or Filename", scale=3, interactive=True, type="text"))
                upload_btn = gr.UploadButton("Upload", file_types=[".safetensors"], scale=1)
                files.append(gr.State(None))
                upload_buttons.append(upload_btn)
                embedding_rows.append(row)

        with gr.Row():
            components[key('add_embedding_button')] = gr.Button("✚ Add Embedding")
            components[key('delete_embedding_button')] = gr.Button("➖ Delete Embedding", visible=False)
        components[key('embedding_count_state')] = gr.State(1)

        all_embedding_components_flat = []
        for i in range(MAX_EMBEDDINGS):
            all_embedding_components_flat.extend([sources[i], ids[i], files[i]])
        components[key('all_embedding_components_flat')] = all_embedding_components_flat

    return components

def create_conditioning_ui(prefix: str):
    components = {}
    key = lambda name: f"{name}_{prefix}"
    
    with gr.Accordion("Conditioning Settings", open=False) as accordion:
        components[key('conditioning_accordion')] = accordion
        gr.Markdown("💡 **Tip:** Define rectangular areas and assign specific prompts to them. Coordinates (X, Y) start from the top-left corner.")
        
        cond_rows, prompts, widths, heights, xs, ys, strengths = [], [], [], [], [], [], []
        components.update({
            key('conditioning_rows'): cond_rows,
            key('conditioning_prompts'): prompts,
            key('conditioning_widths'): widths,
            key('conditioning_heights'): heights,
            key('conditioning_xs'): xs,
            key('conditioning_ys'): ys,
            key('conditioning_strengths'): strengths
        })
        
        for i in range(MAX_CONDITIONINGS):
            with gr.Column(visible=(i < 1)) as row_wrapper:
                prompts.append(gr.Textbox(label=f"Area Prompt {i+1}", lines=2, interactive=True))
                with gr.Row():
                    xs.append(gr.Number(label="X", value=0, interactive=True, step=8, scale=1))
                    ys.append(gr.Number(label="Y", value=0, interactive=True, step=8, scale=1))
                    widths.append(gr.Number(label="Width", value=512, interactive=True, step=8, scale=1))
                    heights.append(gr.Number(label="Height", value=512, interactive=True, step=8, scale=1))
                    strengths.append(gr.Slider(label="Strength", minimum=0.1, maximum=2.0, step=0.05, value=1.0, interactive=True, scale=2))
                cond_rows.append(row_wrapper)

        with gr.Row():
            components[key('add_conditioning_button')] = gr.Button("✚ Add Area")
            components[key('delete_conditioning_button')] = gr.Button("➖ Delete Area", visible=False)
        components[key('conditioning_count_state')] = gr.State(1)

        all_cond_components_flat = prompts + widths + heights + xs + ys + strengths
        components[key('all_conditioning_components_flat')] = all_cond_components_flat
    
    return components
    
def create_vae_override_ui(prefix: str):
    components = {}
    key = lambda name: f"{name}_{prefix}"
    source_choices = ["None"] + LORA_SOURCE_CHOICES

    with gr.Accordion("VAE Settings (Override)", open=False) as accordion:
        components[key('vae_accordion')] = accordion
        gr.Markdown("💡 **Tip:** When downloading from Civitai, please use the **Version ID**, not the Model ID. You can find the Version ID in the URL (e.g., `civitai.com/models/123?modelVersionId=456`) or under the model's download button.")
        with gr.Row():
            components[key('vae_source')] = gr.Dropdown(
                label="VAE Source", 
                choices=source_choices, 
                value="None", 
                scale=1, 
                interactive=True
            )
            components[key('vae_id')] = gr.Textbox(
                label="Civitai Version ID / File", 
                placeholder="Civitai Version ID or Filename", 
                scale=3, 
                interactive=True,
                type="text"
            )
            upload_btn = gr.UploadButton(
                "Upload", 
                file_types=[".safetensors"], 
                scale=1
            )
            components[key('vae_upload_button')] = upload_btn
            components[key('vae_file')] = gr.State(None)
            
    return components