sooh-j commited on
Commit
7d9abb7
1 Parent(s): 3324d89

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +1 -1
handler.py CHANGED
@@ -18,7 +18,7 @@ class EndpointHandler():
18
  # self.base_model = Blip2ForConditionalGeneration.from_pretrained(self.model_base, load_in_8bit=True)
19
  # self.pipe = Blip2ForConditionalGeneration.from_pretrained(self.model_base, load_in_8bit=True, torch_dtype=torch.float16)
20
  self.processor = Blip2Processor.from_pretrained(self.model_name)
21
- self.model = Blip2ForConditionalGeneration.from_pretrained(self.model_name).to(self.device)
22
  # self.model = PeftModel.from_pretrained(self.model_name, self.base_model_name).to(self.device)
23
 
24
  self.device = "cuda" if torch.cuda.is_available() else "cpu"
 
18
  # self.base_model = Blip2ForConditionalGeneration.from_pretrained(self.model_base, load_in_8bit=True)
19
  # self.pipe = Blip2ForConditionalGeneration.from_pretrained(self.model_base, load_in_8bit=True, torch_dtype=torch.float16)
20
  self.processor = Blip2Processor.from_pretrained(self.model_name)
21
+ self.model = BlipForQuestionAnswering.from_pretrained(self.model_name).to(self.device)
22
  # self.model = PeftModel.from_pretrained(self.model_name, self.base_model_name).to(self.device)
23
 
24
  self.device = "cuda" if torch.cuda.is_available() else "cpu"