feat: loading from HF hub

app.py

@@ -1,20 +1,19 @@
+from huggingface_hub.keras_mixin import from_pretrained_keras
 import gradio as gr
 import tensorflow as tf
-import tensorflow_hub as hub
 from PIL import Image
 
 import utils
 
 _RESOLUTION = 224
-_MODEL_PATH = "gs://cait-tf/cait_xxs24_224"
 
 
 def get_model() -> tf.keras.Model:
-    """Initiates a tf.keras.Model from TF-Hub."""
+    """Initiates a tf.keras.Model from HF Hub."""
     inputs = tf.keras.Input((_RESOLUTION, _RESOLUTION, 3))
-    hub_module = hub.KerasLayer(_MODEL_PATH)
+    hub_module = from_pretrained_keras("probing-vits/cait_xxs24_224_classification")
 
-    logits, sa_atn_score_dict, ca_atn_score_dict = hub_module(inputs)
+    logits, sa_atn_score_dict, ca_atn_score_dict = hub_module(inputs, training=False)
 
     return tf.keras.Model(
         inputs, [logits, sa_atn_score_dict, ca_atn_score_dict]