Using hf_hub_download.

app.py

@@ -1,3 +1,4 @@
+import os 
 import matplotlib.pyplot as plt
 import matplotlib.cm as cm
 import numpy as np
@@ -5,7 +6,8 @@ import gradio as gr
 from transformers import AutoModel, AutoImageProcessor
 from PIL import Image
 import torch
-import os 
+from huggingface_hub import hf_hub_download
+
 
 
 os.environ["HF_HUB_OFFLINE"] = "0"
@@ -18,13 +20,6 @@ state = {
     "repo_id": None,
 }
 
-# Predefined supported models (must also exist locally in your Space repo)
-SUPPORTED_MODELS = {
-    "Google ViT-Base (patch16-224)": "./models/vit-base-patch16-224",
-    "Facebook DINO (ViT-S/16)": "./models/dino-vits16",
-    "OpenAI CLIP (ViT-B/32)": "./models/clip-vit-base-patch32",
-}
-
 
 def similarity_heatmap(image): 
     """
@@ -105,19 +100,40 @@ def load_model_dropdown(choice: str):
 
 def load_model(repo_id: str, revision: str = None):
     """
-    Load a Hugging Face model and processor from a repo ID.
+    Load a Hugging Face model + processor from Hub using huggingface_hub.
+    Works with any public repo_id.
     """
     try:
-        model = AutoModel.from_pretrained(repo_id, revision=revision, trust_remote_code=False)
-        processor = AutoImageProcessor.from_pretrained(repo_id, revision=revision, trust_remote_code=False)
-        # Move model to CPU/GPU if needed
+        # Explicitly download model + processor files to local cache
+        model_path = hf_hub_download(
+            repo_id=repo_id,
+            revision=revision,
+            filename="pytorch_model.bin",  # default filename for weights
+            cache_dir="./model_cache"
+        )
+        config_path = hf_hub_download(
+            repo_id=repo_id,
+            revision=revision,
+            filename="config.json",
+            cache_dir="./model_cache"
+        )
+        processor_path = hf_hub_download(
+            repo_id=repo_id,
+            revision=revision,
+            filename="preprocessor_config.json",
+            cache_dir="./model_cache"
+        )
+
+        # Load with transformers (it will reuse the local cache)
+        model = AutoModel.from_pretrained(repo_id, revision=revision, cache_dir="./model_cache")
+        processor = AutoImageProcessor.from_pretrained(repo_id, revision=revision, cache_dir="./model_cache")
+
         if torch.cuda.is_available():
             model.to("cuda")
         else:
             model.to("cpu")
 
-        model.eval() 
-        # Store in global state
+        model.eval()
         state["model"] = model
         state["processor"] = processor
         state["repo_id"] = repo_id
@@ -143,25 +159,34 @@ def visualize_cosine_heatmap(image: Image):
 # Gradio UI
 with gr.Blocks() as demo:
     gr.Markdown("# ViT CLS-Visualizer")
+    gr.Markdown(
+        "Enter the Hugging Face model repo ID (must be public), upload an image, "
+        "and visualize the cosine similarity between the CLS token and patches."
+    )
 
     with gr.Row():
-        model_choice = gr.Dropdown(
-            choices=list(SUPPORTED_MODELS.keys()),
-            label="Choose a Vision Transformer model",
-            value=list(SUPPORTED_MODELS.keys())[0],
+        repo_input = gr.Textbox(
+            label="Hugging Face Model Repo ID",
+            placeholder="e.g. google/vit-base-patch16-224"
+        )
+        revision_input = gr.Textbox(
+            label="Revision (optional)",
+            placeholder="branch, tag, or commit hash"
         )
         load_btn = gr.Button("Load Model")
     load_status = gr.Textbox(label="Model Status", interactive=False)
 
-    image_input = gr.Image(type="pil", label="Upload Image")
-    image_output = gr.Image(label="Uploaded Image")
-    heatmap_output = gr.Image(label="Cosine Similarity Heatmap")
+    with gr.Row():
+        image_input = gr.Image(type="pil", label="Upload Image")
+        image_output = gr.Image(label="Uploaded Image")
+
+    with gr.Row():
+        compute_btn = gr.Button("Compute Heatmap")
+        heatmap_output = gr.Image(label="Cosine Similarity Heatmap")
 
     # Events
-    load_btn.click(fn=load_model_dropdown, inputs=model_choice, outputs=load_status)
+    load_btn.click(fn=load_model, inputs=[repo_input, revision_input], outputs=load_status)
     image_input.change(fn=display_image, inputs=image_input, outputs=image_output)
-
-    compute_btn = gr.Button("Compute Heatmap")
     compute_btn.click(fn=visualize_cosine_heatmap, inputs=image_input, outputs=heatmap_output)
 
-demo.launch()
+demo.launch()