LanguageMachines
/

blip2-flan-t5-xxl

visual-question-answering

image-captioning

Inference Endpoints

Model card Files Files and versions Community

sarang-shrivastava commited on Jun 28, 2023

Commit

6f1c35a

•

1 Parent(s): 91c82aa

Update handler

Files changed (1) hide show

handler.py +1 -1

handler.py CHANGED Viewed

@@ -20,7 +20,7 @@ class EndpointHandler():
     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"

     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"