Merge pull request #19 from corkborg/z3d-e621

Support z3d-e621 models.

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

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',