Rename modules/app (2).py to modules/wdtagger.py

modules/{app (2).py → wdtagger.py}

@@ -8,13 +8,6 @@ import onnxruntime as rt
 import pandas as pd
 from PIL import Image
 
-TITLE = "WaifuDiffusion Tagger"
-DESCRIPTION = """
-Demo for the WaifuDiffusion tagger models
-
-Example image by [ほし☆☆☆](https://www.pixiv.net/en/users/43565085)
-"""
-
 # Dataset v3 series of models:
 SWINV2_MODEL_DSV3_REPO = "SmilingWolf/wd-swinv2-tagger-v3"
 CONV_MODEL_DSV3_REPO = "SmilingWolf/wd-convnext-tagger-v3"
@@ -218,122 +211,21 @@ class Predictor:
         return sorted_general_strings, rating, character_res, general_res
 
 
-def main():
-    args = parse_args()
-
-    predictor = Predictor()
-
-    dropdown_list = [
-        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,
-    ]
-
-    with gr.Blocks(title=TITLE) as demo:
-        with gr.Column():
-            gr.Markdown(
-                value=f"<h1 style='text-align: center; margin-bottom: 1rem'>{TITLE}</h1>"
-            )
-            gr.Markdown(value=DESCRIPTION)
-            with gr.Row():
-                with gr.Column(variant="panel"):
-                    image = gr.Image(type="pil", image_mode="RGBA", label="Input")
-                    model_repo = gr.Dropdown(
-                        dropdown_list,
-                        value=SWINV2_MODEL_DSV3_REPO,
-                        label="Model",
-                    )
-                    with gr.Row():
-                        general_thresh = gr.Slider(
-                            0,
-                            1,
-                            step=args.score_slider_step,
-                            value=args.score_general_threshold,
-                            label="General Tags Threshold",
-                            scale=3,
-                        )
-                        general_mcut_enabled = gr.Checkbox(
-                            value=False,
-                            label="Use MCut threshold",
-                            scale=1,
-                        )
-                    with gr.Row():
-                        character_thresh = gr.Slider(
-                            0,
-                            1,
-                            step=args.score_slider_step,
-                            value=args.score_character_threshold,
-                            label="Character Tags Threshold",
-                            scale=3,
-                        )
-                        character_mcut_enabled = gr.Checkbox(
-                            value=False,
-                            label="Use MCut threshold",
-                            scale=1,
-                        )
-                    with gr.Row():
-                        clear = gr.ClearButton(
-                            components=[
-                                image,
-                                model_repo,
-                                general_thresh,
-                                general_mcut_enabled,
-                                character_thresh,
-                                character_mcut_enabled,
-                            ],
-                            variant="secondary",
-                            size="lg",
-                        )
-                        submit = gr.Button(value="Submit", variant="primary", size="lg")
-                with gr.Column(variant="panel"):
-                    sorted_general_strings = gr.Textbox(label="Output (string)")
-                    rating = gr.Label(label="Rating")
-                    character_res = gr.Label(label="Output (characters)")
-                    general_res = gr.Label(label="Output (tags)")
-                    clear.add(
-                        [
-                            sorted_general_strings,
-                            rating,
-                            character_res,
-                            general_res,
-                        ]
-                    )
-
-        submit.click(
-            predictor.predict,
-            inputs=[
-                image,
-                model_repo,
-                general_thresh,
-                general_mcut_enabled,
-                character_thresh,
-                character_mcut_enabled,
-            ],
-            outputs=[sorted_general_strings, rating, character_res, general_res],
-        )
+args = parse_args()
 
-        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()
+predictor = Predictor()
 
+dropdown_list = [
+    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,
+]
 
-if __name__ == "__main__":
-    main()
+