Gabriel commited on
Commit
983cc5c
1 Parent(s): eacc8d0

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +6 -12
handler.py CHANGED
@@ -1,5 +1,4 @@
1
  from typing import Dict, Any
2
- import torch
3
  from transformers import AutoProcessor, Qwen2VLForConditionalGeneration
4
  from PIL import Image
5
  import io
@@ -8,21 +7,16 @@ import requests
8
 
9
  class EndpointHandler():
10
  def __init__(self, path=""):
11
- self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
12
- self.model = Qwen2VLForConditionalGeneration.from_pretrained(path).to(self.device)
13
  self.processor = AutoProcessor.from_pretrained(path)
 
14
 
15
  def __call__(self, data: Any) -> Dict[str, Any]:
16
- default_prompt = "Describe this image."
 
 
17
 
18
- if isinstance(data, (bytes, bytearray)):
19
- image = Image.open(io.BytesIO(data)).convert('RGB')
20
- text_input = default_prompt
21
- elif isinstance(data, dict):
22
- image_input = data.get('image', None)
23
- text_input = data.get('text', default_prompt)
24
- if image_input is None:
25
- return {"error": "No image provided."}
26
  if image_input.startswith('http'):
27
  image = Image.open(requests.get(image_input, stream=True).raw).convert('RGB')
28
  else:
 
1
  from typing import Dict, Any
 
2
  from transformers import AutoProcessor, Qwen2VLForConditionalGeneration
3
  from PIL import Image
4
  import io
 
7
 
8
  class EndpointHandler():
9
  def __init__(self, path=""):
 
 
10
  self.processor = AutoProcessor.from_pretrained(path)
11
+ self.model = Qwen2VLForConditionalGeneration.from_pretrained(path)
12
 
13
  def __call__(self, data: Any) -> Dict[str, Any]:
14
+
15
+ image_input = data.get('image', None)
16
+ text_input = data.get('text', None)
17
 
18
+
19
+ if isinstance(data, dict):
 
 
 
 
 
 
20
  if image_input.startswith('http'):
21
  image = Image.open(requests.get(image_input, stream=True).raw).convert('RGB')
22
  else: