k3-node 1.0.0__py3-none-any.whl

k3_node/export/cross_backend.py

@@ -0,0 +1,172 @@
+"""Cross-backend export bridge for PyTorch and JAX Keras backends."""
+
+import importlib
+import json
+import os
+from pathlib import Path
+import subprocess
+import sys
+import tempfile
+from typing import Any, Dict, List, Optional, Union
+import numpy as np
+
+import keras
+from keras import ops
+
+
+def is_tensorflow_backend() -> bool:
+    r"""Checks whether the currently active Keras backend is TensorFlow."""
+    return keras.backend.backend() == "tensorflow"
+
+
+def export_via_tf_subprocess(
+    exporter_type: str,  # "onnx" or "tflite"
+    model_or_task: Any,
+    output_path: Union[str, Path],
+    dummy_inputs: Optional[Any] = None,
+    **kwargs: Any,
+) -> Path:
+    r"""Bridges model export to a TensorFlow worker subprocess when running under PyTorch or JAX.
+
+    Args:
+        exporter_type: "onnx" or "tflite".
+        model_or_task: Model or task instance.
+        output_path: Target path for the exported model file.
+        dummy_inputs: Optional dummy inputs.
+        **kwargs: Extra arguments forwarded to the exporter.
+
+    Returns:
+        Path to the exported model file.
+    """
+    out_file = Path(output_path).resolve()
+    out_file.parent.mkdir(parents=True, exist_ok=True)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        tmp_path = Path(tmpdir)
+        model_dir = tmp_path / "model"
+        model_dir.mkdir(parents=True, exist_ok=True)
+
+        # 1. Ensure model is built before saving weights
+        raw_model = getattr(model_or_task, "model", None) or model_or_task
+        if hasattr(raw_model, "built") and not raw_model.built:
+            from k3_node.export.onnx_exporter import _extract_model_and_inputs
+            try:
+                _, extracted_inputs, _ = _extract_model_and_inputs(model_or_task, dummy_inputs)
+                raw_model(*extracted_inputs)
+            except Exception:
+                pass
+
+        # 2. Save model or task weights & config
+        if hasattr(model_or_task, "save_pretrained"):
+            model_or_task.save_pretrained(str(model_dir))
+            cls_module = model_or_task.__class__.__module__
+            cls_name = model_or_task.__class__.__name__
+        elif hasattr(model_or_task, "model") and hasattr(model_or_task.model, "save_pretrained"):
+            model_or_task.model.save_pretrained(str(model_dir))
+            cls_module = model_or_task.model.__class__.__module__
+            cls_name = model_or_task.model.__class__.__name__
+        else:
+            raise TypeError(
+                f"Model or task of type {type(model_or_task)} must support `save_pretrained` "
+                "for cross-backend export."
+            )
+
+        # 2. Serialize dummy inputs if provided
+        dummy_path = tmp_path / "dummy.npz"
+        dummy_meta_path = tmp_path / "dummy_meta.json"
+        has_dummy = False
+
+        if dummy_inputs is not None:
+            arrays: Dict[str, np.ndarray] = {}
+            meta: Dict[str, Any] = {}
+
+            if hasattr(dummy_inputs, "x") and hasattr(dummy_inputs, "edge_index"):
+                meta["type"] = "data"
+                arrays["x"] = np.asarray(ops.convert_to_numpy(dummy_inputs.x))
+                arrays["edge_index"] = np.asarray(ops.convert_to_numpy(dummy_inputs.edge_index))
+                if hasattr(dummy_inputs, "batch") and dummy_inputs.batch is not None:
+                    arrays["batch"] = np.asarray(ops.convert_to_numpy(dummy_inputs.batch))
+            elif hasattr(dummy_inputs, "z") and hasattr(dummy_inputs, "pos"):
+                meta["type"] = "data"
+                arrays["z"] = np.asarray(ops.convert_to_numpy(dummy_inputs.z))
+                arrays["pos"] = np.asarray(ops.convert_to_numpy(dummy_inputs.pos))
+                if hasattr(dummy_inputs, "batch") and dummy_inputs.batch is not None:
+                    arrays["batch"] = np.asarray(ops.convert_to_numpy(dummy_inputs.batch))
+            elif isinstance(dummy_inputs, (tuple, list)):
+                meta["type"] = "tuple"
+                for i, elem in enumerate(dummy_inputs):
+                    arrays[f"arr_{i}"] = np.asarray(ops.convert_to_numpy(elem))
+            elif isinstance(dummy_inputs, dict):
+                meta["type"] = "dict"
+                for k, v in dummy_inputs.items():
+                    if v is not None:
+                        arrays[str(k)] = np.asarray(ops.convert_to_numpy(v))
+            else:
+                meta["type"] = "single"
+                arrays["arr_0"] = np.asarray(ops.convert_to_numpy(dummy_inputs))
+
+            np.savez(str(dummy_path), **arrays)
+            with open(dummy_meta_path, "w", encoding="utf-8") as f:
+                json.dump(meta, f)
+            has_dummy = True
+
+        # 3. Build worker command
+        worker_code = f"""
+import os
+os.environ["KERAS_BACKEND"] = "tensorflow"
+import importlib
+import json
+from pathlib import Path
+import numpy as np
+
+import k3_node
+from k3_node.data import Data
+from k3_node.export.onnx_exporter import export_onnx
+from k3_node.export.tflite_exporter import export_tflite
+
+# Load model / task
+mod = importlib.import_module("{cls_module}")
+cls = getattr(mod, "{cls_name}")
+model_or_task = cls.from_pretrained(r"{model_dir}")
+
+# Load dummy inputs
+dummy = None
+has_dummy = {has_dummy}
+if has_dummy:
+    with open(r"{dummy_meta_path}", "r") as f:
+        meta = json.load(f)
+    npz = np.load(r"{dummy_path}")
+    t = meta["type"]
+    if t == "data":
+        kwargs = {{k: npz[k] for k in npz.files}}
+        dummy = Data(**kwargs)
+    elif t == "tuple":
+        dummy = tuple(npz[f"arr_{{i}}"] for i in range(len(npz.files)))
+    elif t == "dict":
+        dummy = {{k: npz[k] for k in npz.files}}
+    else:
+        dummy = npz["arr_0"]
+
+extra_kwargs = json.loads(r'''{json.dumps(kwargs)}''')
+if "{exporter_type}" == "onnx":
+    export_onnx(model_or_task, r"{out_file}", dummy_inputs=dummy, **extra_kwargs)
+elif "{exporter_type}" == "tflite":
+    export_tflite(model_or_task, r"{out_file}", dummy_inputs=dummy, **extra_kwargs)
+"""
+
+        env = dict(os.environ, KERAS_BACKEND="tensorflow")
+        res = subprocess.run(
+            [sys.executable, "-c", worker_code],
+            capture_output=True,
+            text=True,
+            env=env,
+        )
+
+        if res.returncode != 0:
+            raise RuntimeError(
+                f"Cross-backend export to {exporter_type.upper()} failed with exit code {res.returncode}:\n"
+                f"STDOUT: {res.stdout}\n"
+                f"STDERR: {res.stderr}"
+            )
+
+    return out_file

k3_node/export/onnx_exporter.py

@@ -0,0 +1,190 @@
+"""Turn-key ONNX exporter for K3-Node models and tasks."""
+
+import os
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple, Union
+import numpy as np
+
+import keras
+from keras import ops
+
+
+def _extract_model_and_inputs(
+    model_or_task: Any,
+    dummy_inputs: Optional[Any] = None,
+) -> Tuple[Any, Tuple[Any, ...], List[Any]]:
+    r"""Extracts the underlying neural network model and resolves dummy inputs / signatures."""
+    # 1. Resolve model
+    if hasattr(model_or_task, "model") and model_or_task.model is not None:
+        raw_model = model_or_task.model
+    else:
+        raw_model = model_or_task
+
+    # 2. Extract inputs from dummy_inputs if provided
+    if dummy_inputs is not None:
+        if hasattr(dummy_inputs, "x") and hasattr(dummy_inputs, "edge_index"):
+            inputs = [dummy_inputs.x, dummy_inputs.edge_index]
+            names = ["x", "edge_index"]
+            if hasattr(dummy_inputs, "batch") and dummy_inputs.batch is not None:
+                inputs.append(dummy_inputs.batch)
+                names.append("batch")
+            return raw_model, tuple(inputs), names
+        elif hasattr(dummy_inputs, "z") and hasattr(dummy_inputs, "pos"):
+            inputs = [dummy_inputs.z, dummy_inputs.pos]
+            names = ["z", "pos"]
+            if hasattr(dummy_inputs, "batch") and dummy_inputs.batch is not None:
+                inputs.append(dummy_inputs.batch)
+                names.append("batch")
+            return raw_model, tuple(inputs), names
+        elif isinstance(dummy_inputs, (tuple, list)):
+            return raw_model, tuple(dummy_inputs), None
+        elif isinstance(dummy_inputs, dict):
+            return raw_model, (dummy_inputs,), None
+        else:
+            return raw_model, (dummy_inputs,), None
+
+    # 3. Auto-infer dummy inputs from model hyperparameters
+    cls_name = model_or_task.__class__.__name__
+    in_channels = (
+        getattr(model_or_task, "in_channels", None)
+        or getattr(raw_model, "in_channels", None)
+        or 16
+    )
+
+    if cls_name == "NodeClassifier":
+        x = ops.zeros((4, in_channels), dtype="float32")
+        edge_index = ops.convert_to_tensor([[0, 1, 2, 3], [1, 2, 3, 0]], dtype="int64")
+        return raw_model, (x, edge_index), ["x", "edge_index"]
+    elif cls_name in ("GraphClassifier", "GraphRegressor"):
+        x = ops.zeros((4, in_channels), dtype="float32")
+        edge_index = ops.convert_to_tensor([[0, 1, 2, 3], [1, 2, 3, 0]], dtype="int64")
+        batch = ops.convert_to_tensor([0, 0, 1, 1], dtype="int64")
+        if hasattr(raw_model, "num_graphs"):
+            raw_model.num_graphs = 2
+        return raw_model, (x, edge_index, batch), ["x", "edge_index", "batch"]
+    elif cls_name == "LinkPredictor":
+        x = ops.zeros((4, in_channels), dtype="float32")
+        edge_index = ops.convert_to_tensor([[0, 1, 2, 3], [1, 2, 3, 0]], dtype="int64")
+        label_idx = ops.convert_to_tensor([[0, 1], [1, 2]], dtype="int64")
+        return raw_model, ((x, edge_index), label_idx), ["edge_tuple", "edge_label_index"]
+    elif cls_name in ("SchNet", "DimeNet", "DimeNetPlusPlus", "ViSNet", "GNNFF"):
+        z = ops.convert_to_tensor([1, 6, 8, 1], dtype="int32")
+        pos = ops.zeros((4, 3), dtype="float32")
+        return raw_model, (z, pos), ["z", "pos"]
+    else:
+        # Default standard GNN: (x, edge_index)
+        x = ops.zeros((4, in_channels), dtype="float32")
+        edge_index = ops.convert_to_tensor([[0, 1, 2, 3], [1, 2, 3, 0]], dtype="int64")
+        return raw_model, (x, edge_index), ["x", "edge_index"]
+
+
+def export_onnx(
+    model_or_task: Any,
+    output_path: Union[str, Path],
+    dummy_inputs: Optional[Any] = None,
+    opset: int = 17,
+    dynamic_axes: bool = True,
+    input_names: Optional[List[str]] = None,
+    output_names: Optional[List[str]] = None,
+    verbose: bool = False,
+) -> Path:
+    r"""Exports a K3-Node GNN model or task to high-performance ONNX format.
+
+    Supports arbitrary Graph Neural Networks (GCN, GAT, GraphSAGE, GIN, SchNet,
+    materials models, and task estimators) with dynamic graph sizing (varying numbers
+    of nodes and edges).
+
+    Args:
+        model_or_task: A K3-Node task instance (e.g. `NodeClassifier`, `GraphClassifier`)
+            or model instance (e.g. `GCN`, `SchNet`, `CHGNet`).
+        output_path: Target path for the `.onnx` file.
+        dummy_inputs: Optional sample input data (e.g., PyG `Data` object, tuple of tensors).
+            If `None`, automatically generated based on model topology.
+        opset: ONNX operator set version. (default: `17`)
+        dynamic_axes: Whether node and edge dimensions should be dynamic. (default: `True`)
+        input_names: Optional custom names for input tensors.
+        output_names: Optional custom names for output tensors.
+        verbose: Whether to print verbose export progress. (default: `False`)
+
+    Returns:
+        Path object pointing to the generated `.onnx` file.
+    """
+    out_file = Path(output_path)
+    out_file.parent.mkdir(parents=True, exist_ok=True)
+
+    from k3_node.export.cross_backend import is_tensorflow_backend, export_via_tf_subprocess
+
+    if not is_tensorflow_backend():
+        return export_via_tf_subprocess(
+            exporter_type="onnx",
+            model_or_task=model_or_task,
+            output_path=output_path,
+            dummy_inputs=dummy_inputs,
+            opset=opset,
+            dynamic_axes=dynamic_axes,
+            input_names=input_names,
+            output_names=output_names,
+            verbose=verbose,
+        )
+
+    try:
+        import tf2onnx
+        import tensorflow as tf
+    except ImportError:
+        raise ImportError(
+            "The `tf2onnx` and `tensorflow` packages are required for ONNX export. "
+            "Install them via `pip install tf2onnx tensorflow onnx`."
+        )
+
+    model, inputs, inferred_names = _extract_model_and_inputs(model_or_task, dummy_inputs)
+    names = input_names or inferred_names or [f"input_{i}" for i in range(len(inputs))]
+
+    # Build input signature with dynamic axes if requested
+    signature = []
+    for inp, name in zip(inputs, names):
+        inp_np = ops.convert_to_numpy(inp)
+        dtype = tf.as_dtype(inp_np.dtype)
+        if dynamic_axes:
+            if inp_np.ndim == 2 and inp_np.shape[0] == 2 and (
+                np.issubdtype(inp_np.dtype, np.integer) or "edge" in name.lower()
+            ):
+                # Edge index: (2, num_edges) -> dynamic num_edges
+                shape = (2, None)
+            elif inp_np.ndim == 2:
+                # Node feature: (num_nodes, in_channels) -> dynamic num_nodes
+                shape = (None, inp_np.shape[1])
+            elif inp_np.ndim == 1:
+                # Vector (batch or z): (num_nodes,) -> dynamic num_nodes
+                shape = (None,)
+            else:
+                shape = tuple(None if i == 0 else s for i, s in enumerate(inp_np.shape))
+        else:
+            shape = inp_np.shape
+        signature.append(tf.TensorSpec(shape=shape, dtype=dtype, name=name))
+
+    # Define trace function
+    @tf.function(input_signature=signature)
+    def forward_fn(*tensors):
+        return model(*tensors)
+
+    if verbose:
+        print(f"Exporting model to ONNX with input signature: {signature}")
+
+    # Convert using tf2onnx
+    model_proto, _ = tf2onnx.convert.from_function(
+        forward_fn,
+        input_signature=signature,
+        output_path=str(out_file),
+        opset=opset,
+    )
+
+    # Validate ONNX graph
+    try:
+        import onnx
+        onnx_model = onnx.load(str(out_file))
+        onnx.checker.check_model(onnx_model)
+    except Exception as e:
+        if verbose:
+            print(f"Warning: ONNX validation check returned: {e}")
+
+    return out_file

k3_node/export/runtime.py

@@ -0,0 +1,254 @@
+"""Lightweight inference runtime engines for serving ONNX and TFLite models."""
+
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Union
+import numpy as np
+
+
+class ONNXModel:
+    r"""High-performance serving wrapper for exported ONNX GNN models.
+
+    Requires only `onnxruntime` and `numpy`. Completely decoupled from Keras,
+    PyTorch, and TensorFlow for lightweight production microservices.
+
+    Example:
+        ```python
+        from k3_node.export import ONNXModel
+        model = ONNXModel("cora_gcn.onnx")
+        preds = model.predict(graph_data)
+        ```
+    """
+
+    def __init__(
+        self,
+        model_path: Union[str, Path],
+        providers: Optional[List[str]] = None,
+        session_options: Optional[Any] = None,
+    ):
+        r"""Initializes the ONNX runtime inference session.
+
+        Args:
+            model_path: Path to the `.onnx` model file.
+            providers: Execution providers list (e.g. `['CUDAExecutionProvider', 'CPUExecutionProvider']`).
+                If `None`, automatically picks the fastest available provider.
+            session_options: Optional custom ONNX Runtime SessionOptions.
+        """
+        try:
+            import onnxruntime as ort
+        except ImportError:
+            raise ImportError(
+                "The `onnxruntime` package is required to load and serve ONNX models. "
+                "Install it via `pip install onnxruntime` (or `onnxruntime-gpu` for CUDA/TensorRT)."
+            )
+
+        self.model_path = Path(model_path)
+        if not self.model_path.exists():
+            raise FileNotFoundError(f"ONNX model file not found at: {self.model_path}")
+
+        if providers is None:
+            available = ort.get_available_providers()
+            # Prioritize TensorRT, CUDA, CoreML, DirectML, CPU
+            priority = [
+                "TensorrtExecutionProvider",
+                "CUDAExecutionProvider",
+                "CoreMLExecutionProvider",
+                "DmlExecutionProvider",
+                "CPUExecutionProvider",
+            ]
+            providers = [p for p in priority if p in available]
+
+        self.session = ort.InferenceSession(
+            str(self.model_path),
+            sess_options=session_options,
+            providers=providers,
+        )
+        self.input_names = [inp.name for inp in self.session.get_inputs()]
+        self.output_names = [out.name for out in self.session.get_outputs()]
+
+    def predict(self, data: Any = None, *args: Any, **kwargs: Any) -> np.ndarray:
+        r"""Runs low-latency inference on graph data.
+
+        Accepts:
+            - PyG / K3-Node `Data` object
+            - Dictionary of tensors
+            - Positional numpy arrays
+
+        Args:
+            data: Input graph Data, dictionary, or array.
+            *args: Additional positional inputs.
+
+        Returns:
+            Numpy array containing model predictions or logits.
+        """
+        feed_dict = {}
+
+        # 1. PyG/K3 Data object
+        if hasattr(data, "x") and hasattr(data, "edge_index"):
+            feed_dict = self._match_inputs({
+                "x": np.asarray(data.x, dtype=np.float32),
+                "edge_index": np.asarray(data.edge_index, dtype=np.int64),
+                "batch": np.asarray(getattr(data, "batch", None), dtype=np.int64) if getattr(data, "batch", None) is not None else None,
+            })
+        elif hasattr(data, "z") and hasattr(data, "pos"):
+            feed_dict = self._match_inputs({
+                "z": np.asarray(data.z, dtype=np.int32),
+                "pos": np.asarray(data.pos, dtype=np.float32),
+                "batch": np.asarray(getattr(data, "batch", None), dtype=np.int32) if getattr(data, "batch", None) is not None else None,
+            })
+        elif isinstance(data, dict):
+            feed_dict = self._match_inputs(data)
+        elif isinstance(data, (tuple, list)):
+            for name, val in zip(self.input_names, data):
+                feed_dict[name] = np.asarray(val)
+        elif data is not None:
+            all_args = [data, *args]
+            for name, val in zip(self.input_names, all_args):
+                feed_dict[name] = np.asarray(val)
+
+        outputs = self.session.run(self.output_names, feed_dict)
+        return outputs[0] if len(outputs) == 1 else tuple(outputs)
+
+    def _match_inputs(self, named_inputs: Dict[str, Any]) -> Dict[str, np.ndarray]:
+        matched = {}
+        valid_inputs = {k: np.asarray(v) for k, v in named_inputs.items() if v is not None}
+        lower_inputs = {k.lower(): v for k, v in valid_inputs.items()}
+
+        used_keys = set()
+        session_inputs = self.session.get_inputs()
+
+        # Step 1: Match exact or known semantic names
+        for sess_inp in session_inputs:
+            name = sess_inp.name
+            clean = name.split(":")[0].lower()
+
+            target_val = None
+            matched_key = None
+
+            if clean in lower_inputs:
+                target_val = lower_inputs[clean]
+                matched_key = clean
+            elif "edge" in clean and "edge_index" in lower_inputs:
+                target_val = lower_inputs["edge_index"]
+                matched_key = "edge_index"
+            elif clean == "x" and "x" in lower_inputs:
+                target_val = lower_inputs["x"]
+                matched_key = "x"
+            elif clean in ("batch", "batch_idx") and "batch" in lower_inputs:
+                target_val = lower_inputs["batch"]
+                matched_key = "batch"
+            elif clean in ("z", "atomic_numbers") and "z" in lower_inputs:
+                target_val = lower_inputs["z"]
+                matched_key = "z"
+            elif clean in ("pos", "positions", "coord") and "pos" in lower_inputs:
+                target_val = lower_inputs["pos"]
+                matched_key = "pos"
+
+            if target_val is not None:
+                matched[name] = target_val
+                used_keys.add(matched_key)
+
+        # Step 2: Positional fallback for remaining inputs
+        unmatched_session = [inp for inp in session_inputs if inp.name not in matched]
+        unmatched_keys = [k for k in valid_inputs.keys() if k.lower() not in used_keys]
+
+        if unmatched_session and unmatched_keys:
+            for sess_inp, k in zip(unmatched_session, unmatched_keys):
+                matched[sess_inp.name] = valid_inputs[k]
+
+        # Step 3: Align data types with expected ONNX session tensor types
+        type_map = {
+            "tensor(float)": np.float32,
+            "tensor(float16)": np.float16,
+            "tensor(double)": np.float64,
+            "tensor(int64)": np.int64,
+            "tensor(int32)": np.int32,
+            "tensor(int8)": np.int8,
+            "tensor(uint8)": np.uint8,
+            "tensor(bool)": np.bool_,
+        }
+        for sess_inp in session_inputs:
+            if sess_inp.name in matched:
+                expected_np_type = type_map.get(sess_inp.type)
+                if expected_np_type and matched[sess_inp.name].dtype != expected_np_type:
+                    matched[sess_inp.name] = matched[sess_inp.name].astype(expected_np_type)
+
+        return matched
+
+
+class TFLiteModel:
+    r"""Lightweight serving wrapper for TensorFlow Lite flatbuffer GNN models.
+
+    Requires only standard `tensorflow` or `tflite_runtime`. Ideal for mobile,
+    Raspberry Pi, and edge embedded devices.
+
+    Example:
+        ```python
+        from k3_node.export import TFLiteModel
+        model = TFLiteModel("model.tflite")
+        preds = model.predict(graph_data)
+        ```
+    """
+
+    def __init__(self, model_path: Union[str, Path]):
+        r"""Initializes the TFLite interpreter.
+
+        Args:
+            model_path: Path to the `.tflite` model file.
+        """
+        self.model_path = Path(model_path)
+        if not self.model_path.exists():
+            raise FileNotFoundError(f"TFLite model file not found at: {self.model_path}")
+
+        try:
+            import tensorflow as tf
+            self.interpreter = tf.lite.Interpreter(model_path=str(self.model_path))
+        except ImportError:
+            try:
+                import tflite_runtime.interpreter as tflite
+                self.interpreter = tflite.Interpreter(model_path=str(self.model_path))
+            except ImportError:
+                raise ImportError(
+                    "Either `tensorflow` or `tflite_runtime` is required to run TFLite models. "
+                    "Install via `pip install tflite-runtime` or `pip install tensorflow`."
+                )
+
+        self.interpreter.allocate_tensors()
+        self.input_details = self.interpreter.get_input_details()
+        self.output_details = self.interpreter.get_output_details()
+
+    def predict(self, data: Any = None, *args: Any, **kwargs: Any) -> np.ndarray:
+        r"""Runs inference using the TFLite interpreter.
+
+        Args:
+            data: Input graph Data, dict, or numpy array.
+            *args: Additional positional inputs.
+
+        Returns:
+            Numpy array containing prediction results.
+        """
+        inputs = []
+        if hasattr(data, "x") and hasattr(data, "edge_index"):
+            inputs = [np.asarray(data.x, dtype=np.float32), np.asarray(data.edge_index, dtype=np.int64)]
+            if hasattr(data, "batch") and data.batch is not None:
+                inputs.append(np.asarray(data.batch, dtype=np.int64))
+        elif hasattr(data, "z") and hasattr(data, "pos"):
+            inputs = [np.asarray(data.z, dtype=np.int32), np.asarray(data.pos, dtype=np.float32)]
+            if hasattr(data, "batch") and data.batch is not None:
+                inputs.append(np.asarray(data.batch, dtype=np.int32))
+        elif isinstance(data, (tuple, list)):
+            inputs = [np.asarray(x) for x in data]
+        elif isinstance(data, dict):
+            inputs = [np.asarray(v) for v in data.values()]
+        elif data is not None:
+            inputs = [np.asarray(data)] + [np.asarray(a) for a in args]
+
+        # Feed tensors
+        for detail, inp in zip(self.input_details, inputs):
+            # Cast dtype to match interpreter expectation
+            target_dtype = detail["dtype"]
+            inp_cast = inp.astype(target_dtype) if inp.dtype != target_dtype else inp
+            self.interpreter.set_tensor(detail["index"], inp_cast)
+
+        self.interpreter.invoke()
+        outputs = [self.interpreter.get_tensor(d["index"]) for d in self.output_details]
+        return outputs[0] if len(outputs) == 1 else tuple(outputs)