onnx.md 4.09 KB
Newer Older
Patrick von Platen's avatar
Patrick von Platen committed
1
<!--Copyright 2023 The HuggingFace Team. All rights reserved.
Nathan Lambert's avatar
Nathan Lambert committed
2
3
4
5
6
7
8
9
10
11
12

Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
the License. You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
specific language governing permissions and limitations under the License.
-->

Patrick von Platen's avatar
Patrick von Platen committed
13

Ella Charlaix's avatar
Ella Charlaix committed
14
# How to use ONNX Runtime for inference
Patrick von Platen's avatar
Patrick von Platen committed
15

16
🤗 [Optimum](https://github.com/huggingface/optimum) provides a Stable Diffusion pipeline compatible with ONNX Runtime. 
Patrick von Platen's avatar
Patrick von Platen committed
17

18
## Installation
Patrick von Platen's avatar
Patrick von Platen committed
19

20
21
22
23
24
Install 🤗 Optimum with the following command for ONNX Runtime support:

```
pip install optimum["onnxruntime"]
```
Patrick von Platen's avatar
Patrick von Platen committed
25

26
## Stable Diffusion
Patrick von Platen's avatar
Patrick von Platen committed
27

28
29
### Inference

Ella Charlaix's avatar
Ella Charlaix committed
30
To load an ONNX model and run inference with ONNX Runtime, you need to replace [`StableDiffusionPipeline`] with `ORTStableDiffusionPipeline`. In case you want to load a PyTorch model and convert it to the ONNX format on-the-fly, you can set `export=True`.
Patrick von Platen's avatar
Patrick von Platen committed
31

32
```python
33
from optimum.onnxruntime import ORTStableDiffusionPipeline
34

35
model_id = "runwayml/stable-diffusion-v1-5"
36
37
38
39
pipeline = ORTStableDiffusionPipeline.from_pretrained(model_id, export=True)
prompt = "sailing ship in storm by Leonardo da Vinci"
image = pipeline(prompt).images[0]
pipeline.save_pretrained("./onnx-stable-diffusion-v1-5")
40
```
Patrick von Platen's avatar
Patrick von Platen committed
41

42
43
If you want to export the pipeline in the ONNX format offline and later use it for inference,
you can use the [`optimum-cli export`](https://huggingface.co/docs/optimum/main/en/exporters/onnx/usage_guides/export_a_model#exporting-a-model-to-onnx-using-the-cli) command: 
44

45
46
47
48
49
50
51
52
```bash
optimum-cli export onnx --model runwayml/stable-diffusion-v1-5 sd_v15_onnx/
```

Then perform inference:

```python 
from optimum.onnxruntime import ORTStableDiffusionPipeline
53

54
model_id = "sd_v15_onnx"
55
56
57
pipeline = ORTStableDiffusionPipeline.from_pretrained(model_id)
prompt = "sailing ship in storm by Leonardo da Vinci"
image = pipeline(prompt).images[0]
58
59
```

60
61
Notice that we didn't have to specify `export=True` above.

62
63
64
65
<div class="flex justify-center">
    <img src="https://huggingface.co/datasets/optimum/documentation-images/resolve/main/onnxruntime/stable_diffusion_v1_5_ort_sail_boat.png">
</div>

66
67
You can find more examples in [optimum documentation](https://huggingface.co/docs/optimum/).

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

### Supported tasks

| Task                                 | Loading Class                        |
|--------------------------------------|--------------------------------------|
| `text-to-image`                      | `ORTStableDiffusionPipeline`         |
| `image-to-image`                     | `ORTStableDiffusionImg2ImgPipeline`  |
| `inpaint`                            | `ORTStableDiffusionInpaintPipeline`  |

## Stable Diffusion XL

### Export

To export your model to ONNX, you can use the [Optimum CLI](https://huggingface.co/docs/optimum/main/en/exporters/onnx/usage_guides/export_a_model#exporting-a-model-to-onnx-using-the-cli) as follows :

```bash
optimum-cli export onnx --model stabilityai/stable-diffusion-xl-base-1.0 --task stable-diffusion-xl sd_xl_onnx/
```

### Inference

To load an ONNX model and run inference with ONNX Runtime, you need to replace `StableDiffusionPipelineXL` with `ORTStableDiffusionPipelineXL` :

```python
from optimum.onnxruntime import ORTStableDiffusionXLPipeline

pipeline = ORTStableDiffusionXLPipeline.from_pretrained("sd_xl_onnx")
prompt = "sailing ship in storm by Leonardo da Vinci"
image = pipeline(prompt).images[0]
```

### Supported tasks

| Task                                 | Loading Class                        |
|--------------------------------------|--------------------------------------|
| `text-to-image`                      | `ORTStableDiffusionXLPipeline`       |
| `image-to-image`                     | `ORTStableDiffusionXLImg2ImgPipeline`|

106
## Known Issues
Patrick von Platen's avatar
Patrick von Platen committed
107

108
- Generating multiple prompts in a batch seems to take too much memory. While we look into it, you may need to iterate instead of batching.