Refactor

app.py

@@ -59,7 +59,7 @@ def get_image_features(model, image_dir):
 
     loader = torch.utils.data.DataLoader(
         dataset,
-        batch_size=32,
+        batch_size=16,
         shuffle=False,
         num_workers=4,
         drop_last=False,
@@ -103,7 +103,8 @@ def text_encoder(text, tokenizer):
     return jnp.expand_dims(embedding, axis=0)
 
 
-def precompute_image_features(loader):
+@st.cache
+def precompute_image_features(model, loader):
     image_features = []
     for i, (images) in enumerate(tqdm(loader)):
         images = images.permute(0, 2, 3, 1).numpy()
@@ -145,8 +146,32 @@ if query:
             "dbmdz/bert-base-italian-xxl-uncased", cache_dir=None, use_fast=True
         )
 
-        image_features, dataset = get_image_features(model, "photos")
+        image_size = model.config.vision_config.image_size
+
+        val_preprocess = transforms.Compose(
+            [
+                Resize([image_size], interpolation=InterpolationMode.BICUBIC),
+                CenterCrop(image_size),
+                ToTensor(),
+                Normalize(
+                    (0.48145466, 0.4578275, 0.40821073),
+                    (0.26862954, 0.26130258, 0.27577711),
+                ),
+            ]
+        )
+
+        dataset = CustomDataSet("photos/", transform=val_preprocess)
+
+        loader = torch.utils.data.DataLoader(
+            dataset,
+            batch_size=16,
+            shuffle=False,
+            num_workers=2,
+            drop_last=False,
+        )
+
+        image_features = precompute_image_features(model, loader)
 
-        image_paths = find_image(query, dataset, tokenizer, image_features, n=3)
+        image_paths = find_image(query, dataset, tokenizer, image_features, n=2)
 
     st.image(image_paths)