Spaces:
Running on Zero
Running on Zero
File size: 1,935 Bytes
036940b 9a5065c 036940b 9a5065c 036940b 6d862f4 036940b 9a5065c 036940b | 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 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | """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]
|