summaryrefslogtreecommitdiff
path: root/src/mcp_sd15/preprocessors/openpose.py
blob: 05c879bbf52e4c8a018b964b142d47fe1c119515 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
"""OpenPose body or body+hand+face preprocessing."""

from __future__ import annotations

from PIL import Image


class OpenPosePreprocessor:
    """Lazily loads the fixed OpenPose annotator from the local cache."""

    def __init__(self) -> None:
        self._detector: object | None = None

    def __call__(self, image: Image.Image, *, full: bool) -> Image.Image:
        if self._detector is None:
            from controlnet_aux import OpenposeDetector

            self._detector = OpenposeDetector.from_pretrained(
                "lllyasviel/Annotators", local_files_only=True
            )
        result = self._detector(  # type: ignore[operator]
            image,
            include_body=True,
            include_hand=full,
            include_face=full,
        )
        return result.convert("RGB")