| |
| |
| |
| |
| |
| |
| |
| """Build the Z-Image-Turbo text encoder ONNX graph directly from a Qwen3 HF checkpoint.""" |
|
|
| import argparse |
| import json |
| import os |
|
|
| import numpy as np |
| import onnx |
| import onnx_ir as ir |
| import torch |
| from onnx import TensorProto, helper, numpy_helper |
| from onnxruntime.quantization.matmul_nbits_quantizer import MatMulNBitsQuantizer, QuantFormat |
| from transformers import AutoModelForCausalLM |
|
|
| from external_data_utils import ( |
| INLINE_SIZE_THRESHOLD_BYTES, |
| MAX_SHARD_SIZE_BYTES, |
| save_ir_model_sharded, |
| ) |
|
|
| DEFAULT_OUTPUT_DIR = "z-image-turbo-onnx" |
|
|
| PRECISION_CONFIGS = { |
| "f16": {"dtype": "f16", "quantize": False, "suffix": "f16"}, |
| "f32": {"dtype": "f32", "quantize": False, "suffix": "f32"}, |
| "f16_int4_quant": {"dtype": "f16", "quantize": True, "suffix": "q4f16"}, |
| "f32_int4_quant": {"dtype": "f32", "quantize": True, "suffix": "q4f32"}, |
| } |
|
|
|
|
| def parse_extra_options(pairs): |
| """Parse `key=value` strings (as passed via `--extra_options`) into a dict.""" |
| extra_options = {} |
| for kv_str in pairs or []: |
| key, _, value = kv_str.partition("=") |
| extra_options[key.strip()] = value.strip() |
| return extra_options |
|
|
|
|
| def _read_qwen3_config(text_encoder_dir): |
| with open(os.path.join(text_encoder_dir, "config.json"), encoding="utf-8") as f: |
| cfg = json.load(f) |
| return { |
| "num_layers": cfg["num_hidden_layers"], |
| "hidden_size": cfg["hidden_size"], |
| "num_attn_heads": cfg["num_attention_heads"], |
| "num_kv_heads": cfg["num_key_value_heads"], |
| "head_size": cfg["head_dim"], |
| "rope_theta": float(cfg["rope_theta"]), |
| "rms_norm_eps": float(cfg["rms_norm_eps"]), |
| "max_position_embeddings": cfg["max_position_embeddings"], |
| } |
|
|
|
|
| _TORCH_DTYPES = {"f16": torch.float16, "f32": torch.float32} |
| _ONNX_DTYPES = {"f16": TensorProto.FLOAT16, "f32": TensorProto.FLOAT} |
| _NP_DTYPES = {"f16": np.float16, "f32": np.float32} |
|
|
|
|
| def _rotary_cos_sin_tables(head_size, rope_theta, max_position_embeddings, np_dtype): |
| dim = head_size |
| inv_freq = 1.0 / (rope_theta ** (np.arange(0, dim, 2, dtype=np.int64).astype(np.float64) / dim)) |
| t = np.arange(max_position_embeddings, dtype=np.float64) |
| freqs = np.outer(t, inv_freq) |
| emb = np.concatenate([freqs, freqs], axis=-1) |
| |
| cos_cache = np.cos(emb)[:, : dim // 2].astype(np_dtype) |
| sin_cache = np.sin(emb)[:, : dim // 2].astype(np_dtype) |
| return cos_cache, sin_cache |
|
|
|
|
| class _GraphBuilder: |
| """Accumulates onnx NodeProtos and initializers for the encoder graph.""" |
|
|
| def __init__(self): |
| self.nodes = [] |
| self.initializers = [] |
|
|
| def initializer(self, name, array): |
| array = np.asarray(array) |
| |
| if array.ndim: |
| array = np.ascontiguousarray(array) |
| self.initializers.append(numpy_helper.from_array(array, name=name)) |
| return name |
|
|
| def const_i64(self, name, values): |
| return self.initializer(name, np.array(values, dtype=np.int64)) |
|
|
| def node(self, op_type, inputs, outputs, name, domain="", **attrs): |
| self.nodes.append(helper.make_node(op_type, inputs, outputs, name=name, domain=domain, **attrs)) |
| return outputs[0] |
|
|
|
|
| def _attention_mask_reformat(gb, attention_mask_name, enable_webgpu_graph=False): |
| """Derive GroupQueryAttention's seqlens_k/total_seq_len from the 2D attention_mask.""" |
| if enable_webgpu_graph: |
| mask_i32 = gb.node("Cast", [attention_mask_name], ["attn_mask_reformat/mask_i32"], |
| "attn_mask_reformat/Cast", to=TensorProto.INT32) |
| axis1 = gb.const_i64("attn_mask_reformat/axis1", [1]) |
| mask_sum = gb.node("ReduceSum", [mask_i32, axis1], ["attn_mask_reformat/mask_sum"], |
| "attn_mask_reformat/ReduceSum", keepdims=0) |
| one_i32 = gb.initializer("attn_mask_reformat/one_i32", np.array([1], dtype=np.int32)) |
| seqlens_k = gb.node("Sub", [mask_sum, one_i32], ["attn_mask_reformat/seqlens_k"], |
| "attn_mask_reformat/Sub") |
| total_seq_len = gb.node("ReduceMax", [mask_sum], ["attn_mask_reformat/total_seq_len"], |
| "attn_mask_reformat/ReduceMax", keepdims=0) |
| return seqlens_k, total_seq_len |
|
|
| axis1 = gb.const_i64("attn_mask_reformat/axis1", [1]) |
| mask_sum = gb.node("ReduceSum", [attention_mask_name, axis1], ["attn_mask_reformat/mask_sum"], |
| "attn_mask_reformat/ReduceSum", keepdims=0) |
| one_i64 = gb.const_i64("attn_mask_reformat/one_i64", [1]) |
| sub = gb.node("Sub", [mask_sum, one_i64], ["attn_mask_reformat/sub"], "attn_mask_reformat/Sub") |
| seqlens_k = gb.node("Cast", [sub], ["attn_mask_reformat/seqlens_k"], |
| "attn_mask_reformat/Sub/Cast", to=TensorProto.INT32) |
|
|
| shape = gb.node("Shape", [attention_mask_name], ["attn_mask_reformat/shape"], "attn_mask_reformat/Shape") |
| idx1 = gb.const_i64("attn_mask_reformat/idx1", 1) |
| gathered = gb.node("Gather", [shape, idx1], ["attn_mask_reformat/gathered"], |
| "attn_mask_reformat/Gather", axis=0) |
| total_seq_len = gb.node("Cast", [gathered], ["attn_mask_reformat/total_seq_len"], |
| "attn_mask_reformat/Gather/Cast", to=TensorProto.INT32) |
| return seqlens_k, total_seq_len |
|
|
|
|
| def _simplified_layernorm(gb, x, weight_name, eps, name_prefix, skip=None, need_sum=False, sum_output_name=None): |
| """skip=None: plain norm. skip=<name>: SkipSimplifiedLayerNormalization.""" |
| inputs = [x, weight_name] if skip is None else [x, skip, weight_name] |
| op_type = ("Skip" if skip is not None else "") + "SimplifiedLayerNormalization" |
| domain = "com.microsoft" if skip is not None else "" |
| y = f"{name_prefix}/Y" |
| sum_name = sum_output_name or f"{name_prefix}/sum" |
| outputs = [y] if skip is None else [y, "", "", (sum_name if need_sum else "")] |
| attrs = {"epsilon": eps} |
| if skip is None: |
| attrs.update(axis=-1, stash_type=1) |
| gb.node(op_type, inputs, outputs, name_prefix, domain=domain, **attrs) |
| return y, (outputs[3] if skip is not None else None) |
|
|
|
|
| def _qk_head_norm(gb, x, weight_name, num_heads, head_size, eps, name_prefix): |
| |
| shape1 = gb.const_i64(f"{name_prefix}/shape1", [0, -1, head_size]) |
| r1 = gb.node("Reshape", [x, shape1], [f"{name_prefix}/Reshape_1/out"], f"{name_prefix}/Reshape_1") |
| normed, _ = _simplified_layernorm(gb, r1, weight_name, eps, f"{name_prefix}/SimplifiedLayerNormalization") |
| shape2 = gb.const_i64(f"{name_prefix}/shape2", [0, -1, num_heads * head_size]) |
| return gb.node("Reshape", [normed, shape2], [f"{name_prefix}/Reshape_2/out"], f"{name_prefix}/Reshape_2") |
|
|
|
|
| def _matmul(gb, x, weight_param, name_prefix, torch_dtype): |
| |
| weight = weight_param.detach().to(torch_dtype).numpy().T |
| weight_name = gb.initializer(f"{name_prefix}.weight", weight) |
| return gb.node("MatMul", [x, weight_name], [f"{name_prefix}/out"], f"{name_prefix}/MatMul") |
|
|
|
|
| def _build_decoder_layer(gb, layer_id, layer, root_residual, input_ln_skip, dims, |
| seqlens_k_name, total_seq_len_name, cos_cache_name, sin_cache_name, torch_dtype): |
| """Emits one full Qwen3 decoder layer (attention + MLP). Returns (resid_before_mlp, mlp_output).""" |
| num_attn_heads, num_kv_heads = dims["num_attn_heads"], dims["num_kv_heads"] |
| head_size, eps = dims["head_size"], dims["rms_norm_eps"] |
|
|
| ln1_w = gb.initializer( |
| f"model.layers.{layer_id}.input_layernorm.weight", |
| layer.input_layernorm.weight.detach().to(torch_dtype).numpy(), |
| ) |
| normed, resid_before_attn = _simplified_layernorm( |
| gb, root_residual, ln1_w, eps, f"layer{layer_id}/input_layernorm", |
| skip=input_ln_skip, need_sum=(input_ln_skip is not None), |
| ) |
| if input_ln_skip is None: |
| resid_before_attn = root_residual |
|
|
| attn = layer.self_attn |
| q = _matmul(gb, normed, attn.q_proj.weight, f"layer{layer_id}/attn/q_proj", torch_dtype) |
| k = _matmul(gb, normed, attn.k_proj.weight, f"layer{layer_id}/attn/k_proj", torch_dtype) |
| v = _matmul(gb, normed, attn.v_proj.weight, f"layer{layer_id}/attn/v_proj", torch_dtype) |
|
|
| qn_w = gb.initializer(f"model.layers.{layer_id}.attn.q_norm.weight", |
| attn.q_norm.weight.detach().to(torch_dtype).numpy()) |
| kn_w = gb.initializer(f"model.layers.{layer_id}.attn.k_norm.weight", |
| attn.k_norm.weight.detach().to(torch_dtype).numpy()) |
| q = _qk_head_norm(gb, q, qn_w, num_attn_heads, head_size, eps, f"layer{layer_id}/attn/q_norm") |
| k = _qk_head_norm(gb, k, kn_w, num_kv_heads, head_size, eps, f"layer{layer_id}/attn/k_norm") |
|
|
| scale = float(1.0 / np.sqrt(head_size)) |
| attn_out = gb.node( |
| "GroupQueryAttention", |
| [q, k, v, "", "", seqlens_k_name, total_seq_len_name, cos_cache_name, sin_cache_name, "", "", ""], |
| [f"layer{layer_id}/attn/gqa_out", f"layer{layer_id}/attn/present_k", f"layer{layer_id}/attn/present_v"], |
| f"layer{layer_id}/attn/GQA", domain="com.microsoft", |
| num_heads=num_attn_heads, kv_num_heads=num_kv_heads, scale=scale, |
| local_window_size=-1, do_rotary=1, rotary_interleaved=0, |
| ) |
| o = _matmul(gb, attn_out, attn.o_proj.weight, f"layer{layer_id}/attn/o_proj", torch_dtype) |
|
|
| post_ln_w = gb.initializer( |
| f"model.layers.{layer_id}.post_attention_layernorm.weight", |
| layer.post_attention_layernorm.weight.detach().to(torch_dtype).numpy(), |
| ) |
| mlp_in, resid_before_mlp = _simplified_layernorm( |
| gb, resid_before_attn, post_ln_w, eps, f"layer{layer_id}/post_attention_layernorm", |
| skip=o, need_sum=True, |
| ) |
|
|
| mlp = layer.mlp |
| gate = _matmul(gb, mlp_in, mlp.gate_proj.weight, f"layer{layer_id}/mlp/gate_proj", torch_dtype) |
| up = _matmul(gb, mlp_in, mlp.up_proj.weight, f"layer{layer_id}/mlp/up_proj", torch_dtype) |
| sig = gb.node("Sigmoid", [gate], [f"layer{layer_id}/mlp/sigmoid"], f"layer{layer_id}/mlp/Sigmoid") |
| silu = gb.node("Mul", [gate, sig], [f"layer{layer_id}/mlp/silu"], f"layer{layer_id}/mlp/Mul_silu") |
| gated = gb.node("Mul", [silu, up], [f"layer{layer_id}/mlp/gated"], f"layer{layer_id}/mlp/Mul_gated") |
| down = _matmul(gb, gated, mlp.down_proj.weight, f"layer{layer_id}/mlp/down_proj", torch_dtype) |
|
|
| return resid_before_mlp, down |
|
|
|
|
| def _quantize_int4(onnx_model): |
| """Quantize the encoder graph to int4 using MatMulNBitsQuantizer. |
| |
| Applies int4 weight-only quantization to MatMul and Gather ops: |
| - MatMul -> MatMulNBits |
| - Gather (embedding) -> GatherBlockQuantized |
| """ |
| quantizer = MatMulNBitsQuantizer( |
| model=onnx_model, |
| bits=4, |
| block_size=32, |
| is_symmetric=True, |
| accuracy_level=4, |
| quant_format=QuantFormat.QOperator, |
| op_types_to_quantize=("MatMul", "Gather"), |
| ) |
| quantizer.process() |
| quantized = quantizer.model.model |
| quantized.ir_version = 10 |
| return quantized |
|
|
|
|
| def _build_encoder_graph(checkpoint_dir, dtype="f16", enable_webgpu_graph=False): |
| torch_dtype = _TORCH_DTYPES[dtype] |
| onnx_dtype = _ONNX_DTYPES[dtype] |
| np_dtype = _NP_DTYPES[dtype] |
|
|
| dims = _read_qwen3_config(checkpoint_dir) |
| num_layers = dims["num_layers"] |
|
|
| model = AutoModelForCausalLM.from_pretrained(checkpoint_dir, torch_dtype=torch_dtype, low_cpu_mem_usage=True) |
| model.eval() |
| layers = model.model.layers |
|
|
| gb = _GraphBuilder() |
| input_ids_name, attn_mask_name = "input_ids", "attention_mask" |
|
|
| cos_cache, sin_cache = _rotary_cos_sin_tables( |
| dims["head_size"], dims["rope_theta"], dims["max_position_embeddings"], np_dtype |
| ) |
| cos_cache_name = gb.initializer("cos_cache", cos_cache) |
| sin_cache_name = gb.initializer("sin_cache", sin_cache) |
|
|
| embed_w = gb.initializer("model.embed_tokens.weight", |
| model.model.embed_tokens.weight.detach().to(torch_dtype).numpy()) |
| embeddings = gb.node("Gather", [embed_w, input_ids_name], ["embeddings"], "embed/Gather", axis=0) |
|
|
| seqlens_k_name, total_seq_len_name = _attention_mask_reformat(gb, attn_mask_name, enable_webgpu_graph) |
|
|
| root_residual, input_ln_skip = embeddings, None |
| for layer_id in range(num_layers - 1): |
| resid_before_mlp, mlp_out = _build_decoder_layer( |
| gb, layer_id, layers[layer_id], root_residual, input_ln_skip, dims, |
| seqlens_k_name, total_seq_len_name, cos_cache_name, sin_cache_name, torch_dtype, |
| ) |
| root_residual, input_ln_skip = resid_before_mlp, mlp_out |
|
|
| |
| tap_id = num_layers - 1 |
| tap_ln_w = gb.initializer( |
| f"model.layers.{tap_id}.input_layernorm.weight", |
| layers[tap_id].input_layernorm.weight.detach().to(torch_dtype).numpy(), |
| ) |
| _, encoder_hidden_states = _simplified_layernorm( |
| gb, root_residual, tap_ln_w, dims["rms_norm_eps"], f"layer{tap_id}/input_layernorm", |
| skip=input_ln_skip, need_sum=True, sum_output_name="encoder_hidden_states", |
| ) |
| assert encoder_hidden_states is not None |
|
|
| graph_inputs = [ |
| helper.make_tensor_value_info(input_ids_name, TensorProto.INT64, [1, "sequence_length"]), |
| helper.make_tensor_value_info(attn_mask_name, TensorProto.INT64, [1, "total_sequence_length"]), |
| ] |
| graph_output = helper.make_tensor_value_info( |
| encoder_hidden_states, onnx_dtype, [1, "sequence_length", dims["hidden_size"]] |
| ) |
| graph = helper.make_graph(gb.nodes, "zimage_text_encoder", graph_inputs, [graph_output], initializer=gb.initializers) |
| onnx_model = helper.make_model( |
| graph, opset_imports=[helper.make_opsetid("", 17), helper.make_opsetid("com.microsoft", 1)] |
| ) |
| onnx_model.ir_version = 9 |
| return onnx_model |
|
|
|
|
| def export_qwen3_text_encoder(checkpoint_dir, output_onnx_path, quantize, dtype="f16", enable_webgpu_graph=False): |
| onnx_model = _build_encoder_graph(checkpoint_dir, dtype=dtype, enable_webgpu_graph=enable_webgpu_graph) |
| if quantize: |
| onnx_model = _quantize_int4(onnx_model) |
|
|
| out_dir = os.path.dirname(os.path.abspath(output_onnx_path)) or "." |
| os.makedirs(out_dir, exist_ok=True) |
| save_ir_model_sharded( |
| ir.from_proto(onnx_model), out_dir, os.path.basename(output_onnx_path), |
| size_threshold_bytes=INLINE_SIZE_THRESHOLD_BYTES, |
| max_shard_size_bytes=MAX_SHARD_SIZE_BYTES, |
| ) |
| print(f"Saved text encoder: {output_onnx_path}") |
| return output_onnx_path |
|
|
|
|
| def build(input_path, output_dir, precision="f16_int4_quant", extra_options=None): |
| if precision not in PRECISION_CONFIGS: |
| raise ValueError(f"Unknown precision '{precision}'; choose from {sorted(PRECISION_CONFIGS)}") |
| config = PRECISION_CONFIGS[precision] |
|
|
| extra_options = extra_options or {} |
| enable_webgpu_graph = str(extra_options.get("enable_webgpu_graph", "false")).lower() not in ("false", "0", "") |
|
|
| onnx_dir = os.path.join(output_dir, "onnx") |
| os.makedirs(onnx_dir, exist_ok=True) |
|
|
| suffix = config["suffix"] |
| output_onnx = os.path.join(onnx_dir, f"text_encoder_model_{suffix}.onnx") |
| export_qwen3_text_encoder( |
| input_path, output_onnx, quantize=config["quantize"], dtype=config["dtype"], |
| enable_webgpu_graph=enable_webgpu_graph, |
| ) |
| return onnx_dir |
|
|
|
|
| def get_args(): |
| parser = argparse.ArgumentParser(description="Export the Z-Image-Turbo text encoder (Qwen3) to ONNX.") |
| parser.add_argument("input_path", help="Path to the Z-Image-Turbo `text_encoder/` folder (or its parent).") |
| parser.add_argument( |
| "-o", "--output_dir", default=DEFAULT_OUTPUT_DIR, |
| help=f"Directory to write the ONNX model to (under an `onnx/` subdir). Default: {DEFAULT_OUTPUT_DIR}", |
| ) |
| parser.add_argument( |
| "-p", "--precision", default="f16_int4_quant", choices=sorted(PRECISION_CONFIGS), |
| help="Output precision. Default: f16_int4_quant.", |
| ) |
| parser.add_argument( |
| "--extra_options", nargs="*", default=[], |
| help="Extra key=value options. Supported: enable_webgpu_graph=true|false (default false). " |
| "true emits the GPU-only graph-capture attention-mask reformat (WebGPU graph capture); " |
| "false emits the standard Shape-based form preferred by WebNN EP.", |
| ) |
| return parser.parse_args() |
|
|
|
|
| def main(): |
| args = get_args() |
| input_path = args.input_path |
| if os.path.isdir(os.path.join(input_path, "text_encoder")): |
| input_path = os.path.join(input_path, "text_encoder") |
| onnx_dir = build(input_path, args.output_dir, args.precision, parse_extra_options(args.extra_options)) |
| print(f"\nSuccess: text_encoder exported to {onnx_dir}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
