From d2d9f7caeaa0176233da808a64c9640bc181b6c9 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Fri, 29 May 2026 18:11:12 +0200 Subject: [PATCH] captioning improvements and cleanup Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 13 +- modules/attention.py | 10 -- modules/caption/attention.py | 39 ++++++ modules/caption/deepseek.py | 4 +- modules/caption/helpers.py | 50 +++++++ modules/caption/joycaption.py | 2 +- modules/caption/joytag.py | 2 +- modules/caption/models_def.py | 148 +++++++++++---------- modules/caption/moondream3.py | 50 +++---- modules/caption/vqa.py | 238 ++++++++++++++++++++-------------- modules/devices.py | 2 +- modules/ui_caption.py | 34 ++--- modules/ui_definitions.py | 1 + scripts/prompt_enhance.py | 52 +------- 14 files changed, 370 insertions(+), 275 deletions(-) create mode 100644 modules/caption/attention.py create mode 100644 modules/caption/helpers.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 51dedbb14..01ca1ec4d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -40,11 +40,15 @@ Plus continued work on modernization of codebase: UI is now fully TypeScript bas *for example*, `child` can match `kid`, `girl`, `boy` and it expands the functionality with customizable semantic matching: *for example*, `young ...` will match before next word appears in the prompt and steer away from it towards desired choices - - add *wildcards* (if used) info to image metadata - if wildcards or styles modify prompt, add original prompt to image metadata as *template* - - **Captioning** new feature: analyze existing images for prompt adherence + - **Captioning** + new feature: analyze existing images for prompt adherence *tip*: image analysis requires larger VLM model to produce quality output new api endpoint: `/sdapi/v1/analyze` + cleanup list of predefined models, new models added and some old removed + improved default values plus some new params like min length and `custom args` so you can pass anything to an llm model + improved system prompts + - add *wildcards* (if used) info to image **metadata** + if wildcards or styles modify prompt, add original prompt to image metadata as *template* - **Masking** updated interface and capabilities you can now also select mask type instead of focing alpha mask with all models - **HF download** use `XET` by default @@ -75,7 +79,7 @@ Plus continued work on modernization of codebase: UI is now fully TypeScript bas - improve `kanvas` typing - additional strong typing in core, thanks @awsr - **Fixes** - - *hidream-o1* prequant loading + - `hidream-o1` prequant loading - `gradio` initial hijack - `SmolVLM` captioning - `gradio` temp files guard against large image @@ -94,6 +98,7 @@ Plus continued work on modernization of codebase: UI is now fully TypeScript bas - `styles` loader exception handling - `kanvas` image change notification - `reinstall` force reinstal of transformers and diffusers + - `ipex` torch install error, thanks @liutyi ## Update for 2026-05-13 diff --git a/modules/attention.py b/modules/attention.py index c9136693a..282718e58 100644 --- a/modules/attention.py +++ b/modules/attention.py @@ -226,16 +226,6 @@ def set_diffusers_attention(pipe, quiet = False): pass else: log.error(f'Torch attention: type="{name}" cls={attention.__class__.__name__} pipe={pipe.__class__.__name__} {e}') - """ # each transformer typically has its own attention processor - if getattr(pipe, "transformer", None) is not None and hasattr(pipe.transformer, "set_attn_processor"): - try: - pipe.transformer.set_attn_processor(attention) - except Exception as e: - if 'Nunchaku' in pipe.transformer.__class__.__name__: - pass - else: - log.error(f'Torch attention: type="{name}" cls={attention.__class__.__name__} pipe={pipe.__class__.__name__} {e}') - """ log.quiet(quiet, f'Setting model: attention="{shared.opts.cross_attention_optimization}"') if shared.opts.cross_attention_optimization == "Disabled": diff --git a/modules/caption/attention.py b/modules/caption/attention.py new file mode 100644 index 000000000..175081817 --- /dev/null +++ b/modules/caption/attention.py @@ -0,0 +1,39 @@ +from modules.logger import log + + +ATTENTIONS = ['eager', 'sdpa', 'flash_attention_2', 'flash_attention_3', 'flex_attention', 'paged|eager', 'paged|sdpa', 'paged|flash_attention_2', 'paged|flash_attention_3'] + + +def get_first_attention_block(model): # for models still using transformers==4 internal architecture + m = model.model + if hasattr(m, "layers"): + return m.layers[0].self_attn if hasattr(m.layers[0], "self_attn") else m.layers[0].attention + if hasattr(m, "transformer") and hasattr(m.transformer, "layers"): + return m.transformer.layers[0].attention + if hasattr(m, "block"): + return m.block[0].attention + return None + + +def set_attention(model): + if not hasattr(model, 'set_attn_implementation'): + return + supported = [] + unsupported = [] + if hasattr(model, "_attn_implementation"): + default = model._attn_implementation # pylint: disable=protected-access + elif hasattr(model, "_get_attn_implementation"): + default = model._get_attn_implementation() # pylint: disable=protected-access + else: + default = get_first_attention_block(model) + log.debug(f"LLM attention: cls={model.__class__.__name__} default={default} fixed") + return + + for name in ATTENTIONS: # type: ignore + try: + model.set_attn_implementation(name) + supported.append(name) + except Exception: + unsupported.append(name) + log.debug(f"LLM attention: cls={model.__class__.__name__} default={default} supported={supported} unsupported={unsupported}") + model.set_attn_implementation(default) # restore default diff --git a/modules/caption/deepseek.py b/modules/caption/deepseek.py index 6d36f7db9..b22730dd8 100644 --- a/modules/caption/deepseek.py +++ b/modules/caption/deepseek.py @@ -34,11 +34,11 @@ def load(repo: str): """Load DeepSeek VL2 model (experimental).""" global vl_gpt, vl_chat_processor, loaded_repo # pylint: disable=global-statement if not shared.cmd_opts.experimental: - log.error(f'Caption: type=vlm model="DeepSeek VL2" repo="{repo}" is experimental-only') + log.error(f'LLM: type=vlm model="DeepSeek VL2" repo="{repo}" is experimental-only') return False folder = os.path.join(paths.script_path, 'repositories', 'deepseek-vl2') if not os.path.exists(folder): - log.error(f'Caption: type=vlm model="DeepSeek VL2" repo="{repo}" deepseek-vl2 repo not found') + log.error(f'LLM: type=vlm model="DeepSeek VL2" repo="{repo}" deepseek-vl2 repo not found') return False if vl_gpt is None or loaded_repo != repo: # GLOBAL PATCHES (not reverted): DeepSeek VL2 requires attrdict and uses LlamaFlashAttention2 diff --git a/modules/caption/helpers.py b/modules/caption/helpers.py new file mode 100644 index 000000000..34967964f --- /dev/null +++ b/modules/caption/helpers.py @@ -0,0 +1,50 @@ +import re +import transformers +from modules.logger import log + + +def get_default_args(model): + to_remove = ['_from_model_config', 'transformers_version'] + config = {} + for k, v in transformers.GenerationConfig._get_default_generation_params().items(): # pylint: disable=protected-access + if v is not None: + config[k] = v + for k, v in transformers.GenerationConfig.from_model_config(model.config).to_dict().items(): + if v is not None: + config[k] = v + for k, v in model.generation_config.to_dict().items(): + if v is not None: + config[k] = v + config = {k: v for k, v in config.items() if k not in to_remove} + return config + +def get_custom_args(model, args_str): + args = {} + if args_str is not None and len(args_str) > 0: + default_args = get_default_args(model) + pairs = re.split(r'[;\n]+', args_str) + for pair in pairs: + if '=' in pair or ':' in pair: + key, value = re.split(r'[=:]', pair, maxsplit=1) + key = key.strip() + value = value.strip() + if key not in default_args: + log.warning(f'Prompt enhance: key="{key}" invalid') + continue + default_value = default_args[key] + try: + if isinstance(default_value, bool): + value = value.lower() in ['true', '1', 'yes'] + elif isinstance(default_value, int): + value = int(value) + elif isinstance(default_value, float): + value = float(value) + elif isinstance(default_value, list): + value = [v.strip() for v in value.split(',')] + elif isinstance(default_value, str): + pass + except ValueError: + log.warning(f'Prompt enhance: key="{key}" value="{value}" typecast failed') + if key and value: + args[key] = value + return args diff --git a/modules/caption/joycaption.py b/modules/caption/joycaption.py index a18b46800..ee527bc70 100644 --- a/modules/caption/joycaption.py +++ b/modules/caption/joycaption.py @@ -65,7 +65,7 @@ def load(repo: str | None = None): if llava_model is None or opts.repo != repo: opts.repo = repo llava_model = None - log.info(f'Caption: type=vlm model="JoyCaption" {str(opts)}') + log.info(f'LLM: type=vlm model="JoyCaption" {str(opts)}') processor = AutoProcessor.from_pretrained(repo, max_pixels=1024*1024, cache_dir=shared.opts.hfcache_dir) quant_args = model_quant.create_config(module='LLM') llava_model = LlavaForConditionalGeneration.from_pretrained( diff --git a/modules/caption/joytag.py b/modules/caption/joytag.py index 710c3b7e3..e105eb2fe 100644 --- a/modules/caption/joytag.py +++ b/modules/caption/joytag.py @@ -1051,7 +1051,7 @@ def load(): with open(os.path.join(folder, 'top_tags.txt'), encoding='utf8') as f: tags = [line.strip() for line in f.readlines() if line.strip()] register_aux('joytag', model) - log.info(f'Caption: type=vlm model="JoyTag" repo="{MODEL_REPO}" tags={len(tags)}') + log.info(f'LLM: type=vlm model="JoyTag" repo="{MODEL_REPO}" tags={len(tags)}') move_aux_to_gpu('joytag') diff --git a/modules/caption/models_def.py b/modules/caption/models_def.py index 21ecbc410..d4ec7b018 100644 --- a/modules/caption/models_def.py +++ b/modules/caption/models_def.py @@ -2,86 +2,82 @@ from modules import ui_symbols vlm_models = { - "Google Gemma 3 4B": "google/gemma-3-4b-it", - "Google Gemma 3n E2B": "google/gemma-3n-E2B-it", # 1.5GB - "Google Gemma 3n E4B": "google/gemma-3n-E4B-it", # 1.5GB - "Google Gemma 4 E2B": "google/gemma-4-E2B-it", + # primary "Google Gemma 4 E4B": "google/gemma-4-E4B-it", - "Heretic Gemma 4 E4B": "p-e-w/gemma-4-E2B-it-heretic-ara", - "Nidum Gemma 3 4B Uncensored": "nidum/Nidum-Gemma-3-4B-it-Uncensored", - "Allura Gemma 3 Glitter 4B": "allura-org/Gemma-3-Glitter-4B", - # Qwen3.5 - "Alibaba Qwen 3.5 0.8B": "Qwen/Qwen3.5-0.8B", - "Alibaba Qwen 3.5 2B": "Qwen/Qwen3.5-2B", - "Alibaba Qwen 3.5 4B": "Qwen/Qwen3.5-4B", + "Google Gemma 4 E4B Trohrbaugh Heretic": "trohrbaugh/gemma-4-E4B-it-heretic-ara", + "Google Gemma 4 E2B": "google/gemma-4-E2B-it", + "Google Gemma 3n E4B": "google/gemma-3n-E4B-it", + "Google Gemma 3n E2B": "google/gemma-3n-E2B-it", + "Google Gemma 3 4B": "google/gemma-3-4b-it", + "Google Gemma 3 4B Allura Glitter": "allura-org/Gemma-3-Glitter-4B", + "Google Gemma 3 4B Nidum Uncensored": "nidum/Nidum-Gemma-3-4B-it-Uncensored", "Alibaba Qwen 3.5 9B": "Qwen/Qwen3.5-9B", - "Alibaba Qwen 3.5 27B": "Qwen/Qwen3.5-27B", - "Alibaba Qwen 3.5 35B-A3B": "Qwen/Qwen3.5-35B-A3B", - "Qwen 3.5 27B Heretic": "coder3101/Qwen3.5-27B-heretic", - "Alibaba Qwen 2.0 VL 2B": "Qwen/Qwen2-VL-2B-Instruct", - "Alibaba Qwen 2.5 Omni 3B": "Qwen/Qwen2.5-Omni-3B", - "Alibaba Qwen 2.5 VL 3B": "Qwen/Qwen2.5-VL-3B-Instruct", - # Qwen2.5-VL Finetunes - "Qwen 2.5 VL 3B Heretic": "coder3101/Qwen2.5-VL-3B-Instruct-heretic", - "Qwen 2.5 VL 7B Heretic": "coder3101/Qwen2.5-VL-7B-Instruct-heretic", - "Qwen 2.5 VL 32B Heretic": "coder3101/Qwen2.5-VL-32B-Instruct-heretic", - "Qwen 2.5 VL 72B Heretic": "coder3101/Qwen2.5-VL-72B-Instruct-heretic", - "Alibaba Qwen 3 VL 2B": "Qwen/Qwen3-VL-2B-Instruct", - f"Alibaba Qwen 3 VL 2B Thinking {ui_symbols.reasoning}": "Qwen/Qwen3-VL-2B-Thinking", - "Alibaba Qwen 3 VL 4B": "Qwen/Qwen3-VL-4B-Instruct", - f"Alibaba Qwen 3 VL 4B Thinking {ui_symbols.reasoning}": "Qwen/Qwen3-VL-4B-Thinking", - "Alibaba Qwen 3 VL 8B": "Qwen/Qwen3-VL-8B-Instruct", - f"Alibaba Qwen 3 VL 8B Thinking {ui_symbols.reasoning}": "Qwen/Qwen3-VL-8B-Thinking", - # Qwen3-VL Finetunes - "Qwen 3 VL 2B Heretic": "coder3101/Qwen3-VL-2B-Instruct-heretic", - f"Qwen 3 VL 2B Thinking Heretic {ui_symbols.reasoning}": "coder3101/Qwen3-VL-2B-Thinking-heretic", - "Qwen 3 VL 4B Heretic": "coder3101/Qwen3-VL-4B-Instruct-heretic", - f"Qwen 3 VL 4B Thinking Heretic {ui_symbols.reasoning}": "coder3101/Qwen3-VL-4B-Thinking-heretic", - "Qwen 3 VL 8B Heretic": "coder3101/Qwen3-VL-8B-Instruct-heretic", - "Qwen 3 VL 32B Heretic v2": "coder3101/Qwen3-VL-32B-Instruct-heretic-v2", - f"Qwen 3 VL 32B Thinking Heretic v2 {ui_symbols.reasoning}": "coder3101/Qwen3-VL-32B-Thinking-heretic-v2", - "Qwen 3 VL 8B Abliterated Caption": "prithivMLmods/Qwen3-VL-8B-Abliterated-Caption-it", - "XiaomiMiMo MiMo VL 7B RL": "XiaomiMiMo/MiMo-VL-7B-RL-2508", # 8.3GB - "Huggingface Smol VL2 0.5B": "HuggingFaceTB/SmolVLM-500M-Instruct", - "Huggingface Smol VL2 2B": "HuggingFaceTB/SmolVLM-Instruct", - "Apple FastVLM 0.5B": "apple/FastVLM-0.5B", - "Apple FastVLM 1.5B": "apple/FastVLM-1.5B", - "Apple FastVLM 7B": "apple/FastVLM-7B", - "Microsoft Florence 2 Base": "florence-community/Florence-2-base-ft", # 0.5GB - "Microsoft Florence 2 Large": "florence-community/Florence-2-large-ft", # 1.5GB - "MiaoshouAI PromptGen 1.5 Base": "Disty0/Florence-2-base-PromptGen-v1.5", # 0.5GB - "MiaoshouAI PromptGen 1.5 Large": "Disty0/Florence-2-large-PromptGen-v1.5", # 1.5GB - "MiaoshouAI PromptGen 2.0 Base": "Disty0/Florence-2-base-PromptGen-v2.0", # 0.5GB - "MiaoshouAI PromptGen 2.0 Large": "Disty0/Florence-2-large-PromptGen-v2.0", # 1.5GB - "CogFlorence 2.0 Large": "thwri/CogFlorence-2-Large-Freeze", # 1.6GB + "Alibaba Qwen 3.5 9B Trohrbaugh Heretic": "trohrbaugh/Qwen3.5-9B-heretic-v2", + "Alibaba Qwen 3.5 4B": "Qwen/Qwen3.5-4B", + "Alibaba Qwen 3.5 2B": "Qwen/Qwen3.5-2B", + "Alibaba Qwen 3.5 0.8B": "Qwen/Qwen3.5-0.8B", + "JoyTag": "fancyfeast/joytag", + "JoyCaption Beta": "fancyfeast/llama-joycaption-beta-one-hf-llava", + "JoyCaption Alpha": "fancyfeast/llama-joycaption-alpha-two-hf-llava", + f"Moondream 3 Preview {ui_symbols.reasoning}": "moondream/moondream3-preview", + f"Moondream 2 {ui_symbols.reasoning}": "vikhyatk/moondream2", + "Microsoft Florence 2 Large": "florence-community/Florence-2-large-ft", + "Microsoft Florence 2 Base": "florence-community/Florence-2-base-ft", + "MiaoshouAI PromptGen 2.0 Large": "Disty0/Florence-2-large-PromptGen-v2.0", + "MiaoshouAI PromptGen 2.0 Base": "Disty0/Florence-2-base-PromptGen-v2.0", "CogFlorence 2.2 Large": "thwri/CogFlorence-2.2-Large", # 1.6GB - f"Moondream 2 {ui_symbols.reasoning}": "vikhyatk/moondream2", # 3.7GB - f"Moondream 3 Preview {ui_symbols.reasoning}": "moondream/moondream3-preview", # 9.3GB (gated) - "Google Pix Textcaps": "google/pix2struct-textcaps-base", # 1.1GB - "Google PaliGemma 2 3B": "google/paligemma2-3b-pt-224", - "Salesforce BLIP Base": "Salesforce/blip-vqa-base", # 1.5GB - "Salesforce BLIP Large": "Salesforce/blip-vqa-capfilt-large", # 1.5GB - "Microsoft GIT TextCaps Base": "microsoft/git-base-textcaps", # 0.7GB - "Microsoft GIT VQA Base": "microsoft/git-base-vqav2", # 0.7GB - "Microsoft GIT VQA Large": "microsoft/git-large-vqav2", # 1.6GB - "ToriiGate 0.4 2B": "Minthy/ToriiGate-v0.4-2B", + # secondary or older + "XiaomiMiMo MiMo VL 7B RL": "XiaomiMiMo/MiMo-VL-7B-RL-2508", + "ViLT Base": "dandelin/vilt-b32-finetuned-vqa", "ToriiGate 0.4 7B": "Minthy/ToriiGate-v0.4-7B", - "ViLT Base": "dandelin/vilt-b32-finetuned-vqa", # 0.5GB - "JoyCaption Alpha": "fancyfeast/llama-joycaption-alpha-two-hf-llava", # 17.4GB - "JoyCaption Beta": "fancyfeast/llama-joycaption-beta-one-hf-llava", # 17.4GB - "JoyTag": "fancyfeast/joytag", # 0.7GB - "AIDC Ovis2 1B": "AIDC-AI/Ovis2-1B", - "AIDC Ovis2 2B": "AIDC-AI/Ovis2-2B", - "AIDC Ovis2 4B": "AIDC-AI/Ovis2-4B", - "ByteDance Sa2VA 1B": "ByteDance/Sa2VA-1B", + "ToriiGate 0.4 2B": "Minthy/ToriiGate-v0.4-2B", + "Salesforce BLIP Large": "Salesforce/blip-vqa-capfilt-large", + "Salesforce BLIP Base": "Salesforce/blip-vqa-base", + "Mistral Small 3.2 24B Coder Heretic": "coder3101/Mistral-Small-3.2-24B-Instruct-2506-heretic", + "Microsoft GIT VQA Large": "microsoft/git-large-vqav2", + "Microsoft GIT VQA Base": "microsoft/git-base-vqav2", + "Microsoft GIT TextCaps Base": "microsoft/git-base-textcaps", + "MiaoshouAI PromptGen 1.5 Large": "Disty0/Florence-2-large-PromptGen-v1.5", + "MiaoshouAI PromptGen 1.5 Base": "Disty0/Florence-2-base-PromptGen-v1.5", + "Huggingface Smol VL2 2B": "HuggingFaceTB/SmolVLM-Instruct", + "Huggingface Smol VL2 0.5B": "HuggingFaceTB/SmolVLM-500M-Instruct", + "Google Pix Textcaps": "google/pix2struct-textcaps-base", + "Google PaliGemma 2 3B": "google/paligemma2-3b-pt-224", + "CogFlorence 2.0 Large": "thwri/CogFlorence-2-Large-Freeze", # 1.6GB "ByteDance Sa2VA 4B": "ByteDance/Sa2VA-4B", - # Mistral Finetunes - "Mistral Small 3.2 24B Heretic": "coder3101/Mistral-Small-3.2-24B-Instruct-2506-heretic", + "ByteDance Sa2VA 1B": "ByteDance/Sa2VA-1B", + "Apple FastVLM 7B": "apple/FastVLM-7B", + "Apple FastVLM 1.5B": "apple/FastVLM-1.5B", + "Apple FastVLM 0.5B": "apple/FastVLM-0.5B", + f"Alibaba Qwen 3 VL 8B Thinking {ui_symbols.reasoning}": "Qwen/Qwen3-VL-8B-Thinking", + "Alibaba Qwen 3 VL 8B Coder Heretic": "coder3101/Qwen3-VL-8B-Instruct-heretic", + "Alibaba Qwen 3 VL 8B Abliterated Caption": "prithivMLmods/Qwen3-VL-8B-Abliterated-Caption-it", + "Alibaba Qwen 3 VL 8B": "Qwen/Qwen3-VL-8B-Instruct", + f"Alibaba Qwen 3 VL 4B Thinking {ui_symbols.reasoning}": "Qwen/Qwen3-VL-4B-Thinking", + f"Alibaba Qwen 3 VL 4B Thinking Coder Heretic {ui_symbols.reasoning}": "coder3101/Qwen3-VL-4B-Thinking-heretic", + "Alibaba Qwen 3 VL 4B Coder Heretic": "coder3101/Qwen3-VL-4B-Instruct-heretic", + "Alibaba Qwen 3 VL 4B": "Qwen/Qwen3-VL-4B-Instruct", + f"Alibaba Qwen 3 VL 32B Thinking Coder Heretic v2 {ui_symbols.reasoning}": "coder3101/Qwen3-VL-32B-Thinking-heretic-v2", + f"Alibaba Qwen 3 VL 2B Thinking {ui_symbols.reasoning}": "Qwen/Qwen3-VL-2B-Thinking", + f"Alibaba Qwen 3 VL 2B Thinking Coder Heretic {ui_symbols.reasoning}": "coder3101/Qwen3-VL-2B-Thinking-heretic", + "Alibaba Qwen 3 VL 2B Coder Heretic": "coder3101/Qwen3-VL-2B-Instruct-heretic", + "Alibaba Qwen 3 VL 2B": "Qwen/Qwen3-VL-2B-Instruct", + "Alibaba Qwen 2.5 VL 7B Coder Heretic": "coder3101/Qwen2.5-VL-7B-Instruct-heretic", + "Alibaba Qwen 2.5 VL 3B Coder Heretic": "coder3101/Qwen2.5-VL-3B-Instruct-heretic", + "Alibaba Qwen 2.5 VL 3B": "Qwen/Qwen2.5-VL-3B-Instruct", + "Alibaba Qwen 2.5 Omni 3B": "Qwen/Qwen2.5-Omni-3B", + "Alibaba Qwen 2.0 VL 2B": "Qwen/Qwen2-VL-2B-Instruct", + "AIDC Ovis2 4B": "AIDC-AI/Ovis2-4B", + "AIDC Ovis2 2B": "AIDC-AI/Ovis2-2B", + "AIDC Ovis2 1B": "AIDC-AI/Ovis2-1B", + # cloud + f"Google Gemini 3.5 Flash {ui_symbols.cloud}": "google/gemini-3.5-flash", f"Google Gemini 3.1 Pro {ui_symbols.cloud}": "gemini-3.1-pro-preview", - f"Google Gemini 3.1 Flash Lite {ui_symbols.cloud}": "gemini-3.1-flash-lite-preview", - f"Google Gemini 3.0 Flash {ui_symbols.cloud}": "gemini-3-flash-preview", + f"Google Gemini 3.1 Flash Lite {ui_symbols.cloud}": "gemini-3.1-flash-lite", + f"Google Gemini 3.1 Flash Lite Preview {ui_symbols.cloud}": "gemini-3.1-flash-lite-preview", f"Google Gemini 2.5 Pro {ui_symbols.cloud}": "gemini-2.5-pro", f"Google Gemini 2.5 Flash {ui_symbols.cloud}": "gemini-2.5-flash", + f"Google Gemini 2.5 Flash Lite {ui_symbols.cloud}": "gemini-2.5-flash-lite", } # Default model @@ -149,6 +145,14 @@ vlm_prompt_mapping = { "Detect Gaze": "DETECT_GAZE", } +# Reverse prompt-mapping +vlm_base_text_prompt_suffix = "Describe the main subject, scene layout, perspective, background, colors, lighting, textures, visible objects and their relationships, composition, and any clearly visible art style or medium. Keep the language vivid but objective, avoid subjective commentary, and output only a single caption text suitable for image generation." +vlm_prompt_reverse_mapping = { + 'CAPTION': f"Provide a short and concise, visually descriptive caption suitable for image generation. {vlm_base_text_prompt_suffix}", + 'DETAILED CAPTION': f"Provide a detailed visual caption suitable for image generation. {vlm_base_text_prompt_suffix}", + 'MORE DETAILED CAPTION': f"Provide a very detailed extended visual caption suitable for image generation. {vlm_base_text_prompt_suffix}", +} + # Placeholder hints for prompt field based on selected question vlm_prompt_placeholders = { "Use Prompt": "Enter your question or instruction for the model", diff --git a/modules/caption/moondream3.py b/modules/caption/moondream3.py index 7859f0f06..399836100 100644 --- a/modules/caption/moondream3.py +++ b/modules/caption/moondream3.py @@ -49,7 +49,7 @@ def load_model(repo: str): global moondream3_model, loaded # pylint: disable=global-statement if moondream3_model is None or loaded != repo: - log.debug(f'Caption load: vlm="{repo}"') + log.debug(f'LLM load: vlm="{repo}"') moondream3_model = None moondream3_model = transformers.AutoModelForCausalLM.from_pretrained( @@ -100,7 +100,7 @@ def encode_image(image: Image.Image, cache_key: str | None = None): """ if cache_key and cache_key in image_cache: image_cache.move_to_end(cache_key) # LRU: mark as recently used - debug(f'VQA caption: handler=moondream3 using cached encoding for cache_key="{cache_key}"') + debug(f'LLM: handler=moondream3 using cached encoding for cache_key="{cache_key}"') return image_cache[cache_key] model = load_model(loaded) @@ -112,8 +112,8 @@ def encode_image(image: Image.Image, cache_key: str | None = None): image_cache[cache_key] = encoded while len(image_cache) > IMAGE_CACHE_MAX: evicted_key, _ = image_cache.popitem(last=False) # Evict oldest - debug(f'VQA caption: handler=moondream3 evicted cache_key="{evicted_key}" cache_size={len(image_cache)}') - debug(f'VQA caption: handler=moondream3 cached encoding cache_key="{cache_key}" cache_size={len(image_cache)}') + debug(f'LLM: handler=moondream3 evicted cache_key="{evicted_key}" cache_size={len(image_cache)}') + debug(f'LLM: handler=moondream3 cached encoding cache_key="{cache_key}" cache_size={len(image_cache)}') return encoded @@ -148,7 +148,7 @@ def query(image: Image.Image, question: str, repo: str, stream: bool = False, if max_tokens is not None: settings['max_tokens'] = max_tokens - debug(f'VQA caption: handler=moondream3 method=query question="{question}" stream={stream} settings={settings}') + debug(f'LLM: handler=moondream3 method=query question="{question}" stream={stream} settings={settings}') # Use cached encoding if requested if use_cache: @@ -169,12 +169,12 @@ def query(image: Image.Image, question: str, repo: str, stream: bool = False, # Log response structure (for non-streaming) if not stream: if isinstance(response, dict): - debug(f'VQA caption: handler=moondream3 response_type=dict keys={list(response.keys())}') + debug(f'LLM: handler=moondream3 response_type=dict keys={list(response.keys())}') if 'reasoning' in response: reasoning_text = response['reasoning'].get('text', '')[:100] + '...' if len(response['reasoning'].get('text', '')) > 100 else response['reasoning'].get('text', '') - debug(f'VQA caption: handler=moondream3 reasoning="{reasoning_text}"') + debug(f'LLM: handler=moondream3 reasoning="{reasoning_text}"') if 'answer' in response: - debug(f'VQA caption: handler=moondream3 answer="{response["answer"]}"') + debug(f'LLM: handler=moondream3 answer="{response["answer"]}"') return response @@ -207,7 +207,7 @@ def caption(image: Image.Image, repo: str, length: str = 'normal', stream: bool if max_tokens is not None: settings['max_tokens'] = max_tokens - debug(f'VQA caption: handler=moondream3 method=caption length={length} stream={stream} settings={settings}') + debug(f'LLM: handler=moondream3 method=caption length={length} stream={stream} settings={settings}') with devices.inference_context(): response = model.caption( @@ -219,7 +219,7 @@ def caption(image: Image.Image, repo: str, length: str = 'normal', stream: bool # Log response structure (for non-streaming) if not stream and isinstance(response, dict): - debug(f'VQA caption: handler=moondream3 response_type=dict keys={list(response.keys())}') + debug(f'LLM: handler=moondream3 response_type=dict keys={list(response.keys())}') return response @@ -239,21 +239,21 @@ def point(image: Image.Image, object_name: str, repo: str): """ model = load_model(repo) - debug(f'VQA caption: handler=moondream3 method=point object_name="{object_name}"') + debug(f'LLM: handler=moondream3 method=point object_name="{object_name}"') with devices.inference_context(): result = model.point(image, object_name) - debug(f'VQA caption: handler=moondream3 point_raw_result="{result}" type={type(result)}') + debug(f'LLM: handler=moondream3 point_raw_result="{result}" type={type(result)}') if isinstance(result, dict): - debug(f'VQA caption: handler=moondream3 point_raw_result_keys={list(result.keys())}') + debug(f'LLM: handler=moondream3 point_raw_result_keys={list(result.keys())}') points = vqa_detection.parse_points(result) if points: - debug(f'VQA caption: handler=moondream3 point_result={len(points)} points found') + debug(f'LLM: handler=moondream3 point_result={len(points)} points found') return points - debug('VQA caption: handler=moondream3 point_result=not found') + debug('LLM: handler=moondream3 point_result=not found') return None @@ -276,17 +276,17 @@ def detect(image: Image.Image, object_name: str, repo: str, max_objects: int = 1 """ model = load_model(repo) - debug(f'VQA caption: handler=moondream3 method=detect object_name="{object_name}" max_objects={max_objects}') + debug(f'LLM: handler=moondream3 method=detect object_name="{object_name}" max_objects={max_objects}') with devices.inference_context(): result = model.detect(image, object_name) - debug(f'VQA caption: handler=moondream3 detect_raw_result="{result}" type={type(result)}') + debug(f'LLM: handler=moondream3 detect_raw_result="{result}" type={type(result)}') if isinstance(result, dict): - debug(f'VQA caption: handler=moondream3 detect_raw_result_keys={list(result.keys())}') + debug(f'LLM: handler=moondream3 detect_raw_result_keys={list(result.keys())}') detections = vqa_detection.parse_detections(result, object_name, max_objects) - debug(f'VQA caption: handler=moondream3 detect_result={len(detections)} objects found') + debug(f'LLM: handler=moondream3 detect_result={len(detections)} objects found') return detections @@ -310,7 +310,7 @@ def predict(question: str, image: Image.Image, repo: str, model_name: str | None Response string (detection data stored on VQA singleton instance.last_detection_data) (or generator if stream=True for query/caption modes) """ - debug(f'VQA caption: handler=moondream3 model_name="{model_name}" repo="{repo}" question="{question}" image_size={image.size if image else None} mode={mode} stream={stream}') + debug(f'LLM: handler=moondream3 model_name="{model_name}" repo="{repo}" question="{question}" image_size={image.size if image else None} mode={mode} stream={stream}') # Clean question question = question.replace('<', '').replace('>', '').replace('_', ' ') if question else '' @@ -350,7 +350,7 @@ def predict(question: str, image: Image.Image, repo: str, model_name: str | None else: mode = 'query' - debug(f'VQA caption: handler=moondream3 mode_selected={mode}') + debug(f'LLM: handler=moondream3 mode_selected={mode}') # Dispatch to appropriate method try: @@ -367,7 +367,7 @@ def predict(question: str, image: Image.Image, repo: str, model_name: str | None object_name = re.sub(rf'\b{phrase}\b', '', object_name, flags=re.IGNORECASE) object_name = re.sub(r'[?.!,]', '', object_name).strip() object_name = re.sub(r'^\s*the\s+', '', object_name, flags=re.IGNORECASE) - debug(f'VQA caption: handler=moondream3 point_extracted_object="{object_name}"') + debug(f'LLM: handler=moondream3 point_extracted_object="{object_name}"') result = point(image, object_name, repo) if result: from modules.caption import vqa @@ -383,7 +383,7 @@ def predict(question: str, image: Image.Image, repo: str, model_name: str | None object_name = re.sub(r'^\s*the\s+', '', object_name, flags=re.IGNORECASE) if ' and ' in object_name.lower(): object_name = re.split(r'\s+and\s+', object_name, flags=re.IGNORECASE)[0].strip() - debug(f'VQA caption: handler=moondream3 detect_extracted_object="{object_name}"') + debug(f'LLM: handler=moondream3 detect_extracted_object="{object_name}"') results = detect(image, object_name, repo, max_objects=kwargs.get('max_objects', 10)) if results: @@ -396,7 +396,7 @@ def predict(question: str, image: Image.Image, repo: str, model_name: str | None question = "Describe this image." response = query(image, question, repo, stream=stream, use_cache=use_cache, reasoning=thinking_mode) - debug(f'VQA caption: handler=moondream3 response_before_clean="{response}"') + debug(f'LLM: handler=moondream3 response_before_clean="{response}"') return response except Exception as e: @@ -411,7 +411,7 @@ def clear_cache(): """Clear image encoding cache.""" cache_size = len(image_cache) image_cache.clear() - debug(f'VQA caption: handler=moondream3 cleared image cache cache_size_was={cache_size}') + debug(f'LLM: handler=moondream3 cleared image cache cache_size_was={cache_size}') log.debug(f'Moondream3: Cleared image cache ({cache_size} entries)') diff --git a/modules/caption/vqa.py b/modules/caption/vqa.py index 0e80ccbd2..ec7c63d4e 100644 --- a/modules/caption/vqa.py +++ b/modules/caption/vqa.py @@ -12,10 +12,11 @@ from PIL import Image from modules import shared, devices, errors, model_quant, sd_models, sd_models_compile from modules.sd_offload_aux import register_aux, deregister_aux, move_aux_to_gpu, offload_aux from modules.logger import log, console -from modules.caption import vqa_detection -from modules.caption.models_def import vlm_models, vlm_system, vlm_analyze, vlm_default, vlm_prefill, vlm_prompts, vlm_prompt_mapping, vlm_prompt_placeholders, vlm_prompts_common, vlm_prompts_florence, vlm_prompts_moondream, vlm_prompts_moondream2, vlm_prompts_promptgen, analyze_question, get_vlm_repo # pylint: disable=unused-import +from modules.caption import vqa_detection, helpers +from modules.caption.attention import set_attention +from modules.caption.models_def import vlm_models, vlm_system, vlm_analyze, vlm_default, vlm_prefill, vlm_prompts, vlm_prompt_mapping, vlm_prompt_reverse_mapping, vlm_prompt_placeholders, vlm_prompts_common, vlm_prompts_florence, vlm_prompts_moondream, vlm_prompts_moondream2, vlm_prompts_promptgen, analyze_question, get_vlm_repo # pylint: disable=unused-import + -# Debug logging - function-based to avoid circular import debug_enabled = os.environ.get('SD_CAPTION_DEBUG', None) is not None @@ -180,7 +181,7 @@ def keep_think_block_open(text_prompt: str) -> str: while end_close < len(text_prompt) and text_prompt[end_close] in ('\r', '\n'): end_close += 1 trimmed_prompt = text_prompt[:close_index] + text_prompt[end_close:] - debug('VQA caption: keep_think_block_open applied to prompt segment near assistant reply') + debug('LLM: keep_think_block_open applied to prompt segment near assistant reply') return trimmed_prompt @@ -195,7 +196,7 @@ def b64(image): def clean(response, question, prefill=None): - strip = ['---', '\r', '\t', '**', '"', '"', '"', 'Assistant:', 'Caption:', '<|im_end|>', ''] + strip = ['---', '\r', '\t', '**', '"', '"', '"', 'Assistant:', 'LLM:', '<|im_end|>', ''] if isinstance(response, str): response = response.strip() elif isinstance(response, dict): @@ -306,7 +307,7 @@ def get_keep_prefill(): return shared.opts.caption_vlm_keep_prefill -def get_kwargs(): +def get_kwargs(model): """Build generation kwargs from settings with per-request overrides from VQA instance. Checks the singleton VQA instance's generation_overrides for per-request overrides. @@ -323,6 +324,7 @@ def get_kwargs(): temperature = overrides.get('temperature') if overrides.get('temperature') is not None else shared.opts.caption_vlm_temperature top_k = overrides.get('top_k') if overrides.get('top_k') is not None else shared.opts.caption_vlm_top_k top_p = overrides.get('top_p') if overrides.get('top_p') is not None else shared.opts.caption_vlm_top_p + custom_args = overrides.get('custom_args') if overrides.get('custom_args') is not None else shared.opts.caption_vlm_custom_args kwargs = { 'max_new_tokens': max_tokens, @@ -336,6 +338,11 @@ def get_kwargs(): kwargs['top_k'] = top_k if top_p > 0: kwargs['top_p'] = top_p + + custom = helpers.get_custom_args(model, custom_args) + for k, v in custom.items(): + kwargs[k] = v + return kwargs @@ -460,7 +467,7 @@ class VQA: def _load_fastvlm(self, repo: str): """Load FastVLM model and tokenizer.""" if self.model is None or self.loaded != repo: - log.debug(f'Caption load: vlm="{repo}"') + log.debug(f'LLM load: vlm="{repo}"') self._unload_current() quant_args = model_quant.create_config(module='LLM') self.processor = transformers.AutoTokenizer.from_pretrained(repo, trust_remote_code=True, cache_dir=shared.opts.hfcache_dir) @@ -475,11 +482,12 @@ class VQA: ) self.model.eval() register_aux('vqa', self.model) + set_attention(self.model) self.loaded = repo devices.torch_gc() def _fastvlm(self, question: str, image: Image.Image, repo: str, model_name: str | None = None): - debug(f'VQA caption: handler=fastvlm model_name="{model_name}" repo="{repo}" question="{question}" image_size={image.size if image else None}') + debug(f'LLM: handler=fastvlm model_name="{model_name}" repo="{repo}" question="{question}" image_size={image.size if image else None}') self._load_fastvlm(repo) move_aux_to_gpu('vqa') if len(question) < 2: @@ -519,7 +527,7 @@ class VQA: def _load_qwen(self, repo: str): """Load Qwen VL model and processor.""" if self.model is None or self.loaded != repo: - log.debug(f'Caption load: vlm="{repo}"') + log.debug(f'LLM load: vlm="{repo}"') self._unload_current() if 'Qwen3.5' in repo and re.search(r'-A\d+B', repo): cls_name = transformers.Qwen3_5MoeForConditionalGeneration @@ -552,6 +560,7 @@ class VQA: if 'LLM' in shared.opts.cuda_compile: self.model = sd_models_compile.compile_torch(self.model, apply_to_components=False, op="VQA") register_aux('vqa', self.model) + set_attention(self.model) self.loaded = repo devices.torch_gc() @@ -560,9 +569,14 @@ class VQA: move_aux_to_gpu('vqa') # Get model class name for logging cls_name = self.model.__class__.__name__ - debug(f'VQA caption: handler=qwen model_name="{model_name}" model_class="{cls_name}" repo="{repo}" question="{question}" system_prompt="{system_prompt}" image_size={image.size if image else None}') + debug(f'LLM: handler=qwen model_name="{model_name}" model_class="{cls_name}" repo="{repo}" question="{question}" system_prompt="{system_prompt}" image_size={image.size if image else None}') question = question.replace('<', '').replace('>', '').replace('_', ' ') + if question is not None and len(question) > 4: + if question in vlm_prompt_reverse_mapping: + debug(f'LLM: handler=gemma mapping friendly question="{question}" to internal="{vlm_prompt_reverse_mapping[question]}"') + question = vlm_prompt_reverse_mapping[question] + system_prompt = system_prompt or shared.opts.caption_vlm_system conversation = [ { @@ -591,9 +605,9 @@ class VQA: use_prefill = len(prefill_text) > 0 if debug_enabled: - debug(f'VQA caption: handler=qwen conversation_roles={[msg["role"] for msg in conversation]}') - debug(f'VQA caption: handler=qwen full_conversation={truncate_b64_in_conversation(conversation)}') - debug(f'VQA caption: handler=qwen is_thinking={is_thinking} thinking_mode={thinking_mode} prefill="{prefill_text}"') + debug(f'LLM: handler=qwen conversation_roles={[msg["role"] for msg in conversation]}') + debug(f'LLM: handler=qwen full_conversation={truncate_b64_in_conversation(conversation)}') + debug(f'LLM: handler=qwen is_thinking={is_thinking} thinking_mode={thinking_mode} prefill="{prefill_text}"') # Qwen3.5 uses native enable_thinking parameter in the chat template is_qwen35 = 'qwen3.5' in (model_name or '').lower() or 'qwen3.5' in repo.lower() @@ -608,7 +622,7 @@ class VQA: **template_kwargs, ) except (TypeError, ValueError) as e: - debug(f'VQA caption: handler=qwen chat_template fallback add_generation_prompt=True: {e}') + debug(f'LLM: handler=qwen chat_template fallback add_generation_prompt=True: {e}') text_prompt = self.processor.apply_chat_template(conversation, add_generation_prompt=True) # Manual think handling - skip for Qwen3.5 (template handles it natively) @@ -630,24 +644,28 @@ class VQA: text_prompt += prefill_text if debug_enabled: - debug(f'VQA caption: handler=qwen text_prompt="{text_prompt}"') + debug(f'LLM: handler=qwen text_prompt="{text_prompt}"') inputs = self.processor(text=[text_prompt], images=[image], padding=True, return_tensors="pt") inputs = inputs.to(devices.device, devices.dtype) - gen_kwargs = get_kwargs() - debug(f'VQA caption: handler=qwen generation_kwargs={gen_kwargs} input_ids_shape={inputs.input_ids.shape}') + + gen_kwargs = get_kwargs(self.model) + log.debug(f'LLM: args={gen_kwargs} ids={inputs.input_ids.shape}') + defaults = {k: v for k, v in helpers.get_default_args(self.model).items() if k not in gen_kwargs} + log.debug(f'LLM: defaults={defaults}') + with devices.inference_context(): output_ids = self.model.generate( **inputs, **gen_kwargs, ) - debug(f'VQA caption: handler=qwen output_ids_shape={output_ids.shape}') + debug(f'LLM: handler=qwen output_ids_shape={output_ids.shape}') generated_ids = [ output_ids[len(input_ids):] for input_ids, output_ids in zip(inputs.input_ids, output_ids, strict=False) ] response = self.processor.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True) if debug_enabled: - debug(f'VQA caption: handler=qwen response_before_clean="{response}"') + debug(f'LLM: handler=qwen response_before_clean="{response}"') if len(response) > 0: response[0] = strip_think_xml_tags(response[0], keep=get_keep_thinking()) return response @@ -655,7 +673,7 @@ class VQA: def _load_gemma(self, repo: str): """Load Gemma model and processor.""" if self.model is None or self.loaded != repo: - log.debug(f'Caption load: vlm="{repo}"') + log.debug(f'LLM load: vlm="{repo}"') self._unload_current() if 'gemma-4' in repo.lower(): cls = transformers.Gemma4ForConditionalGeneration @@ -677,6 +695,7 @@ class VQA: self.model = sd_models_compile.compile_torch(self.model, apply_to_components=False, op="VQA") self.processor = transformers.AutoProcessor.from_pretrained(repo, max_pixels=1024*1024, cache_dir=shared.opts.hfcache_dir) register_aux('vqa', self.model) + set_attention(self.model) self.loaded = repo devices.torch_gc() @@ -685,7 +704,7 @@ class VQA: move_aux_to_gpu('vqa') # Get model class name for logging cls_name = self.model.__class__.__name__ - debug(f'VQA caption: handler=gemma model_name="{model_name}" model_class="{cls_name}" repo="{repo}" question="{question}" system_prompt="{system_prompt}" image_size={image.size if image else None}') + debug(f'LLM: handler=gemma model_name="{model_name}" model_class="{cls_name}" repo="{repo}" question="{question}" system_prompt="{system_prompt}" image_size={image.size if image else None}') question = question.replace('<', '').replace('>', '').replace('_', ' ') system_prompt = system_prompt or shared.opts.caption_vlm_system @@ -696,6 +715,9 @@ class VQA: user_content = [] if question is not None and len(question) > 4: + if question in vlm_prompt_reverse_mapping: + debug(f'LLM: handler=gemma mapping friendly question="{question}" to internal="{vlm_prompt_reverse_mapping[question]}"') + question = vlm_prompt_reverse_mapping[question] user_content.append({"type": "text", "text": question}) if image is not None: user_content.append({"type": "image", "image": b64(image)}) @@ -715,14 +737,14 @@ class VQA: "role": "assistant", "content": [{"type": "text", "text": prefill_text}], }) - debug(f'VQA caption: handler=gemma prefill="{prefill_text}"') + debug(f'LLM: handler=gemma prefill="{prefill_text}"') else: - debug('VQA caption: handler=gemma prefill disabled (empty), relying on add_generation_prompt') + debug('LLM: handler=gemma prefill disabled (empty), relying on add_generation_prompt') if debug_enabled: - debug(f'VQA caption: handler=gemma conversation_roles={[msg["role"] for msg in conversation]}') - debug(f'VQA caption: handler=gemma full_conversation={truncate_b64_in_conversation(conversation)}') + debug(f'LLM: handler=gemma conversation_roles={[msg["role"] for msg in conversation]}') + debug(f'LLM: handler=gemma full_conversation={truncate_b64_in_conversation(conversation)}') debug_prefill_mode = 'add_generation_prompt=False continue_final_message=True' if use_prefill else 'add_generation_prompt=True' - debug(f'VQA caption: handler=gemma template_mode={debug_prefill_mode}') + debug(f'LLM: handler=gemma template_mode={debug_prefill_mode}') try: if use_prefill: text_prompt = self.processor.apply_chat_template( @@ -738,7 +760,7 @@ class VQA: tokenize=False, ) except (TypeError, ValueError) as e: - debug(f'VQA caption: handler=gemma chat_template fallback add_generation_prompt=True: {e}') + debug(f'LLM: handler=gemma chat_template fallback add_generation_prompt=True: {e}') text_prompt = self.processor.apply_chat_template( conversation, add_generation_prompt=True, @@ -747,7 +769,7 @@ class VQA: if use_prefill and use_thinking: text_prompt = keep_think_block_open(text_prompt) if debug_enabled: - debug(f'VQA caption: handler=gemma text_prompt="{text_prompt}"') + debug(f'LLM: handler=gemma text_prompt="{text_prompt}"') inputs = self.processor( text=[text_prompt], images=[image], @@ -755,18 +777,22 @@ class VQA: return_tensors="pt", ).to(device=devices.device, dtype=devices.dtype) input_len = inputs["input_ids"].shape[-1] - gen_kwargs = get_kwargs() - debug(f'VQA caption: handler=gemma generation_kwargs={gen_kwargs} input_len={input_len}') + + gen_kwargs = get_kwargs(self.model) + log.debug(f'LLM: args={gen_kwargs} ids={inputs.input_ids.shape}') + defaults = {k: v for k, v in helpers.get_default_args(self.model).items() if k not in gen_kwargs} + log.debug(f'LLM: defaults={defaults}') + with devices.inference_context(): generation = self.model.generate( **inputs, **gen_kwargs, ) - debug(f'VQA caption: handler=gemma output_ids_shape={generation.shape}') + debug(f'LLM: handler=gemma output_ids_shape={generation.shape}') generation = generation[0][input_len:] response = self.processor.decode(generation, skip_special_tokens=True) if debug_enabled: - debug(f'VQA caption: handler=gemma response_before_clean="{response}"') + debug(f'LLM: handler=gemma response_before_clean="{response}"') response = strip_think_xml_tags(response, keep=get_keep_thinking()) return response @@ -774,7 +800,7 @@ class VQA: def _load_mistral(self, repo: str): """Load Mistral3 vision model and processor.""" if self.model is None or self.loaded != repo: - log.debug(f'Caption load: vlm="{repo}"') + log.debug(f'LLM load: vlm="{repo}"') self._unload_current() quant_args = model_quant.create_config(module='LLM') self.model = transformers.Mistral3ForConditionalGeneration.from_pretrained( @@ -790,6 +816,8 @@ class VQA: self.model = sd_models_compile.compile_torch(self.model, apply_to_components=False, op="VQA") self.processor = transformers.AutoProcessor.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir) register_aux('vqa', self.model) + set_attention(self.model) + set_attention(self.model) self.loaded = repo devices.torch_gc() @@ -797,7 +825,7 @@ class VQA: self._load_mistral(repo) move_aux_to_gpu('vqa') cls_name = self.model.__class__.__name__ - debug(f'VQA caption: handler=mistral model_name="{model_name}" model_class="{cls_name}" repo="{repo}" question="{question}" system_prompt="{system_prompt}" image_size={image.size if image else None}') + debug(f'LLM: handler=mistral model_name="{model_name}" model_class="{cls_name}" repo="{repo}" question="{question}" system_prompt="{system_prompt}" image_size={image.size if image else None}') question = question.replace('<', '').replace('>', '').replace('_', ' ') system_prompt = system_prompt or shared.opts.caption_vlm_system @@ -820,8 +848,8 @@ class VQA: conversation.append({"role": "assistant", "content": [{"type": "text", "text": prefill_text}]}) if debug_enabled: - debug(f'VQA caption: handler=mistral conversation_roles={[msg["role"] for msg in conversation]}') - debug(f'VQA caption: handler=mistral full_conversation={truncate_b64_in_conversation(conversation)}') + debug(f'LLM: handler=mistral conversation_roles={[msg["role"] for msg in conversation]}') + debug(f'LLM: handler=mistral full_conversation={truncate_b64_in_conversation(conversation)}') try: if use_prefill: @@ -829,27 +857,31 @@ class VQA: else: text_prompt = self.processor.apply_chat_template(conversation, add_generation_prompt=True, tokenize=False) except (TypeError, ValueError) as e: - debug(f'VQA caption: handler=mistral chat_template fallback: {e}') + debug(f'LLM: handler=mistral chat_template fallback: {e}') text_prompt = self.processor.apply_chat_template(conversation, add_generation_prompt=True, tokenize=False) if debug_enabled: - debug(f'VQA caption: handler=mistral text_prompt="{text_prompt}"') + debug(f'LLM: handler=mistral text_prompt="{text_prompt}"') inputs = self.processor(text=[text_prompt], images=[image], padding=True, return_tensors="pt").to(device=devices.device, dtype=devices.dtype) input_len = inputs["input_ids"].shape[-1] - gen_kwargs = get_kwargs() - debug(f'VQA caption: handler=mistral generation_kwargs={gen_kwargs} input_len={input_len}') + + gen_kwargs = get_kwargs(self.model) + log.debug(f'LLM: args={gen_kwargs} ids={inputs.input_ids.shape}') + defaults = {k: v for k, v in helpers.get_default_args(self.model).items() if k not in gen_kwargs} + log.debug(f'LLM: defaults={defaults}') + with devices.inference_context(): generation = self.model.generate(**inputs, **gen_kwargs) generation = generation[0][input_len:] response = self.processor.decode(generation, skip_special_tokens=True) if debug_enabled: - debug(f'VQA caption: handler=mistral response_before_clean="{response}"') + debug(f'LLM: handler=mistral response_before_clean="{response}"') return response def _load_paligemma(self, repo: str): """Load PaliGemma model and processor.""" if self.model is None or self.loaded != repo: - log.debug(f'Caption load: vlm="{repo}"') + log.debug(f'LLM load: vlm="{repo}"') self._unload_current() self.processor = transformers.PaliGemmaProcessor.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir) self.model = transformers.PaliGemmaForConditionalGeneration.from_pretrained( @@ -861,6 +893,7 @@ class VQA: ) self.model.eval() register_aux('vqa', self.model) + set_attention(self.model) self.loaded = repo devices.torch_gc() @@ -873,7 +906,7 @@ class VQA: with devices.inference_context(): generation = self.model.generate( **model_inputs, - **get_kwargs(), + **get_kwargs(self.model), ) generation = generation[0][input_len:] response = self.processor.decode(generation, skip_special_tokens=True) @@ -882,7 +915,7 @@ class VQA: def _load_ovis(self, repo: str): """Load Ovis model (requires flash-attn).""" if self.model is None or self.loaded != repo: - log.debug(f'Caption load: vlm="{repo}"') + log.debug(f'LLM load: vlm="{repo}"') self._unload_current() # Ovis remote code calls AutoConfig.register("aimv2", ...) at module scope # without exist_ok=True, which fails on reload or when the type is already @@ -903,6 +936,7 @@ class VQA: transformers.AutoConfig.register = _orig self.model.eval() register_aux('vqa', self.model) + set_attention(self.model) self.loaded = repo devices.torch_gc() @@ -910,7 +944,7 @@ class VQA: try: pass # pylint: disable=unused-import except Exception: - log.error(f'Caption: vlm="{repo}" flash-attn is not available') + log.error(f'LLM: vlm="{repo}" flash-attn is not available') return '' self._load_ovis(repo) move_aux_to_gpu('vqa') @@ -935,14 +969,14 @@ class VQA: eos_token_id=self.model.generation_config.eos_token_id, pad_token_id=text_tokenizer.pad_token_id, use_cache=True, - **get_kwargs()) + **get_kwargs(self.model)) response = text_tokenizer.decode(output_ids[0], skip_special_tokens=True) return response def _load_smol(self, repo: str): """Load SmolVLM model and processor.""" if self.model is None or self.loaded != repo: - log.debug(f'Caption load: vlm="{repo}"') + log.debug(f'LLM load: vlm="{repo}"') self._unload_current() quant_args = model_quant.create_config(module='LLM') self.model = transformers.AutoModelForImageTextToText.from_pretrained( @@ -958,6 +992,7 @@ class VQA: if 'LLM' in shared.opts.cuda_compile: self.model = sd_models_compile.compile_torch(self.model, apply_to_components=False, op="VQA") register_aux('vqa', self.model) + set_attention(self.model) self.loaded = repo devices.torch_gc() @@ -966,7 +1001,7 @@ class VQA: move_aux_to_gpu('vqa') # Get model class name for logging cls_name = self.model.__class__.__name__ - debug(f'VQA caption: handler=smol model_name="{model_name}" model_class="{cls_name}" repo="{repo}" question="{question}" system_prompt="{system_prompt}" image_size={image.size if image else None}') + debug(f'LLM: handler=smol model_name="{model_name}" model_class="{cls_name}" repo="{repo}" question="{question}" system_prompt="{system_prompt}" image_size={image.size if image else None}') question = question.replace('<', '').replace('>', '').replace('_', ' ') system_prompt = system_prompt or shared.opts.caption_vlm_system @@ -995,14 +1030,14 @@ class VQA: "role": "assistant", "content": [{"type": "text", "text": prefill_text}], }) - debug(f'VQA caption: handler=smol prefill="{prefill_text}"') + debug(f'LLM: handler=smol prefill="{prefill_text}"') else: - debug('VQA caption: handler=smol prefill disabled (empty), relying on add_generation_prompt') + debug('LLM: handler=smol prefill disabled (empty), relying on add_generation_prompt') if debug_enabled: - debug(f'VQA caption: handler=smol conversation_roles={[msg["role"] for msg in conversation]}') - debug(f'VQA caption: handler=smol full_conversation={truncate_b64_in_conversation(conversation)}') + debug(f'LLM: handler=smol conversation_roles={[msg["role"] for msg in conversation]}') + debug(f'LLM: handler=smol full_conversation={truncate_b64_in_conversation(conversation)}') debug_prefill_mode = 'add_generation_prompt=False continue_final_message=True' if use_prefill else 'add_generation_prompt=True' - debug(f'VQA caption: handler=smol template_mode={debug_prefill_mode}') + debug(f'LLM: handler=smol template_mode={debug_prefill_mode}') try: if use_prefill: text_prompt = self.processor.apply_chat_template( @@ -1013,25 +1048,29 @@ class VQA: else: text_prompt = self.processor.apply_chat_template(conversation, add_generation_prompt=True) except (TypeError, ValueError) as e: - debug(f'VQA caption: handler=smol chat_template fallback add_generation_prompt=True: {e}') + debug(f'LLM: handler=smol chat_template fallback add_generation_prompt=True: {e}') text_prompt = self.processor.apply_chat_template(conversation, add_generation_prompt=True) if use_prefill and use_thinking: text_prompt = keep_think_block_open(text_prompt) if debug_enabled: - debug(f'VQA caption: handler=smol text_prompt="{text_prompt}"') + debug(f'LLM: handler=smol text_prompt="{text_prompt}"') inputs = self.processor(text=text_prompt, images=[image], padding=True, return_tensors="pt") inputs = inputs.to(devices.device, devices.dtype) - gen_kwargs = get_kwargs() - debug(f'VQA caption: handler=smol generation_kwargs={gen_kwargs}') + + gen_kwargs = get_kwargs(self.model) + log.debug(f'LLM: args={gen_kwargs} ids={inputs.input_ids.shape}') + defaults = {k: v for k, v in helpers.get_default_args(self.model).items() if k not in gen_kwargs} + log.debug(f'LLM: defaults={defaults}') + with devices.inference_context(): output_ids = self.model.generate( **inputs, **gen_kwargs, ) - debug(f'VQA caption: handler=smol output_ids_shape={output_ids.shape}') + debug(f'LLM: handler=smol output_ids_shape={output_ids.shape}') response = self.processor.batch_decode(output_ids, skip_special_tokens=True) if debug_enabled: - debug(f'VQA caption: handler=smol response_before_clean="{response}"') + debug(f'LLM: handler=smol response_before_clean="{response}"') if len(response) > 0: response[0] = strip_think_xml_tags(response[0], keep=get_keep_thinking()) @@ -1040,7 +1079,7 @@ class VQA: def _load_git(self, repo: str): """Load Microsoft GIT model and processor.""" if self.model is None or self.loaded != repo: - log.debug(f'Caption load: vlm="{repo}"') + log.debug(f'LLM load: vlm="{repo}"') self._unload_current() self.model = transformers.GitForCausalLM.from_pretrained( repo, @@ -1052,6 +1091,7 @@ class VQA: self.model.eval() self.processor = transformers.GitProcessor.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir) register_aux('vqa', self.model) + set_attention(self.model) self.loaded = repo devices.torch_gc() @@ -1074,7 +1114,7 @@ class VQA: def _load_blip(self, repo: str): """Load Salesforce BLIP model and processor.""" if self.model is None or self.loaded != repo: - log.debug(f'Caption load: vlm="{repo}"') + log.debug(f'LLM load: vlm="{repo}"') self._unload_current() self.model = transformers.BlipForQuestionAnswering.from_pretrained( repo, @@ -1086,6 +1126,7 @@ class VQA: self.model.eval() self.processor = transformers.BlipProcessor.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir) register_aux('vqa', self.model) + set_attention(self.model) self.loaded = repo devices.torch_gc() @@ -1102,7 +1143,7 @@ class VQA: def _load_vilt(self, repo: str): """Load ViLT model and processor.""" if self.model is None or self.loaded != repo: - log.debug(f'Caption load: vlm="{repo}"') + log.debug(f'LLM load: vlm="{repo}"') self._unload_current() self.model = transformers.ViltForQuestionAnswering.from_pretrained( repo, @@ -1114,6 +1155,7 @@ class VQA: self.model.eval() self.processor = transformers.ViltProcessor.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir) register_aux('vqa', self.model) + set_attention(self.model) self.loaded = repo devices.torch_gc() @@ -1132,7 +1174,7 @@ class VQA: def _load_pix(self, repo: str): """Load Pix2Struct model and processor.""" if self.model is None or self.loaded != repo: - log.debug(f'Caption load: vlm="{repo}"') + log.debug(f'LLM load: vlm="{repo}"') self._unload_current() self.model = transformers.Pix2StructForConditionalGeneration.from_pretrained( repo, @@ -1144,6 +1186,7 @@ class VQA: self.model.eval() self.processor = transformers.Pix2StructProcessor.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir) register_aux('vqa', self.model) + set_attention(self.model) self.loaded = repo devices.torch_gc() @@ -1163,7 +1206,7 @@ class VQA: def _load_moondream(self, repo: str): """Load Moondream 2 model and tokenizer.""" if self.model is None or self.loaded != repo: - log.debug(f'Caption load: vlm="{repo}"') + log.debug(f'LLM load: vlm="{repo}"') self._unload_current() self.model = transformers.AutoModelForCausalLM.from_pretrained( repo, @@ -1178,10 +1221,11 @@ class VQA: self.loaded = repo self.model.eval() # required: trust_remote_code model register_aux('vqa', self.model) + set_attention(self.model) devices.torch_gc() def _moondream(self, question: str, image: Image.Image, repo: str, model_name: str | None = None, thinking_mode: bool = False): - debug(f'VQA caption: handler=moondream model_name="{model_name}" repo="{repo}" question="{question}" thinking_mode={thinking_mode}') + debug(f'LLM: handler=moondream model_name="{model_name}" repo="{repo}" question="{question}" thinking_mode={thinking_mode}') self._load_moondream(repo) move_aux_to_gpu('vqa') question = question.replace('<', '').replace('>', '').replace('_', ' ') @@ -1196,9 +1240,9 @@ class VQA: target = question[9:].strip() if question.lower().startswith('point at ') else '' if not target: return "Please specify an object to locate" - debug(f'VQA caption: handler=moondream method=point target="{target}"') + debug(f'LLM: handler=moondream method=point target="{target}"') result = self.model.point(image, target) - debug(f'VQA caption: handler=moondream point_raw_result={result}') + debug(f'LLM: handler=moondream point_raw_result={result}') points = vqa_detection.parse_points(result) if points: self.last_detection_data = {'points': points} @@ -1206,13 +1250,13 @@ class VQA: return "Object not found" elif question == 'DETECT_GAZE' or question.lower() == 'detect gaze': # Must be checked before generic 'detect ' prefix to avoid matching as detect target="Gaze" - debug('VQA caption: handler=moondream method=detect_gaze') + debug('LLM: handler=moondream method=detect_gaze') faces = self.model.detect(image, "face") - debug(f'VQA caption: handler=moondream detect_gaze faces={faces}') + debug(f'LLM: handler=moondream detect_gaze faces={faces}') if faces.get('objects'): eye_x, eye_y = vqa_detection.calculate_eye_position(faces['objects'][0]) result = self.model.detect_gaze(image, eye=(eye_x, eye_y)) - debug(f'VQA caption: handler=moondream detect_gaze result={result}') + debug(f'LLM: handler=moondream detect_gaze result={result}') if result.get('gaze'): gaze = result['gaze'] self.last_detection_data = {'points': [(gaze['x'], gaze['y'])]} @@ -1222,22 +1266,22 @@ class VQA: target = question[7:].strip() if question.lower().startswith('detect ') else '' if not target: return "Please specify an object to detect" - debug(f'VQA caption: handler=moondream method=detect target="{target}"') + debug(f'LLM: handler=moondream method=detect target="{target}"') result = self.model.detect(image, target) - debug(f'VQA caption: handler=moondream detect_raw_result={result}') + debug(f'LLM: handler=moondream detect_raw_result={result}') detections = vqa_detection.parse_detections(result, target) if detections: self.last_detection_data = {'detections': detections} return vqa_detection.format_detections_text(detections, include_confidence=False) return "No objects detected" else: - debug(f'VQA caption: handler=moondream method=query question="{question}" reasoning={thinking_mode}') + debug(f'LLM: handler=moondream method=query question="{question}" reasoning={thinking_mode}') result = self.model.query(image, question, reasoning=thinking_mode) response = result['answer'] - debug(f'VQA caption: handler=moondream query_result keys={list(result.keys()) if isinstance(result, dict) else "not dict"}') + debug(f'LLM: handler=moondream query_result keys={list(result.keys()) if isinstance(result, dict) else "not dict"}') if thinking_mode and 'reasoning' in result: reasoning_text = result['reasoning'].get('text', '') if isinstance(result['reasoning'], dict) else str(result['reasoning']) - debug(f'VQA caption: handler=moondream reasoning_text="{reasoning_text[:100]}..."') + debug(f'LLM: handler=moondream reasoning_text="{reasoning_text[:100]}..."') if get_keep_thinking(): response = f"Reasoning:\n{reasoning_text}\n\nAnswer:\n{response}" # When keep_thinking is False, just use the answer (reasoning is discarded) @@ -1263,7 +1307,7 @@ class VQA: effective_revision = revision_from_repo if self.model is None or self.loaded != cache_key: - log.debug(f'Caption load: vlm="{repo_name}" revision="{effective_revision}" path="{shared.opts.hfcache_dir}"') + log.debug(f'LLM load: vlm="{repo_name}" revision="{effective_revision}" path="{shared.opts.hfcache_dir}"') self._unload_current() transformers.dynamic_module_utils.get_imports = get_imports quant_args = model_quant.create_config(module='LLM') @@ -1280,6 +1324,7 @@ class VQA: self.processor = transformers.AutoProcessor.from_pretrained(repo_name, max_pixels=1024*1024, trust_remote_code=True, revision=effective_revision, cache_dir=shared.opts.hfcache_dir) transformers.dynamic_module_utils.get_imports = _get_imports register_aux('vqa', self.model) + set_attention(self.model) self.loaded = cache_key devices.torch_gc() @@ -1290,11 +1335,11 @@ class VQA: task = question.split('>', 1)[0] + '>' else: task = '' - debug(f'VQA caption: handler=florence model_name="{model_name}" repo="{repo}" task="{task}" question="{question}" image_size={image.size}') + debug(f'LLM: handler=florence model_name="{model_name}" repo="{repo}" task="{task}" question="{question}" image_size={image.size}') inputs = self.processor(text=task, images=image, return_tensors="pt") input_ids = inputs['input_ids'].to(devices.device) pixel_values = inputs['pixel_values'].to(devices.device, devices.dtype) - debug(f'VQA caption: handler=florence input_ids={input_ids.shape} pixel_values={pixel_values.shape} dtype={pixel_values.dtype}') + debug(f'LLM: handler=florence input_ids={input_ids.shape} pixel_values={pixel_values.shape} dtype={pixel_values.dtype}') # Florence-2 requires beam search, not sampling - sampling causes probability tensor errors overrides = _get_overrides() max_tokens = overrides.get('max_tokens') if overrides.get('max_tokens') is not None else shared.opts.caption_vlm_max_length @@ -1303,8 +1348,8 @@ class VQA: if getattr(self.model.config, 'decoder_start_token_id', None) is None: bos_token_id = getattr(self.processor.tokenizer, 'bos_token_id', None) or 0 gen_kwargs['decoder_start_token_id'] = bos_token_id - debug(f'VQA caption: handler=florence setting decoder_start_token_id={bos_token_id}') - debug(f'VQA caption: handler=florence generation_kwargs={gen_kwargs}') + debug(f'LLM: handler=florence setting decoder_start_token_id={bos_token_id}') + debug(f'LLM: handler=florence generation_kwargs={gen_kwargs}') with devices.inference_context(), devices.bypass_sdpa_hijacks(): generated_ids = self.model.generate( input_ids=input_ids, @@ -1312,11 +1357,11 @@ class VQA: **gen_kwargs, ) generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=False)[0] - debug(f'VQA caption: handler=florence generated_text="{generated_text}"') + debug(f'LLM: handler=florence generated_text="{generated_text}"') # task="task" is intentional: produces {'task': text} which both parse_florence_detections and # format_florence_response handle via explicit 'task' key fallbacks, avoiding task-token-specific keys response = self.processor.post_process_generation(generated_text, task="task", image_size=(image.width, image.height)) - debug(f'VQA caption: handler=florence raw_response={response}') + debug(f'LLM: handler=florence raw_response={response}') return response def _load_sa2(self, repo: str): @@ -1339,6 +1384,7 @@ class VQA: cache_dir=shared.opts.hfcache_dir, ) register_aux('vqa', self.model) + set_attention(self.model) self.loaded = repo devices.torch_gc() @@ -1371,6 +1417,7 @@ class VQA: prefill: str | None = None, thinking_mode: bool | None = None, quiet: bool = False, + custom_args: str | None = None, generation_kwargs: dict | None = None, ) -> str: """ @@ -1407,7 +1454,7 @@ class VQA: if image.mode != 'RGB': image = image.convert('RGB') if image is None: - log.error(f'VQA caption: model="{model_name}" error="No input image provided"') + log.error(f'LLM: model="{model_name}" error="No input image provided"') self._generation_overrides = None shared.state.end(jobid) return 'Error: No input image provided. Please upload or select an image.' @@ -1416,7 +1463,7 @@ class VQA: if question.lower() == "use prompt": # Use content from Prompt field directly - requires user input if not prompt or len(prompt.strip()) < 2: - log.error(f'VQA caption: model="{model_name}" error="Please enter a prompt"') + log.error(f'LLM: model="{model_name}" error="Please enter a prompt"') self._generation_overrides = None shared.state.end(jobid) return 'Error: Please enter a question or instruction in the Prompt field.' @@ -1427,7 +1474,7 @@ class VQA: if raw_mapping in ("POINT_MODE", "DETECT_MODE"): # These modes require user input in the prompt field if not prompt or len(prompt.strip()) < 2: - log.error(f'VQA caption: model="{model_name}" error="Please specify what to find in the prompt field"') + log.error(f'LLM: model="{model_name}" error="Please specify what to find in the prompt field"') self._generation_overrides = None shared.state.end(jobid) return 'Error: Please specify what to find in the prompt field (e.g., "the red car" or "faces").' @@ -1436,12 +1483,12 @@ class VQA: # else: question is already an internal token or custom text if model_name is None: - log.error(f'Caption: type=vlm model="{model_name}" no model selected') + log.error(f'LLM: type=vlm model="{model_name}" no model selected') shared.state.end(jobid) return '' vqa_model = get_vlm_repo(model_name) if vqa_model == model_name and model_name not in vlm_models.values(): - log.error(f'Caption: type=vlm model="{model_name}" unknown') + log.error(f'LLM: type=vlm model="{model_name}" unknown') shared.state.end(jobid) return '' if self.model is None or self.loaded != vqa_model: @@ -1477,7 +1524,7 @@ class VQA: florence_detections = vqa_detection.parse_florence_detections(answer, image.size if image else None) if florence_detections: self.last_detection_data = {'detections': florence_detections} - debug(f'VQA caption: handler=florence parsed {len(florence_detections)} detections') + debug(f'LLM: handler=florence parsed {len(florence_detections)} detections') # Format dict answer as readable string (string answers pass through unchanged) if isinstance(answer, dict): answer = vqa_detection.format_florence_response(answer) @@ -1519,7 +1566,7 @@ class VQA: answer = self._fastvlm(question, image, vqa_model, model_name) elif 'gemini' in vqa_model.lower(): handler = 'gemini' - gen_kwargs = get_kwargs() + gen_kwargs = get_kwargs(self.model) from modules.caption import gemini answer = gemini.predict(question, image, vqa_model, system_prompt, model_name, prefill, thinking_mode, gen_kwargs) else: @@ -1542,13 +1589,13 @@ class VQA: points = self.last_detection_data.get('points', None) if detections or points: self.last_annotated_image = vqa_detection.draw_bounding_boxes(image, detections or [], points) - debug(f'VQA caption: handler={handler} created annotated image detections={len(detections) if detections else 0} points={len(points) if points else 0}') + debug(f'LLM: handler={handler} created annotated image detections={len(detections) if detections else 0} points={len(points) if points else 0}') - debug(f'VQA caption: handler={handler} response="{answer}" annotation={self.last_annotated_image is not None}') + debug(f'LLM: handler={handler} response="{answer}" annotation={self.last_annotated_image is not None}') t1 = time.time() if not quiet: model_name = model_name.split(' ')[0] if model_name else 'None' - log.debug(f'Caption: type=vlm model="{model_name}" repo="{vqa_model}" args={get_kwargs()} time={t1-t0:.2f}') + log.debug(f'LLM: type=vlm model="{model_name}" repo="{vqa_model}" args={get_kwargs(self.model)} time={t1-t0:.2f}') self._generation_overrides = None # Clear per-request overrides shared.state.end(jobid) return answer @@ -1563,6 +1610,7 @@ class VQA: prefill: str | None = None, # pylint: disable=unused-argument thinking_mode: bool | None = None, quiet: bool = False, + custom_args: str | None = None, generation_kwargs: dict | None = None, ) -> str: if question is None or len(question.strip()) < 2: @@ -1599,9 +1647,9 @@ class VQA: from modules.files_cache import list_files files += list(list_files(batch_str, ext_filter=['.png', '.jpg', '.jpeg', '.webp', '.jxl'], recursive=recursive)) if len(files) == 0: - log.warning('Caption batch: type=vlm no images') + log.warning('LLM batch: type=vlm no images') return '' - jobid = shared.state.begin('Caption batch') + jobid = shared.state.begin('LLM batch') prompts = [] if save_txt: mode = 'w' if not append_txt else 'a' @@ -1639,7 +1687,7 @@ class VQA: if save_json: writer_json.add(file, result) except Exception as e: - log.error(f'Caption batch: {e}') + log.error(f'LLM batch: {e}') if save_txt: writer_txt.close() if save_json: diff --git a/modules/devices.py b/modules/devices.py index eb4bbb969..543b64065 100644 --- a/modules/devices.py +++ b/modules/devices.py @@ -536,7 +536,7 @@ def set_sdpa_params(): log.debug(f'Torch attention installed: flashattn={flash} sageattention={sage}') from diffusers.models import attention_dispatch as a - log.debug(f'Torch attention status: flash={a._CAN_USE_FLASH_ATTN} flash3={a._CAN_USE_FLASH_ATTN_3} aiter={a._CAN_USE_AITER_ATTN} sage={a._CAN_USE_SAGE_ATTN} flex={a._CAN_USE_FLEX_ATTN} npu={a._CAN_USE_NPU_ATTN} xla={a._CAN_USE_XLA_ATTN} xformers={a._CAN_USE_XFORMERS_ATTN}') # pylint: disable=protected-access + log.debug(f'Torch attention status: flash={a._CAN_USE_FLASH_ATTN} flash3={a._CAN_USE_FLASH_ATTN_3} aiter={a._CAN_USE_AITER_ATTN} sage={a._CAN_USE_SAGE_ATTN} flex={a._CAN_USE_FLEX_ATTN} npu={a._CAN_USE_NPU_ATTN} xla={a._CAN_USE_XLA_ATTN} xformers={a._CAN_USE_XFORMERS_ATTN} kernels={a.is_kernels_available()}') # pylint: disable=protected-access except Exception as e: log.warning(f'Torch SDPA: {e}') diff --git a/modules/ui_caption.py b/modules/ui_caption.py index 3501496eb..cfbf39d9c 100644 --- a/modules/ui_caption.py +++ b/modules/ui_caption.py @@ -8,7 +8,7 @@ default_task = "Normal Caption" def caption_wrapper(tab, image, - vlm_question, vlm_system, vlm_prompt, vlm_model, vlm_prefill, vlm_thinking_mode, + 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, @@ -17,7 +17,7 @@ def caption_wrapper(tab, image, 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) + 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) @@ -25,7 +25,7 @@ def caption_wrapper(tab, image, 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) + 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"') @@ -56,7 +56,7 @@ def update_vlm_prompt_placeholder(question): 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 = 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) @@ -66,6 +66,7 @@ def update_vlm_params(*args): 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() @@ -204,20 +205,23 @@ def create_ui(): 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], 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], 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], 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], 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], 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], 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], 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], 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], outputs=[]) + 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') @@ -375,7 +379,7 @@ def create_ui(): _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_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 @@ -393,7 +397,7 @@ def create_ui(): _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_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 diff --git a/modules/ui_definitions.py b/modules/ui_definitions.py index eb6cb14a7..fd37a6abd 100644 --- a/modules/ui_definitions.py +++ b/modules/ui_definitions.py @@ -711,6 +711,7 @@ def create_settings(cmd_opts): "caption_vlm_keep_prefill": OptionInfo(False, "VLM: keep prefill text in output", gr.Checkbox, {"visible": False}), "caption_vlm_keep_thinking": OptionInfo(False, "VLM: keep reasoning trace in output", gr.Checkbox, {"visible": False}), "caption_vlm_thinking_mode": OptionInfo(False, "VLM: enable thinking/reasoning mode", gr.Checkbox, {"visible": False}), + "caption_vlm_custom_args": OptionInfo("", "VLM: custom arguments", gr.Textbox, {"visible": False}), "tagger_threshold": OptionInfo(0.50, "Tagger: general tag threshold", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01, "visible": False}), "tagger_include_rating": OptionInfo(False, "Tagger: include rating tags", gr.Checkbox, {"visible": False}), "tagger_max_tags": OptionInfo(74, "Tagger: max tags", gr.Slider, {"minimum": 1, "maximum": 512, "step": 1, "visible": False}), diff --git a/scripts/prompt_enhance.py b/scripts/prompt_enhance.py index befcabc94..0390937a5 100644 --- a/scripts/prompt_enhance.py +++ b/scripts/prompt_enhance.py @@ -15,6 +15,7 @@ from modules import ui_control_helpers from modules.sd_offload_aux import register_aux, deregister_aux, move_aux_to_gpu, offload_aux from modules.logger import log from modules.caption.logits import LogitsParser +from modules.caption import helpers debug_enabled = os.environ.get('SD_LLM_DEBUG', None) is not None @@ -132,7 +133,6 @@ class Options: 'mistralai/Ministral-3-3B-Reasoning-2512', 'mistralai/Ministral-3-8B-Reasoning-2512', # Finetunes - 'p-e-w/gemma-4-E2B-it-heretic-ara', 'trohrbaugh/gemma-4-E4B-it-heretic-ara', 'trohrbaugh/Qwen3.5-9B-heretic-v2', ] @@ -538,52 +538,6 @@ class PromptEnhanceScript(scripts_manager.Script): current_image = None return current_image - def get_default_args(self): - to_remove = ['_from_model_config', 'transformers_version'] - config = {} - for k, v in transformers.GenerationConfig._get_default_generation_params().items(): # pylint: disable=protected-access - if v is not None: - config[k] = v - for k, v in transformers.GenerationConfig.from_model_config(self.llm.config).to_dict().items(): - if v is not None: - config[k] = v - for k, v in self.llm.generation_config.to_dict().items(): - if v is not None: - config[k] = v - config = {k: v for k, v in config.items() if k not in to_remove} - return config - - def get_custom_args(self, args_str): - args = {} - if args_str is not None and len(args_str) > 0: - default_args = self.get_default_args() - pairs = re.split(r'[;\n]+', args_str) - for pair in pairs: - if '=' in pair: - key, value = pair.split('=', maxsplit=1) - key = key.strip() - value = value.strip() - if key not in default_args: - log.warning(f'Prompt enhance: key="{key}" invalid') - continue - default_value = default_args[key] - try: - if isinstance(default_value, bool): - value = value.lower() in ['true', '1', 'yes'] - elif isinstance(default_value, int): - value = int(value) - elif isinstance(default_value, float): - value = float(value) - elif isinstance(default_value, list): - value = [v.strip() for v in value.split(',')] - elif isinstance(default_value, str): - pass - except ValueError: - log.warning(f'Prompt enhance: key="{key}" value="{value}" typecast failed') - if key and value: - args[key] = value - return args - def enhance(self, model: str | None=None, prompt:str | None=None, @@ -845,12 +799,12 @@ class PromptEnhanceScript(scripts_manager.Script): logits_processor = LogitsParser(self.tokenizer, process_words, semantic_threshold=semantic_threshold, embedding_similarity=embedding_similarity) gen_kwargs['logits_processor'] = [logits_processor] - custom = self.get_custom_args(custom_args) + custom = helpers.get_custom_args(self.llm, custom_args) for k, v in custom.items(): gen_kwargs[k] = v log.debug(f'Prompt enhance: cls={self.llm.__class__.__name__} model="{model}" tokens={input_len} args={gen_kwargs} custom={custom}') - defaults = {k: v for k, v in self.get_default_args().items() if k not in gen_kwargs} + defaults = {k: v for k, v in helpers.get_default_args(self.llm).items() if k not in gen_kwargs} log.debug(f'Prompt enhance: defaults={defaults}') outputs = self.llm.generate(**inputs, **gen_kwargs)