ura23 commited on
Commit
836fab2
·
verified ·
1 Parent(s): 51fba29

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +25 -25
app.py CHANGED
@@ -13,25 +13,25 @@ DESCRIPTION = """
13
  Demo for the WaifuDiffusion tagger models
14
  """
15
 
16
- HF_TOKEN1 = os.environ.get("HF_TOKEN1", "")
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=HF_TOKEN1)
60
- model_path = huggingface_hub.hf_hub_download(model_repo, MODEL_FILENAME, use_auth_token=HF_TOKEN1)
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",