SmilingWolf commited on
Commit
db51843
1 Parent(s): 4015d2a

Add experimental model

Browse files
Files changed (1) hide show
  1. app.py +10 -1
app.py CHANGED
@@ -15,6 +15,8 @@ Demo for the WaifuDiffusion tagger models
15
  Example image by [ほし☆☆☆](https://www.pixiv.net/en/users/43565085)
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"
@@ -29,6 +31,9 @@ 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
  # Files to download from the repos
33
  MODEL_FILENAME = "model.onnx"
34
  LABEL_FILENAME = "selected_tags.csv"
@@ -62,7 +67,6 @@ def parse_args() -> argparse.Namespace:
62
  parser.add_argument("--score-slider-step", type=float, default=0.05)
63
  parser.add_argument("--score-general-threshold", type=float, default=0.35)
64
  parser.add_argument("--score-character-threshold", type=float, default=0.85)
65
- parser.add_argument("--share", action="store_true")
66
  return parser.parse_args()
67
 
68
 
@@ -102,10 +106,12 @@ class Predictor:
102
  csv_path = huggingface_hub.hf_hub_download(
103
  model_repo,
104
  LABEL_FILENAME,
 
105
  )
106
  model_path = huggingface_hub.hf_hub_download(
107
  model_repo,
108
  MODEL_FILENAME,
 
109
  )
110
  return csv_path, model_path
111
 
@@ -229,11 +235,14 @@ def main():
229
  VIT_MODEL_DSV3_REPO,
230
  VIT_LARGE_MODEL_DSV3_REPO,
231
  EVA02_LARGE_MODEL_DSV3_REPO,
 
232
  MOAT_MODEL_DSV2_REPO,
233
  SWIN_MODEL_DSV2_REPO,
234
  CONV_MODEL_DSV2_REPO,
235
  CONV2_MODEL_DSV2_REPO,
236
  VIT_MODEL_DSV2_REPO,
 
 
237
  ]
238
 
239
  with gr.Blocks(title=TITLE) as demo:
 
15
  Example image by [ほし☆☆☆](https://www.pixiv.net/en/users/43565085)
16
  """
17
 
18
+ HF_TOKEN = os.environ.get("HF_TOKEN", "")
19
+
20
  # Dataset v3 series of models:
21
  SWINV2_MODEL_DSV3_REPO = "SmilingWolf/wd-swinv2-tagger-v3"
22
  CONV_MODEL_DSV3_REPO = "SmilingWolf/wd-convnext-tagger-v3"
 
31
  CONV2_MODEL_DSV2_REPO = "SmilingWolf/wd-v1-4-convnextv2-tagger-v2"
32
  VIT_MODEL_DSV2_REPO = "SmilingWolf/wd-v1-4-vit-tagger-v2"
33
 
34
+ # IdolSankaku series of models:
35
+ EVA02_LARGE_MODEL_IS_DSV1_REPO = "deepghs/idolsankaku-eva02-large-tagger-v1"
36
+
37
  # Files to download from the repos
38
  MODEL_FILENAME = "model.onnx"
39
  LABEL_FILENAME = "selected_tags.csv"
 
67
  parser.add_argument("--score-slider-step", type=float, default=0.05)
68
  parser.add_argument("--score-general-threshold", type=float, default=0.35)
69
  parser.add_argument("--score-character-threshold", type=float, default=0.85)
 
70
  return parser.parse_args()
71
 
72
 
 
106
  csv_path = huggingface_hub.hf_hub_download(
107
  model_repo,
108
  LABEL_FILENAME,
109
+ use_auth_token=HF_TOKEN,
110
  )
111
  model_path = huggingface_hub.hf_hub_download(
112
  model_repo,
113
  MODEL_FILENAME,
114
+ use_auth_token=HF_TOKEN,
115
  )
116
  return csv_path, model_path
117
 
 
235
  VIT_MODEL_DSV3_REPO,
236
  VIT_LARGE_MODEL_DSV3_REPO,
237
  EVA02_LARGE_MODEL_DSV3_REPO,
238
+ # ---
239
  MOAT_MODEL_DSV2_REPO,
240
  SWIN_MODEL_DSV2_REPO,
241
  CONV_MODEL_DSV2_REPO,
242
  CONV2_MODEL_DSV2_REPO,
243
  VIT_MODEL_DSV2_REPO,
244
+ # ---
245
+ EVA02_LARGE_MODEL_IS_DSV1_REPO,
246
  ]
247
 
248
  with gr.Blocks(title=TITLE) as demo: