saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Update optimize_pipeline for SDXL (#17536)

Browse source 
- [x] Optimize SDXL models exported by optimum.
- [x] Enable it to run locally instead of using module.
- [x] Detect external data file in original model, and save with same
format by default.
- [x]  Add tests

### Example
```
pip install optimum transformers diffusers onnx onnxruntime-gpu>=1.16
optimum-cli export onnx --model stabilityai/stable-diffusion-xl-base-1.0 --task stable-diffusion-xl ./sd_xl_base_onnx
python -m  onnxruntime.transformers.models.stable_diffusion.optimize_pipeline -i ./sd_xl_base_onnx -o ./sd_xl_base_fp16 --float16
```

### Known issues
(1) VAE decoder cannot be converted to float16. Otherwise, there will be
black image in output.
(2) To use the float16 models, need a minor change in optimum to convert
the inputs for VAE decoder from float16 to float32 since we keep VAE
decoder as float32. The change is to append a line like the following
after [this
line](afd2b5a366/optimum/pipelines/diffusers/pipeline_stable_diffusion_xl.py (L483))
```
latents = latents.astype(np.float32)
```
This commit is contained in:
Tianlei Wu 2023-09-15 10:17:20 -07:00 committed by GitHub
parent 377f959c69
commit af80542e65
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 226 additions and 69 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
onnxruntime/python/tools/transformers/models/stable_diffusion/README.md
View file

@ -154,6 +154,12 @@ curl https://raw.githubusercontent.com/huggingface/diffusers/v0.15.1/scripts/con
python convert_sd_onnx.py --model_path runwayml/stable-diffusion-v1-5 --output_path ./sd_v1_5/fp32
```
For SDXL, use optimum to export the model:
```
pip install optimum diffusers onnx onnxruntime-gpu
optimum-cli export onnx --model stabilityai/stable-diffusion-xl-base-1.0 --task stable-diffusion-xl ./sd_xl_base_onnx
```
### Optimize ONNX Pipeline
Example to optimize the exported float32 ONNX models, and save to float16 models:
@ -161,7 +167,10 @@ Example to optimize the exported float32 ONNX models, and save to float16 models
python -m onnxruntime.transformers.models.stable_diffusion.optimize_pipeline -i ./sd_v1_5/fp32 -o ./sd_v1_5/fp16 --float16
```
If you installed ONNX Runtime v1.14, some optimizations (packed QKV and BiasAdd) will be disabled automatically since they are not available in v1.14.
For SDXL model, it is recommended to use a machine with 32 GB or more memory to optimize.
```
python optimize_pipeline.py -i ./sd_xl_base_onnx -o ./sd_xl_base_fp16 --float16
```
### Run Benchmark

161
onnxruntime/python/tools/transformers/models/stable_diffusion/optimize_pipeline.py
View file

@ -5,19 +5,15 @@
#
# This script converts stable diffusion onnx models from float to half (mixed) precision for GPU inference.
#
# Before running this script, follow README.md to setup python environment and convert stable diffusion checkpoint to float32 onnx models.
# Before running this script, follow README.md to setup python environment and convert stable diffusion checkpoint
# to float32 onnx models.
#
# For example, the float32 ONNX pipeline is saved to ./sd-v1-5 directory, you can optimize and convert it to float16 like the following:
# For example, the float32 ONNX pipeline is saved to ./sd-v1-5 directory, you can optimize and convert it to float16
# like the following:
# python optimize_pipeline.py -i ./sd-v1-5 -o ./sd-v1-5-fp16 --float16
#
# Note that the optimizations are carried out for CUDA Execution Provider at first, other EPs may not have the support for the fused opeartors.
# In this case, the users should disable the operator fusion manually to workaround.
#
# Stable diffusion 2.1 model will get black images using float16 Attention. A walkaround is to force Attention to run in float32 like the following:
# python optimize_pipeline.py -i ./sd-v2-1 -o ./sd-v2-1-fp16 --float16 --force_fp32_ops unet:Attention
#
# If you are using nightly package (or built from source), you can force MultiHeadAttention to run in float32:
# python optimize_pipeline.py -i ./sd-v2-1 -o ./sd-v2-1-fp16 --float16 --force_fp32_ops unet:MultiHeadAttention
# Note that the optimizations are carried out for CUDA Execution Provider at first, other EPs may not have the support
# for the fused opeartors. The users could disable the operator fusion manually to workaround.
import argparse
import logging
@ -25,8 +21,9 @@ import os
import shutil
import tempfile
from pathlib import Path
from typing import List
from typing import List, Optional
import __init__ # noqa: F401. Walk-around to run this script directly
import coloredlogs
import onnx
from fusion_options import FusionOptions
@ -41,11 +38,19 @@ import onnxruntime
logger = logging.getLogger(__name__)
def optimize_sd_pipeline(
def has_external_data(onnx_model_path):
original_model = onnx.load_model(str(onnx_model_path), load_external_data=False)
for initializer in original_model.graph.initializer:
if initializer.HasField("data_location") and initializer.data_location == onnx.TensorProto.EXTERNAL:
return True
return False
def _optimize_sd_pipeline(
source_dir: Path,
target_dir: Path,
overwrite: bool,
use_external_data_format: bool,
use_external_data_format: Optional[bool],
float16: bool,
force_fp32_ops: List[str],
enable_runtime_optimization: bool,
@ -57,7 +62,7 @@ def optimize_sd_pipeline(
source_dir (Path): Root of input directory of stable diffusion onnx pipeline with float32 models.
target_dir (Path): Root of output directory of stable diffusion onnx pipeline with optimized models.
overwrite (bool): Overwrite files if exists.
use_external_data_format (bool): save onnx model to two files: one for onnx graph, another for weights
use_external_data_format (Optional[bool]): use external data format.
float16 (bool): use half precision
force_fp32_ops(List[str]): operators that are forced to run in float32.
enable_runtime_optimization(bool): run graph optimization using Onnx Runtime.
@ -71,6 +76,7 @@ def optimize_sd_pipeline(
"vae_encoder": "vae",
"vae_decoder": "vae",
"text_encoder": "clip",
"text_encoder_2": "clip",
"safety_checker": "unet",
}
@ -85,9 +91,12 @@ def optimize_sd_pipeline(
"vae_encoder": [],
"vae_decoder": [],
"text_encoder": [],
"text_encoder_2": [],
"safety_checker": [],
}
is_xl = (source_dir / "text_encoder_2").exists()
if force_fp32_ops:
for fp32_operator in force_fp32_ops:
parts = fp32_operator.split(":")
@ -100,26 +109,21 @@ def optimize_sd_pipeline(
for name, model_type in model_type_mapping.items():
onnx_model_path = source_dir / name / "model.onnx"
if not os.path.exists(onnx_model_path):
message = f"input onnx model does not exist: {onnx_model_path}."
if name not in ["safety_checker"]:
raise RuntimeError(message)
if name != "safety_checker":
logger.info("input onnx model does not exist: %s", onnx_model_path)
# some model are optional so we do not raise error here.
continue
# Prepare output directory
optimized_model_path = target_dir / name / "model.onnx"
output_dir = optimized_model_path.parent
if optimized_model_path.exists():
if not overwrite:
raise RuntimeError(f"output onnx model path existed: {optimized_model_path}")
if output_dir.exists():
shutil.rmtree(output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
if use_external_data_format is None:
use_external_data_format = has_external_data(onnx_model_path)
# Graph fusion before fp16 conversion, otherwise they cannot be fused later.
# Right now, onnxruntime does not save >2GB model so we use script to optimize unet instead.
logger.info(f"Optimize {onnx_model_path}...")
args.model_type = model_type
@ -143,12 +147,15 @@ def optimize_sd_pipeline(
)
if float16:
logger.info("Convert %s to float16 ...", name)
op_block_list = ["RandomNormalLike"]
m.convert_float_to_float16(
keep_io_types=False,
op_block_list=op_block_list + force_fp32_operators[name],
)
# For SD-XL, use FP16 in VAE decoder will cause NaN and black image so we keep it in FP32.
if is_xl and name == "vae_decoder":
logger.info("Skip converting %s to float16 to avoid NaN", name)
else:
logger.info("Convert %s to float16 ...", name)
m.convert_float_to_float16(
keep_io_types=False,
op_block_list=force_fp32_operators[name],
)
if enable_runtime_optimization:
# Use this step to see the final graph that executed by Onnx Runtime.
@ -174,35 +181,24 @@ def optimize_sd_pipeline(
logger.info("*" * 20)
def copy_extra_directory(source_dir: Path, target_dir: Path, overwrite: bool):
def _copy_extra_directory(source_dir: Path, target_dir: Path):
"""Copy extra directory that does not have onnx model
Args:
source_dir (Path): source directory
target_dir (Path): target directory
overwrite (bool): overwrite if exists
Raises:
RuntimeError: source path does not exist
RuntimeError: output path exists but overwrite is false.
"""
extra_dirs = ["scheduler", "tokenizer", "feature_extractor"]
extra_dirs = ["scheduler", "tokenizer", "tokenizer_2", "feature_extractor"]
for name in extra_dirs:
source_path = source_dir / name
if not os.path.exists(source_path):
message = f"source path does not exist: {source_path}"
if name not in ["feature_extractor"]:
raise RuntimeError(message)
continue
target_path = target_dir / name
if target_path.exists():
if not overwrite:
raise RuntimeError(f"output path existed: {target_path}")
shutil.rmtree(target_path)
shutil.copytree(source_path, target_path)
logger.info("%s => %s", source_path, target_path)
@ -213,15 +209,54 @@ def copy_extra_directory(source_dir: Path, target_dir: Path, overwrite: bool):
raise RuntimeError(f"source path does not exist: {source_path}")
target_path = target_dir / name
if target_path.exists():
if not overwrite:
raise RuntimeError(f"output path existed: {target_path}")
os.remove(target_path)
shutil.copyfile(source_path, target_path)
logger.info("%s => %s", source_path, target_path)
# Some directory are optional
onnx_model_dirs = ["text_encoder", "text_encoder_2", "unet", "vae_encoder", "vae_decoder", "safety_checker"]
for onnx_model_dir in onnx_model_dirs:
source_path = source_dir / onnx_model_dir / "config.json"
target_path = target_dir / onnx_model_dir / "config.json"
if source_path.exists():
target_path.parent.mkdir(parents=True, exist_ok=True)
shutil.copyfile(source_path, target_path)
logger.info("%s => %s", source_path, target_path)
def parse_arguments():
def optimize_stable_diffusion_pipeline(
input_dir: str,
output_dir: str,
overwrite: bool,
use_external_data_format: Optional[bool],
float16: bool,
enable_runtime_optimization: bool,
args,
):
if os.path.exists(output_dir):
if args.overwrite:
shutil.rmtree(output_dir, ignore_errors=True)
else:
raise RuntimeError("output directory existed:{output_dir}. Add --overwrite to empty the directory.")
source_dir = Path(input_dir)
target_dir = Path(output_dir)
target_dir.mkdir(parents=True, exist_ok=True)
_copy_extra_directory(source_dir, target_dir)
_optimize_sd_pipeline(
source_dir,
target_dir,
overwrite,
use_external_data_format,
float16,
args.force_fp32_ops,
enable_runtime_optimization,
args,
)
def parse_arguments(argv: Optional[List[str]] = None):
"""Parse arguments
Returns:
@ -265,7 +300,8 @@ def parse_arguments():
"--inspect",
required=False,
action="store_true",
help="Inspect the optimized graph from Onnx Runtime for debugging purpose. This option has no impact on model performance.",
help="Save the optimized graph from Onnx Runtime. "
"This option has no impact on inference performance except it might reduce session creation time.",
)
parser.set_defaults(inspect=False)
@ -283,32 +319,25 @@ def parse_arguments():
required=False,
action="store_true",
help="Onnx model larger than 2GB need to use external data format. "
"Save onnx model to two files: one for onnx graph, another for large weights.",
"If specifed, save each onnx model to two files: one for onnx graph, another for weights. "
"If not specified, use same format as original model by default. ",
)
parser.set_defaults(use_external_data_format=False)
parser.set_defaults(use_external_data_format=None)
FusionOptions.add_arguments(parser)
args = parser.parse_args()
args = parser.parse_args(argv)
return args
def main():
coloredlogs.install(fmt="%(funcName)20s: %(message)s")
args = parse_arguments()
def main(argv: Optional[List[str]] = None):
args = parse_arguments(argv)
logger.info("Arguments: %s", str(args))
copy_extra_directory(Path(args.input), Path(args.output), args.overwrite)
optimize_sd_pipeline(
Path(args.input),
Path(args.output),
args.overwrite,
args.use_external_data_format,
args.float16,
args.force_fp32_ops,
args.inspect,
args,
optimize_stable_diffusion_pipeline(
args.input, args.output, args.overwrite, args.use_external_data_format, args.float16, args.inspect, args
)
if __name__ == "__main__":
coloredlogs.install(fmt="%(funcName)20s: %(message)s")
main()

8
onnxruntime/python/tools/transformers/models/stable_diffusion/requirements.txt
View file

@ -1,5 +1,5 @@
diffusers>=0.15.1
transformers>=4.26.0
diffusers>=0.19.3
transformers>=4.31.0
numpy>=1.24.1
accelerate
onnx>=1.13.0
@ -8,3 +8,7 @@ packaging
protobuf==3.20.3
psutil
sympy
# The following are for SDXL
optimum>=1.11.1
safetensors
invisible_watermark

115
onnxruntime/test/python/transformers/test_optimizer_stable_diffusion.py
View file

@ -7,6 +7,7 @@ import os
import shutil
import unittest
import numpy as np
import pytest
from parity_utilities import find_transformers_source
@ -19,6 +20,12 @@ else:
from onnxruntime.transformers.fusion_options import FusionOptions
from onnxruntime.transformers.optimizer import optimize_model
if find_transformers_source(["models", "stable_diffusion"]):
from optimize_pipeline import main as optimize_stable_diffusion
else:
from onnxruntime.transformers.models.stable_diffusion.optimize_pipeline import main as optimize_stable_diffusion
TINY_MODELS = {
"stable-diffusion": "hf-internal-testing/tiny-stable-diffusion-torch",
"stable-diffusion-xl": "echarlaix/tiny-random-stable-diffusion-xl",
@ -151,6 +158,114 @@ class TestStableDiffusionOptimization(unittest.TestCase):
},
)
@pytest.mark.slow
def test_optimize_sdxl_fp32(self):
save_directory = "tiny-random-stable-diffusion-xl"
if os.path.exists(save_directory):
shutil.rmtree(save_directory, ignore_errors=True)
model_type = "stable-diffusion-xl"
model_name = TINY_MODELS[model_type]
from optimum.onnxruntime import ORTStableDiffusionXLPipeline
baseline = ORTStableDiffusionXLPipeline.from_pretrained(model_name, export=True)
if not os.path.exists(save_directory):
baseline.save_pretrained(save_directory)
batch_size, num_images_per_prompt, height, width = 2, 2, 64, 64
latents = baseline.prepare_latents(
batch_size * num_images_per_prompt,
baseline.unet.config["in_channels"],
height,
width,
dtype=np.float32,
generator=np.random.RandomState(0),
)
optimized_directory = "tiny-random-stable-diffusion-xl-optimized"
argv = [
"--input",
save_directory,
"--output",
optimized_directory,
"--disable_group_norm",
"--disable_bias_splitgelu",
"--overwrite",
]
optimize_stable_diffusion(argv)
treatment = ORTStableDiffusionXLPipeline.from_pretrained(optimized_directory, provider="CUDAExecutionProvider")
inputs = {
"prompt": ["starry night by van gogh"] * batch_size,
"num_inference_steps": 3,
"num_images_per_prompt": num_images_per_prompt,
"height": height,
"width": width,
"guidance_rescale": 0.1,
"output_type": "np",
}
ort_outputs_1 = baseline(latents=latents, **inputs)
ort_outputs_2 = treatment(latents=latents, **inputs)
self.assertTrue(np.allclose(ort_outputs_1.images[0], ort_outputs_2.images[0], atol=1e-3))
@pytest.mark.slow
def test_optimize_sdxl_fp16(self):
"""This tests optimized fp16 pipeline, and result is deterministic for a given seed"""
save_directory = "tiny-random-stable-diffusion-xl"
if os.path.exists(save_directory):
shutil.rmtree(save_directory, ignore_errors=True)
model_type = "stable-diffusion-xl"
model_name = TINY_MODELS[model_type]
from optimum.onnxruntime import ORTStableDiffusionXLPipeline
baseline = ORTStableDiffusionXLPipeline.from_pretrained(model_name, export=True)
if not os.path.exists(save_directory):
baseline.save_pretrained(save_directory)
optimized_directory = "tiny-random-stable-diffusion-xl-optimized-fp16"
argv = [
"--input",
save_directory,
"--output",
optimized_directory,
"--disable_group_norm",
"--disable_bias_splitgelu",
"--float16",
"--overwrite",
]
optimize_stable_diffusion(argv)
fp16_pipeline = ORTStableDiffusionXLPipeline.from_pretrained(
optimized_directory, provider="CUDAExecutionProvider"
)
batch_size, num_images_per_prompt, height, width = 1, 1, 64, 64
inputs = {
"prompt": ["starry night by van gogh"] * batch_size,
"num_inference_steps": 3,
"num_images_per_prompt": num_images_per_prompt,
"height": height,
"width": width,
"guidance_rescale": 0.1,
"output_type": "latent",
}
seed = 123
np.random.seed(seed)
ort_outputs_1 = fp16_pipeline(**inputs)
np.random.seed(seed)
ort_outputs_2 = fp16_pipeline(**inputs)
np.random.seed(seed)
ort_outputs_3 = fp16_pipeline(**inputs)
self.assertTrue(np.array_equal(ort_outputs_1.images[0], ort_outputs_2.images[0]))
self.assertTrue(np.array_equal(ort_outputs_1.images[0], ort_outputs_3.images[0]))
if __name__ == "__main__":
unittest.main()