h3_pose.py

"""Turn a motion reference into a pose render.

A reference clip drives how the new subject moves, and the ComfyUI workflows fed two kinds: the original footage, and
a pose render of it. The raw clip carries more — Nekodificador's best result used it, and pose renders transferred the
motion but flattened the acting — so this is offered rather than imposed. It earns its place when the reference shows
the wrong *subject*: a skeleton says "move like this" without also saying "look like this".

mediapipe rather than DWPose, for one measured reason: DWPose on CPU onnxruntime takes about 43 seconds per frame
(the yolox bounding-box pass is almost all of it), which is an hour and a half for a five-second clip. mediapipe's
pose landmarker is three orders of magnitude cheaper and needs no GPU.

Everything here fails soft. `available()` reports why it cannot run, and a caller that gets `None` from `to_pose_clip`
should use the reference as it came.
"""

from __future__ import annotations

import os
import tempfile

import numpy as np

MODEL_URL = (
    "https://storage.googleapis.com/mediapipe-models/pose_landmarker/pose_landmarker_full/float16/1/"
    "pose_landmarker_full.task"
)
# The limb pairs of mediapipe's 33-point pose, and a colour per limb group. Drawn by hand rather than with
# `mediapipe.solutions.drawing_utils`, because that module is the legacy API and imports the proto path that a modern
# protobuf breaks.
_EDGES = [
    (11, 12), (11, 23), (12, 24), (23, 24),                      # torso
    (11, 13), (13, 15), (12, 14), (14, 16),                      # arms
    (23, 25), (25, 27), (24, 26), (26, 28),                      # legs
    (27, 31), (28, 32), (15, 17), (16, 18),                      # feet and hands
    (0, 11), (0, 12),                                            # neck
]
_COLOURS = [
    (0, 255, 0), (0, 200, 120), (0, 200, 120), (0, 160, 200),
    (255, 170, 0), (255, 120, 0), (170, 255, 0), (120, 255, 0),
    (0, 120, 255), (0, 80, 255), (120, 0, 255), (170, 0, 255),
    (255, 0, 170), (255, 0, 120), (255, 220, 0), (220, 255, 0),
    (255, 60, 60), (60, 60, 255),
]
_model_path: str | None = None


def available() -> tuple[bool, str]:
    """Whether a pose render can be produced here, and why not when it cannot."""
    try:
        import mediapipe  # noqa: F401
        from mediapipe.tasks.python import vision  # noqa: F401
    except Exception as error:
        # The usual cause is a protobuf newer than mediapipe accepts (`protobuf<5`), which breaks its generated code.
        return False, f"mediapipe is unusable here ({type(error).__name__}: {error})"
    return True, ""


def _fetch_model() -> str:
    global _model_path
    if _model_path and os.path.exists(_model_path):
        return _model_path
    import requests

    path = os.path.join(tempfile.gettempdir(), "pose_landmarker_full.task")
    if not os.path.exists(path):
        response = requests.get(MODEL_URL, timeout=120)
        response.raise_for_status()
        with open(path, "wb") as handle:
            handle.write(response.content)
    _model_path = path
    return path


def to_pose_clip(frames: np.ndarray, fps: int = 24, progress=None) -> str | None:
    r"""
    Render `frames` as a pose skeleton on black and write it as a clip.

    Args:
        frames (`np.ndarray` of shape `(num_frames, height, width, 3)`): The reference clip, `uint8` RGB.
        fps (`int`, defaults to 24): The rate to write the result at.
        progress (`callable`, *optional*): Gradio progress callback.

    Returns:
        `str` or None: path to the pose clip, or None when no pose was found in any frame.
    """
    import av
    import mediapipe as mp
    from mediapipe.tasks import python as tasks
    from mediapipe.tasks.python import vision

    landmarker = vision.PoseLandmarker.create_from_options(
        vision.PoseLandmarkerOptions(
            base_options=tasks.BaseOptions(model_asset_path=_fetch_model()),
            running_mode=vision.RunningMode.VIDEO,
        )
    )

    num_frames, height, width = frames.shape[:3]
    canvases = np.zeros_like(frames)
    found = 0
    for index, frame in enumerate(frames):
        image = mp.Image(image_format=mp.ImageFormat.SRGB, data=np.ascontiguousarray(frame))
        result = landmarker.detect_for_video(image, int(index * 1000 / fps))
        if not result.pose_landmarks:
            continue
        found += 1
        _draw(canvases[index], result.pose_landmarks[0], width, height)
        if progress is not None:
            progress(min(1.0, (index + 1) / num_frames), desc=f"Reading the motion — frame {index + 1}/{num_frames}")
    if not found:
        return None

    path = os.path.join(tempfile.mkdtemp(), "pose.mp4")
    container = av.open(path, mode="w")
    stream = container.add_stream("libx264", rate=fps)
    stream.width, stream.height, stream.pix_fmt = width, height, "yuv420p"
    stream.options = {"crf": "18", "preset": "veryfast"}
    for canvas in canvases:
        container.mux(stream.encode(av.VideoFrame.from_ndarray(canvas, format="rgb24")))
    for packet in stream.encode():
        container.mux(packet)
    container.close()
    print(f"[pose] {found}/{num_frames} frames had a pose", flush=True)
    return path


def _draw(canvas: np.ndarray, landmarks, width: int, height: int) -> None:
    """One skeleton onto one frame, thick enough to survive the video encoder and the VAE."""
    import cv2

    points = [(int(point.x * width), int(point.y * height)) for point in landmarks]
    visible = [point.visibility > 0.4 for point in landmarks]
    thickness = max(2, round(min(width, height) / 160))
    for (start, end), colour in zip(_EDGES, _COLOURS):
        if start < len(points) and end < len(points) and visible[start] and visible[end]:
            cv2.line(canvas, points[start], points[end], colour, thickness, cv2.LINE_AA)
    for index, (point, seen) in enumerate(zip(points, visible)):
        if seen and any(index in edge for edge in _EDGES):
            cv2.circle(canvas, point, thickness + 1, (255, 255, 255), -1, cv2.LINE_AA)