summaryrefslogtreecommitdiff
path: root/src/mcp_sd15/preprocessors
diff options
context:
space:
mode:
Diffstat (limited to 'src/mcp_sd15/preprocessors')
-rw-r--r--src/mcp_sd15/preprocessors/__init__.py1
-rw-r--r--src/mcp_sd15/preprocessors/depth.py24
-rw-r--r--src/mcp_sd15/preprocessors/openpose.py27
3 files changed, 52 insertions, 0 deletions
diff --git a/src/mcp_sd15/preprocessors/__init__.py b/src/mcp_sd15/preprocessors/__init__.py
new file mode 100644
index 0000000..60e22ba
--- /dev/null
+++ b/src/mcp_sd15/preprocessors/__init__.py
@@ -0,0 +1 @@
+"""Fixed ControlNet preprocessors."""
diff --git a/src/mcp_sd15/preprocessors/depth.py b/src/mcp_sd15/preprocessors/depth.py
new file mode 100644
index 0000000..861b889
--- /dev/null
+++ b/src/mcp_sd15/preprocessors/depth.py
@@ -0,0 +1,24 @@
+"""MiDaS depth-map preprocessing."""
+
+from __future__ import annotations
+
+from typing import cast
+
+from PIL import Image
+
+
+class DepthPreprocessor:
+ """Lazily loads the fixed MiDaS annotator from the local cache."""
+
+ def __init__(self) -> None:
+ self._detector: object | None = None
+
+ def __call__(self, image: Image.Image) -> Image.Image:
+ if self._detector is None:
+ from controlnet_aux import MidasDetector
+
+ self._detector = MidasDetector.from_pretrained(
+ "lllyasviel/Annotators", local_files_only=True
+ )
+ result = cast(Image.Image, self._detector(image)) # type: ignore[operator]
+ return result.convert("RGB")
diff --git a/src/mcp_sd15/preprocessors/openpose.py b/src/mcp_sd15/preprocessors/openpose.py
new file mode 100644
index 0000000..05c879b
--- /dev/null
+++ b/src/mcp_sd15/preprocessors/openpose.py
@@ -0,0 +1,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")