inkling_mlx/processing.py

"""Image + audio preprocessing for Inkling (MLX), ported from the reference
`InklingImageProcessor` / `InklingFeatureExtractor` / `InklingProcessor`.

  * image  -> pixel_values [num_patches, T=2, 40, 40, 3]  (feeds `VisionModel`)
  * audio  -> audio_input_ids [num_frames, 80] dMel bins  (feeds `AudioModel`)

`InklingProcessor.apply` builds the full multimodal input (input_ids + features)
from a chat message list, inserting the right number of placeholder soft-tokens.
Uses numpy/PIL + transformers' mel filterbank; no torch needed at inference.
"""

from __future__ import annotations

import math

import numpy as np

# CLIP normalization (OPENAI_CLIP_MEAN / STD), per processor_config.json
CLIP_MEAN = np.array([0.48145466, 0.4578275, 0.40821073], dtype=np.float32)
CLIP_STD = np.array([0.26862954, 0.26130258, 0.27577711], dtype=np.float32)

PATCH = 40           # image patch size (== vision patch_size)
TEMPORAL = 2         # temporal_patch_size (images duplicated across 2 frames)

# audio (processor_config.json / feature_extraction_inkling.py)
SR = 16000
HOP = 800            # audio_token_duration_s (0.05) * SR
WIN = 1600           # * window_size_multiplier (2.0)
N_FFT = 1600
N_MEL = 80
DMEL_BINS = 16
DMEL_MIN, DMEL_MAX = -7.0, 2.0

# special tokens
IMAGE_TOKEN_ID = 200054      # <|unused_200054|>  (soft-token slot)
AUDIO_TOKEN_ID = 200053      # <|unused_200053|>
IMAGE_BOS = "<|content_image|>"
AUDIO_BOS = "<|content_audio_input|>"


# ------------------------------- image -------------------------------

def preprocess_image(image) -> tuple[np.ndarray, int]:
    """PIL.Image or HxWx3 uint8 array -> (pixel_values [N,2,40,40,3] float32, N)."""
    if hasattr(image, "convert"):
        image = np.asarray(image.convert("RGB"))
    image = np.asarray(image)
    if image.ndim == 2:
        image = np.stack([image] * 3, axis=-1)
    img = image[..., :3].astype(np.float32).transpose(2, 0, 1)  # -> [C, H, W]
    C, H, W = img.shape

    num_rows = (H + PATCH - 1) // PATCH
    num_cols = W // PATCH + 1                       # reference: W//P + 1
    patches = []
    for i in range(num_rows):
        for j in range(num_cols):
            p = img[:, i * PATCH:(i + 1) * PATCH, j * PATCH:(j + 1) * PATCH]  # may be < 40
            padded = np.full((C, PATCH, PATCH), -1.0, dtype=np.float32)       # pad value -1.0
            padded[:, : p.shape[1], : p.shape[2]] = p
            patches.append(padded)
    patches = np.stack(patches, axis=0)             # [N, C, 40, 40]

    # rescale (1/255) + CLIP normalize per channel
    patches = patches / 255.0
    patches = (patches - CLIP_MEAN[None, :, None, None]) / CLIP_STD[None, :, None, None]

    # add temporal dim, duplicate x2, then -> [N, T, H, W, C]
    patches = np.repeat(patches[..., None], TEMPORAL, axis=-1)   # [N, C, 40, 40, 2]
    pixel_values = patches.transpose(0, 4, 2, 3, 1)              # [N, 2, 40, 40, C]
    return pixel_values.astype(np.float32), pixel_values.shape[0]


# ------------------------------- audio -------------------------------

_mel_fb = None
def _mel_filters() -> np.ndarray:
    global _mel_fb
    if _mel_fb is None:
        from transformers.audio_utils import mel_filter_bank
        fb = mel_filter_bank(num_frequency_bins=N_FFT // 2 + 1, num_mel_filters=N_MEL,
                             min_frequency=0.0, max_frequency=SR / 2.0, sampling_rate=SR,
                             norm="slaney", mel_scale="slaney")          # [801, 80]
        _mel_fb = np.ascontiguousarray(fb.T, dtype=np.float32)           # [80, 801]
    return _mel_fb


def _log_mel(waveform: np.ndarray) -> np.ndarray:
    """raw mono waveform -> log10-mel spectrogram [num_frames, 80]."""
    wav = np.asarray(waveform, dtype=np.float32).reshape(-1)
    right = math.ceil(wav.shape[0] / HOP) * HOP - wav.shape[0]
    left = max(N_FFT - HOP, 0)
    wav = np.pad(wav, (left, right))
    window = np.hanning(WIN + 1)[:-1].astype(np.float32)                 # periodic Hann
    n_frames = 1 + (wav.shape[0] - N_FFT) // HOP                          # center=False
    frames = np.stack([wav[i * HOP: i * HOP + N_FFT] * window for i in range(n_frames)])  # [T, N_FFT]
    mag = np.abs(np.fft.rfft(frames, n=N_FFT, axis=-1))                   # [T, 801]
    mag = np.maximum(mag, 1e-10)
    mel = _mel_filters() @ mag.T                                         # [80, T]
    mel = np.log10(np.maximum(mel, 1e-10))
    return mel.T                                                         # [T, 80]


def preprocess_audio(waveform: np.ndarray, sampling_rate: int = SR) -> np.ndarray:
    """raw 16 kHz mono waveform -> dMel bin ids [num_frames, 80] (int32, 0..15)."""
    if sampling_rate != SR:
        raise ValueError(f"Inkling audio expects {SR} Hz, got {sampling_rate}")
    mel = _log_mel(waveform)                                             # [T, 80] log10
    n_valid = math.ceil(len(np.asarray(waveform).reshape(-1)) / HOP)
    mel = mel[:n_valid]                                                  # drop trailing pad frames
    centers = np.linspace(DMEL_MIN, DMEL_MAX, DMEL_BINS)                 # 16 bin centers
    clamped = np.clip(mel.astype(np.float64), DMEL_MIN, DMEL_MAX)
    bins = np.abs(clamped[..., None] - centers).argmin(-1)              # nearest center
    return bins.astype(np.int32)                                        # [T, 80]


# --------------------------- prompt assembly ---------------------------

class InklingProcessor:
    """Assembles multimodal model inputs from chat messages with image/audio parts.

    Content parts: {"type":"text","text":...}, {"type":"image","image":PIL/array},
    {"type":"audio","audio":waveform, "sampling_rate":16000}.
    """

    def __init__(self, tokenizer, chat_template: str):
        self.tok = tokenizer
        self.chat_template = chat_template
        self.image_bos_id = tokenizer.encode(IMAGE_BOS, add_special_tokens=False)[0]
        self.audio_bos_id = tokenizer.encode(AUDIO_BOS, add_special_tokens=False)[0]

    def apply(self, messages, reasoning_effort: str = "none"):
        import mlx.core as mx
        pixel_values, audio_ids = [], []
        # Render text via the chat template with placeholders stripped to a sentinel,
        # then splice media spans in. We build ids directly for robustness.
        ids: list[int] = []

        def emit_text(s):
            ids.extend(self.tok.encode(s, add_special_tokens=False))

        # header: thinking-effort system message (matches chat_template)
        eff = {"none": 0.0, "minimal": 0.1, "low": 0.2, "medium": 0.7, "high": 0.9, "max": 0.99}[reasoning_effort]
        emit_text(f"<|message_system|><|content_text|>Thinking effort level: {0 if eff == 0 else eff}<|end_message|>")

        for msg in messages:
            role = {"user": "<|message_user|>", "assistant": "<|message_model|>",
                    "system": "<|message_system|>"}[msg["role"]]
            content = msg["content"]
            if isinstance(content, str):
                content = [{"type": "text", "text": content}]
            for part in content:
                t = part.get("type", "text")
                if t == "text":
                    emit_text(role + "<|content_text|>" + part["text"] + "<|end_message|>")
                elif t == "image":
                    pv, n = preprocess_image(part["image"])
                    pixel_values.append(pv)
                    ids.append(self.tok.encode(role, add_special_tokens=False)[0])
                    ids.append(self.image_bos_id)
                    ids.extend([IMAGE_TOKEN_ID] * n)
                    ids.extend(self.tok.encode("<|end_message|>", add_special_tokens=False))
                elif t == "audio":
                    aid = preprocess_audio(part["audio"], part.get("sampling_rate", SR))
                    audio_ids.append(aid)
                    ids.append(self.tok.encode(role, add_special_tokens=False)[0])
                    ids.append(self.audio_bos_id)
                    ids.extend([AUDIO_TOKEN_ID] * aid.shape[0])
                    ids.extend(self.tok.encode("<|end_message|>", add_special_tokens=False))
        emit_text("<|message_model|>")   # generation prompt

        out = {"input_ids": ids}
        if pixel_values:
            out["pixel_values"] = mx.array(np.concatenate(pixel_values, axis=0))
        if audio_ids:
            out["audio_input_ids"] = mx.array(np.concatenate(audio_ids, axis=0))
        return out