Added HF_HUB_OFFLINE env variable.

app.py

@@ -5,8 +5,12 @@ import gradio as gr
 from transformers import AutoModel, AutoImageProcessor
 from PIL import Image
 import torch
+import os 
 
-# Global state to store loaded model + processor
+
+os.environ["HF_HUB_OFFLINE"] = "0"
+
+# Global state to store loaded model + processors
 state = {
     "model_type": None, 
     "model": None,
@@ -14,6 +18,13 @@ 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): 
     """
@@ -72,6 +83,26 @@ def overlay_cosine_grid_on_image(cos_grid: np.ndarray, image: Image.Image, alpha
 
 
 
+def load_model_dropdown(choice: str):
+    """
+    Load one of the predefined models.
+    """
+    repo_path = SUPPORTED_MODELS[choice]
+    try:
+        model = AutoModel.from_pretrained(repo_path)
+        processor = AutoImageProcessor.from_pretrained(repo_path)
+        model.to("cuda" if torch.cuda.is_available() else "cpu")
+        model.eval()
+
+        state["model"] = model
+        state["processor"] = processor
+        state["repo_id"] = choice
+        return f"Successfully loaded model: {choice}"
+    except Exception as e:
+        return f"Error loading model {choice}: {e}"
+
+
+
 def load_model(repo_id: str, revision: str = None):
     """
     Load a Hugging Face model and processor from a repo ID.
@@ -108,30 +139,29 @@ def visualize_cosine_heatmap(image: Image):
     blended = overlay_cosine_grid_on_image(cos_grid, image)
     return blended
 
-# Build the Gradio interface
+
+# Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("# Dynamic ViT Loader Template")
-    
-    # TODO: Add drop-down menu (or something else) for user to allow choosing model type (e.g. DINOv2, Google ViT-Base etc.) 
-    # ...
+    gr.Markdown("# ViT CLS-Visualizer")
 
     with gr.Row():
-        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")
+        model_choice = gr.Dropdown(
+            choices=list(SUPPORTED_MODELS.keys()),
+            label="Choose a Vision Transformer model",
+            value=list(SUPPORTED_MODELS.keys())[0],
+        )
         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="Displayed Image")
 
-    # cos-sim visualization: 
+    image_input = gr.Image(type="pil", label="Upload Image")
+    image_output = gr.Image(label="Uploaded Image")
     heatmap_output = gr.Image(label="Cosine Similarity Heatmap")
 
-    # Button clicks / image upload handlers
-    load_btn.click(fn=load_model, inputs=[repo_input, revision_input], outputs=load_status)
+    # Events
+    load_btn.click(fn=load_model_dropdown, inputs=model_choice, 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()