[Chore] Use upstream Transformers InklingProcessor (#61883)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Harry Mellor
co-authored by
Claude Opus 5.5
Claude Opus 5.5
parent
f019e3e105
commit
f8112704d5
@@ -4,12 +4,15 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from collections.abc import Iterable, Mapping, Sequence
|
||||
from typing import Any, cast
|
||||
|
||||
import numpy as np
|
||||
import regex as re
|
||||
import torch
|
||||
from PIL import Image
|
||||
from transformers import InklingProcessor
|
||||
from transformers.feature_extraction_utils import BatchFeature
|
||||
|
||||
from vllm.config.multimodal import (
|
||||
@@ -29,18 +32,19 @@ from vllm.multimodal.processing import (
|
||||
PromptUpdate,
|
||||
PromptUpdateDetails,
|
||||
)
|
||||
from vllm.transformers_utils.processors.inkling import (
|
||||
AUDIO_MARKER_ID,
|
||||
AUDIO_TOKEN_ID,
|
||||
IMAGE_MARKER_ID,
|
||||
IMAGE_TOKEN_ID,
|
||||
InklingAudioFeatureExtractor,
|
||||
InklingImageProcessor,
|
||||
InklingProcessor,
|
||||
)
|
||||
|
||||
from ..configs import InklingMMConfig
|
||||
|
||||
# Block-start markers (<|content_image|>, <|content_audio_input|>), kept verbatim
|
||||
IMAGE_MARKER_ID = 200005
|
||||
AUDIO_MARKER_ID = 200020
|
||||
# Per-patch / per-frame placeholders (<|unused_200054|>, <|unused_200053|>)
|
||||
IMAGE_TOKEN_ID = 200054
|
||||
AUDIO_TOKEN_ID = 200053
|
||||
|
||||
# Long-edge upscale applied before patchifying (upstream defaults to no rescale)
|
||||
DEFAULT_RESCALE_IMAGE_FRAC = 2.0
|
||||
|
||||
# Maximum audio tokens accepted per clip. At the dMel rate of 20 tokens/s
|
||||
# (50 ms hop) this is ~10 minutes of audio. It bounds the persistent per-request
|
||||
# buffers and the encoder/memory budget; longer clips are rejected up front.
|
||||
@@ -57,6 +61,24 @@ class InklingMultiModalDataParser(MultiModalDataParser):
|
||||
return super()._parse_audio_data(data)
|
||||
|
||||
|
||||
def _rescale_image(
|
||||
image: Image.Image, frac: float | None, max_upscaled_long_edge: int | None
|
||||
) -> Image.Image:
|
||||
"""Scale the long edge by `frac` with PIL LANCZOS, as the reference does."""
|
||||
image = image.convert("RGB")
|
||||
if frac is None:
|
||||
return image
|
||||
long_edge = max(image.size)
|
||||
target_long_edge = long_edge * frac
|
||||
if max_upscaled_long_edge is not None:
|
||||
target_long_edge = min(target_long_edge, max(max_upscaled_long_edge, long_edge))
|
||||
ratio = target_long_edge / long_edge
|
||||
if ratio == 1.0:
|
||||
return image
|
||||
size = tuple(max(1, math.floor(d * ratio + 0.5)) for d in image.size)
|
||||
return image.resize(size, resample=Image.Resampling.LANCZOS)
|
||||
|
||||
|
||||
def inkling_vision_enabled(config: InklingMMConfig) -> bool:
|
||||
return getattr(config.vision_config, "decoder_dmodel", None) is not None
|
||||
|
||||
@@ -70,38 +92,13 @@ class InklingProcessingInfo(BaseProcessingInfo):
|
||||
return self.ctx.get_hf_config(InklingMMConfig)
|
||||
|
||||
def get_hf_processor(self, **kwargs: object) -> InklingProcessor:
|
||||
config = self.get_hf_config()
|
||||
vision_config = config.vision_config
|
||||
audio_config = config.audio_config
|
||||
|
||||
image_processor = InklingImageProcessor(
|
||||
patch_size=getattr(vision_config, "patch_size", None) or 40,
|
||||
)
|
||||
|
||||
if inkling_audio_enabled(config):
|
||||
audio_params = {
|
||||
"n_mels": audio_config.n_mel_bins,
|
||||
"num_dmel_bins": audio_config.mel_vocab_size,
|
||||
"dmel_min_value": audio_config.dmel_min_value,
|
||||
"dmel_max_value": audio_config.dmel_max_value,
|
||||
}
|
||||
else:
|
||||
audio_params = {}
|
||||
audio_extractor = InklingAudioFeatureExtractor(params=audio_params)
|
||||
|
||||
return InklingProcessor(
|
||||
image_processor=image_processor,
|
||||
audio_feature_extractor=audio_extractor,
|
||||
tokenizer=self.get_tokenizer(),
|
||||
)
|
||||
return self.ctx.get_hf_processor(InklingProcessor, **kwargs)
|
||||
|
||||
def get_data_parser(self) -> MultiModalDataParser:
|
||||
# Audio inputs must be resampled to the dMel feature extractor's rate
|
||||
# before `process_audios` (see InklingAudioFeatureExtractor._decode_one).
|
||||
# Without a target_sr the default parser raises on any audio input.
|
||||
extractor = self.get_hf_processor().audio_feature_extractor
|
||||
# The feature extractor requires audio at its sampling rate
|
||||
feature_extractor = self.get_hf_processor().feature_extractor
|
||||
return InklingMultiModalDataParser(
|
||||
target_sr=extractor.params.sample_rate,
|
||||
target_sr=feature_extractor.sampling_rate,
|
||||
target_channels=1,
|
||||
expected_hidden_size=self._get_expected_hidden_size(),
|
||||
)
|
||||
@@ -163,9 +160,8 @@ class InklingDummyInputsBuilder(BaseDummyInputsBuilder[InklingProcessingInfo]):
|
||||
if num_audios:
|
||||
# Size the dummy at the maximum allowed audio so memory/encoder
|
||||
# budgeting reflects the largest clip we accept (MAX_AUDIO_TOKENS).
|
||||
params = self.info.get_hf_processor().audio_feature_extractor.params
|
||||
hop = int(round(params.audio_token_duration_s * params.sample_rate))
|
||||
audio_len = MAX_AUDIO_TOKENS * hop
|
||||
feature_extractor = self.info.get_hf_processor().feature_extractor
|
||||
audio_len = MAX_AUDIO_TOKENS * feature_extractor.hop_length
|
||||
mm_data["audio"] = self._get_dummy_audios(
|
||||
length=audio_len,
|
||||
num_audios=num_audios,
|
||||
@@ -203,31 +199,37 @@ class InklingMultiModalProcessor(BaseMultiModalProcessor[InklingProcessingInfo])
|
||||
data: dict[str, Any] = {"input_ids": [prompt_ids]}
|
||||
|
||||
if images:
|
||||
img_feat = processor.process_images(images)
|
||||
data["pixel_values"] = img_feat["vision_patches_bthwc"]
|
||||
data["num_patches"] = torch.tensor(
|
||||
img_feat["num_patches"], dtype=torch.int64
|
||||
image_processor = processor.image_processor
|
||||
frac = image_processor.rescale_image_frac or DEFAULT_RESCALE_IMAGE_FRAC
|
||||
max_long_edge = image_processor.rescale_image_max_upscaled_long_edge
|
||||
# Resize with PIL: torchvision LANCZOS does not match the reference
|
||||
images = [_rescale_image(img, frac, max_long_edge) for img in images]
|
||||
img_feat = image_processor(
|
||||
images, rescale_image_frac=None, return_tensors="pt"
|
||||
)
|
||||
data["pixel_values"] = img_feat["pixel_values"].to(torch.bfloat16)
|
||||
data["num_patches"] = img_feat["num_patches"].to(torch.int64)
|
||||
|
||||
if audios:
|
||||
aud_feat = processor.process_audios(audios)
|
||||
per_clip = aud_feat["dmel_bins"]
|
||||
num_audio_tokens = aud_feat["num_audio_tokens"]
|
||||
for i, n in enumerate(num_audio_tokens):
|
||||
if int(n) > MAX_AUDIO_TOKENS:
|
||||
feature_extractor = processor.feature_extractor
|
||||
aud_feat = feature_extractor(
|
||||
audios,
|
||||
sampling_rate=feature_extractor.sampling_rate,
|
||||
return_tensors="pt",
|
||||
)
|
||||
audio_input_ids = processor._extract_dmel_bins(aud_feat["input_features"])
|
||||
num_audio_tokens = aud_feat["input_features_mask"].sum(-1)
|
||||
for i, n in enumerate(num_audio_tokens.tolist()):
|
||||
if n > MAX_AUDIO_TOKENS:
|
||||
raise ValueError(
|
||||
f"Audio clip {i} produces {int(n)} tokens, exceeding the "
|
||||
f"Audio clip {i} produces {n} tokens, exceeding the "
|
||||
f"maximum of {MAX_AUDIO_TOKENS} (~10 min at 20 tokens/s). "
|
||||
"Provide a shorter clip."
|
||||
)
|
||||
if per_clip:
|
||||
input_audio_features = torch.cat(
|
||||
[torch.as_tensor(c) for c in per_clip], dim=0
|
||||
)
|
||||
else:
|
||||
input_audio_features = torch.empty(0)
|
||||
data["input_audio_features"] = input_audio_features
|
||||
data["num_audio_tokens"] = torch.tensor(num_audio_tokens, dtype=torch.int64)
|
||||
data["input_audio_features"] = torch.cat(
|
||||
[ids[:n] for ids, n in zip(audio_input_ids, num_audio_tokens)]
|
||||
)
|
||||
data["num_audio_tokens"] = num_audio_tokens.to(torch.int64)
|
||||
|
||||
processed_data = BatchFeature(data=data, tensor_type=None)
|
||||
return self._finalize_hf_mm_data(
|
||||
|
||||
@@ -39,9 +39,6 @@ __all__ = [
|
||||
"OvisProcessor",
|
||||
"Ovis2_5Processor",
|
||||
"Step3VLProcessor",
|
||||
"InklingProcessor",
|
||||
"InklingImageProcessor",
|
||||
"InklingAudioFeatureExtractor",
|
||||
]
|
||||
|
||||
_CLASS_TO_MODULE: dict[str, str] = {
|
||||
@@ -74,9 +71,6 @@ _CLASS_TO_MODULE: dict[str, str] = {
|
||||
"OvisProcessor": "vllm.transformers_utils.processors.ovis",
|
||||
"Ovis2_5Processor": "vllm.transformers_utils.processors.ovis2_5",
|
||||
"Step3VLProcessor": "vllm.transformers_utils.processors.step3_vl",
|
||||
"InklingProcessor": "vllm.transformers_utils.processors.inkling",
|
||||
"InklingImageProcessor": "vllm.transformers_utils.processors.inkling",
|
||||
"InklingAudioFeatureExtractor": "vllm.transformers_utils.processors.inkling",
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -1,504 +0,0 @@
|
||||
# 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",
|
||||
]
|
||||
Reference in New Issue
Block a user