pi-ux-ce/wd14-tagger-standalone
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Support multi models.

Browse Source
This commit is contained in:
corkborg
2023-12-28 02:28:15 +09:00
parent 57bc257b94
commit ba6c3b640a
5 changed files with 176 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+21 -5
View File
@@ -5,16 +5,16 @@ forked from [https://github.com/picobyte/stable-diffusion-webui-wd14-tagger]()
## install
```
venv/bin/python -m venv venv
python -m venv venv
source venv/bin/activate
pip install -r requirements.txt
```
## usage
```
usage: run.py [-h] [--dir DIR | --file FILE] [--threthold THRETHOLD] [--ext EXT]
usage: run.py [-h] (--dir DIR | --file FILE) [--threthold THRETHOLD] [--ext EXT]
[--model {wd14-vit.v1,wd14-vit.v2,wd14-convnext.v1,wd14-convnext.v2,wd14-convnextv2.v1,wd14-swinv2-v1,wd-v1-4-moat-tagger.v2,mld-caformer.dec-5-97527,mld-tresnetd.6-30000}]
options:
-h, --help show this help message and exit
@@ -23,16 +23,32 @@ options:
--threthold THRETHOLD
予測値の足切り確率（デフォルトは0.35）
--ext EXT dirの場合にキャプションファイルにつける拡張子
--model {wd14-vit.v1,wd14-vit.v2,wd14-convnext.v1,wd14-convnext.v2,wd14-convnextv2.v1,wd14-swinv2-v1,wd-v1-4-moat-tagger.v2,mld-caformer.dec-5-97527,mld-tresnetd.6-30000}
予測に使用するモデル名
```
single file
```
python3 run.py --file image.jpg
python run.py --file image.jpg
```
batch execution
```
python3 run.py --dir dir/dir
python run.py --dir dir/dir
```
## Support Models
```
python run.py --file image.jpg --model wd14-vit.v1
python run.py --file image.jpg --model wd14-vit.v2
python run.py --file image.jpg --model wd14-convnext.v1
python run.py --file image.jpg --model wd14-convnext.v2
python run.py --file image.jpg --model wd14-convnextv2.v1
python run.py --file image.jpg --model wd14-swinv2-v1
python run.py --file image.jpg --model wd-v1-4-moat-tagger.v2
python run.py --file image.jpg --model mld-caformer.dec-5-97527
python run.py --file image.jpg --model mld-tresnetd.6-30000
```
run.py
+10 -8
View File
@@ -6,9 +6,11 @@ from PIL import Image
from pathlib import Path
import argparse
from tagger.interrogators import interrogators
parser = argparse.ArgumentParser()
group = parser.add_mutually_exclusive_group()
group = parser.add_mutually_exclusive_group(required=True)
group.add_argument('--dir', help='ディレクトリ内の画像すべてに予測を行う')
group.add_argument('--file', help='ファイルに対して予測を行う')
@@ -21,16 +23,16 @@ parser.add_argument(
'--ext',
default='.txt',
help='dirの場合にキャプションファイルにつける拡張子')
parser.add_argument(
'--model',
default='wd14-convnextv2.v1',
choices=list(interrogators.keys()),
help='予測に使用するモデル名')
args = parser.parse_args()
# 使用するモデル
interrogator = WaifuDiffusionInterrogator(
'wd14-convnextv2-v2',
repo_id='SmilingWolf/wd-v1-4-convnextv2-tagger-v2',
revision='v2.0'
)
# 使用するinterrogatorを読み出す
interrogator = interrogators[args.model]
def image_interrogate(image_path: Path):
"""
tagger/dbimutils.py
+20
View File
@@ -4,6 +4,26 @@ import cv2
import numpy as np
from PIL import Image
def fill_transparent(image: Image.Image, color='WHITE'):
image = image.convert('RGBA')
new_image = Image.new('RGBA', image.size, color)
new_image.paste(image, mask=image)
image = new_image.convert('RGB')
return image
def resize(pic: Image.Image, size: int, keep_ratio=True) -> Image.Image:
if not keep_ratio:
target_size = (size, size)
else:
min_edge = min(pic.size)
target_size = (
int(pic.size[0] / min_edge * size),
int(pic.size[1] / min_edge * size),
)
target_size = (target_size[0] & ~3, target_size[1] & ~3)
return pic.resize(target_size, resample=Image.Resampling.LANCZOS)
def smart_imread(img, flag=cv2.IMREAD_UNCHANGED):
if img.endswith(".gif"):
tagger/interrogator.py
+79 -4
View File
@@ -9,16 +9,15 @@ from PIL import Image
from pathlib import Path
from huggingface_hub import hf_hub_download
import re
import json
from numpy import asarray, float32, expand_dims, exp
tag_escape_pattern = re.compile(r'([\\()])')
import tagger.dbimutils as dbimutils
use_cpu = True
if use_cpu:
tf_device_name = '/cpu:0'
else:
tf_device_name = '/gpu:0'
class Interrogator:
@staticmethod
@@ -187,3 +186,79 @@ class WaifuDiffusionInterrogator(Interrogator):
tags = dict(tags[4:].values)
return ratings, tags
class MLDanbooruInterrogator(Interrogator):
""" Interrogator for the MLDanbooru model. """
def __init__(
self,
name: str,
repo_id: str,
model_path: str,
tags_path='classes.json',
) -> None:
super().__init__(name)
self.model_path = model_path
self.tags_path = tags_path
self.repo_id = repo_id
self.tags = None
self.model = None
def download(self) -> Tuple[str, str]:
print(f"Loading {self.name} model file from {self.repo_id}")
model_path = hf_hub_download(
repo_id=self.repo_id,
filename=self.model_path
)
tags_path = hf_hub_download(
repo_id=self.repo_id,
filename=self.tags_path,
)
return model_path, tags_path
def load(self) -> None:
model_path, tags_path = self.download()
from onnxruntime import InferenceSession
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
if use_cpu:
providers.pop(0)
self.model = InferenceSession(model_path,
providers=providers)
print(f'Loaded {self.name} model from {model_path}')
with open(tags_path, 'r', encoding='utf-8') as filen:
self.tags = json.load(filen)
def interrogate(
self,
image: Image
) -> Tuple[
Dict[str, float], # rating confidents
Dict[str, float] # tag confidents
]:
# init model
if self.model is None:
self.load()
image = dbimutils.fill_transparent(image)
image = dbimutils.resize(image, 448) # TODO CUSTOMIZE
x = asarray(image, dtype=float32) / 255
# HWC -> 1CHW
x = x.transpose((2, 0, 1))
x = expand_dims(x, 0)
input_ = self.model.get_inputs()[0]
output = self.model.get_outputs()[0]
# evaluate model
y, = self.model.run([output.name], {input_.name: x})
# Softmax
y = 1 / (1 + exp(-y))
tags = {tag: float(conf) for tag, conf in zip(self.tags, y.flatten())}
return {}, tags
def large_batch_interrogate(self, images: List, dry_run=False) -> str:
raise NotImplementedError()
tagger/interrogators.py
+46
View File
@@ -0,0 +1,46 @@
from typing import List, Dict
from tagger.interrogator import Interrogator, WaifuDiffusionInterrogator, MLDanbooruInterrogator
interrogators: Dict[str, Interrogator] = {
'wd14-vit.v1': WaifuDiffusionInterrogator(
'WD14 ViT v1',
repo_id='SmilingWolf/wd-v1-4-vit-tagger'
),
'wd14-vit.v2': WaifuDiffusionInterrogator(
'WD14 ViT v2',
repo_id='SmilingWolf/wd-v1-4-vit-tagger-v2',
),
'wd14-convnext.v1': WaifuDiffusionInterrogator(
'WD14 ConvNeXT v1',
repo_id='SmilingWolf/wd-v1-4-convnext-tagger'
),
'wd14-convnext.v2': WaifuDiffusionInterrogator(
'WD14 ConvNeXT v2',
repo_id='SmilingWolf/wd-v1-4-convnext-tagger-v2',
),
'wd14-convnextv2.v1': WaifuDiffusionInterrogator(
'WD14 ConvNeXTV2 v1',
# the name is misleading, but it's v1
repo_id='SmilingWolf/wd-v1-4-convnextv2-tagger-v2',
),
'wd14-swinv2-v1': WaifuDiffusionInterrogator(
'WD14 SwinV2 v1',
# again misleading name
repo_id='SmilingWolf/wd-v1-4-swinv2-tagger-v2',
),
'wd-v1-4-moat-tagger.v2': WaifuDiffusionInterrogator(
'WD14 moat tagger v2',
repo_id='SmilingWolf/wd-v1-4-moat-tagger-v2'
),
'mld-caformer.dec-5-97527': MLDanbooruInterrogator(
'ML-Danbooru Caformer dec-5-97527',
repo_id='deepghs/ml-danbooru-onnx',
model_path='ml_caformer_m36_dec-5-97527.onnx'
),
'mld-tresnetd.6-30000': MLDanbooruInterrogator(
'ML-Danbooru TResNet-D 6-30000',
repo_id='deepghs/ml-danbooru-onnx',
model_path='TResnet-D-FLq_ema_6-30000.onnx'
),
}