From 691ddf434e24df0917aaa413ff58220a3cc3cc1d Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Wed, 15 Jul 2026 07:02:33 -0600 Subject: [PATCH] Add Qwen3.6 VL captioner. --- .../captioner/BaseCaptioner.py | 1 + .../captioner/Qwen3VLCaptioner.py | 66 +++++++++++++++---- ui/src/components/CaptionSimpleJob.tsx | 7 ++ ui/src/helpers/captionOptions.ts | 6 +- ui/src/types.ts | 1 + 5 files changed, 68 insertions(+), 13 deletions(-) diff --git a/extensions_built_in/captioner/BaseCaptioner.py b/extensions_built_in/captioner/BaseCaptioner.py index cb63002..032b5f8 100644 --- a/extensions_built_in/captioner/BaseCaptioner.py +++ b/extensions_built_in/captioner/BaseCaptioner.py @@ -40,6 +40,7 @@ class CaptionConfig: self.recaption = kwargs.get("recaption", False) self.max_res = kwargs.get("max_res", 512) self.max_new_tokens = kwargs.get("max_new_tokens", 128) + self.thinking = kwargs.get("thinking", False) self.caption_prompt = kwargs.get( "caption_prompt", "Describe this image in detail." ) diff --git a/extensions_built_in/captioner/Qwen3VLCaptioner.py b/extensions_built_in/captioner/Qwen3VLCaptioner.py index 9888c7f..88bf41b 100644 --- a/extensions_built_in/captioner/Qwen3VLCaptioner.py +++ b/extensions_built_in/captioner/Qwen3VLCaptioner.py @@ -1,7 +1,8 @@ from transformers import ( - Qwen3VLForConditionalGeneration, - Qwen3VLMoeForConditionalGeneration, + AutoModelForImageTextToText, AutoProcessor, + StoppingCriteria, + StoppingCriteriaList, ) from collections import OrderedDict @@ -44,6 +45,27 @@ def patch_qwen_vl_patch_embed(model): warnings.filterwarnings("ignore") logging.disable(logging.WARNING) +# hard cap on reasoning tokens so a runaway think block cannot generate forever +MAX_THINKING_TOKENS = 4096 + + +class ThinkingBudgetCriteria(StoppingCriteria): + """For thinking models: lets the model reason freely, then counts + max_new_tokens starting from the token after so the visible answer + gets the full budget regardless of how long the reasoning ran.""" + + def __init__(self, think_end_token_id: int, max_new_tokens: int): + self.think_end_token_id = think_end_token_id + self.max_new_tokens = max_new_tokens + self.answer_start = None + + def __call__(self, input_ids, scores, **kwargs): + if self.answer_start is None: + if input_ids[0, -1].item() == self.think_end_token_id: + self.answer_start = input_ids.shape[1] + return False + return (input_ids.shape[1] - self.answer_start) >= self.max_new_tokens + class Qwen3VLCaptioner(BaseCaptioner): def __init__(self, process_id: int, job, config: OrderedDict, **kwargs): @@ -51,12 +73,7 @@ class Qwen3VLCaptioner(BaseCaptioner): def load_model(self): self.print_and_status_update("Loading Qwen3VL model") - ModelClass = ( - Qwen3VLMoeForConditionalGeneration - if "B-A" in self.caption_config.model_name_or_path - else Qwen3VLForConditionalGeneration - ) - self.model = ModelClass.from_pretrained( + self.model = AutoModelForImageTextToText.from_pretrained( self.caption_config.model_name_or_path, dtype=self.torch_dtype, device_map="cpu", @@ -100,13 +117,33 @@ class Qwen3VLCaptioner(BaseCaptioner): add_generation_prompt=True, return_dict=True, return_tensors="pt", + enable_thinking=self.caption_config.thinking, ) inputs = inputs.to(self.device_torch) + gen_kwargs = {"max_new_tokens": self.caption_config.max_new_tokens} + if self.caption_config.thinking: + think_end_token_id = self.processor.tokenizer.convert_tokens_to_ids( + "" + ) + if think_end_token_id is not None: + # give the model room to think, but start the max_new_tokens + # budget only once the think block closes + gen_kwargs = { + "max_new_tokens": MAX_THINKING_TOKENS + + self.caption_config.max_new_tokens, + "stopping_criteria": StoppingCriteriaList( + [ + ThinkingBudgetCriteria( + think_end_token_id, + self.caption_config.max_new_tokens, + ) + ] + ), + } + # Inference: Generation of the output - generated_ids = self.model.generate( - **inputs, max_new_tokens=self.caption_config.max_new_tokens - ) + generated_ids = self.model.generate(**inputs, **gen_kwargs) generated_ids_trimmed = [ out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids) @@ -117,7 +154,12 @@ class Qwen3VLCaptioner(BaseCaptioner): clean_up_tokenization_spaces=False, ) - return output_text[0].strip() + caption = output_text[0] + # thinking models (e.g. Qwen3.6) may still emit reasoning before the + # answer; keep only what follows the think block + if "" in caption: + caption = caption.split("")[-1] + return caption.strip() except Exception as e: print(f"Error processing {file_path}: {e}") traceback.print_exc() diff --git a/ui/src/components/CaptionSimpleJob.tsx b/ui/src/components/CaptionSimpleJob.tsx index c875f9f..269d4b5 100644 --- a/ui/src/components/CaptionSimpleJob.tsx +++ b/ui/src/components/CaptionSimpleJob.tsx @@ -183,6 +183,13 @@ const CaptionSimpleJob: React.FC = ({ jobConfig, setJobConfig, gpuIDs, se checked={jobConfig.config.process[0].caption.compile || false} onChange={value => setJobConfig(value, 'config.process[0].caption.compile')} /> + {additionalSections.includes('caption.thinking') && ( + setJobConfig(value, 'config.process[0].caption.thinking')} + /> + )} diff --git a/ui/src/helpers/captionOptions.ts b/ui/src/helpers/captionOptions.ts index d3f2fa6..223bb29 100644 --- a/ui/src/helpers/captionOptions.ts +++ b/ui/src/helpers/captionOptions.ts @@ -1,7 +1,7 @@ import { GroupedSelectOption, SelectOption } from "@/types"; type CaptionGroup = 'image' | 'music'; -type AdditionalSections = 'caption.model_name_or_path2' | 'caption.caption_prompt' | 'caption.max_res' | 'caption.max_new_tokens' | 'caption.fixed_caption'; +type AdditionalSections = 'caption.model_name_or_path2' | 'caption.caption_prompt' | 'caption.max_res' | 'caption.max_new_tokens' | 'caption.fixed_caption' | 'caption.thinking'; export interface CaptionOption { name: string; @@ -66,12 +66,16 @@ export const captionerTypes: CaptionOption[] = [ { value: 'Qwen/Qwen3-VL-2B-Instruct', label: 'Qwen/Qwen3-VL-2B-Instruct' }, { value: 'Qwen/Qwen3-VL-4B-Instruct', label: 'Qwen/Qwen3-VL-4B-Instruct' }, { value: 'Qwen/Qwen3-VL-8B-Instruct', label: 'Qwen/Qwen3-VL-8B-Instruct' }, + { value: 'huihui-ai/Huihui-Qwen3-VL-8B-Instruct-abliterated', label: 'huihui-ai/Huihui-Qwen3-VL-8B-Instruct-abliterated' }, { value: 'Qwen/Qwen3-VL-30B-A3B-Instruct', label: 'Qwen/Qwen3-VL-30B-A3B-Instruct' }, + { value: 'Qwen/Qwen3.6-27B', label: 'Qwen/Qwen3.6-27B' }, + { value: 'huihui-ai/Huihui-Qwen3.6-27B-abliterated', label: 'huihui-ai/Huihui-Qwen3.6-27B-abliterated' }, ], additionalSections: [ 'caption.caption_prompt', 'caption.max_res', 'caption.max_new_tokens', + 'caption.thinking', ], }, { diff --git a/ui/src/types.ts b/ui/src/types.ts index 125d6de..0633cc5 100644 --- a/ui/src/types.ts +++ b/ui/src/types.ts @@ -281,6 +281,7 @@ export interface CaptionProcessConfig { max_new_tokens?: number; fixed_caption?: string; caption_extension?: string; + thinking?: boolean; } }