Update handler.py

handler.py

@@ -1,9 +1,15 @@
-from typing import Dict, Any, List
+from typing import Dict, Any
+from fastapi import FastAPI, File, UploadFile
+from fastapi.responses import StreamingResponse
 from PIL import Image
 import torch
 from transformers import AutoModelForCausalLM, AutoProcessor
 from transformers.image_utils import to_numpy_array, PILImageResampling, ChannelDimension
 from transformers.image_transforms import resize, to_channel_dimension_format
+import json
+import io
+
+app = FastAPI()
 
 class EndpointHandler:
     def __init__(self, model_path: str):
@@ -36,17 +42,7 @@ class EndpointHandler:
         image = to_channel_dimension_format(image, ChannelDimension.FIRST)
         return torch.tensor(image)
 
-    def generate_responses(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
-        results = []
-        image = data.get("inputs")
-
-        if isinstance(image, str):
-            try:
-                image = Image.open(image)
-            except Exception as e:
-                results.append({"error": f"Failed to open image: {e}"})
-                return results
-
+    async def generate_responses(self, image: Image.Image):
         try:
             inputs = self.processor.tokenizer(
                 f"{self.bos_token}<fake_token_around_image>{'<image>' * self.image_seq_len}<fake_token_around_image>",
@@ -58,14 +54,20 @@ class EndpointHandler:
 
             generated_ids = self.model.generate(**inputs, bad_words_ids=self.bad_words_ids, max_length=2048, early_stopping=True)
             generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-            results.append({"label": generated_text, "score": 1.0})
+            yield json.dumps({"label": generated_text, "score": 1.0}) + '\n'
         
         except torch.cuda.CudaError as e:
-            results.append({"error": f"CUDA error: {e}"})
+            yield json.dumps({"error": f"CUDA error: {e}"}) + '\n'
         except Exception as e:
-            results.append({"error": f"Unexpected error: {e}"})
+            yield json.dumps({"error": f"Unexpected error: {e}"}) + '\n'
+
+handler = EndpointHandler(model_path="path/to/your/model")
 
-        return results
+@app.post("/")
+async def handle_request(file: UploadFile = File(...)):
+    image = Image.open(io.BytesIO(await file.read()))
+    return StreamingResponse(handler.generate_responses(image), media_type="application/json")
 
-    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
-        return self.generate_responses(data)
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8080)