main
py 539 lines 18.7 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Custom processor for Moondream3 model."""
4
5 import math
6
7 import numpy as np
8 import torch
9 from PIL import Image
10 from transformers import AutoProcessor, BatchFeature
11 from transformers.image_utils import ImageInput
12 from transformers.processing_utils import ProcessingKwargs, ProcessorMixin, Unpack
13 from transformers.tokenization_utils_base import (
14 PreTokenizedInput,
15 PreTrainedTokenizerBase,
16 TextInput,
17 )
18
19 from vllm.multimodal.image import convert_image_mode
20
21 __all__ = ["Moondream3Processor"]
22
23
24 class Moondream3ProcessorKwargs(ProcessingKwargs, total=False): # type: ignore[call-arg]
25 _defaults = {
26 "text_kwargs": {
27 "padding": False,
28 },
29 "images_kwargs": {
30 "max_crops": 12,
31 "overlap_margin": 4,
32 "crop_size": 378,
33 "patch_size": 14,
34 "convert_to_rgb": True,
35 "return_tensors": "pt",
36 },
37 }
38
39
40 def select_tiling(
41 height: int, width: int, crop_size: int, max_crops: int
42 ) -> tuple[int, int]:
43 """Determine the optimal number of tiles to cover an image."""
44 if height <= crop_size or width <= crop_size:
45 return (1, 1)
46
47 min_h = math.ceil(height / crop_size)
48 min_w = math.ceil(width / crop_size)
49
50 if min_h * min_w > max_crops:
51 ratio = math.sqrt(max_crops / (min_h * min_w))
52 return (max(1, math.floor(min_h * ratio)), max(1, math.floor(min_w * ratio)))
53
54 h_tiles = math.floor(math.sqrt(max_crops * height / width))
55 w_tiles = math.floor(math.sqrt(max_crops * width / height))
56
57 h_tiles = max(h_tiles, min_h)
58 w_tiles = max(w_tiles, min_w)
59
60 if h_tiles * w_tiles > max_crops:
61 if w_tiles > h_tiles:
62 w_tiles = math.floor(max_crops / h_tiles)
63 else:
64 h_tiles = math.floor(max_crops / w_tiles)
65
66 return (max(1, h_tiles), max(1, w_tiles))
67
68
69 class Moondream3Processor(ProcessorMixin):
70 """Constructs a Moondream3 processor which handles image preprocessing
71 and tokenization for the Moondream3 multimodal model.
72
73 Args:
74 tokenizer: The tokenizer to use for text processing.
75 chat_template: Optional chat template string.
76 crop_size: Size of each image crop.
77 max_crops: Maximum number of crops per image.
78 overlap_margin: Margin for overlapping crops in patches.
79 patch_size: Size of each patch.
80
81 """
82
83 attributes = ["tokenizer"]
84 valid_kwargs = [
85 "chat_template",
86 "crop_size",
87 "max_crops",
88 "overlap_margin",
89 "patch_size",
90 ]
91
92 tokenizer_class = "AutoTokenizer"
93 # Use separate tokenizer repo
94 _tokenizer_repo = "moondream/starmie-v1"
95
96 # Default chat template for Moondream3
97 # Moondream uses special tokens for prompting:
98 # - Token 0 (<|endoftext|>): BOS token (ALWAYS present at position 0)
99 # - Token 1 (<|md_reserved_0|>): Start of instruction
100 # - Token 2 (<|md_reserved_1|>): Separator before question
101 # - Token 3 (<|md_reserved_2|>): End of question / start of answer
102 #
103 # Task routing based on text prefix:
104 # "caption [short|normal|long]" → describe<|md_reserved_1|>{length}
105 # "describe [short|normal|long]" → describe<|md_reserved_1|>{length}
106 # otherwise → query<|md_reserved_1|><text>
107 #
108 # Format with image:
109 # <|endoftext|><image><|md_reserved_0|>{task}<|md_reserved_1|>{q}<|md_reserved_2|>
110 # Format without image:
111 # <|endoftext|><|md_reserved_0|>{task}<|md_reserved_1|>{q}<|md_reserved_2|>
112 _default_chat_template = (
113 "{% for message in messages %}"
114 "{% if message['role'] == 'user' %}"
115 "{% if message['content'] is string %}"
116 # Simple string content (with image assumed) - route by prefix
117 "<|endoftext|><image><|md_reserved_0|>"
118 "{% if message['content'] == 'caption' %}"
119 "describe<|md_reserved_1|>normal<|md_reserved_2|>"
120 "{% elif message['content'].startswith('caption ') %}"
121 "describe<|md_reserved_1|>{{ message['content'][8:] }}<|md_reserved_2|>"
122 "{% elif message['content'] == 'describe' %}"
123 "describe<|md_reserved_1|>normal<|md_reserved_2|>"
124 "{% elif message['content'].startswith('describe ') %}"
125 "describe<|md_reserved_1|>{{ message['content'][9:] }}<|md_reserved_2|>"
126 "{% else %}"
127 "query<|md_reserved_1|>{{ message['content'] }}<|md_reserved_2|>"
128 "{% endif %}"
129 "{% else %}"
130 # List content - build Moondream's image prefix independently of
131 # OpenAI-style content part order, then render the text task.
132 "<|endoftext|>"
133 "{% for content in message['content'] %}"
134 "{% if content['type'] in ['image', 'image_url', 'input_image', 'image_pil'] %}" # noqa: E501
135 "<image>"
136 "{% endif %}"
137 "{% endfor %}"
138 "{% for content in message['content'] %}"
139 "{% if content['type'] == 'text' %}"
140 "<|md_reserved_0|>"
141 "{% if content['text'] == 'caption' %}"
142 "describe<|md_reserved_1|>normal<|md_reserved_2|>"
143 "{% elif content['text'].startswith('caption ') %}"
144 "describe<|md_reserved_1|>{{ content['text'][8:] }}<|md_reserved_2|>"
145 "{% elif content['text'] == 'describe' %}"
146 "describe<|md_reserved_1|>normal<|md_reserved_2|>"
147 "{% elif content['text'].startswith('describe ') %}"
148 "describe<|md_reserved_1|>{{ content['text'][9:] }}<|md_reserved_2|>"
149 "{% else %}"
150 "query<|md_reserved_1|>{{ content['text'] }}<|md_reserved_2|>"
151 "{% endif %}"
152 "{% endif %}"
153 "{% endfor %}"
154 "{% endif %}"
155 "{% elif message['role'] == 'assistant' %}"
156 "{{ message['content'] }}"
157 "{% endif %}"
158 "{% endfor %}"
159 )
160
161 def __init__(
162 self,
163 tokenizer: PreTrainedTokenizerBase | None = None,
164 chat_template: str | None = None,
165 crop_size: int = 378,
166 max_crops: int = 12,
167 overlap_margin: int = 4,
168 patch_size: int = 14,
169 **kwargs,
170 ):
171 self.image_token = "<image>"
172 self.crop_size = crop_size
173 self.max_crops = max_crops
174 self.overlap_margin = overlap_margin
175 self.patch_size = patch_size
176
177 # Number of patches per crop (27x27 = 729 for 378/14)
178 self.patches_per_crop = (crop_size // patch_size) ** 2
179
180 # Use default chat template if none provided
181 if chat_template is None:
182 chat_template = self._default_chat_template
183
184 super().__init__(tokenizer, chat_template=chat_template)
185
186 @classmethod
187 def from_pretrained(
188 cls,
189 pretrained_model_name_or_path,
190 **kwargs,
191 ):
192 """Load the processor, using a separate tokenizer repo.
193
194 The moondream3 model uses a custom tokenizer from 'moondream/starmie-v1'
195 instead of having tokenizer files in the model repo.
196 """
197 from transformers import AutoTokenizer, TokenizersBackend
198 from transformers.utils import cached_file
199
200 tokenizer = kwargs.pop("tokenizer", None)
201
202 tokenizer_kwargs = {
203 "trust_remote_code": kwargs.get("trust_remote_code", False),
204 }
205 for key in (
206 "cache_dir",
207 "force_download",
208 "local_files_only",
209 "revision",
210 "subfolder",
211 "token",
212 ):
213 if key in kwargs:
214 tokenizer_kwargs[key] = kwargs[key]
215
216 cached_file_kwargs = {
217 key: tokenizer_kwargs[key]
218 for key in (
219 "cache_dir",
220 "force_download",
221 "local_files_only",
222 "revision",
223 "subfolder",
224 "token",
225 )
226 if key in tokenizer_kwargs
227 }
228
229 def load_tokenizer(repo_or_path):
230 try:
231 return AutoTokenizer.from_pretrained(repo_or_path, **tokenizer_kwargs)
232 except Exception:
233 tokenizer_file = cached_file(
234 repo_or_path,
235 "tokenizer.json",
236 **cached_file_kwargs,
237 )
238 return TokenizersBackend(
239 tokenizer_file=tokenizer_file,
240 clean_up_tokenization_spaces=False,
241 )
242
243 if isinstance(tokenizer, str):
244 tokenizer = load_tokenizer(tokenizer)
245
246 if tokenizer is None:
247 # Prefer model-local tokenizer files first. If unavailable, fall
248 # back to moondream's dedicated tokenizer repository.
249 try:
250 tokenizer = load_tokenizer(pretrained_model_name_or_path)
251 except Exception:
252 tokenizer = load_tokenizer(cls._tokenizer_repo)
253
254 # Configure special tokens for Moondream3
255 # BOS and EOS are both token 0 (<|endoftext|>), matching the native
256 # config (TokenizerConfig.bos_id=0, eos_id=0). This is standard for
257 # GPT-2 style models where <|endoftext|> signals both start and end.
258 # Token 1 (<|md_reserved_0|>) is a template delimiter, NOT the EOS.
259 tokenizer.bos_token = "<|endoftext|>"
260 tokenizer.bos_token_id = 0
261 tokenizer.eos_token = "<|endoftext|>"
262 tokenizer.eos_token_id = 0
263
264 # Extract processor-specific kwargs
265 crop_size = kwargs.pop("crop_size", 378)
266 max_crops = kwargs.pop("max_crops", 12)
267 overlap_margin = kwargs.pop("overlap_margin", 4)
268 patch_size = kwargs.pop("patch_size", 14)
269 chat_template = kwargs.pop("chat_template", None)
270
271 # Set default chat template on tokenizer if not already set
272 if chat_template is None:
273 chat_template = cls._default_chat_template
274 if tokenizer.chat_template is None:
275 tokenizer.chat_template = chat_template
276
277 return cls(
278 tokenizer=tokenizer,
279 chat_template=chat_template,
280 crop_size=crop_size,
281 max_crops=max_crops,
282 overlap_margin=overlap_margin,
283 patch_size=patch_size,
284 )
285
286 def __call__(
287 self,
288 images: ImageInput = None,
289 text: TextInput
290 | PreTokenizedInput
291 | list[TextInput]
292 | list[PreTokenizedInput] = None,
293 **kwargs: Unpack[Moondream3ProcessorKwargs],
294 ) -> BatchFeature:
295 """Process images and text for Moondream3 model.
296
297 Args:
298 images: Input images (PIL Image, numpy array, or list thereof).
299 text: Input text or list of texts.
300 **kwargs: Additional processing arguments.
301
302 Returns:
303 BatchFeature with processed inputs.
304
305 """
306 output_kwargs = self._merge_kwargs(
307 Moondream3ProcessorKwargs,
308 tokenizer_init_kwargs=self.tokenizer.init_kwargs,
309 **kwargs,
310 )
311
312 # Process images
313 image_features = {}
314 if images is not None:
315 processed_images = []
316 tilings = []
317
318 images_list = images if isinstance(images, list) else [images]
319 for image in images_list:
320 pixel_values, tiling = self.preprocess_image(
321 image, **output_kwargs["images_kwargs"]
322 )
323 processed_images.append(pixel_values)
324 tilings.append(tiling)
325
326 if processed_images:
327 image_features["pixel_values"] = processed_images
328 image_features["tilings"] = tilings
329
330 # Process text
331 if text is not None:
332 if not isinstance(text, list):
333 text = [text]
334
335 # Get text kwargs, remove keys we set ourselves
336 text_kwargs = output_kwargs.get("text_kwargs", {}).copy()
337 text_kwargs.pop("return_tensors", None)
338 text_kwargs.pop("add_special_tokens", None)
339
340 # Tokenize text
341 tokenized = self.tokenizer(
342 text,
343 add_special_tokens=True,
344 return_tensors="pt",
345 **text_kwargs,
346 )
347
348 output = BatchFeature(data=dict(tokenized))
349
350 # Add image features
351 if image_features:
352 output["pixel_values"] = image_features["pixel_values"]
353 output["tilings"] = image_features["tilings"]
354
355 return output
356
357 # If only images were provided
358 return BatchFeature(data=image_features)
359
360 @staticmethod
361 def _image_array_to_uint8(array: np.ndarray) -> np.ndarray:
362 if array.dtype == np.uint8:
363 return np.ascontiguousarray(array)
364
365 if array.dtype == np.bool_:
366 return np.ascontiguousarray(array.astype(np.uint8) * 255)
367
368 if np.issubdtype(array.dtype, np.floating):
369 array = np.nan_to_num(array, nan=0.0, posinf=255.0, neginf=0.0)
370 if array.size > 0 and array.max() <= 1.0:
371 array = array * 255.0
372 array = np.rint(array)
373
374 return np.ascontiguousarray(np.clip(array, 0, 255).astype(np.uint8))
375
376 @staticmethod
377 def _to_pil_image(image: ImageInput) -> Image.Image:
378 if isinstance(image, Image.Image):
379 return image
380
381 if isinstance(image, torch.Tensor):
382 tensor = image.detach().cpu()
383 if tensor.dtype == torch.bfloat16:
384 tensor = tensor.to(torch.float32)
385 image_array = tensor.numpy()
386 elif isinstance(image, np.ndarray):
387 image_array = image
388 else:
389 raise TypeError(
390 "Moondream3 images must be PIL images, numpy arrays, "
391 f"or torch tensors, got {type(image)!r}."
392 )
393
394 if image_array.ndim == 2:
395 image_array = Moondream3Processor._image_array_to_uint8(image_array)
396 return Image.fromarray(image_array)
397
398 if image_array.ndim != 3:
399 raise ValueError(
400 "Moondream3 image arrays must have 2 or 3 dimensions, "
401 f"got shape {image_array.shape}."
402 )
403
404 channel_dims = (1, 3, 4)
405 if image_array.shape[-1] not in channel_dims:
406 if image_array.shape[0] not in channel_dims:
407 raise ValueError(
408 "Moondream3 image arrays must be HWC or CHW with 1, 3, "
409 f"or 4 channels, got shape {image_array.shape}."
410 )
411 image_array = np.transpose(image_array, (1, 2, 0))
412
413 image_array = Moondream3Processor._image_array_to_uint8(image_array)
414 if image_array.shape[-1] == 1:
415 image_array = image_array[..., 0]
416
417 return Image.fromarray(image_array)
418
419 def preprocess_image(
420 self,
421 image: ImageInput,
422 max_crops: int = 12,
423 overlap_margin: int = 4,
424 crop_size: int = 378,
425 patch_size: int = 14,
426 convert_to_rgb: bool = True,
427 return_tensors: str = "pt",
428 ) -> tuple[torch.Tensor, tuple[int, int]]:
429 """Preprocess an image using overlap-and-resize cropping strategy.
430
431 Args:
432 image: Input PIL image, numpy array, or torch tensor.
433 max_crops: Maximum number of crops.
434 overlap_margin: Margin for overlapping in patches.
435 crop_size: Size of each crop.
436 patch_size: Size of each patch.
437 convert_to_rgb: Whether to convert to RGB.
438 return_tensors: Return type ("pt" for PyTorch).
439
440 Returns:
441 Tuple of (pixel_values tensor, tiling tuple).
442
443 """
444 image = self._to_pil_image(image)
445 if convert_to_rgb:
446 image = convert_image_mode(image, "RGB")
447
448 # Convert to numpy array
449 image_array = np.array(image)
450 original_h, original_w = image_array.shape[:2]
451
452 margin_pixels = patch_size * overlap_margin
453 total_margin_pixels = margin_pixels * 2
454
455 crop_patches = crop_size // patch_size
456 crop_window_patches = crop_patches - (2 * overlap_margin)
457 crop_window_size = crop_window_patches * patch_size
458
459 tiling = select_tiling(
460 original_h - total_margin_pixels,
461 original_w - total_margin_pixels,
462 crop_window_size,
463 max_crops,
464 )
465
466 n_crops = tiling[0] * tiling[1] + 1
467 crops = np.zeros((n_crops, crop_size, crop_size, 3), dtype=np.uint8)
468
469 target_size = (
470 tiling[0] * crop_window_size + total_margin_pixels,
471 tiling[1] * crop_window_size + total_margin_pixels,
472 )
473
474 # Resize image
475 pil_img = Image.fromarray(image_array)
476 resized = pil_img.resize(
477 (int(target_size[1]), int(target_size[0])),
478 resample=Image.Resampling.LANCZOS,
479 )
480 resized_array = np.asarray(resized)
481
482 # Create global crop
483 global_pil = pil_img.resize(
484 (crop_size, crop_size), resample=Image.Resampling.LANCZOS
485 )
486 crops[0] = np.asarray(global_pil)
487
488 # Create local crops
489 for i in range(tiling[0]):
490 for j in range(tiling[1]):
491 y0 = i * crop_window_size
492 x0 = j * crop_window_size
493 y_end = min(y0 + crop_size, resized_array.shape[0])
494 x_end = min(x0 + crop_size, resized_array.shape[1])
495
496 crop_region = resized_array[y0:y_end, x0:x_end]
497 crop_idx = 1 + i * tiling[1] + j
498 h_slice = slice(None, crop_region.shape[0])
499 w_slice = slice(None, crop_region.shape[1])
500 crops[crop_idx, h_slice, w_slice] = crop_region
501
502 # Convert to tensor: (n_crops, H, W, C) -> (n_crops, C, H, W)
503 pixel_values = np.transpose(crops, (0, 3, 1, 2))
504
505 if return_tensors == "pt":
506 # Match HF reference preprocessing exactly: convert uint8 crops to
507 # bfloat16 before in-place normalization.
508 pixel_values = (
509 torch.from_numpy(pixel_values)
510 .to(dtype=torch.bfloat16)
511 .div_(255.0)
512 .sub_(0.5)
513 .div_(0.5)
514 )
515 else:
516 pixel_values = pixel_values.astype(np.float32) / 255.0
517 pixel_values = (pixel_values - 0.5) / 0.5
518
519 return pixel_values, tiling
520
521 def get_num_image_tokens(self) -> int:
522 """Return the number of image tokens (729 = 27x27 patches)."""
523 return self.patches_per_crop
524
525 def batch_decode(self, *args, **kwargs):
526 """Forward to tokenizer's batch_decode."""
527 return self.tokenizer.batch_decode(*args, **kwargs)
528
529 def decode(self, *args, **kwargs):
530 """Forward to tokenizer's decode."""
531 return self.tokenizer.decode(*args, **kwargs)
532
533 @property
534 def model_input_names(self):
535 tokenizer_input_names = self.tokenizer.model_input_names
536 return tokenizer_input_names + ["pixel_values", "tilings"]
537
538
539 AutoProcessor.register("Moondream3Processor", Moondream3Processor)