dev(narugo): add vit large model into list

deepghs · Jul 27, 2024 · f886d62 · f886d62
1 parent 1321b9d
commit f886d62
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/imgutils/tagging/wd14.py b/imgutils/tagging/wd14.py
@@ -26,6 +26,7 @@
 CONV_V3_MODEL_REPO = 'SmilingWolf/wd-convnext-tagger-v3'
 SWIN_V3_MODEL_REPO = 'SmilingWolf/wd-swinv2-tagger-v3'
 VIT_V3_MODEL_REPO = 'SmilingWolf/wd-vit-tagger-v3'
+VIT_LARGE_MODEL_REPO = 'SmilingWolf/wd-vit-large-tagger-v3'
 MODEL_FILENAME = "model.onnx"
 LABEL_FILENAME = "selected_tags.csv"
 
@@ -41,6 +42,8 @@
     "SwinV2_v3": SWIN_V3_MODEL_REPO,
     "ConvNext_v3": CONV_V3_MODEL_REPO,
     "ViT_v3": VIT_V3_MODEL_REPO,
+
+    "Vit_Large": VIT_LARGE_MODEL_REPO,
 }
 _DEFAULT_MODEL_NAME = 'SwinV2_v3'