Hugging Face's logo

Spaces:
LinzhouLi
/
RealDenseFace
Running on Zero

App Files Community
RealDenseFace / cuda_extensions /flame_solver /src /flame_tracking.h
LinzhouLi's picture
LinzhouLi
Update RealDenseFace image reconstruction demo
8448f43 verified
Raw History Blame Contribute Delete
8.05 kB
#pragma once
#include <functional>
#include <cuda.h>
#include <cuda_runtime.h>
#include <device_launch_parameters.h>
#include <cublas_v2.h>
#include <cusolverDn.h>
#include <glm/glm.hpp>
namespace AssembleJacobian {
// Host launcher keeps templated CUDA kernel stubs in their defining TU.
// CUDA 13 gives these stubs internal linkage by default.
void LaunchExpPoseRegJacobian(
const int num_expressions,
const bool calc_jacobian,
const int batch_size,
const int num_offsets,
const int num_residuals,
const float* __restrict__ unknows,
const float exp_reg_weight,
const float pose_reg_weight,
float* __restrict__ residual,
float* __restrict__ jacobian
);
__global__ void AssembleIdenRegJacobianCuda(
const bool calc_jacobian,
const int num_offsets,
const float* __restrict__ unknows,
const float iden_reg_weight,
float* __restrict__ residual,
float* __restrict__ jacobian
);
}
namespace FlameTracking {
void flame_calc_canonical(
const int num_vertices,
const float* identity,
const glm::vec3* v_template,
const glm::vec3* shape_dirs,
const int num_J_regressor_nonzero,
const int* J_regressor_row,
const int* J_regressor_col,
const float* J_regressor_values,
glm::vec3* v_canonical,
glm::vec3* j_canonical
);
void flame_expression_jacobian(
const int num_vertices,
const int num_expressions,
const int batch_size,
const float* batched_x,
const glm::vec3* v_canonical,
const glm::vec3* j_canonical,
const glm::vec3* expr_dirs, // [B, V, 3]
const glm::vec3* pose_dirs, // [36, V, 3]
const float* lbs_weights, // [V, J]
float* batched_pose_feat,
glm::mat4* batched_affine_mats,
glm::mat3* batched_right_jacobian,
glm::mat4* batched_B_masked,
glm::mat3* batched_Ar,
glm::vec3* batched_v_output,
glm::vec3* batched_j_output,
float* batched_v_jacobian,
float* batched_j_jacobian
);
void flame_expression_forward(
const int num_vertices,
const int num_expressions,
const int batch_size,
const float* batched_x,
const glm::vec3* v_canonical,
const glm::vec3* j_canonical,
const glm::vec3* expr_dirs, // [B, V, 3]
const glm::vec3* pose_dirs, // [36, V, 3]
const float* lbs_weights, // [V, J]
float* batched_pose_feat,
glm::mat4* batched_affine_mats,
glm::vec3* batched_v_output,
glm::vec3* batched_j_output
);
void flame_identity_jacobian(
const int num_vertices,
const int num_expressions,
const int batch_size,
const float* batched_x,
const glm::vec3* v_canonical,
const glm::vec3* joints,
const glm::vec3* ref_joints_dirs,
const glm::vec3* shape_dirs,
const glm::vec3* expr_dirs,
const glm::vec3* pose_dirs,
const float* lbs_weights,
float* batched_pose_feat,
glm::mat4* batched_rel_affine_mats,
glm::mat4* batched_affine_mats,
glm::vec3* batched_v_output,
glm::vec3* batched_j_output,
float* batched_v_jacobian,
float* batched_j_jacobian
);
void assemble_x_jacobian(
const int num_vertices,
const int num_expressions,
const int batch_size,
const int num_views,
const int num_items,
const int camera_view_stride,
const int64_t* vertex_indices,
const float alignment_weight,
const float rel_depth_weight,
const float* batched_x,
const float exp_reg_weight,
const float pose_reg_weight,
const glm::mat4* view_mat,
const glm::mat4* viewproj_mat,
const glm::vec3* batched_vertices, // [T, V, 3]
const glm::vec3* batched_joints, // [T, J, 3]
const glm::vec2* gt_projected_vertices, // [T, C, V, 2]
const float* gt_vertex_rel_depths, // [T, C, V, 1]
const float* gt_projected_log_vars, // [T, C, V, 1], nullptr means disabled
const float* gt_vertex_rel_depth_log_vars, // [T, C, V, 1], nullptr means disabled
const glm::vec3* batched_v_jacobian, // [T, V, X, 3]
const glm::vec3* batched_j_jacobian, // [T, J, X, 3]
float* residual, // [T, R], per-item obs dim is 3 if rel_depth_weight > 0, else 2
float* jacobian // [T, R, X], per-item obs dim is 3 if rel_depth_weight > 0, else 2
);
void assemble_identity_jacobian(
const int num_vertices,
const int batch_size,
const int num_views,
const int num_items,
const int camera_view_stride,
const int64_t* vertex_indices,
const float alignment_weight,
const float rel_depth_weight,
const float* identity,
const float identity_reg_weight,
const glm::mat4* view_mat,
const glm::mat4* viewproj_mat,
const glm::vec3* batched_vertices, // [T, V, 3]
const glm::vec3* batched_joints, // [T, J, 3]
const glm::vec2* gt_projected_vertices, // [T, C, V, 2]
const float* gt_vertex_rel_depths, // [T, C, V, 1]
const float* gt_projected_log_vars, // [T, C, V, 1], nullptr means disabled
const float* gt_vertex_rel_depth_log_vars, // [T, C, V, 1], nullptr means disabled
const glm::vec3* batched_v_jacobian, // [T, V, X, 3]
const glm::vec3* batched_j_jacobian, // [T, J, X, 3]
float* residual, // [R], per-item obs dim is 3 if rel_depth_weight > 0, else 2
float* jacobian // [R, X], per-item obs dim is 3 if rel_depth_weight > 0, else 2
);
void assemble_x_direct_jacobian(
const int num_vertices,
const int num_expressions,
const int batch_size,
const int num_items,
const int64_t* vertex_indices,
const float vertex_weight,
const float* batched_x,
const float exp_reg_weight,
const float pose_reg_weight,
const glm::vec3* batched_vertices,
const glm::vec3* gt_vertices,
const glm::vec3* batched_v_jacobian,
float* residual,
float* jacobian
);
void assemble_identity_direct_jacobian(
const int num_vertices,
const int batch_size,
const int num_items,
const int64_t* vertex_indices,
const float vertex_weight,
const float* identity,
const float identity_reg_weight,
const glm::vec3* batched_vertices,
const glm::vec3* gt_vertices,
const glm::vec3* batched_v_jacobian,
float* residual,
float* jacobian
);
void cholesky_solve(
cublasHandle_t cublas_handle,
cusolverDnHandle_t cusolver_handle,
std::function<void*(size_t N)> mem_func,
const float* d_J, // [R, X]
const float* d_r, // [R]
float* d_x, // [X]
const int m, // num_residuals
const int n, // num_unknowns
const float lambda
);
void cholesky_solve_batched(
cublasHandle_t cublas_handle,
cusolverDnHandle_t cusolver_handle,
std::function<void*(size_t N)> mem_func,
const float* d_J, // [B, R, X]
const float* d_r, // [B, R]
float* d_x, // [B, X]
const int m, // num_residuals
const int n, // num_unknowns
const int batch_size,
const float lambda
);
}