switch default to ViT-Base, keep multi_fps_k32 on ViT-Small, pre-load both

app.py

@@ -9,7 +9,7 @@ import os
 from typing import Tuple
 
 # --- Model Configuration ---
-DEFAULT_MODEL_NAME = "dino-vits-mae-100epoch-1217-1220-e50"
+DEFAULT_MODEL_NAME = "dino-vitb-mae-100epoch-1217-1220-e50"
 MODEL_CONFIGS = {
     "mars-ctx-vitb-0217": {
         "path": "models/0217-checkpoint-300.pth",
@@ -36,6 +36,12 @@ MODEL_CONFIGS = {
         "in_chans": 1,
         "description": "ViT-Small/16 DINO+MAE (Grayscale Input)"
     },
+    "dino-vitb-mae-100epoch-1217-1220-e50": {
+        "path": "models/dino-vitb-mae-100epoch-1217-1220-e50.pth",
+        "timm_id": "vit_base_patch16_224",
+        "in_chans": 1,
+        "description": "ViT-Base/16 DINO+MAE (Grayscale Input)"
+    },
 }
 
 # Global dictionary to store loaded models
@@ -90,15 +96,16 @@ def load_model(model_name: str):
     model.eval() # Set model to evaluation mode
     return model
 
-# --- Pre-load Default Model --- (Or load on demand in get_embedding)
-try:
-    print(f"Pre-loading default model: {DEFAULT_MODEL_NAME}...")
-    LOADED_MODELS[DEFAULT_MODEL_NAME] = load_model(DEFAULT_MODEL_NAME)
-    print(f"Default model {DEFAULT_MODEL_NAME} loaded successfully.")
-except Exception as e:
-    print(f"ERROR: Failed to pre-load default model {DEFAULT_MODEL_NAME}: {e}")
-    # Decide how to handle this - exit, or let Gradio fail later?
-    # For now, we'll print the error and continue; the app might fail if the default model is needed.
+# --- Pre-load Default Models ---
+MULTI_FPS_MODEL_NAME = "dino-vits-mae-100epoch-1217-1220-e50"
+
+for _name in [DEFAULT_MODEL_NAME, MULTI_FPS_MODEL_NAME]:
+    try:
+        print(f"Pre-loading model: {_name}...")
+        LOADED_MODELS[_name] = load_model(_name)
+        print(f"Model {_name} loaded successfully.")
+    except Exception as e:
+        print(f"ERROR: Failed to pre-load model {_name}: {e}")
 
 # --- Image Preprocessing --- (Now depends on model input channels)
 def get_preprocess(model_name: str):
@@ -306,14 +313,21 @@ def get_embedding(image_pil: Image.Image, model_name: str, embedding_method: str
 
             normalized_embedding = torch.nn.functional.normalize(embedding, p=2, dim=1)
 
-            # Compute multi-token FPS aggregation (32 tokens)
+            # Compute multi-token FPS aggregation (32 tokens) using ViT-Small model
             multi_fps_data = None
-            if len(features.shape) == 3 and features.shape[1] > 1:
-                patch_tokens = features[:, 1:]  # (B, num_patches, D)
-                num_patches = patch_tokens.shape[1]
-                k = min(32, num_patches)
+            if MULTI_FPS_MODEL_NAME not in LOADED_MODELS:
+                LOADED_MODELS[MULTI_FPS_MODEL_NAME] = load_model(MULTI_FPS_MODEL_NAME)
+            fps_model = LOADED_MODELS[MULTI_FPS_MODEL_NAME]
+            fps_preprocess = get_preprocess(MULTI_FPS_MODEL_NAME)
+            fps_tensor = fps_preprocess(image_pil).unsqueeze(0)
+            fps_features = fps_model.forward_features(fps_tensor)
+            if isinstance(fps_features, tuple):
+                fps_features = fps_features[0]
+            if len(fps_features.shape) == 3 and fps_features.shape[1] > 1:
+                fps_patch_tokens = fps_features[:, 1:]  # (B, num_patches, D)
+                k = min(32, fps_patch_tokens.shape[1])
                 if k > 0:
-                    agg_tokens = compute_multi_fps(patch_tokens, k=k)  # (B, K, D)
+                    agg_tokens = compute_multi_fps(fps_patch_tokens, k=k)  # (B, K, D)
                     multi_fps_data = agg_tokens.squeeze(0).cpu().numpy().tolist()
 
         embedding_list = normalized_embedding.squeeze().cpu().numpy().tolist()