| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import torch |
| 5 | import torchvision.transforms as T |
| 6 | from PIL import Image |
| 7 | |
| 8 | from vllm.multimodal.image import convert_image_mode |
| 9 | from vllm.tokenizers.hf import HfTokenizer |
| 10 | |
| 11 | from .internvl import InternVLImageProcessor, InternVLProcessor |
| 12 | |
| 13 | |
| 14 | def build_transform(input_size: int): |
| 15 | return T.Compose( |
| 16 | [ |
| 17 | T.Lambda(lambda img: convert_image_mode(img, "RGB")), |
| 18 | T.Resize( |
| 19 | (input_size, input_size), interpolation=T.InterpolationMode.BICUBIC |
| 20 | ), |
| 21 | T.ToTensor(), |
| 22 | ] |
| 23 | ) |
| 24 | |
| 25 | |
| 26 | # adapted from https://huggingface.co/nvidia/Llama-3.1-Nemotron-Nano-VL-8B-V1 |
| 27 | def find_closest_aspect_ratio( |
| 28 | aspect_ratio: float, |
| 29 | target_ratios: list[tuple[int, int]], |
| 30 | *, |
| 31 | width: int, |
| 32 | height: int, |
| 33 | image_size: int, |
| 34 | ) -> tuple[int, int]: |
| 35 | best_factor = float("-inf") |
| 36 | best_ratio = (1, 1) |
| 37 | area = width * height |
| 38 | |
| 39 | for rw, rh in target_ratios: |
| 40 | target_aspect_ratio = rw / rh |
| 41 | size_factor = min((rw * rh * image_size * image_size) / area, 0.6) |
| 42 | ratio_closeness = min( |
| 43 | target_aspect_ratio / aspect_ratio, aspect_ratio / target_aspect_ratio |
| 44 | ) |
| 45 | factor = size_factor * ratio_closeness |
| 46 | |
| 47 | if factor > best_factor: |
| 48 | best_factor = factor |
| 49 | best_ratio = (rw, rh) |
| 50 | |
| 51 | return best_ratio |
| 52 | |
| 53 | |
| 54 | def calculate_nemotron_vl_targets( |
| 55 | *, |
| 56 | orig_width: int, |
| 57 | orig_height: int, |
| 58 | target_ratios: list[tuple[int, int]], |
| 59 | image_size: int, |
| 60 | use_thumbnail: bool, |
| 61 | ) -> tuple[int, int, int]: |
| 62 | aspect_ratio = orig_width / orig_height |
| 63 | |
| 64 | # find the closest aspect ratio to the target |
| 65 | target_aspect_ratio = find_closest_aspect_ratio( |
| 66 | aspect_ratio, |
| 67 | target_ratios, |
| 68 | width=orig_width, |
| 69 | height=orig_height, |
| 70 | image_size=image_size, |
| 71 | ) |
| 72 | |
| 73 | # calculate the target width and height |
| 74 | target_width = image_size * target_aspect_ratio[0] |
| 75 | target_height = image_size * target_aspect_ratio[1] |
| 76 | blocks = target_aspect_ratio[0] * target_aspect_ratio[1] |
| 77 | |
| 78 | # add thumbnail image if num_blocks != 1 |
| 79 | if use_thumbnail and blocks != 1: |
| 80 | blocks += 1 |
| 81 | |
| 82 | return blocks, target_width, target_height |
| 83 | |
| 84 | |
| 85 | def dynamic_preprocess_nemotron_vl( |
| 86 | image: Image.Image, |
| 87 | *, |
| 88 | target_ratios: list[tuple[int, int]], |
| 89 | image_size: int, |
| 90 | use_thumbnail: bool, |
| 91 | ) -> list[Image.Image]: |
| 92 | orig_width, orig_height = image.size |
| 93 | |
| 94 | # calculate the number of blocks without thumbnail |
| 95 | blocks, target_width, target_height = calculate_nemotron_vl_targets( |
| 96 | orig_width=orig_width, |
| 97 | orig_height=orig_height, |
| 98 | target_ratios=target_ratios, |
| 99 | image_size=image_size, |
| 100 | use_thumbnail=False, |
| 101 | ) |
| 102 | |
| 103 | # resize the image |
| 104 | resized_img = image.resize((target_width, target_height)) |
| 105 | processed_images = [] |
| 106 | for i in range(blocks): |
| 107 | box = ( |
| 108 | (i % (target_width // image_size)) * image_size, |
| 109 | (i // (target_width // image_size)) * image_size, |
| 110 | ((i % (target_width // image_size)) + 1) * image_size, |
| 111 | ((i // (target_width // image_size)) + 1) * image_size, |
| 112 | ) |
| 113 | # split the image |
| 114 | split_img = resized_img.crop(box) |
| 115 | processed_images.append(split_img) |
| 116 | |
| 117 | assert len(processed_images) == blocks |
| 118 | |
| 119 | if use_thumbnail and len(processed_images) != 1: |
| 120 | thumbnail_img = image.resize((image_size, image_size)) |
| 121 | processed_images.append(thumbnail_img) |
| 122 | |
| 123 | return processed_images |
| 124 | |
| 125 | |
| 126 | def get_nemotron_vl_target_ratios( |
| 127 | min_num: int, |
| 128 | max_num: int, |
| 129 | ) -> list[tuple[int, int]]: |
| 130 | target_ratios = { |
| 131 | (i, j) |
| 132 | for n in range(min_num, max_num + 1) |
| 133 | for i in range(1, n + 1) |
| 134 | for j in range(1, n + 1) |
| 135 | if min_num <= i * j <= max_num |
| 136 | } |
| 137 | return sorted(target_ratios, key=lambda x: x[0] * x[1]) |
| 138 | |
| 139 | |
| 140 | def image_to_pixel_values_nemotron_vl( |
| 141 | image: Image.Image, |
| 142 | *, |
| 143 | input_size: int, |
| 144 | min_num: int, |
| 145 | max_num: int, |
| 146 | use_thumbnail: bool, |
| 147 | transform: T.Compose | None = None, |
| 148 | ) -> torch.Tensor: |
| 149 | target_ratios = get_nemotron_vl_target_ratios(min_num, max_num) |
| 150 | |
| 151 | if transform is None: |
| 152 | transform = build_transform(input_size=input_size) |
| 153 | |
| 154 | images = dynamic_preprocess_nemotron_vl( |
| 155 | image, |
| 156 | target_ratios=target_ratios, |
| 157 | image_size=input_size, |
| 158 | use_thumbnail=use_thumbnail, |
| 159 | ) |
| 160 | |
| 161 | pixel_values = torch.stack([transform(image) for image in images]) |
| 162 | return pixel_values |
| 163 | |
| 164 | |
| 165 | class LlamaNemotronNanoVLImageProcessor(InternVLImageProcessor): |
| 166 | def _images_to_pixel_values_lst( |
| 167 | self, |
| 168 | images: list[Image.Image], |
| 169 | min_dynamic_patch: int | None = None, |
| 170 | max_dynamic_patch: int | None = None, |
| 171 | dynamic_image_size: bool | None = None, |
| 172 | ) -> list[torch.Tensor]: |
| 173 | min_num, max_num = self.resolve_min_max_num( |
| 174 | min_dynamic_patch=min_dynamic_patch, |
| 175 | max_dynamic_patch=max_dynamic_patch, |
| 176 | dynamic_image_size=dynamic_image_size, |
| 177 | use_thumbnail=False, # Applied in image_to_pixel_values |
| 178 | ) |
| 179 | |
| 180 | return [ |
| 181 | image_to_pixel_values_nemotron_vl( |
| 182 | image, |
| 183 | input_size=self.image_size, |
| 184 | min_num=min_num, |
| 185 | max_num=max_num, |
| 186 | use_thumbnail=self.use_thumbnail, |
| 187 | transform=build_transform(self.image_size), |
| 188 | ) |
| 189 | for image in images |
| 190 | ] |
| 191 | |
| 192 | |
| 193 | class LlamaNemotronNanoVLProcessor(InternVLProcessor): |
| 194 | """This model doesn't define its own HF processor, |
| 195 | so we implement our own one here. |
| 196 | |
| 197 | The image processor is given by: |
| 198 | https://huggingface.co/nvidia/Llama-3.1-Nemotron-Nano-VL-8B-V1/blob/main/image_processing.py |
| 199 | """ |
| 200 | |
| 201 | def __init__( |
| 202 | self, |
| 203 | image_processor: LlamaNemotronNanoVLImageProcessor, |
| 204 | tokenizer: HfTokenizer, |
| 205 | *, |
| 206 | image_seq_length: int, |
| 207 | start_image_token: str = "<img>", |
| 208 | end_image_token: str = "</img>", |
| 209 | ctx_image_token: str = "<image>", |
| 210 | ) -> None: |
| 211 | super().__init__( |
| 212 | image_processor=image_processor, |
| 213 | tokenizer=tokenizer, |
| 214 | image_seq_length=image_seq_length, |
| 215 | start_image_token=start_image_token, |
| 216 | end_image_token=end_image_token, |
| 217 | ctx_image_token=ctx_image_token, |
| 218 | ) |
| 219 | |
| 220 | def get_num_image_tokens( |
| 221 | self, |
| 222 | *, |
| 223 | image_width: int, |
| 224 | image_height: int, |
| 225 | ) -> int: |
| 226 | image_processor = self.image_processor |
| 227 | target_ratios = self.resolve_target_ratios( |
| 228 | use_thumbnail=False, # Applied in calculate_targets |
| 229 | ) |
| 230 | |
| 231 | num_patches, _, _ = calculate_nemotron_vl_targets( |
| 232 | orig_width=image_width, |
| 233 | orig_height=image_height, |
| 234 | image_size=image_processor.image_size, |
| 235 | target_ratios=target_ratios, |
| 236 | use_thumbnail=image_processor.use_thumbnail, |
| 237 | ) |
| 238 | |
| 239 | return num_patches * self.image_seq_length |
| 240 | |
| 241 | |
| 242 | # SigLIP normalization constants |
| 243 | SIGLIP_MEAN = (0.5, 0.5, 0.5) |
| 244 | SIGLIP_STD = (0.5, 0.5, 0.5) |
| 245 | |
| 246 | |
| 247 | def build_siglip_transform( |
| 248 | input_size: int, |
| 249 | *, |
| 250 | normalize_on_cpu: bool = True, |
| 251 | ): |
| 252 | """Build transform for SigLIP vision encoder with normalization. |
| 253 | |
| 254 | Extends the base transform from nemotron_vl with SigLIP-specific normalization. |
| 255 | """ |
| 256 | transforms = [ |
| 257 | T.Lambda(lambda img: convert_image_mode(img, "RGB")), |
| 258 | T.Resize((input_size, input_size), interpolation=T.InterpolationMode.BICUBIC), |
| 259 | T.ToTensor() if normalize_on_cpu else T.PILToTensor(), |
| 260 | ] |
| 261 | if normalize_on_cpu: |
| 262 | transforms.append(T.Normalize(mean=SIGLIP_MEAN, std=SIGLIP_STD)) |
| 263 | |
| 264 | return T.Compose(transforms) |
| 265 | |
| 266 | |
| 267 | class LlamaNemotronVLEmbedImageProcessor(InternVLImageProcessor): |
| 268 | def __init__( |
| 269 | self, |
| 270 | image_size: int, |
| 271 | min_dynamic_patch: int, |
| 272 | max_dynamic_patch: int, |
| 273 | dynamic_image_size: bool, |
| 274 | use_thumbnail: bool, |
| 275 | *, |
| 276 | do_rescale: bool = True, |
| 277 | do_normalize: bool = True, |
| 278 | ) -> None: |
| 279 | if do_rescale != do_normalize: |
| 280 | raise ValueError( |
| 281 | "LlamaNemotronVLEmbedImageProcessor requires do_rescale and " |
| 282 | "do_normalize to be both enabled or both disabled." |
| 283 | ) |
| 284 | |
| 285 | super().__init__( |
| 286 | image_size=image_size, |
| 287 | min_dynamic_patch=min_dynamic_patch, |
| 288 | max_dynamic_patch=max_dynamic_patch, |
| 289 | dynamic_image_size=dynamic_image_size, |
| 290 | use_thumbnail=use_thumbnail, |
| 291 | ) |
| 292 | self.normalize_on_cpu = do_normalize |
| 293 | |
| 294 | def _images_to_pixel_values_lst( |
| 295 | self, |
| 296 | images: list[Image.Image], |
| 297 | min_dynamic_patch: int | None = None, |
| 298 | max_dynamic_patch: int | None = None, |
| 299 | dynamic_image_size: bool | None = None, |
| 300 | ) -> list[torch.Tensor]: |
| 301 | min_num, max_num = self.resolve_min_max_num( |
| 302 | min_dynamic_patch=min_dynamic_patch, |
| 303 | max_dynamic_patch=max_dynamic_patch, |
| 304 | dynamic_image_size=dynamic_image_size, |
| 305 | use_thumbnail=False, # Applied in image_to_pixel_values |
| 306 | ) |
| 307 | |
| 308 | return [ |
| 309 | image_to_pixel_values_nemotron_vl( |
| 310 | image, |
| 311 | input_size=self.image_size, |
| 312 | min_num=min_num, |
| 313 | max_num=max_num, |
| 314 | use_thumbnail=self.use_thumbnail, |
| 315 | transform=build_siglip_transform( |
| 316 | self.image_size, |
| 317 | normalize_on_cpu=self.normalize_on_cpu, |
| 318 | ), |
| 319 | ) |
| 320 | for image in images |
| 321 | ] |
| 322 | |
| 323 | |
| 324 | class LlamaNemotronVLEmbedProcessor(InternVLProcessor): |
| 325 | """Processor for LlamaNemotronVL embedding model. |
| 326 | |
| 327 | Inherits from NemotronVLProcessor and specializes it for embedding tasks: |
| 328 | - Uses SigLIP transform with normalization instead of base transform |
| 329 | - Uses different image context token (<IMG_CONTEXT> vs <image>) |
| 330 | """ |
| 331 | |
| 332 | def __init__( |
| 333 | self, |
| 334 | image_processor: LlamaNemotronVLEmbedImageProcessor, |
| 335 | tokenizer: HfTokenizer, |
| 336 | *, |
| 337 | image_seq_length: int, |
| 338 | start_image_token: str = "<img>", |
| 339 | end_image_token: str = "</img>", |
| 340 | ctx_image_token: str = "<IMG_CONTEXT>", |
| 341 | ) -> None: |
| 342 | super().__init__( |
| 343 | image_processor=image_processor, |
| 344 | tokenizer=tokenizer, |
| 345 | image_seq_length=image_seq_length, |
| 346 | start_image_token=start_image_token, |
| 347 | end_image_token=end_image_token, |
| 348 | ctx_image_token=ctx_image_token, |
| 349 | ) |
| 350 | |
| 351 | self.image_processor: LlamaNemotronVLEmbedImageProcessor |
| 352 | |
| 353 | def get_num_image_tokens( |
| 354 | self, |
| 355 | *, |
| 356 | image_width: int, |
| 357 | image_height: int, |
| 358 | ) -> int: |
| 359 | image_processor = self.image_processor |
| 360 | target_ratios = self.resolve_target_ratios( |
| 361 | use_thumbnail=False, # Applied in calculate_targets |
| 362 | ) |
| 363 | |
| 364 | num_patches, _, _ = calculate_nemotron_vl_targets( |
| 365 | orig_width=image_width, |
| 366 | orig_height=image_height, |
| 367 | image_size=image_processor.image_size, |
| 368 | target_ratios=target_ratios, |
| 369 | use_thumbnail=image_processor.use_thumbnail, |
| 370 | ) |
| 371 | |
| 372 | return num_patches * self.image_seq_length |