Skip to content

Commit

Permalink
[docs] Update ONNX doc to use optimum (open-mmlab#2702)
Browse files Browse the repository at this point in the history
* minor edits to onnx and openvino docs.

* Apply suggestions from code review

Co-authored-by: Ella Charlaix <80481427+echarlaix@users.noreply.github.com>

---------

Co-authored-by: Ella Charlaix <80481427+echarlaix@users.noreply.github.com>
  • Loading branch information
sayakpaul and echarlaix committed Mar 17, 2023
1 parent f4bbcb2 commit a169571
Show file tree
Hide file tree
Showing 2 changed files with 31 additions and 39 deletions.
68 changes: 30 additions & 38 deletions docs/source/en/optimization/onnx.mdx
Original file line number Diff line number Diff line change
Expand Up @@ -13,61 +13,53 @@ specific language governing permissions and limitations under the License.

# How to use the ONNX Runtime for inference

🤗 Diffusers provides a Stable Diffusion pipeline compatible with the ONNX Runtime. This allows you to run Stable Diffusion on any hardware that supports ONNX (including CPUs), and where an accelerated version of PyTorch is not available.
🤗 [Optimum](https://github.com/huggingface/optimum) provides a Stable Diffusion pipeline compatible with ONNX Runtime.

## Installation

- TODO
Install 🤗 Optimum with the following command for ONNX Runtime support:

```
pip install optimum["onnxruntime"]
```

## Stable Diffusion Inference

The snippet below demonstrates how to use the ONNX runtime. You need to use `OnnxStableDiffusionPipeline` instead of `StableDiffusionPipeline`. You also need to download the weights from the `onnx` branch of the repository, and indicate the runtime provider you want to use.
To load an ONNX model and run inference with the 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`.

```python
# make sure you're logged in with `huggingface-cli login`
from diffusers import OnnxStableDiffusionPipeline

pipe = OnnxStableDiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5",
revision="onnx",
provider="CUDAExecutionProvider",
)
from optimum.onnxruntime import ORTStableDiffusionPipeline

model_id = "runwayml/stable-diffusion-v1-5"
pipe = ORTStableDiffusionPipeline.from_pretrained(model_id, export=True)
prompt = "a photo of an astronaut riding a horse on mars"
image = pipe(prompt).images[0]
images = pipe(prompt).images[0]
pipe.save_pretrained("./onnx-stable-diffusion-v1-5")
```

The snippet below demonstrates how to use the ONNX runtime with the Stable Diffusion upscaling pipeline.
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:

```python
from diffusers import OnnxStableDiffusionPipeline, OnnxStableDiffusionUpscalePipeline
```bash
optimum-cli export onnx --model runwayml/stable-diffusion-v1-5 sd_v15_onnx/
```

Then perform inference:

```python
from optimum.onnxruntime import ORTStableDiffusionPipeline

model_id = "sd_v15_onnx"
pipe = ORTStableDiffusionPipeline.from_pretrained(model_id)
prompt = "a photo of an astronaut riding a horse on mars"
steps = 50

txt2img = OnnxStableDiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5",
revision="onnx",
provider="CUDAExecutionProvider",
)
small_image = txt2img(
prompt,
num_inference_steps=steps,
).images[0]

generator = torch.manual_seed(0)
upscale = OnnxStableDiffusionUpscalePipeline.from_pretrained(
"ssube/stable-diffusion-x4-upscaler-onnx",
provider="CUDAExecutionProvider",
)
large_image = upscale(
prompt,
small_image,
generator=generator,
num_inference_steps=steps,
).images[0]
images = pipe(prompt).images[0]
```

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

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

## Known Issues

- 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.
2 changes: 1 addition & 1 deletion docs/source/en/optimization/open_vino.mdx
Original file line number Diff line number Diff line change
Expand Up @@ -36,4 +36,4 @@ prompt = "a photo of an astronaut riding a horse on mars"
images = pipe(prompt).images[0]
```

You can find more examples in [optimum documentation](https://huggingface.co/docs/optimum/intel/inference#export-and-inference-of-stable-diffusion-models).
You can find more examples (such as static reshaping and model compilation) in [optimum documentation](https://huggingface.co/docs/optimum/intel/inference#export-and-inference-of-stable-diffusion-models).

0 comments on commit a169571

Please sign in to comment.