From f886d62d05e6fdabc19bb82e73328d974305f19b Mon Sep 17 00:00:00 2001 From: narugo1992 Date: Sun, 28 Jul 2024 05:37:54 +0800 Subject: [PATCH] dev(narugo): add vit large model into list --- imgutils/tagging/wd14.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/imgutils/tagging/wd14.py b/imgutils/tagging/wd14.py index 35c3aac5fb..fe1696dafa 100644 --- 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'