Skip to content

Commit 8d9c4a5

Browse files
anton-lpatrickvonplatenpatil-suraj
authored
[ONNX] Stable Diffusion exporter and pipeline (huggingface#399)
* initial export and design * update imports * custom prover, import fixes * Update src/diffusers/onnx_utils.py Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com> * Update src/diffusers/onnx_utils.py Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com> * remove push_to_hub * Update src/diffusers/onnx_utils.py Co-authored-by: Suraj Patil <surajp815@gmail.com> * remove torch_device * numpify the rest of the pipeline * torchify the safety checker * revert tensor * Code review suggestions + quality * fix tests * fix provider, add an end-to-end test * style Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com> Co-authored-by: Suraj Patil <surajp815@gmail.com>
1 parent 7bcc873 commit 8d9c4a5

13 files changed

Lines changed: 657 additions & 6 deletions

‎scripts/convert_stable_diffusion_checkpoint_to_onnx.py‎

Lines changed: 196 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,196 @@
1+
# Copyright 2022 The HuggingFace Team. All rights reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import argparse
16+
from pathlib import Path
17+
18+
import torch
19+
from torch.onnx import export
20+
21+
from diffusers import StableDiffusionOnnxPipeline, StableDiffusionPipeline
22+
from diffusers.onnx_utils import OnnxRuntimeModel
23+
from packaging import version
24+
25+
26+
is_torch_less_than_1_11 = version.parse(version.parse(torch.__version__).base_version) < version.parse("1.11")
27+
28+
29+
def onnx_export(
30+
model,
31+
model_args: tuple,
32+
output_path: Path,
33+
ordered_input_names,
34+
output_names,
35+
dynamic_axes,
36+
opset,
37+
use_external_data_format=False,
38+
):
39+
output_path.parent.mkdir(parents=True, exist_ok=True)
40+
# PyTorch deprecated the `enable_onnx_checker` and `use_external_data_format` arguments in v1.11,
41+
# so we check the torch version for backwards compatibility
42+
if is_torch_less_than_1_11:
43+
export(
44+
model,
45+
model_args,
46+
f=output_path.as_posix(),
47+
input_names=ordered_input_names,
48+
output_names=output_names,
49+
dynamic_axes=dynamic_axes,
50+
do_constant_folding=True,
51+
use_external_data_format=use_external_data_format,
52+
enable_onnx_checker=True,
53+
opset_version=opset,
54+
)
55+
else:
56+
export(
57+
model,
58+
model_args,
59+
f=output_path.as_posix(),
60+
input_names=ordered_input_names,
61+
output_names=output_names,
62+
dynamic_axes=dynamic_axes,
63+
do_constant_folding=True,
64+
opset_version=opset,
65+
)
66+
67+
68+
@torch.no_grad()
69+
def convert_models(model_path: str, output_path: str, opset: int):
70+
pipeline = StableDiffusionPipeline.from_pretrained(model_path, use_auth_token=True)
71+
output_path = Path(output_path)
72+
73+
# TEXT ENCODER
74+
text_input = pipeline.tokenizer(
75+
"A sample prompt",
76+
padding="max_length",
77+
max_length=pipeline.tokenizer.model_max_length,
78+
truncation=True,
79+
return_tensors="pt",
80+
)
81+
onnx_export(
82+
pipeline.text_encoder,
83+
# casting to torch.int32 until the CLIP fix is released: https://github.com/huggingface/transformers/pull/18515/files
84+
model_args=(text_input.input_ids.to(torch.int32)),
85+
output_path=output_path / "text_encoder" / "model.onnx",
86+
ordered_input_names=["input_ids"],
87+
output_names=["last_hidden_state", "pooler_output"],
88+
dynamic_axes={
89+
"input_ids": {0: "batch", 1: "sequence"},
90+
},
91+
opset=opset,
92+
)
93+
94+
# UNET
95+
onnx_export(
96+
pipeline.unet,
97+
model_args=(torch.randn(2, 4, 64, 64), torch.LongTensor([0, 1]), torch.randn(2, 77, 768), False),
98+
output_path=output_path / "unet" / "model.onnx",
99+
ordered_input_names=["sample", "timestep", "encoder_hidden_states", "return_dict"],
100+
output_names=["out_sample"], # has to be different from "sample" for correct tracing
101+
dynamic_axes={
102+
"sample": {0: "batch", 1: "channels", 2: "height", 3: "width"},
103+
"timestep": {0: "batch"},
104+
"encoder_hidden_states": {0: "batch", 1: "sequence"},
105+
},
106+
opset=opset,
107+
use_external_data_format=True, # UNet is > 2GB, so the weights need to be split
108+
)
109+
110+
# VAE ENCODER
111+
vae_encoder = pipeline.vae
112+
# need to get the raw tensor output (sample) from the encoder
113+
vae_encoder.forward = lambda sample, return_dict: vae_encoder.encode(sample, return_dict)[0].sample()
114+
onnx_export(
115+
vae_encoder,
116+
model_args=(torch.randn(1, 3, 512, 512), False),
117+
output_path=output_path / "vae_encoder" / "model.onnx",
118+
ordered_input_names=["sample", "return_dict"],
119+
output_names=["latent_sample"],
120+
dynamic_axes={
121+
"sample": {0: "batch", 1: "channels", 2: "height", 3: "width"},
122+
},
123+
opset=opset,
124+
)
125+
126+
# VAE DECODER
127+
vae_decoder = pipeline.vae
128+
# forward only through the decoder part
129+
vae_decoder.forward = vae_encoder.decode
130+
onnx_export(
131+
vae_decoder,
132+
model_args=(torch.randn(1, 4, 64, 64), False),
133+
output_path=output_path / "vae_decoder" / "model.onnx",
134+
ordered_input_names=["latent_sample", "return_dict"],
135+
output_names=["sample"],
136+
dynamic_axes={
137+
"latent_sample": {0: "batch", 1: "channels", 2: "height", 3: "width"},
138+
},
139+
opset=opset,
140+
)
141+
142+
# SAFETY CHECKER
143+
safety_checker = pipeline.safety_checker
144+
safety_checker.forward = safety_checker.forward_onnx
145+
onnx_export(
146+
pipeline.safety_checker,
147+
model_args=(torch.randn(1, 3, 224, 224), torch.randn(1, 512, 512, 3)),
148+
output_path=output_path / "safety_checker" / "model.onnx",
149+
ordered_input_names=["clip_input", "images"],
150+
output_names=["out_images", "has_nsfw_concepts"],
151+
dynamic_axes={
152+
"clip_input": {0: "batch", 1: "channels", 2: "height", 3: "width"},
153+
"images": {0: "batch", 1: "channels", 2: "height", 3: "width"},
154+
},
155+
opset=opset,
156+
)
157+
158+
onnx_pipeline = StableDiffusionOnnxPipeline(
159+
vae_decoder=OnnxRuntimeModel.from_pretrained(output_path / "vae_decoder"),
160+
text_encoder=OnnxRuntimeModel.from_pretrained(output_path / "text_encoder"),
161+
tokenizer=pipeline.tokenizer,
162+
unet=OnnxRuntimeModel.from_pretrained(output_path / "unet"),
163+
scheduler=pipeline.scheduler,
164+
safety_checker=OnnxRuntimeModel.from_pretrained(output_path / "safety_checker"),
165+
feature_extractor=pipeline.feature_extractor,
166+
)
167+
168+
onnx_pipeline.save_pretrained(output_path)
169+
print("ONNX pipeline saved to", output_path)
170+
171+
_ = StableDiffusionOnnxPipeline.from_pretrained(output_path, provider="CPUExecutionProvider")
172+
print("ONNX pipeline is loadable")
173+
174+
175+
if __name__ == "__main__":
176+
parser = argparse.ArgumentParser()
177+
178+
parser.add_argument(
179+
"--model_path",
180+
type=str,
181+
required=True,
182+
help="Path to the `diffusers` checkpoint to convert (either a local directory or on the Hub).",
183+
)
184+
185+
parser.add_argument("--output_path", type=str, required=True, help="Path to the output model.")
186+
187+
parser.add_argument(
188+
"--opset",
189+
default=14,
190+
type=str,
191+
help="The version of the ONNX operator set to use.",
192+
)
193+
194+
args = parser.parse_args()
195+
196+
convert_models(args.model_path, args.output_path, args.opset)

‎setup.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -170,7 +170,7 @@ def run(self):
170170
extras["quality"] = ["black==22.3", "isort>=5.5.4", "flake8>=3.8.3", "hf-doc-builder"]
171171
extras["docs"] = ["hf-doc-builder"]
172172
extras["training"] = ["accelerate", "datasets", "tensorboard", "modelcards"]
173-
extras["test"] = ["datasets", "pytest", "pytest-timeout", "pytest-xdist", "scipy", "transformers"]
173+
extras["test"] = ["datasets", "onnxruntime", "pytest", "pytest-timeout", "pytest-xdist", "scipy", "transformers"]
174174
extras["dev"] = extras["quality"] + extras["test"] + extras["training"] + extras["docs"]
175175

176176
install_requires = [

‎src/diffusers/__init__.py‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,17 @@
1-
from .utils import is_inflect_available, is_scipy_available, is_transformers_available, is_unidecode_available
1+
from .utils import (
2+
is_inflect_available,
3+
is_onnx_available,
4+
is_scipy_available,
5+
is_transformers_available,
6+
is_unidecode_available,
7+
)
28

39

410
__version__ = "0.3.0.dev0"
511

612
from .modeling_utils import ModelMixin
713
from .models import AutoencoderKL, UNet2DConditionModel, UNet2DModel, VQModel
14+
from .onnx_utils import OnnxRuntimeModel
815
from .optimization import (
916
get_constant_schedule,
1017
get_constant_schedule_with_warmup,
@@ -44,3 +51,9 @@
4451
)
4552
else:
4653
from .utils.dummy_transformers_objects import * # noqa F403
54+
55+
56+
if is_transformers_available() and is_onnx_available():
57+
from .pipelines import StableDiffusionOnnxPipeline
58+
else:
59+
from .utils.dummy_transformers_and_onnx_objects import * # noqa F403

0 commit comments

Comments
 (0)