comfyui/ComfyUI-anima-pose-control/node.py

"""AnimaPoseControl ComfyUI node: detect -> edit (JS) -> styled skeleton -> control IMAGE."""
import numpy as np

try:                                   # keep the helper importable without the package context (tests)
    from .pose_render import render, STYLES
    from .keypoints import from_json, to_json
    from .detect import detect
except ImportError:                    # flat import when running tests inside the package dir
    from pose_render import render, STYLES
    from keypoints import from_json, to_json
    from detect import detect


def compose(image, pose_json, style, hands, face, feet, redetect, resolution, estimator=None):
    """Pure branch logic. Returns (skeleton uint8 [H,W,3] RGB, pose_json_out str)."""
    if redetect and image is not None:
        kp, sc = detect(image, int(resolution), estimator=estimator)
    elif pose_json and pose_json.strip():
        kp, sc, _ = from_json(pose_json)
    elif image is not None:
        kp, sc = detect(image, int(resolution), estimator=estimator)   # no pose yet -> detect
    else:
        raise ValueError("AnimaPoseControl: no pose to render — connect an image, or Queue once with redetect on first")
    img = render(kp, sc, style, int(resolution), hands=hands, face=face, feet=feet)
    return img, to_json(kp, sc, int(resolution))


class AnimaPoseControl:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "style": (list(STYLES),),
                "hands": ("BOOLEAN", {"default": True}),
                "face": ("BOOLEAN", {"default": True}),
                "feet": ("BOOLEAN", {"default": True}),
                "redetect": ("BOOLEAN", {"default": True}),
                "resolution": ("INT", {"default": 768, "min": 256, "max": 2048, "step": 64}),
                "pose_json": ("STRING", {"multiline": True, "default": ""}),
            },
            "optional": {"image": ("IMAGE",)},
        }

    RETURN_TYPES = ("IMAGE", "POSE_KEYPOINT")
    RETURN_NAMES = ("skeleton", "keypoints")
    FUNCTION = "run"
    OUTPUT_NODE = True
    CATEGORY = "Anima/pose"

    def run(self, style, hands, face, feet, redetect, resolution, pose_json, image=None):
        import torch
        img_in = None
        if image is not None:
            img_in = (image[0].cpu().numpy() * 255).astype(np.uint8)   # ComfyUI IMAGE -> HWC uint8
        skel, pj = compose(img_in, pose_json, style, hands, face, feet, redetect, resolution)
        t = torch.from_numpy(skel.astype(np.float32) / 255.0)[None, ...]
        import json
        kpts = json.loads(pj)
        # Show the styled render inline via ComfyUI's standard image preview.
        return {"ui": {"images": _save_preview(skel), "pose_json": [pj]}, "result": (t, kpts)}


def _save_preview(skel_uint8):
    """Write the styled skeleton to ComfyUI's temp dir; return a /view descriptor for the JS to load."""
    import folder_paths, os, uuid
    from PIL import Image
    d = folder_paths.get_temp_directory(); os.makedirs(d, exist_ok=True)
    fn = "anima_pose_%s.png" % uuid.uuid4().hex[:8]
    Image.fromarray(skel_uint8).save(os.path.join(d, fn))
    return [{"filename": fn, "subfolder": "", "type": "temp"}]