Upload 25 files

.gitignore

@@ -0,0 +1,6 @@
+__pycache__/
+.vscode/
+.venv/
+.env
+
+presets/

README.ko.md

@@ -0,0 +1,61 @@
+[Automatic1111 웹UI](https://github.com/AUTOMATIC1111/stable-diffusion-webui)를 위한 태깅(라벨링) 확장 기능
+---
+DeepDanbooru 와 같은 모델을 통해 단일 또는 여러 이미지로부터 부루에서 사용하는 태그를 알아냅니다.
+
+[You don't know how to read Korean? Read it in English here!](README.md)
+
+## 들어가기 앞서
+모델과 대부분의 코드는 제가 만들지 않았고 [DeepDanbooru](https://github.com/KichangKim/DeepDanbooru) 와 MrSmillingWolf 의 태거에서 가져왔습니다.
+
+## 설치하기
+1. *확장기능* -> *URL로부터 확장기능 설치* -> 이 레포지토리 주소 입력 -> *설치*
+   - 또는 이 레포지토리를 `extensions/` 디렉터리 내에 클론합니다.
+      ```sh
+      $ git clone https://github.com/toriato/stable-diffusion-webui-wd14-tagger.git extensions/tagger
+      ```
+
+1. 모델 추가하기
+   - #### *MrSmilingWolf's model (a.k.a. Waifu Diffusion 1.4 tagger)*
+      처음 실행할 때 [HuggingFace 레포지토리](https://huggingface.co/SmilingWolf/wd-v1-4-vit-tagger)로부터 자동으로 받아옵니다.
+
+      모델과 관련된 또는 추가 학습에 대한 질문은 원작자인 MrSmilingWolf#5991 으로 물어봐주세요.
+
+   - #### *DeepDanbooru*
+      1. 다양한 모델 파일은 아래 주소에서 찾을 수 있습니다.
+         - [DeepDanbooru model](https://github.com/KichangKim/DeepDanbooru/releases)
+         - [e621 model by 🐾Zack🐾#1984](https://discord.gg/BDFpq9Yb7K)
+            *(NSFW 주의!)*
+
+      1. 모델과 설정 파일이 포함된 프로젝트 폴더를 `models/deepdanbooru` 경로로 옮깁니다.
+
+      1. 파일 구조는 다음과 같습니다:
+         ```
+         models/
+         └╴deepdanbooru/
+           ├╴deepdanbooru-v3-20211112-sgd-e28/
+           │ ├╴project.json
+           │ └╴...
+           │
+           ├╴deepdanbooru-v4-20200814-sgd-e30/
+           │ ├╴project.json
+           │ └╴...
+           │
+           ├╴e621-v3-20221117-sgd-e32/
+           │ ├╴project.json
+           │ └╴...
+           │
+           ...
+         ```
+
+1. 웹UI 를 시작하거나 재시작합니다.
+   - 또는 *Interrogator* 드롭다운 상자 우측에 있는 새로고침 버튼을 누릅니다.
+
+
+## 스크린샷
+![Screenshot](docs/screenshot.png)
+
+Artwork made by [hecattaart](https://vk.com/hecattaart?w=wall-89063929_3767)
+
+## 저작권
+
+빌려온 코드(예: `dbimutils.py`)를 제외하고 모두 Public domain 

README.md

@@ -0,0 +1,118 @@
+Tagger for [Automatic1111's WebUI](https://github.com/AUTOMATIC1111/stable-diffusion-webui)
+---
+Interrogate booru style tags for single or multiple image files using various models, such as DeepDanbooru.
+
+[한국어를 사용하시나요? 여기에 한국어 설명서가 있습니다!](README.ko.md)
+
+## Disclaimer
+I didn't make any models, and most of the code was heavily borrowed from the [DeepDanbooru](https://github.com/KichangKim/DeepDanbooru) and MrSmillingWolf's tagger.
+
+## Installation
+1. *Extensions* -> *Install from URL* -> Enter URL of this repository -> Press *Install* button
+   - or clone this repository under `extensions/`
+      ```sh
+      $ git clone https://github.com/toriato/stable-diffusion-webui-wd14-tagger.git extensions/tagger
+      ```
+
+1. Add interrogate model
+   - #### *MrSmilingWolf's model (a.k.a. Waifu Diffusion 1.4 tagger)*
+      Downloads automatically from the [HuggingFace repository](https://huggingface.co/SmilingWolf/wd-v1-4-vit-tagger) the first time you run it.
+
+      Please ask the original author MrSmilingWolf#5991 for questions related to model or additional training.
+
+      ##### ViT vs Convnext
+      > To make it clear: the ViT model is the one used to tag images for WD 1.4. That's why the repo was originally called like that. This one has been trained on the same data and tags, but has got no other relation to WD 1.4, aside from stemming from the same coordination effort. They were trained in parallel, and the best one at the time was selected for WD 1.4
+
+      > This particular model was trained later and might actually be slightly better than the ViT one. Difference is in the noise range tho 
+
+      — [SmilingWolf](https://github.com/SmilingWolf) from [this thread](https://discord.com/channels/930499730843250783/1052283314997837955) in the [東方Project AI server](https://discord.com/invite/touhouai) 
+
+   - #### *DeepDanbooru*
+      1. Various model files can be found below.
+         - [DeepDanbooru models](https://github.com/KichangKim/DeepDanbooru/releases)
+         - [e621 model by 🐾Zack🐾#1984](https://discord.gg/BDFpq9Yb7K)
+            *(link contains NSFW contents!)*
+
+      1. Move the project folder containing the model and config to `models/deepdanbooru`
+
+      1. The file structure should look like:
+         ```
+         models/
+         └╴deepdanbooru/
+           ├╴deepdanbooru-v3-20211112-sgd-e28/
+           │ ├╴project.json
+           │ └╴...
+           │
+           ├╴deepdanbooru-v4-20200814-sgd-e30/
+           │ ├╴project.json
+           │ └╴...
+           │
+           ├╴e621-v3-20221117-sgd-e32/
+           │ ├╴project.json
+           │ └╴...
+           │
+           ...
+         ```
+
+1. Start or restart the WebUI.
+   - or you can press refresh button after *Interrogator* dropdown box.
+
+
+## Model comparison
+
+* Used image: [hecattaart's artwork](https://vk.com/hecattaart?w=wall-89063929_3767)
+* Threshold: `0.5`
+
+### DeepDanbooru
+Used the same image as the one used in the Screenshot item
+
+#### [`deepdanbooru-v3-20211112-sgd-e28`](https://github.com/KichangKim/DeepDanbooru/releases/tag/v3-20211112-sgd-e28)
+```
+1girl, animal ears, cat ears, cat tail, clothes writing, full body, rating:safe, shiba inu, shirt, shoes, simple background, sneakers, socks, solo, standing, t-shirt, tail, white background, white shirt
+```
+
+#### [`deepdanbooru-v4-20200814-sgd-e30`](https://github.com/KichangKim/DeepDanbooru/releases/tag/v4-20200814-sgd-e30)
+```
+1girl, animal, animal ears, bottomless, clothes writing, full body, rating:safe, shirt, shoes, short sleeves, sneakers, solo, standing, t-shirt, tail, white background, white shirt
+```
+
+#### `e621-v3-20221117-sgd-e32`
+```
+anthro, bottomwear, clothing, footwear, fur, hi res, mammal, shirt, shoes, shorts, simple background, sneakers, socks, solo, standing, text on clothing, text on topwear, topwear, white background
+```
+
+### Waifu Diffusion Tagger
+
+#### [`wd14-vit`](https://huggingface.co/SmilingWolf/wd-v1-4-vit-tagger)
+```
+1boy, animal ears, dog, furry, leg hair, male focus, shirt, shoes, simple background, socks, solo, tail, white background
+```
+
+#### [`wd14-convnext`](https://huggingface.co/SmilingWolf/wd-v1-4-convnext-tagger)
+```
+full body, furry, shirt, shoes, simple background, socks, solo, tail, white background
+```
+
+#### [`wd14-vit-v2`](https://huggingface.co/SmilingWolf/wd-v1-4-vit-tagger-v2)
+```
+1boy, animal ears, cat, furry, male focus, shirt, shoes, simple background, socks, solo, tail, white background
+```
+
+#### [`wd14-convnext-v2`](https://huggingface.co/SmilingWolf/wd-v1-4-convnext-tagger-v2)
+```
+animal focus, clothes writing, earrings, full body, meme, shirt, shoes, simple background, socks, solo, sweat, tail, white background, white shirt
+```
+
+#### [`wd14-swinv2-v2`](https://huggingface.co/SmilingWolf/wd-v1-4-swinv2-tagger-v2)
+```
+1boy, arm hair, black footwear, cat, dirty, full body, furry, leg hair, male focus, shirt, shoes, simple background, socks, solo, standing, tail, white background, white shirt
+```
+
+## Screenshot
+![Screenshot](docs/screenshot.png)
+
+Artwork made by [hecattaart](https://vk.com/hecattaart?w=wall-89063929_3767)
+
+## Copyright
+
+Public domain, except borrowed parts (e.g. `dbimutils.py`)

javascript/tagger.js

@@ -0,0 +1,108 @@
+/**
+ * wait until element is loaded and returns
+ * @param {string} selector
+ * @param {number} timeout 
+ * @param {Element} $rootElement
+ * @returns {Promise<HTMLElement>}
+ */
+function waitQuerySelector(selector, timeout = 5000, $rootElement = gradioApp()) {
+    return new Promise((resolve, reject) => {
+        const element = $rootElement.querySelector(selector)
+        if (document.querySelector(element)) {
+            return resolve(element)
+        }
+
+        let timeoutId
+
+        const observer = new MutationObserver(() => {
+            const element = $rootElement.querySelector(selector)
+            if (!element) {
+                return
+            }
+
+            if (timeoutId) {
+                clearInterval(timeoutId)
+            }
+
+            observer.disconnect()
+            resolve(element)
+        })
+
+        timeoutId = setTimeout(() => {
+            observer.disconnect()
+            reject(new Error(`timeout, cannot find element by '${selector}'`))
+        }, timeout)
+
+        observer.observe($rootElement, {
+            childList: true,
+            subtree: true
+        })
+    })
+}
+
+document.addEventListener('DOMContentLoaded', () => {
+    Promise.all([
+        // option texts
+        waitQuerySelector('#additioanl-tags'),
+        waitQuerySelector('#exclude-tags'),
+
+        // tag-confident labels
+        waitQuerySelector('#rating-confidents'),
+        waitQuerySelector('#tag-confidents')
+    ]).then(elements => {
+
+        const $additionalTags = elements[0].querySelector('textarea')
+        const $excludeTags = elements[1].querySelector('textarea')
+        const $ratingConfidents = elements[2]
+        const $tagConfidents = elements[3]
+
+        let $selectedTextarea = $additionalTags
+
+        /**
+         * @this {HTMLElement}
+         * @param {MouseEvent} e
+         * @listens document#click
+         */
+        function onClickTextarea(e) {
+            $selectedTextarea = this
+        }
+
+        $additionalTags.addEventListener('click', onClickTextarea)
+        $excludeTags.addEventListener('click', onClickTextarea)
+
+        /**
+         * @this {HTMLElement}
+         * @param {MouseEvent} e
+         * @listens document#click
+         */
+        function onClickLabels(e) {
+            // find clicked label item's wrapper element
+            const $tag = e.target.closest('.output-label > div:not(:first-child)')
+            if (!$tag) {
+                return
+            }
+
+            /** @type {string} */
+            const tag = $tag.querySelector('.leading-snug').textContent
+
+            // ignore if tag is already exist in textbox
+            const escapedTag = tag.replace(/[.*+?^${}()|[\]\\]/g, '\\$&')
+            const pattern = new RegExp(`(^|,)\\s{0,}${escapedTag}\\s{0,}($|,)`)
+            if (pattern.test($selectedTextarea.value)) {
+                return
+            }
+
+            if ($selectedTextarea.value !== '') {
+                $selectedTextarea.value += ', '
+            }
+
+            $selectedTextarea.value += tag
+        }
+
+        $ratingConfidents.addEventListener('click', onClickLabels)
+        $tagConfidents.addEventListener('click', onClickLabels)
+
+    }).catch(err => {
+        console.error(err)
+    })
+})

preload.py

@@ -0,0 +1,27 @@
+from pathlib import Path
+from argparse import ArgumentParser
+
+from modules.shared import models_path
+
+default_ddp_path = Path(models_path, 'deepdanbooru')
+default_onnx_path = Path(models_path, 'TaggerOnnx')
+
+
+def preload(parser: ArgumentParser):
+    # default deepdanbooru use different paths:
+    # models/deepbooru and models/torch_deepdanbooru
+    # https://github.com/AUTOMATIC1111/stable-diffusion-webui/commit/c81d440d876dfd2ab3560410f37442ef56fc6632
+
+    parser.add_argument(
+        '--deepdanbooru-projects-path',
+        type=str,
+        help='Path to directory with DeepDanbooru project(s).',
+        default=default_ddp_path
+    )
+    parser.add_argument(
+        '--onnxtagger-path',
+        type=str,
+        help='Path to directory with DeepDanbooru project(s).',
+        default=default_onnx_path
+    )
+

scripts/tagger.py

@@ -0,0 +1,18 @@
+from PIL import Image, ImageFile
+
+from modules import script_callbacks
+from tagger.api import on_app_started
+from tagger.ui import on_ui_tabs
+
+
+# if you do not initialize the Image object
+# Image.registered_extensions() returns only PNG
+Image.init()
+
+# PIL spits errors when loading a truncated image by default
+# https://pillow.readthedocs.io/en/stable/reference/ImageFile.html#PIL.ImageFile.LOAD_TRUNCATED_IMAGES
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+
+script_callbacks.on_app_started(on_app_started)
+script_callbacks.on_ui_tabs(on_ui_tabs)

style.css

@@ -0,0 +1,7 @@
+#rating-confidents .output-label>div:not(:first-child) {
+    cursor: pointer;
+}
+
+#tag-confidents .output-label>div:not(:first-child) {
+    cursor: pointer;
+}

tagger/api.py

@@ -0,0 +1,90 @@
+from typing import Callable
+from threading import Lock
+from secrets import compare_digest
+
+from modules import shared
+from modules.api.api import decode_base64_to_image
+from modules.call_queue import queue_lock
+from fastapi import FastAPI, Depends, HTTPException
+from fastapi.security import HTTPBasic, HTTPBasicCredentials
+
+from tagger import utils
+from tagger import api_models as models
+
+
+class Api:
+    def __init__(self, app: FastAPI, queue_lock: Lock, prefix: str = None) -> None:
+        if shared.cmd_opts.api_auth:
+            self.credentials = dict()
+            for auth in shared.cmd_opts.api_auth.split(","):
+                user, password = auth.split(":")
+                self.credentials[user] = password
+
+        self.app = app
+        self.queue_lock = queue_lock
+        self.prefix = prefix
+
+        self.add_api_route(
+            'interrogate',
+            self.endpoint_interrogate,
+            methods=['POST'],
+            response_model=models.TaggerInterrogateResponse
+        )
+
+        self.add_api_route(
+            'interrogators',
+            self.endpoint_interrogators,
+            methods=['GET'],
+            response_model=models.InterrogatorsResponse
+        )
+
+    def auth(self, creds: HTTPBasicCredentials = Depends(HTTPBasic())):
+        if creds.username in self.credentials:
+            if compare_digest(creds.password, self.credentials[creds.username]):
+                return True
+
+        raise HTTPException(
+            status_code=401,
+            detail="Incorrect username or password",
+            headers={
+                "WWW-Authenticate": "Basic"
+            })
+
+    def add_api_route(self, path: str, endpoint: Callable, **kwargs):
+        if self.prefix:
+            path = f'{self.prefix}/{path}'
+
+        if shared.cmd_opts.api_auth:
+            return self.app.add_api_route(path, endpoint, dependencies=[Depends(self.auth)], **kwargs)
+        return self.app.add_api_route(path, endpoint, **kwargs)
+
+    def endpoint_interrogate(self, req: models.TaggerInterrogateRequest):
+        if req.image is None:
+            raise HTTPException(404, 'Image not found')
+
+        if req.model not in utils.interrogators.keys():
+            raise HTTPException(404, 'Model not found')
+
+        image = decode_base64_to_image(req.image)
+        interrogator = utils.interrogators[req.model]
+
+        with self.queue_lock:
+            ratings, tags = interrogator.interrogate(image)
+
+        return models.TaggerInterrogateResponse(
+            caption={
+                **ratings,
+                **interrogator.postprocess_tags(
+                    tags,
+                    req.threshold
+                )
+            })
+
+    def endpoint_interrogators(self):
+        return models.InterrogatorsResponse(
+            models=list(utils.interrogators.keys())
+        )
+
+
+def on_app_started(_, app: FastAPI):
+    Api(app, queue_lock, '/tagger/v1')

tagger/api_models.py

@@ -0,0 +1,33 @@
+from typing import List, Dict
+
+from modules.api import models as sd_models
+from pydantic import BaseModel, Field
+
+
+class TaggerInterrogateRequest(sd_models.InterrogateRequest):
+    model: str = Field(
+        title='Model',
+        description='The interrogate model used.'
+    )
+
+    threshold: float = Field(
+        default=0.35,
+        title='Threshold',
+        description='',
+        ge=0,
+        le=1
+    )
+
+
+class TaggerInterrogateResponse(BaseModel):
+    caption: Dict[str, float] = Field(
+        title='Caption',
+        description='The generated caption for the image.'
+    )
+
+
+class InterrogatorsResponse(BaseModel):
+    models: List[str] = Field(
+        title='Models',
+        description=''
+    )

tagger/dbimutils.py

@@ -0,0 +1,54 @@
+# DanBooru IMage Utility functions
+
+import cv2
+import numpy as np
+from PIL import Image
+
+
+def smart_imread(img, flag=cv2.IMREAD_UNCHANGED):
+    if img.endswith(".gif"):
+        img = Image.open(img)
+        img = img.convert("RGB")
+        img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
+    else:
+        img = cv2.imread(img, flag)
+    return img
+
+
+def smart_24bit(img):
+    if img.dtype is np.dtype(np.uint16):
+        img = (img / 257).astype(np.uint8)
+
+    if len(img.shape) == 2:
+        img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
+    elif img.shape[2] == 4:
+        trans_mask = img[:, :, 3] == 0
+        img[trans_mask] = [255, 255, 255, 255]
+        img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
+    return img
+
+
+def make_square(img, target_size):
+    old_size = img.shape[:2]
+    desired_size = max(old_size)
+    desired_size = max(desired_size, target_size)
+
+    delta_w = desired_size - old_size[1]
+    delta_h = desired_size - old_size[0]
+    top, bottom = delta_h // 2, delta_h - (delta_h // 2)
+    left, right = delta_w // 2, delta_w - (delta_w // 2)
+
+    color = [255, 255, 255]
+    new_im = cv2.copyMakeBorder(
+        img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color
+    )
+    return new_im
+
+
+def smart_resize(img, size):
+    # Assumes the image has already gone through make_square
+    if img.shape[0] > size:
+        img = cv2.resize(img, (size, size), interpolation=cv2.INTER_AREA)
+    elif img.shape[0] < size:
+        img = cv2.resize(img, (size, size), interpolation=cv2.INTER_CUBIC)
+    return img

tagger/format.py

@@ -0,0 +1,46 @@
+import re
+import hashlib
+
+from typing import Dict, Callable, NamedTuple
+from pathlib import Path
+
+
+class Info(NamedTuple):
+    path: Path
+    output_ext: str
+
+
+def hash(i: Info, algo='sha1') -> str:
+    try:
+        hash = hashlib.new(algo)
+    except ImportError:
+        raise ValueError(f"'{algo}' is invalid hash algorithm")
+
+    # TODO: is okay to hash large image?
+    with open(i.path, 'rb') as file:
+        hash.update(file.read())
+
+    return hash.hexdigest()
+
+
+pattern = re.compile(r'\[([\w:]+)\]')
+
+# all function must returns string or raise TypeError or ValueError
+# other errors will cause the extension error
+available_formats: Dict[str, Callable] = {
+    'name': lambda i: i.path.stem,
+    'extension': lambda i: i.path.suffix[1:],
+    'hash': hash,
+
+    'output_extension': lambda i: i.output_ext
+}
+
+
+def format(match: re.Match, info: Info) -> str:
+    matches = match[1].split(':')
+    name, args = matches[0], matches[1:]
+
+    if name not in available_formats:
+        return match[0]
+
+    return available_formats[name](info, *args)

tagger/interrogator.py

@@ -0,0 +1,320 @@
+import os
+import gc
+import pandas as pd
+import numpy as np
+
+from typing import Tuple, List, Dict
+from io import BytesIO
+from PIL import Image
+
+from pathlib import Path
+from huggingface_hub import hf_hub_download
+
+from modules import shared
+from modules.deepbooru import re_special as tag_escape_pattern
+
+# i'm not sure if it's okay to add this file to the repository
+from . import dbimutils
+
+# select a device to process
+use_cpu = ('all' in shared.cmd_opts.use_cpu) or (
+    'interrogate' in shared.cmd_opts.use_cpu)
+
+if use_cpu:
+    tf_device_name = '/cpu:0'
+else:
+    tf_device_name = '/gpu:0'
+
+    if shared.cmd_opts.device_id is not None:
+        try:
+            tf_device_name = f'/gpu:{int(shared.cmd_opts.device_id)}'
+        except ValueError:
+            print('--device-id is not a integer')
+
+
+class Interrogator:
+    @staticmethod
+    def postprocess_tags(
+        tags: Dict[str, float],
+
+        threshold=0.35,
+        additional_tags: List[str] = [],
+        exclude_tags: List[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]:
+
+        tags = {
+            **{t: 1.0 for t in additional_tags},
+            **tags
+        }
+
+        # 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
+
+    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 interrogate(
+        self,
+        image: Image
+    ) -> Tuple[
+        Dict[str, float],  # rating confidents
+        Dict[str, float]  # tag confidents
+    ]:
+        raise NotImplementedError()
+
+
+class DeepDanbooruInterrogator(Interrogator):
+    def __init__(self, name: str, project_path: os.PathLike) -> None:
+        super().__init__(name)
+        self.project_path = project_path
+
+    def load(self) -> None:
+        print(f'Loading {self.name} from {str(self.project_path)}')
+
+        # deepdanbooru package is not include in web-sd anymore
+        # https://github.com/AUTOMATIC1111/stable-diffusion-webui/commit/c81d440d876dfd2ab3560410f37442ef56fc663
+        from launch import is_installed, run_pip
+        if not is_installed('deepdanbooru'):
+            package = os.environ.get(
+                'DEEPDANBOORU_PACKAGE',
+                'git+https://github.com/KichangKim/DeepDanbooru.git@d91a2963bf87c6a770d74894667e9ffa9f6de7ff'
+            )
+
+            run_pip(
+                f'install {package} tensorflow tensorflow-io', 'deepdanbooru')
+
+        import tensorflow as tf
+
+        # tensorflow maps nearly all vram by default, so we limit this
+        # https://www.tensorflow.org/guide/gpu#limiting_gpu_memory_growth
+        # TODO: only run on the first run
+        for device in tf.config.experimental.list_physical_devices('GPU'):
+            tf.config.experimental.set_memory_growth(device, True)
+
+        with tf.device(tf_device_name):
+            import deepdanbooru.project as ddp
+
+            self.model = ddp.load_model_from_project(
+                project_path=self.project_path,
+                compile_model=False
+            )
+
+            print(f'Loaded {self.name} model from {str(self.project_path)}')
+
+            self.tags = ddp.load_tags_from_project(
+                project_path=self.project_path
+            )
+
+    def unload(self) -> bool:
+        # unloaded = super().unload()
+
+        # if unloaded:
+        #     # tensorflow suck
+        #     # https://github.com/keras-team/keras/issues/2102
+        #     import tensorflow as tf
+        #     tf.keras.backend.clear_session()
+        #     gc.collect()
+
+        # return unloaded
+
+        # There is a bug in Keras where it is not possible to release a model that has been loaded into memory.
+        # Downgrading to keras==2.1.6 may solve the issue, but it may cause compatibility issues with other packages.
+        # Using subprocess to create a new process may also solve the problem, but it can be too complex (like Automatic1111 did).
+        # It seems that for now, the best option is to keep the model in memory, as most users use the Waifu Diffusion model with onnx.
+
+        return False
+
+    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()
+
+        import deepdanbooru.data as ddd
+
+        # convert an image to fit the model
+        image_bufs = BytesIO()
+        image.save(image_bufs, format='PNG')
+        image = ddd.load_image_for_evaluate(
+            image_bufs,
+            self.model.input_shape[2],
+            self.model.input_shape[1]
+        )
+
+        image = image.reshape((1, *image.shape[0:3]))
+
+        # evaluate model
+        result = self.model.predict(image)
+
+        confidents = result[0].tolist()
+        ratings = {}
+        tags = {}
+
+        for i, tag in enumerate(self.tags):
+            tags[tag] = confidents[i]
+
+        return ratings, tags
+
+
+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]:
+        #if model_path exists, skip download
+        print(self.model_path, self.tags_path)
+        if os.path.exists(self.model_path) and os.path.exists(self.tags_path):
+            return self.model_path, self.tags_path
+        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()
+
+        # only one of these packages should be installed at a time in any one environment
+        # https://onnxruntime.ai/docs/get-started/with-python.html#install-onnx-runtime
+        # TODO: remove old package when the environment changes?
+        from launch import is_installed, run_pip
+        if not is_installed('onnxruntime'):
+            package = os.environ.get(
+                'ONNXRUNTIME_PACKAGE',
+                'onnxruntime-gpu'
+            )
+
+            run_pip(f'install {package}', 'onnxruntime')
+
+        from onnxruntime import InferenceSession
+
+        # https://onnxruntime.ai/docs/execution-providers/
+        # https://github.com/toriato/stable-diffusion-webui-wd14-tagger/commit/e4ec460122cf674bbf984df30cdb10b4370c1224#r92654958
+        providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
+        if use_cpu:
+            providers.pop(0)
+
+        self.model = InferenceSession(str(model_path), providers=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

tagger/preset.py

@@ -0,0 +1,108 @@
+import os
+import json
+
+from typing import Tuple, List, Dict
+from pathlib import Path
+from modules.images import sanitize_filename_part
+
+PresetDict = Dict[str, Dict[str, any]]
+
+
+class Preset:
+    base_dir: Path
+    default_filename: str
+    default_values: PresetDict
+    components: List[object]
+
+    def __init__(
+        self,
+        base_dir: os.PathLike,
+        default_filename='default.json'
+    ) -> None:
+        self.base_dir = Path(base_dir)
+        self.default_filename = default_filename
+        self.default_values = self.load(default_filename)[1]
+        self.components = []
+
+    def component(self, component_class: object, **kwargs) -> object:
+        # find all the top components from the Gradio context and create a path
+        from gradio.context import Context
+        parent = Context.block
+        paths = [kwargs['label']]
+
+        while parent is not None:
+            if hasattr(parent, 'label'):
+                paths.insert(0, parent.label)
+
+            parent = parent.parent
+
+        path = '/'.join(paths)
+
+        component = component_class(**{
+            **kwargs,
+            **self.default_values.get(path, {})
+        })
+
+        setattr(component, 'path', path)
+
+        self.components.append(component)
+        return component
+
+    def load(self, filename: str) -> Tuple[str, PresetDict]:
+        if not filename.endswith('.json'):
+            filename += '.json'
+
+        path = self.base_dir.joinpath(sanitize_filename_part(filename))
+        configs = {}
+
+        if path.is_file():
+            configs = json.loads(path.read_text())
+
+        return path, configs
+
+    def save(self, filename: str, *values) -> Tuple:
+        path, configs = self.load(filename)
+
+        for index, component in enumerate(self.components):
+            config = configs.get(component.path, {})
+            config['value'] = values[index]
+
+            for attr in ['visible', 'min', 'max', 'step']:
+                if hasattr(component, attr):
+                    config[attr] = config.get(attr, getattr(component, attr))
+
+            configs[component.path] = config
+
+        self.base_dir.mkdir(0o777, True, True)
+        path.write_text(
+            json.dumps(configs, indent=4)
+        )
+
+        return 'successfully saved the preset'
+
+    def apply(self, filename: str) -> Tuple:
+        values = self.load(filename)[1]
+        outputs = []
+
+        for component in self.components:
+            config = values.get(component.path, {})
+
+            if 'value' in config and hasattr(component, 'choices'):
+                if config['value'] not in component.choices:
+                    config['value'] = None
+
+            outputs.append(component.update(**config))
+
+        return (*outputs, 'successfully loaded the preset')
+
+    def list(self) -> List[str]:
+        presets = [
+            p.name
+            for p in self.base_dir.glob('*.json')
+            if p.is_file()
+        ]
+
+        if len(presets) < 1:
+            presets.append(self.default_filename)
+
+        return presets

tagger/ui.py

@@ -0,0 +1,464 @@
+import os
+import json
+import gradio as gr
+
+from pathlib import Path
+from glob import glob
+from PIL import Image, UnidentifiedImageError
+
+from webui import wrap_gradio_gpu_call
+from modules import ui
+from modules import generation_parameters_copypaste as parameters_copypaste
+
+from tagger import format, utils
+from tagger.utils import split_str
+from tagger.interrogator import Interrogator
+
+
+def unload_interrogators():
+    unloaded_models = 0
+
+    for i in utils.interrogators.values():
+        if i.unload():
+            unloaded_models = unloaded_models + 1
+
+    return [f'Successfully unload {unloaded_models} model(s)']
+
+
+def on_interrogate(
+    image: Image,
+    batch_input_glob: str,
+    batch_input_recursive: bool,
+    batch_output_dir: str,
+    batch_output_filename_format: str,
+    batch_output_action_on_conflict: str,
+    batch_output_save_json: bool,
+
+    interrogator: str,
+    threshold: float,
+    additional_tags: str,
+    exclude_tags: str,
+    sort_by_alphabetical_order: bool,
+    add_confident_as_weight: bool,
+    replace_underscore: bool,
+    replace_underscore_excludes: str,
+    escape_tag: bool,
+
+    unload_model_after_running: bool
+):
+    if interrogator not in utils.interrogators:
+        return ['', None, None, f"'{interrogator}' is not a valid interrogator"]
+
+    interrogator: Interrogator = utils.interrogators[interrogator]
+
+    postprocess_opts = (
+        threshold,
+        split_str(additional_tags),
+        split_str(exclude_tags),
+        sort_by_alphabetical_order,
+        add_confident_as_weight,
+        replace_underscore,
+        split_str(replace_underscore_excludes),
+        escape_tag
+    )
+
+    # single process
+    if image is not None:
+        ratings, tags = interrogator.interrogate(image)
+        processed_tags = Interrogator.postprocess_tags(
+            tags,
+            *postprocess_opts
+        )
+
+        if unload_model_after_running:
+            interrogator.unload()
+
+        return [
+            ', '.join(processed_tags),
+            ratings,
+            tags,
+            ''
+        ]
+
+    # batch process
+    batch_input_glob = batch_input_glob.strip()
+    batch_output_dir = batch_output_dir.strip()
+    batch_output_filename_format = batch_output_filename_format.strip()
+
+    if batch_input_glob != '':
+        # if there is no glob pattern, insert it automatically
+        if not batch_input_glob.endswith('*'):
+            if not batch_input_glob.endswith('/'):
+                batch_input_glob += '/'
+            batch_input_glob += '*'
+
+        # get root directory of input glob pattern
+        base_dir = batch_input_glob.replace('?', '*')
+        base_dir = base_dir.split('/*').pop(0)
+
+        # check the input directory path
+        if not os.path.isdir(base_dir):
+            return ['', None, None, 'input path is not a directory']
+
+        # this line is moved here because some reason
+        # PIL.Image.registered_extensions() returns only PNG if you call too early
+        supported_extensions = [
+            e
+            for e, f in Image.registered_extensions().items()
+            if f in Image.OPEN
+        ]
+
+        paths = [
+            Path(p)
+            for p in glob(batch_input_glob, recursive=batch_input_recursive)
+            if '.' + p.split('.').pop().lower() in supported_extensions
+        ]
+
+        print(f'found {len(paths)} image(s)')
+
+        for path in paths:
+            try:
+                image = Image.open(path)
+            except UnidentifiedImageError:
+                # just in case, user has mysterious file...
+                print(f'${path} is not supported image type')
+                continue
+
+            # guess the output path
+            base_dir_last = Path(base_dir).parts[-1]
+            base_dir_last_idx = path.parts.index(base_dir_last)
+            output_dir = Path(
+                batch_output_dir) if batch_output_dir else Path(base_dir)
+            output_dir = output_dir.joinpath(
+                *path.parts[base_dir_last_idx + 1:]).parent
+
+            output_dir.mkdir(0o777, True, True)
+
+            # format output filename
+            format_info = format.Info(path, 'txt')
+
+            try:
+                formatted_output_filename = format.pattern.sub(
+                    lambda m: format.format(m, format_info),
+                    batch_output_filename_format
+                )
+            except (TypeError, ValueError) as error:
+                return ['', None, None, str(error)]
+
+            output_path = output_dir.joinpath(
+                formatted_output_filename
+            )
+
+            output = []
+
+            if output_path.is_file():
+                output.append(output_path.read_text())
+
+                if batch_output_action_on_conflict == 'ignore':
+                    print(f'skipping {path}')
+                    continue
+
+            ratings, tags = interrogator.interrogate(image)
+            processed_tags = Interrogator.postprocess_tags(
+                tags,
+                *postprocess_opts
+            )
+
+            # TODO: switch for less print
+            print(
+                f'found {len(processed_tags)} tags out of {len(tags)} from {path}'
+            )
+
+            plain_tags = ', '.join(processed_tags)
+
+            if batch_output_action_on_conflict == 'copy':
+                output = [plain_tags]
+            elif batch_output_action_on_conflict == 'prepend':
+                output.insert(0, plain_tags)
+            else:
+                output.append(plain_tags)
+
+            output_path.write_text(' '.join(output))
+
+            if batch_output_save_json:
+                output_path.with_suffix('.json').write_text(
+                    json.dumps([ratings, tags])
+                )
+
+        print('all done :)')
+
+    if unload_model_after_running:
+        interrogator.unload()
+
+    return ['', None, None, '']
+
+
+def on_ui_tabs():
+    with gr.Blocks(analytics_enabled=False) as tagger_interface:
+        with gr.Row().style(equal_height=False):
+            with gr.Column(variant='panel'):
+
+                # input components
+                with gr.Tabs():
+                    with gr.TabItem(label='Single process'):
+                        image = gr.Image(
+                            label='Source',
+                            source='upload',
+                            interactive=True,
+                            type="pil"
+                        )
+
+                    with gr.TabItem(label='Batch from directory'):
+                        batch_input_glob = utils.preset.component(
+                            gr.Textbox,
+                            label='Input directory',
+                            placeholder='/path/to/images or /path/to/images/**/*'
+                        )
+                        batch_input_recursive = utils.preset.component(
+                            gr.Checkbox,
+                            label='Use recursive with glob pattern'
+                        )
+
+                        batch_output_dir = utils.preset.component(
+                            gr.Textbox,
+                            label='Output directory',
+                            placeholder='Leave blank to save images to the same path.'
+                        )
+
+                        batch_output_filename_format = utils.preset.component(
+                            gr.Textbox,
+                            label='Output filename format',
+                            placeholder='Leave blank to use same filename as original.',
+                            value='[name].[output_extension]'
+                        )
+
+                        import hashlib
+                        with gr.Accordion(
+                            label='Output filename formats',
+                            open=False
+                        ):
+                            gr.Markdown(
+                                value=f'''
+                                ### Related to original file
+                                - `[name]`: Original filename without extension
+                                - `[extension]`: Original extension
+                                - `[hash:<algorithms>]`: Original extension
+                                    Available algorithms: `{', '.join(hashlib.algorithms_available)}`
+
+                                ### Related to output file
+                                - `[output_extension]`: Output extension (has no dot)
+
+                                ## Examples
+                                ### Original filename without extension
+                                `[name].[output_extension]`
+
+                                ### Original file's hash (good for deleting duplication)
+                                `[hash:sha1].[output_extension]`
+                                '''
+                            )
+
+                        batch_output_action_on_conflict = utils.preset.component(
+                            gr.Dropdown,
+                            label='Action on exiting caption',
+                            value='ignore',
+                            choices=[
+                                'ignore',
+                                'copy',
+                                'append',
+                                'prepend'
+                            ]
+                        )
+
+                        batch_output_save_json = utils.preset.component(
+                            gr.Checkbox,
+                            label='Save with JSON'
+                        )
+
+                submit = gr.Button(
+                    value='Interrogate',
+                    variant='primary'
+                )
+
+                info = gr.HTML()
+
+                # preset selector
+                with gr.Row(variant='compact'):
+                    available_presets = utils.preset.list()
+                    selected_preset = gr.Dropdown(
+                        label='Preset',
+                        choices=available_presets,
+                        value=available_presets[0]
+                    )
+
+                    save_preset_button = gr.Button(
+                        value=ui.save_style_symbol
+                    )
+
+                    ui.create_refresh_button(
+                        selected_preset,
+                        lambda: None,
+                        lambda: {'choices': utils.preset.list()},
+                        'refresh_preset'
+                    )
+
+                # option components
+
+                # interrogator selector
+                with gr.Column():
+                    with gr.Row(variant='compact'):
+                        interrogator_names = utils.refresh_interrogators()
+                        interrogator = utils.preset.component(
+                            gr.Dropdown,
+                            label='Interrogator',
+                            choices=interrogator_names,
+                            value=(
+                                None
+                                if len(interrogator_names) < 1 else
+                                interrogator_names[-1]
+                            )
+                        )
+
+                        ui.create_refresh_button(
+                            interrogator,
+                            lambda: None,
+                            lambda: {'choices': utils.refresh_interrogators()},
+                            'refresh_interrogator'
+                        )
+
+                    unload_all_models = gr.Button(
+                        value='Unload all interrogate models'
+                    )
+
+                threshold = utils.preset.component(
+                    gr.Slider,
+                    label='Threshold',
+                    minimum=0,
+                    maximum=1,
+                    value=0.35
+                )
+
+                additional_tags = utils.preset.component(
+                    gr.Textbox,
+                    label='Additional tags (split by comma)',
+                    elem_id='additioanl-tags'
+                )
+
+                exclude_tags = utils.preset.component(
+                    gr.Textbox,
+                    label='Exclude tags (split by comma)',
+                    elem_id='exclude-tags'
+                )
+
+                sort_by_alphabetical_order = utils.preset.component(
+                    gr.Checkbox,
+                    label='Sort by alphabetical order',
+                )
+                add_confident_as_weight = utils.preset.component(
+                    gr.Checkbox,
+                    label='Include confident of tags matches in results'
+                )
+                replace_underscore = utils.preset.component(
+                    gr.Checkbox,
+                    label='Use spaces instead of underscore',
+                    value=True
+                )
+                replace_underscore_excludes = utils.preset.component(
+                    gr.Textbox,
+                    label='Excudes (split by comma)',
+                    # kaomoji from WD 1.4 tagger csv. thanks, Meow-San#5400!
+                    value='0_0, (o)_(o), +_+, +_-, ._., <o>_<o>, <|>_<|>, =_=, >_<, 3_3, 6_9, >_o, @_@, ^_^, o_o, u_u, x_x, |_|, ||_||'
+                )
+                escape_tag = utils.preset.component(
+                    gr.Checkbox,
+                    label='Escape brackets',
+                )
+
+                unload_model_after_running = utils.preset.component(
+                    gr.Checkbox,
+                    label='Unload model after running',
+                )
+
+            # output components
+            with gr.Column(variant='panel'):
+                tags = gr.Textbox(
+                    label='Tags',
+                    placeholder='Found tags',
+                    interactive=False
+                )
+
+                with gr.Row():
+                    parameters_copypaste.bind_buttons(
+                        parameters_copypaste.create_buttons(
+                            ["txt2img", "img2img"],
+                        ),
+                        None,
+                        tags
+                    )
+
+                rating_confidents = gr.Label(
+                    label='Rating confidents',
+                    elem_id='rating-confidents'
+                )
+                tag_confidents = gr.Label(
+                    label='Tag confidents',
+                    elem_id='tag-confidents'
+                )
+
+        # register events
+        selected_preset.change(
+            fn=utils.preset.apply,
+            inputs=[selected_preset],
+            outputs=[*utils.preset.components, info]
+        )
+
+        save_preset_button.click(
+            fn=utils.preset.save,
+            inputs=[selected_preset, *utils.preset.components],  # values only
+            outputs=[info]
+        )
+
+        unload_all_models.click(
+            fn=unload_interrogators,
+            outputs=[info]
+        )
+
+        for func in [image.change, submit.click]:
+            func(
+                fn=wrap_gradio_gpu_call(on_interrogate),
+                inputs=[
+                    # single process
+                    image,
+
+                    # batch process
+                    batch_input_glob,
+                    batch_input_recursive,
+                    batch_output_dir,
+                    batch_output_filename_format,
+                    batch_output_action_on_conflict,
+                    batch_output_save_json,
+
+                    # options
+                    interrogator,
+                    threshold,
+                    additional_tags,
+                    exclude_tags,
+                    sort_by_alphabetical_order,
+                    add_confident_as_weight,
+                    replace_underscore,
+                    replace_underscore_excludes,
+                    escape_tag,
+
+                    unload_model_after_running
+                ],
+                outputs=[
+                    tags,
+                    rating_confidents,
+                    tag_confidents,
+
+                    # contains execution time, memory usage and other stuffs...
+                    # generated from modules.ui.wrap_gradio_call
+                    info
+                ]
+            )
+
+    return [(tagger_interface, "Tagger", "tagger")]

tagger/utils.py

@@ -0,0 +1,103 @@
+import os
+
+from typing import List, Dict
+from pathlib import Path
+
+from modules import shared, scripts
+from preload import default_ddp_path, default_onnx_path
+from tagger.preset import Preset
+from tagger.interrogator import Interrogator, DeepDanbooruInterrogator, WaifuDiffusionInterrogator
+
+preset = Preset(Path(scripts.basedir(), 'presets'))
+
+interrogators: Dict[str, Interrogator] = {}
+
+
+def refresh_interrogators() -> List[str]:
+    global interrogators
+    interrogators = {
+        'wd14-vit-v2': WaifuDiffusionInterrogator(
+            'wd14-vit-v2',
+            repo_id='SmilingWolf/wd-v1-4-vit-tagger-v2',
+            revision='v2.0'
+        ),
+        'wd14-convnext-v2': WaifuDiffusionInterrogator(
+            'wd14-convnext-v2',
+            repo_id='SmilingWolf/wd-v1-4-convnext-tagger-v2',
+            revision='v2.0'
+        ),
+        'wd14-swinv2-v2': WaifuDiffusionInterrogator(
+            'wd14-swinv2-v2',
+            repo_id='SmilingWolf/wd-v1-4-swinv2-tagger-v2',
+            revision='v2.0'
+        ),
+        'wd14-vit-v2-git': WaifuDiffusionInterrogator(
+            'wd14-vit-v2-git',
+            repo_id='SmilingWolf/wd-v1-4-vit-tagger-v2'
+        ),
+        'wd14-convnext-v2-git': WaifuDiffusionInterrogator(
+            'wd14-convnext-v2-git',
+            repo_id='SmilingWolf/wd-v1-4-convnext-tagger-v2'
+        ),
+        'wd14-swinv2-v2-git': WaifuDiffusionInterrogator(
+            'wd14-swinv2-v2-git',
+            repo_id='SmilingWolf/wd-v1-4-swinv2-tagger-v2'
+        ),
+        'wd14-vit': WaifuDiffusionInterrogator(
+            'wd14-vit',
+            repo_id='SmilingWolf/wd-v1-4-vit-tagger'),
+        'wd14-convnext': WaifuDiffusionInterrogator(
+            'wd14-convnext',
+            repo_id='SmilingWolf/wd-v1-4-convnext-tagger'
+        ),
+        #'Z3D-E621-convnext': WaifuDiffusionInterrogator(
+        #    'Z3D-E621-convnext',
+        #    model_path=r'SmilingWolf/wd-v1-4-convnext-tagger',
+        #    tags_path=r''
+        #),
+    }
+
+    # load deepdanbooru project
+    os.makedirs(
+        getattr(shared.cmd_opts, 'deepdanbooru_projects_path', default_ddp_path),
+        exist_ok=True
+    )
+    os.makedirs(
+        getattr(shared.cmd_opts, 'onnxtagger_path', default_onnx_path),
+        exist_ok=True
+    )
+
+    for path in os.scandir(shared.cmd_opts.deepdanbooru_projects_path):
+        if not path.is_dir():
+            continue
+
+        if not Path(path, 'project.json').is_file():
+            continue
+
+        interrogators[path.name] = DeepDanbooruInterrogator(path.name, path)
+    #scan for onnx models as well
+    for path in os.scandir(shared.cmd_opts.onnxtagger_path):
+        if not path.is_dir():
+            continue
+        
+        #if no file with the extension .onnx is found, skip. If there is more than one file with that name, warn. Else put it in model_path
+        onnx_files = [x for x in os.scandir(path) if x.name.endswith('.onnx')]
+        if len(onnx_files) == 0:
+            print(f"Warning: {path} has no model, skipping")
+            continue
+        elif len(onnx_files) > 1:
+            print(f"Warning: {path} has multiple models, skipping")
+            continue
+        model_path = Path(path, onnx_files[0].name)
+
+        if not Path(path, 'tags-selected.csv').is_file():
+            print(f"Warning: {path} has a model but no tags-selected.csv file, skipping")
+            continue
+
+        interrogators[path.name] = WaifuDiffusionInterrogator(path.name,model_path=model_path, tags_path=Path(path, 'tags-selected.csv'))
+
+    return sorted(interrogators.keys())
+
+
+def split_str(s: str, separator=',') -> List[str]:
+    return [x.strip() for x in s.split(separator) if x]