File size: 2,422 Bytes
5fab480 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | # -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------
# Modifications Copyright (C) 2026 Intel Corporation. All rights reserved.
# --------------------------------------------------------------------------
"""Shared ONNX external-data save helper for the Z-Image-Turbo exporters."""
import glob
import os
import re
import onnx
import onnx_ir as ir
INLINE_SIZE_THRESHOLD_BYTES = 1 * 1024**2
MAX_SHARD_SIZE_BYTES = int(1.9 * 1024**3) # stays under the browser's 2 GiB ArrayBuffer ceiling
_SHARD_RE = re.compile(r"-(\d{5})-of-(\d{5})") # onnx_ir's "<stem>-000i-of-000N<ext>" naming
def save_ir_model_sharded(
model, out_dir, onnx_filename, *, size_threshold_bytes, max_shard_size_bytes, callback=None
):
"""Save with small weights inline and external data sharded; returns the `.onnx` path."""
out_path = os.path.join(out_dir, onnx_filename)
data_base = onnx_filename + "_data"
stem, ext = os.path.splitext(data_base)
# onnx_ir refuses to overwrite a pre-existing shard file, so clear stale outputs first.
stale = [out_path, *glob.glob(os.path.join(out_dir, data_base + "*")),
*glob.glob(os.path.join(out_dir, f"{stem}-*-of-*{ext}"))]
for path in stale:
if os.path.exists(path):
os.remove(path)
ir.save(
model, out_path, external_data=data_base,
size_threshold_bytes=size_threshold_bytes, max_shard_size_bytes=max_shard_size_bytes,
callback=callback,
)
temp_shards = sorted(glob.glob(os.path.join(out_dir, f"{stem}-*-of-*{ext}")))
if not temp_shards:
return out_path
remap = {}
for path in temp_shards:
base = os.path.basename(path)
idx = int(_SHARD_RE.search(base).group(1))
final = data_base if idx == 1 else f"{data_base}_{idx - 1}"
os.replace(path, os.path.join(out_dir, final))
remap[base] = final
proto = onnx.load(out_path, load_external_data=False)
for init in proto.graph.initializer:
for kv in init.external_data:
if kv.key == "location" and kv.value in remap:
kv.value = remap[kv.value]
onnx.save(proto, out_path)
return out_path
|