| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Vendored HuggingFace-convention processors for the Inkling Titan model. |
| 4 | |
| 5 | Implements the image processor (numba patchifier), the audio feature extractor |
| 6 | (STFT/dMel path), the composite processor, and the MM token-id constants. |
| 7 | |
| 8 | Besides raw bytes / file paths, the extractors also accept the dummy inputs |
| 9 | vLLM generates during profiling (PIL images / numpy audio arrays). |
| 10 | """ |
| 11 | |
| 12 | from __future__ import annotations |
| 13 | |
| 14 | import io |
| 15 | import math |
| 16 | from collections.abc import Sequence |
| 17 | from dataclasses import dataclass |
| 18 | |
| 19 | import numpy as np |
| 20 | import torch |
| 21 | import torch.nn.functional as F |
| 22 | from numba import njit |
| 23 | from transformers.feature_extraction_utils import ( |
| 24 | BatchFeature, |
| 25 | FeatureExtractionMixin, |
| 26 | ) |
| 27 | from transformers.image_processing_utils import BaseImageProcessor |
| 28 | from transformers.image_utils import ImageInput |
| 29 | |
| 30 | # --------------------------------------------------------------------------- |
| 31 | # MM token-id constants |
| 32 | # --------------------------------------------------------------------------- |
| 33 | |
| 34 | # Block-start marker token ids. These are real tokens the model was trained on |
| 35 | # (``<|content_image|>`` / ``<|content_audio_input|>``) that mark the start of an |
| 36 | # image/audio embedding block; they are kept verbatim in ``input_ids``. |
| 37 | IMAGE_MARKER_ID = 200005 # <|content_image|> |
| 38 | AUDIO_MARKER_ID = 200020 # <|content_audio_input|> |
| 39 | |
| 40 | # Per-patch / per-frame placeholder ids marking where tower embeddings are |
| 41 | # scattered in. These are unused slots in the padded vocabulary and the |
| 42 | # corresponding positions are always overwritten by tower embeddings. |
| 43 | IMAGE_TOKEN_ID = 200054 # <|unused_200054|> |
| 44 | AUDIO_TOKEN_ID = 200053 # <|unused_200053|> |
| 45 | |
| 46 | # --------------------------------------------------------------------------- |
| 47 | # Image processing |
| 48 | # --------------------------------------------------------------------------- |
| 49 | |
| 50 | IMAGE_MEAN = np.array([0.48145466, 0.4578275, 0.40821073], dtype=np.float32) |
| 51 | IMAGE_STD = np.array([0.26862954, 0.2613026, 0.2757771], dtype=np.float32) |
| 52 | PAD_RAW_VALUE = np.float32(-1.0 / 255.0) |
| 53 | PAD_NORM = (np.full((3,), PAD_RAW_VALUE, dtype=np.float32) - IMAGE_MEAN) / IMAGE_STD |
| 54 | |
| 55 | |
| 56 | def _validate_image_rescale( |
| 57 | rescale_image_frac: float | None, |
| 58 | rescale_image_max_upscaled_long_edge: int | None, |
| 59 | ) -> None: |
| 60 | if rescale_image_frac is not None and ( |
| 61 | not math.isfinite(rescale_image_frac) or rescale_image_frac <= 0 |
| 62 | ): |
| 63 | raise ValueError( |
| 64 | "rescale_image_frac must be positive and finite or None, " |
| 65 | f"got {rescale_image_frac}" |
| 66 | ) |
| 67 | if rescale_image_max_upscaled_long_edge is None: |
| 68 | return |
| 69 | if rescale_image_max_upscaled_long_edge <= 0: |
| 70 | raise ValueError( |
| 71 | "rescale_image_max_upscaled_long_edge must be positive or None, " |
| 72 | f"got {rescale_image_max_upscaled_long_edge}" |
| 73 | ) |
| 74 | if rescale_image_frac is None or rescale_image_frac <= 1.0: |
| 75 | raise ValueError( |
| 76 | "rescale_image_max_upscaled_long_edge requires rescale_image_frac > 1, " |
| 77 | f"got {rescale_image_frac}" |
| 78 | ) |
| 79 | |
| 80 | |
| 81 | def _scaled_image_dimensions( |
| 82 | width: int, |
| 83 | height: int, |
| 84 | rescale_image_frac: float | None, |
| 85 | rescale_image_max_upscaled_long_edge: int | None, |
| 86 | ) -> tuple[int, int]: |
| 87 | """Return the long-edge-scaled ``(width, height)``.""" |
| 88 | if rescale_image_frac is None: |
| 89 | return width, height |
| 90 | |
| 91 | long_edge = max(width, height) |
| 92 | if long_edge == 0: |
| 93 | return width, height |
| 94 | |
| 95 | target_long_edge = float(long_edge) * rescale_image_frac |
| 96 | if rescale_image_max_upscaled_long_edge is not None: |
| 97 | effective_cap = max(rescale_image_max_upscaled_long_edge, long_edge) |
| 98 | target_long_edge = min(target_long_edge, float(effective_cap)) |
| 99 | |
| 100 | ratio = target_long_edge / float(long_edge) |
| 101 | if ratio == 1.0: |
| 102 | return width, height |
| 103 | |
| 104 | def scale(value: int) -> int: |
| 105 | return max(1, math.floor(float(value) * ratio + 0.5)) |
| 106 | |
| 107 | return scale(width), scale(height) |
| 108 | |
| 109 | |
| 110 | def _load_image_bytes(image) -> bytes: |
| 111 | """Encode a PIL image as raw PNG bytes for preprocessing. |
| 112 | |
| 113 | The HF processor is always handed ``PIL.Image`` instances by vLLM, so no |
| 114 | other input types need to be supported here. |
| 115 | """ |
| 116 | if image.mode != "RGB": |
| 117 | image = image.convert("RGB") |
| 118 | buf = io.BytesIO() |
| 119 | image.save(buf, format="PNG") |
| 120 | return buf.getvalue() |
| 121 | |
| 122 | |
| 123 | @njit(cache=True) |
| 124 | def _fill_patches_numba( |
| 125 | arr: np.ndarray, |
| 126 | patch_size: int, |
| 127 | patches: np.ndarray, |
| 128 | mean: np.ndarray, |
| 129 | std: np.ndarray, |
| 130 | pad_norm: np.ndarray, |
| 131 | ) -> None: |
| 132 | h = arr.shape[0] |
| 133 | w = arr.shape[1] |
| 134 | nph = (h + patch_size - 1) // patch_size |
| 135 | npw = w // patch_size + 1 |
| 136 | inv255 = np.float32(1.0 / 255.0) |
| 137 | |
| 138 | for k in range(nph * npw): |
| 139 | i = k // npw |
| 140 | j = k - i * npw |
| 141 | y_base = i * patch_size |
| 142 | x_base = j * patch_size |
| 143 | |
| 144 | for y in range(patch_size): |
| 145 | iy = y_base + y |
| 146 | for x in range(patch_size): |
| 147 | ix = x_base + x |
| 148 | if iy < h and ix < w: |
| 149 | for c in range(3): |
| 150 | raw = np.float32(arr[iy, ix, c]) * inv255 |
| 151 | patches[k, y, x, c] = (raw - mean[c]) / std[c] |
| 152 | else: |
| 153 | for c in range(3): |
| 154 | patches[k, y, x, c] = pad_norm[c] |
| 155 | |
| 156 | |
| 157 | def _encode_image_bytes( |
| 158 | image_bytes: bytes, |
| 159 | *, |
| 160 | patch_size: int, |
| 161 | rescale_image_frac: float | None, |
| 162 | rescale_image_max_upscaled_long_edge: int | None, |
| 163 | ) -> torch.Tensor: |
| 164 | if patch_size <= 0: |
| 165 | raise ValueError("patch_size must be greater than zero") |
| 166 | _validate_image_rescale( |
| 167 | rescale_image_frac, |
| 168 | rescale_image_max_upscaled_long_edge, |
| 169 | ) |
| 170 | |
| 171 | from PIL import Image |
| 172 | |
| 173 | image = Image.open(io.BytesIO(image_bytes)).convert("RGB") |
| 174 | scaled_size = _scaled_image_dimensions( |
| 175 | image.width, |
| 176 | image.height, |
| 177 | rescale_image_frac=rescale_image_frac, |
| 178 | rescale_image_max_upscaled_long_edge=rescale_image_max_upscaled_long_edge, |
| 179 | ) |
| 180 | if scaled_size != image.size: |
| 181 | image = image.resize(scaled_size, resample=Image.Resampling.LANCZOS) |
| 182 | arr = np.array(image, dtype=np.uint8, copy=True) |
| 183 | height, width, _ = arr.shape |
| 184 | |
| 185 | nph = (height + patch_size - 1) // patch_size |
| 186 | npw = width // patch_size + 1 |
| 187 | num_patches = nph * npw |
| 188 | |
| 189 | patches = np.empty((num_patches, patch_size, patch_size, 3), dtype=np.float32) |
| 190 | _fill_patches_numba(arr, patch_size, patches, IMAGE_MEAN, IMAGE_STD, PAD_NORM) |
| 191 | |
| 192 | return ( |
| 193 | torch.from_numpy(patches) |
| 194 | .to(torch.bfloat16) |
| 195 | .view(num_patches, 1, patch_size, patch_size, 3) |
| 196 | .expand(num_patches, 2, patch_size, patch_size, 3) |
| 197 | ) |
| 198 | |
| 199 | |
| 200 | class InklingImageProcessor(BaseImageProcessor): |
| 201 | r"""Turn raw images into ``vision_patches_bthwc`` for Inkling hMLP. |
| 202 | |
| 203 | ``rescale_image_frac`` scales the long edge while preserving aspect ratio. |
| 204 | ``rescale_image_max_upscaled_long_edge`` optionally caps only upscaling and |
| 205 | therefore requires a scale factor greater than one. The defaults, ``2.0`` and |
| 206 | ``2048``, grow images toward a 2048-pixel long edge by at most 2x, while leaving |
| 207 | images already at or above 2048 unchanged. |
| 208 | """ |
| 209 | |
| 210 | model_input_names = ["vision_patches_bthwc"] |
| 211 | |
| 212 | def __init__( |
| 213 | self, |
| 214 | patch_size: int = 40, |
| 215 | rescale_image_frac: float | None = 2.0, |
| 216 | rescale_image_max_upscaled_long_edge: int | None = 2048, |
| 217 | **kwargs, |
| 218 | ): |
| 219 | if patch_size <= 0: |
| 220 | raise ValueError("patch_size must be greater than zero") |
| 221 | _validate_image_rescale( |
| 222 | rescale_image_frac, |
| 223 | rescale_image_max_upscaled_long_edge, |
| 224 | ) |
| 225 | super().__init__(**kwargs) |
| 226 | self.patch_size = patch_size |
| 227 | self.rescale_image_frac = rescale_image_frac |
| 228 | self.rescale_image_max_upscaled_long_edge = rescale_image_max_upscaled_long_edge |
| 229 | |
| 230 | def _encode_one(self, image) -> torch.Tensor: |
| 231 | return _encode_image_bytes( |
| 232 | _load_image_bytes(image), |
| 233 | patch_size=self.patch_size, |
| 234 | rescale_image_frac=self.rescale_image_frac, |
| 235 | rescale_image_max_upscaled_long_edge=self.rescale_image_max_upscaled_long_edge, |
| 236 | ) |
| 237 | |
| 238 | def preprocess( |
| 239 | self, |
| 240 | images: ImageInput | list, |
| 241 | return_tensors: str | None = "pt", |
| 242 | **kwargs, |
| 243 | ) -> BatchFeature: |
| 244 | del return_tensors, kwargs |
| 245 | if not isinstance(images, (list, tuple)): |
| 246 | images = [images] |
| 247 | |
| 248 | per_image_patches: list[torch.Tensor] = [] |
| 249 | num_patches: list[int] = [] |
| 250 | num_tokens: list[int] = [] |
| 251 | for img in images: |
| 252 | vp = self._encode_one(img) |
| 253 | n_patches = int(vp.shape[0]) |
| 254 | per_image_patches.append(vp) |
| 255 | num_patches.append(n_patches) |
| 256 | num_tokens.append(n_patches) |
| 257 | |
| 258 | if len(per_image_patches) == 1: |
| 259 | vision_patches_bthwc = per_image_patches[0] |
| 260 | elif per_image_patches: |
| 261 | vision_patches_bthwc = torch.cat(per_image_patches, dim=0) |
| 262 | else: |
| 263 | vision_patches_bthwc = torch.empty(0) |
| 264 | |
| 265 | data = { |
| 266 | "vision_patches_bthwc": vision_patches_bthwc, |
| 267 | "num_patches": num_patches, |
| 268 | "num_tokens": num_tokens, |
| 269 | } |
| 270 | return BatchFeature(data=data, tensor_type=None) |
| 271 | |
| 272 | |
| 273 | # --------------------------------------------------------------------------- |
| 274 | # Audio feature extraction |
| 275 | # --------------------------------------------------------------------------- |
| 276 | |
| 277 | |
| 278 | @dataclass |
| 279 | class InklingAudioEncoderParams: |
| 280 | """Audio preprocessing parameters used to convert raw audio into dMel bins.""" |
| 281 | |
| 282 | sample_rate: int = 16_000 |
| 283 | window_size_multiplier: float = 2.0 |
| 284 | n_fft: int | None = None |
| 285 | n_mels: int = 80 |
| 286 | num_dmel_bins: int = 16 |
| 287 | dmel_min_value: float = -7.0 |
| 288 | dmel_max_value: float = 2.0 |
| 289 | audio_token_duration_s: float = 0.05 |
| 290 | |
| 291 | |
| 292 | def _to_exact_int(value: float, name: str, tolerance: float = 1e-6) -> int: |
| 293 | rounded = round(value) |
| 294 | if abs(value - rounded) > tolerance: |
| 295 | raise ValueError(f"{name} must resolve to an integer sample count, got {value}") |
| 296 | return int(rounded) |
| 297 | |
| 298 | |
| 299 | def _hz_to_mel(frequencies: np.ndarray) -> np.ndarray: |
| 300 | """Slaney mel scale, matching the librosa/torchaudio convention.""" |
| 301 | frequencies = np.asarray(frequencies, dtype=np.float64) |
| 302 | f_sp = 200.0 / 3.0 |
| 303 | min_log_hz = 1000.0 |
| 304 | min_log_mel = min_log_hz / f_sp |
| 305 | logstep = np.log(6.4) / 27.0 |
| 306 | linear = frequencies / f_sp |
| 307 | log = ( |
| 308 | min_log_mel + np.log(np.maximum(frequencies, min_log_hz) / min_log_hz) / logstep |
| 309 | ) |
| 310 | return np.where(frequencies >= min_log_hz, log, linear) |
| 311 | |
| 312 | |
| 313 | def _mel_to_hz(mels: np.ndarray) -> np.ndarray: |
| 314 | mels = np.asarray(mels, dtype=np.float64) |
| 315 | f_sp = 200.0 / 3.0 |
| 316 | min_log_hz = 1000.0 |
| 317 | min_log_mel = min_log_hz / f_sp |
| 318 | logstep = np.log(6.4) / 27.0 |
| 319 | linear = mels * f_sp |
| 320 | log = min_log_hz * np.exp(logstep * (mels - min_log_mel)) |
| 321 | return np.where(mels >= min_log_mel, log, linear) |
| 322 | |
| 323 | |
| 324 | _MEL_BASIS_CACHE: dict[tuple[int, int, int], torch.Tensor] = {} |
| 325 | |
| 326 | |
| 327 | def _mel_basis(sample_rate: int, n_fft: int, n_mels: int) -> torch.Tensor: |
| 328 | key = (sample_rate, n_fft, n_mels) |
| 329 | cached = _MEL_BASIS_CACHE.get(key) |
| 330 | if cached is not None: |
| 331 | return cached |
| 332 | |
| 333 | fft_bins = n_fft // 2 + 1 |
| 334 | fft_freqs = np.arange(fft_bins, dtype=np.float64) * sample_rate / n_fft |
| 335 | mel_edges = _mel_to_hz( |
| 336 | np.linspace( |
| 337 | _hz_to_mel(np.array([0.0]))[0], |
| 338 | _hz_to_mel(np.array([sample_rate / 2.0]))[0], |
| 339 | n_mels + 2, |
| 340 | dtype=np.float64, |
| 341 | ) |
| 342 | ) |
| 343 | mel_widths = np.diff(mel_edges) |
| 344 | lower = (fft_freqs[None, :] - mel_edges[:-2, None]) / mel_widths[:-1, None] |
| 345 | upper = (mel_edges[2:, None] - fft_freqs[None, :]) / mel_widths[1:, None] |
| 346 | weights = np.maximum(0.0, np.minimum(lower, upper)) |
| 347 | |
| 348 | # Slaney area normalization. |
| 349 | weights *= (2.0 / (mel_edges[2:] - mel_edges[:-2]))[:, None] |
| 350 | basis = torch.from_numpy(weights.astype(np.float32, copy=False)).contiguous() |
| 351 | _MEL_BASIS_CACHE[key] = basis |
| 352 | return basis |
| 353 | |
| 354 | |
| 355 | def _dmel_bins(audio: torch.Tensor, params: InklingAudioEncoderParams) -> torch.Tensor: |
| 356 | hop_length = _to_exact_int( |
| 357 | params.audio_token_duration_s * params.sample_rate, |
| 358 | "audio_token_duration_s * sample_rate", |
| 359 | ) |
| 360 | window_size = _to_exact_int( |
| 361 | params.audio_token_duration_s |
| 362 | * params.window_size_multiplier |
| 363 | * params.sample_rate, |
| 364 | "audio_token_duration_s * window_size_multiplier * sample_rate", |
| 365 | ) |
| 366 | n_fft = params.n_fft or window_size |
| 367 | if hop_length <= 0 or window_size <= 0 or n_fft <= 0: |
| 368 | raise ValueError("audio hop length, window size, and n_fft must be positive") |
| 369 | if audio.numel() == 0: |
| 370 | return torch.empty((0, params.n_mels), dtype=torch.int32) |
| 371 | |
| 372 | right_pad = math.ceil(audio.numel() / hop_length) * hop_length - audio.numel() |
| 373 | left_pad = max(n_fft - hop_length, 0) |
| 374 | audio = F.pad(audio, (left_pad, right_pad)) |
| 375 | |
| 376 | window = torch.hann_window(window_size, periodic=True, dtype=torch.float32) |
| 377 | spec = torch.stft( |
| 378 | audio.unsqueeze(0), |
| 379 | n_fft=n_fft, |
| 380 | hop_length=hop_length, |
| 381 | win_length=window_size, |
| 382 | window=window, |
| 383 | center=False, |
| 384 | normalized=False, |
| 385 | onesided=True, |
| 386 | return_complex=True, |
| 387 | ) |
| 388 | spec_ri = torch.view_as_real(spec) |
| 389 | magnitude = ( |
| 390 | (spec_ri[..., 0].square() + spec_ri[..., 1].square()) |
| 391 | .clamp_min(1e-10) |
| 392 | .sqrt() |
| 393 | .squeeze(0) |
| 394 | ) |
| 395 | |
| 396 | mel = ( |
| 397 | _mel_basis(params.sample_rate, n_fft, params.n_mels) |
| 398 | .matmul(magnitude) |
| 399 | .clamp_min(1e-10) |
| 400 | .log10() |
| 401 | ) |
| 402 | mel = mel.to(torch.float64).clamp( |
| 403 | min=params.dmel_min_value, max=params.dmel_max_value |
| 404 | ) |
| 405 | bin_centers = torch.linspace( |
| 406 | params.dmel_min_value, |
| 407 | params.dmel_max_value, |
| 408 | params.num_dmel_bins, |
| 409 | dtype=torch.float64, |
| 410 | ) |
| 411 | dmel_bins = (mel.unsqueeze(-1) - bin_centers).abs().argmin(dim=-1) |
| 412 | return dmel_bins.to(torch.int32).T.contiguous() |
| 413 | |
| 414 | |
| 415 | class InklingAudioFeatureExtractor(FeatureExtractionMixin): |
| 416 | """Convert raw audio into Inkling dMel bins in the HF feature-extractor API.""" |
| 417 | |
| 418 | model_input_names = ["dmel_bins"] |
| 419 | |
| 420 | def __init__(self, params: dict | None = None, **kwargs): |
| 421 | super().__init__(**kwargs) |
| 422 | merged = InklingAudioEncoderParams() |
| 423 | if params: |
| 424 | for k, v in params.items(): |
| 425 | if hasattr(merged, k): |
| 426 | setattr(merged, k, v) |
| 427 | # also accept flat kwargs (HF config style) |
| 428 | for k in list(kwargs.keys()): |
| 429 | if hasattr(merged, k): |
| 430 | setattr(merged, k, kwargs[k]) |
| 431 | self.params = merged |
| 432 | |
| 433 | def _decode_one(self, audio) -> torch.Tensor: |
| 434 | # vLLM hands the feature extractor numpy arrays (the dummy-input builder |
| 435 | # during profiling, and MultiModalDataParser after resampling to the |
| 436 | # target sample rate), so no other input types need to be supported. |
| 437 | return torch.from_numpy( |
| 438 | np.ascontiguousarray(audio.astype(np.float32, copy=False)) |
| 439 | ).flatten() |
| 440 | |
| 441 | def _encode_one(self, audio) -> torch.Tensor: |
| 442 | return _dmel_bins(self._decode_one(audio), self.params) |
| 443 | |
| 444 | def __call__( |
| 445 | self, |
| 446 | audios: Sequence | None, |
| 447 | return_tensors: str | None = None, |
| 448 | **kwargs, |
| 449 | ) -> BatchFeature: |
| 450 | del return_tensors, kwargs |
| 451 | if audios is None: |
| 452 | audios = [] |
| 453 | if not isinstance(audios, (list, tuple)): |
| 454 | audios = [audios] |
| 455 | |
| 456 | dmel_bins = [self._encode_one(a) for a in audios] |
| 457 | data = { |
| 458 | # per-clip feature: dmel bins as float32 [T, n_mels] |
| 459 | "dmel_bins": [bins.to(torch.float32) for bins in dmel_bins], |
| 460 | "num_audio_tokens": [int(bins.shape[0]) for bins in dmel_bins], |
| 461 | } |
| 462 | # return_tensors intentionally ignored: per-clip features have ragged T. |
| 463 | return BatchFeature(data=data, tensor_type=None) |
| 464 | |
| 465 | |
| 466 | # --------------------------------------------------------------------------- |
| 467 | # Composite processor |
| 468 | # --------------------------------------------------------------------------- |
| 469 | |
| 470 | |
| 471 | class InklingProcessor: |
| 472 | """Bundle Inkling image + audio preprocessing with the MM token ids.""" |
| 473 | |
| 474 | def __init__( |
| 475 | self, |
| 476 | image_processor: InklingImageProcessor | None = None, |
| 477 | audio_feature_extractor: InklingAudioFeatureExtractor | None = None, |
| 478 | tokenizer=None, |
| 479 | ): |
| 480 | self.image_processor = image_processor or InklingImageProcessor() |
| 481 | self.audio_feature_extractor = ( |
| 482 | audio_feature_extractor or InklingAudioFeatureExtractor() |
| 483 | ) |
| 484 | self.tokenizer = tokenizer |
| 485 | |
| 486 | def process_images(self, images: list): |
| 487 | """Raw images -> BatchFeature(vision_patches_bthwc, num_patches, num_tokens).""" |
| 488 | return self.image_processor.preprocess(images, return_tensors="pt") |
| 489 | |
| 490 | def process_audios(self, audios: list): |
| 491 | """Raw audios -> BatchFeature(dmel_bins, num_audio_tokens).""" |
| 492 | return self.audio_feature_extractor(audios) |
| 493 | |
| 494 | |
| 495 | __all__ = [ |
| 496 | "InklingImageProcessor", |
| 497 | "InklingAudioFeatureExtractor", |
| 498 | "InklingAudioEncoderParams", |
| 499 | "InklingProcessor", |
| 500 | "IMAGE_MARKER_ID", |
| 501 | "AUDIO_MARKER_ID", |
| 502 | "IMAGE_TOKEN_ID", |
| 503 | "AUDIO_TOKEN_ID", |
| 504 | ] |