Upload folder using huggingface_hub

config.json

@@ -0,0 +1,16 @@
+{
+  "architectures": [
+    "EdgeFaceModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_edgeface.EdgeFaceConfig",
+    "AutoModel": "modeling_edgeface.EdgeFaceModel"
+  },
+  "dtype": "float32",
+  "featdim": 512,
+  "low_rank_ratio": 0.6,
+  "model_type": "edgeface",
+  "timm_model": "edgenext_base",
+  "transformers_version": "5.12.1",
+  "use_low_rank": false
+}

configuration_edgeface.py

@@ -0,0 +1,36 @@
+from transformers import PretrainedConfig
+
+
+class EdgeFaceConfig(PretrainedConfig):
+    """
+    Configuration for EdgeFace face-recognition models.
+
+    EdgeFace is a `timm` edgenext backbone with the classifier reset to output a
+    `featdim`-dimensional embedding. Some variants additionally replace their
+    nn.Linear layers with a static low-rank factorization (two smaller linears)
+    to cut parameters -- this is EdgeFace's "gamma" trick and is baked into the
+    weights. It is NOT PEFT/LoRA adapters; you can still train real LoRA on top
+    of the resulting model.
+
+    The four published variants map to:
+        edgeface_base        -> timm_model="edgenext_base",    use_low_rank=False
+        edgeface_s_gamma_05  -> timm_model="edgenext_small",   use_low_rank=True, low_rank_ratio=0.5
+        edgeface_xs_gamma_06 -> timm_model="edgenext_x_small", use_low_rank=True, low_rank_ratio=0.6
+        edgeface_xxs         -> timm_model="edgenext_xx_small", use_low_rank=False
+    """
+
+    model_type = "edgeface"
+
+    def __init__(
+        self,
+        timm_model: str = "edgenext_x_small",
+        featdim: int = 512,
+        use_low_rank: bool = False,
+        low_rank_ratio: float = 0.6,
+        **kwargs,
+    ):
+        self.timm_model = timm_model
+        self.featdim = featdim
+        self.use_low_rank = use_low_rank
+        self.low_rank_ratio = low_rank_ratio
+        super().__init__(**kwargs)

image_processing_edgeface.py

@@ -0,0 +1,305 @@
+"""
+Image processor for EdgeFace.
+
+Faithful port of the alignment used by the Idiap EdgeFace Space (utils.py):
+MediaPipe FaceMesh landmarks -> 5 points -> reflective similarity transform onto
+the ArcFace 112x112 template (custom MATLAB cp2tform-style solver).
+
+Works with both MediaPipe backends:
+  * "tasks"     -> latest API (mp.tasks.vision.FaceLandmarker + .task bundle)
+  * "solutions" -> legacy mp.solutions.face_mesh.FaceMesh (older installs)
+The default backend="auto" tries tasks first and falls back to solutions.
+
+Pipeline: (optional) align -> rescale to [0,1] -> normalize mean/std=0.5.
+If do_align=False the input is treated as an already-aligned crop and only
+resized to image_size.
+"""
+
+import os
+import weakref
+from typing import List, Optional, Union
+
+import numpy as np
+from numpy.linalg import inv, lstsq, matrix_rank, norm
+from transformers.image_processing_utils import BaseImageProcessor, BatchFeature
+from transformers.image_utils import ImageInput, make_list_of_images, to_numpy_array
+
+# ArcFace 5-point reference template for a 112x112 crop.
+# order matches the 5 source points: [reye, leye, nose, mouthright, mouthleft]
+REFERENCE_FACIAL_POINTS = np.array(
+    [
+        [38.2946, 51.6963],
+        [73.5318, 51.5014],
+        [56.0252, 71.7366],
+        [41.5493, 92.3655],
+        [70.7299, 92.2041],
+    ],
+    dtype=np.float32,
+)
+
+# MediaPipe FaceMesh indices (from the Space's utils.py). Valid for both the
+# 468-point legacy mesh and the 478-point tasks mesh (extra points are irises).
+IDX_REYE = (362, 263)   # eye on the image-left  (subject's right)
+IDX_LEYE = (33, 243)    # eye on the image-right (subject's left)
+IDX_NOSE = 1
+IDX_MOUTH_RIGHT = 287   # mouth corner on the image-left
+IDX_MOUTH_LEFT = 57     # mouth corner on the image-right
+
+# Official Google model bundle for the tasks API.
+_TASK_MODEL_URL = (
+    "https://storage.googleapis.com/mediapipe-models/face_landmarker/"
+    "face_landmarker/float16/1/face_landmarker.task"
+)
+
+# Live MediaPipe detectors are not JSON/deepcopy-safe, so keep them off the
+# instance __dict__ (which save_pretrained serializes) via a weak cache.
+_RUNTIME: "weakref.WeakKeyDictionary" = weakref.WeakKeyDictionary()
+
+
+# --------------------------------------------------------------------------
+# Similarity transform utilities (ported from the Space's utils.py)
+# --------------------------------------------------------------------------
+def _tformfwd(trans, uv):
+    uv_h = np.hstack((uv, np.ones((uv.shape[0], 1))))
+    xy = uv_h @ trans
+    return xy[:, :-1]
+
+
+def _find_nonreflective_similarity(uv, xy, K=2):
+    M = xy.shape[0]
+    x, y = xy[:, 0:1], xy[:, 1:2]
+    u, v = uv[:, 0:1], uv[:, 1:2]
+
+    X = np.vstack((
+        np.hstack((x, y, np.ones((M, 1)), np.zeros((M, 1)))),
+        np.hstack((y, -x, np.zeros((M, 1)), np.ones((M, 1)))),
+    ))
+    U = np.vstack((u, v))
+
+    if matrix_rank(X) >= 2 * K:
+        r, _, _, _ = lstsq(X, U, rcond=None)
+    else:
+        raise ValueError("cp2tform:twoUniquePointsReq")
+
+    sc, ss, tx, ty = r.flatten()
+    Tinv = np.array([[sc, -ss, 0], [ss, sc, 0], [tx, ty, 1]])
+    T = inv(Tinv)
+    T[:, 2] = [0, 0, 1]
+    return T, Tinv
+
+
+def _find_similarity(uv, xy):
+    trans1, trans1_inv = _find_nonreflective_similarity(uv, xy)
+
+    xyR = xy.copy()
+    xyR[:, 0] *= -1
+    trans2r, _ = _find_nonreflective_similarity(uv, xyR)
+    TreflectY = np.array([[-1, 0, 0], [0, 1, 0], [0, 0, 1]])
+    trans2 = trans2r @ TreflectY
+
+    norm1 = norm(_tformfwd(trans1, uv) - xy)
+    norm2 = norm(_tformfwd(trans2, uv) - xy)
+    return (trans1, trans1_inv) if norm1 <= norm2 else (trans2, inv(trans2))
+
+
+def _get_cv2_affine(src_pts, dst_pts):
+    trans, _ = _find_similarity(src_pts, dst_pts)
+    return trans[:, :2].T  # 2x3 for cv2.warpAffine
+
+
+def _warp_and_crop_face(src_img, facial_pts, reference_pts=REFERENCE_FACIAL_POINTS,
+                        crop_size=(112, 112), scale=1):
+    import cv2
+
+    ref_pts = reference_pts * scale
+    ref_pts = ref_pts + (np.mean(reference_pts, axis=0) - np.mean(ref_pts, axis=0))
+
+    src_pts = np.array(facial_pts, dtype=np.float32)
+    if src_pts.shape != ref_pts.shape:
+        raise ValueError("facial_pts and reference_pts must have the same shape")
+
+    tfm = _get_cv2_affine(src_pts, ref_pts)
+    return cv2.warpAffine(src_img, tfm, crop_size)
+
+
+class EdgeFaceImageProcessor(BaseImageProcessor):
+    model_input_names = ["pixel_values"]
+
+    def __init__(
+        self,
+        do_align: bool = True,
+        image_size: int = 112,
+        do_rescale: bool = True,
+        rescale_factor: float = 1 / 255,
+        do_normalize: bool = True,
+        image_mean: Optional[List[float]] = None,
+        image_std: Optional[List[float]] = None,
+        mp_backend: str = "auto",          # "auto" | "tasks" | "solutions"
+        mp_model_path: Optional[str] = None,  # path to a .task bundle (tasks backend)
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.do_align = do_align
+        self.image_size = image_size
+        self.do_rescale = do_rescale
+        self.rescale_factor = rescale_factor
+        self.do_normalize = do_normalize
+        self.image_mean = image_mean if image_mean is not None else [0.5, 0.5, 0.5]
+        self.image_std = image_std if image_std is not None else [0.5, 0.5, 0.5]
+        self.mp_backend = mp_backend
+        self.mp_model_path = mp_model_path
+
+    # -- runtime (non-serialized) cache ------------------------------------
+    def _runtime(self):
+        d = _RUNTIME.get(self)
+        if d is None:
+            d = {}
+            _RUNTIME[self] = d
+        return d
+
+    # -- model bundle for the tasks backend --------------------------------
+    def _resolve_model_path(self) -> str:
+        if self.mp_model_path:
+            return self.mp_model_path
+        env = os.environ.get("EDGEFACE_MP_MODEL")
+        if env:
+            return env
+        cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "edgeface")
+        os.makedirs(cache_dir, exist_ok=True)
+        path = os.path.join(cache_dir, "face_landmarker.task")
+        if not os.path.exists(path):
+            import urllib.request
+            urllib.request.urlretrieve(_TASK_MODEL_URL, path)
+        return path
+
+    # -- backend builders: each returns fn(rgb_uint8) -> (N,2) norm or None -
+    def _build_tasks_detector(self):
+        import mediapipe as mp
+        from mediapipe.tasks import python as mp_python
+        from mediapipe.tasks.python import vision as mp_vision
+
+        options = mp_vision.FaceLandmarkerOptions(
+            base_options=mp_python.BaseOptions(model_asset_path=self._resolve_model_path()),
+            running_mode=mp_vision.RunningMode.IMAGE,
+            num_faces=1,
+        )
+        landmarker = mp_vision.FaceLandmarker.create_from_options(options)
+
+        def detect(rgb):
+            mp_img = mp.Image(image_format=mp.ImageFormat.SRGB, data=np.ascontiguousarray(rgb))
+            res = landmarker.detect(mp_img)
+            if not res.face_landmarks:
+                return None
+            return np.array([[p.x, p.y] for p in res.face_landmarks[0]], dtype=np.float32)
+
+        return detect
+
+    def _build_solutions_detector(self):
+        import mediapipe as mp
+
+        face_mesh = mp.solutions.face_mesh.FaceMesh(
+            static_image_mode=True, refine_landmarks=True, min_detection_confidence=0.5,
+        )
+
+        def detect(rgb):
+            res = face_mesh.process(rgb)
+            if not res.multi_face_landmarks:
+                return None
+            return np.array([[p.x, p.y] for p in res.multi_face_landmarks[0].landmark],
+                            dtype=np.float32)
+
+        return detect
+
+    def _get_detect_fn(self):
+        runtime = self._runtime()
+        if "detect_fn" in runtime:
+            return runtime["detect_fn"]
+
+        order = {
+            "auto": ["tasks", "solutions"],
+            "tasks": ["tasks"],
+            "solutions": ["solutions"],
+        }.get(self.mp_backend)
+        if order is None:
+            raise ValueError(f"Unknown mp_backend={self.mp_backend!r}")
+
+        errors = []
+        for backend in order:
+            try:
+                fn = (self._build_tasks_detector() if backend == "tasks"
+                      else self._build_solutions_detector())
+                runtime["detect_fn"] = fn
+                return fn
+            except Exception as e:  # noqa: BLE001 - try next backend
+                errors.append(f"{backend}: {type(e).__name__}: {e}")
+
+        raise ImportError(
+            "Could not initialize a MediaPipe face detector. Install mediapipe "
+            "(`pip install mediapipe`) and ensure network access for the .task "
+            "bundle, or pass do_align=False / precomputed landmarks.\n"
+            + "\n".join(errors)
+        )
+
+    # -- landmark extraction -----------------------------------------------
+    def _detect_landmarks(self, image_rgb: np.ndarray) -> Optional[np.ndarray]:
+        """Return the 5 source points in [reye, leye, nose, mouthright, mouthleft] order."""
+        h, w = image_rgb.shape[:2]
+        norm_pts = self._get_detect_fn()(image_rgb)
+        if norm_pts is None or len(norm_pts) <= max(*IDX_REYE, *IDX_LEYE, IDX_MOUTH_RIGHT):
+            return None
+
+        px = norm_pts * np.array([w, h], dtype=np.float32)
+
+        reye = (px[IDX_REYE[0]] + px[IDX_REYE[1]]) / 2.0
+        leye = (px[IDX_LEYE[0]] + px[IDX_LEYE[1]]) / 2.0
+        return np.stack([reye, leye, px[IDX_NOSE], px[IDX_MOUTH_RIGHT], px[IDX_MOUTH_LEFT]]).astype(np.float32)
+
+    def _align_one(self, image_rgb: np.ndarray, landmarks: Optional[np.ndarray]) -> np.ndarray:
+        if landmarks is None:
+            landmarks = self._detect_landmarks(image_rgb)
+        if landmarks is None:
+            import cv2  # detection failed -> plain resize so the batch still runs
+            return cv2.resize(image_rgb, (self.image_size, self.image_size))
+        return _warp_and_crop_face(image_rgb, landmarks, crop_size=(self.image_size, self.image_size))
+
+    # -- main entry point --------------------------------------------------
+    def preprocess(
+        self,
+        images: ImageInput,
+        do_align: Optional[bool] = None,
+        landmarks: Optional[Union[np.ndarray, List[np.ndarray]]] = None,
+        return_tensors: Optional[str] = "pt",
+        **kwargs,
+    ) -> BatchFeature:
+        do_align = self.do_align if do_align is None else do_align
+        images = make_list_of_images(images)
+
+        if landmarks is not None and not isinstance(landmarks, list):
+            landmarks = [landmarks]
+
+        processed = []
+        for i, img in enumerate(images):
+            arr = to_numpy_array(img)  # RGB, HxWxC
+            if arr.ndim == 2:
+                arr = np.stack([arr] * 3, axis=-1)
+            if arr.shape[-1] == 4:
+                arr = arr[..., :3]
+            arr = arr.astype(np.uint8)
+
+            if do_align:
+                lmk = landmarks[i] if landmarks is not None else None
+                arr = self._align_one(arr, lmk)
+            else:
+                import cv2
+                arr = cv2.resize(arr, (self.image_size, self.image_size))
+
+            arr = arr.astype(np.float32)
+            if self.do_rescale:
+                arr = arr * self.rescale_factor
+            if self.do_normalize:
+                arr = (arr - np.array(self.image_mean)) / np.array(self.image_std)
+
+            processed.append(arr.transpose(2, 0, 1))  # CxHxW
+
+        pixel_values = np.stack(processed, axis=0).astype(np.float32)
+        return BatchFeature(data={"pixel_values": pixel_values}, tensor_type=return_tensors)

modeling_edgeface.py

@@ -0,0 +1,99 @@
+from dataclasses import dataclass
+from typing import Optional
+
+import timm
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+from transformers.utils import ModelOutput
+
+try:  # relative import works as Hub remote code; absolute works for local scripts
+    from .configuration_edgeface import EdgeFaceConfig
+except ImportError:
+    from configuration_edgeface import EdgeFaceConfig
+
+
+# ---------------------------------------------------------------------------
+# Static low-rank linear factorization (EdgeFace's "gamma" trick).
+# This is baked into the pretrained weights and is unrelated to PEFT/LoRA
+# adapters -- naming it LowRankLinear keeps "lora" free for real adapters.
+# The submodule attribute names (linear1, linear2) are kept so the original
+# published checkpoints load unchanged.
+# ---------------------------------------------------------------------------
+class LowRankLinear(nn.Module):
+    def __init__(self, in_features, out_features, rank, bias=True):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.rank = rank
+        self.linear1 = nn.Linear(in_features, rank, bias=False)
+        self.linear2 = nn.Linear(rank, out_features, bias=bias)
+
+    def forward(self, x):
+        return self.linear2(self.linear1(x))
+
+
+def _factorize_recursive(module, ratio):
+    for name, child in module.named_children():
+        if isinstance(child, nn.Linear) and "head" not in name:
+            rank = max(2, int(min(child.in_features, child.out_features) * ratio))
+            bias = child.bias is not None
+            setattr(module, name, LowRankLinear(child.in_features, child.out_features, rank, bias))
+        else:
+            _factorize_recursive(child, ratio)
+
+
+def factorize_linear_layers(module, ratio=0.2):
+    """Replace eligible nn.Linear layers with LowRankLinear, in place."""
+    _factorize_recursive(module, ratio)
+    return module
+
+
+@dataclass
+class EdgeFaceOutput(ModelOutput):
+    embeddings: Optional[torch.FloatTensor] = None
+
+
+class EdgeFaceModel(PreTrainedModel):
+    config_class = EdgeFaceConfig
+    main_input_name = "pixel_values"
+
+    def __init__(self, config: EdgeFaceConfig):
+        super().__init__(config)
+
+        # Keep the attribute named `self.model` so the original published
+        # checkpoints (keys prefixed with "model.") load unchanged.
+        self.model = timm.create_model(config.timm_model)
+        self.model.reset_classifier(config.featdim)
+
+        if config.use_low_rank:
+            factorize_linear_layers(self.model, ratio=config.low_rank_ratio)
+
+        self.post_init()
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            nn.init.trunc_normal_(module.weight, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+
+    def forward(
+        self,
+        pixel_values: torch.FloatTensor,
+        normalize: bool = False,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+        embeddings = self.model(pixel_values)
+
+        # Reference code does not normalize inside the model (it normalizes at
+        # comparison time via cosine similarity). Off by default for parity.
+        if normalize:
+            embeddings = F.normalize(embeddings, dim=-1)
+
+        if not return_dict:
+            return (embeddings,)
+        return EdgeFaceOutput(embeddings=embeddings)

preprocessor_config.json

@@ -0,0 +1,22 @@
+{
+  "auto_map": {
+    "AutoImageProcessor": "image_processing_edgeface.EdgeFaceImageProcessor"
+  },
+  "do_align": true,
+  "do_normalize": true,
+  "do_rescale": true,
+  "image_mean": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "image_processor_type": "EdgeFaceImageProcessor",
+  "image_size": 112,
+  "image_std": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "mp_backend": "auto",
+  "rescale_factor": 0.00392156862745098
+}