Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-18 16:54:33 +02:00
Code Issues Packages Projects Releases Wiki Activity

fix(caption): safetensors-only downloads, model load fixes, UI default, prefill tests

Browse Source 
- Add use_safetensors=True to all 16 model from_pretrained calls to
  avoid downloading redundant .bin files alongside safetensors
- Add device property to JoyTag VisionModel so move_model can relocate
  it to CUDA (fixes 'ViT object has no attribute device')
- Fix Pix2Struct dtype mismatch by casting float inputs to model dtype
  while preserving integer tensor types
- Patch AutoConfig.register with exist_ok=True during Ovis loading to
  handle duplicate aimv2 registration on model reload
- Detect Qwen VL fine-tune architecture from config model_type instead
  of repo name, fixing ToriiGate and similar third-party fine-tunes
- Change UI default task from Short Caption to Normal Caption, and
  preserve it on model switch instead of resetting to Use Prompt
- Add dual-prefill testing across 5 VQA test methods using a shared
  _check_prefill helper
- Fix pre-existing ruff W605 in strip_think_xml_tags docstring
This commit is contained in:
CalamitousFelicitousness
2026-01-29 04:21:01 +00:00
parent 57659ab642
commit e2cdbe47fa
7 changed files with 88 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cli/test-caption-api.py
+32
View File
@@ -44,6 +44,9 @@ OCR_TEST_IMAGE = 'models/Reference/HiDream-ai--HiDream-I1-Fast.jpg'
# Bracket test image (must produce tags with parentheses, e.g. pokemon_(creature))
BRACKET_TEST_IMAGE = 'models/Reference/SDXL-Flash_Mini.jpg'
# Custom prefill text used for dual-prefill verification across tests
CUSTOM_PREFILL = "Vlado is the best, and I'm looking at his robot which"
class CaptionAPITest:
"""Test harness for Caption API endpoints."""
@@ -257,6 +260,20 @@ class CaptionAPITest:
return False
return True
def _check_prefill(self, base_request: dict, test_label: str):
"""Re-run a VQA request with custom prefill and verify it appears in output."""
req = {**base_request, 'prefill': CUSTOM_PREFILL, 'keep_prefill': True}
data = self.post('/sdapi/v1/vqa', req)
if 'error' in data:
self.log_skip(f"{test_label} prefill: API error")
elif data.get('answer') and not self.is_error_answer(data['answer']):
if data['answer'].startswith(CUSTOM_PREFILL):
self.log_pass(f"{test_label} prefill: output starts with custom prefill")
else:
self.log_fail(f"{test_label} prefill: expected '{CUSTOM_PREFILL[:30]}...' but got '{data['answer'][:30]}...'")
else:
self.log_fail(f"{test_label} prefill: empty/error")
def get_model_family(self, model_name):
"""Determine model family from model name."""
name_lower = model_name.lower()
@@ -1138,6 +1155,9 @@ class CaptionAPITest:
if results['Long Caption'] < results['Normal Caption']:
self.log_info(f"NOTE: Long ({results['Long Caption']}) < Normal ({results['Normal Caption']}); LLM may interpret length prompts differently per run")
# Dual prefill: re-run 'Normal Caption' with custom prefill
self._check_prefill({'image': self.image_b64, 'question': 'Normal Caption'}, "different_prompts")
# =========================================================================
# TEST: POST /sdapi/v1/vqa - Annotated Image
# =========================================================================
@@ -1256,6 +1276,9 @@ class CaptionAPITest:
else:
self.log_fail("Custom system prompt returned empty answer")
# Dual prefill: re-run with custom system prompt + prefill
self._check_prefill({'image': self.image_b64, 'question': 'describe the image', 'system': custom_system}, "system_prompt")
# =========================================================================
# TEST: POST /sdapi/v1/vqa - Invalid Inputs
# =========================================================================
@@ -1328,6 +1351,9 @@ class CaptionAPITest:
else:
self.log_skip("Detection prompt may require specific model")
# Dual prefill: re-run 'Use Prompt' with custom prefill
self._check_prefill({'image': self.image_b64, 'question': 'Use Prompt', 'prompt': custom_prompt}, "prompt_field")
# =========================================================================
# TEST: POST /sdapi/v1/vqa - Generation Parameters
# =========================================================================
@@ -1434,6 +1460,9 @@ class CaptionAPITest:
else:
self.log_fail("top_k/top_p returned empty/error")
# Dual prefill: re-run temp=0 request with custom prefill
self._check_prefill({'image': self.image_b64, 'question': 'describe the image briefly', 'temperature': 0.0}, "generation_params")
# =========================================================================
# TEST: POST /sdapi/v1/vqa - Sampling Controls
# =========================================================================
@@ -1500,6 +1529,9 @@ class CaptionAPITest:
else:
self.log_fail("num_beams=4 returned empty/error")
# Dual prefill: re-run greedy request with custom prefill
self._check_prefill({'image': self.image_b64, 'question': 'describe the image', 'do_sample': False}, "sampling")
# =========================================================================
# TEST: POST /sdapi/v1/vqa - Thinking Mode
# =========================================================================
modules/caption/deepseek.py
+1
View File
@@ -51,6 +51,7 @@ def load(repo: str):
vl_gpt = AutoModelForCausalLM.from_pretrained(
repo,
trust_remote_code=True,
use_safetensors=True,
cache_dir=shared.opts.hfcache_dir,
)
vl_gpt.to(dtype=devices.dtype)
modules/caption/joycaption.py
+1
View File
@@ -69,6 +69,7 @@ def load(repo: str = None):
llava_model = LlavaForConditionalGeneration.from_pretrained(
repo,
torch_dtype=devices.dtype,
use_safetensors=True,
cache_dir=shared.opts.hfcache_dir,
**quant_args,
)
modules/caption/joytag.py
+4
View File
@@ -120,6 +120,10 @@ class VisionModel(nn.Module):
self.image_size = image_size
self.n_tags = n_tags
@property
def device(self):
return next(self.parameters()).device
@staticmethod
def load_model(path: str) -> 'VisionModel':
with open(Path(path) / 'config.json', 'r', encoding='utf8') as f:
modules/caption/moondream3.py
+1
View File
@@ -54,6 +54,7 @@ def load_model(repo: str):
repo,
trust_remote_code=True,
torch_dtype=devices.dtype,
use_safetensors=True,
cache_dir=shared.opts.hfcache_dir,
)
modules/caption/vqa.py
+46 -12
View File
@@ -366,7 +366,7 @@ def get_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>
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.
@@ -374,7 +374,7 @@ def strip_think_xml_tags(text: str, keep: bool = False) -> str:
response may only contain </think> without a matching <think>.
Args:
text: Model output text potentially containing <think>/<\/think> tags.
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.
"""
@@ -550,6 +550,7 @@ class VQA:
repo,
torch_dtype=devices.dtype,
trust_remote_code=True,
use_safetensors=True,
cache_dir=shared.opts.hfcache_dir,
**quant_args,
)
@@ -585,6 +586,13 @@ class VQA:
answer = self.processor.decode(outputs[0], skip_special_tokens=True)
return answer
# Map Qwen VL config model_type strings to their model classes.
_QWEN_VL_MODEL_TYPE_MAP = {
'qwen3_vl': 'Qwen3VLForConditionalGeneration',
'qwen2_5_vl': 'Qwen2_5_VLForConditionalGeneration',
'qwen2_vl': 'Qwen2VLForConditionalGeneration',
}
def _load_qwen(self, repo: str):
"""Load Qwen VL model and processor."""
if self.model is None or self.loaded != repo:
@@ -597,11 +605,17 @@ class VQA:
elif 'Qwen2-VL' in repo or 'Qwen2VL' in repo:
cls_name = transformers.Qwen2VLForConditionalGeneration
else:
cls_name = transformers.AutoModelForCausalLM
# Fine-tunes (e.g. ToriiGate) may not have "Qwen" in the repo name.
# Detect the correct class from the config's model_type.
config = transformers.AutoConfig.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir)
model_type = getattr(config, 'model_type', '')
cls_attr = self._QWEN_VL_MODEL_TYPE_MAP.get(model_type)
cls_name = getattr(transformers, cls_attr) if cls_attr else transformers.AutoModelForCausalLM
quant_args = model_quant.create_config(module='LLM')
self.model = cls_name.from_pretrained(
repo,
torch_dtype=devices.dtype,
use_safetensors=True,
cache_dir=shared.opts.hfcache_dir,
**quant_args,
)
@@ -716,6 +730,7 @@ class VQA:
self.model = cls.from_pretrained(
repo,
torch_dtype=devices.dtype,
use_safetensors=True,
cache_dir=shared.opts.hfcache_dir,
**quant_args,
)
@@ -826,6 +841,7 @@ class VQA:
repo,
cache_dir=shared.opts.hfcache_dir,
torch_dtype=devices.dtype,
use_safetensors=True,
)
self.loaded = repo
devices.torch_gc()
@@ -850,13 +866,22 @@ class VQA:
if self.model is None or self.loaded != repo:
shared.log.debug(f'Caption load: vlm="{repo}"')
self.model = None
self.model = transformers.AutoModelForCausalLM.from_pretrained(
repo,
torch_dtype=devices.dtype,
multimodal_max_length=32768,
trust_remote_code=True,
cache_dir=shared.opts.hfcache_dir,
)
# Ovis remote code calls AutoConfig.register("aimv2", ...) at module scope
# without exist_ok=True, which fails on reload or when the type is already
# registered by a newer transformers version.
_orig = transformers.AutoConfig.register.__func__ if hasattr(transformers.AutoConfig.register, '__func__') else transformers.AutoConfig.register
transformers.AutoConfig.register = staticmethod(lambda model_type, config, exist_ok=False: _orig(model_type, config, exist_ok=True))
try:
self.model = transformers.AutoModelForCausalLM.from_pretrained(
repo,
torch_dtype=devices.dtype,
multimodal_max_length=32768,
trust_remote_code=True,
use_safetensors=True,
cache_dir=shared.opts.hfcache_dir,
)
finally:
transformers.AutoConfig.register = _orig
self.loaded = repo
devices.torch_gc()
@@ -903,6 +928,7 @@ class VQA:
repo,
cache_dir=shared.opts.hfcache_dir,
torch_dtype=devices.dtype,
use_safetensors=True,
**quant_args,
)
self.processor = transformers.AutoProcessor.from_pretrained(repo, max_pixels=1024*1024, cache_dir=shared.opts.hfcache_dir)
@@ -995,6 +1021,7 @@ class VQA:
self.model = transformers.GitForCausalLM.from_pretrained(
repo,
torch_dtype=devices.dtype,
use_safetensors=True,
cache_dir=shared.opts.hfcache_dir,
)
self.processor = transformers.GitProcessor.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir)
@@ -1025,6 +1052,7 @@ class VQA:
self.model = transformers.BlipForQuestionAnswering.from_pretrained(
repo,
torch_dtype=devices.dtype,
use_safetensors=True,
cache_dir=shared.opts.hfcache_dir,
)
self.processor = transformers.BlipProcessor.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir)
@@ -1049,6 +1077,7 @@ class VQA:
self.model = transformers.ViltForQuestionAnswering.from_pretrained(
repo,
torch_dtype=devices.dtype,
use_safetensors=True,
cache_dir=shared.opts.hfcache_dir,
)
self.processor = transformers.ViltProcessor.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir)
@@ -1075,6 +1104,7 @@ class VQA:
self.model = transformers.Pix2StructForConditionalGeneration.from_pretrained(
repo,
torch_dtype=devices.dtype,
use_safetensors=True,
cache_dir=shared.opts.hfcache_dir,
)
self.processor = transformers.Pix2StructProcessor.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir)
@@ -1085,9 +1115,10 @@ class VQA:
self._load_pix(repo)
sd_models.move_model(self.model, devices.device)
if len(question) > 0:
inputs = self.processor(images=image, text=question, return_tensors="pt").to(devices.device)
inputs = self.processor(images=image, text=question, return_tensors="pt")
else:
inputs = self.processor(images=image, return_tensors="pt").to(devices.device)
inputs = self.processor(images=image, return_tensors="pt")
inputs = {k: v.to(devices.device, devices.dtype) if v.is_floating_point() else v.to(devices.device) for k, v in inputs.items()}
with devices.inference_context():
outputs = self.model.generate(**inputs)
response = self.processor.decode(outputs[0], skip_special_tokens=True)
@@ -1103,6 +1134,7 @@ class VQA:
revision="2025-06-21",
trust_remote_code=True,
torch_dtype=devices.dtype,
use_safetensors=True,
cache_dir=shared.opts.hfcache_dir,
)
self.processor = transformers.AutoTokenizer.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir)
@@ -1201,6 +1233,7 @@ class VQA:
repo_name,
revision=effective_revision,
torch_dtype=devices.dtype,
use_safetensors=True,
cache_dir=shared.opts.hfcache_dir,
**quant_args,
)
@@ -1254,6 +1287,7 @@ class VQA:
torch_dtype=devices.dtype,
low_cpu_mem_usage=True,
use_flash_attn=False,
use_safetensors=True,
trust_remote_code=True)
self.model = self.model.eval() # required: trust_remote_code model
self.processor = transformers.AutoTokenizer.from_pretrained(
modules/ui_caption.py
+3 -2
View File
@@ -3,7 +3,7 @@ from modules import shared, ui_common, generation_parameters_copypaste
from modules.caption import openclip
default_task = "Short Caption"
default_task = "Normal Caption"
def vlm_caption_wrapper(question, system_prompt, prompt, image, model_name, prefill, thinking_mode):
"""Wrapper for vqa.caption that handles annotated image display."""
@@ -19,7 +19,8 @@ def update_vlm_prompts_for_model(model_name):
"""Update the task dropdown choices based on selected model."""
from modules.caption import vqa
prompts = vqa.get_prompts_for_model(model_name)
return gr.update(choices=prompts, value=prompts[0] if prompts else default_task)
value = default_task if default_task in prompts else (prompts[0] if prompts else default_task)
return gr.update(choices=prompts, value=value)
def update_vlm_prompt_placeholder(question):