6570c9800c
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Signed-off-by: Bugen Zhao <i@bugenzhao.com> Co-authored-by: Bugen Zhao <i@bugenzhao.com> Co-authored-by: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Isotr0py <mozf@inferact.ai> Co-authored-by: Jee Jee Li <jeejeelee@inferact.ai> Co-authored-by: Roger Wang <hey@rogerw.io> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: OpenAI Codex <codex@openai.com>
505 lines
17 KiB
Python
505 lines
17 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Vendored HuggingFace-convention processors for the Inkling Titan model.
|
|
|
|
Implements the image processor (numba patchifier), the audio feature extractor
|
|
(STFT/dMel path), the composite processor, and the MM token-id constants.
|
|
|
|
Besides raw bytes / file paths, the extractors also accept the dummy inputs
|
|
vLLM generates during profiling (PIL images / numpy audio arrays).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import io
|
|
import math
|
|
from collections.abc import Sequence
|
|
from dataclasses import dataclass
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from numba import njit
|
|
from transformers.feature_extraction_utils import (
|
|
BatchFeature,
|
|
FeatureExtractionMixin,
|
|
)
|
|
from transformers.image_processing_utils import BaseImageProcessor
|
|
from transformers.image_utils import ImageInput
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# MM token-id constants
|
|
# ---------------------------------------------------------------------------
|
|
|
|
# Block-start marker token ids. These are real tokens the model was trained on
|
|
# (``<|content_image|>`` / ``<|content_audio_input|>``) that mark the start of an
|
|
# image/audio embedding block; they are kept verbatim in ``input_ids``.
|
|
IMAGE_MARKER_ID = 200005 # <|content_image|>
|
|
AUDIO_MARKER_ID = 200020 # <|content_audio_input|>
|
|
|
|
# Per-patch / per-frame placeholder ids marking where tower embeddings are
|
|
# scattered in. These are unused slots in the padded vocabulary and the
|
|
# corresponding positions are always overwritten by tower embeddings.
|
|
IMAGE_TOKEN_ID = 200054 # <|unused_200054|>
|
|
AUDIO_TOKEN_ID = 200053 # <|unused_200053|>
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Image processing
|
|
# ---------------------------------------------------------------------------
|
|
|
|
IMAGE_MEAN = np.array([0.48145466, 0.4578275, 0.40821073], dtype=np.float32)
|
|
IMAGE_STD = np.array([0.26862954, 0.2613026, 0.2757771], dtype=np.float32)
|
|
PAD_RAW_VALUE = np.float32(-1.0 / 255.0)
|
|
PAD_NORM = (np.full((3,), PAD_RAW_VALUE, dtype=np.float32) - IMAGE_MEAN) / IMAGE_STD
|
|
|
|
|
|
def _validate_image_rescale(
|
|
rescale_image_frac: float | None,
|
|
rescale_image_max_upscaled_long_edge: int | None,
|
|
) -> None:
|
|
if rescale_image_frac is not None and (
|
|
not math.isfinite(rescale_image_frac) or rescale_image_frac <= 0
|
|
):
|
|
raise ValueError(
|
|
"rescale_image_frac must be positive and finite or None, "
|
|
f"got {rescale_image_frac}"
|
|
)
|
|
if rescale_image_max_upscaled_long_edge is None:
|
|
return
|
|
if rescale_image_max_upscaled_long_edge <= 0:
|
|
raise ValueError(
|
|
"rescale_image_max_upscaled_long_edge must be positive or None, "
|
|
f"got {rescale_image_max_upscaled_long_edge}"
|
|
)
|
|
if rescale_image_frac is None or rescale_image_frac <= 1.0:
|
|
raise ValueError(
|
|
"rescale_image_max_upscaled_long_edge requires rescale_image_frac > 1, "
|
|
f"got {rescale_image_frac}"
|
|
)
|
|
|
|
|
|
def _scaled_image_dimensions(
|
|
width: int,
|
|
height: int,
|
|
rescale_image_frac: float | None,
|
|
rescale_image_max_upscaled_long_edge: int | None,
|
|
) -> tuple[int, int]:
|
|
"""Return the long-edge-scaled ``(width, height)``."""
|
|
if rescale_image_frac is None:
|
|
return width, height
|
|
|
|
long_edge = max(width, height)
|
|
if long_edge == 0:
|
|
return width, height
|
|
|
|
target_long_edge = float(long_edge) * rescale_image_frac
|
|
if rescale_image_max_upscaled_long_edge is not None:
|
|
effective_cap = max(rescale_image_max_upscaled_long_edge, long_edge)
|
|
target_long_edge = min(target_long_edge, float(effective_cap))
|
|
|
|
ratio = target_long_edge / float(long_edge)
|
|
if ratio == 1.0:
|
|
return width, height
|
|
|
|
def scale(value: int) -> int:
|
|
return max(1, math.floor(float(value) * ratio + 0.5))
|
|
|
|
return scale(width), scale(height)
|
|
|
|
|
|
def _load_image_bytes(image) -> bytes:
|
|
"""Encode a PIL image as raw PNG bytes for preprocessing.
|
|
|
|
The HF processor is always handed ``PIL.Image`` instances by vLLM, so no
|
|
other input types need to be supported here.
|
|
"""
|
|
if image.mode != "RGB":
|
|
image = image.convert("RGB")
|
|
buf = io.BytesIO()
|
|
image.save(buf, format="PNG")
|
|
return buf.getvalue()
|
|
|
|
|
|
@njit(cache=True)
|
|
def _fill_patches_numba(
|
|
arr: np.ndarray,
|
|
patch_size: int,
|
|
patches: np.ndarray,
|
|
mean: np.ndarray,
|
|
std: np.ndarray,
|
|
pad_norm: np.ndarray,
|
|
) -> None:
|
|
h = arr.shape[0]
|
|
w = arr.shape[1]
|
|
nph = (h + patch_size - 1) // patch_size
|
|
npw = w // patch_size + 1
|
|
inv255 = np.float32(1.0 / 255.0)
|
|
|
|
for k in range(nph * npw):
|
|
i = k // npw
|
|
j = k - i * npw
|
|
y_base = i * patch_size
|
|
x_base = j * patch_size
|
|
|
|
for y in range(patch_size):
|
|
iy = y_base + y
|
|
for x in range(patch_size):
|
|
ix = x_base + x
|
|
if iy < h and ix < w:
|
|
for c in range(3):
|
|
raw = np.float32(arr[iy, ix, c]) * inv255
|
|
patches[k, y, x, c] = (raw - mean[c]) / std[c]
|
|
else:
|
|
for c in range(3):
|
|
patches[k, y, x, c] = pad_norm[c]
|
|
|
|
|
|
def _encode_image_bytes(
|
|
image_bytes: bytes,
|
|
*,
|
|
patch_size: int,
|
|
rescale_image_frac: float | None,
|
|
rescale_image_max_upscaled_long_edge: int | None,
|
|
) -> torch.Tensor:
|
|
if patch_size <= 0:
|
|
raise ValueError("patch_size must be greater than zero")
|
|
_validate_image_rescale(
|
|
rescale_image_frac,
|
|
rescale_image_max_upscaled_long_edge,
|
|
)
|
|
|
|
from PIL import Image
|
|
|
|
image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
|
|
scaled_size = _scaled_image_dimensions(
|
|
image.width,
|
|
image.height,
|
|
rescale_image_frac=rescale_image_frac,
|
|
rescale_image_max_upscaled_long_edge=rescale_image_max_upscaled_long_edge,
|
|
)
|
|
if scaled_size != image.size:
|
|
image = image.resize(scaled_size, resample=Image.Resampling.LANCZOS)
|
|
arr = np.array(image, dtype=np.uint8, copy=True)
|
|
height, width, _ = arr.shape
|
|
|
|
nph = (height + patch_size - 1) // patch_size
|
|
npw = width // patch_size + 1
|
|
num_patches = nph * npw
|
|
|
|
patches = np.empty((num_patches, patch_size, patch_size, 3), dtype=np.float32)
|
|
_fill_patches_numba(arr, patch_size, patches, IMAGE_MEAN, IMAGE_STD, PAD_NORM)
|
|
|
|
return (
|
|
torch.from_numpy(patches)
|
|
.to(torch.bfloat16)
|
|
.view(num_patches, 1, patch_size, patch_size, 3)
|
|
.expand(num_patches, 2, patch_size, patch_size, 3)
|
|
)
|
|
|
|
|
|
class InklingImageProcessor(BaseImageProcessor):
|
|
r"""Turn raw images into ``vision_patches_bthwc`` for Inkling hMLP.
|
|
|
|
``rescale_image_frac`` scales the long edge while preserving aspect ratio.
|
|
``rescale_image_max_upscaled_long_edge`` optionally caps only upscaling and
|
|
therefore requires a scale factor greater than one. The defaults, ``2.0`` and
|
|
``2048``, grow images toward a 2048-pixel long edge by at most 2x, while leaving
|
|
images already at or above 2048 unchanged.
|
|
"""
|
|
|
|
model_input_names = ["vision_patches_bthwc"]
|
|
|
|
def __init__(
|
|
self,
|
|
patch_size: int = 40,
|
|
rescale_image_frac: float | None = 2.0,
|
|
rescale_image_max_upscaled_long_edge: int | None = 2048,
|
|
**kwargs,
|
|
):
|
|
if patch_size <= 0:
|
|
raise ValueError("patch_size must be greater than zero")
|
|
_validate_image_rescale(
|
|
rescale_image_frac,
|
|
rescale_image_max_upscaled_long_edge,
|
|
)
|
|
super().__init__(**kwargs)
|
|
self.patch_size = patch_size
|
|
self.rescale_image_frac = rescale_image_frac
|
|
self.rescale_image_max_upscaled_long_edge = rescale_image_max_upscaled_long_edge
|
|
|
|
def _encode_one(self, image) -> torch.Tensor:
|
|
return _encode_image_bytes(
|
|
_load_image_bytes(image),
|
|
patch_size=self.patch_size,
|
|
rescale_image_frac=self.rescale_image_frac,
|
|
rescale_image_max_upscaled_long_edge=self.rescale_image_max_upscaled_long_edge,
|
|
)
|
|
|
|
def preprocess(
|
|
self,
|
|
images: ImageInput | list,
|
|
return_tensors: str | None = "pt",
|
|
**kwargs,
|
|
) -> BatchFeature:
|
|
del return_tensors, kwargs
|
|
if not isinstance(images, (list, tuple)):
|
|
images = [images]
|
|
|
|
per_image_patches: list[torch.Tensor] = []
|
|
num_patches: list[int] = []
|
|
num_tokens: list[int] = []
|
|
for img in images:
|
|
vp = self._encode_one(img)
|
|
n_patches = int(vp.shape[0])
|
|
per_image_patches.append(vp)
|
|
num_patches.append(n_patches)
|
|
num_tokens.append(n_patches)
|
|
|
|
if len(per_image_patches) == 1:
|
|
vision_patches_bthwc = per_image_patches[0]
|
|
elif per_image_patches:
|
|
vision_patches_bthwc = torch.cat(per_image_patches, dim=0)
|
|
else:
|
|
vision_patches_bthwc = torch.empty(0)
|
|
|
|
data = {
|
|
"vision_patches_bthwc": vision_patches_bthwc,
|
|
"num_patches": num_patches,
|
|
"num_tokens": num_tokens,
|
|
}
|
|
return BatchFeature(data=data, tensor_type=None)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Audio feature extraction
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@dataclass
|
|
class InklingAudioEncoderParams:
|
|
"""Audio preprocessing parameters used to convert raw audio into dMel bins."""
|
|
|
|
sample_rate: int = 16_000
|
|
window_size_multiplier: float = 2.0
|
|
n_fft: int | None = None
|
|
n_mels: int = 80
|
|
num_dmel_bins: int = 16
|
|
dmel_min_value: float = -7.0
|
|
dmel_max_value: float = 2.0
|
|
audio_token_duration_s: float = 0.05
|
|
|
|
|
|
def _to_exact_int(value: float, name: str, tolerance: float = 1e-6) -> int:
|
|
rounded = round(value)
|
|
if abs(value - rounded) > tolerance:
|
|
raise ValueError(f"{name} must resolve to an integer sample count, got {value}")
|
|
return int(rounded)
|
|
|
|
|
|
def _hz_to_mel(frequencies: np.ndarray) -> np.ndarray:
|
|
"""Slaney mel scale, matching the librosa/torchaudio convention."""
|
|
frequencies = np.asarray(frequencies, dtype=np.float64)
|
|
f_sp = 200.0 / 3.0
|
|
min_log_hz = 1000.0
|
|
min_log_mel = min_log_hz / f_sp
|
|
logstep = np.log(6.4) / 27.0
|
|
linear = frequencies / f_sp
|
|
log = (
|
|
min_log_mel + np.log(np.maximum(frequencies, min_log_hz) / min_log_hz) / logstep
|
|
)
|
|
return np.where(frequencies >= min_log_hz, log, linear)
|
|
|
|
|
|
def _mel_to_hz(mels: np.ndarray) -> np.ndarray:
|
|
mels = np.asarray(mels, dtype=np.float64)
|
|
f_sp = 200.0 / 3.0
|
|
min_log_hz = 1000.0
|
|
min_log_mel = min_log_hz / f_sp
|
|
logstep = np.log(6.4) / 27.0
|
|
linear = mels * f_sp
|
|
log = min_log_hz * np.exp(logstep * (mels - min_log_mel))
|
|
return np.where(mels >= min_log_mel, log, linear)
|
|
|
|
|
|
_MEL_BASIS_CACHE: dict[tuple[int, int, int], torch.Tensor] = {}
|
|
|
|
|
|
def _mel_basis(sample_rate: int, n_fft: int, n_mels: int) -> torch.Tensor:
|
|
key = (sample_rate, n_fft, n_mels)
|
|
cached = _MEL_BASIS_CACHE.get(key)
|
|
if cached is not None:
|
|
return cached
|
|
|
|
fft_bins = n_fft // 2 + 1
|
|
fft_freqs = np.arange(fft_bins, dtype=np.float64) * sample_rate / n_fft
|
|
mel_edges = _mel_to_hz(
|
|
np.linspace(
|
|
_hz_to_mel(np.array([0.0]))[0],
|
|
_hz_to_mel(np.array([sample_rate / 2.0]))[0],
|
|
n_mels + 2,
|
|
dtype=np.float64,
|
|
)
|
|
)
|
|
mel_widths = np.diff(mel_edges)
|
|
lower = (fft_freqs[None, :] - mel_edges[:-2, None]) / mel_widths[:-1, None]
|
|
upper = (mel_edges[2:, None] - fft_freqs[None, :]) / mel_widths[1:, None]
|
|
weights = np.maximum(0.0, np.minimum(lower, upper))
|
|
|
|
# Slaney area normalization.
|
|
weights *= (2.0 / (mel_edges[2:] - mel_edges[:-2]))[:, None]
|
|
basis = torch.from_numpy(weights.astype(np.float32, copy=False)).contiguous()
|
|
_MEL_BASIS_CACHE[key] = basis
|
|
return basis
|
|
|
|
|
|
def _dmel_bins(audio: torch.Tensor, params: InklingAudioEncoderParams) -> torch.Tensor:
|
|
hop_length = _to_exact_int(
|
|
params.audio_token_duration_s * params.sample_rate,
|
|
"audio_token_duration_s * sample_rate",
|
|
)
|
|
window_size = _to_exact_int(
|
|
params.audio_token_duration_s
|
|
* params.window_size_multiplier
|
|
* params.sample_rate,
|
|
"audio_token_duration_s * window_size_multiplier * sample_rate",
|
|
)
|
|
n_fft = params.n_fft or window_size
|
|
if hop_length <= 0 or window_size <= 0 or n_fft <= 0:
|
|
raise ValueError("audio hop length, window size, and n_fft must be positive")
|
|
if audio.numel() == 0:
|
|
return torch.empty((0, params.n_mels), dtype=torch.int32)
|
|
|
|
right_pad = math.ceil(audio.numel() / hop_length) * hop_length - audio.numel()
|
|
left_pad = max(n_fft - hop_length, 0)
|
|
audio = F.pad(audio, (left_pad, right_pad))
|
|
|
|
window = torch.hann_window(window_size, periodic=True, dtype=torch.float32)
|
|
spec = torch.stft(
|
|
audio.unsqueeze(0),
|
|
n_fft=n_fft,
|
|
hop_length=hop_length,
|
|
win_length=window_size,
|
|
window=window,
|
|
center=False,
|
|
normalized=False,
|
|
onesided=True,
|
|
return_complex=True,
|
|
)
|
|
spec_ri = torch.view_as_real(spec)
|
|
magnitude = (
|
|
(spec_ri[..., 0].square() + spec_ri[..., 1].square())
|
|
.clamp_min(1e-10)
|
|
.sqrt()
|
|
.squeeze(0)
|
|
)
|
|
|
|
mel = (
|
|
_mel_basis(params.sample_rate, n_fft, params.n_mels)
|
|
.matmul(magnitude)
|
|
.clamp_min(1e-10)
|
|
.log10()
|
|
)
|
|
mel = mel.to(torch.float64).clamp(
|
|
min=params.dmel_min_value, max=params.dmel_max_value
|
|
)
|
|
bin_centers = torch.linspace(
|
|
params.dmel_min_value,
|
|
params.dmel_max_value,
|
|
params.num_dmel_bins,
|
|
dtype=torch.float64,
|
|
)
|
|
dmel_bins = (mel.unsqueeze(-1) - bin_centers).abs().argmin(dim=-1)
|
|
return dmel_bins.to(torch.int32).T.contiguous()
|
|
|
|
|
|
class InklingAudioFeatureExtractor(FeatureExtractionMixin):
|
|
"""Convert raw audio into Inkling dMel bins in the HF feature-extractor API."""
|
|
|
|
model_input_names = ["dmel_bins"]
|
|
|
|
def __init__(self, params: dict | None = None, **kwargs):
|
|
super().__init__(**kwargs)
|
|
merged = InklingAudioEncoderParams()
|
|
if params:
|
|
for k, v in params.items():
|
|
if hasattr(merged, k):
|
|
setattr(merged, k, v)
|
|
# also accept flat kwargs (HF config style)
|
|
for k in list(kwargs.keys()):
|
|
if hasattr(merged, k):
|
|
setattr(merged, k, kwargs[k])
|
|
self.params = merged
|
|
|
|
def _decode_one(self, audio) -> torch.Tensor:
|
|
# vLLM hands the feature extractor numpy arrays (the dummy-input builder
|
|
# during profiling, and MultiModalDataParser after resampling to the
|
|
# target sample rate), so no other input types need to be supported.
|
|
return torch.from_numpy(
|
|
np.ascontiguousarray(audio.astype(np.float32, copy=False))
|
|
).flatten()
|
|
|
|
def _encode_one(self, audio) -> torch.Tensor:
|
|
return _dmel_bins(self._decode_one(audio), self.params)
|
|
|
|
def __call__(
|
|
self,
|
|
audios: Sequence | None,
|
|
return_tensors: str | None = None,
|
|
**kwargs,
|
|
) -> BatchFeature:
|
|
del return_tensors, kwargs
|
|
if audios is None:
|
|
audios = []
|
|
if not isinstance(audios, (list, tuple)):
|
|
audios = [audios]
|
|
|
|
dmel_bins = [self._encode_one(a) for a in audios]
|
|
data = {
|
|
# per-clip feature: dmel bins as float32 [T, n_mels]
|
|
"dmel_bins": [bins.to(torch.float32) for bins in dmel_bins],
|
|
"num_audio_tokens": [int(bins.shape[0]) for bins in dmel_bins],
|
|
}
|
|
# return_tensors intentionally ignored: per-clip features have ragged T.
|
|
return BatchFeature(data=data, tensor_type=None)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Composite processor
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class InklingProcessor:
|
|
"""Bundle Inkling image + audio preprocessing with the MM token ids."""
|
|
|
|
def __init__(
|
|
self,
|
|
image_processor: InklingImageProcessor | None = None,
|
|
audio_feature_extractor: InklingAudioFeatureExtractor | None = None,
|
|
tokenizer=None,
|
|
):
|
|
self.image_processor = image_processor or InklingImageProcessor()
|
|
self.audio_feature_extractor = (
|
|
audio_feature_extractor or InklingAudioFeatureExtractor()
|
|
)
|
|
self.tokenizer = tokenizer
|
|
|
|
def process_images(self, images: list):
|
|
"""Raw images -> BatchFeature(vision_patches_bthwc, num_patches, num_tokens)."""
|
|
return self.image_processor.preprocess(images, return_tensors="pt")
|
|
|
|
def process_audios(self, audios: list):
|
|
"""Raw audios -> BatchFeature(dmel_bins, num_audio_tokens)."""
|
|
return self.audio_feature_extractor(audios)
|
|
|
|
|
|
__all__ = [
|
|
"InklingImageProcessor",
|
|
"InklingAudioFeatureExtractor",
|
|
"InklingAudioEncoderParams",
|
|
"InklingProcessor",
|
|
"IMAGE_MARKER_ID",
|
|
"AUDIO_MARKER_ID",
|
|
"IMAGE_TOKEN_ID",
|
|
"AUDIO_TOKEN_ID",
|
|
]
|