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

[ROCm] Update benchmark for stable diffusion (#15602)

Browse source 
1. update scripts for ROCm memory measurement.
2. update README to contain ROCm result.
3. address some minor issue in the README
This commit is contained in:
cloudhan 2023-04-23 11:49:40 +08:00 committed by GitHub
parent 9e44248bf9
commit 8297148bde
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 231 additions and 122 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -1,10 +1,10 @@
# Stable Diffusion CUDA Optimization
# Stable Diffusion GPU Optimization
## Overview
[Stable Diffusion](https://stability.ai/blog/stable-diffusion-announcement) is a text-to-image latent diffusion model for image generation. Explanation of the Stable Diffusion can be found in [Stable Diffusion with Diffusers](https://huggingface.co/blog/stable_diffusion).
## CUDA Optimizations for Stable Diffusion
## Optimizations for Stable Diffusion
ONNX Runtime uses the following optimizations to speed up Stable Diffusion in CUDA:
@ -17,7 +17,9 @@ ONNX Runtime uses the following optimizations to speed up Stable Diffusion in CU
* BiasAdd fuses Add bias and residual.
* Reduce Transpose nodes by graph transformation.
To show the impact of each optimization on latency and GPU memory, we did an experiment on RTX 3060 GPU:
These optimizations are firstly carried out on CUDA EP. They may not work on other EP. To show the impact of each optimization on latency and GPU memory, we did some experiments:
### Results on RTX 3060 GPU:
| Optimizations | Average Latency (batch_size=1) | Memory in MB (batch_size=1) | Average Latency (batch_size=8) | Memory in MB (batch_size=8) |
| ---------------------------------------------------------------------------------- | ------------------------------ | --------------------------- | ------------------------------ | --------------------------- |
@ -36,12 +38,28 @@ Here FMHA means Attention and MultiHeadAttention operators with Flash Attention
The last two optimizations (Packed QKV and BiasAdd) are only available in nightly package. Compared to 1.14.1, nightly package has slight improvement in performance.
### Results on MI250X with 1 GCD
With runtime tuning enabled, we get following performance number on one GCD of a MI250X GPU:
| Optimizations | Average Latency (batch_size=1) | Memory in MB (batch_size=1) | Average Latency (batch_size=8) | Memory in MB (batch_size=8) |
| --------------------------------------------------------------------- | ------------------------------ | --------------------------- | ------------------------------ | --------------------------- |
| Raw FP32 models | 6.7 | 17,319 | 36.4 * | 33,787 |
| FP16 baseline | 4.1 | 8,945 | 24.0 * | 34,493 |
| FP16 baseline + FMHA | 2.6 | 4,886 | 15.0 | 10,146 |
| FP16 baseline + FMHA + NhwcConv | 2.4 | 4,952 | 14.8 | 9,632 |
| FP16 baseline + FMHA + NhwcConv + GroupNorm | 2.3 | 4,906 | 13.6 | 9,774 |
| FP16 baseline + FMHA + NhwcConv + GroupNorm + BiasSplitGelu | 2.2 | 4,910 | 12.5 | 9,646 |
| FP16 baseline + FMHA + NhwcConv + GroupNorm + BiasSplitGelu + BiasAdd | 2.2 | 4,910 | 12.5 | 9,778 |
The entries marked with `*` produce suspicious output images. The might be numerical stability or correctness issue for the pipeline. The performance number is for reference only.
## Scripts:
| Script | Description |
| ---------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------- |
| [optimize_pipeline.py](https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/python/tools/transformers/models/stable_diffusion/optimize_pipeline.py) | Optimize Stable Diffusion ONNX models |
| [benchmark.py](https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/python/tools/transformers/models/stable_diffusion/benchmark.py) | Benchmark latency and memory of OnnxRuntime, xFormers or PyTorch 2.0 on stable diffusion. |
| Script | Description |
| ---------------------------------------------- | ----------------------------------------------------------------------------------------- |
| [optimize_pipeline.py](./optimize_pipeline.py) | Optimize Stable Diffusion ONNX models |
| [benchmark.py](./benchmark.py) | Benchmark latency and memory of OnnxRuntime, xFormers or PyTorch 2.0 on stable diffusion. |
In below example, we run the scripts in source code directory. You can get source code like the following:
@ -54,17 +72,16 @@ cd onnxruntime/python/tools/transformers/models/stable_diffusion
Below is an example to optimize Stable Diffusion 1.5 in Linux. For Windows OS, please change the format of path to be like `.\sd-v1-5` instead of `./sd-v1-5`.
### Setup Environment
First, Let's create a python environment using [AnaConda](https://www.anaconda.com/products/distribution#Downloads), then install packages in [requirements.txt](https://raw.githubusercontent.com/microsoft/onnxruntime/main/onnxruntime/python/tools/transformers/models/stable_diffusion/requirements.txt):
### Setup Environment (CUDA)
```
conda create -n py310 python=3.10
conda activate py310
pip install -r requirements.txt
pip install -r requirements-cuda.txt
```
For Windows, torch installed from pypi is CPU only. Need install PyTorch 1.13.1+cu117 instead like the following to support GPU:
For Windows, the torch package installed from PyPI is CPU only. To enable support for GPU, it is necessary to install PyTorch version 1.13.1+cu117 or above using the following method:
```
pip install torch==1.13.1+cu117 --extra-index-url https://download.pytorch.org/whl/cu117
```
@ -82,19 +99,32 @@ pip uninstall onnxruntime-gpu
pip install ort-nightly-gpu -i https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ORT-Nightly/pypi/simple/
```
Need install diffusers from source until v0.15 is released.
### Setup Environment (ROCm)
It is recommended that the users should run the model with ROCm 5.4 or newer and Python 3.9. Note that Windows is not
supported for ROCm at the moment.
```
git clone https://github.com/huggingface/diffusers
cd diffusers
pip install -e .
conda create -n py39 python=3.9
conda activate py39
pip install -r requirements.txt
pip install -r requirements-rocm.txt
```
AMD GPU version of torch build can be installed from [AMD Radeon repo](https://repo.radeon.com/rocm/manylinux/rocm-rel-5.4/),
user need to download the whl file and `pip install <file.whl>` manually. Or directly from https://download.pytorch.org/whl/rocm5.4.2 via
```
pip install torch==2.0.0 --index-url https://download.pytorch.org/whl/rocm5.4.2
```
Please follow the [official docs](https://onnxruntime.ai/docs/build/eps.html#amd-rocm) to build ONNXRuntime from source.
### Export ONNX pipeline
This step will export stable diffusion 1.5 to ONNX model in float32 using script from diffusers. Before running the script, you need to be logged in via `huggingface-cli login`.
```
curl https://raw.githubusercontent.com/huggingface/diffusers/v0.13.0/scripts/convert_stable_diffusion_checkpoint_to_onnx.py > convert_sd_onnx.py
curl https://raw.githubusercontent.com/huggingface/diffusers/v0.15.1/scripts/convert_stable_diffusion_checkpoint_to_onnx.py > convert_sd_onnx.py
python convert_sd_onnx.py --model_path runwayml/stable-diffusion-v1-5 --output_path ./sd-v1-5
```
@ -103,12 +133,13 @@ python convert_sd_onnx.py --model_path runwayml/stable-diffusion-v1-5 --output_
Example to optimize the exported float32 ONNX models, and save to float16 models:
```
python optimize_pipeline.py -i ./sd-v1-5 -o ./sd-v1-5-fp16 --float16
python -m onnxruntime.transformers.models.stable_diffusion.optimize_pipeline -i ./sd-v1-5 -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 Stable Diffusion 2.1 model, you will need force Attention to run in float32 to avoid black image by appending `--force_fp32_ops unet:Attention` to the command line. If you are using nightly package, append `--force_fp32_ops unet:MultiHeadAttention` instead.
For Stable Diffusion 2.1 model with CUDA EP, you will need force Attention to run in float32 to avoid black image by appending `--force_fp32_ops unet:Attention` to the command line.
If you are using nightly package, append `--force_fp32_ops unet:MultiHeadAttention` instead.
### Run Benchmark
@ -116,11 +147,21 @@ The benchmark.py script will run a warm-up prompt twice, and measure the peak GP
Note that the first run might need more time and memory: For example, cuDNN convolution algorithm search or model compile happens in the first run.
Example to benchmark the optimized pipeline with batch size 1:
Example to benchmark the optimized pipeline with batch size 1 on CUDA EP:
```
python benchmark.py -p ./sd-v1-5-fp16/ -b 1
python -m onnxruntime.transformers.models.stable_diffusion.benchmark -p ./sd-v1-5-fp16/ -b 1
```
On ROCm EP, use the following command instead:
```
python -m onnxruntime.transformers.models.stable_diffusion.benchmark -p ./sd-v1-5-fp16/ -b 1 --tuning --provider=rocm
```
Note: you can substitute `python -m onnxruntime.transformers.models.stable_diffusion.benchmark` with `python benchmark.py` if your current working directory is this files directory.
In the following, we will use it interchangeably.
For ROCm EP, the `--tuning` is mandatory because we heavily rely on tuning to find the runable kernels for ORT `OpKernel`s.
The default parameters are stable diffusion version=1.5, height=512, width=512, steps=50, batch_count=5. Run `python benchmark.py --help` for more information.
### Run Benchmark with xFormers
@ -140,6 +181,7 @@ conda create -n pt2 python=3.10
conda activate pt2
pip install torch --index-url https://download.pytorch.org/whl/cu117
pip install -r requirements.txt
pip install -r requirements_cuda.txt # or requirements_rocm.txt
python benchmark.py -e torch -b 1 --enable_torch_compile
```
@ -155,7 +197,7 @@ Common settings for below test results:
| ------------------------------ | ---------------------- | ------ | ----- | ----- | ----------- | ----------- |
| runwayml/stable-diffusion-v1-5 | TRUE | 512 | 512 | 50 | 5 | 1 |
#### Results of RTX 3060 (in Windows 11)
#### Results of RTX 3060 (Windows 11)
| engine | version | provider | batch size | average latency | first run memory MB | second run memory MB |
| ----------- | ----------------------- | --------------------- | ---------- | --------------- | ------------------- | -------------------- |
@ -170,7 +212,7 @@ Common settings for below test results:
| torch | 1.13.1+cu117 | xformers | 8 | 41.1 | 10,825 | 9,255 |
#### Results of A100-SXM4-40GB (in Ubuntu 20.04)
#### Results of A100-SXM4-40GB (Ubuntu 20.04)
| engine | version | provider | batch size | average latency | first run memory MB | second run memory MB |
| ----------- | ----------------------- | --------------------- | ---------- | --------------- | ------------------- | -------------------- |
| onnxruntime | 1.14.1 | CUDAExecutionProvider | 1 | 1.1 | 6,883 | 7,395 |
@ -186,7 +228,7 @@ Common settings for below test results:
| torch | 2.0.0+cu117 | default | 16 | 14.8 | 32,306 | 16,520 |
| torch | 2.0.0+cu117 | compile | 16 | 12.6 | 32,636 | 16,898 |
#### Results of V100-PCIE-16GB (in Ubuntu 20.04)
#### Results of V100-PCIE-16GB (Ubuntu 20.04)
Results from Standard_NC6s_v3 Azure virtual machine:
@ -205,10 +247,9 @@ Results from Standard_NC6s_v3 Azure virtual machine:
| torch | 2.0.0+cu117 | default | 8 | 17.0 | 14,469 | 9,657 |
| torch | 1.13.1+cu117 | xformers | 8 | 17.4 | 15,593 | 9,133 |
#### Results of T4 (Ubuntu 20.04)
#### Results of T4 (in Ubuntu 20.04)
To make the result stable, we lock the frequency of T4 GPU like
To make the result stable, we lock the frequency of T4 GPU like
`sudo nvidia-smi --lock-gpu-clocks=990` for fair comparison. See [nvidia blog](https://developer.nvidia.com/blog/advanced-api-performance-setstablepowerstate/) for more information. Note that performance might be slightly better without locking frequency.
Results are from Standard_NC4as_T4_v3 Azure virtual machine:
@ -228,12 +269,32 @@ Results are from Standard_NC4as_T4_v3 Azure virtual machine:
| torch | 2.0.0+cu117 | compile | 8 | 46.5 | 12,595 | 10,171 |
| torch | 2.0.0+cu117 | default | 8 | 50.7 | 11,955 | 9,531 |
#### Results of MI250X, 1 GCD (Ubuntu 20.04)
| engine | version | provider | batch size | average latency | first run memory MB | second run memory MB |
| ----------- | ----------------------- | --------------------- | ---------- | --------------- | ------------------- | -------------------- |
| onnxruntime | dev | ROCMExecutionProvider | 1 | 2.2 | 5,548 | 4,908 |
| torch | 1.12.1+rocm5.4 | - | 1 | 3.4 | 6,653 | 4,613 |
| torch | 2.0.0+rocm5.4.2 | default | 1 | 3.2 | 5,977 | 4,368 |
| torch | 2.0.0a0+git254afb8 | compile | 1 | 3.0 | 5,869 | 4,266 |
| onnxruntime | dev | ROCMExecutionProvider | 4 | 6.6 | 5,546 | 4,906 |
| torch | 1.12.1+rocm5.4 | - | 4 | 10.1 | 19,477 | 11,325 |
| torch | 2.0.0+rocm5.4.2 | default | 4 | 10.5 | 13,051 | 7,300 |
| torch | 2.0.0+rocm5.4.2 | compile | 4 | 9.2 | 12,879 | 7,190 |
| onnxruntime | dev | ROCMExecutionProvider | 8 | 12.5 | 9,778 | 9,006 |
| torch | 1.12.1+rocm5.4 | - | 8 | 19.3 | 55,851 | 20,014 |
| torch | 2.0.0+rocm5.4.2 | default | 8 | 20.3 | 23,551 | 11,930 |
| torch | 2.0.0+rocm5.4.2 | compile | 8 | 17.8 | 23,303 | 11,800 |
### Credits
Some CUDA kernels (Flash Attention, GroupNorm, SplitGelu and BiasAdd etc.) were originally implemented in [TensorRT](https://github.com/nviDIA/TensorRT) by Nvidia.
We use Memory efficient attention from [CUTLASS](https://github.com/NVIDIA/cutlass). The kernels were developed by Meta xFormers.
The ONNX export script and pipeline for stable diffusion was developed by Huggingface [diffusers](https://github.com/huggingface/diffusers) library.
Most ROCm kernel optimizations are from [composable kernel](https://github.com/ROCmSoftwarePlatform/composable_kernel).
Some kernels are enabled by MIOpen. We hereby thank for the AMD developers' collaboration.
### Future Works
There are other optimizations might improve the performance or reduce memory footprint:

199
onnxruntime/python/tools/transformers/models/stable_diffusion/benchmark.py
View file

@ -39,57 +39,101 @@ def example_prompts():
return prompts
def measure_gpu_memory(func, start_memory=None):
class MemoryMonitor:
def __init__(self, keep_measuring=True):
self.keep_measuring = keep_measuring
class CudaMemoryMonitor:
def __init__(self, keep_measuring=True):
self.keep_measuring = keep_measuring
def measure_gpu_usage(self):
from py3nvml.py3nvml import (
NVMLError,
nvmlDeviceGetCount,
nvmlDeviceGetHandleByIndex,
nvmlDeviceGetMemoryInfo,
nvmlDeviceGetName,
nvmlInit,
nvmlShutdown,
)
def measure_gpu_usage(self):
from py3nvml.py3nvml import (
NVMLError,
nvmlDeviceGetCount,
nvmlDeviceGetHandleByIndex,
nvmlDeviceGetMemoryInfo,
nvmlDeviceGetName,
nvmlInit,
nvmlShutdown,
)
max_gpu_usage = []
gpu_name = []
try:
nvmlInit()
device_count = nvmlDeviceGetCount()
if not isinstance(device_count, int):
print(f"nvmlDeviceGetCount result is not integer: {device_count}")
return None
max_gpu_usage = [0 for i in range(device_count)]
gpu_name = [nvmlDeviceGetName(nvmlDeviceGetHandleByIndex(i)) for i in range(device_count)]
while True:
for i in range(device_count):
info = nvmlDeviceGetMemoryInfo(nvmlDeviceGetHandleByIndex(i))
if isinstance(info, str):
print(f"nvmlDeviceGetMemoryInfo returns str: {info}")
return None
max_gpu_usage[i] = max(max_gpu_usage[i], info.used / 1024**2)
time.sleep(0.002) # 2ms
if not self.keep_measuring:
break
nvmlShutdown()
return [
{
"device_id": i,
"name": gpu_name[i],
"max_used_MB": max_gpu_usage[i],
}
for i in range(device_count)
]
except NVMLError as error:
print("Error fetching GPU information using nvml: %s", error)
max_gpu_usage = []
gpu_name = []
try:
nvmlInit()
device_count = nvmlDeviceGetCount()
if not isinstance(device_count, int):
print(f"nvmlDeviceGetCount result is not integer: {device_count}")
return None
monitor = MemoryMonitor(False)
max_gpu_usage = [0 for i in range(device_count)]
gpu_name = [nvmlDeviceGetName(nvmlDeviceGetHandleByIndex(i)) for i in range(device_count)]
while True:
for i in range(device_count):
info = nvmlDeviceGetMemoryInfo(nvmlDeviceGetHandleByIndex(i))
if isinstance(info, str):
print(f"nvmlDeviceGetMemoryInfo returns str: {info}")
return None
max_gpu_usage[i] = max(max_gpu_usage[i], info.used / 1024**2)
time.sleep(0.002) # 2ms
if not self.keep_measuring:
break
nvmlShutdown()
return [
{
"device_id": i,
"name": gpu_name[i],
"max_used_MB": max_gpu_usage[i],
}
for i in range(device_count)
]
except NVMLError as error:
print("Error fetching GPU information using nvml: %s", error)
return None
class RocmMemoryMonitor:
def __init__(self, keep_measuring=True):
self.keep_measuring = keep_measuring
rocm_smi_path = "/opt/rocm/libexec/rocm_smi"
if os.path.exists(rocm_smi_path):
if rocm_smi_path not in sys.path:
sys.path.append(rocm_smi_path)
try:
import rocm_smi
self.rocm_smi = rocm_smi
self.rocm_smi.initializeRsmi()
except ImportError:
self.rocm_smi = None
def get_used_memory(self, dev):
if self.rocm_smi is None:
return -1
return self.rocm_smi.getMemInfo(dev, "VRAM")[0] / 1024 / 1024
def measure_gpu_usage(self):
device_count = len(self.rocm_smi.listDevices()) if self.rocm_smi is not None else 0
max_gpu_usage = [0 for i in range(device_count)]
gpu_name = [f"GPU{i}" for i in range(device_count)]
while True:
for i in range(device_count):
max_gpu_usage[i] = max(max_gpu_usage[i], self.get_used_memory(i))
time.sleep(0.002) # 2ms
if not self.keep_measuring:
break
return [
{
"device_id": i,
"name": gpu_name[i],
"max_used_MB": max_gpu_usage[i],
}
for i in range(device_count)
]
def measure_gpu_memory(monitor_type, func, start_memory=None):
if monitor_type is None:
return None
monitor = monitor_type(False)
memory_before_test = monitor.measure_gpu_usage()
if start_memory is None:
@ -102,7 +146,7 @@ def measure_gpu_memory(func, start_memory=None):
from concurrent.futures import ThreadPoolExecutor
with ThreadPoolExecutor() as executor:
monitor = MemoryMonitor()
monitor = monitor_type()
mem_thread = executor.submit(monitor.measure_gpu_usage)
try:
fn_thread = executor.submit(func)
@ -127,7 +171,7 @@ def measure_gpu_memory(func, start_memory=None):
return None
def get_ort_pipeline(model_name: str, directory: str, provider: str, disable_safety_checker: bool):
def get_ort_pipeline(model_name: str, directory: str, provider, disable_safety_checker: bool):
from diffusers import DPMSolverMultistepScheduler, OnnxStableDiffusionPipeline
import onnxruntime
@ -201,7 +245,7 @@ def run_ort_pipeline(
num_prompts,
batch_count,
start_memory,
enable_mem_measure,
memory_monitor_type,
):
from diffusers import OnnxStableDiffusionPipeline
@ -212,11 +256,12 @@ def run_ort_pipeline(
def warmup():
pipe("warm up", height, width, num_inference_steps=steps, num_images_per_prompt=batch_size)
# Run warm up, and measure GPU memory of two runs (The first run has cuDNN algo search so it might need more memory)
first_run_memory = measure_gpu_memory(warmup, start_memory) if enable_mem_measure else -1
second_run_memory = measure_gpu_memory(warmup, start_memory) if enable_mem_measure else -1
# Run warm up, and measure GPU memory of two runs
# cuDNN/MIOpen The first run has algo search so it might need more memory)
first_run_memory = measure_gpu_memory(memory_monitor_type, warmup, start_memory)
second_run_memory = measure_gpu_memory(memory_monitor_type, warmup, start_memory)
if not enable_mem_measure:
if memory_monitor_type is None:
warmup()
latency_list = []
@ -269,7 +314,7 @@ def run_torch_pipeline(
num_prompts,
batch_count,
start_memory,
enable_mem_measure,
memory_monitor_type,
):
import torch
@ -280,10 +325,10 @@ def run_torch_pipeline(
pipe("warm up", height, width, num_inference_steps=steps, num_images_per_prompt=batch_size)
# Run warm up, and measure GPU memory of two runs (The first run has cuDNN algo search so it might need more memory)
first_run_memory = measure_gpu_memory(warmup, start_memory) if enable_mem_measure else -1
second_run_memory = measure_gpu_memory(warmup, start_memory) if enable_mem_measure else -1
first_run_memory = measure_gpu_memory(memory_monitor_type, warmup, start_memory)
second_run_memory = measure_gpu_memory(memory_monitor_type, warmup, start_memory)
if not enable_mem_measure:
if memory_monitor_type is None:
warmup()
torch.set_grad_enabled(False)
@ -342,10 +387,15 @@ def run_ort(
num_prompts,
batch_count,
start_memory,
enable_mem_measure,
memory_monitor_type,
tuning,
):
provider_and_options = provider
if tuning and provider in ["CUDAExecutionProvider", "ROCMExecutionProvider"]:
provider_and_options = (provider, {"tunable_op_enable": 1, "tunable_op_tuning_enable": 1})
load_start = time.time()
pipe = get_ort_pipeline(model_name, directory, provider, disable_safety_checker)
pipe = get_ort_pipeline(model_name, directory, provider_and_options, disable_safety_checker)
load_end = time.time()
print(f"Model loading took {load_end - load_start} seconds")
@ -360,7 +410,7 @@ def run_ort(
num_prompts,
batch_count,
start_memory,
enable_mem_measure,
memory_monitor_type,
)
result.update(
@ -386,7 +436,7 @@ def run_torch(
num_prompts,
batch_count,
start_memory,
enable_mem_measure,
memory_monitor_type,
):
import torch
@ -414,7 +464,7 @@ def run_torch(
num_prompts,
batch_count,
start_memory,
enable_mem_measure,
memory_monitor_type,
)
else:
result = run_torch_pipeline(
@ -427,7 +477,7 @@ def run_torch(
num_prompts,
batch_count,
start_memory,
enable_mem_measure,
memory_monitor_type,
)
result.update(
@ -464,6 +514,14 @@ def parse_arguments():
help="Provider to benchmark. Default is CUDAExecutionProvider.",
)
parser.add_argument(
"-t",
"--tuning",
action="store_true",
help="Enable TunableOp and tuning. "
"This will incur longer warmup latency, and is mandatory for some operators of ROCm EP.",
)
parser.add_argument(
"-v",
"--version",
@ -568,9 +626,13 @@ def main():
args = parse_arguments()
print(args)
enable_mem_measure = args.provider == "cuda"
memory_monitor_type = None
if args.provider == "cuda":
memory_monitor_type = CudaMemoryMonitor
elif args.provider == "rocm":
memory_monitor_type = RocmMemoryMonitor
start_memory = measure_gpu_memory(None) if enable_mem_measure else -1
start_memory = measure_gpu_memory(memory_monitor_type, None)
print("GPU memory used before loading models:", start_memory)
sd_model = SD_MODELS[args.version]
@ -590,7 +652,8 @@ def main():
args.num_prompts,
args.batch_count,
start_memory,
enable_mem_measure,
memory_monitor_type,
args.tuning,
)
else:
result = run_torch(
@ -605,7 +668,7 @@ def main():
args.num_prompts,
args.batch_count,
start_memory,
enable_mem_measure,
memory_monitor_type,
)
print(result)

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

@ -10,16 +10,14 @@
# 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 optimized models are for CUDA Execution Provider. It might not run in other execution provider.
# 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
#
# ROCm EP doesn't support MultiHeadAttention, add --disable_attention to disable attention fusion:
# python optimize_pipeline.py -i ./sd-v1-5 -o ./sd-v1-5-fp16 --float16 --disable_attention
import argparse
import logging

7
onnxruntime/python/tools/transformers/models/stable_diffusion/requirements-cuda.txt Normal file
View file

@ -0,0 +1,7 @@
# Install the following package in python 3.10
onnxruntime-gpu>=1.14
py3nvml==0.2.7
xformers==0.0.16
#For Windows, need install PyTorch 1.13.1+cu117 since torch in pypi is CPU version
#--extra-index-url https://download.pytorch.org/whl/cu117
#torch==1.13.1+cu117

19
onnxruntime/python/tools/transformers/models/stable_diffusion/requirements-rocm.txt
View file

@ -1,17 +1,4 @@
transformers==4.26.0
numpy==1.24.1
accelerate==0.15.0
onnx==1.13.0
coloredlogs
packaging==23.0
protobuf==3.20.3
psutil==5.9.4
sympy==1.11.1
# Install diffusers from source
# git clone https://github.com/huggingface/diffusers.git
# cd diffusers && git checkout c4892f1855097a68703ca2e949aca15829526958
# pip install -e .
# Install onnxruntime-rocm or onnxruntime_training
# Build onnxruntime-rocm from source or install lastest onnxruntime_training rocm nightly python package
# Build onnxruntime-rocm from source
# Directly install pre-built onnxruntime/onnxruntime-training rocm python package is not possible at the moment.
# TODO: update once we have public pre-built packages

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

@ -1,17 +1,10 @@
# Install the following package in python 3.10
diffusers==0.13.0
diffusers==0.15.1
transformers==4.26.0
numpy==1.24.1
accelerate==0.15.0
onnxruntime-gpu>=1.14
onnx==1.13.0
coloredlogs
packaging==23.0
protobuf==3.20.3
psutil==5.9.4
sympy==1.11.1
py3nvml==0.2.7
xformers==0.0.16
#For Windows, need install PyTorch 1.13.1+cu117 since torch in pypi is CPU version
#--extra-index-url https://download.pytorch.org/whl/cu117
#torch==1.13.1+cu117