Update handler.py
Browse files- handler.py +4 -3
handler.py
CHANGED
@@ -15,10 +15,11 @@ class EndpointHandler():
|
|
15 |
self.device = "cuda" if torch.cuda.is_available() else "cpu"
|
16 |
self.model_base = "Salesforce/blip2-opt-2.7b"
|
17 |
self.model_name = "sooh-j/blip2-vizwizqa"
|
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.
|
21 |
-
self.model =
|
|
|
22 |
|
23 |
self.device = "cuda" if torch.cuda.is_available() else "cpu"
|
24 |
self.model.to(self.device)
|
|
|
15 |
self.device = "cuda" if torch.cuda.is_available() else "cpu"
|
16 |
self.model_base = "Salesforce/blip2-opt-2.7b"
|
17 |
self.model_name = "sooh-j/blip2-vizwizqa"
|
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"
|
25 |
self.model.to(self.device)
|