apextoaster
/
onnx-web
1
0
Fork 0
Code Actions

fix(api): enable optimizations during model conversion

This commit is contained in:
Sean Sube 2023-02-23 18:43:49 -06:00
parent 2edb3c6199
commit fd013c88ef
Signed by: ssube
GPG Key ID: 3EED7B957D362AF1
1 changed files with 4 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
api/onnx_web/convert/diffusion_stable.py
View File

@ -25,6 +25,8 @@ from diffusers import (
from onnx import load, save_model from onnx import load, save_model
from torch.onnx import export from torch.onnx import export
from onnx_web.diffusion.load import optimize_pipeline
from ..diffusion.pipeline_onnx_stable_diffusion_upscale import ( from ..diffusion.pipeline_onnx_stable_diffusion_upscale import (
OnnxStableDiffusionUpscalePipeline, OnnxStableDiffusionUpscalePipeline,
) )
@ -96,6 +98,8 @@ def convert_diffusion_stable(
).to(ctx.training_device) ).to(ctx.training_device)
output_path = Path(dest_path) output_path = Path(dest_path)
optimize_pipeline(ctx, pipeline)
# TEXT ENCODER # TEXT ENCODER
num_tokens = pipeline.text_encoder.config.max_position_embeddings num_tokens = pipeline.text_encoder.config.max_position_embeddings
text_hidden_size = pipeline.text_encoder.config.hidden_size text_hidden_size = pipeline.text_encoder.config.hidden_size