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

Merge pull request #23 from corkborg/camie-tagger

Browse Source 
Add Camie tagger.
This commit is contained in:
corkborg
2025-04-10 14:01:29 +09:00
committed by GitHub
parent dd2a01431f c4d410be52
commit a341652c0a
10 changed files with 472 additions and 286 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+21 -20
View File
@@ -12,8 +12,7 @@ pip install -r requirements.txt
## usage
```
usage: run.py [-h] (--dir DIR | --file FILE) [--threshold THRESHOLD] [--ext EXT] [--overwrite] [--cpu] [--rawtag] [--recursive] [--exclude-tag t1,t2,t3]
[--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,wd-v1-4-vit-tagger.v3,wd-v1-4-convnext-tagger.v3,wd-v1-4-swinv2-tagger.v3,mld-caformer.dec-5-97527,mld-tresnetd.6-30000}]
usage: run.py [-h] (--dir DIR | --file FILE) [--threshold THRESHOLD] [--ext EXT] [--overwrite] [--cpu] [--rawtag] [--recursive] [--exclude-tag t1,t2,t3] [--model MODELNAME]
options:
-h, --help show this help message and exit
@@ -28,8 +27,7 @@ options:
--recursive Enable recursive file search
--exclude-tag t1,t2,t3
Specify tags to exclude (Need comma-separated list)
--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,wd-v1-4-vit-tagger.v3,wd-v1-4-convnext-tagger.v3,wd-v1-4-swinv2-tagger.v3,mld-caformer.dec-5-97527,mld-tresnetd.6-30000}
modelname to use for prediction (default is wd14-convnextv2.v1)
--model MODELNAME modelname to use for prediction (default is wd14-convnextv2.v1)
```
single file
@@ -47,33 +45,36 @@ python run.py --dir dir/dir
## Support Models
```
# Camie Tagger (released 2025)
python run.py --model camie-tagger --file image.jpg
# SmilingWolf large model. (released 2024)
python run.py --file image.jpg --model wd-vit-large-tagger-v3
python run.py --file image.jpg --model wd-eva02-large-tagger-v3
python run.py --model wd-vit-large-tagger-v3 --file image.jpg
python run.py --model wd-eva02-large-tagger-v3 --file image.jpg
# SmilingWolf v3 model. (released 2024)
python run.py --file image.jpg --model wd-v1-4-vit-tagger.v3
python run.py --file image.jpg --model wd-v1-4-convnext-tagger.v3
python run.py --file image.jpg --model wd-v1-4-swinv2-tagger.v3
python run.py --model wd-v1-4-vit-tagger.v3 --file image.jpg
python run.py --model wd-v1-4-convnext-tagger.v3 --file image.jpg
python run.py --model wd-v1-4-swinv2-tagger.v3 --file image.jpg
# SmilingWolf v2 model. (released 2023)
python run.py --file image.jpg --model wd-v1-4-moat-tagger.v2
python run.py --file image.jpg --model wd14-vit.v2
python run.py --file image.jpg --model wd14-convnext.v2
python run.py --model wd-v1-4-moat-tagger.v2 --file image.jpg
python run.py --model wd14-vit.v2 --file image.jpg
python run.py --model wd14-convnext.v2 --file image.jpg
# SmilingWolf v1 model. (released 2022)
python run.py --file image.jpg --model wd14-vit.v1
python run.py --file image.jpg --model wd14-convnext.v1
python run.py --file image.jpg --model wd14-convnextv2.v1
python run.py --file image.jpg --model wd14-swinv2-v1
python run.py --model wd14-vit.v1 --file image.jpg
python run.py --model wd14-convnext.v1 --file image.jpg
python run.py --model wd14-convnextv2.v1 --file image.jpg
python run.py --model wd14-swinv2-v1 --file image.jpg
# Z3D-E621-Convnext
python run.py --file image.jpg --model z3d-e621-convnext-toynya
python run.py --file image.jpg --model z3d-e621-convnext-silveroxides
python run.py --model z3d-e621-convnext-toynya --file image.jpg
python run.py --model z3d-e621-convnext-silveroxides --file image.jpg
# kiriyamaX model.
python run.py --file image.jpg --model mld-caformer.dec-5-97527
python run.py --file image.jpg --model mld-tresnetd.6-30000
python run.py --model mld-caformer.dec-5-97527 --file image.jpg
python run.py --model mld-tresnetd.6-30000 --file image.jpg
```
## Using GPU
requirements.txt
+4 -2
View File
@@ -1,7 +1,9 @@
numpy==2.2.2
pandas==2.2.3
huggingface-hub==0.28.1
torchvision==0.21.0
Pillow==11.1.0
huggingface-hub==0.28.1
deepdanbooru==1.0.3
onnxruntime==1.20.1
numpy==2.2.2
opencv-python==4.11.0.86
run.py
+6 -4
View File
@@ -1,5 +1,7 @@
import sys
from typing import Generator, Iterable
from tagger.interrogator import Interrogator
from tagger.interrogator.interrogator import AbsInterrogator
from PIL import Image
from pathlib import Path
import argparse
@@ -46,7 +48,7 @@ parser.add_argument(
parser.add_argument(
'--model',
default='wd14-convnextv2.v1',
choices=list(interrogators.keys()),
metavar='MODELNAME',
help='modelname to use for prediction (default is wd14-convnextv2.v1)')
args = parser.parse_args()
@@ -79,7 +81,7 @@ def image_interrogate(image_path: Path, tag_escape: bool, exclude_tags: Iterable
im = Image.open(image_path)
result = interrogator.interrogate(im)
return Interrogator.postprocess_tags(
return AbsInterrogator.postprocess_tags(
result[1],
threshold=args.threshold,
escape_tag=tag_escape,
@@ -116,7 +118,7 @@ if args.dir:
if args.file:
tags = image_interrogate(Path(args.file), not args.rawtag, parse_exclude_tags())
print()
print(file=sys.stderr)
tags_str = ', '.join(tags.keys())
print(tags_str)
tagger/interrogator.py
-256
View File
@@ -1,256 +0,0 @@
import os
import pandas as pd
import numpy as np
from typing import Iterable, Tuple, List, Dict
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
class Interrogator:
@staticmethod
def postprocess_tags(
tags: Dict[str, float],
threshold=0.35,
additional_tags: List[str] = [],
exclude_tags: Iterable[str] = [],
sort_by_alphabetical_order=False,
add_confident_as_weight=False,
replace_underscore=False,
replace_underscore_excludes: List[str] = [],
escape_tag=False
) -> Dict[str, float]:
for t in additional_tags:
tags[t] = 1.0
# those lines are totally not "pythonic" but looks better to me
tags = {
t: c
# sort by tag name or confident
for t, c in sorted(
tags.items(),
key=lambda i: i[0 if sort_by_alphabetical_order else 1],
reverse=not sort_by_alphabetical_order
)
# filter tags
if (
c >= threshold
and t not in exclude_tags
)
}
new_tags = []
for tag in list(tags):
new_tag = tag
if replace_underscore and tag not in replace_underscore_excludes:
new_tag = new_tag.replace('_', ' ')
if escape_tag:
new_tag = tag_escape_pattern.sub(r'\\\1', new_tag)
if add_confident_as_weight:
new_tag = f'({new_tag}:{tags[tag]})'
new_tags.append((new_tag, tags[tag]))
tags = dict(new_tags)
return tags
def __init__(self, name: str) -> None:
self.name = name
self.providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
def load(self):
raise NotImplementedError()
def unload(self) -> bool:
unloaded = False
if hasattr(self, 'model') and self.model is not None:
del self.model
unloaded = True
print(f'Unloaded {self.name}')
if hasattr(self, 'tags'):
del self.tags
return unloaded
def use_cpu(self) -> None:
self.providers = ['CPUExecutionProvider']
def interrogate(
self,
image: Image
) -> Tuple[
Dict[str, float], # rating confidents
Dict[str, float] # tag confidents
]:
raise NotImplementedError()
class WaifuDiffusionInterrogator(Interrogator):
def __init__(
self,
name: str,
model_path='model.onnx',
tags_path='selected_tags.csv',
**kwargs
) -> None:
super().__init__(name)
self.model_path = model_path
self.tags_path = tags_path
self.kwargs = kwargs
def download(self) -> Tuple[os.PathLike, os.PathLike]:
print(f"Loading {self.name} model file from {self.kwargs['repo_id']}")
model_path = Path(hf_hub_download(
**self.kwargs, filename=self.model_path))
tags_path = Path(hf_hub_download(
**self.kwargs, filename=self.tags_path))
return model_path, tags_path
def load(self) -> None:
model_path, tags_path = self.download()
from onnxruntime import InferenceSession
self.model = InferenceSession(str(model_path), providers=self.providers)
print(f'Loaded {self.name} model from {model_path}')
self.tags = pd.read_csv(tags_path)
def interrogate(
self,
image: Image
) -> Tuple[
Dict[str, float], # rating confidents
Dict[str, float] # tag confidents
]:
# init model
if not hasattr(self, 'model') or self.model is None:
self.load()
# code for converting the image and running the model is taken from the link below
# thanks, SmilingWolf!
# https://huggingface.co/spaces/SmilingWolf/wd-v1-4-tags/blob/main/app.py
# convert an image to fit the model
_, height, _, _ = self.model.get_inputs()[0].shape
# alpha to white
image = image.convert('RGBA')
new_image = Image.new('RGBA', image.size, 'WHITE')
new_image.paste(image, mask=image)
image = new_image.convert('RGB')
image = np.asarray(image)
# PIL RGB to OpenCV BGR
image = image[:, :, ::-1]
image = dbimutils.make_square(image, height)
image = dbimutils.smart_resize(image, height)
image = image.astype(np.float32)
image = np.expand_dims(image, 0)
# evaluate model
input_name = self.model.get_inputs()[0].name
label_name = self.model.get_outputs()[0].name
confidents = self.model.run([label_name], {input_name: image})[0]
tags = self.tags[:][['name']]
tags['confidents'] = confidents[0]
# first 4 items are for rating (general, sensitive, questionable, explicit)
ratings = dict(tags[:4].values)
# rest are regular tags
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
self.model = InferenceSession(model_path,
providers=self.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/interrogator/__init__.py
View File
tagger/interrogator/camietaggerinterrogator.py
+143
View File
@@ -0,0 +1,143 @@
"""
Interrogator for the Camie-Tagger model.
reference:
https://huggingface.co/Camais03/camie-tagger/blob/main/onnx_inference.py
"""
import sys
import json
from typing import cast
import numpy as np
import torch
from PIL import Image
from huggingface_hub import hf_hub_download
from onnxruntime import InferenceSession
from tagger.interrogator.interrogator import AbsInterrogator
class CamieTaggerInterrogator(AbsInterrogator):
repo_id: str
model_path: str
tags_path: str
def __init__(
self,
name: str,
repo_id: str,
model_path: str,
tags_path='metadata.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}", file=sys.stderr)
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()
self.model = InferenceSession(model_path,
providers=self.providers)
print(f'Loaded {self.name} model from {model_path}', file=sys.stderr)
with open(tags_path, 'r', encoding='utf-8') as filen:
self.metadata = json.load(filen)
def interrogate(
self,
image: Image.Image
) -> tuple[
dict[str, float], # rating confidents
dict[str, float] # tag confidents
]:
# init model
if self.model is None:
self.load()
if self.model is None:
raise Exception("Model not loading.")
img_tensor = preprocess_image(image)
img_numpy = img_tensor.unsqueeze(0).numpy()
input_ = self.model.get_inputs()[0]
# evaluate model
outputs = self.model.run(None, {input_.name: img_numpy})
# Process outputs
initial_probs: np.ndarray = 1.0 / (1.0 + np.exp(-outputs[0])) # Apply sigmoid
refined_probs: np.ndarray = 1.0 / (1.0 + np.exp(-outputs[1])) if len(outputs) > 1 else initial_probs
# Get top tags
indices = np.atleast_1d(refined_probs[0]).nonzero()[0]
# Group by category
tags_by_category = {}
for idx in indices:
idx_str = str(idx)
tag_name = self.metadata['idx_to_tag'].get(idx_str, f"unknown-{idx}")
category = self.metadata['tag_to_category'].get(tag_name, "general")
if category not in tags_by_category:
tags_by_category[category] = []
prob = float(refined_probs[0, idx])
tags_by_category[category].append((tag_name, prob))
# 'year', 'rating', 'general', 'character', 'copyright', 'artist', 'meta'
return dict(tags_by_category['rating']), dict(tags_by_category['general'])
def preprocess_image(img: Image.Image, image_size=512) -> torch.Tensor:
"""Process an image for inference"""
# Convert RGBA or Palette images to RGB
if img.mode in ('RGBA', 'P'):
img = img.convert('RGB')
# Get original dimensions
width, height = img.size
aspect_ratio = width / height
# Calculate new dimensions to maintain aspect ratio
if aspect_ratio > 1:
new_width = image_size
new_height = int(new_width / aspect_ratio)
else:
new_height = image_size
new_width = int(new_height * aspect_ratio)
# Resize with LANCZOS filter
img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
# Create new image with padding
new_image = Image.new('RGB', (image_size, image_size), (0, 0, 0))
paste_x = (image_size - new_width) // 2
paste_y = (image_size - new_height) // 2
new_image.paste(img, (paste_x, paste_y))
import torchvision.transforms as transforms
# Initialize transform
transform = transforms.Compose([
transforms.ToTensor(),
])
# Apply transforms
img_tensor= transform(new_image)
img_tensor = cast(torch.Tensor, img_tensor)
return img_tensor
tagger/interrogator/interrogator.py
+97
View File
@@ -0,0 +1,97 @@
import re
import sys
import pandas as pd
from typing import Iterable, Tuple, List, Dict
from PIL import Image
from onnxruntime import InferenceSession
tag_escape_pattern = re.compile(r'([\\()])')
class AbsInterrogator:
model: InferenceSession | None
tags: pd.DataFrame | None
@staticmethod
def postprocess_tags(
tags: Dict[str, float],
threshold=0.35,
additional_tags: List[str] = [],
exclude_tags: Iterable[str] = [],
sort_by_alphabetical_order=False,
add_confident_as_weight=False,
replace_underscore=False,
replace_underscore_excludes: List[str] = [],
escape_tag=False
) -> Dict[str, float]:
for t in additional_tags:
tags[t] = 1.0
# those lines are totally not "pythonic" but looks better to me
tags = {
t: c
# sort by tag name or confident
for t, c in sorted(
tags.items(),
key=lambda i: i[0 if sort_by_alphabetical_order else 1],
reverse=not sort_by_alphabetical_order
)
# filter tags
if (
c >= threshold
and t not in exclude_tags
)
}
new_tags = []
for tag in list(tags):
new_tag = tag
if replace_underscore and tag not in replace_underscore_excludes:
new_tag = new_tag.replace('_', ' ')
if escape_tag:
new_tag = tag_escape_pattern.sub(r'\\\1', new_tag)
if add_confident_as_weight:
new_tag = f'({new_tag}:{tags[tag]})'
new_tags.append((new_tag, tags[tag]))
tags = dict(new_tags)
return tags
def __init__(self, name: str) -> None:
self.name = name
self.providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
def load(self):
raise NotImplementedError()
def unload(self) -> bool:
unloaded = False
if hasattr(self, 'model') and self.model is not None:
del self.model
unloaded = True
print(f'Unloaded {self.name}', file=sys.stderr)
if hasattr(self, 'tags'):
del self.tags
return unloaded
def use_cpu(self) -> None:
self.providers = ['CPUExecutionProvider']
def interrogate(
self,
image: Image.Image
) -> Tuple[
Dict[str, float], # rating confidents
Dict[str, float] # tag confidents
]:
raise NotImplementedError()
tagger/interrogator/mldanbooruinterrogator.py
+87
View File
@@ -0,0 +1,87 @@
import sys
import json
from PIL import Image
from numpy import asarray, float32, expand_dims, exp
from huggingface_hub import hf_hub_download
from tagger.interrogator.interrogator import AbsInterrogator
import tagger.dbimutils as dbimutils
class MLDanbooruInterrogator(AbsInterrogator):
""" 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}", file=sys.stderr)
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
self.model = InferenceSession(model_path,
providers=self.providers)
print(f'Loaded {self.name} model from {model_path}', file=sys.stderr)
with open(tags_path, 'r', encoding='utf-8') as filen:
self.tags = json.load(filen)
def interrogate(
self,
image: 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))
if self.tags is None:
raise Exception("Tags not loading.")
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/interrogator/waifudiffusioninterrogator.py
+103
View File
@@ -0,0 +1,103 @@
import os
import sys
import pandas as pd
import numpy as np
from typing import Tuple, Dict
from PIL import Image
from pathlib import Path
from huggingface_hub import hf_hub_download
import re
tag_escape_pattern = re.compile(r'([\\()])')
from tagger.interrogator.interrogator import AbsInterrogator
import tagger.dbimutils as dbimutils
class WaifuDiffusionInterrogator(AbsInterrogator):
def __init__(
self,
name: str,
model_path='model.onnx',
tags_path='selected_tags.csv',
**kwargs
) -> None:
super().__init__(name)
self.model_path = model_path
self.tags_path = tags_path
self.kwargs = kwargs
def download(self) -> Tuple[os.PathLike, os.PathLike]:
print(f"Loading {self.name} model file from {self.kwargs['repo_id']}", file=sys.stderr)
model_path = Path(hf_hub_download(
**self.kwargs, filename=self.model_path))
tags_path = Path(hf_hub_download(
**self.kwargs, filename=self.tags_path))
return model_path, tags_path
def load(self) -> None:
model_path, tags_path = self.download()
from onnxruntime import InferenceSession
self.model = InferenceSession(str(model_path), providers=self.providers)
print(f'Loaded {self.name} model from {model_path}', file=sys.stderr)
self.tags = pd.read_csv(tags_path)
def interrogate(
self,
input_image: Image.Image
) -> tuple[
dict[str, float], # rating confidents
dict[str, float] # tag confidents
]:
# init model
if not hasattr(self, 'model') or self.model is None:
self.load()
if self.model is None:
raise Exception("Model not loading.")
# code for converting the image and running the model is taken from the link below
# thanks, SmilingWolf!
# https://huggingface.co/spaces/SmilingWolf/wd-v1-4-tags/blob/main/app.py
# convert an image to fit the model
_, height, _, _ = self.model.get_inputs()[0].shape
# alpha to white
image = input_image.convert('RGBA')
new_image = Image.new('RGBA', image.size, 'WHITE')
new_image.paste(image, mask=image)
image = new_image.convert('RGB')
image = np.asarray(image)
# PIL RGB to OpenCV BGR
image = image[:, :, ::-1]
image = dbimutils.make_square(image, height)
image = dbimutils.smart_resize(image, height)
image = image.astype(np.float32)
image = np.expand_dims(image, 0)
# evaluate model
input_name = self.model.get_inputs()[0].name
label_name = self.model.get_outputs()[0].name
confidents = self.model.run([label_name], {input_name: image})[0]
if self.tags is None:
raise Exception("Tags not loading.")
tags = self.tags[:][['name']]
tags['confidents'] = confidents[0]
# first 4 items are for rating (general, sensitive, questionable, explicit)
ratings = dict(tags[:4].values)
# rest are regular tags
tags = dict(tags[4:].values)
return ratings, tags
tagger/interrogators.py
+11 -4
View File
@@ -1,8 +1,9 @@
from typing import List, Dict
from tagger.interrogator.interrogator import AbsInterrogator
from tagger.interrogator.waifudiffusioninterrogator import WaifuDiffusionInterrogator
from tagger.interrogator.mldanbooruinterrogator import MLDanbooruInterrogator
from tagger.interrogator.camietaggerinterrogator import CamieTaggerInterrogator
from tagger.interrogator import Interrogator, WaifuDiffusionInterrogator, MLDanbooruInterrogator
interrogators: Dict[str, Interrogator] = {
interrogators: dict[str, AbsInterrogator] = {
'wd14-vit.v1': WaifuDiffusionInterrogator(
'WD14 ViT v1',
repo_id='SmilingWolf/wd-v1-4-vit-tagger'
@@ -73,4 +74,10 @@ interrogators: Dict[str, Interrogator] = {
repo_id='deepghs/ml-danbooru-onnx',
model_path='TResnet-D-FLq_ema_6-30000.onnx'
),
'camie-tagger': CamieTaggerInterrogator(
'Camie Tagger',
repo_id='Camais03/camie-tagger',
model_path='model_initial.onnx',
tags_path='model_initial_metadata.json'
),
}