scripts/export_surya_split_coreml.py

#!/usr/bin/env python3
import argparse
import json
from pathlib import Path

import coremltools as ct
import numpy as np
import torch
from PIL import Image, ImageDraw
from qwen_vl_utils import process_vision_info
from transformers import AutoConfig, AutoProcessor, Qwen3_5ForConditionalGeneration
from transformers.models.qwen3_5.modeling_qwen3_5 import (
    get_vision_bilinear_indices_and_weights,
    get_vision_cu_seqlens,
    get_vision_position_ids,
)


PROMPT = (
    "OCR this image. Return only the exact text visible in the image, preserving "
    "Persian, numbers, line breaks, and punctuation. Do not explain."
)


def layer_type(layer: torch.nn.Module) -> str:
    """Bridge the Qwen 3.5 decoder API used by Transformers 5.3 and 5.13+."""
    value = getattr(layer, "layer_type", None)
    if value is None:
        value = getattr(layer, "block_type", None)
    if value not in {"linear_attention", "full_attention"}:
        raise ValueError(f"Unsupported Qwen 3.5 decoder layer type: {value!r}")
    return value


class VisionCoreMLWrapper(torch.nn.Module):
    def __init__(self, visual, image_grid_thw, pixel_values_shape):
        super().__init__()
        self.visual = visual
        self.register_buffer("image_grid_thw", image_grid_thw)
        bilinear_indices, bilinear_weights = get_vision_bilinear_indices_and_weights(
            image_grid_thw,
            num_grid_per_side=visual.num_grid_per_side,
            spatial_merge_size=visual.config.spatial_merge_size,
            kwargs={},
        )
        position_ids = get_vision_position_ids(image_grid_thw, visual.spatial_merge_size, kwargs={})
        cu_seqlens = get_vision_cu_seqlens(image_grid_thw, kwargs={})
        self.seq_len = int(position_ids.shape[0])
        self.patch_rows = int(pixel_values_shape[0])
        self.patch_embed_dim = int(visual.patch_embed.embed_dim)
        self.merger_hidden_size = int(visual.merger.hidden_size)
        self.merger_rows = int(self.seq_len // (visual.spatial_merge_size**2))
        self.register_buffer("bilinear_indices", bilinear_indices.to(torch.long))
        self.register_buffer("bilinear_weights", bilinear_weights.to(torch.float32))
        self.register_buffer("position_ids", position_ids.to(torch.long))
        self.register_buffer("cu_seqlens", cu_seqlens.to(torch.int32))

    def patch_embed_forward(self, pixel_values):
        patch_embed = self.visual.patch_embed
        hidden_states = pixel_values.reshape(
            self.patch_rows,
            patch_embed.in_channels,
            patch_embed.temporal_patch_size,
            patch_embed.patch_size,
            patch_embed.patch_size,
        )
        hidden_states = patch_embed.proj(hidden_states.to(dtype=patch_embed.proj.weight.dtype))
        return hidden_states.reshape(self.seq_len, self.patch_embed_dim)

    def attention_forward(self, attn, hidden_states, position_embeddings):
        query_states, key_states, value_states = (
            attn.qkv(hidden_states)
            .reshape(self.seq_len, 3, attn.num_heads, -1)
            .permute(1, 0, 2, 3)
            .unbind(0)
        )
        cos, sin = position_embeddings
        head_dim = attn.qkv.out_features // (3 * attn.num_heads)
        half_dim = head_dim // 2
        cos = cos.unsqueeze(-2).float()
        sin = sin.unsqueeze(-2).float()

        query_float = query_states.float()
        key_float = key_states.float()
        query_rot = torch.cat(
            (-query_float.narrow(-1, half_dim, half_dim), query_float.narrow(-1, 0, half_dim)),
            dim=-1,
        )
        key_rot = torch.cat(
            (-key_float.narrow(-1, half_dim, half_dim), key_float.narrow(-1, 0, half_dim)),
            dim=-1,
        )
        query_states = ((query_float * cos) + (query_rot * sin)).to(query_states.dtype)
        key_states = ((key_float * cos) + (key_rot * sin)).to(key_states.dtype)

        query_states = query_states.transpose(0, 1).unsqueeze(0)
        key_states = key_states.transpose(0, 1).unsqueeze(0)
        value_states = value_states.transpose(0, 1).unsqueeze(0)

        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) * attn.scaling
        attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
        attn_output = torch.matmul(attn_weights, value_states)
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.reshape(self.seq_len, -1).contiguous()
        return attn.proj(attn_output)

    def forward(self, pixel_values):
        hidden_states = self.patch_embed_forward(pixel_values)
        pos_embeds = (self.visual.pos_embed(self.bilinear_indices) * self.bilinear_weights[:, :, None]).sum(0)
        hidden_states = hidden_states + pos_embeds.to(hidden_states.dtype)
        rotary_pos_emb = self.visual.rotary_pos_emb(self.position_ids)

        hidden_states = hidden_states.reshape(self.seq_len, -1)
        rotary_pos_emb = rotary_pos_emb.reshape(self.seq_len, -1)
        emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1)
        position_embeddings = (emb.cos(), emb.sin())

        for block in self.visual.blocks:
            hidden_states = hidden_states + self.attention_forward(
                block.attn,
                block.norm1(hidden_states),
                position_embeddings=position_embeddings,
            )
            hidden_states = hidden_states + block.mlp(block.norm2(hidden_states))

        merger = self.visual.merger
        if merger.use_postshuffle_norm:
            hidden_states = hidden_states.reshape(self.merger_rows, self.merger_hidden_size)
        hidden_states = merger.norm(hidden_states).reshape(self.merger_rows, self.merger_hidden_size)
        return merger.linear_fc2(merger.act_fn(merger.linear_fc1(hidden_states)))


class FullLastLogitsCoreMLWrapper(torch.nn.Module):
    def __init__(self, model, image_grid_thw):
        super().__init__()
        self.model = model
        self.register_buffer("image_grid_thw", image_grid_thw)

    def forward(self, input_ids, attention_mask, pixel_values, mm_token_type_ids):
        out = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            pixel_values=pixel_values,
            image_grid_thw=self.image_grid_thw,
            mm_token_type_ids=mm_token_type_ids,
            use_cache=False,
            return_dict=True,
        )
        return out.logits[:, -1:, :]


class LanguageLastLogitsCoreMLWrapper(torch.nn.Module):
    def __init__(self, language_model, lm_head, position_embeddings, seq_len):
        super().__init__()
        self.language_model = language_model
        self.lm_head = lm_head
        cos, sin = position_embeddings
        self.register_buffer("cos", cos)
        self.register_buffer("sin", sin)
        mask = torch.full((1, 1, seq_len, seq_len), torch.finfo(torch.float32).min)
        mask = torch.triu(mask, diagonal=1)
        self.register_buffer("causal_mask", mask)

    @staticmethod
    def rotate_half_static(x, rotary_dim):
        half_dim = rotary_dim // 2
        x1 = x.narrow(-1, 0, half_dim)
        x2 = x.narrow(-1, half_dim, half_dim)
        return torch.cat((-x2, x1), dim=-1)

    def apply_rotary_static(self, q, k):
        cos = self.cos.unsqueeze(1)
        sin = self.sin.unsqueeze(1)
        rotary_dim = int(self.cos.shape[-1])
        q_rot = q.narrow(-1, 0, rotary_dim)
        q_pass = q.narrow(-1, rotary_dim, int(q.shape[-1]) - rotary_dim)
        k_rot = k.narrow(-1, 0, rotary_dim)
        k_pass = k.narrow(-1, rotary_dim, int(k.shape[-1]) - rotary_dim)
        q_embed = (q_rot * cos) + (self.rotate_half_static(q_rot, rotary_dim) * sin)
        k_embed = (k_rot * cos) + (self.rotate_half_static(k_rot, rotary_dim) * sin)
        return torch.cat((q_embed, q_pass), dim=-1), torch.cat((k_embed, k_pass), dim=-1)

    @staticmethod
    def repeat_kv_static(hidden_states, n_rep):
        if n_rep == 1:
            return hidden_states
        batch = int(hidden_states.shape[0])
        num_key_value_heads = int(hidden_states.shape[1])
        seq_len = int(hidden_states.shape[2])
        head_dim = int(hidden_states.shape[3])
        hidden_states = hidden_states[:, :, None, :, :].expand(
            batch,
            num_key_value_heads,
            n_rep,
            seq_len,
            head_dim,
        )
        return hidden_states.reshape(batch, num_key_value_heads * n_rep, seq_len, head_dim)

    def full_attention_forward(self, attn, hidden_states):
        query_states, gate = torch.chunk(
            attn.q_proj(hidden_states).view(1, 300, -1, attn.head_dim * 2),
            2,
            dim=-1,
        )
        gate = gate.reshape(1, 300, -1)

        query_states = attn.q_norm(query_states.view(1, 300, -1, attn.head_dim)).transpose(1, 2)
        key_states = attn.k_norm(attn.k_proj(hidden_states).view(1, 300, -1, attn.head_dim)).transpose(1, 2)
        value_states = attn.v_proj(hidden_states).view(1, 300, -1, attn.head_dim).transpose(1, 2)

        query_states, key_states = self.apply_rotary_static(query_states, key_states)
        key_states = self.repeat_kv_static(key_states, attn.num_key_value_groups)
        value_states = self.repeat_kv_static(value_states, attn.num_key_value_groups)

        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) * attn.scaling
        attn_weights = attn_weights + self.causal_mask.to(attn_weights.dtype)
        attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
        attn_output = torch.matmul(attn_weights, value_states)
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.reshape(1, 300, -1).contiguous()
        attn_output = attn_output * torch.sigmoid(gate)
        return attn.o_proj(attn_output)

    @staticmethod
    def l2norm_static(x):
        return x * torch.rsqrt((x * x).sum(dim=-1, keepdim=True) + 1e-6)

    def chunk_gated_delta_rule_static(self, query, key, value, g, beta):
        chunk_size = 64
        sequence_length = 300
        total_sequence_length = 320
        pad_size = 20

        initial_dtype = query.dtype
        query = self.l2norm_static(query)
        key = self.l2norm_static(key)
        query = query.transpose(1, 2).contiguous().float()
        key = key.transpose(1, 2).contiguous().float()
        value = value.transpose(1, 2).contiguous().float()
        beta = beta.transpose(1, 2).contiguous().float()
        g = g.transpose(1, 2).contiguous().float()

        query = torch.nn.functional.pad(query, (0, 0, 0, pad_size))
        key = torch.nn.functional.pad(key, (0, 0, 0, pad_size))
        value = torch.nn.functional.pad(value, (0, 0, 0, pad_size))
        beta = torch.nn.functional.pad(beta, (0, pad_size))
        g = torch.nn.functional.pad(g, (0, pad_size))

        scale = 1 / (128**0.5)
        query = query * scale

        v_beta = value * beta.unsqueeze(-1)
        k_beta = key * beta.unsqueeze(-1)
        query = query.reshape(1, 16, 5, chunk_size, 128)
        key = key.reshape(1, 16, 5, chunk_size, 128)
        value = value.reshape(1, 16, 5, chunk_size, 128)
        k_beta = k_beta.reshape(1, 16, 5, chunk_size, 128)
        v_beta = v_beta.reshape(1, 16, 5, chunk_size, 128)
        g = g.reshape(1, 16, 5, chunk_size)

        tri0 = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=0)
        tri1 = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=1)
        g_cum = g.cumsum(dim=-1)
        decay_mask = ((g_cum.unsqueeze(-1) - g_cum.unsqueeze(-2)).tril().exp().float()).tril()
        attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(tri0, 0)
        # Fixed-size triangular recurrence from torch_chunk_gated_delta_rule.
        for i in range(1, chunk_size):
            row = attn[..., i, :i].clone()
            sub = attn[..., :i, :i].clone()
            update = row + (row.unsqueeze(-1) * sub).sum(-2)
            suffix = attn[..., i, i:]
            new_row = torch.cat((update, suffix), dim=-1)
            before = attn[..., :i, :]
            after = attn[..., i + 1 :, :]
            attn = torch.cat((before, new_row.unsqueeze(-2), after), dim=-2)
        attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=query.device)
        value = attn @ v_beta
        k_cumdecay = attn @ (k_beta * g_cum.exp().unsqueeze(-1))

        last_recurrent_state = torch.zeros(1, 16, 128, 128, dtype=value.dtype, device=value.device)
        outs = []
        for i in range(5):
            q_i = query[:, :, i]
            k_i = key[:, :, i]
            v_i = value[:, :, i]
            attn_i = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill(tri1, 0)
            v_prime = k_cumdecay[:, :, i] @ last_recurrent_state
            v_new = v_i - v_prime
            attn_inter = (q_i * g_cum[:, :, i, :, None].exp()) @ last_recurrent_state
            outs.append(attn_inter + attn_i @ v_new)
            last_recurrent_state = (
                last_recurrent_state * g_cum[:, :, i, -1, None, None].exp()
                + (k_i * (g_cum[:, :, i, -1, None] - g_cum[:, :, i]).exp()[..., None]).transpose(-1, -2)
                @ v_new
            )

        core_attn_out = torch.stack(outs, dim=2)
        core_attn_out = core_attn_out.reshape(1, 16, total_sequence_length, 128)
        core_attn_out = core_attn_out[:, :, :sequence_length]
        return core_attn_out.transpose(1, 2).contiguous().to(initial_dtype)

    def gated_delta_forward_static(self, linear_attn, hidden_states):
        batch_size = 1
        seq_len = 300
        mixed_qkv = linear_attn.in_proj_qkv(hidden_states).transpose(1, 2)
        z = linear_attn.in_proj_z(hidden_states).reshape(batch_size, seq_len, -1, linear_attn.head_v_dim)
        b = linear_attn.in_proj_b(hidden_states)
        a = linear_attn.in_proj_a(hidden_states)

        mixed_qkv = torch.nn.functional.silu(linear_attn.conv1d(mixed_qkv)[:, :, :seq_len])
        mixed_qkv = mixed_qkv.transpose(1, 2)
        query, key, value = torch.split(
            mixed_qkv,
            [
                linear_attn.key_dim,
                linear_attn.key_dim,
                linear_attn.value_dim,
            ],
            dim=-1,
        )
        query = query.reshape(batch_size, seq_len, -1, linear_attn.head_k_dim)
        key = key.reshape(batch_size, seq_len, -1, linear_attn.head_k_dim)
        value = value.reshape(batch_size, seq_len, -1, linear_attn.head_v_dim)
        beta = b.sigmoid()
        g = -linear_attn.A_log.float().exp() * torch.nn.functional.softplus(a.float() + linear_attn.dt_bias)

        core_attn_out = self.chunk_gated_delta_rule_static(query, key, value, g, beta)
        core_attn_out = core_attn_out.reshape(-1, linear_attn.head_v_dim)
        z = z.reshape(-1, linear_attn.head_v_dim)
        normed = linear_attn.norm(core_attn_out, z)
        normed = normed.reshape(batch_size, seq_len, -1)
        return linear_attn.out_proj(normed)

    def forward(self, inputs_embeds):
        hidden_states = inputs_embeds
        for layer in self.language_model.layers:
            residual = hidden_states
            hidden_states = layer.input_layernorm(hidden_states)
            if layer_type(layer) == "linear_attention":
                hidden_states = self.gated_delta_forward_static(layer.linear_attn, hidden_states)
            else:
                hidden_states = self.full_attention_forward(layer.self_attn, hidden_states)
            hidden_states = residual + hidden_states

            residual = hidden_states
            hidden_states = layer.post_attention_layernorm(hidden_states)
            hidden_states = layer.mlp(hidden_states)
            hidden_states = residual + hidden_states

        hidden_states = self.language_model.norm(hidden_states)
        hidden = hidden_states[:, -1:, :]
        return self.lm_head(hidden)


def load_model(model_id: str, dtype: torch.dtype):
    config = AutoConfig.from_pretrained(model_id, trust_remote_code=True)
    config._attn_implementation = "eager"
    if hasattr(config, "text_config"):
        config.text_config._attn_implementation = "eager"
    if hasattr(config, "vision_config"):
        config.vision_config._attn_implementation = "eager"
    model = Qwen3_5ForConditionalGeneration.from_pretrained(
        model_id,
        config=config,
        torch_dtype=dtype,
        device_map="cpu",
        low_cpu_mem_usage=True,
        trust_remote_code=True,
        attn_implementation="eager",
    ).eval()
    model.config._attn_implementation = "eager"
    if hasattr(model.config, "text_config"):
        model.config.text_config._attn_implementation = "eager"
    return model


def build_sample(processor):
    image = Image.new("RGB", (512, 512), "white")
    draw = ImageDraw.Draw(image)
    draw.text((40, 80), "Invoice 123\nTotal $42.00", fill="black")
    messages = [{"role": "user", "content": [{"type": "image", "image": image}, {"type": "text", "text": PROMPT}]}]
    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    image_inputs, video_inputs = process_vision_info(messages)
    return processor(text=[text], images=image_inputs, videos=video_inputs, return_tensors="pt")


def main() -> None:
    parser = argparse.ArgumentParser(description="Export Bina 0.1 split/full CoreML canaries from source BF16 weights.")
    parser.add_argument("--model-id", default="Reza2kn/Bina-0.1-Koochik")
    parser.add_argument("--output-dir", required=True)
    parser.add_argument("--mode", choices=["vision", "full", "language"], default="vision")
    parser.add_argument("--dtype", choices=["float16", "float32"], default="float16")
    parser.add_argument(
        "--compute-precision",
        choices=["float16", "float32"],
        default="float16",
        help="CoreML ML Program compute precision; the source checkpoint remains BF16.",
    )
    args = parser.parse_args()

    output_dir = Path(args.output_dir).expanduser().resolve()
    output_dir.mkdir(parents=True, exist_ok=True)
    dtype = torch.float16 if args.dtype == "float16" else torch.float32
    np_dtype = np.float16 if args.dtype == "float16" else np.float32
    compute_precision = (
        ct.precision.FLOAT16 if args.compute_precision == "float16" else ct.precision.FLOAT32
    )
    precision_tag = "fp16" if args.compute_precision == "float16" else "fp32"

    processor = AutoProcessor.from_pretrained(args.model_id, trust_remote_code=True)
    model = load_model(args.model_id, dtype)
    sample = build_sample(processor)
    input_ids = sample["input_ids"].to(torch.int64)
    attention_mask = sample["attention_mask"].to(torch.int64)
    mm_token_type_ids = sample["mm_token_type_ids"].to(torch.int64)
    pixel_values = sample["pixel_values"].to(dtype)
    image_grid_thw = sample["image_grid_thw"].to(torch.int64)

    if args.mode == "vision":
        wrapper = VisionCoreMLWrapper(model.model.visual, image_grid_thw, tuple(pixel_values.shape)).eval()
        example = (pixel_values,)
        traced = torch.jit.trace(wrapper, example, strict=False)
        package_path = output_dir / f"surya_vision_{precision_tag}.mlpackage"
        mlmodel = ct.convert(
            traced,
            convert_to="mlprogram",
            minimum_deployment_target=ct.target.macOS14,
            compute_precision=compute_precision,
            inputs=[ct.TensorType(name="pixel_values", shape=tuple(pixel_values.shape), dtype=np_dtype)],
            outputs=[ct.TensorType(name="image_embeds")],
        )
    elif args.mode == "full":
        wrapper = FullLastLogitsCoreMLWrapper(model, image_grid_thw).eval()
        example = (input_ids, attention_mask, pixel_values, mm_token_type_ids)
        traced = torch.jit.trace(wrapper, example, strict=False)
        package_path = output_dir / f"surya_full_last_logits_{precision_tag}.mlpackage"
        mlmodel = ct.convert(
            traced,
            convert_to="mlprogram",
            minimum_deployment_target=ct.target.macOS14,
            compute_precision=compute_precision,
            inputs=[
                ct.TensorType(name="input_ids", shape=tuple(input_ids.shape), dtype=np.int32),
                ct.TensorType(name="attention_mask", shape=tuple(attention_mask.shape), dtype=np.int32),
                ct.TensorType(name="pixel_values", shape=tuple(pixel_values.shape), dtype=np_dtype),
                ct.TensorType(name="mm_token_type_ids", shape=tuple(mm_token_type_ids.shape), dtype=np.int32),
            ],
            outputs=[ct.TensorType(name="logits")],
        )
    else:
        with torch.no_grad():
            inputs_embeds = model.model.get_input_embeddings()(input_ids)
            image_outputs = model.model.get_image_features(pixel_values, image_grid_thw, return_dict=True)
            image_embeds = torch.cat(image_outputs.pooler_output, dim=0).to(inputs_embeds.device, inputs_embeds.dtype)
            image_mask, _ = model.model.get_placeholder_mask(
                input_ids,
                inputs_embeds=inputs_embeds,
                image_features=image_embeds,
            )
            inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
            position_ids = model.model.compute_3d_position_ids(
                input_ids=input_ids,
                image_grid_thw=image_grid_thw,
                video_grid_thw=None,
                inputs_embeds=inputs_embeds,
                attention_mask=attention_mask,
                past_key_values=None,
                mm_token_type_ids=mm_token_type_ids,
            )

        position_embeddings = model.model.language_model.rotary_emb(inputs_embeds, position_ids)
        wrapper = LanguageLastLogitsCoreMLWrapper(
            model.model.language_model,
            model.lm_head,
            position_embeddings,
            int(inputs_embeds.shape[1]),
        ).eval()
        example = (inputs_embeds,)
        traced = torch.jit.trace(wrapper, example, strict=False)
        package_path = output_dir / f"surya_language_last_logits_{precision_tag}.mlpackage"
        mlmodel = ct.convert(
            traced,
            convert_to="mlprogram",
            minimum_deployment_target=ct.target.macOS14,
            compute_precision=compute_precision,
            inputs=[
                ct.TensorType(name="inputs_embeds", shape=tuple(inputs_embeds.shape), dtype=np_dtype),
            ],
            outputs=[ct.TensorType(name="logits")],
        )

    mlmodel.save(str(package_path))
    processor.save_pretrained(output_dir / "processor")
    (output_dir / "export_config.json").write_text(
        json.dumps(
            {
                "model_id": args.model_id,
                "mode": args.mode,
                "source_dtype": "bf16",
                "coreml_compute_precision": precision_tag,
                "sample_shapes": {
                    "input_ids": list(input_ids.shape),
                    "attention_mask": list(attention_mask.shape),
                    "mm_token_type_ids": list(mm_token_type_ids.shape),
                    "pixel_values": list(pixel_values.shape),
                    "image_grid_thw": list(image_grid_thw.shape),
                    **(
                        {
                            "inputs_embeds": list(inputs_embeds.shape),
                            "position_ids": list(position_ids.shape),
                        }
                        if args.mode == "language"
                        else {}
                    ),
                },
                "package": str(package_path),
            },
            indent=2,
        )
        + "\n",
        encoding="utf-8",
    )
    print(json.dumps({"package": str(package_path), "mode": args.mode}, indent=2), flush=True)


if __name__ == "__main__":
    main()