preprocessors.py

"""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:
    # controlnet_aux's Processor takes "depth_midas", NOT "midas".
    # Plain "midas" is not in its MODELS dict and raises KeyError.
    proc = _get_processor("depth_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]