Spaces:
Paused
Paused
File size: 2,934 Bytes
4db294e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 | """Loading the FaceAnything model.
FaceAnything = Depth-Anything-3 (DA3-GIANT) backbone + a lightweight DPT
``deformation_head`` that predicts, per pixel, a 3D coordinate in a shared
canonical facial space (channels 0-2) plus a confidence (channel 3). The head is
attached at construction time and the finetuned weights are loaded from the
released checkpoint.
"""
from __future__ import annotations
import os
import torch
DEFAULT_BASE_MODEL = "depth-anything/DA3-GIANT-1.1"
GIANT_FEATURE_DIM = 3072
def load_model(checkpoint_path: str,
base_model: str = DEFAULT_BASE_MODEL,
device: str = "cuda",
feature_dim: int = GIANT_FEATURE_DIM,
verbose: bool = True):
"""Build the FaceAnything model and load the finetuned checkpoint.
Args:
checkpoint_path: path to ``checkpoint.pt`` (dict with a ``"model"`` key,
or a bare state-dict).
base_model: HuggingFace id of the DA3 backbone used for the architecture.
The backbone weights are overwritten by the checkpoint, but the
config is needed to build the network. Set ``HF_HOME`` to use a local
cache and avoid a download.
device: torch device string.
feature_dim: backbone feature dimension feeding the deformation head
(3072 for DA3-GIANT).
Returns:
A ``DepthAnything3`` model in eval mode on ``device``.
"""
from depth_anything_3.api import DepthAnything3
from depth_anything_3.model import dpt
if verbose:
print(f"[faceanything] building backbone from '{base_model}' ...", flush=True)
model = DepthAnything3.from_pretrained(base_model)
# Canonical / deformation head (3 coord channels + 1 confidence, no activation).
model.model.deformation_head = dpt.DPT(
feature_dim, output_dim=4, head_name="deformation",
use_sky_head=False, activation="linear",
)
if verbose:
print(f"[faceanything] loading checkpoint '{checkpoint_path}' ...", flush=True)
if not os.path.exists(checkpoint_path):
raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")
checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
state_dict = checkpoint["model"] if isinstance(checkpoint, dict) and "model" in checkpoint else checkpoint
missing, unexpected = model.load_state_dict(state_dict, strict=False)
if verbose:
miss = [m for m in missing if "deformation_head" in m]
print(f"[faceanything] loaded. missing={len(missing)} "
f"(deformation_head missing={len(miss)}), unexpected={len(unexpected)}",
flush=True)
if miss:
print("[faceanything] WARNING: deformation_head weights are missing — "
"canonical predictions will be untrained!", flush=True)
model = model.to(device=device)
model.eval()
return model
|