SlowPacer commited on
Commit
1da5b95
1 Parent(s): d2d706d

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +2 -6
handler.py CHANGED
@@ -15,9 +15,7 @@ class EndpointHandler():
15
  # load the optimized model
16
 
17
  self.processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
18
- self.model = BlipForConditionalGeneration.from_pretrained(
19
- "Salesforce/blip-image-captioning-base"
20
- ).to(device)
21
  self.model.eval()
22
  self.model = self.model.to(device)
23
 
@@ -45,9 +43,7 @@ class EndpointHandler():
45
  processed_images = {**processed_images, **parameters}
46
 
47
  with torch.no_grad():
48
- out = self.model.generate(
49
- **processed_images
50
- )
51
  captions = self.processor.batch_decode(out, skip_special_tokens=True)
52
 
53
  return {"captions": captions}
 
15
  # load the optimized model
16
 
17
  self.processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
18
+ self.model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base").to(device)
 
 
19
  self.model.eval()
20
  self.model = self.model.to(device)
21
 
 
43
  processed_images = {**processed_images, **parameters}
44
 
45
  with torch.no_grad():
46
+ out = self.model.generate(**processed_images)
 
 
47
  captions = self.processor.batch_decode(out, skip_special_tokens=True)
48
 
49
  return {"captions": captions}