dev: add batching to litserve app

Author: ajndkr

litserve-app.py

@@ -17,11 +17,15 @@ def setup(self, device):
     def decode_request(self, request):
         image_bytes = bytes.fromhex(request["image_bytes"])
         image = Image.open(BytesIO(image_bytes))
-        return self.processor(image, return_tensors="pt").to(self.device)
+        image_tensor = self.processor(image, return_tensors="pt")["pixel_values"]
+        return image_tensor.to(self.device)
+
+    def batch(self, inputs):
+        return torch.cat(inputs, dim=0)
 
     def predict(self, inputs):
         with torch.no_grad():
-            logits = self.model(**inputs).logits
+            logits = self.model(inputs).logits
         return logits
 
     def encode_response(self, logits):
@@ -32,5 +36,5 @@ def encode_response(self, logits):
 
 if __name__ == "__main__":
     api = ResNetLitAPI()
-    server = ls.LitServer(api, accelerator="gpu")
+    server = ls.LitServer(api, accelerator="gpu", max_batch_size=8, batch_timeout=0.05)
     server.run(port=8000)