From af80542e65334979d3e37c1cb6c88068c5c65f3b Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Fri, 15 Sep 2023 10:17:20 -0700 Subject: [PATCH] Update optimize_pipeline for SDXL (#17536) - [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](https://github.com/huggingface/optimum/blob/afd2b5a36663bebd1f501486acee065c728947bc/optimum/pipelines/diffusers/pipeline_stable_diffusion_xl.py#L483) ``` latents = latents.astype(np.float32) ``` --- .../models/stable_diffusion/README.md | 11 +- .../stable_diffusion/optimize_pipeline.py | 161 +++++++++++------- .../models/stable_diffusion/requirements.txt | 8 +- .../test_optimizer_stable_diffusion.py | 115 +++++++++++++ 4 files changed, 226 insertions(+), 69 deletions(-) diff --git a/onnxruntime/python/tools/transformers/models/stable_diffusion/README.md b/onnxruntime/python/tools/transformers/models/stable_diffusion/README.md index facbd3bf69..7ffefdd05f 100644 --- a/onnxruntime/python/tools/transformers/models/stable_diffusion/README.md +++ b/onnxruntime/python/tools/transformers/models/stable_diffusion/README.md @@ -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 diff --git a/onnxruntime/python/tools/transformers/models/stable_diffusion/optimize_pipeline.py b/onnxruntime/python/tools/transformers/models/stable_diffusion/optimize_pipeline.py index 22fee4bfea..4512c971ac 100644 --- a/onnxruntime/python/tools/transformers/models/stable_diffusion/optimize_pipeline.py +++ b/onnxruntime/python/tools/transformers/models/stable_diffusion/optimize_pipeline.py @@ -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() diff --git a/onnxruntime/python/tools/transformers/models/stable_diffusion/requirements.txt b/onnxruntime/python/tools/transformers/models/stable_diffusion/requirements.txt index 68947a1618..d4e6c9fa07 100644 --- a/onnxruntime/python/tools/transformers/models/stable_diffusion/requirements.txt +++ b/onnxruntime/python/tools/transformers/models/stable_diffusion/requirements.txt @@ -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 diff --git a/onnxruntime/test/python/transformers/test_optimizer_stable_diffusion.py b/onnxruntime/test/python/transformers/test_optimizer_stable_diffusion.py index cde6b56a66..dca250f39f 100644 --- a/onnxruntime/test/python/transformers/test_optimizer_stable_diffusion.py +++ b/onnxruntime/test/python/transformers/test_optimizer_stable_diffusion.py @@ -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()