Spaces:
Runtime error
Runtime error
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() | |