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, MAX_IPADAPTERS, RESOLUTION_MAP, ARCHITECTURES_CONFIG,
    MODEL_MAP_CHECKPOINT, MODEL_TYPE_MAP, FEATURES_CONFIG, ARCH_CATEGORIES_MAP
)
import yaml
import os
from functools import lru_cache

default_model_name = list(MODEL_MAP_CHECKPOINT.keys())[0] if MODEL_MAP_CHECKPOINT else None
default_m_type = MODEL_TYPE_MAP.get(default_model_name, "SDXL") if default_model_name else "SDXL"
default_architectures_dict = ARCHITECTURES_CONFIG.get('architectures', {})
default_arch_model_type = default_architectures_dict.get(default_m_type, {}).get("model_type", default_m_type.lower().replace(" ", "").replace(".", ""))
default_arch_features = FEATURES_CONFIG.get(default_arch_model_type, FEATURES_CONFIG.get('default', {}))
default_enabled_chains = default_arch_features.get('enabled_chains', [])


@lru_cache(maxsize=1)
def get_ipadapter_config_from_yaml():
    try:
        _PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
        _IPADAPTER_LIST_PATH = os.path.join(_PROJECT_ROOT, 'yaml', 'ipadapter.yaml')
        with open(_IPADAPTER_LIST_PATH, 'r', encoding='utf-8') as f:
            config = yaml.safe_load(f)
        return config
    except Exception as e:
        print(f"Warning: Could not load ipadapter.yaml for UI components: {e}")
        return {}

def get_ipadapter_presets(arch="SDXL"):
    config = get_ipadapter_config_from_yaml()
    presets = []
    if config:
        std_presets = config.get("IPAdapter_presets", {}).get(arch, [])
        face_presets = config.get("IPAdapter_FaceID_presets", {}).get(arch, [])
        if std_presets:
            presets.extend(std_presets)
        if face_presets:
            presets.extend(face_presets)
    return presets if presets else ["STANDARD (medium strength)"]

def create_model_architecture_filter_ui(prefix):
    components = {}
    ordered_architectures = ARCHITECTURES_CONFIG.get("architecture_order", [])
    choices = ["ALL"] + ordered_architectures

    components[f'model_arch_{prefix}'] = gr.Radio(
        label="Model Architecture",
        choices=choices,
        value="ALL",
        interactive=True
    )
    return components

def create_category_filter_ui(prefix):
    valid_cats = list(set(cat for cats in ARCH_CATEGORIES_MAP.values() for cat in cats))
    cat_choices = ["ALL"] + sorted(valid_cats)

    components = {}
    components[f'model_cat_{prefix}'] = gr.Dropdown(
        label="Filter Models",
        choices=cat_choices,
        value="ALL",
        interactive=True,
        scale=1
    )
    return components

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.get('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=SAMPLER_CHOICES[0])
        components[f'scheduler_{prefix}'] = gr.Dropdown(label="Scheduler", choices=SCHEDULER_CHOICES, value='normal' if 'normal' in SCHEDULER_CHOICES else SCHEDULER_CHOICES[0])
    with gr.Row():
        components[f'steps_{prefix}'] = gr.Slider(label="Steps", minimum=1, maximum=100, step=1, value=28)
        components[f'cfg_{prefix}'] = gr.Slider(label="CFG Scale", minimum=1.0, maximum=20.0, step=0.1, value=7.5)
    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'clip_skip_{prefix}'] = gr.Slider(label="Clip Skip", minimum=1, maximum=2, step=1, value=1, visible=False, interactive=True)
        components[f'guidance_{prefix}'] = gr.Slider(label="Guidance (FLUX)", minimum=1.0, maximum=10.0, step=0.1, value=3.5, visible=False, interactive=True)
        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.")
            
    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, visible=('lora' in default_enabled_chains)) 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, visible=('controlnet' in default_enabled_chains)) 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, allow_custom_value=True))
                    series.append(gr.Dropdown(label="Series", choices=[], interactive=True, allow_custom_value=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, visible=('controlnet_model_patch' in default_enabled_chains)) 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, allow_custom_value=True))
                    series.append(gr.Dropdown(label="Series", choices=[], interactive=True, allow_custom_value=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_ipadapter_ui(prefix: str, max_units=MAX_IPADAPTERS):
    components = {}
    key = lambda name: f"{name}_{prefix}"

    sdxl_presets = get_ipadapter_presets("SDXL")
    default_preset = sdxl_presets[0] if sdxl_presets else None

    with gr.Accordion("IPAdapter Settings", open=False, visible=('ipadapter' in default_enabled_chains)) as accordion:
        components[key('ipadapter_accordion')] = accordion
        gr.Markdown("Powered by [cubiq/ComfyUI_IPAdapter_plus](https://github.com/cubiq/ComfyUI_IPAdapter_plus).")
        
        with gr.Row():
            components[key('ipadapter_final_preset')] = gr.Dropdown(
                label="Preset (for all images)", 
                choices=sdxl_presets, 
                value=default_preset, 
                interactive=True,
                allow_custom_value=True
            )
            components[key('ipadapter_embeds_scaling')] = gr.Dropdown(
                label="Embeds Scaling", 
                choices=['V only', 'K+V', 'K+V w/ C penalty', 'K+mean(V) w/ C penalty'],
                value='V only',
                interactive=True
            )
        
        with gr.Row():
            components[key('ipadapter_combine_method')] = gr.Dropdown(
                label="Combine Method",
                choices=["concat", "add", "subtract", "average", "norm average", "max", "min"],
                value="concat",
                interactive=True
            )
            components[key('ipadapter_final_weight')] = gr.Slider(label="Final Weight", minimum=0.0, maximum=2.0, step=0.05, value=1.0, interactive=True)
            components[key('ipadapter_final_lora_strength')] = gr.Slider(label="Final LoRA Strength", minimum=0.0, maximum=2.0, step=0.05, value=0.6, interactive=True, visible=False)

        gr.Markdown("---")
        
        ipa_rows, images, weights, lora_strengths = [], [], [], []
        components.update({
            key('ipadapter_rows'): ipa_rows,
            key('ipadapter_images'): images,
            key('ipadapter_weights'): weights,
            key('ipadapter_lora_strengths'): lora_strengths
        })
        
        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"IPAdapter Image {i+1}", type="pil", sources=["upload"], height=256))
                with gr.Column(scale=2):
                    weights.append(gr.Slider(label="Weight", minimum=0.0, maximum=2.0, step=0.05, value=1.0, interactive=True))
                    lora_strengths.append(gr.Slider(label="LoRA Strength", minimum=0.0, maximum=2.0, step=0.05, value=0.6, interactive=True, visible=False))
                ipa_rows.append(row)

        with gr.Row():
            components[key('add_ipadapter_button')] = gr.Button("✚ Add IPAdapter")
            components[key('delete_ipadapter_button')] = gr.Button("➖ Delete IPAdapter", visible=False)
        components[key('ipadapter_count_state')] = gr.State(1)
        
        all_ipa_components_flat = images + weights + lora_strengths
        all_ipa_components_flat += [
            components[key('ipadapter_final_preset')],
            components[key('ipadapter_final_weight')],
            components[key('ipadapter_final_lora_strength')],
            components[key('ipadapter_embeds_scaling')],
            components[key('ipadapter_combine_method')],
        ]
        components[key('all_ipadapter_components_flat')] = all_ipa_components_flat

    return components

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

    with gr.Accordion("IPAdapter Settings (FLUX.1)", open=False, visible=('flux1_ipadapter' in default_enabled_chains)) as accordion:
        components[key('flux1_ipadapter_accordion')] = accordion
        
        ipa_rows, images, weights, start_percents, end_percents = [], [], [], [], []
        components.update({
            key('flux1_ipadapter_rows'): ipa_rows,
            key('flux1_ipadapter_images'): images,
            key('flux1_ipadapter_weights'): weights,
            key('flux1_ipadapter_start_percents'): start_percents,
            key('flux1_ipadapter_end_percents'): end_percents,
        })
        
        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"IPAdapter Image {i+1}", type="pil", sources=["upload"], height=256))
                with gr.Column(scale=2):
                    weights.append(gr.Slider(label="Weight", minimum=0.0, maximum=2.0, step=0.05, value=0.6, interactive=True))
                    with gr.Row():
                        start_percents.append(gr.Slider(label="Start At", minimum=0.0, maximum=1.0, step=0.01, value=0.0, interactive=True))
                        end_percents.append(gr.Slider(label="End At", minimum=0.0, maximum=1.0, step=0.01, value=0.6, interactive=True))
                ipa_rows.append(row)

        with gr.Row():
            components[key('add_flux1_ipadapter_button')] = gr.Button("✚ Add IPAdapter (FLUX)")
            components[key('delete_flux1_ipadapter_button')] = gr.Button("➖ Delete IPAdapter (FLUX)", visible=False)
        components[key('flux1_ipadapter_count_state')] = gr.State(1)
        
        all_flux1_ipa_components_flat = images + weights + start_percents + end_percents
        components[key('all_flux1_ipadapter_components_flat')] = all_flux1_ipa_components_flat

    return components

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

    with gr.Accordion("IPAdapter Settings (SD3)", open=False, visible=('sd3_ipadapter' in default_enabled_chains)) as accordion:
        components[key('sd3_ipadapter_accordion')] = accordion
        
        ipa_rows, images, weights, start_percents, end_percents = [], [], [], [], []
        components.update({
            key('sd3_ipadapter_rows'): ipa_rows,
            key('sd3_ipadapter_images'): images,
            key('sd3_ipadapter_weights'): weights,
            key('sd3_ipadapter_start_percents'): start_percents,
            key('sd3_ipadapter_end_percents'): end_percents,
        })
        
        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"IPAdapter Image {i+1}", type="pil", sources=["upload"], height=256))
                with gr.Column(scale=2):
                    weights.append(gr.Slider(label="Weight", minimum=0.0, maximum=2.0, step=0.05, value=0.5, interactive=True))
                    with gr.Row():
                        start_percents.append(gr.Slider(label="Start At", minimum=0.0, maximum=1.0, step=0.01, value=0.0, interactive=True))
                        end_percents.append(gr.Slider(label="End At", minimum=0.0, maximum=1.0, step=0.01, value=1.0, interactive=True))
                ipa_rows.append(row)

        with gr.Row():
            components[key('add_sd3_ipadapter_button')] = gr.Button("✚ Add IPAdapter (SD3)")
            components[key('delete_sd3_ipadapter_button')] = gr.Button("➖ Delete IPAdapter (SD3)", visible=False)
        components[key('sd3_ipadapter_count_state')] = gr.State(1)
        
        all_sd3_ipa_components_flat = images + weights + start_percents + end_percents
        components[key('all_sd3_ipadapter_components_flat')] = all_sd3_ipa_components_flat

    return components

def create_style_ui(prefix: str):
    components = {}
    key = lambda name: f"{name}_{prefix}"
    
    with gr.Accordion("Style Settings (FLUX.1)", open=False, visible=('style' in default_enabled_chains)) as accordion:
        components[key('style_accordion')] = accordion
        
        style_rows, images, strengths = [], [], []
        components.update({
            key('style_rows'): style_rows,
            key('style_images'): images,
            key('style_strengths'): strengths
        })
        
        for i in range(5):
            with gr.Row(visible=(i < 1)) as row:
                with gr.Column(scale=1):
                    images.append(gr.Image(label=f"Style Image {i+1}", type="pil", sources=["upload"], height=256))
                with gr.Column(scale=2):
                    strengths.append(gr.Slider(label="Strength", minimum=0.0, maximum=2.0, step=0.05, value=1.0, interactive=True))
                style_rows.append(row)

        with gr.Row():
            components[key('add_style_button')] = gr.Button("✚ Add Style (FLUX)")
            components[key('delete_style_button')] = gr.Button("➖ Delete Style (FLUX)", visible=False)
        components[key('style_count_state')] = gr.State(1)
        
        all_style_components_flat = images + strengths
        components[key('all_style_components_flat')] = all_style_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=('embedding' in default_enabled_chains)) 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, visible=('conditioning' in default_enabled_chains)) 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

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

    with gr.Accordion("Reference Edit Settings", open=False, visible=('reference_latent' in default_enabled_chains)) as ref_accordion:
        components[key('reference_latent_accordion')] = ref_accordion
        gr.Markdown("💡 **Tip:** For multimodal models (like FLUX.2 or OmniGen), this feature enables powerful editing and combining capabilities. In txt2img mode, adding a single reference image performs an **Image Edit**, while adding multiple images performs an **Image Combine**.")
        
        ref_image_groups = []
        ref_image_inputs = []
        with gr.Row():
            for i in range(max_units):
                with gr.Column(visible=(i < 1), min_width=160) as img_col:
                    img_comp = gr.Image(type="pil", label=f"Ref. {i+1}", sources=["upload"], height=150)
                    ref_image_groups.append(img_col)
                    ref_image_inputs.append(img_comp)
        
        components[key('reference_latent_rows')] = ref_image_groups
        components[key('reference_latent_images')] = ref_image_inputs

        with gr.Row():
            components[key('add_reference_latent_button')] = gr.Button("✚ Add Reference Image")
            components[key('delete_reference_latent_button')] = gr.Button("➖ Delete Reference Image", visible=False)
        components[key('reference_latent_count_state')] = gr.State(1)
        
        components[key('all_reference_latent_components_flat')] = ref_image_inputs

    return components