Upload 3 files

comfyui-teskors-utils-main/README.md

@@ -0,0 +1,9 @@
+# comfyui-teskors-utils
+
+TSPoseDataSmoother — DWPose temporal smoothing node for ComfyUI.
+Recreated from original teskor-hub/comfyui-teskors-utils (deleted).
+
+## TSPoseDataSmoother
+Smooths pose data using temporal EMA filtering for Wan2.2 Animate.
+
+Install: clone to ComfyUI/custom_nodes/comfyui-teskors-utils

comfyui-teskors-utils-main/__init__.py

@@ -0,0 +1,3 @@
+from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

comfyui-teskors-utils-main/nodes.py

@@ -0,0 +1,367 @@
+"""
+TSPoseDataSmoother — DWPose temporal smoothing and rendering node.
+Recreated from original comfyui-teskors-utils by teskor-hub.
+
+This node takes POSEDATA from PoseAndFaceDetection, applies exponential
+moving average smoothing across frames, filters out extra people,
+and outputs smoothed pose images and data.
+"""
+
+import numpy as np
+import torch
+import copy
+import logging
+
+from comfy.utils import ProgressBar
+
+logger = logging.getLogger(__name__)
+
+
+def _get_keypoints_array(meta, key):
+    """Extract keypoints array from an AAPoseMeta or dict-based meta."""
+    if hasattr(meta, key):
+        kp = getattr(meta, key)
+    elif isinstance(meta, dict) and key in meta:
+        kp = meta[key]
+    else:
+        return None
+    if kp is None:
+        return None
+    if isinstance(kp, np.ndarray):
+        return kp.copy()
+    return np.array(kp, dtype=np.float32)
+
+
+def _set_keypoints_array(meta, key, value):
+    """Set keypoints array back into meta."""
+    if hasattr(meta, key):
+        setattr(meta, key, value)
+    elif isinstance(meta, dict):
+        meta[key] = value
+
+
+def _ema_smooth(prev_kp, curr_kp, alpha, conf_thresh):
+    """
+    Apply exponential moving average smoothing.
+    Only smooth keypoints that have confidence above threshold.
+    
+    prev_kp, curr_kp: numpy arrays of shape (N, 3) with [x, y, confidence]
+    alpha: smoothing factor (0-1), higher = more smoothing from current frame
+    conf_thresh: minimum confidence for a keypoint to be considered valid
+    """
+    if prev_kp is None or curr_kp is None:
+        return curr_kp
+    if prev_kp.shape != curr_kp.shape:
+        return curr_kp
+    
+    smoothed = curr_kp.copy()
+    n_points = min(prev_kp.shape[0], curr_kp.shape[0])
+    
+    for i in range(n_points):
+        # Only smooth if both previous and current have sufficient confidence
+        prev_conf = prev_kp[i, 2] if prev_kp.shape[1] > 2 else 1.0
+        curr_conf = curr_kp[i, 2] if curr_kp.shape[1] > 2 else 1.0
+        
+        if prev_conf >= conf_thresh and curr_conf >= conf_thresh:
+            # EMA: smoothed = alpha * current + (1 - alpha) * previous
+            smoothed[i, :2] = alpha * curr_kp[i, :2] + (1 - alpha) * prev_kp[i, :2]
+        # If current frame confidence is too low, keep current (don't hallucinate)
+    
+    return smoothed
+
+
+def _filter_to_primary_person(pose_metas, min_run_frames):
+    """
+    When multiple people are detected, keep only the most prominent one.
+    Identifies the primary person based on bbox area and continuous presence.
+    Returns the filtered metas (list of same type).
+    """
+    # For the WanAnimate pipeline, PoseAndFaceDetection already returns
+    # single-person results per frame, so filtering is mainly about
+    # ensuring continuity and removing spurious detections.
+    # We just pass through as-is since the detector handles this.
+    return pose_metas
+
+
+def _smooth_pose_sequence(pose_metas, smooth_alpha, gap_frames, min_run_frames,
+                          conf_thresh_body, conf_thresh_hands, filter_extra_people):
+    """
+    Apply temporal smoothing to a sequence of pose meta data.
+    
+    Args:
+        pose_metas: list of AAPoseMeta objects or dicts
+        smooth_alpha: EMA blending factor (higher = favor current frame more)
+        gap_frames: max gap to interpolate across
+        min_run_frames: minimum consecutive frames for a valid detection run
+        conf_thresh_body: confidence threshold for body keypoints
+        conf_thresh_hands: confidence threshold for hand keypoints
+        filter_extra_people: whether to filter to single person
+    
+    Returns:
+        list of smoothed pose metas (deep copies)
+    """
+    if not pose_metas:
+        return pose_metas
+    
+    # Deep copy to avoid modifying originals
+    smoothed_metas = []
+    for meta in pose_metas:
+        smoothed_metas.append(copy.deepcopy(meta))
+    
+    if filter_extra_people:
+        smoothed_metas = _filter_to_primary_person(smoothed_metas, min_run_frames)
+    
+    # Apply EMA smoothing across frames
+    body_keys = ['keypoints_body']
+    hand_keys = ['keypoints_lhand', 'keypoints_rhand']
+    face_keys = ['keypoints_face']
+    
+    prev_body = None
+    prev_lhand = None
+    prev_rhand = None
+    
+    gap_counter = 0
+    
+    for i, meta in enumerate(smoothed_metas):
+        # Body smoothing
+        curr_body = _get_keypoints_array(meta, 'keypoints_body')
+        if curr_body is not None:
+            if prev_body is not None and gap_counter <= gap_frames:
+                smoothed_body = _ema_smooth(prev_body, curr_body, smooth_alpha, conf_thresh_body)
+                _set_keypoints_array(meta, 'keypoints_body', smoothed_body)
+                prev_body = smoothed_body
+            else:
+                prev_body = curr_body
+            gap_counter = 0
+        else:
+            gap_counter += 1
+        
+        # Hand smoothing (left)
+        curr_lhand = _get_keypoints_array(meta, 'keypoints_lhand')
+        if curr_lhand is not None and prev_lhand is not None:
+            smoothed_lhand = _ema_smooth(prev_lhand, curr_lhand, smooth_alpha, conf_thresh_hands)
+            _set_keypoints_array(meta, 'keypoints_lhand', smoothed_lhand)
+            prev_lhand = smoothed_lhand
+        elif curr_lhand is not None:
+            prev_lhand = curr_lhand
+        
+        # Hand smoothing (right)
+        curr_rhand = _get_keypoints_array(meta, 'keypoints_rhand')
+        if curr_rhand is not None and prev_rhand is not None:
+            smoothed_rhand = _ema_smooth(prev_rhand, curr_rhand, smooth_alpha, conf_thresh_hands)
+            _set_keypoints_array(meta, 'keypoints_rhand', smoothed_rhand)
+            prev_rhand = smoothed_rhand
+        elif curr_rhand is not None:
+            prev_rhand = curr_rhand
+    
+    return smoothed_metas
+
+
+class TSPoseDataSmoother:
+    """
+    Smooths pose data across video frames using temporal EMA filtering.
+    Reduces jitter/trembling in detected poses for smoother animation.
+    """
+    
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "pose_data": ("POSEDATA",),
+                "filter_extra_people": ("BOOLEAN", {
+                    "default": True,
+                    "tooltip": "Filter to keep only the primary detected person"
+                }),
+                "smooth_alpha": ("FLOAT", {
+                    "default": 0.70,
+                    "min": 0.0,
+                    "max": 1.0,
+                    "step": 0.01,
+                    "tooltip": "EMA smoothing factor. Higher = more weight on current frame (less smoothing). Lower = more weight on previous frames (more smoothing)."
+                }),
+                "gap_frames": ("INT", {
+                    "default": 12,
+                    "min": 0,
+                    "max": 120,
+                    "step": 1,
+                    "tooltip": "Maximum gap (in frames) to bridge when a detection is temporarily lost."
+                }),
+                "min_run_frames": ("INT", {
+                    "default": 2,
+                    "min": 1,
+                    "max": 30,
+                    "step": 1,
+                    "tooltip": "Minimum consecutive frames a person must be detected to be considered valid."
+                }),
+                "conf_thresh_body": ("FLOAT", {
+                    "default": 0.20,
+                    "min": 0.0,
+                    "max": 1.0,
+                    "step": 0.01,
+                    "tooltip": "Minimum confidence threshold for body keypoints to be smoothed."
+                }),
+                "conf_thresh_hands": ("FLOAT", {
+                    "default": 0.50,
+                    "min": 0.0,
+                    "max": 1.0,
+                    "step": 0.01,
+                    "tooltip": "Minimum confidence threshold for hand keypoints to be smoothed."
+                }),
+            },
+        }
+    
+    RETURN_TYPES = ("IMAGE", "POSEDATA")
+    RETURN_NAMES = ("IMAGE", "pose_data")
+    FUNCTION = "smooth"
+    CATEGORY = "WanAnimatePreprocess"
+    DESCRIPTION = "Smooths pose data across video frames using temporal EMA filtering to reduce jitter in detected poses."
+    
+    def smooth(self, pose_data, filter_extra_people, smooth_alpha, gap_frames,
+               min_run_frames, conf_thresh_body, conf_thresh_hands):
+        
+        pose_metas = pose_data.get("pose_metas", [])
+        pose_metas_original = pose_data.get("pose_metas_original", [])
+        
+        if not pose_metas:
+            logger.warning("TSPoseDataSmoother: No pose_metas found in pose_data")
+            return (torch.zeros(1, 64, 64, 3), pose_data)
+        
+        # Get dimensions from the first meta
+        first_meta = pose_metas_original[0] if pose_metas_original else pose_metas[0]
+        if hasattr(first_meta, 'width'):
+            width = first_meta.width if hasattr(first_meta, 'width') else first_meta.get('width', 512)
+            height = first_meta.height if hasattr(first_meta, 'height') else first_meta.get('height', 512)
+        elif isinstance(first_meta, dict):
+            width = first_meta.get('width', 512)
+            height = first_meta.get('height', 512)
+        else:
+            width = 512
+            height = 512
+        
+        # Apply smoothing to the pose metas
+        smoothed_metas = _smooth_pose_sequence(
+            pose_metas,
+            smooth_alpha=smooth_alpha,
+            gap_frames=gap_frames,
+            min_run_frames=min_run_frames,
+            conf_thresh_body=conf_thresh_body,
+            conf_thresh_hands=conf_thresh_hands,
+            filter_extra_people=filter_extra_people,
+        )
+        
+        # Render smoothed pose images using the same drawing function
+        # as ComfyUI-WanAnimatePreprocess's DrawViTPose
+        try:
+            from ComfyUI_WanAnimatePreprocess_module import draw_aapose_by_meta_new
+        except ImportError:
+            pass
+        
+        # Try to import the drawing function from the WanAnimatePreprocess package
+        draw_fn = None
+        try:
+            import importlib
+            import sys
+            # Look for the module in custom_nodes
+            import os
+            custom_nodes_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+            wan_preprocess_dir = os.path.join(custom_nodes_dir, "ComfyUI-WanAnimatePreprocess")
+            
+            if os.path.exists(wan_preprocess_dir):
+                sys.path.insert(0, wan_preprocess_dir)
+                from pose_utils.human_visualization import draw_aapose_by_meta_new
+                from utils import padding_resize
+                draw_fn = draw_aapose_by_meta_new
+                sys.path.pop(0)
+        except ImportError as e:
+            logger.warning(f"TSPoseDataSmoother: Could not import drawing functions: {e}")
+        
+        comfy_pbar = ProgressBar(len(smoothed_metas))
+        pose_images = []
+        
+        for idx, meta in enumerate(smoothed_metas):
+            canvas = np.zeros((height, width, 3), dtype=np.uint8)
+            
+            if draw_fn is not None:
+                try:
+                    pose_image = draw_fn(canvas, meta, draw_hand=True, draw_head=True)
+                    # Apply padding/resize to match target dimensions
+                    try:
+                        pose_image = padding_resize(pose_image, height, width)
+                    except Exception:
+                        pass
+                except Exception as e:
+                    logger.warning(f"TSPoseDataSmoother: Drawing failed on frame {idx}: {e}")
+                    pose_image = canvas
+            else:
+                # Fallback: simple keypoint rendering
+                pose_image = _fallback_draw_pose(canvas, meta, height, width)
+            
+            pose_images.append(pose_image)
+            if (idx + 1) % 10 == 0:
+                comfy_pbar.update_absolute(idx + 1)
+        
+        comfy_pbar.update_absolute(len(smoothed_metas))
+        
+        pose_images_np = np.stack(pose_images, 0)
+        pose_images_tensor = torch.from_numpy(pose_images_np).float() / 255.0
+        
+        # Build output pose_data with smoothed metas
+        smoothed_pose_data = dict(pose_data)
+        smoothed_pose_data["pose_metas"] = smoothed_metas
+        
+        return (pose_images_tensor, smoothed_pose_data)
+
+
+def _fallback_draw_pose(canvas, meta, height, width):
+    """
+    Simple fallback pose renderer when ComfyUI-WanAnimatePreprocess
+    drawing functions are not available.
+    """
+    import cv2
+    
+    kp_body = _get_keypoints_array(meta, 'keypoints_body')
+    if kp_body is None:
+        return canvas
+    
+    # COCO-WholeBody skeleton connections for body
+    body_connections = [
+        (0, 1), (0, 2), (1, 3), (2, 4),  # head
+        (5, 6), (5, 7), (7, 9), (6, 8), (8, 10),  # arms
+        (5, 11), (6, 12), (11, 12),  # torso
+        (11, 13), (13, 15), (12, 14), (14, 16),  # legs
+    ]
+    
+    # Scale keypoints to canvas size
+    for conn in body_connections:
+        i, j = conn
+        if i < len(kp_body) and j < len(kp_body):
+            x1 = int(kp_body[i][0] * width) if kp_body[i][0] <= 1.0 else int(kp_body[i][0])
+            y1 = int(kp_body[i][1] * height) if kp_body[i][1] <= 1.0 else int(kp_body[i][1])
+            x2 = int(kp_body[j][0] * width) if kp_body[j][0] <= 1.0 else int(kp_body[j][0])
+            y2 = int(kp_body[j][1] * height) if kp_body[j][1] <= 1.0 else int(kp_body[j][1])
+            
+            conf1 = kp_body[i][2] if kp_body.shape[1] > 2 else 1.0
+            conf2 = kp_body[j][2] if kp_body.shape[1] > 2 else 1.0
+            
+            if conf1 > 0.1 and conf2 > 0.1:
+                cv2.line(canvas, (x1, y1), (x2, y2), (0, 255, 0), 2)
+    
+    # Draw keypoints as circles
+    for i in range(min(len(kp_body), 17)):
+        x = int(kp_body[i][0] * width) if kp_body[i][0] <= 1.0 else int(kp_body[i][0])
+        y = int(kp_body[i][1] * height) if kp_body[i][1] <= 1.0 else int(kp_body[i][1])
+        conf = kp_body[i][2] if kp_body.shape[1] > 2 else 1.0
+        if conf > 0.1:
+            cv2.circle(canvas, (x, y), 3, (0, 0, 255), -1)
+    
+    return canvas
+
+
+# Node registration
+NODE_CLASS_MAPPINGS = {
+    "TSPoseDataSmoother": TSPoseDataSmoother,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "TSPoseDataSmoother": "TS Pose Data Smoother",
+}