Spaces:
Running on Zero
Running on Zero
| import pickle | |
| from typing import Optional, Union | |
| import numpy as np | |
| import torch | |
| from roma import rotvec_to_rotmat | |
| from common.config import FLAMEConfig | |
| from common.utils import load_expression_blendshapes | |
| def _matmul_acc(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: | |
| a = a.unsqueeze(-1) | |
| b = b.unsqueeze(-3) | |
| return (a * b).sum(-2) | |
| def _affine_transform(transform: torch.Tensor, position: torch.Tensor) -> torch.Tensor: | |
| transform_rot = transform[..., :3, :3] | |
| transform_trans = transform[..., :3, 3] | |
| return (transform_rot * position.unsqueeze(-2)).sum(-1) + transform_trans | |
| class FLAME: | |
| def __init__( | |
| self, | |
| config: Union[FLAMEConfig, str], | |
| simplified: bool = False, | |
| device: str | torch.device | None = None, | |
| ): | |
| if isinstance(config, str): | |
| config = FLAMEConfig(flame_model_path=config) | |
| elif not isinstance(config, FLAMEConfig): | |
| raise TypeError(f"FLAME expects FLAMEConfig or flame_model_path string, got {type(config).__name__}.") | |
| if device is None: | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| else: | |
| device = torch.device(device) | |
| with open(config.flame_model_path, "rb") as f: | |
| data = pickle.load(f, encoding="latin1") | |
| self.simplified = simplified | |
| self.device = device | |
| self.dtype = torch.float32 | |
| self.config = config | |
| self.num_expressions = config.resolved_num_expressions() | |
| self.v_template = torch.from_numpy(np.asarray(data["v_template"])).to(device=device, dtype=self.dtype) | |
| self.faces = torch.from_numpy(np.asarray(data["f"]).astype(np.int64)).to(device=device, dtype=torch.int64) | |
| self.pose_dirs = torch.from_numpy(np.asarray(data["posedirs"])).to(device=device, dtype=self.dtype) | |
| self.lbs_weights = torch.from_numpy(np.asarray(data["weights"])).to(device=device, dtype=self.dtype) | |
| self.parents = torch.from_numpy(np.asarray(data["kintree_table"][0])).to(device=device, dtype=torch.int64) | |
| self.J_regressor = torch.from_numpy(np.asarray(data["J_regressor"].todense())).to(device=device, dtype=self.dtype) | |
| shape_dirs_np = np.asarray(data["shapedirs"])[..., :300].astype(np.float32, copy=False) | |
| self.shape_dirs = torch.from_numpy(np.ascontiguousarray(shape_dirs_np)).to( | |
| device=device, dtype=self.dtype | |
| ).view(-1, 300).T.contiguous() | |
| expr_dirs_np = load_expression_blendshapes(data, config) | |
| self.expr_dirs = torch.from_numpy(expr_dirs_np).to(device=device, dtype=self.dtype) \ | |
| .view(-1, self.num_expressions).T.contiguous() | |
| num_pose_basis = self.pose_dirs.shape[-1] | |
| self.pose_dirs = self.pose_dirs.reshape([-1, num_pose_basis]).T.contiguous() | |
| self.joints_dirs = ( | |
| self.J_regressor | |
| ).view(5, 3, 300).permute(2, 0, 1).contiguous() | |
| def calc_canonical(self, identity: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: | |
| if identity.ndim == 1: | |
| v_canonical = self.v_template + torch.matmul(identity, self.shape_dirs).view(-1, 3) | |
| j_canonical = torch.matmul(self.J_regressor, v_canonical) | |
| return v_canonical, j_canonical | |
| v_canonical = self.v_template.unsqueeze(0) + torch.matmul(identity, self.shape_dirs).view(identity.shape[0], -1, 3) | |
| j_canonical = torch.matmul(self.J_regressor.unsqueeze(0).expand(identity.shape[0], -1, -1), v_canonical) | |
| return v_canonical, j_canonical | |
| def forward_from_canonical( | |
| self, | |
| v_canonical: torch.Tensor, | |
| j_canonical: torch.Tensor, | |
| x: torch.Tensor, | |
| return_joints: bool = False, | |
| ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: | |
| add_batch = False | |
| if x.ndim == 1: | |
| add_batch = True | |
| x = x.unsqueeze(0) | |
| batch_size = x.shape[0] | |
| if v_canonical.ndim == 2: | |
| v_canonical = v_canonical.unsqueeze(0).expand(batch_size, -1, -1) | |
| if j_canonical.ndim == 2: | |
| j_canonical = j_canonical.unsqueeze(0).expand(batch_size, -1, -1) | |
| n_expr = self.num_expressions | |
| expression = x[:, :n_expr].reshape(batch_size, n_expr) | |
| pose = x[:, n_expr:n_expr + 15].reshape(batch_size, 5, 3) | |
| translation = x[:, n_expr + 15:n_expr + 18].reshape(batch_size, 3) | |
| v_shaped = v_canonical + torch.matmul(expression, self.expr_dirs).view(batch_size, -1, 3) | |
| if self.simplified: | |
| joints = j_canonical | |
| else: | |
| joints = torch.matmul(self.J_regressor.unsqueeze(0).expand(batch_size, -1, -1), v_shaped) | |
| rot_mats = rotvec_to_rotmat(pose) | |
| if self.simplified: | |
| v_posed = v_shaped | |
| else: | |
| identity_mat = torch.eye(3, device=self.device, dtype=self.dtype).unsqueeze(0).unsqueeze(0) | |
| pose_feat = (rot_mats[:, 1:] - identity_mat).reshape(batch_size, -1) | |
| v_posed = v_shaped + torch.matmul(pose_feat, self.pose_dirs).view(batch_size, -1, 3) | |
| rel_joints = joints.clone() | |
| rel_joints[:, 1] -= joints[:, 0] | |
| rel_joints[:, 2:] -= joints[:, 1:2] | |
| rel_affine_mats = torch.zeros(batch_size, 5, 4, 4, device=self.device, dtype=self.dtype) | |
| rel_affine_mats[..., :3, :3] = rot_mats | |
| rel_affine_mats[..., :3, 3] = rel_joints | |
| rel_affine_mats[..., 3, 3] = 1.0 | |
| affine_mats = rel_affine_mats.clone() | |
| affine_mats[:, 1] = _matmul_acc(affine_mats[:, 0].clone(), rel_affine_mats[:, 1]) | |
| affine_mats[:, 2:] = _matmul_acc(affine_mats[:, 1:2].clone(), rel_affine_mats[:, 2:]) | |
| result_joints = affine_mats[..., :3, 3].clone() | |
| joint_transformed = (affine_mats[..., :3, :3] * joints.unsqueeze(-2)).sum(-1) | |
| affine_mats = affine_mats.clone() | |
| affine_mats[..., :3, 3] -= joint_transformed | |
| transforms = _matmul_acc(self.lbs_weights.unsqueeze(0), affine_mats.view(batch_size, 5, 16)).view(batch_size, -1, 4, 4) | |
| v_skinned = _affine_transform(transforms, v_posed) | |
| v_skinned += translation.unsqueeze(1) | |
| result_joints += translation.unsqueeze(1) | |
| if add_batch: | |
| v_skinned = v_skinned.squeeze(0) | |
| result_joints = result_joints.squeeze(0) | |
| if return_joints: | |
| return v_skinned, result_joints | |
| return v_skinned | |
| def forward( | |
| self, | |
| identity: torch.Tensor, | |
| x: torch.Tensor, | |
| v_offsets: Optional[torch.Tensor] = None, | |
| return_joints: bool = False, | |
| ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: | |
| v_canonical, j_canonical = self.calc_canonical(identity) | |
| vertices, joints = self.forward_from_canonical(v_canonical, j_canonical, x, return_joints=True) | |
| if v_offsets is not None: | |
| if vertices.ndim == 2 and v_offsets.ndim == 2: | |
| vertices = vertices + v_offsets | |
| elif vertices.ndim == 3 and v_offsets.ndim == 2: | |
| vertices = vertices + v_offsets.unsqueeze(0) | |
| else: | |
| vertices = vertices + v_offsets | |
| if return_joints: | |
| return vertices, joints | |
| return vertices | |