Update optimize_pipeline for SDXL

onnxruntime/python/tools/transformers/models/stable_diffusion/README.md

@@ -154,14 +154,23 @@ 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:
 ```
 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
 

onnxruntime/python/tools/transformers/models/stable_diffusion/optimize_pipeline.py

@@ -5,28 +5,25 @@
 #
 # 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
 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 @@
 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()

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