| import os | |
| import re | |
| import torch | |
| from PIL import Image | |
| import numpy as np | |
| from modules import modelloader, paths, deepbooru_model, devices, images, shared | |
| re_special = re.compile(r'([\\()])') | |
| class DeepDanbooru: | |
| def __init__(self): | |
| self.model = None | |
| def load(self): | |
| if self.model is not None: | |
| return | |
| files = modelloader.load_models( | |
| model_path=os.path.join(paths.models_path, "torch_deepdanbooru"), | |
| model_url='https://github.com/AUTOMATIC1111/TorchDeepDanbooru/releases/download/v1/model-resnet_custom_v3.pt', | |
| ext_filter=[".pt"], | |
| download_name='model-resnet_custom_v3.pt', | |
| ) | |
| self.model = deepbooru_model.DeepDanbooruModel() | |
| self.model.load_state_dict(torch.load(files[0], map_location="cpu")) | |
| self.model.eval() | |
| self.model.to(devices.cpu, devices.dtype) | |
| def start(self): | |
| self.load() | |
| self.model.to(devices.device) | |
| def stop(self): | |
| if not shared.opts.interrogate_keep_models_in_memory: | |
| self.model.to(devices.cpu) | |
| devices.torch_gc() | |
| def tag(self, pil_image): | |
| self.start() | |
| res = self.tag_multi(pil_image) | |
| self.stop() | |
| return res | |
| def tag_multi(self, pil_image, force_disable_ranks=False): | |
| threshold = shared.opts.interrogate_deepbooru_score_threshold | |
| use_spaces = shared.opts.deepbooru_use_spaces | |
| use_escape = shared.opts.deepbooru_escape | |
| alpha_sort = shared.opts.deepbooru_sort_alpha | |
| include_ranks = shared.opts.interrogate_return_ranks and not force_disable_ranks | |
| pic = images.resize_image(2, pil_image.convert("RGB"), 512, 512) | |
| a = np.expand_dims(np.array(pic, dtype=np.float32), 0) / 255 | |
| with torch.no_grad(), devices.autocast(): | |
| x = torch.from_numpy(a).to(devices.device) | |
| y = self.model(x)[0].detach().cpu().numpy() | |
| probability_dict = {} | |
| for tag, probability in zip(self.model.tags, y): | |
| if probability < threshold: | |
| continue | |
| if tag.startswith("rating:"): | |
| continue | |
| probability_dict[tag] = probability | |
| if alpha_sort: | |
| tags = sorted(probability_dict) | |
| else: | |
| tags = [tag for tag, _ in sorted(probability_dict.items(), key=lambda x: -x[1])] | |
| res = [] | |
| filtertags = set([x.strip().replace(' ', '_') for x in shared.opts.deepbooru_filter_tags.split(",")]) | |
| for tag in [x for x in tags if x not in filtertags]: | |
| probability = probability_dict[tag] | |
| tag_outformat = tag | |
| if use_spaces: | |
| tag_outformat = tag_outformat.replace('_', ' ') | |
| if use_escape: | |
| tag_outformat = re.sub(re_special, r'\\\1', tag_outformat) | |
| if include_ranks: | |
| tag_outformat = f"({tag_outformat}:{probability:.3f})" | |
| res.append(tag_outformat) | |
| return ", ".join(res) | |
| model = DeepDanbooru() | |