From 3261fbd0c6793cd857b10f0b423ae96e8d252582 Mon Sep 17 00:00:00 2001 From: kijai <40791699+kijai@users.noreply.github.com> Date: Sun, 26 Jul 2026 21:01:25 +0300 Subject: [PATCH] Optimize speed a little bit --- comfy/ldm/sam3d_body/mhr/mhr_head.py | 8 ++- comfy/ldm/sam3d_body/mhr/mhr_rig.py | 71 +++++++++++++++++++++------ comfy/ldm/sam3d_body/model/prompt.py | 28 ++++++++--- comfy_extras/nodes_sam3d_body.py | 2 +- comfy_extras/sam3d_body/rasterizer.py | 19 ++++++- comfy_extras/sam3d_body/utils.py | 2 +- 6 files changed, 104 insertions(+), 26 deletions(-) diff --git a/comfy/ldm/sam3d_body/mhr/mhr_head.py b/comfy/ldm/sam3d_body/mhr/mhr_head.py index fbeac93c5..286b6f171 100644 --- a/comfy/ldm/sam3d_body/mhr/mhr_head.py +++ b/comfy/ldm/sam3d_body/mhr/mhr_head.py @@ -52,6 +52,7 @@ class MHRHead(nn.Module): self.scale_mean = _p(68) self.scale_comps = _p(28, 68) self.register_buffer("faces", torch.empty(36874, 3, dtype=torch.int64)) + self._faces_np = None self.hand_pose_mean = _p(54) self.hand_pose_comps = nn.Parameter(torch.eye(54), requires_grad=False) self.register_buffer("hand_joint_idxs_left", torch.empty(27, dtype=torch.int64)) @@ -93,6 +94,12 @@ class MHRHead(nn.Module): ) # single-tensor shape (1, N_v, 3) in meters return verts[0] + def faces_np(self): + """Static topology — cached so the per-layer pose_output doesn't force a D2H sync.""" + if self._faces_np is None: + self._faces_np = self.faces.cpu().numpy() + return self._faces_np + def replace_hands_in_pose(self, full_pose_params, hand_pose_params): assert full_pose_params.shape[1] == 136 @@ -325,7 +332,6 @@ class MHRHead(nn.Module): "pred_keypoints_3d": j3d.reshape(batch_size, -1, 3), "pred_vertices": verts.reshape(batch_size, -1, 3) if verts is not None else None, "pred_joint_coords": jcoords.reshape(batch_size, -1, 3) if jcoords is not None else None, - "faces": self.faces.cpu().numpy(), "joint_global_rots": joint_global_rots, "mhr_model_params": mhr_model_params, } diff --git a/comfy/ldm/sam3d_body/mhr/mhr_rig.py b/comfy/ldm/sam3d_body/mhr/mhr_rig.py index 90202a6fe..1b7e65683 100644 --- a/comfy/ldm/sam3d_body/mhr/mhr_rig.py +++ b/comfy/ldm/sam3d_body/mhr/mhr_rig.py @@ -15,27 +15,66 @@ from .mhr_utils import batch6DFromXYZ _LN2 = 0.6931471824645996 +# Half-angle cos/sin are computed on the +# whole (..., 3) at once and concatenated to [cr, cp, cy, sr, sp, sy]; _EQ_I then +# picks the three factors of each term, reproducing: +# x = sr*cp*cy - cr*sp*sy z = cr*cp*sy - sr*sp*cy +# y = cr*sp*cy + sr*cp*sy w = cr*cp*cy + sr*sp*sy +_EQ_I = (((3, 1, 2), (0, 4, 5)), + ((0, 4, 2), (3, 1, 5)), + ((0, 1, 5), (3, 4, 2)), + ((0, 1, 2), (3, 4, 5))) +_EQ_S = (-1., 1., -1., 1.) +_eq_tables: dict = {} + + +def _euler_quat_tables(device, dtype): + key = (device, dtype) + cached = _eq_tables.get(key) + if cached is None: + cached = (torch.tensor(_EQ_I, device=device), + torch.tensor(_EQ_S, device=device, dtype=dtype)) + _eq_tables[key] = cached + return cached + + def _euler_xyz_to_quat(angles): """(roll, pitch, yaw) -> quaternion (x, y, z, w). Matches pymomentum.quaternion.euler_xyz_to_quaternion.""" - roll, pitch, yaw = angles.unbind(-1) - cy, sy = torch.cos(yaw * 0.5), torch.sin(yaw * 0.5) - cp, sp = torch.cos(pitch * 0.5), torch.sin(pitch * 0.5) - cr, sr = torch.cos(roll * 0.5), torch.sin(roll * 0.5) - x = sr * cp * cy - cr * sp * sy - y = cr * sp * cy + sr * cp * sy - z = cr * cp * sy - sr * sp * cy - w = cr * cp * cy + sr * sp * sy - return torch.stack([x, y, z, w], dim=-1) + idx, sign = _euler_quat_tables(angles.device, angles.dtype) + half = angles * 0.5 + cs = torch.cat([torch.cos(half), torch.sin(half)], dim=-1) + p = cs[..., idx] # (..., 4, 2, 3) + term = p[..., 0] * p[..., 1] * p[..., 2] # (..., 4, 2) + return term[..., 0] + term[..., 1] * sign + + +# Hamilton product as gather + 3 adds. Each output component is a 4-term sum; +# _QM_P1/_QM_P2 pick the operands and _QM_S the signs, reproducing: +# x = w1*x2 + x1*w2 + y1*z2 - z1*y2 +# y = w1*y2 - x1*z2 + y1*w2 + z1*x2 +# z = w1*z2 + x1*y2 - y1*x2 + z1*w2 +# w = w1*w2 - x1*x2 - y1*y2 - z1*z2 +_QM_P1 = ((3, 0, 1, 2),) * 4 +_QM_P2 = ((0, 3, 2, 1), (1, 2, 3, 0), (2, 1, 0, 3), (3, 0, 1, 2)) +_QM_S = ((1., 1., 1., -1.), (1., -1., 1., 1.), (1., 1., -1., 1.), (1., -1., -1., -1.)) +_qm_tables: dict = {} + + +def _quat_mul_tables(device, dtype): + key = (device, dtype) + cached = _qm_tables.get(key) + if cached is None: + cached = (torch.tensor(_QM_P1, device=device), + torch.tensor(_QM_P2, device=device), + torch.tensor(_QM_S, device=device, dtype=dtype)) + _qm_tables[key] = cached + return cached def _quat_multiply(q1, q2): - x1, y1, z1, w1 = q1.unbind(-1) - x2, y2, z2, w2 = q2.unbind(-1) - x = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2 - y = w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2 - z = w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2 - w = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2 - return torch.stack([x, y, z, w], dim=-1) + p1, p2, s = _quat_mul_tables(q1.device, q1.dtype) + t = q1[..., p1] * q2[..., p2] * s + return ((t[..., 0] + t[..., 1]) + t[..., 2]) + t[..., 3] def _quat_rotate(q, v): diff --git a/comfy/ldm/sam3d_body/model/prompt.py b/comfy/ldm/sam3d_body/model/prompt.py index 2e7fafad6..c968c8449 100644 --- a/comfy/ldm/sam3d_body/model/prompt.py +++ b/comfy/ldm/sam3d_body/model/prompt.py @@ -40,6 +40,7 @@ class PromptEncoder(nn.Module): ) self.not_a_point_embed = operations.Embedding(1, embed_dim, device=device, dtype=dtype) self.invalid_point_embed = operations.Embedding(1, embed_dim, device=device, dtype=dtype) + self._joint_w_cache = None # Mask prompt: 5-stage 2x2 strided conv downscaling to embed_dim. LN2d = LayerNorm2d_op(operations) @@ -77,16 +78,31 @@ class PromptEncoder(nn.Module): # PE compute in fp32 for precision (sin/cos of large coords), then cast back to the embedding weight dtype weight_dtype = self.invalid_point_embed.weight.dtype point_embedding = self.pe_layer._encode(points.to(torch.float)).to(weight_dtype) - point_embedding[labels == -2] = 0.0 # invalid points - point_embedding[labels == -2] += cast_to_input(self.invalid_point_embed.weight, point_embedding) - point_embedding[labels == -1] = 0.0 - point_embedding[labels == -1] += cast_to_input(self.not_a_point_embed.weight, point_embedding) - for i in range(self.num_body_joints): - point_embedding[labels == i] += cast_to_input(self.point_embeddings[i].weight, point_embedding) + + # One gather over the stacked joint table. + joint_w = self._joint_embed_weights(point_embedding) + idx = labels.long().clamp(0, self.num_body_joints - 1) + is_joint = ((labels >= 0) & (labels < self.num_body_joints)).unsqueeze(-1) + point_embedding = point_embedding + joint_w[idx] * is_joint.to(point_embedding.dtype) + + # -2/-1 zero the PE first, so the embedding replaces it outright. + invalid_w = cast_to_input(self.invalid_point_embed.weight, point_embedding, copy=False) + not_a_point_w = cast_to_input(self.not_a_point_embed.weight, point_embedding, copy=False) + point_embedding = torch.where((labels == -2).unsqueeze(-1), invalid_w, point_embedding) + point_embedding = torch.where((labels == -1).unsqueeze(-1), not_a_point_w, point_embedding) point_mask = labels > -2 return point_embedding, point_mask + def _joint_embed_weights(self, ref: torch.Tensor) -> torch.Tensor: + """(num_body_joints, C) stack of the per-joint embeddings, cached per device/dtype.""" + cached = self._joint_w_cache + if cached is not None and cached.device == ref.device and cached.dtype == ref.dtype: + return cached + w = cast_to_input(torch.cat([e.weight for e in self.point_embeddings], dim=0), ref, copy=False) + self._joint_w_cache = w + return w + def _get_batch_size(self, keypoints: Optional[torch.Tensor], boxes: Optional[torch.Tensor], masks: Optional[torch.Tensor]) -> int: if keypoints is not None: return keypoints.shape[0] diff --git a/comfy_extras/nodes_sam3d_body.py b/comfy_extras/nodes_sam3d_body.py index 12cba4e3b..a8584057a 100644 --- a/comfy_extras/nodes_sam3d_body.py +++ b/comfy_extras/nodes_sam3d_body.py @@ -234,7 +234,7 @@ class SAM3DBody_Predict(io.ComfyNode): mhr_pose_data = { "frames": frames_out, - "faces": inner.head_pose.faces.cpu().numpy(), + "faces": inner.head_pose.faces_np(), "image_size": (int(H), int(W)), "canonical_colors": inner.canonical_colors, "hand_vert_mask": inner.hand_vert_mask, diff --git a/comfy_extras/sam3d_body/rasterizer.py b/comfy_extras/sam3d_body/rasterizer.py index b6c590a52..8aa343735 100644 --- a/comfy_extras/sam3d_body/rasterizer.py +++ b/comfy_extras/sam3d_body/rasterizer.py @@ -18,6 +18,23 @@ _CANONICAL_PRESETS = {"rainbow", "rainbow_face_normal", "rainbow_face_semantic"} _rainbow_cache: dict = {} +_faces_cache: dict = {} + +def _faces_to_device(faces, device) -> torch.Tensor: + """Device-side face indices. Topology is static per model, but the render + loop calls this once per frame, so keep the H2D copy out of the loop.""" + key = (id(faces), device) + hit = _faces_cache.get(key) + if hit is not None: + return hit[1] + t = torch.as_tensor(np.asarray(faces, dtype=np.int64), device=device) + # Hold the source array too: id() alone could be recycled after a GC. + _faces_cache[key] = (faces, t) + if len(_faces_cache) > 4: + _faces_cache.pop(next(iter(_faces_cache))) + return t + + def rainbow_colors_from_canonical( positions: np.ndarray, tilt_x_deg: float = 0.0, @@ -372,7 +389,7 @@ def render_pose_data_torch( return bg.clamp(0.0, 1.0) return torch.zeros((H, W, 3), device=device, dtype=torch.float32) - faces = torch.as_tensor(np.asarray(pose_data["faces"], dtype=np.int64), device=device) + faces = _faces_to_device(pose_data["faces"], device) canonical_colors = pose_data.get("canonical_colors") using_canonical = shader_preset in _CANONICAL_PRESETS diff --git a/comfy_extras/sam3d_body/utils.py b/comfy_extras/sam3d_body/utils.py index bcab02a14..c3ab2ad8b 100644 --- a/comfy_extras/sam3d_body/utils.py +++ b/comfy_extras/sam3d_body/utils.py @@ -388,7 +388,7 @@ def compute_canonical_colors(model) -> Dict[str, np.ndarray]: norm (Nv,3 in [0,1]), face_mask, head_mask, and face_region_rgb (per-region painted color from the .safetensors).""" verts = model.head_pose.canonical_vertices().float().cpu().numpy() - faces = model.head_pose.faces.cpu().numpy() + faces = model.head_pose.faces_np() v0 = verts[faces[:, 0]] v1 = verts[faces[:, 1]]