feat(preprocessors): canny/depth/pose via controlnet_aux (lazy imports)

preprocessors.py

@@ -0,0 +1,59 @@
+"""ControlNet preprocessors — lazy imports so an unused mode pays no cost."""
+from __future__ import annotations
+
+from typing import Any
+
+from PIL import Image
+
+MODES: tuple[str, ...] = ("Canny", "Depth", "Pose", "Pre-processed")
+
+
+def run(mode: str, image: Image.Image | None) -> Image.Image:
+    if image is None:
+        raise ValueError("preprocessor needs an input image")
+    if mode == "Canny":
+        return _run_canny(image)
+    if mode == "Depth":
+        return _run_depth(image)
+    if mode == "Pose":
+        return _run_pose(image)
+    if mode == "Pre-processed":
+        return image
+    raise ValueError(f"unknown preprocessor mode: {mode!r}; expected one of {MODES}")
+
+
+def _run_canny(image: Image.Image) -> Image.Image:
+    import cv2
+    import numpy as np
+    arr = np.array(image.convert("RGB"))
+    gray = cv2.cvtColor(arr, cv2.COLOR_RGB2GRAY)
+    edges = cv2.Canny(gray, threshold1=100, threshold2=200)
+    rgb = cv2.cvtColor(edges, cv2.COLOR_GRAY2RGB)
+    return Image.fromarray(rgb)
+
+
+def _run_depth(image: Image.Image) -> Image.Image:
+    proc = _get_processor("midas")
+    out: Any = proc(image)
+    if isinstance(out, Image.Image):
+        return out.convert("RGB")
+    return Image.fromarray(out).convert("RGB")
+
+
+def _run_pose(image: Image.Image) -> Image.Image:
+    proc = _get_processor("openpose")
+    out: Any = proc(image)
+    if isinstance(out, Image.Image):
+        return out.convert("RGB")
+    return Image.fromarray(out).convert("RGB")
+
+
+_PROCESSOR_CACHE: dict[str, Any] = {}
+
+
+def _get_processor(name: str) -> Any:
+    """Lazy-init and cache a controlnet_aux Processor."""
+    if name not in _PROCESSOR_CACHE:
+        from controlnet_aux.processor import Processor
+        _PROCESSOR_CACHE[name] = Processor(name)
+    return _PROCESSOR_CACHE[name]

tests/test_preprocessors.py

@@ -0,0 +1,37 @@
+import numpy as np
+import pytest
+from PIL import Image
+
+import preprocessors
+
+
+@pytest.fixture
+def gradient_image():
+    arr = np.linspace(0, 255, 256 * 256, dtype=np.uint8).reshape(256, 256)
+    return Image.fromarray(arr).convert("RGB")
+
+
+def test_modes_are_listed():
+    assert preprocessors.MODES == ("Canny", "Depth", "Pose", "Pre-processed")
+
+
+def test_canny_returns_rgb_image_of_same_size(gradient_image):
+    out = preprocessors.run("Canny", gradient_image)
+    assert isinstance(out, Image.Image)
+    assert out.size == gradient_image.size
+    assert out.mode == "RGB"
+
+
+def test_passthrough_returns_input_unchanged(gradient_image):
+    out = preprocessors.run("Pre-processed", gradient_image)
+    assert out is gradient_image
+
+
+def test_unknown_mode_raises():
+    with pytest.raises(ValueError):
+        preprocessors.run("Sobel", Image.new("RGB", (32, 32)))
+
+
+def test_run_with_image_none_raises():
+    with pytest.raises(ValueError):
+        preprocessors.run("Canny", None)