Slice register tokens (#2)

- Slice register tokens (221f4602c71290ae688ef1d7856618da24ebbf6e)


Co-authored-by: Pavel Iakubovskii <[email protected]>

app.py

@@ -35,7 +35,8 @@ def extract_features(image, model_name):
     with torch.no_grad():
         outputs = model(**inputs)
         features = outputs.last_hidden_state
-    return features[:, 1:, :].float().cpu(), original_size, model_size
+    num_register_tokens = getattr(model.config, "num_register_tokens", 0)
+    return features[:, 1 + num_register_tokens:, :].float().cpu(), original_size, model_size
 
 def find_correspondences(features1, features2, threshold=0.8):
     device = torch.device("cpu")