Commit f886d62d authored by narugo1992's avatar narugo1992
Browse files

dev(narugo): add vit large model into list

parent 1321b9d2

imgutils/tagging/wd14.py

+3 −0
Original line number Diff line number Diff line
@@ -26,6 +26,7 @@ MOAT_MODEL_REPO = "SmilingWolf/wd-v1-4-moat-tagger-v2" 
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 @@ MODEL_NAMES = { 
    "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'