A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 01:30 UTC
[Model] Nano Nemotron VL - fast media preprocessing (#35657)
Signed-off-by: Natan Bagrov <nbagrov@nvidia.com>
nvnbagrov committed
Mar 8, 2026 at 12:04 UTC
b7332b058c3b0d8533395b49dea9273aa0973b4e
1 file changed
+80
-61
vllm/model_executor/models/nano_nemotron_vl.py
+80
-61
@@ -17,11 +17,11 @@ from functools import cached_property
17
from typing import Annotated, Any, Literal, TypeAlias, TypeVar
18
19
import einops
20
+import numpy as np
21
import numpy.typing as npt
22
import regex as re
23
import torch
24
import torch.nn as nn
24
-import torchvision.transforms as T
25
from PIL import Image
26
from transformers import BatchFeature, PretrainedConfig, TensorType
27
@@ -214,7 +214,12 @@ NanoNemotronVLVideoInputs: TypeAlias = (
214
215
216
def dynamic_preprocess(
217
- image, *, image_size=512, max_num_tiles=12, use_thumbnail=True, idx=0
217
+ image,
218
+ *,
219
+ image_size=512,
220
+ max_num_tiles=12,
221
+ use_thumbnail=True,
222
+ idx=0,
223
):
224
orig_width, orig_height = image.size
225
@@ -227,35 +232,44 @@ def dynamic_preprocess(
232
image_size=image_size,
233
use_thumbnail=False,
234
)
230
- # resize the image
231
- resized_img = image.resize((target_width, target_height))
232
- processed_images = []
233
- for i in range(blocks):
234
- box = (
235
- (i % (target_width // image_size)) * image_size,
236
- (i // (target_width // image_size)) * image_size,
237
- ((i % (target_width // image_size)) + 1) * image_size,
238
- ((i // (target_width // image_size)) + 1) * image_size,
239
- )
240
- # split the image
241
- split_img = resized_img.crop(box)
242
- processed_images.append(split_img)
243
- assert len(processed_images) == blocks
244
- if use_thumbnail and len(processed_images) != 1:
245
- thumbnail_img = image.resize((image_size, image_size))
246
- processed_images.append(thumbnail_img)
247
-
248
- processed_images = [
249
- img.convert("RGB") if img.mode != "RGB" else img for img in processed_images
250
- ]
251
- processed_images = [
252
- T.Resize((image_size, image_size), interpolation=T.InterpolationMode.BICUBIC)(
253
- img
235
+
236
+ image = np.asarray(
237
+ image.convert("RGB") if image.mode != "RGB" else image, dtype=np.uint8
238
+ )
239
+
240
+ image = torch.from_numpy(image).unsqueeze(0) # (1, H, W, 3)
241
+ image = image.permute(0, 3, 1, 2) # (1, 3, H, W)
242
+
243
+ resized_img = torch.nn.functional.interpolate(
244
+ image,
245
+ size=(target_height, target_width),
246
+ mode="bicubic",
247
+ align_corners=False,
248
+ antialias=True,
249
+ )
250
+ B, C, H, W = resized_img.shape
251
+ hp, wp = H // image_size, W // image_size
252
+ patches = (
253
+ resized_img.reshape(B, C, hp, image_size, wp, image_size)
254
+ .permute(0, 2, 4, 1, 3, 5)
255
+ .reshape(B * hp * wp, C, image_size, image_size)
256
+ / 255.0
257
+ )
258
+
259
+ if use_thumbnail and patches.shape[0] > 1:
260
+ thumb = (
261
+ torch.nn.functional.interpolate(
262
+ image,
263
+ size=(image_size, image_size),
264
+ mode="bicubic",
265
+ align_corners=False,
266
+ antialias=True,
267
+ )
268
+ / 255.0
269
)
255
- for img in processed_images
256
- ]
257
- processed_images = [T.ToTensor()(img) for img in processed_images]
258
- return processed_images
270
+ patches = torch.cat([patches, thumb], dim=0)
271
+
272
+ return list(patches)
273
274
275
def image_to_pixel_values(
@@ -287,22 +301,21 @@ def video_to_pixel_values(
301
) -> torch.Tensor:
302
assert max_num_tiles == 1, "Video modality always uses one tile"
303
290
- # Convert each frame to a single resized tile tensor consistent
291
- # with image path
292
- frames_tensors: list[torch.Tensor] = []
293
- for frame in video:
294
- pil_frame = dynamic_preprocess(
295
- Image.fromarray(frame, mode="RGB"),
296
- image_size=input_size,
297
- max_num_tiles=max_num_tiles,
298
- use_thumbnail=use_thumbnail,
299
- idx=0,
304
+ # (num_frames, H, W, C) -> (num_frames, C, H, W)
305
+ video_tensor = torch.from_numpy(video).permute(0, 3, 1, 2)
306
+
307
+ if video_tensor.shape[2] != input_size or video_tensor.shape[3] != input_size:
308
+ video_tensor = torch.nn.functional.interpolate(
309
+ video_tensor,
310
+ size=(input_size, input_size),
311
+ mode="bicubic",
312
+ align_corners=False,
313
+ antialias=True,
314
)
301
- # dynamic_preprocess returns tensors already; take the single tile
302
- assert len(pil_frame) >= 1
303
- frames_tensors.append(pil_frame[-1])
315
305
- return torch.stack(frames_tensors)
316
+ video_tensor = video_tensor / 255.0
317
+
318
+ return video_tensor
319
320
321
def input_conditioner(x, norm_mean, norm_std):
@@ -346,12 +359,6 @@ class DynamicResolutionImageTiler:
359
self._factor_max = factor_max
360
self.norm_mean = torch.tensor(norm_mean).reshape(3, 1, 1)
361
self.norm_std = torch.tensor(norm_std).reshape(3, 1, 1)
349
- self._transform = T.Compose(
350
- [
351
- T.Lambda(lambda img: img.convert("RGB") if img.mode != "RGB" else img),
352
- T.ToTensor(),
353
- ]
354
- )
362
assert downsample_ratio < 1
363
reduction_factor = 1 / downsample_ratio
364
assert reduction_factor == 2.0
@@ -441,15 +448,25 @@ class DynamicResolutionImageTiler:
448
patch_size: tuple[int, int]
449
450
def apply_params(self, params: DynamicResolutionParams) -> list[torch.Tensor]:
444
- resized_img = params.media.resize(
445
- (
446
- params.patch_size[0] * self._patch_size,
447
- params.patch_size[1] * self._patch_size,
451
+ target_size = (
452
+ params.patch_size[1] * self._patch_size,
453
+ params.patch_size[0] * self._patch_size,
454
+ )
455
+ image = np.asarray(
456
+ params.media.convert("RGB") if params.media.mode != "RGB" else params.media,
457
+ dtype=np.uint8,
458
+ )
459
+ resized_img = (
460
+ torch.nn.functional.interpolate(
461
+ torch.from_numpy(image).unsqueeze(0).permute(0, 3, 1, 2),
462
+ size=target_size,
463
+ mode="bicubic",
464
+ align_corners=False,
465
+ antialias=True,
466
)
467
+ / 255.0
468
)
450
- processed_images = [resized_img]
451
-
452
- return [self._transform(img) for img in processed_images]
469
+ return list(resized_img)
470
471
def process_media(
472
self,
@@ -803,6 +820,7 @@ class BaseNanoNemotronVLProcessor(ABC):
820
image_repl = self.get_image_repl(feature_size, num_patches)
821
parts[i] = parts[i].replace("<image>", image_repl.full)
822
text = ["".join(parts)]
823
+
824
return text, image_inputs
825
826
def _make_batch_input(self, input_item: Any | list[Any] | None = None):
@@ -922,14 +940,14 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):
940
frames_indices_lst = [
941
metadata["frames_indices"] for metadata in video_metadata_lst
942
]
925
-
943
+ video_num_patches = torch.tensor(
944
+ [len(item) for item in pixel_values_lst_video]
945
+ )
946
video_inputs = {
947
"pixel_values_flat_video": input_conditioner(
948
torch.cat(pixel_values_lst_video), self.norm_mean, self.norm_std
949
),
930
- "video_num_patches": torch.tensor(
931
- [len(item) for item in pixel_values_lst_video]
932
- ),
950
+ "video_num_patches": video_num_patches,
951
"frames_indices": frames_indices_lst,
952
"frame_duration_ms": torch.tensor(frame_duration_ms_lst),
953
}
@@ -985,6 +1003,7 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):
1003
video_repl.full, skip_special_tokens=False
1004
)
1005
text = [t.replace("<video>", video_repl_text, 1) for t in text]
1006
+
1007
return text, video_inputs
1008
1009
def _preprocess_audio(