ChirathD commited on
Commit
9e0a81f
1 Parent(s): 5f5cce7

Create handler.py

Browse files
Files changed (1) hide show
  1. handler.py +64 -0
handler.py ADDED
@@ -0,0 +1,64 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # +
2
+ from typing import Dict, List, Any
3
+ from PIL import Image
4
+ import torch
5
+ import os
6
+ import io
7
+ import base64
8
+ from io import BytesIO
9
+ # from transformers import BlipForConditionalGeneration, BlipProcessor
10
+ from transformers import Blip2Processor, Blip2ForConditionalGeneration
11
+
12
+
13
+ # -
14
+
15
+ # device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
16
+
17
+ class EndpointHandler():
18
+ def __init__(self, path=""):
19
+ # load the optimized model
20
+ print("####### Start Deploying #####")
21
+ self.processor = Blip2Processor.from_pretrained("ChirathD/Blip-2-test-1")
22
+ self.model = Blip2ForConditionalGeneration.from_pretrained("ChirathD/Blip-2-test-1")
23
+ # self.model.eval()
24
+ # self.model = self.model.to(device)
25
+
26
+
27
+
28
+ def __call__(self, data: Any) -> Dict[str, Any]:
29
+ """
30
+ Args:
31
+ data (:obj:):
32
+ includes the input data and the parameters for the inference.
33
+ Return:
34
+ A :obj:`dict`:. The object returned should be a dict of one list like {"captions": ["A hugging face at the office"]} containing :
35
+ - "caption": A string corresponding to the generated caption.
36
+ """
37
+ print(data)
38
+ inputs = data.pop("inputs", data)
39
+ parameters = data.pop("parameters", {})
40
+ print(input)
41
+ image_bytes = base64.b64decode(inputs)
42
+ image_io = io.BytesIO(image_bytes)
43
+ image = Image.open(image_io)
44
+
45
+ inputs = self.processor(images=image, return_tensors="pt")
46
+ pixel_values = inputs.pixel_values
47
+
48
+ generated_ids = self.model.generate(pixel_values=pixel_values, max_length=25)
49
+ generated_caption = self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
50
+ print(generated_caption)
51
+
52
+ # raw_images = [Image.open(BytesIO(_img)) for _img in inputs]
53
+
54
+ # processed_image = self.processor(images=raw_images, return_tensors="pt")
55
+ # processed_image["pixel_values"] = processed_image["pixel_values"].to(device)
56
+ # processed_image = {**processed_image, **parameters}
57
+
58
+ # with torch.no_grad():
59
+ # out = self.model.generate(
60
+ # **processed_image
61
+ # )
62
+ # captions = self.processor.batch_decode(out, skip_special_tokens=True)
63
+
64
+ return {"captions": generated_caption}