Vivien commited on
Commit
b629ac6
1 Parent(s): da6fb88

Keep only 3 models to avoid RAM problems

Browse files
Files changed (1) hide show
  1. app.py +5 -5
app.py CHANGED
@@ -5,7 +5,7 @@ import pandas as pd, numpy as np
5
  from transformers import CLIPProcessor, CLIPModel
6
  from st_clickable_images import clickable_images
7
 
8
- MODEL_NAMES = ["base-patch32", "base-patch16", "large-patch14", "large-patch14-336"]
9
 
10
 
11
  @st.cache(allow_output_mutation=True)
@@ -165,10 +165,10 @@ def main():
165
  corpus = st.radio("", ["Unsplash", "Movies"])
166
 
167
  models_dict = {
168
- "ViT-B/32 (quickest)": "base-patch32",
169
- "ViT-B/16 (quick)": "base-patch16",
170
- "ViT-L/14 (slow)": "large-patch14",
171
- "ViT-L/14@336px (slowest)": "large-patch14-336",
172
  }
173
 
174
  if "Comparison" in mode:
 
5
  from transformers import CLIPProcessor, CLIPModel
6
  from st_clickable_images import clickable_images
7
 
8
+ MODEL_NAMES = ["base-patch32", "base-patch16", "large-patch14-336"] #, "large-patch14", "large-patch14-336"]
9
 
10
 
11
  @st.cache(allow_output_mutation=True)
 
165
  corpus = st.radio("", ["Unsplash", "Movies"])
166
 
167
  models_dict = {
168
+ "ViT-B/32 (small)": "base-patch32",
169
+ "ViT-B/16 (medium)": "base-patch16",
170
+ #"ViT-L/14 (slow)": "large-patch14",
171
+ "ViT-L/14@336px (large)": "large-patch14-336",
172
  }
173
 
174
  if "Comparison" in mode: