Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 17:24:32 +02:00
Code Issues Packages Projects Releases Wiki Activity

refactor(caption): code review fixes for offload, inference, and maintainability

Browse Source 
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)
This commit is contained in:
CalamitousFelicitousness
2026-01-29 01:28:38 +00:00
parent bf7a72f12e
commit 443a73b740
10 changed files with 205 additions and 178 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/api/endpoints.py
+4
View File
@@ -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()
modules/caption/deepbooru.py
+4 -25
View File
@@ -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}')
modules/caption/deepseek.py
+7 -6
View File
@@ -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
modules/caption/joycaption.py
+22 -23
View File
@@ -1,7 +1,6 @@
# based on <https://huggingface.co/fancyfeast/llama-joycaption-alpha-two-hf-llava>
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)
modules/caption/joytag.py
+6 -2
View File
@@ -1,4 +1,6 @@
# based on <https://huggingface.co/spaces/fancyfeast/joytag>
# 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]
modules/caption/moondream3.py
+21 -3
View File
@@ -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():
modules/caption/openclip.py
+6 -3
View File
@@ -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()
modules/caption/tagger.py
+25
View File
@@ -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
modules/caption/vqa.py
+108 -92
View File
@@ -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 <think>...</think> blocks from model output.
Applies to models that use HuggingFace chat templates with <think>/<\/think>
tokens (Qwen, Gemma, SmolVLM). Models with structured reasoning APIs
(e.g. Moondream) handle their reasoning output separately.
The opening <think> tag is often in the prompt (not the response), so the
response may only contain </think> without a matching <think>.
Args:
text: Model output text potentially containing <think>/<\/think> tags.
keep: If True, reformat tags as human-readable Reasoning/Answer sections.
If False, strip thinking blocks entirely.
"""
if keep:
if '</think>' in text and '<think>' not in text:
text = 'Reasoning:\n' + text.replace('</think>', '\n\nAnswer:')
else:
text = text.replace('<think>', 'Reasoning:\n').replace('</think>', '\n\nAnswer:')
else:
while '</think>' in text:
start = text.find('<think>')
end = text.find('</think>')
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: <think> is in the prompt, not the response - only </think> appears in generated output
if len(response) > 0:
text = response[0]
if get_keep_thinking():
# Handle case where <think> is in prompt (not response) but </think> is in response
if '</think>' in text and '<think>' not in text:
text = 'Reasoning:\n' + text.replace('</think>', '\n\nAnswer:')
else:
text = text.replace('<think>', 'Reasoning:\n').replace('</think>', '\n\nAnswer:')
else:
while '</think>' in text:
start = text.find('<think>')
end = text.find('</think>')
if start != -1 and start < end:
# Standard <think>...content...</think> block
text = text[:start] + text[end+8:]
else:
# Missing <think> (implied at start) or malformed
# Remove from start up to </think>
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('<think>', 'Reasoning:\n').replace('</think>', '\n\nAnswer:')
else:
text = response
while '</think>' in text:
start = text.find('<think>')
end = text.find('</think>')
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('<think>', 'Reasoning:\n').replace('</think>', '\n\nAnswer:')
else:
while '</think>' in text:
start = text.find('<think>')
end = text.find('</think>')
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)
modules/caption/waifudiffusion.py
+2 -24
View File
@@ -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}')