fix(api): change weights filename for new models to match optimum

ssube · Dec 24, 2023 · ebe813d · ebe813d
1 parent 14208de
commit ebe813d
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 4 deletions.
diff --git a/api/onnx_web/constants.py b/api/onnx_web/constants.py
@@ -1,5 +1,5 @@
 ONNX_MODEL = "model.onnx"
-ONNX_WEIGHTS = "weights.pb"
+ONNX_WEIGHTS = "model.onnx_data"
 
 LATENT_FACTOR = 8
 LATENT_CHANNELS = 4
diff --git a/api/onnx_web/convert/diffusion/diffusion.py b/api/onnx_web/convert/diffusion/diffusion.py
@@ -865,7 +865,7 @@ def convert_diffusion_diffusers_optimum(
             unet_path,
             save_as_external_data=True,
             all_tensors_to_one_file=True,
-            location="weights.pb",
+            location=ONNX_WEIGHTS,
         )
 
     return (True, dest_path)
diff --git a/api/onnx_web/convert/diffusion/diffusion_xl.py b/api/onnx_web/convert/diffusion/diffusion_xl.py
@@ -9,7 +9,7 @@
 from onnxruntime.transformers.float16 import convert_float_to_float16
 from optimum.exporters.onnx import main_export
 
-from ...constants import ONNX_MODEL
+from ...constants import ONNX_MODEL, ONNX_WEIGHTS
 from ..client import fetch_model
 from ..utils import RESOLVE_FORMATS, ConversionContext, check_ext
 
@@ -112,7 +112,7 @@ def convert_diffusion_diffusers_xl(
             unet_path,
             save_as_external_data=True,
             all_tensors_to_one_file=True,
-            location="weights.pb",
+            location=ONNX_WEIGHTS,
         )
 
     return (True, dest_path)