FilesExpand file tree

main

/

convert_vae_diff_to_onnx.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

122 lines (107 loc) · 3.9 KB

main

/

convert_vae_diff_to_onnx.py

Copy path

Top

File metadata and controls

• Code
• Blame

122 lines (107 loc) · 3.9 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

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

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

# Copyright 2022 The HuggingFace Team. All rights reserved.

#

# 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.

import argparse

from pathlib import Path

import torch

from packaging import version

from torch.onnx import export

from diffusers import AutoencoderKL

is_torch_less_than_1_11 = version.parse(version.parse(torch.__version__).base_version) < version.parse("1.11")

def onnx_export(

model,

model_args: tuple,

output_path: Path,

ordered_input_names,

output_names,

dynamic_axes,

opset,

use_external_data_format=False,

):

output_path.parent.mkdir(parents=True, exist_ok=True)

# PyTorch deprecated the `enable_onnx_checker` and `use_external_data_format` arguments in v1.11,

# so we check the torch version for backwards compatibility

if is_torch_less_than_1_11:

export(

model,

model_args,

f=output_path.as_posix(),

input_names=ordered_input_names,

output_names=output_names,

dynamic_axes=dynamic_axes,

do_constant_folding=True,

use_external_data_format=use_external_data_format,

enable_onnx_checker=True,

opset_version=opset,

)

else:

export(

model,

model_args,

f=output_path.as_posix(),

input_names=ordered_input_names,

output_names=output_names,

dynamic_axes=dynamic_axes,

do_constant_folding=True,

opset_version=opset,

)

@torch.no_grad()

def convert_models(model_path: str, output_path: str, opset: int, fp16: bool = False):

dtype = torch.float16 if fp16 else torch.float32

if fp16 and torch.cuda.is_available():

device = "cuda"

elif fp16 and not torch.cuda.is_available():

raise ValueError("`float16` model export is only supported on GPUs with CUDA")

else:

device = "cpu"

output_path = Path(output_path)

# VAE DECODER

vae_decoder = AutoencoderKL.from_pretrained(model_path + "/vae")

vae_latent_channels = vae_decoder.config.latent_channels

# forward only through the decoder part

vae_decoder.forward = vae_decoder.decode

onnx_export(

vae_decoder,

model_args=(

torch.randn(1, vae_latent_channels, 25, 25).to(device=device, dtype=dtype),

False,

),

output_path=output_path / "vae_decoder" / "model.onnx",

ordered_input_names=["latent_sample", "return_dict"],

output_names=["sample"],

dynamic_axes={

"latent_sample": {0: "batch", 1: "channels", 2: "height", 3: "width"},

},

opset=opset,

)

del vae_decoder

if __name__ == "__main__":

parser = argparse.ArgumentParser()

parser.add_argument(

"--model_path",

type=str,

required=True,

help="Path to the `diffusers` checkpoint to convert (either a local directory or on the Hub).",

)

parser.add_argument("--output_path", type=str, required=True, help="Path to the output model.")

parser.add_argument(

"--opset",

default=14,

type=int,

help="The version of the ONNX operator set to use.",

)

parser.add_argument("--fp16", action="store_true", default=False, help="Export the models in `float16` mode")

args = parser.parse_args()

print(args.output_path)

convert_models(args.model_path, args.output_path, args.opset, args.fp16)

print("SD: Done: ONNX")