Optimize speed a little bit

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,
         }

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):

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]

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,

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

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]]