From 8fbc1431d394e8b7ab6a6cc450f4f6bf177ce953 Mon Sep 17 00:00:00 2001
From: DevTekVE <devtekve@gmail.com>
Date: Sun, 29 Dec 2024 17:40:45 +0100
Subject: [PATCH] Refactor model inference to use internal state for inputs

Simplified the `run_model` method by removing the requirement to pass inputs as arguments, and instead leveraging an internal `inputs` state. Adjusted `prepare_inputs` methods across model runners to populate this internal state. This refactor improves code clarity and reduces redundancy in managing input data.
---
 selfdrive/modeld/modeld.py               |  4 ++--
 selfdrive/modeld/runners/model_runner.py | 25 ++++++++++++------------
 2 files changed, 15 insertions(+), 14 deletions(-)

diff --git a/selfdrive/modeld/modeld.py b/selfdrive/modeld/modeld.py
index 7eac542969..3d6341dc29 100755
--- a/selfdrive/modeld/modeld.py
+++ b/selfdrive/modeld/modeld.py
@@ -79,13 +79,13 @@ class ModelState:
                'big_input_imgs': self.frames['big_input_imgs'].prepare(wbuf, transform_wide.flatten())}
 
     # Prepare inputs using the model runner
-    prepared_inputs = self.model_runner.prepare_inputs(imgs_cl, self.numpy_inputs)
+    self.model_runner.prepare_inputs(imgs_cl, self.numpy_inputs)
 
     if prepare_only:
       return None
 
     # Run model inference
-    self.output = self.model_runner.run_model(prepared_inputs)
+    self.output = self.model_runner.run_model()
     outputs = self.parser.parse_outputs(self.model_runner.slice_outputs(self.output))
 
     self.full_features_20Hz[:-1] = self.full_features_20Hz[1:]
diff --git a/selfdrive/modeld/runners/model_runner.py b/selfdrive/modeld/runners/model_runner.py
index f0f5360cba..67afe71522 100644
--- a/selfdrive/modeld/runners/model_runner.py
+++ b/selfdrive/modeld/runners/model_runner.py
@@ -31,13 +31,14 @@ class ModelRunner(ABC):
       self.model_metadata = pickle.load(f)
     self.input_shapes = self.model_metadata['input_shapes']
     self.output_slices = self.model_metadata['output_slices']
+    self.inputs = {}
 
   @abstractmethod
   def prepare_inputs(self, imgs_cl: dict[str, any], numpy_inputs: dict[str, np.ndarray]) -> dict[str, any]:
     """Prepare inputs for model inference."""
 
   @abstractmethod
-  def run_model(self, inputs: dict[str, any]) -> np.ndarray:
+  def run_model(self) -> np.ndarray:
     """Run model inference with prepared inputs."""
 
   def slice_outputs(self, model_outputs: np.ndarray) -> dict[str, np.ndarray]:
@@ -56,23 +57,22 @@ class TinyGradRunner(ModelRunner):
     # Load TinyGrad model
     with open(MODEL_PKL_PATH, "rb") as f:
       self.model_run = pickle.load(f)
-    self.tensor_inputs = {}
 
   def prepare_inputs(self, imgs_cl: dict[str, any], numpy_inputs: dict[str, np.ndarray]) -> dict[str, any]:
     # Initialize image tensors if not already done
     for key in imgs_cl:
-      if key not in self.tensor_inputs:
-        self.tensor_inputs[key] = qcom_tensor_from_opencl_address(imgs_cl[key].mem_address, self.input_shapes[key], dtype=dtypes.uint8)
+      if key not in self.inputs:
+        self.inputs[key] = qcom_tensor_from_opencl_address(imgs_cl[key].mem_address, self.input_shapes[key], dtype=dtypes.uint8)
 
     # Update numpy inputs
     for k, v in numpy_inputs.items():
-      if k not in self.tensor_inputs:
-        self.tensor_inputs[k] = Tensor(v, device='NPY').realize()
+      if k not in self.inputs:
+        self.inputs[k] = Tensor(v, device='NPY').realize()
 
-    return self.tensor_inputs
+    return self.inputs
 
-  def run_model(self, inputs: dict[str, any]) -> np.ndarray:
-    return self.model_run(**inputs).numpy().flatten()
+  def run_model(self) -> np.ndarray:
+    return self.model_run(**self.inputs).numpy().flatten()
 
 
 class ONNXRunner(ModelRunner):
@@ -86,7 +86,8 @@ class ONNXRunner(ModelRunner):
   def prepare_inputs(self, imgs_cl: dict[str, any], numpy_inputs: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
     for key in imgs_cl:
       numpy_inputs[key] = self.frames[key].buffer_from_cl(imgs_cl[key]).reshape(self.input_shapes[key])
-    return numpy_inputs
+    self.inputs = numpy_inputs
+    return self.inputs
 
-  def run_model(self, inputs: dict[str, any]) -> np.ndarray:
-    return self.runner.run(None, inputs)[0].flatten()
+  def run_model(self) -> np.ndarray:
+    return self.runner.run(None, self.inputs)[0].flatten()