From 48c5360994085e00f0de86c86993ea7883109292 Mon Sep 17 00:00:00 2001
From: corkborg <62742115+corkborg@users.noreply.github.com>
Date: Wed, 28 Aug 2024 17:05:50 +0900
Subject: [PATCH] Support z3d-e621 models.

---
 README.md               |  4 ++++
 tagger/interrogators.py | 10 ++++++++++
 2 files changed, 14 insertions(+)

diff --git a/README.md b/README.md
index dec6666..4057293 100644
--- a/README.md
+++ b/README.md
@@ -67,6 +67,10 @@ 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
 
+# 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
+
 # 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
diff --git a/tagger/interrogators.py b/tagger/interrogators.py
index 13c2859..01c18ce 100644
--- a/tagger/interrogators.py
+++ b/tagger/interrogators.py
@@ -53,6 +53,16 @@ interrogators: Dict[str, Interrogator] = {
         'WD EVA02-Large Tagger v3',
         repo_id='SmilingWolf/wd-eva02-large-tagger-v3'
     ),
+    'z3d-e621-convnext-toynya': WaifuDiffusionInterrogator(
+        'Z3D-E621-Convnext(toynya)',
+        repo_id='toynya/Z3D-E621-Convnext',
+        tags_path='tags-selected.csv'
+    ),
+    'z3d-e621-convnext-silveroxides': WaifuDiffusionInterrogator(
+        'Z3D-E621-Convnext(silveroxides)',
+        repo_id='silveroxides/Z3D-E621-Convnext',
+        tags_path='tags-selected.csv'
+    ),
     'mld-caformer.dec-5-97527': MLDanbooruInterrogator(
         'ML-Danbooru Caformer dec-5-97527',
         repo_id='deepghs/ml-danbooru-onnx',