Spaces:
Running
Running
Update app.py
Browse files
app.py
CHANGED
@@ -13,25 +13,25 @@ DESCRIPTION = """
|
|
13 |
Demo for the WaifuDiffusion tagger models
|
14 |
"""
|
15 |
|
16 |
-
|
17 |
|
18 |
# Dataset v3 series of models:
|
19 |
-
SWINV2_MODEL_DSV3_REPO = "SmilingWolf/wd-swinv2-tagger-v3"
|
20 |
-
CONV_MODEL_DSV3_REPO = "SmilingWolf/wd-convnext-tagger-v3"
|
21 |
VIT_MODEL_DSV3_REPO = "ura23/wd-vit-tagger-v3"
|
22 |
-
VIT_LARGE_MODEL_DSV3_REPO = "SmilingWolf/wd-vit-large-tagger-v3"
|
23 |
-
EVA02_LARGE_MODEL_DSV3_REPO = "SmilingWolf/wd-eva02-large-tagger-v3"
|
24 |
|
25 |
# Dataset v2 series of models:
|
26 |
-
MOAT_MODEL_DSV2_REPO = "SmilingWolf/wd-v1-4-moat-tagger-v2"
|
27 |
-
SWIN_MODEL_DSV2_REPO = "SmilingWolf/wd-v1-4-swinv2-tagger-v2"
|
28 |
-
CONV_MODEL_DSV2_REPO = "SmilingWolf/wd-v1-4-convnext-tagger-v2"
|
29 |
-
CONV2_MODEL_DSV2_REPO = "SmilingWolf/wd-v1-4-convnextv2-tagger-v2"
|
30 |
-
VIT_MODEL_DSV2_REPO = "SmilingWolf/wd-v1-4-vit-tagger-v2"
|
31 |
|
32 |
# IdolSankaku series of models:
|
33 |
-
EVA02_LARGE_MODEL_IS_DSV1_REPO = "deepghs/idolsankaku-eva02-large-tagger-v1"
|
34 |
-
SWINV2_MODEL_IS_DSV1_REPO = "deepghs/idolsankaku-swinv2-tagger-v1"
|
35 |
|
36 |
# Files to download from the repos
|
37 |
MODEL_FILENAME = "model.onnx"
|
@@ -56,8 +56,8 @@ class Predictor:
|
|
56 |
self.last_loaded_repo = None
|
57 |
|
58 |
def download_model(self, model_repo):
|
59 |
-
csv_path = huggingface_hub.hf_hub_download(model_repo, LABEL_FILENAME, use_auth_token=
|
60 |
-
model_path = huggingface_hub.hf_hub_download(model_repo, MODEL_FILENAME, use_auth_token=
|
61 |
return csv_path, model_path
|
62 |
|
63 |
def load_model(self, model_repo):
|
@@ -123,20 +123,20 @@ def main():
|
|
123 |
predictor = Predictor()
|
124 |
|
125 |
model_repos = [
|
126 |
-
SWINV2_MODEL_DSV3_REPO,
|
127 |
-
CONV_MODEL_DSV3_REPO,
|
128 |
VIT_MODEL_DSV3_REPO,
|
129 |
-
VIT_LARGE_MODEL_DSV3_REPO,
|
130 |
-
EVA02_LARGE_MODEL_DSV3_REPO,
|
131 |
# ---
|
132 |
-
MOAT_MODEL_DSV2_REPO,
|
133 |
-
SWIN_MODEL_DSV2_REPO,
|
134 |
-
CONV_MODEL_DSV2_REPO,
|
135 |
-
CONV2_MODEL_DSV2_REPO,
|
136 |
-
VIT_MODEL_DSV2_REPO,
|
137 |
# ---
|
138 |
-
SWINV2_MODEL_IS_DSV1_REPO,
|
139 |
-
EVA02_LARGE_MODEL_IS_DSV1_REPO,
|
140 |
]
|
141 |
|
142 |
predefined_tags = ["loli",
|
|
|
13 |
Demo for the WaifuDiffusion tagger models
|
14 |
"""
|
15 |
|
16 |
+
HF_TOKEN = os.environ.get("HF_TOKEN", "")
|
17 |
|
18 |
# Dataset v3 series of models:
|
19 |
+
# SWINV2_MODEL_DSV3_REPO = "SmilingWolf/wd-swinv2-tagger-v3"
|
20 |
+
# CONV_MODEL_DSV3_REPO = "SmilingWolf/wd-convnext-tagger-v3"
|
21 |
VIT_MODEL_DSV3_REPO = "ura23/wd-vit-tagger-v3"
|
22 |
+
# VIT_LARGE_MODEL_DSV3_REPO = "SmilingWolf/wd-vit-large-tagger-v3"
|
23 |
+
# EVA02_LARGE_MODEL_DSV3_REPO = "SmilingWolf/wd-eva02-large-tagger-v3"
|
24 |
|
25 |
# Dataset v2 series of models:
|
26 |
+
# MOAT_MODEL_DSV2_REPO = "SmilingWolf/wd-v1-4-moat-tagger-v2"
|
27 |
+
# SWIN_MODEL_DSV2_REPO = "SmilingWolf/wd-v1-4-swinv2-tagger-v2"
|
28 |
+
# CONV_MODEL_DSV2_REPO = "SmilingWolf/wd-v1-4-convnext-tagger-v2"
|
29 |
+
# CONV2_MODEL_DSV2_REPO = "SmilingWolf/wd-v1-4-convnextv2-tagger-v2"
|
30 |
+
# VIT_MODEL_DSV2_REPO = "SmilingWolf/wd-v1-4-vit-tagger-v2"
|
31 |
|
32 |
# IdolSankaku series of models:
|
33 |
+
# EVA02_LARGE_MODEL_IS_DSV1_REPO = "deepghs/idolsankaku-eva02-large-tagger-v1"
|
34 |
+
# SWINV2_MODEL_IS_DSV1_REPO = "deepghs/idolsankaku-swinv2-tagger-v1"
|
35 |
|
36 |
# Files to download from the repos
|
37 |
MODEL_FILENAME = "model.onnx"
|
|
|
56 |
self.last_loaded_repo = None
|
57 |
|
58 |
def download_model(self, model_repo):
|
59 |
+
csv_path = huggingface_hub.hf_hub_download(model_repo, LABEL_FILENAME, use_auth_token=HF_TOKEN)
|
60 |
+
model_path = huggingface_hub.hf_hub_download(model_repo, MODEL_FILENAME, use_auth_token=HF_TOKEN)
|
61 |
return csv_path, model_path
|
62 |
|
63 |
def load_model(self, model_repo):
|
|
|
123 |
predictor = Predictor()
|
124 |
|
125 |
model_repos = [
|
126 |
+
# SWINV2_MODEL_DSV3_REPO,
|
127 |
+
# CONV_MODEL_DSV3_REPO,
|
128 |
VIT_MODEL_DSV3_REPO,
|
129 |
+
# VIT_LARGE_MODEL_DSV3_REPO,
|
130 |
+
# EVA02_LARGE_MODEL_DSV3_REPO,
|
131 |
# ---
|
132 |
+
# MOAT_MODEL_DSV2_REPO,
|
133 |
+
# SWIN_MODEL_DSV2_REPO,
|
134 |
+
# CONV_MODEL_DSV2_REPO,
|
135 |
+
# CONV2_MODEL_DSV2_REPO,
|
136 |
+
# VIT_MODEL_DSV2_REPO,
|
137 |
# ---
|
138 |
+
# SWINV2_MODEL_IS_DSV1_REPO,
|
139 |
+
# EVA02_LARGE_MODEL_IS_DSV1_REPO,
|
140 |
]
|
141 |
|
142 |
predefined_tags = ["loli",
|