Changes in the model loading

.gitignore

@@ -0,0 +1,4 @@
+collected/
+__pycache__/
+.idea/
+tests.ipynb

app.py

@@ -15,8 +15,8 @@ HF_API_TOKEN = os.getenv('HF_API_TOKEN')
 ENC_KEY = os.getenv('ENC_KEY')
 dataset_name = os.getenv('DATASET_NAME')
 ds_manager_queue = Queue(maxsize=1)
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-
+# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+device = torch.device('cpu')
 
 def push_files_async():
     try:
@@ -72,8 +72,10 @@ demo = gr.Interface(
 )
 
 if __name__ == "__main__":
+    vit_model = torch.load('vit_model_complete.pt')
+    vit_processor = torch.load('vit_processor_complete.pt')
     model_cls = load_enc_cls_model('model_scripted.pt_enc', ENC_KEY)
-    feat_extractor = VitLaionFeatureExtractor()
+    feat_extractor = VitLaionFeatureExtractor(vit_model, vit_processor)
     processor = feat_extractor.transforms
     ds_manager = HFPetDatasetManager(dataset_name, hf_token=HF_API_TOKEN, queue=ds_manager_queue)
     ds_manager.daemon = True

feat_ext.py

@@ -1,12 +1,11 @@
 import torch
-from transformers import AutoModel, AutoProcessor
 
 
 class VitLaionPreProcess(torch.nn.Module):
 
-    def __init__(self):
+    def __init__(self, processor):
         super().__init__()
-        self.processor = AutoProcessor.from_pretrained("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k")
+        self.processor = processor
 
     def forward(self, img):
         out = self.processor(images=img, return_tensors="pt")
@@ -14,10 +13,10 @@ class VitLaionPreProcess(torch.nn.Module):
 
 
 class VitLaionFeatureExtractor(torch.nn.Module):
-    def __init__(self):
+    def __init__(self, model, processor):
         super().__init__()
-        self.vit_model = AutoModel.from_pretrained("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k")
-        self.transforms = VitLaionPreProcess()
+        self.vit_model = model
+        self.transforms = VitLaionPreProcess(processor)
 
     def forward(self, x):
         img_a, img_b = x

vit_model_complete.pt

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

vit_processor_complete.pt

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