Reuse existing quantization file

Author: jkosek

model_navigator/commands/convert/converters/onnx2trt.py

@@ -62,6 +62,30 @@ def _get_precisions(precision, precision_mode):
     return tf32, fp16, bf16, fp8, int8, nvfp4
 
 
+def _quantize_model(
+    navigator_workspace: pathlib.Path,
+    batch_dim: int,
+    quantized_onnx_path: pathlib.Path,
+    onnx_path: pathlib.Path,
+    precision: str,
+):
+    import modelopt.onnx.quantization as moq  # pytype: disable=import-error # noqa: F401
+
+    correctness_samples = load_samples("correctness_samples", navigator_workspace, batch_dim)
+    calibration_data = {name: tensor for sample in correctness_samples for name, tensor in sample.items()}
+
+    # Prepare quantization parameters
+    quantize_kwargs = {
+        "onnx_path": onnx_path.as_posix(),
+        "calibration_data": calibration_data,
+        "output_path": quantized_onnx_path.as_posix(),
+        "quantize_mode": precision,
+    }
+
+    moq.quantize(**quantize_kwargs)
+    LOGGER.info("Quantized ONNX model saved in {}", quantized_onnx_path)
+
+
 def convert(
     exported_model_path: str,
     converted_model_path: str,
@@ -111,7 +135,6 @@ def convert(
     exported_model_path = pathlib.Path(exported_model_path)
     if not exported_model_path.is_absolute():
         exported_model_path = navigator_workspace / exported_model_path
-    exported_model_path = exported_model_path.as_posix()
 
     if model_name is None:
         model_name = navigator_workspace.stem
@@ -125,7 +148,6 @@ def convert(
         quantized_onnx_path = pathlib.Path(quantized_onnx_path)
         if not quantized_onnx_path.is_absolute():
             quantized_onnx_path = navigator_workspace / quantized_onnx_path
-        quantized_onnx_path = quantized_onnx_path.as_posix()
 
     custom_args = custom_args or {}
 
@@ -140,9 +162,10 @@ def convert(
 
     # nvfp4 is currently not used as flag for converter, skip it
     tf32, fp16, bf16, fp8, int8, _ = _get_precisions(precision, precision_mode)
+    strongly_typed = False
 
     # Determine the path to use for ONNX model
-    onnx_path = exported_model_path
+    onnx_path = pathlib.Path(exported_model_path)
 
     # Check if we need to perform quantization
     should_quantize = (
@@ -154,31 +177,19 @@ def convert(
     # Use ModelOpt for quantization if needed
     if quantized_onnx_path and should_quantize:
         LOGGER.info("Quantize model through TensorRT ModelOpt with {} precision", precision)
-        import modelopt.onnx.quantization as moq  # pytype: disable=import-error # noqa: F401
 
-        correctness_samples = load_samples("correctness_samples", navigator_workspace, batch_dim)
-        calibration_data = {name: tensor for sample in correctness_samples for name, tensor in sample.items()}
+        if not pathlib.Path(quantized_onnx_path).exists():
+            _quantize_model(navigator_workspace, batch_dim, quantized_onnx_path, onnx_path, precision)
+        else:
+            LOGGER.info("Quantized ONNX model already exists in {}", quantized_onnx_path)
 
-        # Prepare quantization parameters
-        quantize_kwargs = {
-            "onnx_path": onnx_path,
-            "calibration_data": calibration_data,
-            "output_path": quantized_onnx_path,
-            "quantize_mode": precision,
-        }
-
-        moq.quantize(**quantize_kwargs)
-        LOGGER.info("Quantized ONNX model saved in {}", quantized_onnx_path)
-        onnx_path = quantized_onnx_path
+        onnx_path = pathlib.Path(quantized_onnx_path)
     # For NVFP4, always use the quantized path (even if not quantized yet)
     elif quantized_onnx_path and TensorRTPrecision(precision) == TensorRTPrecision.NVFP4:
-        onnx_path = quantized_onnx_path
-
-    if TensorRTPrecision(precision) == TensorRTPrecision.NVFP4:
         strongly_typed = True
-    else:
-        strongly_typed = False
-    network = network_from_onnx_path(onnx_path, flags=onnx_parser_flags, strongly_typed=strongly_typed)
+        onnx_path = pathlib.Path(quantized_onnx_path)
+
+    network = network_from_onnx_path(onnx_path.as_posix(), flags=onnx_parser_flags, strongly_typed=strongly_typed)
 
     config_kwargs = {}
     if optimization_level: