File size: 5,717 Bytes
610d9c2 | 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 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 | # -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------
"""
ONNX Model Optimizer for Stable Diffusion
"""
import gc
import logging
import os
import shutil
import tempfile
from pathlib import Path
import onnx
from packaging import version
from onnxruntime.transformers.fusion_options import FusionOptions
from onnxruntime.transformers.onnx_model_clip import ClipOnnxModel
from onnxruntime.transformers.onnx_model_unet import UnetOnnxModel
from onnxruntime.transformers.onnx_model_vae import VaeOnnxModel
from onnxruntime.transformers.optimizer import optimize_by_onnxruntime, optimize_model
logger = logging.getLogger(__name__)
class OrtStableDiffusionOptimizer:
def __init__(self, model_type: str):
assert model_type in ["vae", "unet", "clip"]
self.model_type = model_type
self.model_type_class_mapping = {
"unet": UnetOnnxModel,
"vae": VaeOnnxModel,
"clip": ClipOnnxModel,
}
def _optimize_by_ort(self, onnx_model, use_external_data_format, tmp_dir):
# Save to a temporary file so that we can load it with Onnx Runtime.
logger.info("Saving a temporary model to run OnnxRuntime graph optimizations...")
tmp_model_path = Path(tmp_dir) / "model.onnx"
onnx_model.save_model_to_file(str(tmp_model_path), use_external_data_format=use_external_data_format)
del onnx_model
gc.collect()
ort_optimized_model_path = Path(tmp_dir) / "optimized.onnx"
optimize_by_onnxruntime(
str(tmp_model_path),
use_gpu=True,
optimized_model_path=str(ort_optimized_model_path),
save_as_external_data=use_external_data_format,
external_data_filename="optimized.onnx_data",
)
model = onnx.load(str(ort_optimized_model_path), load_external_data=True)
return self.model_type_class_mapping[self.model_type](model)
def optimize_by_ort(self, onnx_model, use_external_data_format=False, tmp_dir=None):
# Use this step to see the final graph that executed by Onnx Runtime.
if tmp_dir is None:
with tempfile.TemporaryDirectory() as temp_dir:
return self._optimize_by_ort(onnx_model, use_external_data_format, temp_dir)
else:
os.makedirs(tmp_dir, exist_ok=True)
model = self._optimize_by_ort(onnx_model, use_external_data_format, tmp_dir)
shutil.rmtree(tmp_dir)
return model
def optimize(
self,
input_fp32_onnx_path,
optimized_onnx_path,
float16=True,
keep_io_types=False,
fp32_op_list=None,
keep_outputs=None,
optimize_by_ort=True,
optimize_by_fusion=True,
final_target_float16=True,
tmp_dir=None,
):
"""Optimize onnx model using ONNX Runtime transformers optimizer"""
logger.info(f"Optimize {input_fp32_onnx_path}...")
if optimize_by_fusion:
fusion_options = FusionOptions(self.model_type)
# It is allowed float16=False and final_target_float16=True, for using fp32 as intermediate optimization step.
# For rare fp32 use case, we can disable packed kv/qkv since there is no fp32 TRT fused attention kernel.
if self.model_type in ["unet"] and not final_target_float16:
fusion_options.enable_packed_kv = False
fusion_options.enable_packed_qkv = False
m = optimize_model(
input_fp32_onnx_path,
model_type=self.model_type,
num_heads=0, # will be deduced from graph
hidden_size=0, # will be deduced from graph
opt_level=0,
optimization_options=fusion_options,
use_gpu=True,
)
else:
model = onnx.load_model(input_fp32_onnx_path, load_external_data=True)
m = self.model_type_class_mapping[self.model_type](model)
if keep_outputs:
m.prune_graph(outputs=keep_outputs)
model_size = m.model.ByteSize()
# model size might be negative (overflow?) in Windows.
use_external_data_format = model_size <= 0 or model_size >= onnx.checker.MAXIMUM_PROTOBUF
# Note that ORT < 1.16 could not save model larger than 2GB.
# This step is is optional since it has no impact on inference latency.
# The optimized model is not portable. It could only run in the same execution provider (CUDA EP in this case).
# When the model has been optimized by onnxruntime, we can disable optimization in SessionOption
# to save session creation time. Another benefit is to inspect the final graph for developing purpose.
from onnxruntime import __version__ as ort_version # noqa: PLC0415
if optimize_by_ort and (version.parse(ort_version) >= version.parse("1.16.0") or not use_external_data_format):
m = self.optimize_by_ort(m, use_external_data_format=use_external_data_format, tmp_dir=tmp_dir)
if float16:
logger.info("Convert to float16 ...")
m.convert_float_to_float16(
keep_io_types=keep_io_types,
op_block_list=fp32_op_list,
)
m.get_operator_statistics()
m.get_fused_operator_statistics()
m.save_model_to_file(optimized_onnx_path, use_external_data_format=use_external_data_format)
logger.info("%s is optimized: %s", self.model_type, optimized_onnx_path)
|