handle inference_state properly

2aed3f3518ab7a0609c8f019141a91e7896798b5 · 2025-01-12 03:13:17 +0000 · Alex Cheema

Files touched

M exo/inference/inference_engine.py
M exo/inference/mlx/sharded_inference_engine.py

Diff

commit 2aed3f3518ab7a0609c8f019141a91e7896798b5
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Sun Jan 12 03:13:17 2025 +0000

    handle inference_state properly
---
 exo/inference/inference_engine.py             | 2 +-
 exo/inference/mlx/sharded_inference_engine.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/exo/inference/inference_engine.py b/exo/inference/inference_engine.py
index 97cd6aa5..d7d20ba9 100644
--- a/exo/inference/inference_engine.py
+++ b/exo/inference/inference_engine.py
@@ -23,7 +23,7 @@ class InferenceEngine(ABC):
     pass
 
   @abstractmethod
-  async def infer_tensor(self, request_id: str, shard: Shard, input_data: np.ndarray) -> np.ndarray:
+  async def infer_tensor(self, request_id: str, shard: Shard, input_data: np.ndarray, inference_state: Optional[dict] = None) -> np.ndarray:
     pass
 
   @abstractmethod
diff --git a/exo/inference/mlx/sharded_inference_engine.py b/exo/inference/mlx/sharded_inference_engine.py
index 7af64971..66db2fb2 100644
--- a/exo/inference/mlx/sharded_inference_engine.py
+++ b/exo/inference/mlx/sharded_inference_engine.py
@@ -77,7 +77,7 @@ class MLXDynamicShardInferenceEngine(InferenceEngine):
     await self.ensure_shard(shard)
     await asyncio.get_running_loop().run_in_executor(self.executor, self.model.load_weights, path)
     
-  async def infer_tensor(self, request_id: str, shard: Shard, input_data: np.ndarray, inference_state: Optional[dict] = None) -> np.ndarray:
+  async def infer_tensor(self, request_id: str, shard: Shard, input_data: np.ndarray, inference_state: dict = {}) -> np.ndarray:
     await self.ensure_shard(shard)
     loop = asyncio.get_running_loop()
     state = await self.poll_state(request_id) if self.model.model_type != 'StableDiffusionPipeline' else {}

← 2af5ee02 fix exo folder · back to Exo · ensure dir exists 9b2c01c8 →