#!/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()