Update app.py

app.py

@@ -15,25 +15,25 @@ Demo for the WaifuDiffusion tagger models
 Example image by [ほし☆☆☆](https://www.pixiv.net/en/users/43565085)
 """
 
-HF_TOKEN = os.environ.get("HF_TOKEN", "hf_your_actual_token")
+HF_TOKEN = os.environ.get("HF_TOKEN", "")
 
 # Dataset v3 series of models:
-SWINV2_MODEL_DSV3_REPO = "SmilingWolf/wd-swinv2-tagger-v3"
-CONV_MODEL_DSV3_REPO = "SmilingWolf/wd-convnext-tagger-v3"
-VIT_MODEL_DSV3_REPO = "SmilingWolf/wd-vit-tagger-v3"
-VIT_LARGE_MODEL_DSV3_REPO = "SmilingWolf/wd-vit-large-tagger-v3"
+# SWINV2_MODEL_DSV3_REPO = "SmilingWolf/wd-swinv2-tagger-v3"
+# CONV_MODEL_DSV3_REPO = "SmilingWolf/wd-convnext-tagger-v3"
+# VIT_MODEL_DSV3_REPO = "SmilingWolf/wd-vit-tagger-v3"
+# VIT_LARGE_MODEL_DSV3_REPO = "SmilingWolf/wd-vit-large-tagger-v3"
 EVA02_LARGE_MODEL_DSV3_REPO = "SmilingWolf/wd-eva02-large-tagger-v3"
 
 # Dataset v2 series of models:
-MOAT_MODEL_DSV2_REPO = "SmilingWolf/wd-v1-4-moat-tagger-v2"
-SWIN_MODEL_DSV2_REPO = "SmilingWolf/wd-v1-4-swinv2-tagger-v2"
-CONV_MODEL_DSV2_REPO = "SmilingWolf/wd-v1-4-convnext-tagger-v2"
-CONV2_MODEL_DSV2_REPO = "SmilingWolf/wd-v1-4-convnextv2-tagger-v2"
-VIT_MODEL_DSV2_REPO = "SmilingWolf/wd-v1-4-vit-tagger-v2"
+# MOAT_MODEL_DSV2_REPO = "SmilingWolf/wd-v1-4-moat-tagger-v2"
+# SWIN_MODEL_DSV2_REPO = "SmilingWolf/wd-v1-4-swinv2-tagger-v2"
+# CONV_MODEL_DSV2_REPO = "SmilingWolf/wd-v1-4-convnext-tagger-v2"
+# CONV2_MODEL_DSV2_REPO = "SmilingWolf/wd-v1-4-convnextv2-tagger-v2"
+# VIT_MODEL_DSV2_REPO = "SmilingWolf/wd-v1-4-vit-tagger-v2"
 
 # IdolSankaku series of models:
-EVA02_LARGE_MODEL_IS_DSV1_REPO = "deepghs/idolsankaku-eva02-large-tagger-v1"
-SWINV2_MODEL_IS_DSV1_REPO = "deepghs/idolsankaku-swinv2-tagger-v1"
+# EVA02_LARGE_MODEL_IS_DSV1_REPO = "deepghs/idolsankaku-eva02-large-tagger-v1"
+# SWINV2_MODEL_IS_DSV1_REPO = "deepghs/idolsankaku-swinv2-tagger-v1"
 
 # Files to download from the repos
 MODEL_FILENAME = "model.onnx"
@@ -66,8 +66,8 @@ kaomojis = [
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser()
     parser.add_argument("--score-slider-step", type=float, default=0.05)
-    parser.add_argument("--score-general-threshold", type=float, default=0.35)
-    parser.add_argument("--score-character-threshold", type=float, default=0.85)
+    parser.add_argument("--score-general-threshold", type=float, default=0.1)
+    parser.add_argument("--score-character-threshold", type=float, default=1)
     return parser.parse_args()
 
 
@@ -231,20 +231,20 @@ def main():
     predictor = Predictor()
 
     dropdown_list = [
-        SWINV2_MODEL_DSV3_REPO,
-        CONV_MODEL_DSV3_REPO,
-        VIT_MODEL_DSV3_REPO,
-        VIT_LARGE_MODEL_DSV3_REPO,
+        # SWINV2_MODEL_DSV3_REPO,
+        # CONV_MODEL_DSV3_REPO,
+        # VIT_MODEL_DSV3_REPO,
+        # VIT_LARGE_MODEL_DSV3_REPO,
         EVA02_LARGE_MODEL_DSV3_REPO,
-        # ---
-        MOAT_MODEL_DSV2_REPO,
-        SWIN_MODEL_DSV2_REPO,
-        CONV_MODEL_DSV2_REPO,
-        CONV2_MODEL_DSV2_REPO,
-        VIT_MODEL_DSV2_REPO,
-        # ---
-        SWINV2_MODEL_IS_DSV1_REPO,
-        EVA02_LARGE_MODEL_IS_DSV1_REPO,
+        # # ---
+        # MOAT_MODEL_DSV2_REPO,
+        # SWIN_MODEL_DSV2_REPO,
+        # CONV_MODEL_DSV2_REPO,
+        # CONV2_MODEL_DSV2_REPO,
+        # VIT_MODEL_DSV2_REPO,
+        # # ---
+        # SWINV2_MODEL_IS_DSV1_REPO,
+        # EVA02_LARGE_MODEL_IS_DSV1_REPO,
     ]
 
     with gr.Blocks(title=TITLE) as demo:
@@ -258,7 +258,7 @@ def main():
                     image = gr.Image(type="pil", image_mode="RGBA", label="Input")
                     model_repo = gr.Dropdown(
                         dropdown_list,
-                        value=SWINV2_MODEL_DSV3_REPO,
+                        value=EVA02_LARGE_MODEL_DSV3_REPO,
                         label="Model",
                     )
                     with gr.Row():
@@ -330,17 +330,17 @@ def main():
             outputs=[sorted_general_strings, rating, character_res, general_res],
         )
 
-        gr.Examples(
-            [["power.jpg", SWINV2_MODEL_DSV3_REPO, 0.35, False, 0.85, False]],
-            inputs=[
-                image,
-                model_repo,
-                general_thresh,
-                general_mcut_enabled,
-                character_thresh,
-                character_mcut_enabled,
-            ],
-        )
+        # gr.Examples(
+        #     [["power.jpg", SWINV2_MODEL_DSV3_REPO, 0.35, False, 0.85, False]],
+        #     inputs=[
+        #         image,
+        #         model_repo,
+        #         general_thresh,
+        #         general_mcut_enabled,
+        #         character_thresh,
+        #         character_mcut_enabled,
+        #     ],
+        # )
 
     demo.queue(max_size=10)
     demo.launch()