| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | |
| 5 | import math |
| 6 | |
| 7 | import torch |
| 8 | from torchvision.transforms import InterpolationMode |
| 9 | from transformers import AutoTokenizer |
| 10 | from transformers.feature_extraction_utils import BatchFeature |
| 11 | from transformers.image_utils import ( |
| 12 | ChannelDimension, |
| 13 | PILImageResampling, |
| 14 | SizeDict, |
| 15 | get_image_size, |
| 16 | ) |
| 17 | from transformers.models.qwen3_vl.processing_qwen3_vl import Qwen3VLProcessor |
| 18 | from transformers.models.qwen3_vl.video_processing_qwen3_vl import ( |
| 19 | Qwen3VLVideoProcessor, |
| 20 | Qwen3VLVideoProcessorInitKwargs, |
| 21 | smart_resize, |
| 22 | ) |
| 23 | from transformers.models.siglip2.image_processing_siglip2 import ( |
| 24 | Siglip2ImageProcessor, |
| 25 | convert_image_to_patches, |
| 26 | ) |
| 27 | from transformers.processing_utils import Unpack |
| 28 | from transformers.utils import TensorType |
| 29 | from transformers.video_processing_utils import BaseVideoProcessor |
| 30 | from transformers.video_utils import group_videos_by_shape, reorder_videos |
| 31 | |
| 32 | |
| 33 | def round_by_factor(number: int, factor: int) -> int: |
| 34 | """Returns the closest integer to 'number' that is divisible by 'factor'.""" |
| 35 | return round(number / factor) * factor |
| 36 | |
| 37 | |
| 38 | def ceil_by_factor(number: int, factor: int) -> int: |
| 39 | """Return the smallest multiple of ``factor`` greater than ``number``.""" |
| 40 | return math.ceil(number / factor) * factor |
| 41 | |
| 42 | |
| 43 | def floor_by_factor(number: int, factor: int) -> int: |
| 44 | """Return the largest multiple of ``factor`` less than ``number``.""" |
| 45 | return math.floor(number / factor) * factor |
| 46 | |
| 47 | |
| 48 | class Cosmos3EdgeImagesKwargs( # type: ignore[call-arg] |
| 49 | Siglip2ImageProcessor.valid_kwargs, |
| 50 | total=False, |
| 51 | ): |
| 52 | # global setting for all images, can be overridden by per-image kwargs |
| 53 | max_pixels: int | None |
| 54 | min_pixels: int | None |
| 55 | |
| 56 | per_image_kwargs: list[dict] | None |
| 57 | |
| 58 | |
| 59 | class Cosmos3EdgeImageProcessor(Siglip2ImageProcessor): |
| 60 | resample = PILImageResampling.BICUBIC |
| 61 | valid_kwargs = Cosmos3EdgeImagesKwargs |
| 62 | |
| 63 | def __init__(self, **kwargs: Unpack[Cosmos3EdgeImagesKwargs]): |
| 64 | super().__init__(**kwargs) |
| 65 | |
| 66 | def _resize_image( |
| 67 | self, |
| 68 | image: torch.Tensor, |
| 69 | max_ratio=200, |
| 70 | resample: PILImageResampling | InterpolationMode | int | None = None, |
| 71 | **kwargs, |
| 72 | ) -> torch.Tensor: |
| 73 | """Rescales the image so that the following conditions are met: |
| 74 | |
| 75 | 1. Both dimensions (height and width) are divisible by 'factor'. |
| 76 | 2. The total number of pixels is within the range ['min_pixels', 'max_pixels']. |
| 77 | 3. The aspect ratio of the image is maintained as closely as possible. |
| 78 | """ |
| 79 | image_min_pixels = kwargs.get("min_pixels") or self.size.get( |
| 80 | "shortest_edge", None |
| 81 | ) |
| 82 | image_max_pixels = kwargs.get("max_pixels") or self.size.get( |
| 83 | "longest_edge", None |
| 84 | ) |
| 85 | assert image_min_pixels is not None and image_max_pixels is not None, ( |
| 86 | "When do_resize is True, min_pixels and max_pixels must be provided." |
| 87 | ) |
| 88 | assert image_max_pixels >= image_min_pixels, ( |
| 89 | "The max_pixels of image must be greater than or equal to min_pixels." |
| 90 | ) |
| 91 | |
| 92 | _, height, width = image.shape |
| 93 | aspect_ratio = max(height, width) / min(height, width) |
| 94 | if aspect_ratio > max_ratio: |
| 95 | raise ValueError( |
| 96 | f"absolute aspect ratio must be smaller than {max_ratio}, " |
| 97 | f"got {aspect_ratio}" |
| 98 | ) |
| 99 | factor = self.merge_size * self.patch_size |
| 100 | h_bar = max(factor, round_by_factor(height, factor)) |
| 101 | w_bar = max(factor, round_by_factor(width, factor)) |
| 102 | if h_bar * w_bar > image_max_pixels: |
| 103 | beta = math.sqrt((height * width) / image_max_pixels) |
| 104 | h_bar = floor_by_factor(height / beta, factor) |
| 105 | w_bar = floor_by_factor(width / beta, factor) |
| 106 | elif h_bar * w_bar < image_min_pixels: |
| 107 | beta = math.sqrt(image_min_pixels / (height * width)) |
| 108 | h_bar = ceil_by_factor(height * beta, factor) |
| 109 | w_bar = ceil_by_factor(width * beta, factor) |
| 110 | |
| 111 | image = self.resize( |
| 112 | image, |
| 113 | size=SizeDict(height=h_bar, width=w_bar), |
| 114 | resample=resample, |
| 115 | ) |
| 116 | return image |
| 117 | |
| 118 | def _preprocess( |
| 119 | self, |
| 120 | images: list["torch.Tensor"], |
| 121 | do_resize: bool, |
| 122 | patch_size: int, |
| 123 | max_num_patches: int, |
| 124 | resample: PILImageResampling | InterpolationMode | int | None, |
| 125 | do_rescale: bool, |
| 126 | rescale_factor: float, |
| 127 | do_normalize: bool, |
| 128 | image_mean: float | list[float] | None, |
| 129 | image_std: float | list[float] | None, |
| 130 | return_tensors: str | TensorType | None, |
| 131 | **kwargs, |
| 132 | ) -> BatchFeature: |
| 133 | per_image_kwargs = kwargs.pop("per_image_kwargs", None) |
| 134 | pixel_values = [] |
| 135 | spatial_shapes = [] |
| 136 | for idx, image in enumerate(images): |
| 137 | # merge global kwargs with per-image overrides (if provided) |
| 138 | image_kwargs = kwargs.copy() |
| 139 | if ( |
| 140 | per_image_kwargs is not None |
| 141 | and idx < len(per_image_kwargs) |
| 142 | and per_image_kwargs[idx] |
| 143 | ): |
| 144 | image_kwargs.update(per_image_kwargs[idx]) |
| 145 | if do_resize: |
| 146 | image = self._resize_image( |
| 147 | image, |
| 148 | max_ratio=200, |
| 149 | resample=resample, |
| 150 | **image_kwargs, |
| 151 | ) |
| 152 | |
| 153 | image = self.rescale_and_normalize( |
| 154 | image, do_rescale, rescale_factor, do_normalize, image_mean, image_std |
| 155 | ) |
| 156 | |
| 157 | # (channels, height, width) -> (patches, patch_size**2 * channels) |
| 158 | patches = convert_image_to_patches(image, patch_size) |
| 159 | |
| 160 | num_patches_height = image.shape[1] // patch_size |
| 161 | num_patches_width = image.shape[2] // patch_size |
| 162 | |
| 163 | spatial_shapes.append((num_patches_height, num_patches_width)) |
| 164 | pixel_values.append(patches) |
| 165 | |
| 166 | pixel_values = torch.cat(pixel_values, dim=0) |
| 167 | |
| 168 | spatial_shapes = torch.tensor(spatial_shapes) |
| 169 | t_dim = torch.ones( |
| 170 | (spatial_shapes.shape[0], 1), |
| 171 | dtype=spatial_shapes.dtype, |
| 172 | device=spatial_shapes.device, |
| 173 | ) |
| 174 | image_grid_thw = torch.cat([t_dim, spatial_shapes], dim=1) |
| 175 | |
| 176 | batch_feature = BatchFeature( |
| 177 | data={ |
| 178 | "pixel_values": pixel_values, |
| 179 | "image_grid_thw": image_grid_thw, |
| 180 | }, |
| 181 | tensor_type=return_tensors, |
| 182 | ) |
| 183 | return batch_feature |
| 184 | |
| 185 | |
| 186 | class Cosmos3EdgeVideoProcessor(Qwen3VLVideoProcessor): |
| 187 | def __init__(self, **kwargs: Unpack[Qwen3VLVideoProcessorInitKwargs]): |
| 188 | super().__init__(**kwargs) |
| 189 | |
| 190 | def _preprocess( |
| 191 | self, |
| 192 | videos: list[torch.Tensor], |
| 193 | do_convert_rgb: bool = True, |
| 194 | do_resize: bool = True, |
| 195 | size: SizeDict | None = None, |
| 196 | resample: PILImageResampling |
| 197 | | InterpolationMode |
| 198 | | int |
| 199 | | None = PILImageResampling.BICUBIC, |
| 200 | do_rescale: bool = True, |
| 201 | rescale_factor: float = 1 / 255.0, |
| 202 | do_normalize: bool = True, |
| 203 | image_mean: float | list[float] | None = None, |
| 204 | image_std: float | list[float] | None = None, |
| 205 | patch_size: int | None = None, |
| 206 | return_tensors: str | TensorType | None = None, |
| 207 | **kwargs, |
| 208 | ): |
| 209 | assert patch_size is not None |
| 210 | |
| 211 | merge_size = self.merge_size |
| 212 | grouped_videos, grouped_videos_index = group_videos_by_shape(videos) |
| 213 | resized_videos_grouped = {} |
| 214 | |
| 215 | for shape, stacked_videos in grouped_videos.items(): |
| 216 | B, T, C, H, W = stacked_videos.shape |
| 217 | num_frames, height, width = T, H, W |
| 218 | if do_resize: |
| 219 | assert size is not None |
| 220 | # 1. Both dimensions (height and width) are divisible by 'factor'. |
| 221 | # 2. Total pixels are within the configured min/max range. |
| 222 | resized_height, resized_width = smart_resize( |
| 223 | num_frames=num_frames, |
| 224 | height=height, |
| 225 | width=width, |
| 226 | temporal_factor=1, |
| 227 | factor=patch_size * merge_size, |
| 228 | min_pixels=size.shortest_edge, |
| 229 | max_pixels=size.longest_edge, |
| 230 | ) |
| 231 | stacked_videos = stacked_videos.view(B * T, C, H, W) |
| 232 | # The target size is already computed, so bypass Qwen3-VL's |
| 233 | # dynamic resize, which expects an unflattened video tensor. |
| 234 | stacked_videos = BaseVideoProcessor.resize( |
| 235 | self, |
| 236 | image=stacked_videos, |
| 237 | size=SizeDict(height=resized_height, width=resized_width), |
| 238 | resample=resample, |
| 239 | ) |
| 240 | stacked_videos = stacked_videos.view( |
| 241 | B, T, C, resized_height, resized_width |
| 242 | ) |
| 243 | resized_videos_grouped[shape] = stacked_videos |
| 244 | resized_videos = reorder_videos(resized_videos_grouped, grouped_videos_index) |
| 245 | |
| 246 | # Group videos by size for further processing |
| 247 | # Resize may return videos with different sizes. |
| 248 | grouped_videos, grouped_videos_index = group_videos_by_shape(resized_videos) |
| 249 | processed_videos_grouped = {} |
| 250 | processed_grids = {} |
| 251 | for shape, stacked_videos in grouped_videos.items(): |
| 252 | resized_height, resized_width = get_image_size( |
| 253 | stacked_videos[0], channel_dim=ChannelDimension.FIRST |
| 254 | ) |
| 255 | # Fused rescale and normalize |
| 256 | stacked_videos = self.rescale_and_normalize( |
| 257 | stacked_videos, |
| 258 | do_rescale, |
| 259 | rescale_factor, |
| 260 | do_normalize, |
| 261 | image_mean, |
| 262 | image_std, |
| 263 | ) |
| 264 | patches = stacked_videos |
| 265 | |
| 266 | batch_size, grid_t, channel = patches.shape[:3] |
| 267 | grid_h, grid_w = resized_height // patch_size, resized_width // patch_size |
| 268 | |
| 269 | # [1, 16, 3, 320, 608] -> [1, 16, 3, 20, 16, 38, 16] |
| 270 | patches = patches.view( |
| 271 | batch_size, # 0 |
| 272 | grid_t, # 1 |
| 273 | channel, # 2 |
| 274 | grid_h, # 3 |
| 275 | patch_size, # 4 |
| 276 | grid_w, # 5 |
| 277 | patch_size, # 6 |
| 278 | ) |
| 279 | # [1, 16, 3, 20, 16, 38, 16] -> [1, 16, 20, 38, 16, 16, 3] |
| 280 | # Result: [batch, grid_t, grid_h, grid_w, patch_h, patch_w, channel] |
| 281 | patches = patches.permute(0, 1, 3, 5, 4, 6, 2) |
| 282 | |
| 283 | # [1, 16, 20, 38, 16, 16, 3] -> [1, 5120, 768] |
| 284 | flatten_patches = patches.reshape( |
| 285 | batch_size, |
| 286 | grid_t * grid_h * grid_w, |
| 287 | patch_size * patch_size * channel, |
| 288 | ) |
| 289 | |
| 290 | processed_videos_grouped[shape] = flatten_patches |
| 291 | processed_grids[shape] = [[grid_t, grid_h, grid_w]] * batch_size |
| 292 | |
| 293 | processed_videos = reorder_videos( |
| 294 | processed_videos_grouped, grouped_videos_index |
| 295 | ) |
| 296 | processed_grids = reorder_videos(processed_grids, grouped_videos_index) |
| 297 | pixel_values_videos = torch.cat(processed_videos, dim=0) |
| 298 | video_grid_thw = torch.tensor(processed_grids) |
| 299 | data = { |
| 300 | "pixel_values_videos": pixel_values_videos, |
| 301 | "video_grid_thw": video_grid_thw, |
| 302 | } |
| 303 | |
| 304 | return BatchFeature(data=data, tensor_type=return_tensors) |
| 305 | |
| 306 | |
| 307 | class Cosmos3EdgeProcessor(Qwen3VLProcessor): |
| 308 | """Construct a processor for Cosmos3 Edge multimodal inputs. |
| 309 | |
| 310 | This processor combines :class:`Cosmos3EdgeImageProcessor`, |
| 311 | :class:`Cosmos3EdgeVideoProcessor`, and a tokenizer into a single |
| 312 | interface. It inherits the text, image, and video input handling and the |
| 313 | decoding methods from :class:`Qwen3VLProcessor`, while using the Cosmos3 |
| 314 | Edge-specific image and video preprocessing implementations. |
| 315 | |
| 316 | Args: |
| 317 | image_processor: A :class:`Cosmos3EdgeImageProcessor` used to resize, |
| 318 | normalize, and patch image inputs. |
| 319 | tokenizer: The tokenizer used to encode and decode text. |
| 320 | video_processor: A :class:`Cosmos3EdgeVideoProcessor` used to resize, |
| 321 | normalize, and patch video inputs. |
| 322 | chat_template: An optional Jinja template used to render chat messages |
| 323 | into a tokenizable string. |
| 324 | |
| 325 | """ |
| 326 | |
| 327 | attributes = ["image_processor", "tokenizer", "video_processor"] |
| 328 | |
| 329 | # Transformers' auto loaders cannot resolve these vLLM-native image and |
| 330 | # video processors without remote code or global auto-class registration. |
| 331 | @classmethod |
| 332 | def _get_arguments_from_pretrained( |
| 333 | cls, |
| 334 | pretrained_model_name_or_path, |
| 335 | processor_dict=None, |
| 336 | **kwargs, |
| 337 | ): |
| 338 | component_kwargs = dict(kwargs) |
| 339 | component_kwargs.pop("trust_remote_code", None) |
| 340 | |
| 341 | image_processor = Cosmos3EdgeImageProcessor.from_pretrained( |
| 342 | pretrained_model_name_or_path, |
| 343 | **component_kwargs, |
| 344 | ) |
| 345 | tokenizer = AutoTokenizer.from_pretrained( |
| 346 | pretrained_model_name_or_path, |
| 347 | trust_remote_code=False, |
| 348 | **component_kwargs, |
| 349 | ) |
| 350 | video_processor = Cosmos3EdgeVideoProcessor.from_pretrained( |
| 351 | pretrained_model_name_or_path, |
| 352 | **component_kwargs, |
| 353 | ) |
| 354 | |
| 355 | return [ |
| 356 | image_processor, |
| 357 | tokenizer, |
| 358 | video_processor, |
| 359 | ] |
| 360 | |
| 361 | |
| 362 | __all__ = [ |
| 363 | "Cosmos3EdgeProcessor", |
| 364 | ] |