hughtayloe commited on
Commit
e6b2e2d
1 Parent(s): 0dbf0cf

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +5 -4
handler.py CHANGED
@@ -7,21 +7,22 @@ from transformers import AutoProcessor, LlavaForConditionalGeneration
7
  class EndpointHandler():
8
  def __init__(self, path=""):
9
  model_id = ""
10
- model = LlavaForConditionalGeneration.from_pretrained(
11
  model_id,
12
  torch_dtype=torch.float16,
13
  low_cpu_mem_usage=True,
14
  ).to(0)
15
- processor = AutoProcessor.from_pretrained(model_id)
16
 
17
  def __call__(self, data: Dict[str, Any]):
18
  parameters = data.pop("inputs",data)
19
  inputs = data.pop("inputs", data)
20
  if parameters is not None:
21
  url = "http://images.cocodataset.org/val2017/000000039769.jpg"
22
- image = Image.open(requests.get(url, stream=True).raw)
 
23
  prompt = "USER: <image>\nWhat are these?\nASSISTANT:"
24
- output = model.generate(**inputs, max_new_tokens=200, do_sample=False)
25
  return output
26
 
27
 
 
7
  class EndpointHandler():
8
  def __init__(self, path=""):
9
  model_id = ""
10
+ self.model = LlavaForConditionalGeneration.from_pretrained(
11
  model_id,
12
  torch_dtype=torch.float16,
13
  low_cpu_mem_usage=True,
14
  ).to(0)
15
+ self.processor = AutoProcessor.from_pretrained(model_id)
16
 
17
  def __call__(self, data: Dict[str, Any]):
18
  parameters = data.pop("inputs",data)
19
  inputs = data.pop("inputs", data)
20
  if parameters is not None:
21
  url = "http://images.cocodataset.org/val2017/000000039769.jpg"
22
+ raw_image = Image.open(requests.get(url, stream=True).raw)
23
+ inputs = self.processor(prompt, raw_image, return_tensors='pt').to(0, torch.float16)
24
  prompt = "USER: <image>\nWhat are these?\nASSISTANT:"
25
+ output = self.model.generate(**inputs, max_new_tokens=200, do_sample=False)
26
  return output
27
 
28