mirror of https://github.com/vladmandic/automatic
22 lines
809 B
Python
22 lines
809 B
Python
import numpy as np
|
|
from PIL import Image
|
|
|
|
|
|
class AnylineDetector:
|
|
def __init__(self, model):
|
|
self.model = model
|
|
|
|
@classmethod
|
|
def from_pretrained(cls, pretrained_model_or_path="TheMistoAI/MistoLine", cache_dir=None, local_files_only=False):
|
|
from installer import install
|
|
install('controlnet-aux', quiet=True)
|
|
from controlnet_aux import AnylineDetector as _AnylineDetector
|
|
model = _AnylineDetector.from_pretrained(pretrained_model_or_path, filename="MTEED.pth", subfolder="Anyline", cache_dir=cache_dir)
|
|
return cls(model)
|
|
|
|
def __call__(self, image, output_type="pil", **kwargs):
|
|
if isinstance(image, np.ndarray):
|
|
image = Image.fromarray(image)
|
|
result = self.model(image, output_type=output_type)
|
|
return result
|