mirror of
https://github.com/modelscope/DiffSynth-Studio.git
synced 2026-03-19 06:23:43 +00:00
51 lines
1.9 KiB
Python
51 lines
1.9 KiB
Python
from typing_extensions import Literal, TypeAlias
|
|
import warnings
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter("ignore")
|
|
from controlnet_aux.processor import (
|
|
CannyDetector, MidasDetector, HEDdetector, LineartDetector, LineartAnimeDetector, OpenposeDetector
|
|
)
|
|
|
|
|
|
Processor_id: TypeAlias = Literal[
|
|
"canny", "depth", "softedge", "lineart", "lineart_anime", "openpose", "tile"
|
|
]
|
|
|
|
class Annotator:
|
|
def __init__(self, processor_id: Processor_id, model_path="models/Annotators", detect_resolution=512):
|
|
if processor_id == "canny":
|
|
self.processor = CannyDetector()
|
|
elif processor_id == "depth":
|
|
self.processor = MidasDetector.from_pretrained(model_path)
|
|
elif processor_id == "softedge":
|
|
self.processor = HEDdetector.from_pretrained(model_path)
|
|
elif processor_id == "lineart":
|
|
self.processor = LineartDetector.from_pretrained(model_path)
|
|
elif processor_id == "lineart_anime":
|
|
self.processor = LineartAnimeDetector.from_pretrained(model_path)
|
|
elif processor_id == "openpose":
|
|
self.processor = OpenposeDetector.from_pretrained(model_path)
|
|
elif processor_id == "tile":
|
|
self.processor = None
|
|
else:
|
|
raise ValueError(f"Unsupported processor_id: {processor_id}")
|
|
|
|
self.processor_id = processor_id
|
|
self.detect_resolution = detect_resolution
|
|
|
|
def __call__(self, image):
|
|
width, height = image.size
|
|
if self.processor_id == "openpose":
|
|
kwargs = {
|
|
"include_body": True,
|
|
"include_hand": True,
|
|
"include_face": True
|
|
}
|
|
else:
|
|
kwargs = {}
|
|
if self.processor is not None:
|
|
image = self.processor(image, detect_resolution=self.detect_resolution, image_resolution=min(width, height), **kwargs)
|
|
image = image.resize((width, height))
|
|
return image
|
|
|