zesquirrelnator commited on
Commit
d7456e8
·
verified ·
1 Parent(s): 3efbe83

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +4 -3
handler.py CHANGED
@@ -4,14 +4,15 @@ from PIL import Image
4
  import torch
5
  import base64
6
  from io import BytesIO
7
- from transformers import BlipForConditionalGeneration, BlipProcessor
 
8
 
9
  device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
10
 
11
  class EndpointHandler():
12
  def __init__(self, path=""):
13
- self.processor = BlipProcessor.from_pretrained("zesquirrelnator/idefics2-8b-docvqa-finetuned-tutorial")
14
- self.model = BlipForConditionalGeneration.from_pretrained(
15
  "zesquirrelnator/idefics2-8b-docvqa-finetuned-tutorial"
16
  ).to(device)
17
  self.model.eval()
 
4
  import torch
5
  import base64
6
  from io import BytesIO
7
+ from transformers import AutoProcessor, AutoModelForVision2Seq
8
+ from transformers.image_utils import load_image
9
 
10
  device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
11
 
12
  class EndpointHandler():
13
  def __init__(self, path=""):
14
+ self.processor = AutoProcessor.from_pretrained("zesquirrelnator/idefics2-8b-docvqa-finetuned-tutorial")
15
+ self.model = AutoModelForVision2Seq.from_pretrained(
16
  "zesquirrelnator/idefics2-8b-docvqa-finetuned-tutorial"
17
  ).to(device)
18
  self.model.eval()