import gradio as gr
from modules import shared, ui_common, generation_parameters_copypaste
from modules.logger import log
from modules.caption import openclip


default_task = "Normal Caption"


def caption_wrapper(tab, image,
                    vlm_question, vlm_system, vlm_prompt, vlm_model, vlm_prefill, vlm_thinking_mode, vlm_custom_args,
                    analyze_question, analyze_system, analyze_prompt, analyze_model,
                    clip_model, blip_model, clip_mode,
                    wd_model, wd_general_threshold, wd_character_threshold, wd_include_rating, wd_exclude_tags, wd_max_tags, wd_sort_alpha, wd_use_spaces, wd_escape,
                   ):
    """Wrapper for vqa.caption, vqa.analysis, openclip.caption_image, tagger.tag."""
    if tab <= 0:
        log.debug('Caption: mode="VLM Caption"')
        from modules.caption import vqa
        answer = vqa.caption(vlm_question, vlm_system, vlm_prompt, image, vlm_model, vlm_prefill, vlm_thinking_mode, vlm_custom_args)
        annotated_image = vqa.get_last_annotated_image()
        if annotated_image is not None:
            return answer, gr.update(value=annotated_image, visible=True)
        return answer, gr.update(visible=False)
    elif tab == 1:
        log.debug('Caption: mode="VLM Analyze"')
        from modules.caption import vqa
        answer = vqa.analyze(analyze_question, analyze_system, analyze_prompt, image, analyze_model, vlm_thinking_mode, vlm_custom_args)
        return answer, gr.update(visible=False)
    elif tab == 2:
        log.debug('Caption: mode="OpenCLIP"')
        caption = openclip.caption_image(image, clip_model, blip_model, clip_mode)
        return caption, gr.update(visible=False)
    elif tab == 3:
        log.debug('Caption: mode="Tagger"')
        tags = tagger_tag_wrapper(image, wd_model, wd_general_threshold, wd_character_threshold, wd_include_rating, wd_exclude_tags, wd_max_tags, wd_sort_alpha, wd_use_spaces, wd_escape)
        return tags, gr.update(visible=False)
    else:
        log.error(f'Caption: mode={tab} unknown')
        return 'Unknown caption mode', gr.update(visible=False)


def update_vlm_prompts_for_model(model_name):
    """Update the task dropdown choices based on selected model."""
    from modules.caption import vqa
    prompts = vqa.get_prompts_for_model(model_name)
    value = default_task if default_task in prompts else (prompts[0] if prompts else default_task)
    return gr.update(choices=prompts, value=value)


def update_vlm_prompt_placeholder(question):
    """Update the prompt field placeholder based on selected task."""
    from modules.caption import vqa
    placeholder = vqa.get_prompt_placeholder(question)
    return gr.update(placeholder=placeholder)


def update_vlm_params(*args):
    vlm_max_tokens, vlm_num_beams, vlm_temperature, vlm_do_sample, vlm_top_k, vlm_top_p, vlm_keep_prefill, vlm_keep_thinking, vlm_thinking_mode, vlm_custom_args = args
    shared.opts.caption_vlm_max_length = int(vlm_max_tokens)
    shared.opts.caption_vlm_num_beams = int(vlm_num_beams)
    shared.opts.caption_vlm_temperature = float(vlm_temperature)
    shared.opts.caption_vlm_do_sample = bool(vlm_do_sample)
    shared.opts.caption_vlm_top_k = int(vlm_top_k)
    shared.opts.caption_vlm_top_p = float(vlm_top_p)
    shared.opts.caption_vlm_keep_prefill = bool(vlm_keep_prefill)
    shared.opts.caption_vlm_keep_thinking = bool(vlm_keep_thinking)
    shared.opts.caption_vlm_thinking_mode = bool(vlm_thinking_mode)
    shared.opts.caption_vlm_custom_args = vlm_custom_args
    shared.opts.save()


def tagger_tag_wrapper(image, model_name, general_threshold, character_threshold, include_rating, exclude_tags, max_tags, sort_alpha, use_spaces, escape_brackets):
    """Wrapper for tagger.tag that maps UI inputs to function parameters."""
    from modules.caption import tagger
    return tagger.tag(
        image=image,
        model_name=model_name,
        general_threshold=general_threshold,
        character_threshold=character_threshold,
        include_rating=include_rating,
        exclude_tags=exclude_tags,
        max_tags=int(max_tags),
        sort_alpha=sort_alpha,
        use_spaces=use_spaces,
        escape_brackets=escape_brackets,
    )


def tagger_batch_wrapper(model_name, batch_files, batch_folder, batch_str, save_output, save_append, recursive, general_threshold, character_threshold, include_rating, exclude_tags, max_tags, sort_alpha, use_spaces, escape_brackets):
    """Wrapper for tagger.batch that maps UI inputs to function parameters."""
    from modules.caption import tagger
    return tagger.batch(
        model_name=model_name,
        batch_files=batch_files,
        batch_folder=batch_folder,
        batch_str=batch_str,
        save_output=save_output,
        save_append=save_append,
        recursive=recursive,
        general_threshold=general_threshold,
        character_threshold=character_threshold,
        include_rating=include_rating,
        exclude_tags=exclude_tags,
        max_tags=int(max_tags),
        sort_alpha=sort_alpha,
        use_spaces=use_spaces,
        escape_brackets=escape_brackets,
    )


def update_tagger_ui(model_name):
    """Update UI controls based on selected tagger model.

    When DeepBooru is selected, character_threshold is disabled since DeepBooru
    doesn't support separate character threshold.
    """
    from modules.caption import tagger
    is_db = tagger.is_deepbooru(model_name)
    return [
        gr.update(interactive=not is_db),  # character_threshold
        gr.update(),  # include_rating - now supported by both taggers
    ]


def update_tagger_params(model_name, general_threshold, character_threshold, include_rating, max_tags, sort_alpha, use_spaces, escape_brackets, exclude_tags, show_scores):
    """Save all tagger parameters to shared.opts when UI controls change."""
    shared.opts.waifudiffusion_model = model_name
    shared.opts.tagger_threshold = float(general_threshold)
    shared.opts.waifudiffusion_character_threshold = float(character_threshold)
    shared.opts.tagger_include_rating = bool(include_rating)
    shared.opts.tagger_max_tags = int(max_tags)
    shared.opts.tagger_sort_alpha = bool(sort_alpha)
    shared.opts.tagger_use_spaces = bool(use_spaces)
    shared.opts.tagger_escape_brackets = bool(escape_brackets)
    shared.opts.tagger_exclude_tags = str(exclude_tags)
    shared.opts.tagger_show_scores = bool(show_scores)
    shared.opts.save()


def update_clip_params(*args):
    clip_max_length, clip_chunk_size, clip_min_flavors, clip_max_flavors, clip_flavor_count, clip_num_beams = args
    shared.opts.caption_openclip_max_length = int(clip_max_length)
    shared.opts.caption_openclip_min_flavors = int(clip_min_flavors)
    shared.opts.caption_openclip_max_flavors = int(clip_max_flavors)
    shared.opts.caption_openclip_num_beams = int(clip_num_beams)
    shared.opts.caption_openclip_flavor_count = int(clip_flavor_count)
    shared.opts.caption_openclip_chunk_size = int(clip_chunk_size)
    shared.opts.save()
    openclip.update_caption_params()


def update_clip_model_params(clip_model, blip_model, clip_mode):
    """Save CLiP model settings to shared.opts when UI controls change."""
    shared.opts.caption_openclip_model = str(clip_model)
    shared.opts.caption_openclip_blip_model = str(blip_model)
    shared.opts.caption_openclip_mode = str(clip_mode)
    shared.opts.save()


def update_vlm_model_params(vlm_model, vlm_system):
    """Save VLM model settings to shared.opts when UI controls change."""
    shared.opts.caption_vlm_model = str(vlm_model)
    shared.opts.caption_vlm_system = str(vlm_system)
    shared.opts.save()


def update_default_caption_type(caption_type):
    """Save the default caption type to shared.opts."""
    shared.opts.caption_default_type = str(caption_type)
    shared.opts.save()


def create_ui():
    log.debug('UI initialize: tab=caption')
    with gr.Row(equal_height=False, variant='compact', elem_classes="caption", elem_id="caption_tab"):
        with gr.Column(variant='compact', elem_id='caption_input'):
            with gr.Row():
                image = gr.Image(type='pil', label="Image", height=512, visible=True, image_mode='RGB', elem_id='caption_image')
            with gr.Tabs(elem_id="mode_caption"):

                with gr.Tab("VLM Caption", elem_id="tab_vlm_caption"):
                    from modules.caption import vqa
                    current_vlm_model = shared.opts.caption_vlm_model or vqa.vlm_default
                    initial_prompts = vqa.get_prompts_for_model(current_vlm_model)
                    with gr.Row():
                        vlm_system = gr.Textbox(label="System Prompt", value=vqa.vlm_system, lines=1, elem_id='vlm_system')
                    with gr.Row():
                        vlm_question = gr.Dropdown(label="Task", allow_custom_value=False, choices=initial_prompts, value=default_task, elem_id='vlm_question')
                    with gr.Row():
                        vlm_prompt = gr.Textbox(label="Prompt", placeholder=vqa.get_prompt_placeholder(initial_prompts[0]), lines=2, elem_id='vlm_prompt')
                    with gr.Row(elem_id='caption_buttons_query'):
                        vlm_model = gr.Dropdown(list(vqa.vlm_models), value=current_vlm_model, label='VLM Model', elem_id='vlm_model')
                    with gr.Row():
                        vlm_load_btn = gr.Button(value='Load', elem_id='vlm_load', variant='secondary')
                        vlm_unload_btn = gr.Button(value='Unload', elem_id='vlm_unload', variant='secondary')
                    with gr.Accordion(label='Caption: Advanced Options', open=False, visible=True):
                        with gr.Row():
                            vlm_max_tokens = gr.Slider(label='VLM Max Tokens', value=shared.opts.caption_vlm_max_length, minimum=16, maximum=4096, step=1, elem_id='vlm_max_tokens')
                            vlm_num_beams = gr.Slider(label='VLM Num Beams', value=shared.opts.caption_vlm_num_beams, minimum=1, maximum=16, step=1, elem_id='vlm_num_beams')
                            vlm_temperature = gr.Slider(label='VLM Temperature', value=shared.opts.caption_vlm_temperature, minimum=0.0, maximum=1.0, step=0.01, elem_id='vlm_temperature')
                        with gr.Row():
                            vlm_top_k = gr.Slider(label='Top-K', value=shared.opts.caption_vlm_top_k, minimum=0, maximum=99, step=1, elem_id='vlm_top_k')
                            vlm_top_p = gr.Slider(label='Top-P', value=shared.opts.caption_vlm_top_p, minimum=0.0, maximum=1.0, step=0.01, elem_id='vlm_top_p')
                        with gr.Row():
                            vlm_do_sample = gr.Checkbox(label='Use Samplers', value=shared.opts.caption_vlm_do_sample, elem_id='vlm_do_sample')
                            vlm_thinking_mode = gr.Checkbox(label='Thinking Mode', value=shared.opts.caption_vlm_thinking_mode, elem_id='vlm_thinking_mode')
                        with gr.Row():
                            vlm_custom_args = gr.Textbox(label='Custom Args', value=shared.opts.caption_vlm_custom_args, placeholder='key=value pairs separated by ; or newlines', lines=2, elem_id='vlm_custom_args')
                        with gr.Row():
                            vlm_keep_thinking = gr.Checkbox(label='Keep Thinking Trace', value=shared.opts.caption_vlm_keep_thinking, elem_id='vlm_keep_thinking')
                            vlm_keep_prefill = gr.Checkbox(label='Keep Prefill', value=shared.opts.caption_vlm_keep_prefill, elem_id='vlm_keep_prefill')
                        with gr.Row():
                            vlm_prefill = gr.Textbox(label='Prefill Text', value='', lines=1, elem_id='vlm_prefill', placeholder='Optional prefill text for model to continue from')
                        vlm_max_tokens.change(fn=update_vlm_params, inputs=[vlm_max_tokens, vlm_num_beams, vlm_temperature, vlm_do_sample, vlm_top_k, vlm_top_p, vlm_keep_prefill, vlm_keep_thinking, vlm_thinking_mode, vlm_custom_args], outputs=[])
                        vlm_num_beams.change(fn=update_vlm_params, inputs=[vlm_max_tokens, vlm_num_beams, vlm_temperature, vlm_do_sample, vlm_top_k, vlm_top_p, vlm_keep_prefill, vlm_keep_thinking, vlm_thinking_mode, vlm_custom_args], outputs=[])
                        vlm_temperature.change(fn=update_vlm_params, inputs=[vlm_max_tokens, vlm_num_beams, vlm_temperature, vlm_do_sample, vlm_top_k, vlm_top_p, vlm_keep_prefill, vlm_keep_thinking, vlm_thinking_mode, vlm_custom_args], outputs=[])
                        vlm_do_sample.change(fn=update_vlm_params, inputs=[vlm_max_tokens, vlm_num_beams, vlm_temperature, vlm_do_sample, vlm_top_k, vlm_top_p, vlm_keep_prefill, vlm_keep_thinking, vlm_thinking_mode, vlm_custom_args], outputs=[])
                        vlm_top_k.change(fn=update_vlm_params, inputs=[vlm_max_tokens, vlm_num_beams, vlm_temperature, vlm_do_sample, vlm_top_k, vlm_top_p, vlm_keep_prefill, vlm_keep_thinking, vlm_thinking_mode, vlm_custom_args], outputs=[])
                        vlm_top_p.change(fn=update_vlm_params, inputs=[vlm_max_tokens, vlm_num_beams, vlm_temperature, vlm_do_sample, vlm_top_k, vlm_top_p, vlm_keep_prefill, vlm_keep_thinking, vlm_thinking_mode, vlm_custom_args], outputs=[])
                        vlm_keep_prefill.change(fn=update_vlm_params, inputs=[vlm_max_tokens, vlm_num_beams, vlm_temperature, vlm_do_sample, vlm_top_k, vlm_top_p, vlm_keep_prefill, vlm_keep_thinking, vlm_thinking_mode, vlm_custom_args], outputs=[])
                        vlm_keep_thinking.change(fn=update_vlm_params, inputs=[vlm_max_tokens, vlm_num_beams, vlm_temperature, vlm_do_sample, vlm_top_k, vlm_top_p, vlm_keep_prefill, vlm_keep_thinking, vlm_thinking_mode, vlm_custom_args], outputs=[])
                        vlm_thinking_mode.change(fn=update_vlm_params, inputs=[vlm_max_tokens, vlm_num_beams, vlm_temperature, vlm_do_sample, vlm_top_k, vlm_top_p, vlm_keep_prefill, vlm_keep_thinking, vlm_thinking_mode, vlm_custom_args], outputs=[])
                        vlm_custom_args.change(fn=update_vlm_params, inputs=[vlm_max_tokens, vlm_num_beams, vlm_temperature, vlm_do_sample, vlm_top_k, vlm_top_p, vlm_keep_prefill, vlm_keep_thinking, vlm_thinking_mode, vlm_custom_args], outputs=[])
                    with gr.Accordion(label='Caption: Batch', open=False, visible=True):
                        with gr.Row():
                            vlm_batch_files = gr.File(label="Files", show_label=True, file_count='multiple', file_types=['image'], interactive=True, height=100, elem_id='vlm_batch_files')
                        with gr.Row():
                            vlm_batch_folder = gr.File(label="Folder", show_label=True, file_count='directory', file_types=['image'], interactive=True, height=100, elem_id='vlm_batch_folder')
                        with gr.Row():
                            vlm_batch_str = gr.Textbox(label="Folder", value="", interactive=True, elem_id='vlm_batch_str')
                        with gr.Row():
                            vlm_save_txt = gr.Checkbox(label='Save Caption Files', value=True, elem_id="vlm_save_output")
                            vlm_save_append = gr.Checkbox(label='Append Caption Files', value=False, elem_id="vlm_save_append")
                            vlm_save_json = gr.Checkbox(label='Save Caption JSON', value=True, elem_id="vlm_save_json")
                            vlm_folder_recursive = gr.Checkbox(label='Recursive', value=False, elem_id="vlm_folder_recursive")
                        with gr.Row(elem_id='caption_buttons_batch'):
                            btn_vlm_caption_batch = gr.Button("Batch Caption", variant='primary', elem_id="btn_vlm_caption_batch")
                    with gr.Row():
                        btn_vlm_caption = gr.Button("Caption", variant='primary', elem_id="btn_vlm_caption")

                with gr.Tab("VLM Analyze", elem_id="tab_vlm_analyze"):
                    from modules.caption import vqa
                    analyze_question_placeholder = 'Enter your analysis question or leave blank to use default'
                    analyze_prompt_placeholder = 'Enter your prompt to match with image or leave blank to use image metadata'
                    with gr.Row():
                        analyze_system = gr.Textbox(label="Analyze System Prompt", value=vqa.vlm_analyze, lines=1, elem_id='analyze_system')
                    with gr.Row():
                        analyze_question = gr.Textbox(label="Analyze Question", placeholder=analyze_question_placeholder, lines=1, elem_id='analyze_question')
                    with gr.Row():
                        analyze_prompt = gr.Textbox(label="Prompt", placeholder=analyze_prompt_placeholder, lines=2, elem_id='analyze_prompt')
                    with gr.Row(elem_id='caption_buttons_query'):
                        analyze_model = gr.Dropdown(list(vqa.vlm_models), value=current_vlm_model, label='VLM Model', elem_id='analyze_model')
                    with gr.Row():
                        analyze_load_btn = gr.Button(value='Load', elem_id='analyze_load', variant='secondary')
                        analyze_unload_btn = gr.Button(value='Unload', elem_id='analyze_unload', variant='secondary')
                    with gr.Accordion(label='Analyze: Batch', open=False, visible=True):
                        with gr.Row():
                            analyze_batch_files = gr.File(label="Files", show_label=True, file_count='multiple', file_types=['image'], interactive=True, height=100, elem_id='analyze_batch_files')
                        with gr.Row():
                            analyze_batch_folder = gr.File(label="Folder", show_label=True, file_count='directory', file_types=['image'], interactive=True, height=100, elem_id='analyze_batch_folder')
                        with gr.Row():
                            analyze_batch_str = gr.Textbox(label="Folder", value="", interactive=True, elem_id='analyze_batch_str')
                        with gr.Row():
                            analyze_save_txt = gr.Checkbox(label='Save Analysis Files', value=True, elem_id="analyze_save_txt")
                            analyze_save_append = gr.Checkbox(label='Append Analysis Files', value=False, elem_id="analyze_save_append")
                            analyze_save_json = gr.Checkbox(label='Save Analysis JSON', value=False, elem_id="analyze_save_json")
                            analyze_folder_recursive = gr.Checkbox(label='Recursive', value=False, elem_id="analyze_folder_recursive")
                        with gr.Row(elem_id='caption_buttons_batch'):
                            btn_analyze_caption_batch = gr.Button("Batch Analysis", variant='primary', elem_id="btn_analyze_caption_batch")
                    with gr.Row():
                        btn_analyze_caption = gr.Button("Analyze", variant='primary', elem_id="btn_analyze_caption")

                with gr.Tab("OpenCLiP", elem_id='tab_openclip'):
                    with gr.Row():
                        clip_model = gr.Dropdown([], value=shared.opts.caption_openclip_model, label='CLiP Model', elem_id='clip_clip_model')
                        ui_common.create_refresh_button(clip_model, openclip.refresh_clip_models, lambda: {"choices": openclip.refresh_clip_models()}, 'clip_models_refresh')
                        blip_model = gr.Dropdown(list(openclip.caption_models), value=shared.opts.caption_openclip_blip_model, label='Caption Model', elem_id='btN_clip_blip_model')
                        clip_mode = gr.Dropdown(openclip.caption_types, label='CLiP Mode', value='fast', elem_id='clip_clip_mode')
                    with gr.Accordion(label='Caption: Advanced Options', open=False, visible=True):
                        with gr.Row():
                            clip_max_length = gr.Slider(label='clip: max length', value=shared.opts.caption_openclip_max_length, minimum=16, maximum=1024, step=1, elem_id='clip_caption_max_length')
                            clip_chunk_size = gr.Slider(label='clip: chunk size', value=shared.opts.caption_openclip_chunk_size, minimum=256, maximum=4096, step=8, elem_id='clip_chunk_size')
                        with gr.Row():
                            clip_min_flavors = gr.Slider(label='clip: min flavors', value=shared.opts.caption_openclip_min_flavors, minimum=1, maximum=16, step=1, elem_id='clip_min_flavors')
                            clip_max_flavors = gr.Slider(label='clip: max flavors', value=shared.opts.caption_openclip_max_flavors, minimum=1, maximum=64, step=1, elem_id='clip_max_flavors')
                            clip_flavor_count = gr.Slider(label='clip: intermediates', value=shared.opts.caption_openclip_flavor_count, minimum=256, maximum=4096, step=8, elem_id='clip_flavor_intermediate_count')
                        with gr.Row():
                            clip_num_beams = gr.Slider(label='clip: num beams', value=shared.opts.caption_openclip_num_beams, minimum=1, maximum=16, step=1, elem_id='clip_num_beams')
                        clip_max_length.change(fn=update_clip_params, inputs=[clip_max_length, clip_chunk_size, clip_min_flavors, clip_max_flavors, clip_flavor_count, clip_num_beams], outputs=[])
                        clip_chunk_size.change(fn=update_clip_params, inputs=[clip_max_length, clip_chunk_size, clip_min_flavors, clip_max_flavors, clip_flavor_count, clip_num_beams], outputs=[])
                        clip_min_flavors.change(fn=update_clip_params, inputs=[clip_max_length, clip_chunk_size, clip_min_flavors, clip_max_flavors, clip_flavor_count, clip_num_beams], outputs=[])
                        clip_max_flavors.change(fn=update_clip_params, inputs=[clip_max_length, clip_chunk_size, clip_min_flavors, clip_max_flavors, clip_flavor_count, clip_num_beams], outputs=[])
                        clip_flavor_count.change(fn=update_clip_params, inputs=[clip_max_length, clip_chunk_size, clip_min_flavors, clip_max_flavors, clip_flavor_count, clip_num_beams], outputs=[])
                        clip_num_beams.change(fn=update_clip_params, inputs=[clip_max_length, clip_chunk_size, clip_min_flavors, clip_max_flavors, clip_flavor_count, clip_num_beams], outputs=[])
                    with gr.Accordion(label='Caption: Batch', open=False, visible=True):
                        with gr.Row():
                            clip_batch_files = gr.File(label="Files", show_label=True, file_count='multiple', file_types=['image'], interactive=True, height=100, elem_id='clip_batch_files')
                        with gr.Row():
                            clip_batch_folder = gr.File(label="Folder", show_label=True, file_count='directory', file_types=['image'], interactive=True, height=100, elem_id='clip_batch_folder')
                        with gr.Row():
                            clip_batch_str = gr.Textbox(label="Folder", value="", interactive=True, elem_id='clip_batch_str')
                        with gr.Row():
                            clip_save_output = gr.Checkbox(label='Save Caption Files', value=True, elem_id="clip_save_output")
                            clip_save_append = gr.Checkbox(label='Append Caption Files', value=False, elem_id="clip_save_append")
                            clip_folder_recursive = gr.Checkbox(label='Recursive', value=False, elem_id="clip_folder_recursive")
                        with gr.Row():
                            btn_clip_caption_batch = gr.Button("Batch Caption", variant='primary', elem_id="btn_clip_caption_batch")
                    with gr.Row():
                        btn_clip_caption_img = gr.Button("CLiP Caption", variant='primary', elem_id="btn_clip_caption_img")
                        btn_clip_analyze_img = gr.Button("CLiP Analyze", variant='primary', elem_id="btn_clip_analyze_img")

                with gr.Tab("Tagger", elem_id='tab_tagger'):
                    from modules.caption import tagger
                    with gr.Row():
                        wd_model = gr.Dropdown(tagger.get_models(), value=shared.opts.waifudiffusion_model, label='Tagger Model', elem_id='wd_model')
                        ui_common.create_refresh_button(wd_model, tagger.refresh_models, lambda: {"choices": tagger.get_models()}, 'wd_models_refresh')
                    with gr.Row():
                        wd_load_btn = gr.Button(value='Load', elem_id='wd_load', variant='secondary')
                        wd_unload_btn = gr.Button(value='Unload', elem_id='wd_unload', variant='secondary')
                    with gr.Accordion(label='Caption: Advanced Options', open=True, visible=True):
                        with gr.Row():
                            wd_general_threshold = gr.Slider(label='General threshold', value=shared.opts.tagger_threshold, minimum=0.0, maximum=1.0, step=0.01, elem_id='wd_general_threshold')
                            wd_character_threshold = gr.Slider(label='Character threshold', value=shared.opts.waifudiffusion_character_threshold, minimum=0.0, maximum=1.0, step=0.01, elem_id='wd_character_threshold')
                        with gr.Row():
                            wd_max_tags = gr.Slider(label='Max tags', value=shared.opts.tagger_max_tags, minimum=1, maximum=512, step=1, elem_id='wd_max_tags')
                            wd_include_rating = gr.Checkbox(label='Include rating', value=shared.opts.tagger_include_rating, elem_id='wd_include_rating')
                        with gr.Row():
                            wd_sort_alpha = gr.Checkbox(label='Sort alphabetically', value=shared.opts.tagger_sort_alpha, elem_id='wd_sort_alpha')
                            wd_use_spaces = gr.Checkbox(label='Use spaces', value=shared.opts.tagger_use_spaces, elem_id='wd_use_spaces')
                            wd_escape = gr.Checkbox(label='Escape brackets', value=shared.opts.tagger_escape_brackets, elem_id='wd_escape')
                        with gr.Row():
                            wd_exclude_tags = gr.Textbox(label='Exclude tags', value=shared.opts.tagger_exclude_tags, placeholder='Comma-separated tags to exclude', elem_id='wd_exclude_tags')
                        with gr.Row():
                            wd_show_scores = gr.Checkbox(label='Show confidence scores', value=shared.opts.tagger_show_scores, elem_id='wd_show_scores')
                    gr.HTML('<style>#wd_character_threshold:has(input:disabled), #wd_include_rating:has(input:disabled) { opacity: 0.5; }</style>')
                    gr.HTML('<style>#vlm_batch_files, #vlm_batch_folder, #clip_batch_files, #clip_batch_folder, #wd_batch_files, #wd_batch_folder { max-height: 10em; overflow-y: auto !important; }</style>')
                    with gr.Accordion(label='Caption: Batch', open=False, visible=True):
                        with gr.Row():
                            wd_batch_files = gr.File(label="Files", show_label=True, file_count='multiple', file_types=['image'], interactive=True, height=100, elem_id='wd_batch_files')
                        with gr.Row():
                            wd_batch_folder = gr.File(label="Folder", show_label=True, file_count='directory', file_types=['image'], interactive=True, height=100, elem_id='wd_batch_folder')
                        with gr.Row():
                            wd_batch_str = gr.Textbox(label="Folder", value="", interactive=True, elem_id='wd_batch_str')
                        with gr.Row():
                            wd_save_output = gr.Checkbox(label='Save Caption Files', value=True, elem_id="wd_save_output")
                            wd_save_append = gr.Checkbox(label='Append Caption Files', value=False, elem_id="wd_save_append")
                            wd_folder_recursive = gr.Checkbox(label='Recursive', value=False, elem_id="wd_folder_recursive")
                        with gr.Row():
                            btn_wd_tag_batch = gr.Button("Batch Tag", variant='primary', elem_id="btn_wd_tag_batch")
                    with gr.Row():
                        btn_wd_tag = gr.Button("Tag", variant='primary', elem_id="btn_wd_tag")
                with gr.Tab("Default", elem_id='tab_caption_default'):
                    with gr.Row():
                        default_caption_type = gr.Radio(
                            choices=["VLM", "OpenCLiP", "Tagger"],
                            value=shared.opts.caption_default_type,
                            label="Default Caption Type",
                            elem_id="default_caption_type"
                        )

        with gr.Column(variant='compact', elem_id='caption_output'):
            with gr.Row(elem_id='caption_output_prompt'):
                prompt = gr.Textbox(label="Answer", lines=12, placeholder="VLM generated image description")
            with gr.Row(elem_id='caption_output_image'):
                output_image = gr.Image(type='pil', label="Annotated Image", interactive=False, visible=False, elem_id='caption_output_image_display')
            with gr.Row(elem_id='caption_output_classes'):
                medium = gr.Label(elem_id="caption_label_medium", label="Medium", num_top_classes=5, visible=False)
                artist = gr.Label(elem_id="caption_label_artist", label="Artist", num_top_classes=5, visible=False)
                movement = gr.Label(elem_id="caption_label_movement", label="Movement", num_top_classes=5, visible=False)
                trending = gr.Label(elem_id="caption_label_trending", label="Trending", num_top_classes=5, visible=False)
                flavor = gr.Label(elem_id="caption_label_flavor", label="Flavor", num_top_classes=5, visible=False)
                clip_labels_text = gr.Textbox(elem_id="caption_clip_labels_text", label="CLiP Analysis", lines=15, interactive=False, visible=False, show_label=False)
            with gr.Row(elem_id='copy_buttons_caption'):
                copy_caption_buttons = generation_parameters_copypaste.create_buttons(["txt2img", "img2img", "control", "extras"])

    dummy = gr.Label(value='-1', visible=False)
    btn_vlm_caption.click(
        _js="getCaptionActiveTab", # js to insert current tab name as first argument
        fn=caption_wrapper,
        inputs=[dummy, image,
                vlm_question, vlm_system, vlm_prompt, vlm_model, vlm_prefill, vlm_thinking_mode, vlm_custom_args,
                analyze_question, analyze_system, analyze_prompt, analyze_model,
                clip_model, blip_model, clip_mode,
                wd_model, wd_general_threshold, wd_character_threshold, wd_include_rating, wd_exclude_tags, wd_max_tags, wd_sort_alpha, wd_use_spaces, wd_escape
               ],
        outputs=[prompt, output_image]
    )
    vlm_batch = gr.State('caption')
    btn_vlm_caption_batch.click(
        fn=vqa.batch,
        inputs=[vlm_model, vlm_system, vlm_batch_files, vlm_batch_folder, vlm_batch_str, vlm_question, vlm_prompt, vlm_save_txt, vlm_save_append, vlm_save_json, vlm_folder_recursive, vlm_prefill, vlm_thinking_mode, vlm_batch],
        outputs=[prompt]
    ).then(fn=lambda: gr.update(visible=False), inputs=[], outputs=[output_image])

    btn_analyze_caption.click(
        _js="getCaptionActiveTab", # js to insert current tab name as first argument
        fn=caption_wrapper,
        inputs=[dummy, image,
                vlm_question, vlm_system, vlm_prompt, vlm_model, vlm_prefill, vlm_thinking_mode, vlm_custom_args,
                analyze_question, analyze_system, analyze_prompt, analyze_model,
                clip_model, blip_model, clip_mode,
                wd_model, wd_general_threshold, wd_character_threshold, wd_include_rating, wd_exclude_tags, wd_max_tags, wd_sort_alpha, wd_use_spaces, wd_escape
               ],
        outputs=[prompt, output_image]
    )
    analyze_batch = gr.State('analyze')
    btn_analyze_caption_batch.click(
        fn=vqa.batch,
        inputs=[analyze_model, analyze_system, analyze_batch_files, analyze_batch_folder, analyze_batch_str, analyze_question, analyze_prompt, analyze_save_txt, analyze_save_append, analyze_save_json,analyze_folder_recursive, vlm_prefill, vlm_thinking_mode, analyze_batch],
        outputs=[prompt]
    ).then(fn=lambda: gr.update(visible=False), inputs=[], outputs=[output_image])

    btn_clip_caption_img.click(
        fn=openclip.caption_image,
        inputs=[image, clip_model, blip_model, clip_mode],
        outputs=[prompt]
    ).then(fn=lambda: gr.update(visible=False), inputs=[], outputs=[output_image])
    btn_clip_analyze_img.click(
        fn=openclip.analyze_image,
        inputs=[image, clip_model, blip_model],
        outputs=[medium, artist, movement, trending, flavor, clip_labels_text]
    ).then(fn=lambda: gr.update(visible=False), inputs=[], outputs=[output_image])

    btn_wd_tag.click(
        fn=tagger_tag_wrapper,
        inputs=[image, wd_model, wd_general_threshold, wd_character_threshold, wd_include_rating, wd_exclude_tags, wd_max_tags, wd_sort_alpha, wd_use_spaces, wd_escape],
        outputs=[prompt]
    ).then(fn=lambda: gr.update(visible=False), inputs=[], outputs=[output_image])
    btn_wd_tag_batch.click(
        fn=tagger_batch_wrapper,
        inputs=[wd_model, wd_batch_files, wd_batch_folder, wd_batch_str, wd_save_output, wd_save_append, wd_folder_recursive, wd_general_threshold, wd_character_threshold, wd_include_rating, wd_exclude_tags, wd_max_tags, wd_sort_alpha, wd_use_spaces, wd_escape],
        outputs=[prompt]
    ).then(fn=lambda: gr.update(visible=False), inputs=[], outputs=[output_image])
    btn_clip_caption_batch.click(
        fn=openclip.caption_batch,
        inputs=[dummy, clip_batch_files, clip_batch_folder, clip_batch_str, clip_model, blip_model, clip_mode, clip_save_output, clip_save_append, clip_folder_recursive],
        outputs=[prompt]
    ).then(fn=lambda: gr.update(visible=False), inputs=[], outputs=[output_image])

    # Dynamic UI updates based on selected model and task
    vlm_model.change(fn=update_vlm_prompts_for_model, inputs=[vlm_model], outputs=[vlm_question])
    vlm_question.change(fn=update_vlm_prompt_placeholder, inputs=[vlm_question], outputs=[vlm_prompt])

    # Load/Unload model buttons
    vlm_load_btn.click(fn=vqa.load_model, inputs=[vlm_model], outputs=[])
    vlm_unload_btn.click(fn=vqa.unload_model, inputs=[], outputs=[])
    analyze_load_btn.click(fn=vqa.load_model, inputs=[analyze_model], outputs=[])
    analyze_unload_btn.click(fn=vqa.unload_model, inputs=[], outputs=[])

    def tagger_load_wrapper(model_name):
        from modules.caption import tagger
        return tagger.load_model(model_name)
    def tagger_unload_wrapper():
        from modules.caption import tagger
        return tagger.unload_model()
    wd_load_btn.click(fn=tagger_load_wrapper, inputs=[wd_model], outputs=[])
    wd_unload_btn.click(fn=tagger_unload_wrapper, inputs=[], outputs=[])

    # Dynamic UI update when tagger model changes (disable controls for DeepBooru)
    wd_model.change(fn=update_tagger_ui, inputs=[wd_model], outputs=[wd_character_threshold, wd_include_rating], show_progress=False)

    # Save tagger parameters to shared.opts when UI controls change
    tagger_inputs = [wd_model, wd_general_threshold, wd_character_threshold, wd_include_rating, wd_max_tags, wd_sort_alpha, wd_use_spaces, wd_escape, wd_exclude_tags, wd_show_scores]
    wd_model.change(fn=update_tagger_params, inputs=tagger_inputs, outputs=[], show_progress=False)
    wd_general_threshold.change(fn=update_tagger_params, inputs=tagger_inputs, outputs=[], show_progress=False)
    wd_character_threshold.change(fn=update_tagger_params, inputs=tagger_inputs, outputs=[], show_progress=False)
    wd_include_rating.change(fn=update_tagger_params, inputs=tagger_inputs, outputs=[], show_progress=False)
    wd_max_tags.change(fn=update_tagger_params, inputs=tagger_inputs, outputs=[], show_progress=False)
    wd_sort_alpha.change(fn=update_tagger_params, inputs=tagger_inputs, outputs=[], show_progress=False)
    wd_use_spaces.change(fn=update_tagger_params, inputs=tagger_inputs, outputs=[], show_progress=False)
    wd_escape.change(fn=update_tagger_params, inputs=tagger_inputs, outputs=[], show_progress=False)
    wd_exclude_tags.change(fn=update_tagger_params, inputs=tagger_inputs, outputs=[], show_progress=False)
    wd_show_scores.change(fn=update_tagger_params, inputs=tagger_inputs, outputs=[], show_progress=False)

    # Save CLiP model parameters to shared.opts when UI controls change
    clip_model_inputs = [clip_model, blip_model, clip_mode]
    clip_model.change(fn=update_clip_model_params, inputs=clip_model_inputs, outputs=[], show_progress=False)
    blip_model.change(fn=update_clip_model_params, inputs=clip_model_inputs, outputs=[], show_progress=False)
    clip_mode.change(fn=update_clip_model_params, inputs=clip_model_inputs, outputs=[], show_progress=False)

    # Save VLM model parameters to shared.opts when UI controls change
    vlm_model_inputs = [vlm_model, vlm_system]
    vlm_model.change(fn=update_vlm_model_params, inputs=vlm_model_inputs, outputs=[], show_progress=False)
    vlm_system.change(fn=update_vlm_model_params, inputs=vlm_model_inputs, outputs=[], show_progress=False)

    # Save default caption type to shared.opts when UI control changes
    default_caption_type.change(fn=update_default_caption_type, inputs=[default_caption_type], outputs=[], show_progress=False)

    for tabname, button in copy_caption_buttons.items():
        generation_parameters_copypaste.register_paste_params_button(generation_parameters_copypaste.ParamBinding(paste_button=button, tabname=tabname, source_text_component=prompt, source_image_component=image,))
    generation_parameters_copypaste.add_paste_fields("caption", image, None)