Update app.py

app.py

@@ -11,9 +11,8 @@ import os
 # ----------------------------
 # Configuration
 # ----------------------------
-# The model will be downloaded from the Hugging Face Hub
-# Using the specific revision that works well with transformers AutoModel
-MODEL_ID = "facebook/dinov3-vith16plus"
+# 💡 FIX: Use the full, correct model ID from the Hugging Face Hub.
+MODEL_ID = "facebook/dinov3-vith16plus-pretrain-lvd1689m"
 PATCH_SIZE = 16
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -28,17 +27,22 @@ def load_model_from_hub():
     """Loads the DINOv3 model from the Hugging Face Hub."""
     print(f"Loading model '{MODEL_ID}' from Hugging Face Hub...")
     try:
-        # Use your HF token if the model is gated
-        # You can set this as a secret in your Hugging Face Space settings
+        # 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.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 an error message in the Gradio interface
-        raise gr.Error(f"Could not load model from Hub. If it's a gated model, ensure you have access and have set your HF_TOKEN secret in the Space settings. Error: {e}")
+        # This will display a clear error message in the Gradio interface
+        raise gr.Error(
+            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()
@@ -103,8 +107,7 @@ def generate_pca_visuals(
     
     # 💡 FIX: The model output includes a [CLS] token AND 4 register tokens.
     # We must skip all of them (total 5) to get only the patch embeddings.
-    # The original code only skipped 1, causing the size mismatch.
-    n_special_tokens = 5 # 1 [CLS] token + 4 register tokens
+    n_special_tokens = 5 # 1 [CLS] token + 4 register tokens for ViT-H/16+
     patch_embeddings = outputs.last_hidden_state.squeeze(0)[n_special_tokens:, :]
     
     # 3. PCA Calculation