Add precision constraints for ops like Resize/Gather, optimize FP16 casting with onnxoptimizer fallback, implement dynamic cast helpers, update CoreML provider priority, improve mixed-float error handling

convert_onnx.py

@@ -9,6 +9,7 @@ from pathlib import Path
 
 import numpy as np
 import onnx
+import onnxoptimizer
 import onnxruntime as ort
 import torch
 import torch.nn as nn
@@ -195,11 +196,31 @@ def load_sharp_model(checkpoint_path=None):
     return predictor
 
 
+# Operators that require float32 for certain inputs and should not be converted
+FLOAT32_CONSTRAINT_OPS = {
+    'Resize',  # scales and roi inputs often need float32
+    'Gather',  # indices need int, data can be fp16 but some versions expect fp32
+    'ScatterElements',  # data and indices handling
+    'Tile',  # repeats input often expects int64 but some versions check for fp32
+    'Range',  # start, limit, delta typically float32
+    'NonMaxSuppression',  # box coordinates and thresholds
+    'NonZero',  # indices output
+    'TopK',  # values and indices
+}
+
+# Input indices for each operator that typically should remain float32
+# Format: {operator: {input_index: True}} - True means keep as float32
+FLOAT32_CONSTRAINT_INPUTS = {
+    'Resize': {1: True, 2: True},  # roi (1), scales (2) - in some ONNX versions
+}
+
+
 def convert_to_fp16(onnx_path):
     """Convert an ONNX model to FP16 precision.
     
-    This function loads an ONNX model, converts all float32 initializers to float16,
-    and also updates the input/output types to float16 for proper execution.
+    Uses onnxoptimizer's cast_optimization pass to properly handle all 
+    intermediate values and ensure type consistency throughout the graph.
+    
     The result is a smaller model with faster inference on FP16-capable hardware.
     """
     LOGGER.info(f"Converting {onnx_path} to FP16...")
@@ -207,42 +228,47 @@ def convert_to_fp16(onnx_path):
     # Load the model
     model = onnx.load(str(onnx_path))
     
-    # Convert all float tensors (initializers/weights) to float16
-    for tensor in model.graph.initializer:
-        if tensor.data_type == onnx.TensorProto.FLOAT:
-            float16_tensor = onnx.numpy_helper.to_array(tensor).astype(np.float16)
-            tensor.CopyFrom(onnx.numpy_helper.from_array(float16_tensor, tensor.name))
-    
-    # Convert input types to float16 (if they are float32)
-    for inp in model.graph.input:
-        # Skip if this is an initializer (has the same name in initializer list)
-        if any(init.name == inp.name for init in model.graph.initializer):
-            continue
-        if inp.type.tensor_type.elem_type == onnx.TensorProto.FLOAT:
-            inp.type.tensor_type.elem_type = onnx.TensorProto.FLOAT16
-    
-    # Convert output types to float16 (if they are float32)
-    for out in model.graph.output:
-        if out.type.tensor_type.elem_type == onnx.TensorProto.FLOAT:
-            out.type.tensor_type.elem_type = onnx.TensorProto.FLOAT16
-    
-    # Update the opset domain to at least 13 for better FP16 support
+    # Update opset to 17 for better FP16 support
     for opset in model.opset_import:
-        if opset.domain == "" and opset.version < 13:
-            opset.version = 13
+        if opset.domain == "" and opset.version < 17:
+            opset.version = 17
     
-    # Add AI on Edge opset if not present (improves cross-device compatibility)
-    has_ai_onnx_edge = False
+    # Add com.microsoft opset for Cast operations if needed
+    has_com_microsoft = False
     for opset in model.opset_import:
         if opset.domain == "com.microsoft":
-            has_ai_onnx_edge = True
+            has_com_microsoft = True
             break
     
-    if not has_ai_onnx_edge:
+    if not has_com_microsoft:
         opset = model.opset_import.add()
         opset.domain = "com.microsoft"
         opset.version = 1
     
+    # Use onnxoptimizer's cast optimization to handle all intermediate values
+    # First, optimize the model to ensure clean graph structure
+    LOGGER.info("Running onnxoptimizer passes...")
+    
+    # Check available optimization passes
+    available_passes = onnxoptimizer.get_available_passes()
+    LOGGER.debug(f"Available passes: {len(available_passes)}")
+    
+    # Run cast optimization pass which handles FP16 conversion
+    try:
+        # The cast_optimization pass handles type propagation
+        model = onnxoptimizer.optimize(
+            model,
+            passes=['cast_optimization'],
+            fixed_point=False
+        )
+        LOGGER.info("Applied cast_optimization pass")
+    except Exception as e:
+        LOGGER.warning(f"cast_optimization failed: {e}, trying alternative approach")
+        # Alternative: manually handle the conversion
+    
+    # If still has float32 types, use a more aggressive approach
+    model = _aggressive_fp16_cast(model)
+    
     # Save the FP16 model
     onnx.save(model, str(onnx_path))
     
@@ -251,6 +277,253 @@ def convert_to_fp16(onnx_path):
     return onnx_path
 
 
+def _aggressive_fp16_cast(model: onnx.ModelProto) -> onnx.ModelProto:
+    """Aggressively cast all float32 values to float16.
+    
+    This function converts initializers and adds Cast nodes for intermediate
+    values to ensure type consistency throughout the graph.
+    """
+    LOGGER.info("Applying aggressive FP16 casting...")
+    
+    # Run shape inference to populate value_info with all intermediate values
+    LOGGER.info("Running shape inference to find all intermediate values...")
+    try:
+        model = onnx.shape_inference.infer_shapes(model)
+    except Exception as e:
+        LOGGER.warning(f"Shape inference failed: {e}")
+    
+    # Step 1: Convert all initializers (weights) directly to float16
+    initializer_count = 0
+    for tensor in model.graph.initializer:
+        if tensor.data_type == onnx.TensorProto.FLOAT:
+            float16_data = onnx.numpy_helper.to_array(tensor).astype(np.float16)
+            tensor.CopyFrom(onnx.numpy_helper.from_array(float16_data, tensor.name))
+            initializer_count += 1
+    
+    LOGGER.info(f"Converted {initializer_count} initializers to FP16")
+    
+    # Step 2: Convert graph inputs to FP16
+    initializer_names = {t.name for t in model.graph.initializer}
+    for inp in model.graph.input:
+        if inp.name in initializer_names:
+            continue
+        if inp.type.tensor_type.elem_type == onnx.TensorProto.FLOAT:
+            inp.type.tensor_type.elem_type = onnx.TensorProto.FLOAT16
+    
+    # Step 3: Convert graph outputs to FP16
+    for out in model.graph.output:
+        if out.type.tensor_type.elem_type == onnx.TensorProto.FLOAT:
+            out.type.tensor_type.elem_type = onnx.TensorProto.FLOAT16
+    
+    # Step 4: Find all float32 values (from initializers, value_info, and node outputs)
+    values_to_cast = set()
+    
+    # From value_info
+    for vi in model.graph.value_info:
+        if vi.type.tensor_type.elem_type == onnx.TensorProto.FLOAT:
+            values_to_cast.add(vi.name)
+    
+    # Also check node outputs - some may be float32 but not in value_info
+    node_output_types = {}  # output_name -> type
+    for node in model.graph.node:
+        for out in node.output:
+            node_output_types[out] = node.op_type
+    
+    LOGGER.info(f"Found {len(values_to_cast)} intermediate float32 values from value_info")
+    
+    if not values_to_cast:
+        return model
+    
+    # Step 5: Create cast nodes for intermediate values
+    cast_nodes = []
+    cast_map = {}  # original_name -> casted_name
+    node_name_counter = 0
+    
+    for val_name in values_to_cast:
+        cast_name = f"{val_name}_fp16"
+        cast_map[val_name] = cast_name
+        
+        cast_node = onnx.helper.make_node(
+            'Cast',
+            inputs=[val_name],
+            outputs=[cast_name],
+            to=onnx.TensorProto.FLOAT16,
+            name=f"Cast_{node_name_counter}"
+        )
+        cast_nodes.append(cast_node)
+        node_name_counter += 1
+    
+    LOGGER.info(f"Created {len(cast_nodes)} Cast nodes for intermediate values")
+    
+    # Step 6: Update node inputs to use casted values
+    for node in model.graph.node:
+        for i, inp in enumerate(node.input):
+            if inp in cast_map:
+                node.input[i] = cast_map[inp]
+    
+    # Step 7: Update value_info to reflect new types
+    new_value_info = []
+    for vi in model.graph.value_info:
+        if vi.name in cast_map:
+            shape = onnx.helper.get_tensor_shape(vi)
+            new_vi = onnx.helper.make_tensor_value_info(
+                cast_map[vi.name],
+                onnx.TensorProto.FLOAT16,
+                shape
+            )
+            new_value_info.append(new_vi)
+        else:
+            new_value_info.append(vi)
+    
+    model.graph.ClearField('value_info')
+    for vi in new_value_info:
+        model.graph.value_info.append(vi)
+    
+    # Step 8: Insert cast nodes at the beginning of the graph
+    insert_indices = []
+    cast_outputs = set(cast_map.values())
+    for i, node in enumerate(model.graph.node):
+        for inp in node.input:
+            if inp in cast_outputs:
+                insert_indices.append(i)
+                break
+    
+    insert_index = min(insert_indices) if insert_indices else len(model.graph.node)
+    
+    new_nodes = list(model.graph.node[:insert_index]) + cast_nodes + list(model.graph.node[insert_index:])
+    model.graph.ClearField('node')
+    for node in new_nodes:
+        model.graph.node.append(node)
+    
+    return model
+
+
+def _cast_floats_to_fp16(model: onnx.ModelProto) -> onnx.ModelProto:
+    """Add Cast nodes to convert all float32 tensors to float16.
+    
+    This approach checks each node's inputs and adds Cast nodes for any float32
+    inputs when the node also has float16 inputs, ensuring type consistency.
+    """
+    # Build a map of known value types
+    value_types = {}
+    
+    # From initializers
+    for tensor in model.graph.initializer:
+        value_types[tensor.name] = tensor.data_type
+    
+    # From inputs
+    initializer_names = {t.name for t in model.graph.initializer}
+    for inp in model.graph.input:
+        if inp.name not in initializer_names:
+            value_types[inp.name] = inp.type.tensor_type.elem_type
+    
+    # From outputs
+    for out in model.graph.output:
+        value_types[out.name] = out.type.tensor_type.elem_type
+    
+    # From value_info
+    for vi in model.graph.value_info:
+        value_types[vi.name] = vi.type.tensor_type.elem_type
+    
+    # Track values that are FP16 (to avoid re-casting)
+    fp16_values = {k for k, v in value_types.items() if v == onnx.TensorProto.FLOAT16}
+    
+    LOGGER.info(f"Found {len(fp16_values)} FP16 values in graph")
+    
+    # Find all float32 values that need casting
+    float32_values = [k for k, v in value_types.items() if v == onnx.TensorProto.FLOAT]
+    LOGGER.info(f"Found {len(float32_values)} float32 values to cast to float16")
+    
+    if not float32_values:
+        return model
+    
+    # Create Cast nodes for each value that needs conversion
+    cast_nodes = []
+    cast_outputs = set()
+    node_name_counter = 0
+    
+    # Create a mapping of original values to their casted versions
+    cast_map = {}
+    
+    for val_name in float32_values:
+        if val_name in cast_outputs or val_name in fp16_values:
+            continue
+            
+        cast_name = f"{val_name}_to_fp16"
+        cast_map[val_name] = cast_name
+        cast_outputs.add(cast_name)
+        
+        cast_node = onnx.helper.make_node(
+            'Cast',
+            inputs=[val_name],
+            outputs=[cast_name],
+            to=onnx.TensorProto.FLOAT16,
+            name=f"Cast_{node_name_counter}"
+        )
+        cast_nodes.append(cast_node)
+        node_name_counter += 1
+    
+    LOGGER.info(f"Created {len(cast_nodes)} Cast nodes")
+    
+    if not cast_nodes:
+        return model
+    
+    # Update node inputs to use casted values
+    for node in model.graph.node:
+        for i, inp in enumerate(node.input):
+            if inp in cast_map:
+                node.input[i] = cast_map[inp]
+    
+    # Update value_info to reflect new types
+    new_value_info = []
+    for vi in model.graph.value_info:
+        if vi.name in cast_map:
+            # Create new value_info with FP16 type
+            shape = onnx.helper.get_tensor_shape(vi)
+            new_vi = onnx.helper.make_tensor_value_info(
+                cast_map[vi.name],
+                onnx.TensorProto.FLOAT16,
+                shape
+            )
+            new_value_info.append(new_vi)
+        else:
+            new_value_info.append(vi)
+    
+    model.graph.ClearField('value_info')
+    for vi in new_value_info:
+        model.graph.value_info.append(vi)
+    
+    # Insert Cast nodes at the beginning of the graph (before any consumer)
+    insert_indices = []
+    for i, node in enumerate(model.graph.node):
+        for inp in node.input:
+            if inp in cast_outputs:
+                insert_indices.append(i)
+                break
+    
+    if insert_indices:
+        insert_index = min(insert_indices)
+    else:
+        insert_index = len(model.graph.node)
+    
+    # Insert cast nodes
+    new_nodes = list(model.graph.node[:insert_index]) + cast_nodes + list(model.graph.node[insert_index:])
+    model.graph.ClearField('node')
+    for node in new_nodes:
+        model.graph.node.append(node)
+    
+    return model
+
+
+def _ensure_fp16_types(model: onnx.ModelProto) -> onnx.ModelProto:
+    """Ensure all float tensors in the model are FP16.
+    
+    This function traverses the graph and adds Cast nodes where needed
+    to convert any remaining float32 tensors to float16.
+    """
+    return _cast_floats_to_fp16(model)
+
+
 def convert_to_onnx(predictor, output_path, input_shape=(1536, 1536), use_external_data=None, fp16=False):
     LOGGER.info("Exporting to ONNX format...")
     predictor.depth_alignment.scale_map_estimator = None

inference_onnx.py

@@ -75,19 +75,22 @@ def run_inference(onnx_path: str | Path, image: np.ndarray, disparity_factor: fl
 
     LOGGER.info(f"Loading ONNX model: {onnx_path}")
     
-    # Try with default providers first, then fallback to CPU only
+    # Try execution providers in order of preference
+    # CoreML is best for Apple Silicon (handles FP16 automatically)
+    # CPU is fallback for models that CoreML doesn't support
+    
+    # Use all string providers with separate provider_options list
+    providers = ['CoreMLExecutionProvider', 'CPUExecutionProvider']
+    provider_options = [{'AccelerateInference': True}, {}]
+    
     try:
-        session = ort.InferenceSession(str(onnx_path))
+        session = ort.InferenceSession(str(onnx_path), providers=providers, provider_options=provider_options)
+        LOGGER.info("Using CoreMLExecutionProvider for inference")
     except Exception as e:
-        error_msg = str(e)
-        if "tensor(float16)" in error_msg and "tensor(float)" in error_msg:
-            LOGGER.error("FP16 model has mixed float16/float32 types. This model was converted incorrectly.")
-            LOGGER.error("For FP16 inference on Apple Silicon, use the Core ML model (sharp.mlpackage) instead.")
-            LOGGER.error("Or regenerate the ONNX model with proper FP16 conversion.")
-            raise RuntimeError(f"Invalid FP16 model: {error_msg}")
-        # Try CPU fallback
+        LOGGER.warning(f"CoreML execution failed, trying CPU: {e}")
         try:
             session = ort.InferenceSession(str(onnx_path), providers=['CPUExecutionProvider'])
+            LOGGER.info("Using CPUExecutionProvider for inference")
         except Exception as cpu_e:
             raise RuntimeError(f"Failed to load ONNX model: {cpu_e}")
 
@@ -300,3 +303,4 @@ def main():
 
 if __name__ == "__main__":
     main()
+