Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in / Register
Toggle navigation
Menu
Open sidebar
renzhc
diffusers_dcu
Commits
f38e3626
Commit
f38e3626
authored
Mar 06, 2023
by
Patrick von Platen
Browse files
make style
parent
5f826a35
Changes
1
Show whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
3 additions
and
7 deletions
+3
-7
scripts/convert_vae_diff_to_onnx.py
scripts/convert_vae_diff_to_onnx.py
+3
-7
No files found.
scripts/convert_vae_diff_to_onnx.py
View file @
f38e3626
...
...
@@ -13,16 +13,13 @@
# limitations under the License.
import
argparse
import
os
import
shutil
from
pathlib
import
Path
import
torch
from
packaging
import
version
from
torch.onnx
import
export
import
onnx
from
diffusers
import
OnnxRuntimeModel
,
OnnxStableDiffusionPipeline
,
StableDiffusionPipeline
,
AutoencoderKL
from
packaging
import
version
from
diffusers
import
AutoencoderKL
is_torch_less_than_1_11
=
version
.
parse
(
version
.
parse
(
torch
.
__version__
).
base_version
)
<
version
.
parse
(
"1.11"
)
...
...
@@ -81,7 +78,6 @@ def convert_models(model_path: str, output_path: str, opset: int, fp16: bool = F
# VAE DECODER
vae_decoder
=
AutoencoderKL
.
from_pretrained
(
model_path
+
"/vae"
)
vae_latent_channels
=
vae_decoder
.
config
.
latent_channels
vae_out_channels
=
vae_decoder
.
config
.
out_channels
# forward only through the decoder part
vae_decoder
.
forward
=
vae_decoder
.
decode
onnx_export
(
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment