| |
| |
| |
| |
| |
| |
| |
| """Standalone exporter for the Z-Image-Turbo VAE decoder (`AutoencoderKL.decoder`).""" |
|
|
| import argparse |
| import json |
| import os |
| import types |
| from typing import Any |
|
|
| import onnx_ir as ir |
| import torch |
| from onnxruntime_genai.models.builders.base import Model |
|
|
| PRECISION_CONFIGS = { |
| "f16": {"builder_precision": "fp16"}, |
| "f32": {"builder_precision": "fp32"}, |
| } |
|
|
| DEFAULT_OUTPUT_DIR = "z-image-turbo-onnx" |
|
|
|
|
| |
| |
| |
| def set_io_dtype(precision, execution_provider, extra_options) -> ir.DataType: |
| cpu_quant = precision in {"int4", "int8"} and execution_provider == "cpu" |
| fp32_webgpu = execution_provider == "webgpu" and extra_options.get("use_webgpu_fp32", False) |
| bf16_cuda = precision == "int4" and execution_provider in {"cuda", "trt-rtx"} and extra_options.get("use_cuda_bf16", False) |
|
|
| if precision == "fp32" or cpu_quant or fp32_webgpu: |
| return ir.DataType.FLOAT |
| if precision == "bf16" or bf16_cuda: |
| return ir.DataType.BFLOAT16 |
| return ir.DataType.FLOAT16 |
|
|
|
|
| def set_onnx_dtype(precision: str, extra_options: dict[str, Any]) -> ir.DataType: |
| if precision == "int4": |
| return ir.DataType.INT4 if extra_options.get("is_symmetric", True) else ir.DataType.UINT4 |
| if precision == "int8": |
| return ir.DataType.INT8 if extra_options.get("is_symmetric", True) else ir.DataType.UINT8 |
| to_onnx_dtype = { |
| "fp32": ir.DataType.FLOAT, |
| "fp16": ir.DataType.FLOAT16, |
| "bf16": ir.DataType.BFLOAT16, |
| } |
| return to_onnx_dtype[precision] |
|
|
|
|
| def load_diffusers_config(input_path): |
| config_path = os.path.join(input_path, "config.json") |
| with open(config_path) as f: |
| raw_config = json.load(f) |
| raw_config.setdefault("_name_or_path", input_path) |
| return types.SimpleNamespace(**raw_config) |
|
|
|
|
| def parse_extra_options(pairs): |
| extra_options = {} |
| for kv_str in pairs or []: |
| key, _, value = kv_str.partition("=") |
| extra_options[key.strip()] = value.strip() |
|
|
| for bool_key in ("hf_remote", "use_webgpu_fp32", "fuse_group_norm"): |
| if bool_key in extra_options: |
| extra_options[bool_key] = extra_options[bool_key] not in ("false", "False", "0") |
|
|
| return extra_options |
|
|
|
|
| class ZImageVAEDecoderModel(Model): |
| def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options): |
| mid_channels = int(config.block_out_channels[-1]) |
| fake_config = types.SimpleNamespace( |
| _name_or_path=getattr(config, "_name_or_path", "z-image-vae-decoder"), |
| architectures=["AutoencoderKL"], |
| hidden_size=mid_channels, |
| num_attention_heads=1, |
| num_key_value_heads=1, |
| num_hidden_layers=1, |
| intermediate_size=mid_channels, |
| vocab_size=0, |
| hidden_act="silu", |
| max_position_embeddings=1, |
| rms_norm_eps=1e-6, |
| ) |
| super().__init__(fake_config, io_dtype, onnx_dtype, ep, cache_dir, extra_options) |
|
|
| self.model_type = "z-image-vae-decoder" |
| |
| self.attention_attrs["unidirectional"] = False |
|
|
| self.mid_channels = mid_channels |
| self.latent_channels = int(config.latent_channels) |
| self.out_channels = int(config.out_channels) |
| self.block_out_channels = [int(c) for c in config.block_out_channels] |
| self.norm_num_groups = int(config.norm_num_groups) |
| |
| self.fuse_group_norm = bool(extra_options.get("fuse_group_norm", False)) |
| self._io_dtype = self.io_dtype |
|
|
| |
| |
| |
| def load_weights(self, input_path): |
| from diffusers import AutoencoderKL |
|
|
| |
| model = AutoencoderKL.from_pretrained( |
| input_path, |
| cache_dir=self.cache_dir, |
| token=self.hf_token, |
| trust_remote_code=self.hf_remote, |
| ) |
| model.eval() |
| return model.decoder |
|
|
| |
| |
| |
| def make_inputs_and_outputs(self): |
| self.input_names = {"latent_sample": "latent_sample"} |
| self.output_names = {"sample": "sample"} |
|
|
| latent = self.make_value( |
| "latent_sample", self._io_dtype, |
| shape=[1, self.latent_channels, "latent_height", "latent_width"], |
| ) |
| self.model.graph.inputs.extend([latent]) |
|
|
| sample = self.make_value( |
| "sample", self._io_dtype, shape=[1, self.out_channels, "height", "width"] |
| ) |
| self.model.graph.outputs.append(sample) |
|
|
| |
| |
| |
| def _init_name(self, node_name, suffix): |
| return node_name[1:].replace("/", ".") + suffix |
|
|
| def _const(self, dtype, value): |
| return f"/model/constants/{self.to_str_dtype(dtype)}/{value}" |
|
|
| def _transpose(self, name, root_input, perm): |
| self.make_transpose(name, root_input, self._io_dtype, shape=None, perm=perm) |
| return f"{name}/output_0" |
|
|
| def _add(self, name, inputs): |
| self.make_add(name, inputs, self._io_dtype, shape=None) |
| return f"{name}/output_0" |
|
|
| def _mul(self, name, inputs): |
| self.make_mul(name, inputs, self._io_dtype, shape=None) |
| return f"{name}/output_0" |
|
|
| def _reshape(self, name, root_input, shape_tensor_name): |
| self.make_reshape(name, [root_input, shape_tensor_name], self._io_dtype, shape=None) |
| return f"{name}/output_0" |
|
|
| def _conv(self, name, root_input, conv): |
| weight = self._init_name(name, ".weight") |
| self.make_initializer(conv.weight, weight, to=self._io_dtype) |
| inputs = [root_input, weight] |
| if conv.bias is not None: |
| bias = self._init_name(name, ".bias") |
| self.make_initializer(conv.bias, bias, to=self._io_dtype) |
| inputs.append(bias) |
|
|
| kh, kw = conv.kernel_size |
| ph, pw = conv.padding |
| sh, sw = conv.stride |
| dh, dw = conv.dilation |
| self.make_conv( |
| name, inputs, self._io_dtype, shape=None, |
| strides=[sh, sw], pads=[ph, pw, ph, pw], dilations=[dh, dw], |
| group=conv.groups, kernel_shape=[kh, kw], |
| ) |
| return f"{name}/output_0" |
|
|
| def _linear(self, name, root_input, linear): |
| """Attention Q/K/V/out projection. Goes through `make_matmul` so `-p int4/int8` applies.""" |
| self.make_matmul(linear, name, root_input, seq_dim="attn_seq") |
| output = f"{name}/output_0" |
| if linear.bias is not None: |
| self.make_add_bias(linear.bias, f"{name}/Add", output, seq_dim="attn_seq") |
| output = f"{name}/Add/output_0" |
| return output |
|
|
| |
| |
| |
| def _group_norm(self, name, source, gnmod, swish, expose_skip=True): |
| if self.fuse_group_norm: |
| return self._group_norm_fused(name, source, gnmod, swish, expose_skip) |
|
|
| skip_sum = None |
| if isinstance(source, tuple): |
| conv_branch, skip_branch = source |
| skip_sum = self._add(f"{name}/skip_add", [conv_branch, skip_branch]) |
| x = skip_sum |
| else: |
| x = source |
| normed = self._group_norm_decomposed(name, x, gnmod, swish) |
| return normed, (skip_sum if expose_skip else None) |
|
|
| def _group_norm_decomposed(self, name, x_nchw, gnmod, swish): |
| groups = int(gnmod.num_groups) |
| channels = int(gnmod.num_channels) |
| eps = float(gnmod.eps) |
|
|
| gamma = self._init_name(name, ".weight") |
| beta = self._init_name(name, ".bias") |
| self.make_initializer(gnmod.weight.detach().reshape(channels, 1, 1), gamma, to=self._io_dtype) |
| self.make_initializer(gnmod.bias.detach().reshape(channels, 1, 1), beta, to=self._io_dtype) |
| in_scale = f"{name}/instnorm_scale" |
| in_bias = f"{name}/instnorm_bias" |
| self.make_initializer(torch.ones(groups), in_scale, to=self._io_dtype) |
| self.make_initializer(torch.zeros(groups), in_bias, to=self._io_dtype) |
|
|
| self.make_shape(f"{name}/in_shape", x_nchw, shape=[4]) |
| grouped = self._reshape( |
| f"{name}/to_groups", x_nchw, self._const(ir.DataType.INT64, [0, groups, -1]) |
| ) |
| in_normed = f"{name}/InstanceNormalization/output_0" |
| self.make_node( |
| "InstanceNormalization", inputs=[grouped, in_scale, in_bias], outputs=[in_normed], |
| name=f"{name}/InstanceNormalization", epsilon=eps, |
| ) |
| self.make_value(in_normed, self._io_dtype) |
| normed = self._reshape(f"{name}/from_groups", in_normed, f"{name}/in_shape/output_0") |
| normed = self._mul(f"{name}/Mul", [normed, gamma]) |
| normed = self._add(f"{name}/Add", [normed, beta]) |
| if swish: |
| self.make_sigmoid(f"{name}/Sigmoid", normed, self._io_dtype, shape=None) |
| normed = self._mul(f"{name}/silu", [normed, f"{name}/Sigmoid/output_0"]) |
| return normed |
|
|
| def _group_norm_fused(self, name, source, gnmod, swish, expose_skip): |
| gamma = self._init_name(name, ".weight") |
| beta = self._init_name(name, ".bias") |
| self.make_initializer(gnmod.weight, gamma, to=self._io_dtype) |
| self.make_initializer(gnmod.bias, beta, to=self._io_dtype) |
| groups = int(gnmod.num_groups) |
| eps = float(gnmod.eps) |
| activation = 1 if swish else 0 |
|
|
| if isinstance(source, tuple): |
| conv_branch, skip_branch = source |
| x_nhwc = self._transpose(f"{name}/x_t", conv_branch, [0, 2, 3, 1]) |
| skip_nhwc = self._transpose(f"{name}/skip_t", skip_branch, [0, 2, 3, 1]) |
| normed_nhwc = f"{name}/output_0" |
| outputs = [normed_nhwc] |
| skip_sum_nhwc = None |
| if expose_skip: |
| skip_sum_nhwc = f"{name}/sum" |
| outputs.append(skip_sum_nhwc) |
| self.make_node( |
| "SkipGroupNorm", inputs=[x_nhwc, gamma, beta, skip_nhwc], outputs=outputs, |
| name=name, domain="com.microsoft", |
| activation=activation, channels_last=1, epsilon=eps, groups=groups, |
| ) |
| for output in outputs: |
| self.make_value(output, self._io_dtype) |
| normed = self._transpose(f"{name}/post_t", normed_nhwc, [0, 3, 1, 2]) |
| skip_sum = self._transpose(f"{name}/sum_post_t", skip_sum_nhwc, [0, 3, 1, 2]) if skip_sum_nhwc else None |
| return normed, skip_sum |
|
|
| x_nhwc = self._transpose(f"{name}/x_t", source, [0, 2, 3, 1]) |
| normed_nhwc = f"{name}/output_0" |
| self.make_node( |
| "GroupNorm", inputs=[x_nhwc, gamma, beta], outputs=[normed_nhwc], |
| name=name, domain="com.microsoft", |
| activation=activation, channels_last=1, epsilon=eps, groups=groups, |
| ) |
| self.make_value(normed_nhwc, self._io_dtype) |
| normed = self._transpose(f"{name}/post_t", normed_nhwc, [0, 3, 1, 2]) |
| return normed, None |
|
|
| |
| |
| |
| def _resnet(self, name, source, resnet): |
| normed, skip_sum = self._group_norm(f"{name}/norm1", source, resnet.norm1, swish=True, expose_skip=True) |
| if skip_sum is not None: |
| |
| residual_base = skip_sum |
| elif resnet.conv_shortcut is not None: |
| residual_base = self._conv(f"{name}/conv_shortcut", source, resnet.conv_shortcut) |
| else: |
| residual_base = source |
|
|
| hidden = self._conv(f"{name}/conv1", normed, resnet.conv1) |
| normed2, _ = self._group_norm(f"{name}/norm2", hidden, resnet.norm2, swish=True) |
| hidden = self._conv(f"{name}/conv2", normed2, resnet.conv2) |
| return (hidden, residual_base) |
|
|
| |
| |
| |
| def _mid_attention(self, name, x_nchw, attn): |
| assert attn.heads == 1, f"expected single-head VAE attention, got heads={attn.heads}" |
| normed, _ = self._group_norm(f"{name}/group_norm", x_nchw, attn.group_norm, swish=False) |
|
|
| to_seq_shape = self._const(ir.DataType.INT64, [0, self.mid_channels, -1]) |
| flat = self._reshape(f"{name}/flatten", normed, to_seq_shape) |
| seq = self._transpose(f"{name}/to_seq", flat, [0, 2, 1]) |
|
|
| q = self._linear(f"{name}/to_q", seq, attn.to_q) |
| k = self._linear(f"{name}/to_k", seq, attn.to_k) |
| v = self._linear(f"{name}/to_v", seq, attn.to_v) |
|
|
| scale = self._const(self._io_dtype, self.mid_channels ** -0.25) |
| q = self._mul(f"{name}/q_scale", [q, scale]) |
| k_t = self._transpose(f"{name}/kT", k, [0, 2, 1]) |
| k_t = self._mul(f"{name}/k_scale", [k_t, scale]) |
| scores = f"{name}/scores/output_0" |
| self.make_node("MatMul", inputs=[q, k_t], outputs=[scores], name=f"{name}/scores") |
| self.make_value(scores, self._io_dtype) |
| probs = f"{name}/softmax/output_0" |
| self.make_node("Softmax", inputs=[scores], outputs=[probs], name=f"{name}/softmax", axis=-1) |
| self.make_value(probs, self._io_dtype) |
| ctx = f"{name}/context/output_0" |
| self.make_node("MatMul", inputs=[probs, v], outputs=[ctx], name=f"{name}/context") |
| self.make_value(ctx, self._io_dtype) |
| attn_out = self._linear(f"{name}/to_out.0", ctx, attn.to_out[0]) |
|
|
| back_seq = self._transpose(f"{name}/from_seq", attn_out, [0, 2, 1]) |
| self.make_shape(f"{name}/in_shape", x_nchw, shape=[4]) |
| spatial = self._reshape(f"{name}/unflatten", back_seq, f"{name}/in_shape/output_0") |
|
|
| return self._add(f"{name}/Add", [spatial, x_nchw]) |
|
|
| |
| |
| |
| def _upsample(self, name, x_nchw, conv): |
| if self.fuse_group_norm: |
| x = self._transpose(f"{name}/pre_t", x_nchw, [0, 2, 3, 1]) |
| scales = self._const(ir.DataType.FLOAT, [1.0, 2.0, 2.0, 1.0]) |
| else: |
| x = x_nchw |
| scales = self._const(ir.DataType.FLOAT, [1.0, 1.0, 2.0, 2.0]) |
| resize_out = f"{name}/Resize/output_0" |
| self.make_node( |
| "Resize", inputs=[x, "", scales], outputs=[resize_out], name=f"{name}/Resize", |
| coordinate_transformation_mode="asymmetric", cubic_coeff_a=-0.75, |
| mode="nearest", nearest_mode="floor", |
| ) |
| self.make_value(resize_out, self._io_dtype) |
| y_nchw = self._transpose(f"{name}/post_t", resize_out, [0, 3, 1, 2]) if self.fuse_group_norm else resize_out |
| return self._conv(f"{name}/conv", y_nchw, conv) |
|
|
| |
| |
| |
| def make_model(self, input_path): |
| self.make_inputs_and_outputs() |
| self.weights = self.load_weights(input_path) |
| dec = self.weights |
|
|
| x = self.input_names["latent_sample"] |
|
|
| |
| x = self._conv("/decoder/conv_in", x, dec.conv_in) |
|
|
| |
| mid = dec.mid_block |
| r0 = self._resnet("/decoder/mid_block/resnets.0", x, mid.resnets[0]) |
| x = self._add("/decoder/mid_block/resnets.0/Add", list(r0)) |
| x = self._mid_attention("/decoder/mid_block/attentions.0", x, mid.attentions[0]) |
| source = self._resnet("/decoder/mid_block/resnets.1", x, mid.resnets[1]) |
|
|
| |
| for bi, up in enumerate(dec.up_blocks): |
| n_res = len(up.resnets) |
| for rj in range(n_res): |
| r = self._resnet(f"/decoder/up_blocks.{bi}/resnets.{rj}", source, up.resnets[rj]) |
| if rj < n_res - 1: |
| source = r |
| elif getattr(up, "upsamplers", None): |
| x = self._add(f"/decoder/up_blocks.{bi}/resnets.{rj}/Add", list(r)) |
| x = self._upsample(f"/decoder/up_blocks.{bi}/upsamplers.0", x, up.upsamplers[0].conv) |
| source = x |
| else: |
| source = r |
|
|
| |
| y, _ = self._group_norm( |
| "/decoder/conv_norm_out", source, dec.conv_norm_out, swish=True, expose_skip=False |
| ) |
| out = self._conv("/decoder/conv_out", y, dec.conv_out) |
|
|
| self.make_node("Identity", inputs=[out], outputs=["sample"], name="/decoder/output_identity") |
|
|
| del self.weights |
|
|
| |
| |
| |
| def save_model(self, out_dir): |
| print(f"Saving ONNX model in {out_dir}") |
| already_quantized_in_qdq_format = self.quant_type is not None and self.quant_attrs["use_qdq"] |
| if self.onnx_dtype in {ir.DataType.INT4, ir.DataType.UINT4, ir.DataType.INT8, ir.DataType.UINT8} and not already_quantized_in_qdq_format: |
| model = self.to_nbits() |
| else: |
| model = self.model |
|
|
| model.graph.sort() |
|
|
| out_path = os.path.join(out_dir, self.filename) |
| data_path = out_path + ".data" |
| for stale in (out_path, data_path): |
| if os.path.exists(stale): |
| print(f"Overwriting {stale}") |
| os.remove(stale) |
|
|
| ir.save(model, out_path) |
|
|
| if os.path.isdir(self.cache_dir) and not os.listdir(self.cache_dir): |
| os.rmdir(self.cache_dir) |
|
|
|
|
| |
| |
| |
| |
| def build(input_path, output_dir, precision="f16", extra_options=None): |
| if precision not in PRECISION_CONFIGS: |
| raise ValueError(f"Unknown precision '{precision}'; choose from {sorted(PRECISION_CONFIGS)}") |
|
|
| precision_config = PRECISION_CONFIGS[precision] |
| extra_options = dict(extra_options or {}) |
| extra_options.setdefault("hf_remote", False) |
| extra_options["filename"] = f"vae_decoder_model_{precision}.onnx" |
|
|
| builder_precision = precision_config["builder_precision"] |
| io_dtype = set_io_dtype(builder_precision, "webgpu", extra_options) |
| onnx_dtype = set_onnx_dtype(builder_precision, extra_options) |
|
|
| config = load_diffusers_config(input_path) |
|
|
| onnx_dir = os.path.join(output_dir, "onnx") |
| os.makedirs(onnx_dir, exist_ok=True) |
| cache_dir = os.path.join(output_dir, "cache") |
| os.makedirs(cache_dir, exist_ok=True) |
|
|
| model = ZImageVAEDecoderModel(config, io_dtype, onnx_dtype, "webgpu", cache_dir, extra_options) |
| model.make_model(input_path) |
| model.save_model(onnx_dir) |
| return onnx_dir |
|
|
|
|
| def get_args(): |
| parser = argparse.ArgumentParser(description="Export the Z-Image-Turbo VAE decoder to ONNX.") |
| parser.add_argument("input_path", help="Path to the Z-Image-Turbo `vae/` 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", choices=sorted(PRECISION_CONFIGS), |
| help="Output precision. Default: f16.", |
| ) |
| parser.add_argument( |
| "--extra_options", nargs="*", default=[], |
| help="Extra key=value options, e.g. fuse_group_norm=true.", |
| ) |
| return parser.parse_args() |
|
|
|
|
| def main(): |
| args = get_args() |
| input_path = args.input_path |
| if os.path.isdir(os.path.join(input_path, "vae")): |
| input_path = os.path.join(input_path, "vae") |
| onnx_dir = build(input_path, args.output_dir, args.precision, parse_extra_options(args.extra_options)) |
| print(f"\nSuccess: vae_decoder exported to {onnx_dir}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
