Added model architecture selection

app.py

@@ -26,8 +26,8 @@ def load_image_from_url(url: str) -> Image.Image:
     return Image.open(requests.get(url, stream=True).raw)
 
 @st.cache
-def load_model() -> ClipModel:
-    return ClipModel()
+def load_model(model_architecture: str) -> ClipModel:
+    return ClipModel(model_architecture)
 
 def init_state():
     if "images" not in st.session_state:
@@ -38,6 +38,8 @@ def init_state():
         st.session_state.predictions = None
     if "default_text_input" not in st.session_state:
         st.session_state.default_text_input = None
+    if "model_architecture" not in st.session_state:
+        st.session_state.model_architecture = "RN50"
 
 
 def limit_number_images():
@@ -278,7 +280,7 @@ if __name__ == "__main__":
     task_name: str = col2.selectbox("", options=["Prompt ranking", "Image ranking", "Image classification"])
     st.markdown("<br>", unsafe_allow_html=True)
     init_state()
-    model = load_model()
+    model = load_model(st.session_state.model_architecture)
     if task_name == "Image classification":
         Sections.image_uploader(accept_multiple_files=False)
         if st.session_state.images is None:
@@ -311,6 +313,10 @@ if __name__ == "__main__":
         limit_number_prompts()
         Sections.multiple_images_input_preview()
         Sections.classification_output(model)
+    
+    with st.expander("Advanced settings"):
+        st.session_state.model_architecture = st.selectbox("Model architecture", options=['RN50', 'RN101', 'RN50x4', 'RN50x16', 'RN50x64', 'ViT-B/32',
+         'ViT-B/16', 'ViT-L/14', 'ViT-L/14@336px'], index=0)
 
     st.markdown("<br><br><br><br>Made by [@JavierFnts](https://twitter.com/JavierFnts) | [How was CLIP Playground built?](https://twitter.com/JavierFnts/status/1363522529072214019)"
                 "", unsafe_allow_html=True)

images_mocker.py

@@ -1,31 +0,0 @@
-from typing import List
-import uuid
-from mock import patch
-
-
-class ImagesMocker:
-    """HACK ALERT: I needed a way to call the booste API without storing the images first
-     (as that is not allowed in streamlit sharing). If you have a better idea on hwo to this let me know!"""
-
-    def __init__(self):
-        self.pil_patch = patch('PIL.Image.open', lambda x: self.image_id2image(x))
-        self.path_patch = patch('os.path.exists', lambda x: True)
-        self.image_id2image_lookup = {}
-
-    def start_mocking(self):
-        self.pil_patch.start()
-        self.path_patch.start()
-
-    def stop_mocking(self):
-        self.pil_patch.stop()
-        self.path_patch.stop()
-
-    def image_id2image(self, image_id: str):
-        return self.image_id2image_lookup[image_id]
-
-    def calculate_image_id2image_lookup(self, images: List):
-        self.image_id2image_lookup = {str(uuid.uuid4()) + ".png": image for image in images}
-
-    @property
-    def image_ids(self):
-        return list(self.image_id2image_lookup.keys())

session_state.py

@@ -1,86 +0,0 @@
-# From https://gist.github.com/okld/0aba4869ba6fdc8d49132e6974e2e662
-from streamlit.hashing import _CodeHasher
-
-try:
-    # Before Streamlit 0.65
-    from streamlit.ReportThread import get_report_ctx
-    from streamlit.server.Server import Server
-except ModuleNotFoundError:
-    # After Streamlit 0.65
-    from streamlit.report_thread import get_report_ctx
-    from streamlit.server.server import Server
-
-
-class SessionState:
-    def __init__(self, session, hash_funcs):
-        """Initialize SessionState instance."""
-        self.__dict__["_state"] = {
-            "data": {},
-            "hash": None,
-            "hasher": _CodeHasher(hash_funcs),
-            "is_rerun": False,
-            "session": session,
-        }
-
-    def __call__(self, **kwargs):
-        """Initialize state data once."""
-        for item, value in kwargs.items():
-            if item not in self._state["data"]:
-                self._state["data"][item] = value
-
-    def __getitem__(self, item):
-        """Return a saved state value, None if item is undefined."""
-        return self._state["data"].get(item, None)
-
-    def __getattr__(self, item):
-        """Return a saved state value, None if item is undefined."""
-        return self._state["data"].get(item, None)
-
-    def __setitem__(self, item, value):
-        """Set state value."""
-        self._state["data"][item] = value
-
-    def __setattr__(self, item, value):
-        """Set state value."""
-        self._state["data"][item] = value
-
-    def clear(self):
-        """Clear session state and request a rerun."""
-        self._state["data"].clear()
-        self._state["session"].request_rerun()
-
-    def sync(self):
-        """Rerun the app with all state values up to date from the beginning to fix rollbacks."""
-
-        # Ensure to rerun only once to avoid infinite loops
-        # caused by a constantly changing state value at each run.
-        #
-        # Example: state.value += 1
-        if self._state["is_rerun"]:
-            self._state["is_rerun"] = False
-
-        elif self._state["hash"] is not None:
-            if self._state["hash"] != self._state["hasher"].to_bytes(self._state["data"], None):
-                self._state["is_rerun"] = True
-                self._state["session"].request_rerun()
-
-        self._state["hash"] = self._state["hasher"].to_bytes(self._state["data"], None)
-
-
-def get_session():
-    session_id = get_report_ctx().session_id
-    session_info = Server.get_current()._get_session_info(session_id)
-
-    if session_info is None:
-        raise RuntimeError("Couldn't get your Streamlit Session object.")
-
-    return session_info.session
-
-
-def get_state(hash_funcs=None):
-    session = get_session()
-
-    if not hasattr(session, "_custom_session_state"):
-        session._custom_session_state = SessionState(session, hash_funcs)
-
-    return session._custom_session_state