feat: hub deit model.

README.md

@@ -12,5 +12,4 @@ license: apache-2.0
 
 Attention Rollout was proposed by [Abnar et al.](https://arxiv.org/abs/2005.00928) to quantify the information
 that flows through self-attention layers. In the original ViT paper ([Dosovitskiy et al.](https://arxiv.org/abs/2010.11929)),
-the authors use it to investigate the representations learned by ViTs. The model used in the backend is a ViT B-16 model. For more
-details about it, refer to [this notebook](https://github.com/sayakpaul/probing-vits/blob/main/notebooks/load-jax-weights-vitb16.ipynb).
+the authors use it to investigate the representations learned by ViTs. The model used in the backend is `deit_tiny_patch16_224`. For more details about it, refer [here](https://tfhub.dev/sayakpaul/collections/deit/1). DeiT was proposed by [Touvron et al.](https://arxiv.org/abs/2012.12877)"

app.py

@@ -1,31 +1,51 @@
 import gradio as gr
-from huggingface_hub.keras_mixin import from_pretrained_keras
+import tensorflow as tf
+import tensorflow_hub as hub
 from PIL import Image
 
 import utils
 
-_MODEL = from_pretrained_keras("probing-vits/vit_b16_patch16_224_i21k_i1k")
+_RESOLUTION = 224
+_MODEL_URL = "https://tfhub.dev/sayakpaul/deit_tiny_patch16_224/1"
+
+
+def get_model() -> tf.keras.Model:
+    """Initiates a tf.keras.Model from TF-Hub."""
+    inputs = tf.keras.Input((_RESOLUTION, _RESOLUTION, 3))
+    hub_module = hub.KerasLayer(_MODEL_URL)
+
+    logits, attention_scores_dict = hub_module(
+        inputs
+    )  # Second output in the tuple is a dictionary containing attention scores.
+
+    return tf.keras.Model(inputs, [logits, attention_scores_dict])
+
+
+_MODEL = get_model()
 
 
 def show_rollout(image):
-    _, preprocessed_image = utils.preprocess_image(image, "original_vit")
+    """Function to be called when user hits submit on the UI."""
+    _, preprocessed_image = utils.preprocess_image(
+        image, "deit_tiny_patch16_224"
+    )
     _, attention_scores_dict = _MODEL.predict(preprocessed_image)
     result = utils.attention_rollout_map(
-        image, attention_scores_dict, "original_vit"
+        image, attention_scores_dict, "deit_tiny_patch16_224"
     )
     return Image.fromarray(result)
 
 
 title = "Generate Attention Rollout Plots"
-article = "Attention Rollout was proposed by [Abnar et al.](https://arxiv.org/abs/2005.00928) to quantify the information that flows through self-attention layers. In the original ViT paper ([Dosovitskiy et al.](https://arxiv.org/abs/2010.11929)), the authors use it to investigate the representations learned by ViTs. The model used in the backend is a ViT B-16 model. For more details about it, refer to [this notebook](https://github.com/sayakpaul/probing-vits/blob/main/notebooks/load-jax-weights-vitb16.ipynb)."
+article = "Attention Rollout was proposed by [Abnar et al.](https://arxiv.org/abs/2005.00928) to quantify the information that flows through self-attention layers. In the original ViT paper ([Dosovitskiy et al.](https://arxiv.org/abs/2010.11929)), the authors use it to investigate the representations learned by ViTs. The model used in the backend is `deit_tiny_patch16_224`. For more details about it, refer [here](https://tfhub.dev/sayakpaul/collections/deit/1). DeiT was proposed by [Touvron et al.](https://arxiv.org/abs/2012.12877)"
 
 iface = gr.Interface(
     show_rollout,
-    gr.inputs.Image(type="pil", label="Input Image"),
-    "image",
+    inputs=gr.inputs.Image(type="pil", label="Input Image"),
+    outputs="image",
     title=title,
     article=article,
     allow_flagging="never",
-    # examples=[["car.jpeg", "bulbul.jpeg"]]
+    examples=[["./car.jpeg", "./bulbul.jpeg"]],
 )
 iface.launch()

requirements.txt

@@ -1,4 +1,4 @@
 tensorflow
+tensorflow-hub
 opencv-python
-numpy
-huggingface_hub
+numpy