Blame for tracker/fitting/gn_flame_optimizer.py

from __future__ import annotations

import math
import pickle
from typing import Optional, Sequence

import numpy as np
import torch
from tqdm import tqdm

from camera import Camera, PerspectiveCamera
from camera.camera import focal2fov, fov2focal
from flame_solver import FlameSolverOps, SolverContext
from common.config import INVISIBLE_LOG_VAR, GNFlameOptimizerConfig
from common.utils import load_expression_blendshapes, load_mesh_from_obj
from common.types import FittingOutput, InferenceOutput


EYE_REGION_VERTICES = [
    # eyelid
    807, 814, 822, 824, 827, 991, 994, 995, 1093, 1096, 1113, 1170, 1175, 1195, 1200, 1201, 1202, 1216, 1218, 1329, 1331, 1336, 1340, 1343, 1344, 1354, 1355, 1357, 3833, 3855,
    2264, 2267, 2271, 2273, 2276, 2355, 2358, 2359, 2402, 2403, 2406, 2436, 2437, 2447, 2448, 2449, 2450, 2451, 2453, 2485, 2486, 2487, 2491, 2494, 2495, 2505, 2506, 2508, 3632, 3689,
    # iris
    4478, 4482, 4486, 4490, 4494, 4498, 4502, 4506, 4510, 4514, 4518, 4522, 4526, 4530, 4534, 4538, 4542, 4546, 4550, 4554, 4558, 4562, 4566, 4570, 4574, 4578, 4582, 4586, 4590, 4594, 4599, 4603,
    3932, 3936, 3940, 3944, 3948, 3952, 3956, 3960, 3964, 3968, 3972, 3976, 3980, 3984, 3988, 3992, 3996, 4000, 4004, 4008, 4012, 4016, 4020, 4024, 4028, 4032, 4036, 4040, 4044, 4048, 4053, 4057
]


class GNFlameOptimizer:
    HISTORY_COMPONENTS = ("align_loss", "depth_loss", "reg_expr", "reg_pose", "reg_ident", "total_loss")

    def __init__(self, config: GNFlameOptimizerConfig):
        self.config = config
        self.device = torch.device(config.device)
        self._load_flame_assets(config.flame)
        self.solver_ops = FlameSolverOps(num_expressions=self.num_expressions)
        self.solver_context = SolverContext.create(self.solver_ops.extension)
        self.key_vertex_ids = torch.tensor([], dtype=torch.long, device=self.device)
        # key_vertex_ids = np.unique(load_mesh_from_obj("weights/flame/flame_key_region.obj")[1].flatten())
        # self.key_vertex_ids = torch.from_numpy(key_vertex_ids).to(dtype=torch.int64, device=self.device)
        self._empty_uncertainty = torch.empty((0,), dtype=torch.float32, device=self.device)
        self._eye_region_ids = torch.tensor(EYE_REGION_VERTICES, dtype=torch.long, device=self.device)
        self._eye_boost_log_var_delta = -2.0 * math.log(max(float(config.eye_region_boost_factor), 1e-8))

    def _load_flame_assets(self, flame_config) -> None:
        self.num_expressions = flame_config.resolved_num_expressions()
        self.dim_x = self.num_expressions + 18
        self.pose_start = self.num_expressions
        self.trans_start = self.num_expressions + 15

        with open(flame_config.flame_model_path, "rb") as f:
            data = pickle.load(f, encoding="latin1")

        self.v_template = torch.from_numpy(np.asarray(data["v_template"])).to(dtype=torch.float32, device=self.device)
        self.faces = torch.from_numpy(np.asarray(data["f"]).astype(np.int64)).to(dtype=torch.int64, device=self.device)
        self.pose_dirs = torch.from_numpy(np.asarray(data["posedirs"])).to(dtype=torch.float32, device=self.device)
        self.lbs_weights = torch.from_numpy(np.asarray(data["weights"])).to(dtype=torch.float32, device=self.device)

        j_regressor_coo = data["J_regressor"].tocoo()
        self.J_regressor_row = torch.from_numpy(np.asarray(j_regressor_coo.row)).to(dtype=torch.int32, device=self.device)
        self.J_regressor_col = torch.from_numpy(np.asarray(j_regressor_coo.col)).to(dtype=torch.int32, device=self.device)
        self.J_regressor_values = torch.from_numpy(np.asarray(j_regressor_coo.data)).to(dtype=torch.float32, device=self.device)

        num_pose_basis = self.pose_dirs.shape[-1]
        shapedirs_np = np.asarray(data["shapedirs"]).astype(np.float32, copy=False)
        self.shape_dirs = torch.from_numpy(np.ascontiguousarray(shapedirs_np[..., :300])).to(
            dtype=torch.float32, device=self.device
        ).view(-1, 300).T.contiguous()

        expr_dirs_np = load_expression_blendshapes(data, flame_config)  # (V, 3, num_expressions)
        self.expr_dirs = torch.from_numpy(expr_dirs_np).to(dtype=torch.float32, device=self.device) \
            .view(-1, self.num_expressions).T.contiguous()
        self.pose_dirs = self.pose_dirs.reshape([-1, num_pose_basis]).T.contiguous()

        j_regressor_dense = torch.from_numpy(np.asarray(data["J_regressor"].todense())).to(dtype=torch.float32, device=self.device)
        self.joints_dirs = (
            j_regressor_dense @ self.shape_dirs.view(300, 5023, 3).permute(1, 2, 0).reshape(5023, 300 * 3)
        ).view(5, 3, 300).permute(2, 0, 1).contiguous()

    def _apply_eye_region_boost(self, vertex_coord_log_var: torch.Tensor | None) -> torch.Tensor | None:
        if vertex_coord_log_var is None or not self.config.boost_eye_region:
            return vertex_coord_log_var
        vertex_coord_log_var = vertex_coord_log_var.clone()
        vertex_coord_log_var[..., self._eye_region_ids, :] += self._eye_boost_log_var_delta
        return vertex_coord_log_var

    def _calc_canonical_flame(self, identity: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        return self.solver_ops.calc_canonical(
            identity,
            self.v_template,
            self.shape_dirs,
            self.J_regressor_row,
            self.J_regressor_col,
            self.J_regressor_values,
        )

    def _calc_posed_flame(

        self,

        v_canonical: torch.Tensor,

        j_canonical: torch.Tensor,

        x: torch.Tensor,

    ) -> tuple[torch.Tensor, torch.Tensor]:
        return self.solver_ops.forward_vertices(
            x,
            v_canonical,
            j_canonical,
            self.expr_dirs,
            self.pose_dirs,
            self.lbs_weights,
        )

    @staticmethod
    def _check_finite(tensor: torch.Tensor | None, label: str) -> bool:
        """Return True if *tensor* is fully finite. Otherwise print a NaN/inf diagnostic

        line including stats over the finite slice, and return False. ``None`` is

        treated as finite (nothing to check)."""
        if tensor is None or not torch.is_tensor(tensor):
            return True
        finite_mask = torch.isfinite(tensor)
        if bool(finite_mask.all().item()):
            return True
        bad = (~finite_mask)
        n_bad = int(bad.sum().item())
        n_total = int(tensor.numel())
        n_nan = int(torch.isnan(tensor).sum().item())
        n_pinf = int(torch.isposinf(tensor).sum().item())
        n_ninf = int(torch.isneginf(tensor).sum().item())
        finite_vals = tensor[finite_mask].detach().to(torch.float32)
        if finite_vals.numel() > 0:
            amin = float(finite_vals.min().item())
            amax = float(finite_vals.max().item())
            amean = float(finite_vals.mean().item())
            stats = f"finite[min={amin:.4g} max={amax:.4g} mean={amean:.4g}]"
        else:
            stats = "no finite values"
        print(
            f"[GN][NaN!!] {label}: bad={n_bad}/{n_total} "
            f"(nan={n_nan}, +inf={n_pinf}, -inf={n_ninf}) "
            f"shape={tuple(tensor.shape)} dtype={tensor.dtype}; {stats}",
            flush=True,
        )
        return False

    def _stack_camera_matrices(self, cameras: Sequence[Camera]) -> tuple[torch.Tensor, torch.Tensor]:
        if len(cameras) == 0:
            raise ValueError("At least one camera is required.")
        view_mats = [camera.get_w2v for camera in cameras]
        viewproj_mats = [camera.get_full_proj for camera in cameras]
        return (
            torch.stack(view_mats, dim=0).to(self.device, torch.float32), 
            torch.stack(viewproj_mats, dim=0).to(self.device, torch.float32)
        )

    def _single_camera_matrices(self, camera: Camera) -> tuple[torch.Tensor, torch.Tensor]:
        view_mat = camera.get_w2v.to(self.device, torch.float32).unsqueeze(0)
        viewproj_mat = camera.get_full_proj.to(self.device, torch.float32).unsqueeze(0)
        return view_mat, viewproj_mat

    def _build_config_camera(self, image_width: int, image_height: int, fov_y: float) -> PerspectiveCamera:
        rot = self.config.camera_rotation.reshape(3, 3)
        pos = self.config.camera_position.reshape(3)
        return PerspectiveCamera(
            fov_y=math.radians(float(fov_y)),
            rot=rot,
            pos=pos,
            width=int(image_width),
            height=int(image_height),
        )

    def _fov_y_to_focal(self, image_height: int, fov_y: float) -> float:
        return fov2focal(math.radians(float(fov_y)), float(image_height))

    def _focal_to_fov_y(self, image_height: int, focal_y: float) -> float:
        return math.degrees(focal2fov(float(focal_y), float(image_height)))

    def _evaluate_energy(

        self,

        camera_view_mat: torch.Tensor,

        camera_viewproj_mat: torch.Tensor,

        vertex_coord: torch.Tensor,

        vertex_depth: torch.Tensor,

        vertex_coord_log_var: torch.Tensor | None,

        vertex_depth_log_var: torch.Tensor | None,

        identity: torch.Tensor,

        x: torch.Tensor,

    ) -> torch.Tensor:
        num_views = int(camera_view_mat.shape[0])

        identity = identity.to(device=self.device, dtype=torch.float32).reshape(300)
        x = x.to(device=self.device, dtype=torch.float32)
        vertex_indices = self.key_vertex_ids
        num_items = int(vertex_indices.numel()) if vertex_indices.numel() > 0 else int(vertex_coord.shape[-2])
        # CUDA assemble_jacobian writes obs_dim columns per item: 3 (uv + depth)
        # when depth_weight > 0, else 2 (uv only). See assemble_jacobian.cu:39, 205.
        obs_dim = 3 if float(self.config.relative_depth_weight) > 0.0 else 2
        num_data_rows = num_views * num_items * obs_dim
        chunk_size = int(self.config.energy_chunk_size)

        def energy_from_residual(residual: torch.Tensor) -> torch.Tensor:
            if residual.ndim == 1:
                data_residual = residual[:num_data_rows].reshape(-1, obs_dim)
                return torch.square(data_residual).sum()
            if residual.ndim == 2:
                data_residual = residual[:, :num_data_rows].reshape(residual.shape[0], -1, obs_dim)
                return torch.square(data_residual).sum(dim=(1, 2))
            raise ValueError("x residual must be [R] or [T,R].")

        v_canonical, j_canonical = self._calc_canonical_flame(identity)
        if x.ndim == 1:
            flame_v, flame_j, flame_v_jacobian, flame_j_jacobian = self.solver_ops.expression_jacobian(
                x,
                v_canonical,
                j_canonical,
                self.expr_dirs,
                self.pose_dirs,
                self.lbs_weights,
            )
            residual, _ = self.solver_ops.assemble_x_system(
                camera_view_mat,
                camera_viewproj_mat,
                flame_v,
                flame_j,
                vertex_coord,
                vertex_depth,
                vertex_coord_log_var,
                vertex_depth_log_var,
                x,
                flame_v_jacobian,
                flame_j_jacobian,
                self.config.correspondence_weight,
                self.config.relative_depth_weight,
                self.config.expression_regularization,
                self.config.pose_regularization,
                vertex_indices,
            )
            return energy_from_residual(residual)

        if chunk_size <= 0:
            raise ValueError("sequence_energy_chunk_size must be positive.")

        batch_size = int(x.shape[0])
        chunk_energies: list[torch.Tensor] = []
        for start in range(0, batch_size, chunk_size):
            end = min(start + chunk_size, batch_size)
            x_chunk = x[start:end]
            vertex_coord_chunk = vertex_coord[start:end]
            vertex_depth_chunk = vertex_depth[start:end]
            if vertex_coord_log_var is None or vertex_depth_log_var is None:
                vertex_coord_log_var_chunk = vertex_coord_log_var
                vertex_depth_log_var_chunk = vertex_depth_log_var
            else:
                vertex_coord_log_var_chunk = vertex_coord_log_var[start:end]
                vertex_depth_log_var_chunk = vertex_depth_log_var[start:end]

            flame_v, flame_j, flame_v_jacobian, flame_j_jacobian = self.solver_ops.expression_jacobian(
                x_chunk,
                v_canonical,
                j_canonical,
                self.expr_dirs,
                self.pose_dirs,
                self.lbs_weights,
            )
            residual, _ = self.solver_ops.assemble_x_system(
                camera_view_mat,
                camera_viewproj_mat,
                flame_v,
                flame_j,
                vertex_coord_chunk,
                vertex_depth_chunk,
                vertex_coord_log_var_chunk,
                vertex_depth_log_var_chunk,
                x_chunk,
                flame_v_jacobian,
                flame_j_jacobian,
                self.config.correspondence_weight,
                self.config.relative_depth_weight,
                self.config.expression_regularization,
                self.config.pose_regularization,
                vertex_indices,
            )
            chunk_energies.append(energy_from_residual(residual))
        return torch.cat(chunk_energies, dim=0)

    def _register_wi_fov_search(

        self,

        image_width: int,

        image_height: int,

        vertex_coord: torch.Tensor,

        vertex_depth: torch.Tensor,

        vertex_coord_log_var: torch.Tensor | None,

        vertex_depth_log_var: torch.Tensor | None,

        init_identity: Optional[torch.Tensor] = None,

        init_x: Optional[torch.Tensor] = None,

        pose_init: bool = True,

    ) -> tuple[float, torch.Tensor, torch.Tensor, torch.Tensor]:
        pose_v_canonical, pose_j_canonical = self._calc_canonical_flame(init_identity)

        def evaluate_fov(fov_y: float) -> float:
            camera = self._build_config_camera(image_width, image_height, fov_y)
            camera_view_mat, camera_viewproj_mat = self._single_camera_matrices(camera)
            # Keep the search objective comparable across candidate cameras.

            if pose_init:
                pose_x = self._register_pose(
                    camera_view_mat,
                    camera_viewproj_mat,
                    vertex_coord,
                    vertex_depth,
                    vertex_coord_log_var,
                    vertex_depth_log_var,
                    pose_v_canonical,
                    pose_j_canonical,
                    init_x=init_x,
                )
            else:
                pose_x = init_x

            identity, x = self._register(
                camera_view_mat,
                camera_viewproj_mat,
                vertex_coord,
                vertex_depth,
                vertex_coord_log_var,
                vertex_depth_log_var,
                init_identity=init_identity,
                init_x=pose_x,
                num_iters=self.config.fov_registration_iterations,
            )
            return float(
                self._evaluate_energy(
                    camera_view_mat,
                    camera_viewproj_mat,
                    vertex_coord,
                    vertex_depth,
                    vertex_coord_log_var,
                    vertex_depth_log_var,
                    identity,
                    x,
                ).mean().item()
            )

        l = float(self.config.fov_y_min)
        r = float(self.config.fov_y_max)
        phi = 0.5 * (math.sqrt(5.0) - 1.0)

        x1 = r - phi * (r - l)
        x2 = l + phi * (r - l)
        e1 = evaluate_fov(x1)
        e2 = evaluate_fov(x2)

        for iter_id in range(self.config.fov_search_iterations):
            if r - l < self.config.fov_search_epsilon:
                break
            if e1 < e2:
                r = x2
                x2 = x1
                e2 = e1
                x1 = r - phi * (r - l)
                e1 = evaluate_fov(x1)
            else:
                l = x1
                x1 = x2
                e1 = e2
                x2 = l + phi * (r - l)
                e2 = evaluate_fov(x2)

        best_fov_y = 0.5 * (l + r)
        best_focal = self._fov_y_to_focal(image_height, best_fov_y)
        final_camera = self._build_config_camera(image_width, image_height, best_fov_y)
        final_view_mat, final_viewproj_mat = self._single_camera_matrices(final_camera)

        if pose_init:
            pose_x = self._register_pose(
                final_view_mat,
                final_viewproj_mat,
                vertex_coord,
                vertex_depth,
                vertex_coord_log_var,
                vertex_depth_log_var,
                pose_v_canonical,
                pose_j_canonical,
                init_x=init_x,
            )
        else:
            pose_x = init_x

        final_identity, final_x = self._register(
            final_view_mat,
            final_viewproj_mat,
            vertex_coord,
            vertex_depth,
            vertex_coord_log_var,
            vertex_depth_log_var,
            init_identity=init_identity,
            init_x=pose_x,
        )
        return best_fov_y, final_identity, final_x, pose_x

    def _register(

        self,

        camera_view_mat: torch.Tensor,

        camera_viewproj_mat: torch.Tensor,

        vertex_coord: torch.Tensor,

        vertex_depth: torch.Tensor,

        vertex_coord_log_var: torch.Tensor | None = None,

        vertex_depth_log_var: torch.Tensor | None = None,

        init_identity: Optional[torch.Tensor] = None,

        init_x: Optional[torch.Tensor] = None,

        num_iters: int | None = None,

    ) -> tuple[torch.Tensor, torch.Tensor]:
        num_views = int(camera_view_mat.shape[-3])

        if vertex_coord.ndim == 2 and vertex_depth.ndim == 2:
            shape = (self.dim_x,)
        elif vertex_coord.ndim == 3 and vertex_depth.ndim == 3:
            if vertex_coord.shape[0] == num_views:
                shape = (self.dim_x,)
            else:
                shape = (vertex_coord.shape[0], self.dim_x)
        elif vertex_coord.ndim == 4 and vertex_depth.ndim == 4:
            shape = (vertex_coord.shape[0], self.dim_x)
        else:
            raise ValueError("Invalid shape of vertex_coord or vertex_depth.")

        if init_identity is None:
            identity = torch.zeros((300,), dtype=torch.float32, device=self.device)
        else:
            identity = init_identity.to(device=self.device, dtype=torch.float32).reshape(300).clone()

        if init_x is None:
            x = torch.zeros(shape, dtype=torch.float32, device=self.device)
        else:
            x = init_x.to(device=self.device, dtype=torch.float32).reshape(shape).clone()

        effective_num_iters = self.config.registration_iterations if num_iters is None else int(num_iters)
        num_batch_frames = x.shape[0] if x.ndim == 2 else 1
        for it in range(effective_num_iters):
            v_canonical, j_canonical = self._calc_canonical_flame(identity)
            damping_lambda = self.config.damping

            flame_v, flame_j, flame_v_jacobian, flame_j_jacobian = self.solver_ops.expression_jacobian(
                x, v_canonical, j_canonical, self.expr_dirs, self.pose_dirs, self.lbs_weights
            )
            residual, jacobian = self.solver_ops.assemble_x_system(
                camera_view_mat,
                camera_viewproj_mat,
                flame_v,
                flame_j,
                vertex_coord,
                vertex_depth,
                vertex_coord_log_var,
                vertex_depth_log_var,
                x,
                flame_v_jacobian,
                flame_j_jacobian,
                self.config.correspondence_weight,
                self.config.relative_depth_weight,
                self.config.expression_regularization,
                self.config.pose_regularization,
                self.key_vertex_ids,
            )
            x = x + self.solver_ops.solve_normal_equation(
                self.solver_context,
                jacobian,
                residual,
                damping_lambda,
            )

            flame_v, flame_j, flame_v_jacobian, flame_j_jacobian = self.solver_ops.identity_jacobian(
                x,
                v_canonical,
                j_canonical,
                self.joints_dirs,
                self.shape_dirs,
                self.expr_dirs,
                self.pose_dirs,
                self.lbs_weights,
            )
            residual, jacobian = self.solver_ops.assemble_identity_system(
                camera_view_mat,
                camera_viewproj_mat,
                flame_v,
                flame_j,
                vertex_coord,
                vertex_depth,
                vertex_coord_log_var,
                vertex_depth_log_var,
                identity if init_identity is None else identity - init_identity,
                flame_v_jacobian,
                flame_j_jacobian,
                self.config.correspondence_weight,
                self.config.relative_depth_weight,
                self.config.identity_regularization,
                self.key_vertex_ids,
            )
            identity = identity + self.solver_ops.solve_normal_equation(
                self.solver_context,
                jacobian,
                residual,
                damping_lambda * num_batch_frames,
            )

        return identity, x

    def _register_pose(

        self,

        camera_view_mat: torch.Tensor,

        camera_viewproj_mat: torch.Tensor,

        vertex_coord: torch.Tensor,

        vertex_depth: torch.Tensor,

        vertex_coord_log_var: torch.Tensor | None,

        vertex_depth_log_var: torch.Tensor | None,

        v_canonical: torch.Tensor,

        j_canonical: torch.Tensor,

        init_x: Optional[torch.Tensor] = None,

        num_iters: int | None = None,

    ) -> torch.Tensor:
        num_views = int(camera_view_mat.shape[0])
        if vertex_coord.ndim == 2 and vertex_depth.ndim == 2:
            shape = (self.dim_x,)
        elif vertex_coord.ndim == 3 and vertex_depth.ndim == 3:
            if vertex_coord.shape[0] == num_views:
                shape = (self.dim_x,)
            else:
                shape = (vertex_coord.shape[0], self.dim_x)
        elif vertex_coord.ndim == 4 and vertex_depth.ndim == 4:
            shape = (vertex_coord.shape[0], self.dim_x)
        else:
            raise ValueError("Invalid shape of vertex_coord or vertex_depth.")

        if init_x is None:
            x = torch.zeros(shape, dtype=torch.float32, device=self.device)
        else:
            x = init_x.to(device=self.device, dtype=torch.float32).reshape(shape).clone()

        opt_dims = torch.tensor(
            [
                self.pose_start, self.pose_start + 1, self.pose_start + 2,
                self.trans_start, self.trans_start + 1, self.trans_start + 2,
            ],
            dtype=torch.int64,
            device=self.device,
        )
        effective_num_iters = self.config.pose_registration_iterations if num_iters is None else int(num_iters)
        for _ in range(effective_num_iters):
            flame_v, flame_j, flame_v_jacobian, flame_j_jacobian = self.solver_ops.expression_jacobian(
                x,
                v_canonical,
                j_canonical,
                self.expr_dirs,
                self.pose_dirs,
                self.lbs_weights,
            )
            residual, jacobian = self.solver_ops.assemble_x_system(
                camera_view_mat,
                camera_viewproj_mat,
                flame_v,
                flame_j,
                vertex_coord,
                vertex_depth,
                vertex_coord_log_var,
                vertex_depth_log_var,
                x,
                flame_v_jacobian,
                flame_j_jacobian,
                alignment_weight=1.0,
                rel_depth_weight=0.0,
                exp_reg_weight=0.0,
                pose_reg_weight=0.0,
                vertex_indices=self.key_vertex_ids,
            )
            delta = self.solver_ops.solve_normal_equation(
                self.solver_context,
                jacobian[..., opt_dims],
                residual,
                self.config.pose_damping,
            )
            x[..., opt_dims] = x[..., opt_dims] + delta
        
        return x

    def _track(

        self,

        camera_view_mat: torch.Tensor,

        camera_viewproj_mat: torch.Tensor,

        vertex_coord: torch.Tensor,

        vertex_depth: torch.Tensor,

        vertex_coord_log_var: torch.Tensor | None,

        vertex_depth_log_var: torch.Tensor | None,

        v_canonical: torch.Tensor,

        j_canonical: torch.Tensor,

        x: torch.Tensor,

    ) -> tuple[torch.Tensor, int]:
        # Root-pose pre-step: push bulk rotation into root joint first
        root_pose_dims = torch.tensor(
            [self.pose_start, self.pose_start + 1, self.pose_start + 2,
             self.trans_start, self.trans_start + 1, self.trans_start + 2],
            dtype=torch.int64, device=self.device,
        )
        for _ in range(self.config.root_pose_tracking_iterations):
            flame_v, flame_j, flame_v_jacobian, flame_j_jacobian = self.solver_ops.expression_jacobian(
                x, v_canonical, j_canonical, self.expr_dirs, self.pose_dirs, self.lbs_weights,
            )
            residual, jacobian = self.solver_ops.assemble_x_system(
                camera_view_mat, camera_viewproj_mat,
                flame_v, flame_j, vertex_coord, vertex_depth,
                vertex_coord_log_var, vertex_depth_log_var,
                x, flame_v_jacobian, flame_j_jacobian,
                self.config.correspondence_weight, self.config.relative_depth_weight, 0.0, 0.0,
                self.key_vertex_ids,
            )
            delta = self.solver_ops.solve_normal_equation(
                self.solver_context, jacobian[..., root_pose_dims], residual, self.config.damping,
            )
            x[..., root_pose_dims] = x[..., root_pose_dims] + delta

        # Full optimization with early stopping
        prev_energy = float('inf')
        steps = 0
        for _ in range(self.config.tracking_iterations):
            flame_v, flame_j, flame_v_jacobian, flame_j_jacobian = self.solver_ops.expression_jacobian(
                x, v_canonical, j_canonical, self.expr_dirs, self.pose_dirs, self.lbs_weights,
            )
            residual, jacobian = self.solver_ops.assemble_x_system(
                camera_view_mat, camera_viewproj_mat,
                flame_v, flame_j, vertex_coord, vertex_depth,
                vertex_coord_log_var, vertex_depth_log_var,
                x, flame_v_jacobian, flame_j_jacobian,
                self.config.correspondence_weight, self.config.relative_depth_weight,
                self.config.expression_regularization, self.config.pose_regularization,
                self.key_vertex_ids,
            )

            energy = float(torch.sum(residual ** 2).item())
            if prev_energy - energy < self.config.tracking_early_stop_ratio * prev_energy:
                break
            prev_energy = energy
            steps += 1

            x = x + self.solver_ops.solve_normal_equation(
                self.solver_context, jacobian, residual, self.config.damping,
            )

        return x, steps

    def _track_pass(

        self,

        camera_view_mat: torch.Tensor,

        camera_viewproj_mat: torch.Tensor,

        identity: torch.Tensor,

        init_x: torch.Tensor,

        vertex_coords: torch.Tensor,

        vertex_depths: torch.Tensor,

        vertex_coord_log_vars: torch.Tensor | None,

        vertex_depth_log_vars: torch.Tensor | None,

        desc: str = "track_pass"

    ) -> tuple[torch.Tensor, torch.Tensor]:
        num_frames = vertex_coords.shape[0]
        xs = torch.zeros((num_frames, self.dim_x), dtype=torch.float32, device=self.device)

        per_frame_init = init_x.ndim == 2
        if not per_frame_init:
            x = init_x.clone()

        v_canonical, j_canonical = self._calc_canonical_flame(identity)
        step_counts = []
        for frame_id in tqdm(range(num_frames), desc=desc, leave=False):
            if per_frame_init:
                x = init_x[frame_id].clone()
            if vertex_coord_log_vars is None or vertex_depth_log_vars is None:
                frame_projected_log_var = vertex_coord_log_vars
                frame_rel_depth_log_var = vertex_depth_log_vars
            else:
                frame_projected_log_var = vertex_coord_log_vars[frame_id]
                frame_rel_depth_log_var = vertex_depth_log_vars[frame_id]
            x, steps = self._track(
                camera_view_mat,
                camera_viewproj_mat,
                vertex_coords[frame_id],
                vertex_depths[frame_id],
                frame_projected_log_var,
                frame_rel_depth_log_var,
                v_canonical,
                j_canonical,
                x,
            )
            step_counts.append(steps)
            xs[frame_id] = x
        print(f"  [{desc}] steps: min={min(step_counts)} max={max(step_counts)} mean={np.mean(step_counts):.1f}")

        return xs

    def _select_keyframes(self, identity: torch.Tensor, xs: torch.Tensor, num_keyframes: int) -> torch.Tensor:
        num_frames = xs.shape[0]
        xs_eval = xs.clone()
        xs_eval[:, self.trans_start:self.dim_x] = 0
        v_canonical, j_canonical = self._calc_canonical_flame(identity)
        vertices, _ = self._calc_posed_flame(v_canonical, j_canonical, xs_eval)
        flat_vertices = vertices.flatten(1)

        key_frames = []
        key_frame = int(np.random.randint(num_frames))
        dists = torch.sum((flat_vertices - flat_vertices[key_frame].unsqueeze(0)) ** 2, dim=-1)
        dists[key_frame] = -1.0
        key_frames.append(key_frame)

        for _ in range(num_keyframes - 1):
            key_frame = int(torch.argmax(dists).item())
            dists_to_this = torch.sum((flat_vertices - flat_vertices[key_frame].unsqueeze(0)) ** 2, dim=-1)
            dists = torch.minimum(dists, dists_to_this)
            dists[key_frame] = -1.0
            key_frames.append(key_frame)

        return torch.as_tensor(key_frames, dtype=torch.long, device=self.device)

    def fit_frame(self,

        inference: InferenceOutput,

        camera: Camera | None = None,

        init_identity: np.ndarray | None = None,

    ) -> FittingOutput:
        if inference.vertex_coord.ndim != 2 or inference.vertex_coord.shape[-1] != 2:
            raise ValueError("fit_frame expects inference.vertex_coord with shape [5023, 2].")
        if inference.vertex_depth.ndim != 2 or inference.vertex_depth.shape[-1] != 1:
            raise ValueError("fit_frame expects inference.vertex_depth with shape [5023, 1].")
        if self.config.use_uncertainty_weights:
            if inference.vertex_coord_log_var.ndim != 2 or inference.vertex_coord_log_var.shape[-1] != 1:
                raise ValueError("fit_frame expects inference.vertex_coord_log_var with shape [5023, 1].")
            if inference.vertex_depth_log_var.ndim != 2 or inference.vertex_depth_log_var.shape[-1] != 1:
                raise ValueError("fit_frame expects inference.vertex_depth_log_var with shape [5023, 1].")

        vertex_coord = torch.from_numpy(inference.vertex_coord).to(device=self.device, dtype=torch.float32)
        vertex_depth = torch.from_numpy(inference.vertex_depth).to(device=self.device, dtype=torch.float32)
        vertex_coord[..., 0] /= float(inference.image_width)
        vertex_coord[..., 1] /= float(inference.image_height)
        if self.config.use_uncertainty_weights:
            vertex_coord_log_var: torch.Tensor | None = torch.from_numpy(inference.vertex_coord_log_var).to(device=self.device, dtype=torch.float32)
            vertex_depth_log_var: torch.Tensor | None = torch.from_numpy(inference.vertex_depth_log_var).to(device=self.device, dtype=torch.float32)
        else:
            vertex_coord_log_var = None
            vertex_depth_log_var = None
        vertex_coord_log_var = self._apply_eye_region_boost(vertex_coord_log_var)

        if init_identity is None:
            pose_identity = torch.zeros((300,), dtype=torch.float32, device=self.device)
        else:
            pose_identity = torch.from_numpy(init_identity).to(dtype=torch.float32, device=self.device)

        camera_fov_y = None
        if camera is None:
            camera_fov_y, identity, x, pose_x = self._register_wi_fov_search(
                inference.image_width,
                inference.image_height,
                vertex_coord,
                vertex_depth,
                vertex_coord_log_var,
                vertex_depth_log_var,
                init_identity=pose_identity,
            )
        else:
            camera_view_mat, camera_viewproj_mat = self._single_camera_matrices(camera)
            pose_v_canonical, pose_j_canonical = self._calc_canonical_flame(pose_identity)
            pose_x = self._register_pose(
                camera_view_mat,
                camera_viewproj_mat,
                vertex_coord,
                vertex_depth,
                vertex_coord_log_var,
                vertex_depth_log_var,
                pose_v_canonical,
                pose_j_canonical,
            )
            identity, x = self._register(
                camera_view_mat,
                camera_viewproj_mat,
                vertex_coord,
                vertex_depth,
                vertex_coord_log_var,
                vertex_depth_log_var,
                init_identity=pose_identity,
                init_x=pose_x,
            )

        return FittingOutput(identity=identity, x=x, pose_x=pose_x, camera_fov_y=camera_fov_y)

    def fit_frame_mv(self, inferences: list[InferenceOutput], cameras: list[Camera]) -> FittingOutput:
        if len(inferences) == 0:
            raise ValueError("fit_frame_mv expects at least one InferenceOutput.")
        if len(inferences) != len(cameras):
            raise ValueError("fit_frame_mv expects the same number of inferences and cameras.")
        if any(camera is None for camera in cameras):
            raise ValueError("fit_frame_mv requires an explicit camera for each view.")

        vertex_coords = []
        vertex_depths = []
        vertex_coord_log_vars = []
        vertex_depth_log_vars = []
        for inference in inferences:
            if inference.vertex_coord.ndim != 2 or inference.vertex_coord.shape[-1] != 2:
                raise ValueError("fit_frame_mv expects each inference.vertex_coord with shape [5023, 2].")
            if inference.vertex_depth.ndim != 2 or inference.vertex_depth.shape[-1] != 1:
                raise ValueError("fit_frame_mv expects each inference.vertex_depth with shape [5023, 1].")
            if self.config.use_uncertainty_weights:
                if inference.vertex_coord_log_var.ndim != 2 or inference.vertex_coord_log_var.shape[-1] != 1:
                    raise ValueError("fit_frame_mv expects each inference.vertex_coord_log_var with shape [5023, 1].")
                if inference.vertex_depth_log_var.ndim != 2 or inference.vertex_depth_log_var.shape[-1] != 1:
                    raise ValueError("fit_frame_mv expects each inference.vertex_depth_log_var with shape [5023, 1].")

            vertex_coord = torch.from_numpy(inference.vertex_coord).to(device=self.device, dtype=torch.float32)
            vertex_depth = torch.from_numpy(inference.vertex_depth).to(device=self.device, dtype=torch.float32)
            vertex_coord[..., 0] /= float(inference.image_width)
            vertex_coord[..., 1] /= float(inference.image_height)
            vertex_coords.append(vertex_coord)
            vertex_depths.append(vertex_depth)
            if self.config.use_uncertainty_weights:
                vertex_coord_log_vars.append(
                    torch.from_numpy(inference.vertex_coord_log_var).to(device=self.device, dtype=torch.float32)
                )
                vertex_depth_log_vars.append(
                    torch.from_numpy(inference.vertex_depth_log_var).to(device=self.device, dtype=torch.float32)
                )

        vertex_coords_tensor = torch.stack(vertex_coords, dim=0)  # [C, V, 2]
        vertex_depths_tensor = torch.stack(vertex_depths, dim=0)  # [C, V, 1]
        if self.config.use_uncertainty_weights:
            vertex_coord_log_vars_tensor: torch.Tensor | None = torch.stack(vertex_coord_log_vars, dim=0)  # [C, V, 1]
            vertex_depth_log_vars_tensor: torch.Tensor | None = torch.stack(vertex_depth_log_vars, dim=0)  # [C, V, 1]
        else:
            vertex_coord_log_vars_tensor = None
            vertex_depth_log_vars_tensor = None
        vertex_coord_log_vars_tensor = self._apply_eye_region_boost(vertex_coord_log_vars_tensor)
        camera_view_mat, camera_viewproj_mat = self._stack_camera_matrices(cameras)
        pose_identity = torch.zeros((300,), dtype=torch.float32, device=self.device)
        pose_v_canonical, pose_j_canonical = self._calc_canonical_flame(pose_identity)
        pose_x = self._register_pose(
            camera_view_mat,
            camera_viewproj_mat,
            vertex_coords_tensor,
            vertex_depths_tensor,
            vertex_coord_log_vars_tensor,
            vertex_depth_log_vars_tensor,
            pose_v_canonical,
            pose_j_canonical,
        )
        identity, x = self._register(
            camera_view_mat,
            camera_viewproj_mat,
            vertex_coords_tensor,
            vertex_depths_tensor,
            vertex_coord_log_vars_tensor,
            vertex_depth_log_vars_tensor,
            init_x=pose_x,
        )

        return FittingOutput(identity=identity, x=x)

    def fit_sequence(self,

        inference: InferenceOutput,

        camera: Camera | None = None,

        init_identity: np.ndarray | None = None,

    ) -> FittingOutput:
        if inference.vertex_coord.ndim != 3 or inference.vertex_coord.shape[-1] != 2:
            raise ValueError("fit_sequence expects inference.vertex_coord with shape [T, 5023, 2].")
        if inference.vertex_depth.ndim != 3 or inference.vertex_depth.shape[-1] != 1:
            raise ValueError("fit_sequence expects inference.vertex_depth with shape [T, 5023, 1].")
        if self.config.use_uncertainty_weights:
            if inference.vertex_coord_log_var.ndim != 3 or inference.vertex_coord_log_var.shape[-1] != 1:
                raise ValueError("fit_sequence expects inference.vertex_coord_log_var with shape [T, 5023, 1].")
            if inference.vertex_depth_log_var.ndim != 3 or inference.vertex_depth_log_var.shape[-1] != 1:
                raise ValueError("fit_sequence expects inference.vertex_depth_log_var with shape [T, 5023, 1].")

        vertex_coords = torch.from_numpy(inference.vertex_coord).to(device=self.device, dtype=torch.float32)
        vertex_depths = torch.from_numpy(inference.vertex_depth).to(device=self.device, dtype=torch.float32)
        vertex_coords[..., 0] /= float(inference.image_width)
        vertex_coords[..., 1] /= float(inference.image_height)
        if self.config.use_uncertainty_weights:
            vertex_coord_log_vars: torch.Tensor | None = torch.from_numpy(inference.vertex_coord_log_var).to(device=self.device, dtype=torch.float32)
            vertex_depth_log_vars: torch.Tensor | None = torch.from_numpy(inference.vertex_depth_log_var).to(device=self.device, dtype=torch.float32)
        else:
            vertex_coord_log_vars = None
            vertex_depth_log_vars = None
        vertex_coord_log_vars = self._apply_eye_region_boost(vertex_coord_log_vars)
        if vertex_coord_log_vars is not None:
            frame0_coord_log_var: torch.Tensor | None = vertex_coord_log_vars[0]
            frame0_depth_log_var: torch.Tensor | None = vertex_depth_log_vars[0]
        else:
            frame0_coord_log_var = None
            frame0_depth_log_var = None

        fit_camera = camera is None

        if init_identity is None:
            init_identity = torch.zeros((300,), dtype=torch.float32, device=self.device)
        else:
            init_identity = torch.from_numpy(init_identity).to(dtype=torch.float32, device=self.device)
        init_v_canonical, init_j_canonical = self._calc_canonical_flame(init_identity)

        print("[GN] register first frame")
        camera_fov_y = None
        if fit_camera:
            camera_fov_y, identity, x0, pose_x0 = self._register_wi_fov_search(
                inference.image_width,
                inference.image_height,
                vertex_coords[0],
                vertex_depths[0],
                frame0_coord_log_var,
                frame0_depth_log_var,
                init_identity=init_identity,
            )
            print(f"[GN] optimized camera fov_y after first frame register: {camera_fov_y:.6f} deg")
            camera = self._build_config_camera(inference.image_width, inference.image_height, camera_fov_y)
            camera_view_mat, camera_viewproj_mat = self._single_camera_matrices(camera)
        else:
            camera_view_mat, camera_viewproj_mat = self._single_camera_matrices(camera)
            pose_x0 = self._register_pose(
                camera_view_mat,
                camera_viewproj_mat,
                vertex_coords[0],
                vertex_depths[0],
                frame0_coord_log_var,
                frame0_depth_log_var,
                init_v_canonical,
                init_j_canonical,
            )
            identity, x0 = self._register(
                camera_view_mat,
                camera_viewproj_mat,
                vertex_coords[0],
                vertex_depths[0],
                frame0_coord_log_var,
                frame0_depth_log_var,
                init_identity=init_identity,
                init_x=pose_x0,
            )
        xs = self._track_pass(
            camera_view_mat,
            camera_viewproj_mat,
            identity,
            x0,
            vertex_coords,
            vertex_depths,
            vertex_coord_log_vars,
            vertex_depth_log_vars,
            desc="[GN] track pass 0"
        )

        num_frames = xs.shape[0]
        for round_id in range(self.config.identity_refinement_rounds):
            print(f"[GN] refine round {round_id + 1}/{self.config.identity_refinement_rounds}")
            
            # select keyframes
            if num_frames > self.config.num_keyframes:
                key_frames = self._select_keyframes(identity, xs, self.config.num_keyframes)
                xs_key = xs[key_frames]
                vertex_coords_key = vertex_coords[key_frames]
                vertex_depths_key = vertex_depths[key_frames]
                if vertex_coord_log_vars is None or vertex_depth_log_vars is None:
                    vertex_coord_log_vars_key = None
                    vertex_depth_log_vars_key = None
                else:
                    vertex_coord_log_vars_key = vertex_coord_log_vars[key_frames]
                    vertex_depth_log_vars_key = vertex_depth_log_vars[key_frames]
            else:
                xs_key = xs
                vertex_coords_key = vertex_coords
                vertex_depths_key = vertex_depths
                vertex_coord_log_vars_key = vertex_coord_log_vars
                vertex_depth_log_vars_key = vertex_depth_log_vars
            
            identity, _ = self._register(
                camera_view_mat,
                camera_viewproj_mat,
                vertex_coords_key,
                vertex_depths_key,
                vertex_coord_log_vars_key,
                vertex_depth_log_vars_key,
                init_identity=identity,
                init_x=xs_key
            )
            
            # tracking pass
            xs = self._track_pass(
                camera_view_mat,
                camera_viewproj_mat,
                identity,
                xs,
                vertex_coords,
                vertex_depths,
                vertex_coord_log_vars,
                vertex_depth_log_vars,
                desc=f"[GN] track pass {round_id + 1}"
            )

        return FittingOutput(identity=identity, x=xs, camera_fov_y=camera_fov_y)

    def fit_sequence_mv(self, inferences: list[InferenceOutput], cameras: list[Camera]) -> FittingOutput:
        if len(inferences) == 0:
            raise ValueError("fit_sequence_mv expects at least one InferenceOutput.")
        if len(inferences) != len(cameras):
            raise ValueError("fit_sequence_mv expects the same number of inferences and cameras.")
        if any(camera is None for camera in cameras):
            raise ValueError("fit_sequence_mv requires an explicit camera for each view.")

        first_inference = inferences[0]
        if first_inference.vertex_coord.ndim != 3 or first_inference.vertex_coord.shape[-1] != 2:
            raise ValueError("fit_sequence_mv expects each inference.vertex_coord with shape [T, 5023, 2].")
        if first_inference.vertex_depth.ndim != 3 or first_inference.vertex_depth.shape[-1] != 1:
            raise ValueError("fit_sequence_mv expects each inference.vertex_depth with shape [T, 5023, 1].")
        if self.config.use_uncertainty_weights:
            if first_inference.vertex_coord_log_var.ndim != 3 or first_inference.vertex_coord_log_var.shape[-1] != 1:
                raise ValueError("fit_sequence_mv expects each inference.vertex_coord_log_var with shape [T, 5023, 1].")
            if first_inference.vertex_depth_log_var.ndim != 3 or first_inference.vertex_depth_log_var.shape[-1] != 1:
                raise ValueError("fit_sequence_mv expects each inference.vertex_depth_log_var with shape [T, 5023, 1].")

        num_views = len(inferences)
        num_frames = first_inference.vertex_coord.shape[0]
        num_vertices = first_inference.vertex_coord.shape[1]
        vertex_coords_tensor = torch.empty((num_frames, num_views, num_vertices, 2), dtype=torch.float32, device=self.device)
        vertex_depths_tensor = torch.empty((num_frames, num_views, num_vertices, 1), dtype=torch.float32, device=self.device)
        if self.config.use_uncertainty_weights:
            vertex_coord_log_vars_tensor: torch.Tensor | None = torch.empty(
                (num_frames, num_views, num_vertices, 1), dtype=torch.float32, device=self.device
            )
            vertex_depth_log_vars_tensor: torch.Tensor | None = torch.empty(
                (num_frames, num_views, num_vertices, 1), dtype=torch.float32, device=self.device
            )
        else:
            vertex_coord_log_vars_tensor = None
            vertex_depth_log_vars_tensor = None

        for view_id, inference in enumerate(inferences):
            if inference.vertex_coord.ndim != 3 or inference.vertex_coord.shape[-1] != 2:
                raise ValueError("fit_sequence_mv expects each inference.vertex_coord with shape [T, 5023, 2].")
            if inference.vertex_depth.ndim != 3 or inference.vertex_depth.shape[-1] != 1:
                raise ValueError("fit_sequence_mv expects each inference.vertex_depth with shape [T, 5023, 1].")
            if self.config.use_uncertainty_weights:
                if inference.vertex_coord_log_var.ndim != 3 or inference.vertex_coord_log_var.shape[-1] != 1:
                    raise ValueError("fit_sequence_mv expects each inference.vertex_coord_log_var with shape [T, 5023, 1].")
                if inference.vertex_depth_log_var.ndim != 3 or inference.vertex_depth_log_var.shape[-1] != 1:
                    raise ValueError("fit_sequence_mv expects each inference.vertex_depth_log_var with shape [T, 5023, 1].")
            if inference.vertex_coord.shape[0] != num_frames:
                raise ValueError("fit_sequence_mv expects all views to share the same frame count.")

            vertex_coords = torch.from_numpy(inference.vertex_coord).to(device=self.device, dtype=torch.float32)
            vertex_depths = torch.from_numpy(inference.vertex_depth).to(device=self.device, dtype=torch.float32)
            vertex_coords[..., 0] /= float(inference.image_width)
            vertex_coords[..., 1] /= float(inference.image_height)
            vertex_coords_tensor[:, view_id] = vertex_coords
            vertex_depths_tensor[:, view_id] = vertex_depths
            if self.config.use_uncertainty_weights:
                vertex_coord_log_vars_tensor[:, view_id] = torch.from_numpy(inference.vertex_coord_log_var).to(
                    device=self.device, dtype=torch.float32
                )
                vertex_depth_log_vars_tensor[:, view_id] = torch.from_numpy(inference.vertex_depth_log_var).to(
                    device=self.device, dtype=torch.float32
                )

        vertex_coord_log_vars_tensor = self._apply_eye_region_boost(vertex_coord_log_vars_tensor)
        if self.config.use_uncertainty_weights:
            frame0_coord_log_var = vertex_coord_log_vars_tensor[0]
            frame0_depth_log_var = vertex_depth_log_vars_tensor[0]
        else:
            frame0_coord_log_var = None
            frame0_depth_log_var = None
        camera_view_mat, camera_viewproj_mat = self._stack_camera_matrices(cameras)

        print("[GN] register first frame (multi-view)")
        pose_identity = torch.zeros((300,), dtype=torch.float32, device=self.device)
        pose_v_canonical, pose_j_canonical = self._calc_canonical_flame(pose_identity)
        pose_x0 = self._register_pose(
            camera_view_mat,
            camera_viewproj_mat,
            vertex_coords_tensor[0],
            vertex_depths_tensor[0],
            frame0_coord_log_var,
            frame0_depth_log_var,
            pose_v_canonical,
            pose_j_canonical,
        )
        self._check_finite(pose_x0, "after _register_pose: pose_x0")
        identity, x0 = self._register(
            camera_view_mat,
            camera_viewproj_mat,
            vertex_coords_tensor[0],
            vertex_depths_tensor[0],
            frame0_coord_log_var,
            frame0_depth_log_var,
            init_x=pose_x0,
        )
        self._check_finite(identity, "after first _register: identity")
        self._check_finite(x0, "after first _register: x0")
        xs = self._track_pass(
            camera_view_mat,
            camera_viewproj_mat,
            identity,
            x0,
            vertex_coords_tensor,
            vertex_depths_tensor,
            vertex_coord_log_vars_tensor,
            vertex_depth_log_vars_tensor,
            desc="[GN] track pass 0 (multi-view)"
        )
        self._check_finite(xs, "after track pass 0: xs")

        num_frames = xs.shape[0]
        for round_id in range(self.config.identity_refinement_rounds):
            print(f"[GN] refine round {round_id + 1}/{self.config.identity_refinement_rounds} (multi-view)")

            # select key frames
            if num_frames > self.config.num_keyframes:
                key_frames = self._select_keyframes(identity, xs, self.config.num_keyframes)
                xs_key = xs[key_frames]
                vertex_coords_key = vertex_coords_tensor[key_frames]
                vertex_depths_key = vertex_depths_tensor[key_frames]
                if vertex_coord_log_vars_tensor is None or vertex_depth_log_vars_tensor is None:
                    vertex_coord_log_vars_key = None
                    vertex_depth_log_vars_key = None
                else:
                    vertex_coord_log_vars_key = vertex_coord_log_vars_tensor[key_frames]
                    vertex_depth_log_vars_key = vertex_depth_log_vars_tensor[key_frames]
            else:
                xs_key = xs
                vertex_coords_key = vertex_coords_tensor
                vertex_depths_key = vertex_depths_tensor
                vertex_coord_log_vars_key = vertex_coord_log_vars_tensor
                vertex_depth_log_vars_key = vertex_depth_log_vars_tensor
            self._check_finite(xs_key, f"refine {round_id + 1}: xs_key (after keyframe select)")

            # register pass
            identity, _ = self._register(
                camera_view_mat,
                camera_viewproj_mat,
                vertex_coords_key,
                vertex_depths_key,
                vertex_coord_log_vars_key,
                vertex_depth_log_vars_key,
                init_identity=identity,
                init_x=xs_key,
            )
            self._check_finite(identity, f"refine {round_id + 1}: identity (after _register)")

            # tracking pass
            xs = self._track_pass(
                camera_view_mat,
                camera_viewproj_mat,
                identity,
                xs,
                vertex_coords_tensor,
                vertex_depths_tensor,
                vertex_coord_log_vars_tensor,
                vertex_depth_log_vars_tensor,
                desc=f"[GN] track pass {round_id + 1} (multi-view)"
            )
            self._check_finite(xs, f"refine {round_id + 1}: xs (after track pass)")

        return FittingOutput(identity=identity, x=xs)

    # ---- Robust multi-view tracking (adaptive per-frame camera selection) ----

    def _track_pass_robust(

        self,

        all_cameras: list[Camera],

        identity: torch.Tensor,

        init_x: torch.Tensor,

        per_view_coords: list[torch.Tensor],

        per_view_depths: list[torch.Tensor],

        per_view_coord_log_vars: list[torch.Tensor] | None,

        per_view_depth_log_vars: list[torch.Tensor] | None,

        per_frame_views: list[list[int]],

        frame_reliable: torch.Tensor,

        desc: str = "track_pass_robust",

    ) -> torch.Tensor:
        num_frames = len(per_frame_views)
        xs = torch.zeros((num_frames, self.dim_x), dtype=torch.float32, device=self.device)

        per_frame_init = init_x.ndim == 2
        if not per_frame_init:
            x = init_x.clone()

        v_canonical, j_canonical = self._calc_canonical_flame(identity)
        step_counts = []
        for t in tqdm(range(num_frames), desc=desc, leave=False):
            if per_frame_init:
                x = init_x[t].clone()

            if not frame_reliable[t]:
                xs[t] = x
                continue

            views = per_frame_views[t]
            cam_view, cam_proj = self._stack_camera_matrices([all_cameras[v] for v in views])
            frame_coords = torch.stack([per_view_coords[v][t] for v in views])
            frame_depths = torch.stack([per_view_depths[v][t] for v in views])
            if per_view_coord_log_vars is not None:
                frame_coord_lv = torch.stack([per_view_coord_log_vars[v][t] for v in views])
                frame_depth_lv = torch.stack([per_view_depth_log_vars[v][t] for v in views])
            else:
                frame_coord_lv = None
                frame_depth_lv = None

            x, steps = self._track(
                cam_view, cam_proj,
                frame_coords, frame_depths,
                frame_coord_lv, frame_depth_lv,
                v_canonical, j_canonical, x,
            )
            step_counts.append(steps)
            xs[t] = x

        print(f"  [{desc}] steps: min={min(step_counts)} max={max(step_counts)} mean={np.mean(step_counts):.1f}")
        return xs

    def _gather_register_data(

        self,

        keyframe_indices: torch.Tensor,

        per_frame_views: list[list[int]],

        all_cameras: list[Camera],

        per_view_coords: list[torch.Tensor],

        per_view_depths: list[torch.Tensor],

        per_view_coord_log_vars: list[torch.Tensor] | None,

        per_view_depth_log_vars: list[torch.Tensor] | None,

        num_vertices: int,

    ):
        key_ts = keyframe_indices.cpu().tolist()
        n_keys = len(key_ts)

        per_key_views = [per_frame_views[t] for t in key_ts]
        c_max = max(len(v) for v in per_key_views)

        cam_views_list = []
        cam_projs_list = []
        for ki in range(n_keys):
            views = per_key_views[ki]
            cv, cp = self._stack_camera_matrices([all_cameras[v] for v in views])
            if len(views) < c_max:
                pad = c_max - len(views)
                cv = torch.cat([cv, cv[:1].expand(pad, -1, -1)], dim=0)
                cp = torch.cat([cp, cp[:1].expand(pad, -1, -1)], dim=0)
            cam_views_list.append(cv)
            cam_projs_list.append(cp)
        cam_view = torch.stack(cam_views_list)
        cam_proj = torch.stack(cam_projs_list)

        coords = torch.full((n_keys, c_max, num_vertices, 2), 0.5,
                            dtype=torch.float32, device=self.device)
        depths = torch.zeros((n_keys, c_max, num_vertices, 1),
                             dtype=torch.float32, device=self.device)
        coord_lvs = torch.full((n_keys, c_max, num_vertices, 1), INVISIBLE_LOG_VAR,
                               dtype=torch.float32, device=self.device)
        depth_lvs = torch.full((n_keys, c_max, num_vertices, 1), INVISIBLE_LOG_VAR,
                               dtype=torch.float32, device=self.device)

        for ki, t in enumerate(key_ts):
            for ci, v in enumerate(per_key_views[ki]):
                coords[ki, ci] = per_view_coords[v][t]
                depths[ki, ci] = per_view_depths[v][t]
                if per_view_coord_log_vars is not None:
                    coord_lvs[ki, ci] = per_view_coord_log_vars[v][t]
                    depth_lvs[ki, ci] = per_view_depth_log_vars[v][t]

        if per_view_coord_log_vars is None:
            return cam_view, cam_proj, coords, depths, None, None
        return cam_view, cam_proj, coords, depths, coord_lvs, depth_lvs

    def fit_sequence_mv_robust(

        self,

        inferences: list[InferenceOutput],

        cameras: list[Camera],

        init_x: Optional[torch.Tensor] = None,

        n_max_views: int = 16,

        n_min_views: int = 3,

    ) -> FittingOutput:
        num_all_views = len(inferences)
        num_frames = inferences[0].vertex_coord.shape[0]
        num_vertices = inferences[0].vertex_coord.shape[1]

        # --- Per-view data to GPU tensors ---
        per_view_coords: list[torch.Tensor] = []
        per_view_depths: list[torch.Tensor] = []
        per_view_coord_log_vars: list[torch.Tensor] | None = []
        per_view_depth_log_vars: list[torch.Tensor] | None = []

        for inference in inferences:
            vc = torch.from_numpy(np.asarray(inference.vertex_coord)).to(
                device=self.device, dtype=torch.float32)
            vd = torch.from_numpy(np.asarray(inference.vertex_depth)).to(
                device=self.device, dtype=torch.float32)
            vc[..., 0] /= float(inference.image_width)
            vc[..., 1] /= float(inference.image_height)
            per_view_coords.append(vc)
            per_view_depths.append(vd)
            if self.config.use_uncertainty_weights:
                per_view_coord_log_vars.append(
                    torch.from_numpy(np.asarray(inference.vertex_coord_log_var)).to(
                        device=self.device, dtype=torch.float32))
                per_view_depth_log_vars.append(
                    torch.from_numpy(np.asarray(inference.vertex_depth_log_var)).to(
                        device=self.device, dtype=torch.float32))

        if self.config.use_uncertainty_weights:
            per_view_coord_log_vars = [self._apply_eye_region_boost(lv) for lv in per_view_coord_log_vars]
        else:
            per_view_coord_log_vars = None
            per_view_depth_log_vars = None

        # --- Per-frame view selection ---
        frame_valid_matrix = np.ones((num_frames, num_all_views), dtype=bool)
        for v, inf in enumerate(inferences):
            if inf.frame_valid is not None:
                frame_valid_matrix[:, v] = np.asarray(inf.frame_valid)

        per_frame_views: list[list[int]] = []
        frame_reliable = torch.ones(num_frames, dtype=torch.bool, device=self.device)

        for t in range(num_frames):
            valid = [v for v in range(num_all_views) if frame_valid_matrix[t, v]]
            if len(valid) < n_min_views:
                print(f"[GN][warn] frame {t}: only {len(valid)} valid views (< {n_min_views}), marked unreliable")
                frame_reliable[t] = False
                per_frame_views.append([])
            elif len(valid) > n_max_views:
                per_frame_views.append(valid[:n_max_views])
            else:
                per_frame_views.append(valid)

        n_reliable = int(frame_reliable.sum().item())
        view_counts = [len(v) for v in per_frame_views]
        print(f"[GN] {n_reliable}/{num_frames} reliable frames, "
              f"views per frame: mean={np.mean(view_counts):.1f} min={np.min(view_counts)} max={np.max(view_counts)} "
              f"(n_max_views={n_max_views}, n_min_views={n_min_views})")

        # --- Register first reliable frame ---
        first_t = int(torch.argmax(frame_reliable.int()).item())
        first_views = per_frame_views[first_t]
        first_cam_view, first_cam_proj = self._stack_camera_matrices(
            [cameras[v] for v in first_views])
        first_coords = torch.stack([per_view_coords[v][first_t] for v in first_views])
        first_depths = torch.stack([per_view_depths[v][first_t] for v in first_views])
        if per_view_coord_log_vars is not None:
            first_coord_lv = torch.stack([per_view_coord_log_vars[v][first_t] for v in first_views])
            first_depth_lv = torch.stack([per_view_depth_log_vars[v][first_t] for v in first_views])
        else:
            first_coord_lv = None
            first_depth_lv = None

        print(f"[GN] register frame {first_t} with {len(first_views)} views (robust)")
        pose_identity = torch.zeros((300,), dtype=torch.float32, device=self.device)
        pose_v, pose_j = self._calc_canonical_flame(pose_identity)
        pose_x0 = self._register_pose(
            first_cam_view, first_cam_proj,
            first_coords, first_depths,
            first_coord_lv, first_depth_lv,
            pose_v, pose_j,
            init_x=init_x,
        )
        identity, x0 = self._register(
            first_cam_view, first_cam_proj,
            first_coords, first_depths,
            first_coord_lv, first_depth_lv,
            init_x=pose_x0,
        )

        # --- Track pass 0 ---
        xs = self._track_pass_robust(
            cameras, identity, x0,
            per_view_coords, per_view_depths,
            per_view_coord_log_vars, per_view_depth_log_vars,
            per_frame_views, frame_reliable,
            desc="[GN] track pass 0 (robust)",
        )

        # --- Refine rounds ---
        reliable_indices = torch.where(frame_reliable)[0]
        register_keyframes: list[list[int]] = []
        for round_id in range(self.config.identity_refinement_rounds):
            print(f"[GN] refine round {round_id + 1}/{self.config.identity_refinement_rounds} (robust)")

            xs_reliable = xs[reliable_indices]
            n_keys = min(self.config.num_keyframes, len(reliable_indices))
            if len(reliable_indices) > n_keys:
                key_in_reliable = self._select_keyframes(identity, xs_reliable, n_keys)
                key_frames = reliable_indices[key_in_reliable]
            else:
                key_frames = reliable_indices

            register_keyframes.append(key_frames.cpu().tolist())

            cam_view, cam_proj, reg_coords, reg_depths, reg_coord_lv, reg_depth_lv = \
                self._gather_register_data(
                    key_frames, per_frame_views, cameras,
                    per_view_coords, per_view_depths,
                    per_view_coord_log_vars, per_view_depth_log_vars,
                    num_vertices,
                )
            xs_key = xs[key_frames]
            print(f"[GN]   {len(key_frames)} keyframes, {cam_view.shape[1]} cameras for register pass")

            identity, _ = self._register(
                cam_view, cam_proj,
                reg_coords, reg_depths,
                reg_coord_lv, reg_depth_lv,
                init_identity=identity,
                init_x=xs_key,
            )

            xs = self._track_pass_robust(
                cameras, identity, xs,
                per_view_coords, per_view_depths,
                per_view_coord_log_vars, per_view_depth_log_vars,
                per_frame_views, frame_reliable,
                desc=f"[GN] track pass {round_id + 1} (robust)",
            )

        return FittingOutput(
            identity=identity,
            x=xs,
            per_frame_views=per_frame_views,
            frame_reliable=frame_reliable.cpu(),
            register_keyframes=register_keyframes,
        )

    def decode(self, output: FittingOutput) -> tuple[torch.Tensor, torch.Tensor]:
        identity = output.identity.to(device=self.device, dtype=torch.float32).reshape(300)
        x = output.x.to(device=self.device, dtype=torch.float32)
        squeeze_batch = False
        if x.ndim == 1:
            x = x.unsqueeze(0)
            squeeze_batch = True
        if x.ndim != 2 or x.shape[-1] != self.dim_x:
            raise ValueError(
                f"decode expects output.x with shape [{self.dim_x}] or [T, {self.dim_x}]."
            )

        v_canonical, j_canonical = self._calc_canonical_flame(identity)
        vertices, joints = self._calc_posed_flame(v_canonical, j_canonical, x)
        if squeeze_batch:
            return vertices[0], joints[0]
        return vertices, joints