diff --git a/modules/interrogate/interrogate.py b/modules/interrogate/interrogate.py index 7f7befcf2..ae28fe110 100644 --- a/modules/interrogate/interrogate.py +++ b/modules/interrogate/interrogate.py @@ -12,7 +12,7 @@ def interrogate(image): shared.log.error('Interrogate: no image provided') return '' t0 = time.time() - if shared.opts.interrogate_default_type == 'OpenCLiP': + if shared.opts.interrogate_default_type == 'CLiP': shared.log.info(f'Interrogate: type={shared.opts.interrogate_default_type} clip="{shared.opts.interrogate_clip_model}" blip="{shared.opts.interrogate_blip_model}" mode="{shared.opts.interrogate_clip_mode}"') from modules.interrogate import openclip openclip.load_interrogator(clip_model=shared.opts.interrogate_clip_model, blip_model=shared.opts.interrogate_blip_model) @@ -20,10 +20,21 @@ def interrogate(image): prompt = openclip.interrogate(image, mode=shared.opts.interrogate_clip_mode) shared.log.debug(f'Interrogate: time={time.time()-t0:.2f} answer="{prompt}"') return prompt - elif shared.opts.interrogate_default_type == 'DeepBooru': - shared.log.info(f'Interrogate: type={shared.opts.interrogate_default_type}') - from modules.interrogate import deepbooru - prompt = deepbooru.model.tag(image) + elif shared.opts.interrogate_default_type == 'Tagger': + shared.log.info(f'Interrogate: type={shared.opts.interrogate_default_type} model="{shared.opts.wd14_model}"') + from modules.interrogate import tagger + prompt = tagger.tag( + image=image, + model_name=shared.opts.wd14_model, + general_threshold=shared.opts.wd14_general_threshold, + character_threshold=shared.opts.wd14_character_threshold, + include_rating=shared.opts.wd14_include_rating, + exclude_tags=shared.opts.tagger_exclude_tags, + max_tags=shared.opts.tagger_max_tags, + sort_alpha=shared.opts.tagger_sort_alpha, + use_spaces=shared.opts.tagger_use_spaces, + escape_brackets=shared.opts.tagger_escape, + ) shared.log.debug(f'Interrogate: time={time.time()-t0:.2f} answer="{prompt}"') return prompt elif shared.opts.interrogate_default_type == 'VLM': diff --git a/modules/shared.py b/modules/shared.py index d49005a9a..7b39d88a2 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -673,14 +673,15 @@ options_templates.update(options_section(('postprocessing', "Postprocessing"), { })) options_templates.update(options_section(('interrogate', "Interrogate"), { - "interrogate_default_type": OptionInfo("VLM", "Default caption type", gr.Radio, {"choices": ["OpenCLiP", "VLM", "DeepBooru"]}), + "interrogate_default_type": OptionInfo("VLM", "Default caption type", gr.Radio, {"choices": ["VLM", "CLiP", "Tagger"]}), "interrogate_offload": OptionInfo(True, "Offload models "), - "interrogate_score": OptionInfo(False, "Include scores in results when available"), + "interrogate_score": OptionInfo(False, "Include scores in results when available", gr.Checkbox, {"visible": False}), - "interrogate_clip_sep": OptionInfo("

OpenCLiP

", "", gr.HTML), - "interrogate_clip_model": OptionInfo("ViT-L-14/openai", "CLiP: default model", gr.Dropdown, lambda: {"choices": get_clip_models()}, refresh=refresh_clip_models), - "interrogate_clip_mode": OptionInfo(caption_types[0], "CLiP: default mode", gr.Dropdown, {"choices": caption_types}), - "interrogate_blip_model": OptionInfo(list(caption_models)[0], "CLiP: default captioner", gr.Dropdown, {"choices": list(caption_models)}), + # OpenCLiP settings (hidden - controlled via Caption Tab) + "interrogate_clip_sep": OptionInfo("

OpenCLiP

", "", gr.HTML, {"visible": False}), + "interrogate_clip_model": OptionInfo("ViT-L-14/openai", "CLiP: default model", gr.Dropdown, lambda: {"choices": get_clip_models(), "visible": False}, refresh=refresh_clip_models), + "interrogate_clip_mode": OptionInfo(caption_types[0], "CLiP: default mode", gr.Dropdown, {"choices": caption_types, "visible": False}), + "interrogate_blip_model": OptionInfo(list(caption_models)[0], "CLiP: default captioner", gr.Dropdown, {"choices": list(caption_models), "visible": False}), "interrogate_clip_num_beams": OptionInfo(1, "CLiP: num beams", gr.Slider, {"minimum": 1, "maximum": 16, "step": 1, "visible": False}), "interrogate_clip_min_length": OptionInfo(32, "CLiP: min length", gr.Slider, {"minimum": 1, "maximum": 128, "step": 1, "visible": False}), "interrogate_clip_max_length": OptionInfo(74, "CLiP: max length", gr.Slider, {"minimum": 1, "maximum": 512, "step": 1, "visible": False}), @@ -689,13 +690,14 @@ options_templates.update(options_section(('interrogate', "Interrogate"), { "interrogate_clip_flavor_count": OptionInfo(1024, "CLiP: intermediate flavors", gr.Slider, {"minimum": 256, "maximum": 4096, "step": 64, "visible": False}), "interrogate_clip_chunk_size": OptionInfo(1024, "CLiP: chunk size", gr.Slider, {"minimum": 256, "maximum": 4096, "step": 64, "visible": False}), - "interrogate_vlm_sep": OptionInfo("

VLM

", "", gr.HTML), - "interrogate_vlm_model": OptionInfo(vlm_default, "VLM: default model", gr.Dropdown, {"choices": list(vlm_models)}), - "interrogate_vlm_prompt": OptionInfo(vlm_prompts[2], "VLM: default prompt", DropdownEditable, {"choices": vlm_prompts }), - "interrogate_vlm_system": OptionInfo(vlm_system, "VLM: default prompt"), + # VLM settings (hidden - controlled via Caption Tab) + "interrogate_vlm_sep": OptionInfo("

VLM

", "", gr.HTML, {"visible": False}), + "interrogate_vlm_model": OptionInfo(vlm_default, "VLM: default model", gr.Dropdown, {"choices": list(vlm_models), "visible": False}), + "interrogate_vlm_prompt": OptionInfo(vlm_prompts[2], "VLM: default prompt", DropdownEditable, {"choices": vlm_prompts, "visible": False}), + "interrogate_vlm_system": OptionInfo(vlm_system, "VLM: default prompt", gr.Textbox, {"visible": False}), "interrogate_vlm_num_beams": OptionInfo(1, "VLM: num beams", gr.Slider, {"minimum": 1, "maximum": 16, "step": 1, "visible": False}), "interrogate_vlm_max_length": OptionInfo(512, "VLM: max length", gr.Slider, {"minimum": 1, "maximum": 4096, "step": 1, "visible": False}), - "interrogate_vlm_do_sample": OptionInfo(True, "VLM: use sample method"), + "interrogate_vlm_do_sample": OptionInfo(True, "VLM: use sample method", gr.Checkbox, {"visible": False}), "interrogate_vlm_temperature": OptionInfo(0.8, "VLM: temperature", gr.Slider, {"minimum": 0, "maximum": 1.0, "step": 0.01, "visible": False}), "interrogate_vlm_top_k": OptionInfo(0, "VLM: top-k", gr.Slider, {"minimum": 0, "maximum": 99, "step": 1, "visible": False}), "interrogate_vlm_top_p": OptionInfo(0, "VLM: top-p", gr.Slider, {"minimum": 0, "maximum": 1.0, "step": 0.01, "visible": False}), @@ -703,24 +705,24 @@ options_templates.update(options_section(('interrogate', "Interrogate"), { "interrogate_vlm_keep_thinking": OptionInfo(False, "VLM: keep reasoning trace in output", gr.Checkbox, {"visible": False}), "interrogate_vlm_thinking_mode": OptionInfo(False, "VLM: enable thinking/reasoning mode", gr.Checkbox, {"visible": False}), - # Common tagger settings (shared by DeepBooru and WD14) - "tagger_sep": OptionInfo("

Tagger Settings

", "", gr.HTML), - "tagger_max_tags": OptionInfo(74, "Tagger: max tags", gr.Slider, {"minimum": 1, "maximum": 512, "step": 1}), - "tagger_sort_alpha": OptionInfo(False, "Tagger: sort alphabetically"), - "tagger_use_spaces": OptionInfo(False, "Tagger: use spaces for tags"), - "tagger_escape": OptionInfo(True, "Tagger: escape brackets"), - "tagger_exclude_tags": OptionInfo("", "Tagger: exclude tags"), + # Common tagger settings (hidden - controlled via Caption Tab) + "tagger_sep": OptionInfo("

Tagger Settings

", "", gr.HTML, {"visible": False}), + "tagger_max_tags": OptionInfo(74, "Tagger: max tags", gr.Slider, {"minimum": 1, "maximum": 512, "step": 1, "visible": False}), + "tagger_sort_alpha": OptionInfo(False, "Tagger: sort alphabetically", gr.Checkbox, {"visible": False}), + "tagger_use_spaces": OptionInfo(False, "Tagger: use spaces for tags", gr.Checkbox, {"visible": False}), + "tagger_escape": OptionInfo(True, "Tagger: escape brackets", gr.Checkbox, {"visible": False}), + "tagger_exclude_tags": OptionInfo("", "Tagger: exclude tags", gr.Textbox, {"visible": False}), - # DeepBooru-specific settings - "deepbooru_sep": OptionInfo("

DeepBooru

", "", gr.HTML), - "deepbooru_score_threshold": OptionInfo(0.65, "DeepBooru: score threshold", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01}), + # DeepBooru-specific settings (hidden - controlled via Caption Tab) + "deepbooru_sep": OptionInfo("

DeepBooru

", "", gr.HTML, {"visible": False}), + "deepbooru_score_threshold": OptionInfo(0.65, "DeepBooru: score threshold", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01, "visible": False}), - # WD14-specific settings - "wd14_sep": OptionInfo("

WD14 Tagger

", "", gr.HTML), - "wd14_model": OptionInfo("wd-eva02-large-tagger-v3", "WD14: default model", gr.Dropdown, {"choices": []}), - "wd14_general_threshold": OptionInfo(0.35, "WD14: general tag threshold", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01}), - "wd14_character_threshold": OptionInfo(0.85, "WD14: character tag threshold", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01}), - "wd14_include_rating": OptionInfo(False, "WD14: include rating tags"), + # WD14-specific settings (hidden - controlled via Caption Tab) + "wd14_sep": OptionInfo("

WD14 Tagger

", "", gr.HTML, {"visible": False}), + "wd14_model": OptionInfo("wd-eva02-large-tagger-v3", "WD14: default model", gr.Dropdown, {"choices": [], "visible": False}), + "wd14_general_threshold": OptionInfo(0.35, "WD14: general tag threshold", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01, "visible": False}), + "wd14_character_threshold": OptionInfo(0.85, "WD14: character tag threshold", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01, "visible": False}), + "wd14_include_rating": OptionInfo(False, "WD14: include rating tags", gr.Checkbox, {"visible": False}), })) options_templates.update(options_section(('huggingface', "Huggingface"), { diff --git a/modules/ui_caption.py b/modules/ui_caption.py index 4adf4f9b5..0c962d266 100644 --- a/modules/ui_caption.py +++ b/modules/ui_caption.py @@ -96,6 +96,20 @@ def update_tagger_ui(model_name): ] +def update_tagger_params(model_name, general_threshold, character_threshold, include_rating, max_tags, sort_alpha, use_spaces, escape_brackets, exclude_tags): + """Save all tagger parameters to shared.opts when UI controls change.""" + shared.opts.wd14_model = model_name + shared.opts.wd14_general_threshold = float(general_threshold) + shared.opts.wd14_character_threshold = float(character_threshold) + shared.opts.wd14_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 = bool(escape_brackets) + shared.opts.tagger_exclude_tags = str(exclude_tags) + shared.opts.save() + + def update_clip_params(*args): clip_min_length, clip_max_length, clip_chunk_size, clip_min_flavors, clip_max_flavors, clip_flavor_count, clip_num_beams = args shared.opts.interrogate_clip_min_length = int(clip_min_length) @@ -109,6 +123,21 @@ def update_clip_params(*args): openclip.update_interrogate_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.interrogate_clip_model = str(clip_model) + shared.opts.interrogate_blip_model = str(blip_model) + shared.opts.interrogate_clip_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.interrogate_vlm_model = str(vlm_model) + shared.opts.interrogate_vlm_system = str(vlm_system) + shared.opts.save() + + def create_ui(): shared.log.debug('UI initialize: tab=caption') with gr.Row(equal_height=False, variant='compact', elem_classes="caption", elem_id="caption_tab"): @@ -211,7 +240,7 @@ def create_ui(): with gr.Row(): btn_clip_interrogate_img = gr.Button("Interrogate", variant='primary', elem_id="btn_clip_interrogate_img") btn_clip_analyze_img = gr.Button("Analyze", variant='primary', elem_id="btn_clip_analyze_img") - with gr.Tab("Booru Tags", elem_id='tab_booru_tags'): + with gr.Tab("Tagger", elem_id='tab_tagger'): from modules.interrogate import tagger with gr.Row(): wd_model = gr.Dropdown(tagger.get_models(), value=shared.opts.wd14_model, label='Tagger Model', elem_id='wd_model') @@ -290,6 +319,29 @@ def create_ui(): # 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_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) + + # 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) + for tabname, button in copy_interrogate_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)