Update handler.py

handler.py

@@ -1,21 +1,23 @@
-from typing import Dict, Any
-from fastapi import FastAPI, File, UploadFile
-from fastapi.responses import StreamingResponse
+from typing import Dict, List, Any
 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):
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.processor = AutoProcessor.from_pretrained(model_path)
-        self.model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True, torch_dtype=torch.bfloat16).to(self.device)
+        self.processor = AutoProcessor.from_pretrained(
+            model_path,
+            # token=api_token
+        )
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_path,
+            # token=api_token,
+            trust_remote_code=True,
+            torch_dtype=torch.bfloat16,
+        ).to(self.device)
         self.image_seq_len = self.model.config.perceiver_config.resampler_n_latents
         self.bos_token = self.processor.tokenizer.bos_token
         self.bad_words_ids = self.processor.tokenizer(["<image>", "<fake_token_around_image>"], add_special_tokens=False).input_ids
@@ -42,32 +44,25 @@ class EndpointHandler:
         image = to_channel_dimension_format(image, ChannelDimension.FIRST)
         return torch.tensor(image)
 
-    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>",
-                return_tensors="pt",
-                add_special_tokens=False,
-            )
-            inputs["pixel_values"] = self.processor.image_processor([image], transform=self.custom_transform)
-            inputs = {k: v.to(self.device) for k, v in inputs.items()}
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        image = data.get("inputs")
 
-            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]
-            yield json.dumps({"label": generated_text, "score": 1.0}) + '\n'
-        
-        except torch.cuda.CudaError as e:
-            yield json.dumps({"error": f"CUDA error: {e}"}) + '\n'
-        except Exception as e:
-            yield json.dumps({"error": f"Unexpected error: {e}"}) + '\n'
+        if isinstance(image, str):
+            image = Image.open(image)
 
-handler = EndpointHandler(model_path="Jaykintecblic/Html")
+        inputs = self.processor.tokenizer(
+            f"{self.bos_token}<fake_token_around_image>{'<image>' * self.image_seq_len}<fake_token_around_image>",
+            return_tensors="pt",
+            add_special_tokens=False,
+        )
+        inputs["pixel_values"] = self.processor.image_processor([image], transform=self.custom_transform)
+        inputs = {k: v.to(self.device) for k, v in inputs.items()}
 
-@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")
+        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]
+        # print(generated_text)
+        # return {"text": generated_text}
+        # Format the output as an array of dictionaries with 'label' and 'score'
+        output = [{"label": generated_text, "score": 1.0}]
 
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8080)
+        return output