From 443a73b740fd591db516f76d16735802cb110075 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Thu, 29 Jan 2026 01:28:38 +0000 Subject: [PATCH] refactor(caption): code review fixes for offload, inference, and maintainability Comprehensive review of modules/caption/ addressing memory management, consistency, and code quality: Inference correctness: - Add devices.inference_context() to _qwen(), _smol(), _sa2() handlers - Remove redundant @torch.no_grad() decorator from joycaption predict() - Remove dead dtype=torch.bfloat16 kwarg from Florence loader Memory management: - Bound moondream3 image cache with LRU eviction (max 8 entries) - Replace fragile id(image) cache keys with content-based md5 hash - Add devices.torch_gc() after model loading in deepseek - Move deepbooru model to CPU before dropping reference on unload - Add external handler delegation to VQA.unload() (moondream3, joycaption, joytag, deepseek) - Protect batch offload mutation with try/finally Code deduplication: - Extract strip_think_xml_tags() shared helper for Qwen/Gemma/SmolVLM - Extract save_tags_to_file() into tagger.py from deepbooru and waifudiffusion Documentation and clarity: - Document deepseek global monkey-patches (LlamaFlashAttention2, attrdict) - Document Florence task="task" as intentional design choice - Add vendored-code comment to joytag.py - Document openclip direct .to() usage vs sd_models.move_model - Comment model.eval() calls that are required (trust_remote_code, custom loaders) vs removed where redundant (standard from_pretrained) API robustness: - Add HTTP 422 error response for VQA caption error strings in API endpoints (post_vqa, _dispatch_vlm) --- modules/api/endpoints.py | 4 + modules/caption/deepbooru.py | 29 +---- modules/caption/deepseek.py | 13 +- modules/caption/joycaption.py | 45 ++++--- modules/caption/joytag.py | 8 +- modules/caption/moondream3.py | 24 +++- modules/caption/openclip.py | 9 +- modules/caption/tagger.py | 25 ++++ modules/caption/vqa.py | 200 ++++++++++++++++-------------- modules/caption/waifudiffusion.py | 26 +--- 10 files changed, 205 insertions(+), 178 deletions(-) diff --git a/modules/api/endpoints.py b/modules/api/endpoints.py index b37422479..f8be5f285 100644 --- a/modules/api/endpoints.py +++ b/modules/api/endpoints.py @@ -252,6 +252,8 @@ def post_vqa(req: models.ReqVQA): thinking_mode=req.thinking_mode, generation_kwargs=generation_kwargs if generation_kwargs else None ) + if isinstance(answer, str) and answer.startswith('Error:'): + raise HTTPException(status_code=422, detail=answer) # Return annotated image if requested and available annotated_b64 = None if req.include_annotated: @@ -445,6 +447,8 @@ def _dispatch_vlm(req: models.ReqCaptionVLM) -> models.ResCaptionDispatch: thinking_mode=req.thinking_mode, generation_kwargs=generation_kwargs if generation_kwargs else None ) + if isinstance(answer, str) and answer.startswith('Error:'): + raise HTTPException(status_code=422, detail=answer) annotated_b64 = None if req.include_annotated: annotated_img = vqa.get_last_annotated_image() diff --git a/modules/caption/deepbooru.py b/modules/caption/deepbooru.py index c43d683f1..eb97de4cd 100644 --- a/modules/caption/deepbooru.py +++ b/modules/caption/deepbooru.py @@ -30,7 +30,7 @@ class DeepDanbooru: from modules.caption.deepbooru_model import DeepDanbooruModel self.model = DeepDanbooruModel() self.model.load_state_dict(torch.load(files[0], map_location="cpu")) - self.model.eval() + self.model.eval() # required: loaded via torch.load + load_state_dict self.model.to(devices.cpu, devices.dtype) def start(self): @@ -126,29 +126,6 @@ class DeepDanbooru: model = DeepDanbooru() -def _save_tags_to_file(img_path, tags_str: str, save_append: bool) -> bool: - """Save tags to a text file with error handling. - - Args: - img_path: Path to the image file - tags_str: Tags string to save - save_append: If True, append to existing file; otherwise overwrite - - Returns: - True if save succeeded, False otherwise - """ - try: - txt_path = img_path.with_suffix('.txt') - if save_append and txt_path.exists(): - with open(txt_path, 'a', encoding='utf-8') as f: - f.write(f', {tags_str}') - else: - with open(txt_path, 'w', encoding='utf-8') as f: - f.write(tags_str) - return True - except Exception as e: - shared.log.error(f'DeepBooru batch: failed to save file="{img_path}" error={e}') - return False def get_models() -> list: @@ -170,6 +147,7 @@ def unload_model(): """Unload the DeepBooru model and free memory.""" if model.model is not None: shared.log.debug('DeepBooru unload') + model.model.to(devices.cpu) model.model = None devices.torch_gc(force=True) @@ -312,7 +290,8 @@ def batch( tags_str = model.tag_multi(image, **kwargs) if save_output: - _save_tags_to_file(img_path, tags_str, save_append) + from modules.caption import tagger + tagger.save_tags_to_file(img_path, tags_str, save_append) results.append(f'{img_path.name}: {tags_str[:100]}...' if len(tags_str) > 100 else f'{img_path.name}: {tags_str}') diff --git a/modules/caption/deepseek.py b/modules/caption/deepseek.py index 05f7a7a64..680f2e497 100644 --- a/modules/caption/deepseek.py +++ b/modules/caption/deepseek.py @@ -39,6 +39,9 @@ def load(repo: str): shared.log.error(f'Caption: 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 + # which may not be available. These patches persist for the lifetime of the process and may + # affect other Llama model loads (forcing standard attention instead of flash attention). sys.modules['attrdict'] = fake_attrdict from transformers.models.llama import modeling_llama modeling_llama.LlamaFlashAttention2 = modeling_llama.LlamaAttention @@ -51,8 +54,9 @@ def load(repo: str): cache_dir=shared.opts.hfcache_dir, ) vl_gpt.to(dtype=devices.dtype) - vl_gpt.eval() + vl_gpt.eval() # required: trust_remote_code model loaded_repo = repo + devices.torch_gc() shared.log.info(f'Caption: type=vlm model="DeepSeek VL2" repo="{repo}"') sd_models.move_model(vl_gpt, devices.device) return True @@ -73,7 +77,6 @@ def unload(): def predict(question, image, repo): - global vl_gpt # pylint: disable=global-statement if not load(repo): return '' @@ -97,7 +100,6 @@ def predict(question, image, repo): ).to(device=devices.device, dtype=devices.dtype) inputs_embeds = vl_gpt.prepare_inputs_embeds(**prepare_inputs) inputs_embeds = inputs_embeds.to(device=devices.device, dtype=devices.dtype) - sd_models.move_model(vl_gpt, devices.device) with devices.inference_context(): outputs = vl_gpt.language.generate( inputs_embeds=inputs_embeds, @@ -109,8 +111,7 @@ def predict(question, image, repo): do_sample=False, use_cache=True ) - vl_gpt = vl_gpt.to(devices.cpu) + if shared.opts.caption_offload: + sd_models.move_model(vl_gpt, devices.cpu, force=True) answer = vl_chat_processor.tokenizer.decode(outputs[0].cpu().tolist(), skip_special_tokens=True) - print('inputs', prepare_inputs['sft_format'][0]) - print('answer', answer) return answer diff --git a/modules/caption/joycaption.py b/modules/caption/joycaption.py index 9f3827d47..778d72cf2 100644 --- a/modules/caption/joycaption.py +++ b/modules/caption/joycaption.py @@ -1,7 +1,6 @@ # based on from dataclasses import dataclass -import torch from transformers import AutoProcessor, LlavaForConditionalGeneration from modules import shared, devices, sd_models, model_quant @@ -73,7 +72,6 @@ def load(repo: str = None): cache_dir=shared.opts.hfcache_dir, **quant_args, ) - llava_model.eval() sd_models.move_model(llava_model, devices.device) @@ -90,7 +88,6 @@ def unload(): shared.log.debug('JoyCaption unload: no model loaded') -@torch.no_grad() def predict(question: str, image, vqa_model: str = None) -> str: opts.max_new_tokens = shared.opts.caption_vlm_max_length load(vqa_model) @@ -105,23 +102,25 @@ def predict(question: str, image, vqa_model: str = None) -> str: convo_string = processor.apply_chat_template(convo, tokenize=False, add_generation_prompt=True) inputs = processor(text=[convo_string], images=[image], return_tensors="pt").to(devices.device) inputs['pixel_values'] = inputs['pixel_values'].to(devices.dtype) - with devices.inference_context(): - generate_ids = llava_model.generate( # Generate the captions - **inputs, - # input_ids=inputs['input_ids'], - # pixel_values=inputs['pixel_values'], - # attention_mask=inputs['attention_mask'], - max_new_tokens=opts.max_new_tokens, - suppress_tokens=None, - use_cache=True, - do_sample=opts.sample, - temperature=opts.temp, - top_k=opts.top_k, - top_p=opts.top_p, - )[0] - generate_ids = generate_ids[inputs['input_ids'].shape[1]:] # Trim off the prompt - caption = processor.tokenizer.decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False) # Decode the caption - if shared.opts.caption_offload: - sd_models.move_model(llava_model, devices.cpu, force=True) - caption = caption.replace('\n\n', '\n').strip() - return caption + try: + with devices.inference_context(): + generate_ids = llava_model.generate( # Generate the captions + **inputs, + # input_ids=inputs['input_ids'], + # pixel_values=inputs['pixel_values'], + # attention_mask=inputs['attention_mask'], + max_new_tokens=opts.max_new_tokens, + suppress_tokens=None, + use_cache=True, + do_sample=opts.sample, + temperature=opts.temp, + top_k=opts.top_k, + top_p=opts.top_p, + )[0] + generate_ids = generate_ids[inputs['input_ids'].shape[1]:] # Trim off the prompt + caption = processor.tokenizer.decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False) # Decode the caption + caption = caption.replace('\n\n', '\n').strip() + return caption + finally: + if shared.opts.caption_offload: + sd_models.move_model(llava_model, devices.cpu, force=True) diff --git a/modules/caption/joytag.py b/modules/caption/joytag.py index a47ec2041..960e79ec6 100644 --- a/modules/caption/joytag.py +++ b/modules/caption/joytag.py @@ -1,4 +1,6 @@ -# based on +# Vendored from JoyTag: https://huggingface.co/spaces/fancyfeast/joytag +# Contains full model architecture (ViT, CNN stems, MAE) including training-only code +# retained for update compatibility. Do not modify directly — sync from upstream. import os import math @@ -1041,7 +1043,7 @@ def load(): folder = huggingface_hub.snapshot_download(MODEL_REPO, cache_dir=shared.opts.hfcache_dir) model = VisionModel.load_model(folder) model.to(dtype=devices.dtype) - model.eval() + model.eval() # required: custom loader, not from_pretrained with open(os.path.join(folder, 'top_tags.txt'), 'r', encoding='utf8') as f: tags = [line.strip() for line in f.readlines() if line.strip()] shared.log.info(f'Caption: type=vlm model="JoyTag" repo="{MODEL_REPO}" tags={len(tags)}') @@ -1067,6 +1069,8 @@ def predict(image: Image.Image): with devices.inference_context(): preds = model({'image': image_tensor}) tag_preds = preds['tags'].sigmoid().cpu() + if shared.opts.caption_offload: + sd_models.move_model(model, devices.cpu, force=True) scores = {tags[i]: tag_preds[0][i] for i in range(len(tags))} if shared.opts.tagger_show_scores: predicted_tags = [f'{tag}:{score:.2f}' for tag, score in scores.items() if score > THRESHOLD] diff --git a/modules/caption/moondream3.py b/modules/caption/moondream3.py index 1a86cdf85..e6e6b9a9b 100644 --- a/modules/caption/moondream3.py +++ b/modules/caption/moondream3.py @@ -4,6 +4,8 @@ # Architecture: Mixture-of-Experts (9B total params, 2B active) import os import re +import hashlib +import collections import transformers from PIL import Image from modules import shared, devices, sd_models @@ -21,7 +23,8 @@ def debug(*args, **kwargs): # Global state moondream3_model = None loaded = None -image_cache = {} # Cache encoded images for reuse +image_cache: collections.OrderedDict = collections.OrderedDict() # Bounded LRU cache for encoded image tensors +IMAGE_CACHE_MAX = 8 def get_settings(): @@ -54,7 +57,7 @@ def load_model(repo: str): cache_dir=shared.opts.hfcache_dir, ) - moondream3_model.eval() + moondream3_model.eval() # required: trust_remote_code model # Initialize KV caches before moving to device (they're lazy by default) if hasattr(moondream3_model, '_setup_caches'): @@ -72,6 +75,14 @@ def load_model(repo: str): return moondream3_model +def _image_hash(image: Image.Image) -> str: + """Content-based hash for cache keys using image size and pixel sample.""" + h = hashlib.md5(usedforsecurity=False) + h.update(f"{image.size}{image.mode}".encode()) + h.update(image.tobytes()[:4096]) + return h.hexdigest() + + def encode_image(image: Image.Image, cache_key: str = None): """ Encode image for reuse across multiple queries. @@ -84,6 +95,7 @@ def encode_image(image: Image.Image, cache_key: str = None): Encoded image tensor """ 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}"') return image_cache[cache_key] @@ -94,6 +106,9 @@ def encode_image(image: Image.Image, cache_key: str = None): if cache_key: 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)}') return encoded @@ -133,7 +148,7 @@ def query(image: Image.Image, question: str, repo: str, stream: bool = False, # Use cached encoding if requested if use_cache: - cache_key = f"{id(image)}_{question}" + cache_key = f"{_image_hash(image)}_{question}" image_input = encode_image(image, cache_key) else: image_input = image @@ -384,6 +399,9 @@ def predict(question: str, image: Image.Image, repo: str, model_name: str = None from modules import errors errors.display(e, 'Moondream3') return f"Error: {str(e)}" + finally: + if shared.opts.caption_offload and moondream3_model is not None: + sd_models.move_model(moondream3_model, devices.cpu, force=True) def clear_cache(): diff --git a/modules/caption/openclip.py b/modules/caption/openclip.py index 660bf7d5d..22e6c963f 100644 --- a/modules/caption/openclip.py +++ b/modules/caption/openclip.py @@ -5,7 +5,7 @@ import threading import re import gradio as gr from PIL import Image -from modules import devices, shared, errors, sd_models +from modules import devices, shared, errors debug_enabled = os.environ.get('SD_CAPTION_DEBUG', None) is not None @@ -165,8 +165,11 @@ def load_captioner(clip_model, blip_model): def unload_clip_model(): if ci is not None and shared.opts.caption_offload: shared.log.debug('CLIP unload: offloading models to CPU') - sd_models.move_model(ci.caption_model, devices.cpu) - sd_models.move_model(ci.clip_model, devices.cpu) + # Direct .to() instead of sd_models.move_model — models are from clip_interrogator, not transformers + if ci.caption_model is not None and hasattr(ci.caption_model, 'to'): + ci.caption_model.to(devices.cpu) + if ci.clip_model is not None and hasattr(ci.clip_model, 'to'): + ci.clip_model.to(devices.cpu) ci.caption_offloaded = True ci.clip_offloaded = True devices.torch_gc() diff --git a/modules/caption/tagger.py b/modules/caption/tagger.py index 2b18b3937..e0d3b37d8 100644 --- a/modules/caption/tagger.py +++ b/modules/caption/tagger.py @@ -6,6 +6,31 @@ from modules import shared DEEPBOORU_MODEL = "DeepBooru" +def save_tags_to_file(img_path, tags_str: str, save_append: bool) -> bool: + """Save tags to a text file alongside the image. + + Args: + img_path: Path to the image file (pathlib.Path) + tags_str: Tags string to save + save_append: If True, append to existing file; otherwise overwrite + + Returns: + True if save succeeded, False otherwise + """ + try: + txt_path = img_path.with_suffix('.txt') + if save_append and txt_path.exists(): + with open(txt_path, 'a', encoding='utf-8') as f: + f.write(f', {tags_str}') + else: + with open(txt_path, 'w', encoding='utf-8') as f: + f.write(tags_str) + return True + except Exception as e: + shared.log.error(f'Tagger batch: failed to save file="{img_path}" error={e}') + return False + + def get_models() -> list: """Return combined list: DeepBooru + WaifuDiffusion models.""" from modules.caption import waifudiffusion diff --git a/modules/caption/vqa.py b/modules/caption/vqa.py index 80eb53757..5d6ef6c14 100644 --- a/modules/caption/vqa.py +++ b/modules/caption/vqa.py @@ -350,8 +350,8 @@ def clean(response, question, prefill=None): def _get_overrides(): """Get generation overrides from VQA singleton if available.""" - if _instance is not None and _instance._generation_overrides is not None: - return _instance._generation_overrides + if _instance is not None and _instance.generation_overrides is not None: + return _instance.generation_overrides return {} @@ -363,6 +363,37 @@ def get_keep_thinking(): return shared.opts.caption_vlm_keep_thinking +def strip_think_xml_tags(text: str, keep: bool = False) -> str: + """Strip or reformat XML-style ... blocks from model output. + + Applies to models that use HuggingFace chat templates with /<\/think> + tokens (Qwen, Gemma, SmolVLM). Models with structured reasoning APIs + (e.g. Moondream) handle their reasoning output separately. + + The opening tag is often in the prompt (not the response), so the + response may only contain without a matching . + + Args: + text: Model output text potentially containing /<\/think> tags. + keep: If True, reformat tags as human-readable Reasoning/Answer sections. + If False, strip thinking blocks entirely. + """ + if keep: + if '' in text and '' not in text: + text = 'Reasoning:\n' + text.replace('', '\n\nAnswer:') + else: + text = text.replace('', 'Reasoning:\n').replace('', '\n\nAnswer:') + else: + while '' in text: + start = text.find('') + end = text.find('') + if start != -1 and start < end: + text = text[:start] + text[end + 8:] + else: + text = text[end + 8:] + return text + + def get_keep_prefill(): """Check if prefill should be kept in output, with per-request override support.""" overrides = _get_overrides() @@ -374,7 +405,7 @@ def get_keep_prefill(): def get_kwargs(): """Build generation kwargs from settings with per-request overrides from VQA instance. - Checks the singleton VQA instance's _generation_overrides for per-request overrides. + Checks the singleton VQA instance's generation_overrides for per-request overrides. Override keys: max_tokens, temperature, top_k, top_p, num_beams, do_sample None values are ignored, allowing selective override. """ @@ -415,8 +446,13 @@ class VQA: self.last_detection_data = None self._generation_overrides = None # Per-request generation parameter overrides + @property + def generation_overrides(self): + """Get current per-request generation parameter overrides.""" + return self._generation_overrides + def unload(self): - """Release VLM model from GPU/memory.""" + """Release VLM model from GPU/memory, including external handlers.""" if self.model is not None: model_name = self.loaded shared.log.debug(f'VQA unload: unloading model="{model_name}"') @@ -427,7 +463,13 @@ class VQA: devices.torch_gc(force=True, reason='vqa unload') shared.log.debug(f'VQA unload: model="{model_name}" unloaded') else: - shared.log.debug('VQA unload: no model loaded') + shared.log.debug('VQA unload: no internal model loaded') + # External handlers manage their own module-level globals and are not covered by self.model + from modules.caption import moondream3, joycaption, joytag, deepseek + moondream3.unload() + joycaption.unload() + joytag.unload() + deepseek.unload() def load(self, model_name: str = None): """Load VLM model into memory for the specified model name.""" @@ -644,10 +686,11 @@ class VQA: 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}') - output_ids = self.model.generate( - **inputs, - **gen_kwargs, - ) + with devices.inference_context(): + output_ids = self.model.generate( + **inputs, + **gen_kwargs, + ) debug(f'VQA caption: handler=qwen output_ids_shape={output_ids.shape}') generated_ids = [ output_ids[len(input_ids):] @@ -656,29 +699,8 @@ class VQA: 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}"') - # Clean up thinking tags - # Note: is in the prompt, not the response - only appears in generated output if len(response) > 0: - text = response[0] - if get_keep_thinking(): - # Handle case where is in prompt (not response) but is in response - if '' in text and '' not in text: - text = 'Reasoning:\n' + text.replace('', '\n\nAnswer:') - else: - text = text.replace('', 'Reasoning:\n').replace('', '\n\nAnswer:') - else: - while '' in text: - start = text.find('') - end = text.find('') - - if start != -1 and start < end: - # Standard ...content... block - text = text[:start] + text[end+8:] - else: - # Missing (implied at start) or malformed - # Remove from start up to - text = text[end+8:] - response[0] = text + response[0] = strip_think_xml_tags(response[0], keep=get_keep_thinking()) return response def _load_gemma(self, repo: str): @@ -791,20 +813,7 @@ class VQA: if debug_enabled: debug(f'VQA caption: handler=gemma response_before_clean="{response}"') - # Clean up thinking tags (if any remain) - if get_keep_thinking(): - response = response.replace('', 'Reasoning:\n').replace('', '\n\nAnswer:') - else: - text = response - while '' in text: - start = text.find('') - end = text.find('') - if start != -1 and start < end: - text = text[:start] + text[end+8:] - else: - text = text[end+8:] - response = text - + response = strip_think_xml_tags(response, keep=get_keep_thinking()) return response def _load_paligemma(self, repo: str): @@ -964,30 +973,18 @@ class VQA: inputs = inputs.to(devices.device, devices.dtype) gen_kwargs = get_kwargs() debug(f'VQA caption: handler=smol generation_kwargs={gen_kwargs}') - output_ids = self.model.generate( - **inputs, - **gen_kwargs, - ) + with devices.inference_context(): + output_ids = self.model.generate( + **inputs, + **gen_kwargs, + ) debug(f'VQA caption: 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}"') - # Clean up thinking tags if len(response) > 0: - text = response[0] - if get_keep_thinking(): - text = text.replace('', 'Reasoning:\n').replace('', '\n\nAnswer:') - else: - while '' in text: - start = text.find('') - end = text.find('') - if start != -1 and start < end: - text = text[:start] + text[end+8:] - else: - text = text[end+8:] - response[0] = text - + response[0] = strip_think_xml_tags(response[0], keep=get_keep_thinking()) return response def _load_git(self, repo: str): @@ -1110,7 +1107,7 @@ class VQA: ) self.processor = transformers.AutoTokenizer.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir) self.loaded = repo - self.model.eval() + self.model.eval() # required: trust_remote_code model devices.torch_gc() def _moondream(self, question: str, image: Image.Image, repo: str, model_name: str = None, thinking_mode: bool = False): @@ -1202,7 +1199,6 @@ class VQA: quant_args = model_quant.create_config(module='LLM') self.model = transformers.Florence2ForConditionalGeneration.from_pretrained( repo_name, - dtype=torch.bfloat16, revision=effective_revision, torch_dtype=devices.dtype, cache_dir=shared.opts.hfcache_dir, @@ -1211,7 +1207,6 @@ 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 self.loaded = cache_key - self.model.eval() devices.torch_gc() def _florence(self, question: str, image: Image.Image, repo: str, revision: str = None, model_name: str = None): # pylint: disable=unused-argument @@ -1230,8 +1225,13 @@ class VQA: overrides = _get_overrides() max_tokens = overrides.get('max_tokens') if overrides.get('max_tokens') is not None else shared.opts.caption_vlm_max_length gen_kwargs = {'max_new_tokens': max_tokens, 'num_beams': 3, 'do_sample': False} + # Some Florence fine-tunes (e.g., CogFlorence) don't have decoder_start_token_id set + 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}') - with devices.inference_context(): + with devices.inference_context(), devices.bypass_sdpa_hijacks(): generated_ids = self.model.generate( input_ids=input_ids, pixel_values=pixel_values, @@ -1239,6 +1239,8 @@ class VQA: ) generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=False)[0] debug(f'VQA caption: 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}') return response @@ -1253,7 +1255,7 @@ class VQA: low_cpu_mem_usage=True, use_flash_attn=False, trust_remote_code=True) - self.model = self.model.eval() + self.model = self.model.eval() # required: trust_remote_code model self.processor = transformers.AutoTokenizer.from_pretrained( repo, trust_remote_code=True, @@ -1276,7 +1278,8 @@ class VQA: 'mask_prompts': None, 'tokenizer': self.processor, } - return_dict = self.model.predict_forward(**input_dict) + with devices.inference_context(): + return_dict = self.model.predict_forward(**input_dict) response = return_dict["prediction"] # the text format answer return response @@ -1316,6 +1319,7 @@ class VQA: image = image.convert('RGB') if image is None: shared.log.error(f'VQA caption: 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.' @@ -1324,6 +1328,7 @@ class VQA: # Use content from Prompt field directly - requires user input if not prompt or len(prompt.strip()) < 2: shared.log.error(f'VQA caption: 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.' question = prompt @@ -1334,6 +1339,7 @@ class VQA: # These modes require user input in the prompt field if not prompt or len(prompt.strip()) < 2: shared.log.error(f'VQA caption: 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").' # Convert friendly name to internal token (handles Point/Detect prefix) @@ -1377,6 +1383,14 @@ class VQA: elif 'florence' in vqa_model.lower(): handler = 'florence' answer = self._florence(question, image, vqa_model, None, model_name) + # Parse Florence detection response for annotated image (handles both dict and string formats) + 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') + # Format dict answer as readable string (string answers pass through unchanged) + if isinstance(answer, dict): + answer = vqa_detection.format_florence_response(answer) elif 'qwen' in vqa_model.lower() or 'torii' in vqa_model.lower() or 'mimo' in vqa_model.lower(): handler = 'qwen' answer = self._qwen(question, image, vqa_model, system_prompt, model_name, prefill, thinking_mode) @@ -1477,29 +1491,31 @@ class VQA: writer = BatchWriter(os.path.dirname(files[0]), mode=mode) orig_offload = shared.opts.caption_offload shared.opts.caption_offload = False - import rich.progress as rp - pbar = rp.Progress(rp.TextColumn('[cyan]Caption:'), rp.BarColumn(), rp.MofNCompleteColumn(), rp.TaskProgressColumn(), rp.TimeRemainingColumn(), rp.TimeElapsedColumn(), rp.TextColumn('[cyan]{task.description}'), console=shared.console) - with pbar: - task = pbar.add_task(total=len(files), description='starting...') - for file in files: - pbar.update(task, advance=1, description=file) - try: - if shared.state.interrupted: - break - img = Image.open(file) - caption = self.caption(question, system_prompt, prompt, img, model_name, prefill, thinking_mode, quiet=True) - # Save annotated image if available - if self.last_annotated_image and write: - annotated_path = os.path.splitext(file)[0] + "_annotated.png" - self.last_annotated_image.save(annotated_path) - prompts.append(caption) - if write: - writer.add(file, caption) - except Exception as e: - shared.log.error(f'Caption batch: {e}') - if write: - writer.close() - shared.opts.caption_offload = orig_offload + try: + import rich.progress as rp + pbar = rp.Progress(rp.TextColumn('[cyan]Caption:'), rp.BarColumn(), rp.MofNCompleteColumn(), rp.TaskProgressColumn(), rp.TimeRemainingColumn(), rp.TimeElapsedColumn(), rp.TextColumn('[cyan]{task.description}'), console=shared.console) + with pbar: + task = pbar.add_task(total=len(files), description='starting...') + for file in files: + pbar.update(task, advance=1, description=file) + try: + if shared.state.interrupted: + break + img = Image.open(file) + result = self.caption(question, system_prompt, prompt, img, model_name, prefill, thinking_mode, quiet=True) + # Save annotated image if available + if self.last_annotated_image and write: + annotated_path = os.path.splitext(file)[0] + "_annotated.png" + self.last_annotated_image.save(annotated_path) + prompts.append(result) + if write: + writer.add(file, result) + except Exception as e: + shared.log.error(f'Caption batch: {e}') + if write: + writer.close() + finally: + shared.opts.caption_offload = orig_offload shared.state.end(jobid) return '\n\n'.join(prompts) diff --git a/modules/caption/waifudiffusion.py b/modules/caption/waifudiffusion.py index 5d252f0cd..e0089f233 100644 --- a/modules/caption/waifudiffusion.py +++ b/modules/caption/waifudiffusion.py @@ -337,29 +337,6 @@ class WaifuDiffusionTagger: tagger = WaifuDiffusionTagger() -def _save_tags_to_file(img_path, tags_str: str, save_append: bool) -> bool: - """Save tags to a text file with error handling. - - Args: - img_path: Path to the image file - tags_str: Tags string to save - save_append: If True, append to existing file; otherwise overwrite - - Returns: - True if save succeeded, False otherwise - """ - try: - txt_path = img_path.with_suffix('.txt') - if save_append and txt_path.exists(): - with open(txt_path, 'a', encoding='utf-8') as f: - f.write(f', {tags_str}') - else: - with open(txt_path, 'w', encoding='utf-8') as f: - f.write(tags_str) - return True - except Exception as e: - shared.log.error(f'WaifuDiffusion batch: failed to save file="{img_path}" error={e}') - return False def get_models() -> list: @@ -529,7 +506,8 @@ def batch( tags_str = tagger.predict(image, **kwargs) if save_output: - _save_tags_to_file(img_path, tags_str, save_append) + from modules.caption import tagger + tagger.save_tags_to_file(img_path, tags_str, save_append) results.append(f'{img_path.name}: {tags_str[:100]}...' if len(tags_str) > 100 else f'{img_path.name}: {tags_str}')