#!/usr/bin/env python3 """Export the Z-Image control transformer to ONNX for inference.""" import argparse import logging import os import sys from collections import OrderedDict from typing import Any, Dict, List, Optional, OrderedDict as OrderedDictType, Tuple import numpy as np import torch from omegaconf import OmegaConf from loguru import logger import onnx from onnx import numpy_helper import subprocess REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) if REPO_ROOT not in sys.path: sys.path.insert(0, REPO_ROOT) from videox_fun.models import ZImageControlTransformer2DModel # noqa: E402 from videox_fun.models.z_image_transformer2d import pad_stack # noqa: E402 LOGGER = logging.getLogger("export_transformer_onnx") logging.basicConfig(level=logging.INFO, format="[%(asctime)s] %(levelname)s: %(message)s") SEQ_MULTI_OF = 32 def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Export the Z-Image control transformer to ONNX") parser.add_argument("--config", default="config/z_image/z_image_control.yaml", help="Path to the YAML config used to build the transformer") parser.add_argument("--model-root", default="models/Diffusion_Transformer/Z-Image-Turbo/", help="Directory that stores the original diffusers weights") parser.add_argument("--checkpoint", default="models/Personalized_Model/Z-Image-Turbo-Fun-Controlnet-Union.safetensors", help="Optional fine-tuned checkpoint to load") parser.add_argument("--output", default="onnx-models/z_image_control_transformer.onnx", help="Target ONNX file path") parser.add_argument("--body-output", default="onnx-models/z_image_transformer_body.onnx", help="Path for the body-only ONNX when --split-control is enabled") parser.add_argument("--control-output", default="onnx-models/z_image_controlnet.onnx", help="Path for the control-only ONNX when --split-control is enabled") parser.add_argument("--height", type=int, default=864, help="Target image height used to derive latent resolution") parser.add_argument("--width", type=int, default=496, help="Target image width used to derive latent resolution") parser.add_argument("--batch-size", type=int, default=1, help="Batch size for the exported graph") parser.add_argument("--sequence-length", type=int, default=512, help="Prompt embedding sequence length (must be a multiple of 32)") parser.add_argument("--frames", type=int, default=1, help="Number of frames in the latent tensor") parser.add_argument("--latent-downsample-factor", type=int, default=8, help="Downsampling ratio between spatial image size and latent size") parser.add_argument("--latent-height", type=int, default=None, help="Override latent height (after downsampling)") parser.add_argument("--latent-width", type=int, default=None, help="Override latent width (after downsampling)") parser.add_argument("--dtype", choices=["fp16", "fp32"], default="fp16", help="Export precision") parser.add_argument("--control-scale", type=float, default=0.75, help="Default control context scale input") parser.add_argument("--patch-size", type=int, default=2, help="Spatial patch size used by the transformer") parser.add_argument("--f-patch-size", type=int, default=1, help="Frame patch size used by the transformer") parser.add_argument("--opset", type=int, default=17, help="ONNX opset version") parser.add_argument("--no-external-data", action="store_true", help="Disable external data format even if the model is larger than 2GB") parser.add_argument("--skip-ort-check", action="store_true", help="Skip running an ONNX Runtime correctness check") parser.add_argument("--ort-provider", default="CPUExecutionProvider", help="ONNX Runtime provider used during validation") parser.add_argument("--split-control", action="store_true", help="Export transformer body and ControlNet separately instead of a fused model") parser.add_argument("--save-calib-inputs", action="store_true", help="Dump ONNX input dictionaries as .npy for calibration") parser.add_argument("--calib-dir", default="onnx-calibration", help="Directory for storing calibration npy files") parser.add_argument("--dynamic-axes", action="store_true", help="Export ONNX with dynamic batch/seq/latent dims; default is static shape") parser.add_argument("--skip-slim", action="store_true", help="Skip onnxslim simplification for faster debug export") return parser.parse_args() def run_onnxslim(input_file="vae.onnx", output_file="vae_slim.onnx"): """ 执行 onnxslim 命令压缩 ONNX 模型 """ try: # 使用完整的命令路径（如果知道的话） cmd = ["onnxslim", input_file, output_file] print(f"执行命令: {' '.join(cmd)}") # 执行命令, 实时输出 process = subprocess.Popen( cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, bufsize=1, universal_newlines=True ) # 实时打印输出 for line in process.stdout: print(line, end='') # 等待命令完成 stdout, stderr = process.communicate() if process.returncode != 0: print(f"命令执行失败, 错误信息:\n{stderr}") return False else: print("ONNX模型压缩完成!") return True except FileNotFoundError: print("错误: 未找到 onnxslim 命令, 请确保已安装 onnxslim") print("安装方法: pip install onnx-simplifier") return False except Exception as e: print(f"执行命令时发生错误: {e}") return False def _resolve_path(path: str) -> str: return os.path.abspath(os.path.join(REPO_ROOT, path)) if not os.path.isabs(path) else path def load_transformer(args: argparse.Namespace, torch_dtype: torch.dtype, device: torch.device) -> ZImageControlTransformer2DModel: config_path = _resolve_path(args.config) model_root = _resolve_path(args.model_root) checkpoint_path = _resolve_path(args.checkpoint) if args.checkpoint else None if not os.path.exists(config_path): raise FileNotFoundError(f"Config not found: {config_path}") if not os.path.isdir(model_root): raise FileNotFoundError(f"Model root not found: {model_root}") config = OmegaConf.load(config_path) transformer_kwargs = OmegaConf.to_container(config.get("transformer_additional_kwargs", {}), resolve=True) LOGGER.info("Loading transformer from %s", model_root) transformer = ZImageControlTransformer2DModel.from_pretrained( model_root, subfolder="transformer", low_cpu_mem_usage=True, torch_dtype=torch_dtype, transformer_additional_kwargs=transformer_kwargs, ) transformer.eval() transformer.to(device=device, dtype=torch_dtype) if checkpoint_path and os.path.exists(checkpoint_path): LOGGER.info("Loading checkpoint %s", checkpoint_path) if checkpoint_path.endswith(".safetensors"): from safetensors.torch import load_file # type: ignore state_dict = load_file(checkpoint_path) else: state_dict = torch.load(checkpoint_path, map_location="cpu") state_dict = state_dict.get("state_dict", state_dict) missing, unexpected = transformer.load_state_dict(state_dict, strict=False) LOGGER.info("Checkpoint loaded (missing=%d, unexpected=%d)", len(missing), len(unexpected)) elif checkpoint_path: raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}") return transformer def _tensor_batch_to_list(batch_tensor: torch.Tensor) -> List[torch.Tensor]: return [batch_tensor[i] for i in range(batch_tensor.shape[0])] def _prepare_transformer_state( model: ZImageControlTransformer2DModel, latent_list: List[torch.Tensor], prompt_list: List[torch.Tensor], timestep: torch.Tensor, patch_size: int, f_patch_size: int, ) -> Dict[str, Any]: bsz = len(latent_list) device = latent_list[0].device timestep = timestep.to(device=device, dtype=torch.float32) t = timestep * model.t_scale t = model.t_embedder(t) ( x, cap_feats, x_size, x_pos_ids, cap_pos_ids, x_inner_pad_mask, cap_inner_pad_mask, ) = model.patchify_and_embed(latent_list, prompt_list, patch_size, f_patch_size) # Latent tokens refinement x_item_seqlens = [len(_) for _ in x] assert all(_ % SEQ_MULTI_OF == 0 for _ in x_item_seqlens) x_max_item_seqlen = max(x_item_seqlens) x = torch.cat(x, dim=0) x = model.all_x_embedder[f"{patch_size}-{f_patch_size}"](x) adaln_input = t.type_as(x) mask = torch.cat(x_inner_pad_mask) if torch.onnx.is_in_onnx_export(): if model.x_pad_token.dim() == 1: x_pad_token_2d = model.x_pad_token.unsqueeze(0) else: x_pad_token_2d = model.x_pad_token mask_2d = mask.unsqueeze(1) x_pad_expanded = x_pad_token_2d.expand_as(x) x = torch.where(mask_2d, x_pad_expanded, x) else: x[mask] = model.x_pad_token x = list(x.split(x_item_seqlens, dim=0)) x_freqs_cis = list(model.rope_embedder(torch.cat(x_pos_ids, dim=0)).split(x_item_seqlens, dim=0)) x = pad_stack(x, x_max_item_seqlen, pad_value=0.0) x_freqs_cis = pad_stack(x_freqs_cis, x_max_item_seqlen, pad_value=0.0) x_attn_mask = torch.zeros((bsz, x_max_item_seqlen), dtype=torch.bool, device=device) for i, seq_len in enumerate(x_item_seqlens): x_attn_mask[i, :seq_len] = 1 for layer in model.noise_refiner: x = layer(x, x_attn_mask, x_freqs_cis, adaln_input) # Caption refinement cap_item_seqlens = [len(_) for _ in cap_feats] assert all(_ % SEQ_MULTI_OF == 0 for _ in cap_item_seqlens) cap_max_item_seqlen = max(cap_item_seqlens) cap_feats = torch.cat(cap_feats, dim=0) cap_feats = model.cap_embedder(cap_feats) cap_mask = torch.cat(cap_inner_pad_mask) if torch.onnx.is_in_onnx_export(): if model.cap_pad_token.dim() == 1: cap_pad_token_2d = model.cap_pad_token.unsqueeze(0) else: cap_pad_token_2d = model.cap_pad_token mask_2d = cap_mask.unsqueeze(1) cap_pad_expanded = cap_pad_token_2d.expand_as(cap_feats) cap_feats = torch.where(mask_2d, cap_pad_expanded, cap_feats) else: cap_feats[cap_mask] = model.cap_pad_token cap_feats = list(cap_feats.split(cap_item_seqlens, dim=0)) cap_freqs_cis = list(model.rope_embedder(torch.cat(cap_pos_ids, dim=0)).split(cap_item_seqlens, dim=0)) cap_feats = pad_stack(cap_feats, cap_max_item_seqlen, pad_value=0.0) cap_freqs_cis = pad_stack(cap_freqs_cis, cap_max_item_seqlen, pad_value=0.0) cap_attn_mask = torch.zeros((bsz, cap_max_item_seqlen), dtype=torch.bool, device=device) for i, seq_len in enumerate(cap_item_seqlens): cap_attn_mask[i, :seq_len] = 1 for layer in model.context_refiner: cap_feats = layer(cap_feats, cap_attn_mask, cap_freqs_cis) # Context parallel handling if model.sp_world_size > 1: x = torch.chunk(x, model.sp_world_size, dim=1)[model.sp_world_rank] x_item_seqlens = [len(_) for _ in x] x_max_item_seqlen = max(x_item_seqlens) x_attn_mask = torch.zeros((bsz, x_max_item_seqlen), dtype=torch.bool, device=device) for i, seq_len in enumerate(x_item_seqlens): x_attn_mask[i, :seq_len] = 1 if x_freqs_cis is not None: x_freqs_cis = torch.chunk(x_freqs_cis, model.sp_world_size, dim=1)[model.sp_world_rank] unified = [] unified_freqs_cis = [] for i in range(bsz): x_len = x_item_seqlens[i] cap_len = cap_item_seqlens[i] unified.append(torch.cat([x[i][:x_len], cap_feats[i][:cap_len]])) unified_freqs_cis.append(torch.cat([x_freqs_cis[i][:x_len], cap_freqs_cis[i][:cap_len]])) unified_item_seqlens = [a + b for a, b in zip(cap_item_seqlens, x_item_seqlens)] unified_max_item_seqlen = max(unified_item_seqlens) unified = pad_stack(unified, unified_max_item_seqlen, pad_value=0.0) unified_freqs_cis = pad_stack(unified_freqs_cis, unified_max_item_seqlen, pad_value=0.0) unified_attn_mask = torch.zeros((bsz, unified_max_item_seqlen), dtype=torch.bool, device=device) for i, seq_len in enumerate(unified_item_seqlens): unified_attn_mask[i, :seq_len] = 1 kwargs = dict(attn_mask=unified_attn_mask, freqs_cis=unified_freqs_cis, adaln_input=adaln_input) return dict( x=x, cap_feats=cap_feats, unified=unified, kwargs=kwargs, adaln_input=adaln_input, time_embed=t, x_item_seqlens=x_item_seqlens, cap_item_seqlens=cap_item_seqlens, x_size=x_size, unified_attn_mask=unified_attn_mask, unified_freqs_cis=unified_freqs_cis, bsz=bsz, ) class TransformerOnnxWrapper(torch.nn.Module): """Lightweight wrapper that exposes tensor inputs for ONNX export.""" def __init__(self, model: ZImageControlTransformer2DModel, patch_size: int, f_patch_size: int): super().__init__() self.model = model self.patch_size = patch_size self.f_patch_size = f_patch_size def forward( self, latent_model_input: torch.Tensor, timestep: torch.Tensor, prompt_embeds: torch.Tensor, control_context: torch.Tensor, control_context_scale: torch.Tensor, ) -> torch.Tensor: latents = list(latent_model_input.unbind(dim=0)) prompts = list(prompt_embeds.unbind(dim=0)) scale = control_context_scale.to(device=latent_model_input.device, dtype=latent_model_input.dtype) outputs, _ = self.model( latents, timestep, prompts, patch_size=self.patch_size, f_patch_size=self.f_patch_size, control_context=control_context, control_context_scale=scale, ) return outputs class ControlNetWrapper(torch.nn.Module): """Exports only the control branch so it can be invoked on demand.""" def __init__(self, model: ZImageControlTransformer2DModel, patch_size: int, f_patch_size: int): super().__init__() self.model = model self.patch_size = patch_size self.f_patch_size = f_patch_size def forward( self, latent_model_input: torch.Tensor, timestep: torch.Tensor, prompt_embeds: torch.Tensor, control_context: torch.Tensor, ) -> torch.Tensor: latents = _tensor_batch_to_list(latent_model_input) prompts = _tensor_batch_to_list(prompt_embeds) control_list = _tensor_batch_to_list(control_context) state = _prepare_transformer_state( self.model, latents, prompts, timestep, self.patch_size, self.f_patch_size, ) hints = self.model.forward_control( state["unified"], state["cap_feats"], control_list, state["kwargs"], t=state["time_embed"], patch_size=self.patch_size, f_patch_size=self.f_patch_size, ) return torch.stack(hints, dim=0) class TransformerBodyWrapper(torch.nn.Module): """Exports the transformer body which consumes precomputed control hints.""" def __init__(self, model: ZImageControlTransformer2DModel, patch_size: int, f_patch_size: int): super().__init__() self.model = model self.patch_size = patch_size self.f_patch_size = f_patch_size def forward( self, latent_model_input: torch.Tensor, timestep: torch.Tensor, prompt_embeds: torch.Tensor, control_hints: torch.Tensor, control_context_scale: torch.Tensor, ) -> torch.Tensor: latents = _tensor_batch_to_list(latent_model_input) prompts = _tensor_batch_to_list(prompt_embeds) hints_list = list(torch.unbind(control_hints, dim=0)) scale = control_context_scale.to(device=latent_model_input.device, dtype=latent_model_input.dtype) state = _prepare_transformer_state( self.model, latents, prompts, timestep, self.patch_size, self.f_patch_size, ) unified = state["unified"] for layer in self.model.layers: layer_kwargs = dict( attn_mask=state["unified_attn_mask"], freqs_cis=state["unified_freqs_cis"], adaln_input=state["adaln_input"], hints=hints_list, context_scale=scale, ) unified = layer(unified, **layer_kwargs) if self.model.sp_world_size > 1: unified_out = [] for i in range(state["bsz"]): x_len = state["x_item_seqlens"][i] unified_out.append(unified[i, :x_len]) unified = torch.stack(unified_out) unified = self.model.all_gather(unified, dim=1) final_layer = self.model.all_final_layer[f"{self.patch_size}-{self.f_patch_size}"] unified = final_layer(unified, state["adaln_input"]) unified = list(unified.unbind(dim=0)) x = self.model.unpatchify(unified, state["x_size"], self.patch_size, self.f_patch_size) x = torch.stack(x) return x def _validate_sequence_length(seq_len: int) -> None: if seq_len % 32 != 0: raise ValueError("sequence_length must be a multiple of 32 to satisfy transformer padding rules") def _compute_latent_dims(args: argparse.Namespace) -> Dict[str, int]: if args.latent_height is not None and args.latent_width is not None: latent_h = args.latent_height latent_w = args.latent_width else: if args.height % args.latent_downsample_factor != 0 or args.width % args.latent_downsample_factor != 0: raise ValueError("height and width must be divisible by latent_downsample_factor") latent_h = args.height // args.latent_downsample_factor latent_w = args.width // args.latent_downsample_factor if latent_h % args.patch_size != 0 or latent_w % args.patch_size != 0: raise ValueError("latent dimensions must be divisible by patch_size") if args.frames % args.f_patch_size != 0: raise ValueError("frames must be divisible by f_patch_size") return {"latent_h": latent_h, "latent_w": latent_w} def build_dummy_inputs( args: argparse.Namespace, model: ZImageControlTransformer2DModel, torch_dtype: torch.dtype, device: torch.device, ) -> OrderedDictType[str, torch.Tensor]: _validate_sequence_length(args.sequence_length) dims = _compute_latent_dims(args) batch = args.batch_size in_channels = model.config.in_channels cap_dim = model.config.cap_feat_dim latent = torch.randn( batch, in_channels, args.frames, dims["latent_h"], dims["latent_w"], dtype=torch_dtype, device=device, ) timestep = torch.linspace(0.0, 1.0, steps=batch, dtype=torch.float32, device=device) prompts = torch.randn( batch, args.sequence_length, cap_dim, dtype=torch_dtype, device=device, ) control = torch.randn_like(latent) control_scale = torch.full((1,), args.control_scale, dtype=torch.float32, device=device) return OrderedDict( latent_model_input=latent, timestep=timestep, prompt_embeds=prompts, control_context=control, control_context_scale=control_scale, ) def maybe_save_calibration_inputs(tag: str, inputs: OrderedDictType[str, torch.Tensor], args: argparse.Namespace) -> Optional[str]: if not getattr(args, "save_calib_inputs", False): return None output_dir = _resolve_path(args.calib_dir) os.makedirs(output_dir, exist_ok=True) numpy_dict = {name: tensor.detach().cpu().numpy() for name, tensor in inputs.items()} file_path = os.path.join(output_dir, f"{tag}_inputs.npy") np.save(file_path, numpy_dict, allow_pickle=True) LOGGER.info("Saved calibration inputs (%s) to %s", tag, file_path) return file_path def dump_initializer_parameters(model_path: str) -> str: """Save all ONNX initializers into a standalone .npz file.""" model_proto = onnx.load(model_path, load_external_data=True) param_dict = {} for initializer in model_proto.graph.initializer: param_dict[initializer.name] = numpy_helper.to_array(initializer) param_path = f"{model_path}.params.npz" np.savez(param_path, **param_dict) LOGGER.info("Saved %d parameters to %s", len(param_dict), param_path) return param_path def export_onnx( wrapper: torch.nn.Module, sample_inputs: OrderedDictType[str, torch.Tensor], output_path: str, output_names: List[str], args: argparse.Namespace, ) -> Tuple[str, str]: export_path = _resolve_path(output_path) export_dir = os.path.dirname(export_path) if export_dir: os.makedirs(export_dir, exist_ok=True) input_names = list(sample_inputs.keys()) use_external = not args.no_external_data wrapper.eval() dynamic_axes = None if args.dynamic_axes: dynamic_axes = { "latent_model_input": {0: "batch", 2: "frames", 3: "latent_h", 4: "latent_w"}, "prompt_embeds": {0: "batch", 1: "seq_len"}, "timestep": {0: "batch"}, "control_context": {0: "batch", 2: "frames", 3: "latent_h", 4: "latent_w"}, "control_hints": {0: "batch", 2: "frames", 3: "latent_h", 4: "latent_w"}, "control_context_scale": {0: "scale_batch"}, "sample": {0: "batch", 2: "frames", 3: "latent_h", 4: "latent_w"}, "hints": {0: "batch", 2: "frames", 3: "latent_h", 4: "latent_w"}, } LOGGER.info("Exporting ONNX to %s", export_path) with torch.inference_mode(): torch.onnx.export( wrapper, args=tuple(sample_inputs[name] for name in input_names), f=export_path, input_names=input_names, output_names=output_names, opset_version=args.opset, do_constant_folding=True, export_params=True, dynamic_axes={k: v for k, v in dynamic_axes.items() if k in input_names + output_names} if dynamic_axes else None, # use_external_data_format=use_external, ) LOGGER.info("Raw ONNX export finished") trans_onnx = onnx.load(export_path) simp_onnx_data = os.path.splitext(export_path)[0] + "_simp.onnx" onnx.save( trans_onnx, simp_onnx_data, save_as_external_data=True, all_tensors_to_one_file=True, ) external_weight_file = simp_onnx_data + ".data" LOGGER.info("Saved external-data ONNX to %s (weights -> %s)", simp_onnx_data, external_weight_file) if args.skip_slim: LOGGER.info("Skip onnxslim as requested, using simplified external-data ONNX: %s", simp_onnx_data) final_onnx = simp_onnx_data else: slim_onnx_path = os.path.splitext(simp_onnx_data)[0] + "_slim.onnx" LOGGER.info("Transformer ONNX model exported, start to simplify via onnxslim") success = run_onnxslim(simp_onnx_data, slim_onnx_path) if not success: raise RuntimeError("onnxslim simplification failed, please check logs") final_onnx = slim_onnx_path LOGGER.info("Transformer ONNX model exported successfully: %s", final_onnx) param_path = "" # param_path = dump_initializer_parameters(final_onnx) # 当前看来无必要保存 return final_onnx, param_path def run_ort_validation( wrapper: torch.nn.Module, sample_inputs: OrderedDictType[str, torch.Tensor], onnx_path: str, provider: str, ) -> None: try: import onnxruntime as ort except ImportError: # pragma: no cover LOGGER.warning("onnxruntime not installed, skip validation") return wrapper.eval() with torch.inference_mode(): torch_output = wrapper(*sample_inputs.values()).detach().cpu().numpy() sess_options = ort.SessionOptions() session = ort.InferenceSession(onnx_path, sess_options=sess_options, providers=[provider]) ort_inputs = { name: tensor.detach().cpu().numpy() for name, tensor in sample_inputs.items() } ort_output = session.run(None, ort_inputs)[0] abs_diff = np.max(np.abs(torch_output - ort_output)) rel_diff = abs_diff / (np.maximum(1.0, np.max(np.abs(torch_output)))) LOGGER.info("ONNX Runtime check done (abs=%.6f, rel=%.6f)", abs_diff, rel_diff) def main() -> None: args = parse_args() device = torch.device("cpu") torch_dtype = torch.float16 if args.dtype == "fp16" else torch.float32 torch.set_grad_enabled(False) transformer = load_transformer(args, torch_dtype, device) sample_inputs = build_dummy_inputs(args, transformer, torch_dtype, device) if args.split_control: control_wrapper = ControlNetWrapper(transformer, args.patch_size, args.f_patch_size) body_wrapper = TransformerBodyWrapper(transformer, args.patch_size, args.f_patch_size) control_inputs = OrderedDict( latent_model_input=sample_inputs["latent_model_input"], timestep=sample_inputs["timestep"], prompt_embeds=sample_inputs["prompt_embeds"], control_context=sample_inputs["control_context"], ) maybe_save_calibration_inputs("controlnet", control_inputs, args) with torch.inference_mode(): control_hints_sample = control_wrapper(*control_inputs.values()).detach() control_model_path, _ = export_onnx( control_wrapper, control_inputs, args.control_output, ["hints"], args, ) body_inputs = OrderedDict( latent_model_input=sample_inputs["latent_model_input"], timestep=sample_inputs["timestep"], prompt_embeds=sample_inputs["prompt_embeds"], control_hints=control_hints_sample, control_context_scale=sample_inputs["control_context_scale"], ) maybe_save_calibration_inputs("transformer_body", body_inputs, args) body_model_path, _ = export_onnx( body_wrapper, body_inputs, args.body_output, ["sample"], args, ) if not args.skip_ort_check: try: run_ort_validation(control_wrapper, control_inputs, control_model_path, args.ort_provider) run_ort_validation(body_wrapper, body_inputs, body_model_path, args.ort_provider) except Exception as exc: # pragma: no cover LOGGER.warning("ONNX Runtime validation failed: %s", exc) else: wrapper = TransformerOnnxWrapper(transformer, args.patch_size, args.f_patch_size) maybe_save_calibration_inputs("transformer", sample_inputs, args) transformer_model_path, _ = export_onnx( wrapper, sample_inputs, args.output, ["sample"], args, ) if not args.skip_ort_check: try: run_ort_validation(wrapper, sample_inputs, transformer_model_path, args.ort_provider) except Exception as exc: # pragma: no cover LOGGER.warning("ONNX Runtime validation failed: %s", exc) if __name__ == "__main__": """ python scripts/z_image_fun/export_transformer_onnx.py \ --split-control \ --control-output onnx-models/z_image_controlnet.onnx \ --body-output onnx-models/z_image_transformer_body.onnx \ --save-calib-inputs \ --height 512 \ --width 512 \ --sequence-length 128 \ --latent-downsample-factor 8 \ --skip-slim \ --dtype fp32 """ main()