fix(api): enable optimizations during model conversion

ssube · Feb 24, 2023 · fd013c8 · fd013c8
1 parent 2edb3c6
commit fd013c8
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/api/onnx_web/convert/diffusion_stable.py b/api/onnx_web/convert/diffusion_stable.py
@@ -25,6 +25,8 @@
 from onnx import load, save_model
 from torch.onnx import export
 
+from onnx_web.diffusion.load import optimize_pipeline
+
 from ..diffusion.pipeline_onnx_stable_diffusion_upscale import (
     OnnxStableDiffusionUpscalePipeline,
 )
@@ -96,6 +98,8 @@ def convert_diffusion_stable(
     ).to(ctx.training_device)
     output_path = Path(dest_path)
 
+    optimize_pipeline(ctx, pipeline)
+
     # TEXT ENCODER
     num_tokens = pipeline.text_encoder.config.max_position_embeddings
     text_hidden_size = pipeline.text_encoder.config.hidden_size