Update app.py

app.py

@@ -5,14 +5,19 @@ import gradio as gr
 import numpy as np
 from PIL import Image, ImageDraw
 import torchvision.transforms.functional as TF
-from matplotlib import colormaps
-from transformers import AutoModel # 💡 Import AutoModel
+from matplotlib import colaps
+from transformers import AutoModel
 
 # ----------------------------
 # Configuration
 # ----------------------------
-# 💡 Use the full, correct model ID from the Hugging Face Hub
-MODEL_ID = "facebook/dinov3-vith16plus-pretrain-lvd1689m"
+# ⭐ Define available models, with the smaller one as default
+MODELS = {
+    "DINOv3 ViT-S+ (Small, Default)": "facebook/dinov3-vits16plus-pretrain-lvd1689m",
+    "DINOv3 ViT-H+ (Huge)": "facebook/dinov3-vith16plus-pretrain-lvd1689m",
+}
+DEFAULT_MODEL_NAME = "DINOv3 ViT-S+ (Small, Default)"
+
 PATCH_SIZE = 16
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -20,41 +25,41 @@ DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 IMAGENET_MEAN = (0.485, 0.456, 0.406)
 IMAGENET_STD = (0.229, 0.224, 0.225)
 
+# ⭐ Cache for loaded models to avoid re-downloading
+model_cache = {}
+
 # ----------------------------
 # Model Loading (Hugging Face Hub)
 # ----------------------------
-def load_model_from_hub():
-    """Loads the DINOv3 model from the Hugging Face Hub."""
-    print(f"Loading model '{MODEL_ID}' from Hugging Face Hub...")
+def load_model_from_hub(model_id: str):
+    """Loads a DINOv3 model from the Hugging Face Hub."""
+    print(f"Loading model '{model_id}' from Hugging Face Hub...")
     try:
-        # This will use the HF_TOKEN secret if you set it in your Space settings.
         token = os.environ.get("HF_TOKEN")
-        # trust_remote_code is necessary for DINOv3
-        model = AutoModel.from_pretrained(MODEL_ID, token=token, trust_remote_code=True)
+        model = AutoModel.from_pretrained(model_id, token=token, trust_remote_code=True)
         model.to(DEVICE).eval()
         print(f"✅ Model loaded successfully on device: {DEVICE}")
         return model
     except Exception as e:
         print(f"❌ Failed to load model: {e}")
-        # This will display a clear error message in the Gradio interface
         raise gr.Error(
-            f"Could not load model '{MODEL_ID}'. "
+            f"Could not load model '{model_id}'. "
             "This is a gated model. Please ensure you have accepted the terms on its Hugging Face page "
             "and set your HF_TOKEN as a secret in your Space settings. "
             f"Original error: {e}"
         )
 
-# Load the model globally when the app starts
-model = load_model_from_hub()
+def get_model(model_name: str):
+    """Gets a model from the cache or loads it if not present."""
+    model_id = MODELS[model_name]
+    if model_id not in model_cache:
+        model_cache[model_id] = load_model_from_hub(model_id)
+    return model_cache[model_id]
 
 # ----------------------------
-# Helper Functions (resize, viz)
+# Helper Functions (resize, viz) - No changes here
 # ----------------------------
 def resize_to_grid(img: Image.Image, long_side: int, patch: int) -> torch.Tensor:
-    """
-    Resizes so max(h,w)=long_side (keeping aspect), then rounds each side UP to a multiple of 'patch'.
-    Returns CHW float tensor in [0,1].
-    """
     w, h = img.size
     scale = long_side / max(h, w)
     new_h = max(patch, int(round(h * scale)))
@@ -110,7 +115,8 @@ def patch_neighborhood_box(r: int, c: int, Hp: int, Wp: int, rad: int, patch: in
 # Feature Extraction (using transformers)
 # ----------------------------
 @torch.inference_mode()
-def extract_image_features(image_pil: Image.Image, target_long_side: int):
+# ⭐ Pass the model object as an argument
+def extract_image_features(model, image_pil: Image.Image, target_long_side: int):
     """
     Extracts patch features from an image using the loaded Hugging Face model.
     """
@@ -118,23 +124,19 @@ def extract_image_features(image_pil: Image.Image, target_long_side: int):
     t_norm = TF.normalize(t, IMAGENET_MEAN, IMAGENET_STD).unsqueeze(0).to(DEVICE)
     _, _, H, W = t_norm.shape
     Hp, Wp = H // PATCH_SIZE, W // PATCH_SIZE
-    
-    # 💡 Use the standard forward pass of the transformers model
+
     outputs = model(t_norm)
-    
-    # 💡 The model output includes a [CLS] token AND 4 register tokens.
-    # We must skip all 5 to get only the patch embeddings.
-    n_special_tokens = 5 
+
+    n_special_tokens = 5
     patch_embeddings = outputs.last_hidden_state.squeeze(0)[n_special_tokens:, :]
-    
-    # L2-normalize the features to prepare for cosine similarity
+
     X = F.normalize(patch_embeddings, p=2, dim=-1)
 
     img_resized = TF.to_pil_image(t)
     return {"X": X, "Hp": Hp, "Wp": Wp, "img": img_resized}
 
 # ----------------------------
-# Similarity inside the same image (No changes needed here)
+# Similarity inside the same image - No changes here
 # ----------------------------
 def click_to_similarity_in_same_image(
     state: dict,
@@ -206,15 +208,21 @@ def click_to_similarity_in_same_image(
 # ----------------------------
 with gr.Blocks(theme=gr.themes.Soft(), title="DINOv3 Single-Image Patch Similarity") as demo:
     gr.Markdown("# 🦖 DINOv3 Single-Image Patch Similarity")
-    gr.Markdown("## Running on CPU-only Space, feature extraction after uploading an image can take a moment")
-    gr.Markdown("1. Upload an image. 2. Click **Process Image**. 3. **Click anywhere on the processed image** to find similar regions.")
+    gr.Markdown("## Running on CPU-only Space, feature extraction can take a moment")
+    gr.Markdown("1. **Choose a model**. 2. Upload an image. 3. Click **Process Image**. 4. **Click anywhere on the processed image** to find similar regions.")
 
     app_state = gr.State()
 
     with gr.Row():
         with gr.Column(scale=1):
+            # ⭐ ADDED MODEL DROPDOWN
+            model_name_dd = gr.Dropdown(
+                label="1. Choose a Model",
+                choices=list(MODELS.keys()),
+                value=DEFAULT_MODEL_NAME,
+            )
             input_image = gr.Image(
-                label="1. Upload Image",
+                label="2. Upload Image",
                 type="pil",
                 value="https://images.squarespace-cdn.com/content/v1/607f89e638219e13eee71b1e/1684821560422-SD5V37BAG28BURTLIXUQ/michael-sum-LEpfefQf4rU-unsplash.jpg"
             )
@@ -223,9 +231,8 @@ with gr.Blocks(theme=gr.themes.Soft(), title="DINOv3 Single-Image Patch Similari
                 label="Processing Resolution",
                 info="Higher values = more detail but slower processing",
             )
-            # ⭐️ ADDED BUTTON
-            process_button = gr.Button("2. Process Image", variant="primary")
-            
+            process_button = gr.Button("3. Process Image", variant="primary")
+
             with gr.Row():
                 alpha = gr.Slider(0.0, 1.0, value=0.55, step=0.05, label="Overlay opacity")
                 cmap = gr.Dropdown(
@@ -238,19 +245,24 @@ with gr.Blocks(theme=gr.themes.Soft(), title="DINOv3 Single-Image Patch Similari
             box_radius = gr.Slider(0, 10, value=1, step=1, label="Box radius (patches)")
 
     with gr.Row():
-        marked_image = gr.Image(label="3. Click on this image", interactive=True) # Changed from False to True for clarity
+        marked_image = gr.Image(label="4. Click on this image", interactive=True)
         heatmap_output = gr.Image(label="Similarity heatmap", interactive=False)
     with gr.Row():
         overlay_output = gr.Image(label="Overlay (image ⊕ heatmap)", interactive=False)
         overlay_boxes_output = gr.Image(label="Overlay + top-K similar patch boxes", interactive=False)
 
-    # ⭐️ RENAMED FUNCTION
-    def _process_image(img: Image.Image, long_side: int, progress=gr.Progress(track_tqdm=True)):
+    # ⭐ UPDATED to take model_name as input
+    def _process_image(model_name: str, img: Image.Image, long_side: int, progress=gr.Progress(track_tqdm=True)):
         if img is None:
             gr.Warning("Please upload an image first!")
             return None, None
-        progress(0, desc="Extracting features...")
-        st = extract_image_features(img, int(long_side))
+        
+        progress(0, desc=f"Loading model '{model_name}'...")
+        model = get_model(model_name)
+        
+        progress(0.5, desc="Extracting features...")
+        st = extract_image_features(model, img, int(long_side))
+        
         progress(1, desc="Done! You can now click on the image.")
         return st["img"], st
 
@@ -264,18 +276,16 @@ with gr.Blocks(theme=gr.themes.Soft(), title="DINOv3 Single-Image Patch Similari
             box_radius_patches=int(box_rad),
         )
 
-    # ⭐️ UPDATED EVENT WIRING
-    inputs_for_processing = [input_image, target_long_side]
+    # ⭐ UPDATED EVENT WIRING to include the dropdown
+    inputs_for_processing = [model_name_dd, input_image, target_long_side]
     outputs_for_processing = [marked_image, app_state]
-    
-    # The button now triggers the main processing function
+
     process_button.click(
-        _process_image, 
-        inputs=inputs_for_processing, 
+        _process_image,
+        inputs=inputs_for_processing,
         outputs=outputs_for_processing
     )
-    
-    # The click event on the image remains the same
+
     marked_image.select(
         _on_click,
         inputs=[app_state, alpha, cmap, exclude_r, topk, box_radius],