BuildingVirtuallyAI/RenderingModel

1import gc
2import PIL.Image
3import torch
4from controlnet_aux import LineartDetector
5
6class Preprocessor:
7    MODEL_ID = "lllyasviel/Annotators"
8
9    def __init__(self):
10        self.model = None
11        self.name = ""
12
13    def load(self, name: str) -> None:
14        if name == self.name:
15            return
16        if name == "Lineart":
17            self.model = LineartDetector.from_pretrained(self.MODEL_ID)
18
19        else:
20            raise ValueError
21        torch.cuda.empty_cache()
22        gc.collect()
23        self.name = name
24
25    def __call__(self, image: PIL.Image.Image, **kwargs) -> PIL.Image.Image:
26            return self.model(image, **kwargs)
27
28