fix inference engine

b13e368368d7439d23ed2b08d2fd9bd9f1903a05 · 2024-12-30 19:41:19 -0500 · Pranav Veldurthi

Files touched

M exo/inference/inference_engine.py
M exo/inference/tinygrad/inference.py

Diff

commit b13e368368d7439d23ed2b08d2fd9bd9f1903a05
Author: Pranav Veldurthi <veldurthipranav@gmail.com>
Date:   Mon Dec 30 19:41:19 2024 -0500

    fix inference engine
---
 exo/inference/inference_engine.py   | 3 +--
 exo/inference/tinygrad/inference.py | 6 +++---
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/exo/inference/inference_engine.py b/exo/inference/inference_engine.py
index 9bb21f07..97cd6aa5 100644
--- a/exo/inference/inference_engine.py
+++ b/exo/inference/inference_engine.py
@@ -43,10 +43,9 @@ class InferenceEngine(ABC):
     tokens = await self.encode(shard, prompt)
     if shard.model_id != 'stable-diffusion-2-1-base':
       x = tokens.reshape(1, -1)
-      output_data, inference_state = await self.infer_tensor(request_id, shard, x, inference_state)
     else:
       x = tokens
-      output_data, inference_state = await self.infer_tensor(request_id, shard, x, inference_state)
+    output_data, inference_state = await self.infer_tensor(request_id, shard, x, inference_state)
 
     return output_data, inference_state
 
diff --git a/exo/inference/tinygrad/inference.py b/exo/inference/tinygrad/inference.py
index 86abd76b..214cfd3d 100644
--- a/exo/inference/tinygrad/inference.py
+++ b/exo/inference/tinygrad/inference.py
@@ -15,7 +15,7 @@ from .stateful_model import make_prompt_state
 from .losses import length_masked_ce_loss
 from collections import OrderedDict
 import asyncio
-
+from typing import Optional
 Tensor.no_grad = True 
 # default settings
 TEMPERATURE = int(os.getenv("TEMPERATURE", 0.85))
@@ -104,7 +104,7 @@ class TinygradDynamicShardInferenceEngine(InferenceEngine):
     state_dict = await asyncio.get_running_loop().run_in_executor(self.executor, get_state_dict, self.model)
     safe_save(state_dict, path) 
   
-  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:
     await self.ensure_shard(shard)
     def wrap_infer():
       x = Tensor(input_data)
@@ -114,7 +114,7 @@ class TinygradDynamicShardInferenceEngine(InferenceEngine):
       self.states[request_id].start += x.shape[1]
       return out.realize()
     output_data = await asyncio.get_running_loop().run_in_executor(self.executor, wrap_infer)
-    return output_data.numpy()
+    return output_data.numpy(), inference_state
 
   async def evaluate(self, request_id: str, shard: Shard, inputs, targets, lengths, loss=length_masked_ce_loss):
     def step(x, y, l):

← 9986fb86 remove prints and fix download progress for SD · back to Exo · add trending badge to README.md 178cc4d9 →