lite/python/vision_export.py

"""Trace-friendly Qwen3.5 vision tower for Core ML (from clef-vision-0.8b; d3 weight names): one image per call, patches right-padded to a fixed count.

Inputs (all fp32, host-prepared from the HF processor's ``pixel_values`` / ``image_grid_thw``):
  patches     [N, 1536]  flattened (C, T, P, P) patches, zero padded after the image's n patches
  pos_embeds  [N, 768]   bilinear-resampled learned position grid (host: HF interpolation indices/weights)
  cos, sin    [N, 64]    axial 2D rope tables (host)
  mask        [1, 1, N, N] additive key mask: 0 for real patches, -1e4 for padding
Output: merged tokens [N/4, 1024]; rows [0, n/4) are the image's tokens in the order the language model expects.

Padding is appended after the image, so with the key mask real patches never attend to it; the merger groups four
consecutive patches, and n is always a multiple of 4 (grid h, w are multiples of the merge size), so no group mixes
real and padded patches. Parity vs ``Qwen3_5VisionModel`` is checked in ``vision_parity``.
"""
from __future__ import annotations

import argparse
import json
import math
import time
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F
from torch import nn

VISION_PREFIX = "visual."


class VisionConfig:
    def __init__(self, cfg: dict):
        self.depth = cfg["depth"]
        self.hidden = cfg["hidden_size"]
        self.heads = cfg["num_heads"]
        self.head_dim = self.hidden // self.heads
        self.intermediate = cfg["intermediate_size"]
        self.out_hidden = cfg["out_hidden_size"]
        self.patch = cfg["patch_size"]
        self.temporal = cfg["temporal_patch_size"]
        self.channels = cfg["in_channels"]
        self.merge = cfg["spatial_merge_size"]
        self.num_pos = cfg["num_position_embeddings"]
        self.act = cfg["hidden_act"]  # gelu_pytorch_tanh for the blocks; the merger uses exact GELU
        self.rope_theta = cfg.get("rope_parameters", {}).get("rope_theta", 10000.0)
        self.patch_dim = self.channels * self.temporal * self.patch * self.patch


def rotate_half(x: torch.Tensor) -> torch.Tensor:
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat([-x2, x1], dim=-1)


class Block(nn.Module):
    def __init__(self, cfg: VisionConfig):
        super().__init__()
        self.cfg = cfg
        self.norm1 = nn.LayerNorm(cfg.hidden, eps=1e-6)
        self.norm2 = nn.LayerNorm(cfg.hidden, eps=1e-6)
        self.qkv = nn.Linear(cfg.hidden, cfg.hidden * 3)
        self.proj = nn.Linear(cfg.hidden, cfg.hidden)
        self.fc1 = nn.Linear(cfg.hidden, cfg.intermediate)
        self.fc2 = nn.Linear(cfg.intermediate, cfg.hidden)
        self.scale = cfg.head_dim**-0.5

    def forward(self, x, cos, sin, mask):
        n = x.shape[0]
        h = self.norm1(x)
        qkv = self.qkv(h).reshape(n, 3, self.cfg.heads, self.cfg.head_dim)
        q, k, v = qkv[:, 0], qkv[:, 1], qkv[:, 2]  # [N, H, d]
        c, s = cos[:, None, :], sin[:, None, :]
        q = q * c + rotate_half(q) * s
        k = k * c + rotate_half(k) * s
        q = q.permute(1, 0, 2)[None]  # [1, H, N, d]
        k = k.permute(1, 0, 2)[None]
        v = v.permute(1, 0, 2)[None]
        scores = torch.matmul(q, k.transpose(-1, -2)) * self.scale + mask
        attn = torch.matmul(torch.softmax(scores, dim=-1), v)  # [1, H, N, d]
        attn = attn[0].permute(1, 0, 2).reshape(n, self.cfg.hidden)
        x = x + self.proj(attn)
        h = self.norm2(x)
        act = F.gelu(self.fc1(h), approximate="tanh" if self.cfg.act == "gelu_pytorch_tanh" else "none")
        return x + self.fc2(act)


class VisionTower(nn.Module):
    def __init__(self, cfg: VisionConfig, n_patches: int):
        super().__init__()
        assert n_patches % (cfg.merge**2) == 0
        self.cfg = cfg
        self.n_patches = n_patches
        self.patch_embed = nn.Linear(cfg.patch_dim, cfg.hidden)  # Conv3d with kernel == stride == patch
        self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.depth)])
        merged = cfg.hidden * cfg.merge**2
        self.merger_norm = nn.LayerNorm(cfg.hidden, eps=1e-6)
        self.merger_fc1 = nn.Linear(merged, merged)
        self.merger_fc2 = nn.Linear(merged, cfg.out_hidden)

    def forward(self, patches, pos_embeds, cos, sin, mask):
        x = self.patch_embed(patches) + pos_embeds
        for block in self.blocks:
            x = block(x, cos, sin, mask)
        x = self.merger_norm(x).reshape(self.n_patches // self.cfg.merge**2, -1)
        return self.merger_fc2(F.gelu(self.merger_fc1(x)))

    def load_hf(self, state: dict[str, torch.Tensor]) -> None:
        """``state`` holds the HF vision weights with the ``visual.`` prefix stripped."""
        own = {
            "patch_embed.weight": state["patch_embed.proj.weight"].reshape(self.cfg.hidden, -1),
            "patch_embed.bias": state["patch_embed.proj.bias"],
            "merger_norm.weight": state["merger.norm.weight"],
            "merger_norm.bias": state["merger.norm.bias"],
            "merger_fc1.weight": state["merger.linear_fc1.weight"],
            "merger_fc1.bias": state["merger.linear_fc1.bias"],
            "merger_fc2.weight": state["merger.linear_fc2.weight"],
            "merger_fc2.bias": state["merger.linear_fc2.bias"],
        }
        for i in range(self.cfg.depth):
            src, dst = f"blocks.{i}.", f"blocks.{i}."
            for a, b in (("norm1", "norm1"), ("norm2", "norm2"), ("attn.qkv", "qkv"), ("attn.proj", "proj"),
                         ("mlp.linear_fc1", "fc1"), ("mlp.linear_fc2", "fc2")):
                own[f"{dst}{b}.weight"] = state[f"{src}{a}.weight"]
                own[f"{dst}{b}.bias"] = state[f"{src}{a}.bias"]
        missing, unexpected = self.load_state_dict({k: v.float() for k, v in own.items()}, strict=False)
        if missing or unexpected:
            raise RuntimeError(f"vision weight mismatch: missing={missing[:6]} unexpected={unexpected[:6]}")


# ----------------------------------------------------------------------------------------------------------------------
# Host-side preparation (numpy/torch; mirrors transformers.vision_utils + Qwen3_5VisionRotaryEmbedding)

def axial_rope(cfg: VisionConfig, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """``position_ids`` [n, 2] (h, w) -> cos/sin [n, head_dim]: freqs over head_dim//4 per axis, concat (h, w), x2."""
    spatial = cfg.head_dim // 2
    inv_freq = 1.0 / (cfg.rope_theta ** (torch.arange(0, spatial, 2, dtype=torch.float32) / spatial))
    freqs = position_ids[..., None].float() * inv_freq  # [n, 2, spatial/2]
    hw = torch.cat([freqs[:, 0], freqs[:, 1]], dim=-1)
    emb = torch.cat([hw, hw], dim=-1)
    return emb.cos(), emb.sin()


def prepare_inputs(cfg: VisionConfig, pos_table: torch.Tensor, pixel_values: torch.Tensor, grid_thw: torch.Tensor,
                   n_patches: int):
    """One image. Uses the HF helpers for position ids and pos-embed interpolation so the host matches the reference."""
    from transformers.vision_utils import get_vision_interpolation_indices_and_weights, get_vision_position_ids

    n = int(pixel_values.shape[0])
    if n > n_patches:
        raise ValueError(f"image has {n} patches; bucket holds {n_patches}")
    grid = grid_thw.reshape(1, 3)
    idx, w = get_vision_interpolation_indices_and_weights(
        grid, num_grid_per_side=int(math.isqrt(cfg.num_pos)), mode="bilinear", align_corners=True,
        spatial_merge_size=cfg.merge, kwargs={})
    pos_embeds = (pos_table[idx] * w[:, :, None]).sum(1)  # [n, hidden]
    position_ids = get_vision_position_ids(grid, cfg.merge, kwargs={})  # [n, 2]
    cos, sin = axial_rope(cfg, position_ids)
    patches = torch.zeros(n_patches, cfg.patch_dim)
    patches[:n] = pixel_values.float()
    pe = torch.zeros(n_patches, cfg.hidden)
    pe[:n] = pos_embeds.float()
    cs = torch.zeros(n_patches, cfg.head_dim)
    sn = torch.zeros(n_patches, cfg.head_dim)
    cs[:n], sn[:n] = cos, sin
    mask = torch.zeros(1, 1, n_patches, n_patches)
    mask[..., n:] = -1e4
    return patches, pe, cs, sn, mask, n


def load_vision(merged_backbone: Path, n_patches: int):
    from safetensors.torch import load_file

    config = json.loads((merged_backbone / "config.json").read_text())
    cfg = VisionConfig(config["vision_config"])
    state = {}
    for path in sorted(merged_backbone.glob("model*.safetensors")):  # one file or shards
        state.update({k[len(VISION_PREFIX):]: v for k, v in load_file(str(path)).items() if k.startswith(VISION_PREFIX)})
    tower = VisionTower(cfg, n_patches)
    tower.load_hf(state)
    return tower.eval(), cfg, state["pos_embed.weight"].float()


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--merged", type=Path, required=True, help="d3 checkpoint dir")
    ap.add_argument("--patches", type=int, default=784, help="patch bucket (784 = 448x448 cap -> 196 tokens)")
    ap.add_argument("--precision", choices=("fp16", "fp32", "mixed"), default="fp32",
                    help="fp32 is the only exact option: fp16 (0.27 max err) and mixed with fp32 norms/softmax "
                    "(0.33) both break on this ViT's activation range; fp32 runs 35 ms on the M5 Pro GPU")
    ap.add_argument("--out", type=Path, default=Path("build/coreml"))
    args = ap.parse_args()
    import coremltools as ct

    tower, cfg, pos_table = load_vision(args.merged, args.patches)
    N = args.patches
    example = (torch.zeros(N, cfg.patch_dim), torch.zeros(N, cfg.hidden), torch.zeros(N, cfg.head_dim),
               torch.zeros(N, cfg.head_dim), torch.zeros(1, 1, N, N))
    started = time.time()
    with torch.no_grad():
        traced = torch.jit.trace(tower, example, check_trace=False)
    names = ["patches", "pos_embeds", "cos", "sin", "mask"]
    if args.precision == "fp32":
        precision = ct.precision.FLOAT32
    elif args.precision == "fp16":
        precision = ct.precision.FLOAT16
    else:
        keep_fp32 = {"layer_norm", "softmax", "reduce_mean", "reduce_sum", "pow", "sqrt", "rsqrt", "real_div"}
        precision = ct.transform.FP16ComputePrecision(op_selector=lambda op: op.op_type not in keep_fp32)
    model = ct.convert(
        traced, convert_to="mlprogram", minimum_deployment_target=ct.target.iOS17,
        compute_precision=precision,
        compute_units=ct.ComputeUnit.CPU_ONLY,
        inputs=[ct.TensorType(name=n, shape=tuple(t.shape), dtype=np.float32) for n, t in zip(names, example)],
        outputs=[ct.TensorType(name="tokens", dtype=np.float32)],
    )
    model.short_description = "d3 vision tower (vLLM-SR Decision 3.0): one image, fixed patch bucket"
    model.license = "Apache-2.0"
    model.user_defined_metadata.update({"patches": str(N), "tokens": str(N // cfg.merge**2)})
    out = args.out / f"Vision_P{N}"
    out.mkdir(parents=True, exist_ok=True)
    package = out / f"VisionTower_{args.precision}.mlpackage"
    model.save(str(package))
    pos_table.numpy().astype(np.float32).tofile(out / "pos_embed.f32")  # the host resamples this per image
    (out / "config.json").write_text(json.dumps({"patches": N, "tokens": N // cfg.merge**2, "patch_dim": cfg.patch_dim,
                                                  "hidden": cfg.hidden, "head_dim": cfg.head_dim,
                                                  "out_hidden": cfg.out_hidden, "precision": args.precision},
                                                 indent=2) + "\n")
    print(f"saved {package} in {time.time() - started:.0f} s")


if __name__ == "__main__":
    main()