Update handler.py
Browse files- 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 =
|
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"
|