Add processing.py to bundled loader

inkling_mlx/processing.py

@@ -0,0 +1,177 @@
+"""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