From e2cdbe47fad3a28a921d12a700d29ba1ab84c2ca Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Thu, 29 Jan 2026 04:21:01 +0000 Subject: [PATCH] fix(caption): safetensors-only downloads, model load fixes, UI default, prefill tests - 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 --- cli/test-caption-api.py | 32 +++++++++++++++++++ modules/caption/deepseek.py | 1 + modules/caption/joycaption.py | 1 + modules/caption/joytag.py | 4 +++ modules/caption/moondream3.py | 1 + modules/caption/vqa.py | 58 +++++++++++++++++++++++++++-------- modules/ui_caption.py | 5 +-- 7 files changed, 88 insertions(+), 14 deletions(-) diff --git a/cli/test-caption-api.py b/cli/test-caption-api.py index f5ce205ee..7c9e3726a 100755 --- a/cli/test-caption-api.py +++ b/cli/test-caption-api.py @@ -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 # ========================================================================= diff --git a/modules/caption/deepseek.py b/modules/caption/deepseek.py index 680f2e497..e7d3eac0c 100644 --- a/modules/caption/deepseek.py +++ b/modules/caption/deepseek.py @@ -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) diff --git a/modules/caption/joycaption.py b/modules/caption/joycaption.py index 778d72cf2..3fc990ba8 100644 --- a/modules/caption/joycaption.py +++ b/modules/caption/joycaption.py @@ -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, ) diff --git a/modules/caption/joytag.py b/modules/caption/joytag.py index 960e79ec6..97a05b8b5 100644 --- a/modules/caption/joytag.py +++ b/modules/caption/joytag.py @@ -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: diff --git a/modules/caption/moondream3.py b/modules/caption/moondream3.py index e6e6b9a9b..12c28b1b6 100644 --- a/modules/caption/moondream3.py +++ b/modules/caption/moondream3.py @@ -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, ) diff --git a/modules/caption/vqa.py b/modules/caption/vqa.py index 5d6ef6c14..b0905ef82 100644 --- a/modules/caption/vqa.py +++ b/modules/caption/vqa.py @@ -366,7 +366,7 @@ def get_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> + Applies to models that use HuggingFace chat templates with / 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 without a matching . Args: - text: Model output text potentially containing /<\/think> tags. + text: Model output text potentially containing / 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( diff --git a/modules/ui_caption.py b/modules/ui_caption.py index 90e94387f..6498d463b 100644 --- a/modules/ui_caption.py +++ b/modules/ui_caption.py @@ -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):