LanguageMachines
/

blip2-flan-t5-xl

visual-question-answering

image-captioning

Model card Files Files and versions Community

sarang-shrivastava commited on Jun 28, 2023

Commit

1a9b1a2

·

1 Parent(s): 01e3d98

Update handler

Files changed (1) hide show

handler.py +2 -2

handler.py CHANGED Viewed

@@ -20,11 +20,11 @@ class EndpointHandler():
     def __init__(self, path=""):
         self.processor = Blip2Processor.from_pretrained(path)
-        self.model = Blip2ForConditionalGeneration.from_pretrained(path, load_in_8bit=True, device_map="auto")
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        # self.model.to(self.device)
         logging.info('Model moved to device-' + self.device)

     def __init__(self, path=""):
         self.processor = Blip2Processor.from_pretrained(path)
+        self.model = Blip2ForConditionalGeneration.from_pretrained(path, torch_dtype=torch.float16 , device_map="auto")
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model.to(self.device)
         logging.info('Model moved to device-' + self.device)