revert torch_onnx.py

Signed-off-by: Zhiyu Cheng <[email protected]>

Author: Edwardf0t1

modelopt/torch/_deploy/utils/torch_onnx.py

@@ -45,15 +45,14 @@
     get_node_names,
     get_output_names,
     get_output_shapes,
-    infer_shapes,
     remove_node_training_mode,
 )
 from modelopt.torch.quantization.export_onnx import configure_linear_module_onnx_quantizers
 from modelopt.torch.utils import flatten_tree, standardize_named_model_args
 from modelopt.torch.utils._pytree import TreeSpec
 
 from ..utils.onnx_optimizer import Optimizer
-from .onnx_utils import check_model_uses_external_data
+from .onnx_utils import _get_onnx_external_data_tensors, check_model_uses_external_data
 
 ModelMetadata = dict[str, Any]
 ModelType = Any
@@ -84,8 +83,15 @@ def __init__(self, onnx_load_path: str) -> None:
         self.onnx_load_path = os.path.abspath(onnx_load_path)
         self.onnx_model = {}
         self.model_name = ""
+        onnx_model = onnx.load(self.onnx_load_path, load_external_data=False)
 
-        if has_external_data(onnx_load_path):
+        # Check for external data
+        external_data_format = False
+        for initializer in onnx_model.graph.initializer:
+            if initializer.external_data:
+                external_data_format = True
+
+        if external_data_format:
             onnx_model_dir = os.path.dirname(self.onnx_load_path)
             for onnx_model_file in os.listdir(onnx_model_dir):
                 with open(os.path.join(onnx_model_dir, onnx_model_file), "rb") as f:
@@ -413,7 +419,9 @@ def get_onnx_bytes_and_metadata(
     # Export onnx model from pytorch model
     # As the maximum size of protobuf is 2GB, we cannot use io.BytesIO() buffer during export.
     model_name = model.__class__.__name__
-    onnx_path = tempfile.mkdtemp(prefix=f"modelopt_{model_name}_")
+    onnx_build_folder = os.path.join(tempfile.gettempdir(), "modelopt_build/onnx/")
+    onnx_path = os.path.join(onnx_build_folder, model_name)
+    os.makedirs(onnx_path, exist_ok=True)
     onnx_save_path = os.path.join(onnx_path, f"{model_name}.onnx")
 
     # Configure quantizers if the model is quantized in NVFP4 or MXFP8 mode
@@ -444,7 +452,7 @@ def get_onnx_bytes_and_metadata(
     onnx_graph = onnx.load(onnx_save_path, load_external_data=True)
 
     try:
-        onnx_graph = infer_shapes(onnx_graph)
+        onnx_graph = onnx.shape_inference.infer_shapes(onnx_graph)
     except Exception as e:
         print(f"Shape inference failed: {e}")
 
@@ -494,37 +502,28 @@ def get_onnx_bytes_and_metadata(
 
     # If the onnx model contains external data store the external tensors in one file and save the onnx model
     if has_external_data(onnx_save_path):
-        tensor_paths = get_external_tensor_paths(onnx_path)
+        tensor_paths = _get_onnx_external_data_tensors(onnx_opt_graph)
         onnx.save_model(
             onnx_opt_graph,
             onnx_save_path,
             save_as_external_data=True,
             all_tensors_to_one_file=True,
             location=f"{model_name}.onnx_data",
             size_threshold=1024,
-            convert_attribute=False,
         )
-        for path in tensor_paths:
-            os.remove(path)
+        for tensor in tensor_paths:
+            tensor_path = os.path.join(onnx_path, tensor)
+            os.remove(tensor_path)
     else:
         onnx.save_model(onnx_opt_graph, onnx_save_path)
 
     onnx_bytes = OnnxBytes(onnx_save_path)
 
     if remove_exported_model:
-        shutil.rmtree(onnx_path)
+        shutil.rmtree(os.path.dirname(onnx_build_folder))
     return onnx_bytes.to_bytes(), model_metadata
 
 
-def get_external_tensor_paths(model_dir: str) -> list[str]:
-    """Get the paths of the external data tensors in the model."""
-    return [
-        os.path.join(model_dir, file)
-        for file in os.listdir(model_dir)
-        if not file.endswith(".onnx")
-    ]
-
-
 def has_external_data(onnx_model_path: str):
     """Check if the onnx model has external data."""
     onnx_model = onnx.load(onnx_model_path, load_external_data=False)