160 lines
6.1 KiB
Python
160 lines
6.1 KiB
Python
from transformers import (
|
|
AutoModelForImageTextToText,
|
|
AutoProcessor,
|
|
StoppingCriteria,
|
|
StoppingCriteriaList,
|
|
)
|
|
from collections import OrderedDict
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from optimum.quanto import freeze
|
|
from toolkit.basic import flush
|
|
from toolkit.util.quantize import quantize, get_qtype
|
|
|
|
from toolkit.models.v2.text_encoders.qwen3_vl import patch_qwen_vl_patch_embed
|
|
|
|
from .BaseCaptioner import BaseCaptioner
|
|
import transformers
|
|
import logging
|
|
import traceback
|
|
import warnings
|
|
|
|
|
|
# transformers.logging.set_verbosity_error()
|
|
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 </think> 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):
|
|
super(Qwen3VLCaptioner, self).__init__(process_id, job, config, **kwargs)
|
|
|
|
def load_model(self):
|
|
self.print_and_status_update("Loading Qwen3VL model")
|
|
self.model = AutoModelForImageTextToText.from_pretrained(
|
|
self.caption_config.model_name_or_path,
|
|
dtype=self.torch_dtype,
|
|
device_map="cpu",
|
|
)
|
|
# swap the slow bf16 Conv3d patch_embed for an equivalent fast linear
|
|
patch_qwen_vl_patch_embed(self.model)
|
|
if not self.caption_config.low_vram:
|
|
self.model.to(self.device_torch)
|
|
if self.caption_config.quantize:
|
|
self.print_and_status_update("Quantizing Qwen3VL model")
|
|
# in low vram mode the model stays on cpu; quantize each layer on the
|
|
# gpu and move it back so the math is fast without holding the whole
|
|
# model in vram
|
|
# lm_head is huge (vocab x hidden) and quality-critical; quantizing it
|
|
# needs a ~4x transient allocation that can OOM, so keep it in full
|
|
# precision
|
|
quantize(
|
|
self.model,
|
|
weights=get_qtype(self.caption_config.qtype),
|
|
exclude=["lm_head", "*.lm_head"],
|
|
quantize_device=self.device_torch
|
|
if self.caption_config.low_vram
|
|
else None,
|
|
)
|
|
freeze(self.model)
|
|
flush()
|
|
self.processor = AutoProcessor.from_pretrained(
|
|
self.caption_config.model_name_or_path
|
|
)
|
|
if self.caption_config.low_vram:
|
|
self.model.to(self.device_torch)
|
|
flush()
|
|
|
|
def get_caption_for_file(self, file_path: str) -> str:
|
|
img = self.load_pil_image(file_path, max_res=self.caption_config.max_res)
|
|
try:
|
|
messages = [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"type": "image",
|
|
"image": img,
|
|
},
|
|
{"type": "text", "text": self.resolve_prompt(self.caption_config.caption_prompt, file_path)},
|
|
],
|
|
}
|
|
]
|
|
|
|
# Preparation for inference
|
|
inputs = self.processor.apply_chat_template(
|
|
messages,
|
|
tokenize=True,
|
|
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(
|
|
"</think>"
|
|
)
|
|
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, **gen_kwargs)
|
|
generated_ids_trimmed = [
|
|
out_ids[len(in_ids) :]
|
|
for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
|
|
]
|
|
output_text = self.processor.batch_decode(
|
|
generated_ids_trimmed,
|
|
skip_special_tokens=True,
|
|
clean_up_tokenization_spaces=False,
|
|
)
|
|
|
|
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 "</think>" in caption:
|
|
caption = caption.split("</think>")[-1]
|
|
return caption.strip()
|
|
except Exception as e:
|
|
print(f"Error processing {file_path}: {e}")
|
|
traceback.print_exc()
|
|
return None
|