Changing weights and fixes

model_large_caption.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0d79b3b7c41478b5fe55c35b73ca6f3525a09708289371c6c0fac641e588287e
+size 1785411505

models/blip_decoder.py

@@ -8,8 +8,8 @@
 import warnings
 warnings.filterwarnings("ignore")
 
-from vit import VisionTransformer, interpolate_pos_embed
-from med import BertConfig, BertModel, BertLMHeadModel
+from models.vit import VisionTransformer, interpolate_pos_embed
+from models.med import BertConfig, BertModel, BertLMHeadModel
 from transformers import BertTokenizer
 
 import torch

pipeline.py

@@ -10,12 +10,11 @@ from torchvision import transforms
 from torchvision.transforms.functional import InterpolationMode
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-print(device)
 
 class PreTrainedPipeline():
-    def __init__(self):
+    def __init__(self, path=""):
         # load the optimized model
-        self.model_path = 'model_base_capfilt_large.pth'
+        self.model_path = 'model_large_caption.pth'
         self.model = blip_decoder(
             pretrained=self.model_path, 
             image_size=384, 
@@ -34,7 +33,7 @@ class PreTrainedPipeline():
      
 
 
-    def __call__(self, data: Any) -> Dict[str]:
+    def __call__(self, data: Any) -> Dict[str, Any]:
         """
         Args:
             data (:obj:):