mirror of
https://github.com/vladmandic/automatic
synced 2026-08-27 07:31:01 +02:00
vladmandic
22 lines
779 B
Python
22 lines
779 B
Python
import numpy as np
|
|
from PIL import Image
|
|
|
|
|
|
class TEEDDetector:
|
|
def __init__(self, model):
|
|
self.model = model
|
|
|
|
@classmethod
|
|
def from_pretrained(cls, pretrained_model_or_path="fal/teed", cache_dir=None, local_files_only=False): # pylint: disable=unused-argument
|
|
from installer import install
|
|
install('controlnet-aux', quiet=True)
|
|
from controlnet_aux import TEEDdetector as _TEEDdetector
|
|
model = _TEEDdetector.from_pretrained(pretrained_model_or_path, filename="5_model.pth")
|
|
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
|