Update handler.py
Browse files- handler.py +4 -4
handler.py
CHANGED
@@ -12,16 +12,16 @@ from transformers import Blip2Processor, Blip2ForConditionalGeneration
|
|
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 |
|
@@ -42,7 +42,7 @@ class EndpointHandler():
|
|
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)
|
|
|
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 |
|
|
|
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)
|