Upload DiskForKeypointDetection

config.json

@@ -2,6 +2,9 @@
   "architectures": [
     "DiskForKeypointDetection"
   ],
+  "auto_map": {
+    "AutoConfig": "configuration_disk.DiskConfig"
+  },
   "descriptor_decoder_dim": 128,
   "detection_threshold": 0.0,
   "max_num_keypoints": null,

configuration_disk.py

@@ -0,0 +1,30 @@
+from typing import Optional
+
+from transformers import PretrainedConfig
+
+
+class DiskConfig(PretrainedConfig):
+    model_type = "disk"
+
+    def __init__(
+        self,
+        weights: str = "depth",
+        max_num_keypoints: Optional[int] = None,
+        descriptor_decoder_dim: int = 128,
+        nms_window_size: int = 5,
+        detection_threshold: float = 0.0,
+        pad_if_not_divisible: bool = True,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.weights = weights
+        self.max_num_keypoints = max_num_keypoints
+        self.descriptor_decoder_dim = descriptor_decoder_dim
+        self.nms_window_size = nms_window_size
+        self.detection_threshold = detection_threshold
+        self.pad_if_not_divisible = pad_if_not_divisible
+
+
+if __name__ == "__main__":
+    config = DiskConfig()
+    config.save_pretrained("stevenbucaille/disk", push_to_hub=True)