Update app.py

app.py

@@ -16,7 +16,22 @@ Demo for the WaifuDiffusion tagger models
 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 = "ura23/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"
+
+# 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"
 
 # Files to download from the repos
 MODEL_FILENAME = "model.onnx"
@@ -108,64 +123,214 @@ def main():
     predictor = Predictor()
 
     model_repos = [
+        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,
     ]
 
-    predefined_tags = ["2025",
+    predefined_tags = ["loli",
+                       "oppai_loli",
                        "2024",
-                       "2020",
                        "2023",
-                       "2025"]
+                       "2025",
+                       "genderswap",
+                       "genderswap_(otm)",
+                       "genderswap_(otf)",
+                       "genderswap_(mtf)",
+                       "genderswap_(ftm)",
+                       "respirator",
+                       "head-mounted_display",
+                       "2022",
+                       "muscular_female",
+                       "muscular",
+                       "abs",
+                       "2021",
+                       "peeing",
+                       "pee",
+                       "round_eyewear",
+                       "yellow-framed_eyewear",
+                       "hetero",
+                       "vaginal",
+                       "straddling",
+                       "girl_on_top",
+                       "male_pubic_hair",
+                       "cowgirl_position",
+                       "happy_sex",
+                       "vibrator_under_panties",
+                       "vibrator_in_thighhighs",
+                       "anal_beads",
+                       "butt_plug",
+                       "sex_toy",
+                       "anal",
+                       "object_insertion",
+                       "dildo",
+                       "anal_object_insertion",
+                       "vaginal_object_insertion",
+                       "semi-rimless_eyewear",
+                       "red-framed_eyewear",
+                       "under-rim_eyewear",
+                       "3d_background",
+                       "sample_watermark",
+                       "onee-shota",
+                       "incest",
+                       "furry",
+                       "can",
+                       "drinking_can",
+                       "holding_can",
+                       "twitter_strip_game_(meme)",
+                       "like_and_retweet",
+                       "furry_female",
+                       "realistic",
+                       "egg_vibrator",
+                       "tongue_piercing",
+                       "handheld_game_console",
+                       "game_controller",
+                       "nintendo_switch",
+                       "talking",
+                       "swastika",
+                       "character_name",
+                       "vibrator",
+                       "black-framed_eyewear",
+                       "heterochromia",
+                       "chibi",
+                       "mini_person",
+                       "controller",
+                       "remote_control_vibrator",
+                       "vibrator_under_clothes",
+                       "thank_you",
+                       "vibrator_cord",
+                       "shota",
+                       "cropped_legs",
+                       "cropped_torso",
+                       "traditional_media",
+                       "color_guide",
+                       "photorealistic",
+                       "male_focus",
+                       "black_babydoll",
+                       "signature",
+                       "web_address",
+                       "censored_nipples",
+                       "rhodes_island_logo_(arknights)",
+                       "gothic_lolita",
+                       "glasses",
+                       "reference_inset",
+                       "twitter_logo",
+                       "mother_and_daughter",
+                       "holding_controller",
+                       "holding_game_controller",
+                       "baby",
+                       "heart_censor",
+                       "pixiv_username",
+                       "korean_text",
+                       "pixiv_logo",
+                       "greyscale_with_colored_background",
+                       "water_bottle",
+                       "body_writing",
+                       "used_condom",
+                       "multiple_condoms",
+                       "condom_belt", 
+                       "holding_phone",
+                       "multiple_views",
+                       "phone",
+                       "cellphone",
+                       "zoom_layer",
+                       "smartphone",
+                       "lolita_hairband",
+                       "lactation",
+                       "otoko_no_ko",
+                       "minigirl",
+                       "babydoll",
+                       "domino_mask",
+                       "pixiv_id",
+                       "qr_code",
+                       "monochrome",
+                       "trick_or_treat",
+                       "happy_birthday",
+                       "lolita_fashion",
+                       "arrow_(symbol)",
+                       "happy_new_year", 
+                       "dated", 
+                       "thought_bubble",
+                       "greyscale",
+                       "speech_bubble",
+                       "mask",
+                       "comic",
+                       "bottle",
+                       "holding_bottle",
+                       "milk",
+                       "milk_bottle",
+                       "english_text",
+                       "copyright_name",
+                       "twitter_username",
+                       "fanbox_username",
+                       "patreon_username",
+                       "patreon_logo",
+                       "cover",
+                       "weibo_logo",
+                       "weibo_username",
+                       "signature",
+                       "content_rating",
+                       "cover_page", 
+                       "doujin_cover", 
+                       "sex",
+                       "artist_name",
+                       "watermark",
+                       "censored",
+                       "bar_censor",
+                       "blank_censor",
+                       "blur_censor",
+                       "light_censor",
+                       "mosaic_censoring"]
 
     with gr.Blocks(title=TITLE) as demo:
-    gr.Markdown(f"<h1 style='text-align: center;'>{TITLE}</h1>")
-    gr.Markdown(DESCRIPTION)
-
-    with gr.Blocks(title=TITLE) as demo:
-    gr.Markdown(f"<h1 style='text-align: center;'>{TITLE}</h1>")
-    gr.Markdown(DESCRIPTION)
-
-    with gr.Row():
-        with gr.Column():
-            submit = gr.Button(
-                value="Process Images", variant="primary"
-            )
-
-            image_files = gr.File(
-                file_types=["image"], label="Upload Images", file_count="multiple",
-            )
-
-            # Wrap the model selection and sliders in an Accordion
-            with gr.Accordion("Advanced Settings", open=False):  # Collapsible by default
-                model_repo = gr.Dropdown(
-                    model_repos,
-                    value=VIT_MODEL_DSV3_REPO,
-                    label="Select Model",
-                )
+        gr.Markdown(f"<h1 style='text-align: center;'>{TITLE}</h1>")
+        gr.Markdown(DESCRIPTION)
 
-                # Fallbacks for possibly invalid args
-                step = 0.1 if isinstance(args.score_slider_step, bool) else float(args.score_slider_step)
-                general_val = 0.25 if isinstance(args.score_general_threshold, bool) else float(args.score_general_threshold)
-                char_val = 1.0 if isinstance(args.score_character_threshold, bool) else float(args.score_character_threshold)
+        with gr.Row():
+            with gr.Column():
 
-                general_thresh = gr.Slider(
-                    0, 1, step=step, value=general_val, label="General Tags Threshold"
+                submit = gr.Button(
+                    value="Process Images", variant="primary"
                 )
-
-                character_thresh = gr.Slider(
-                    0, 1, step=step, value=char_val, label="Character Tags Threshold"
-                )
-
-                filter_tags = gr.Textbox(
-                    value=", ".join(predefined_tags),
-                    label="Filter Tags (comma-separated)",
-                    placeholder="Add tags to filter out (e.g., winter, red, from above)",
-                    lines=9
+                
+                image_files = gr.File(
+                    file_types=["image"], label="Upload Images", file_count="multiple",
                 )
-
-        with gr.Column():
-            output = gr.Textbox(label="Output", lines=10)
-
+                
+                # Wrap the model selection and sliders in an Accordion
+                with gr.Accordion("Advanced Settings", open=False):  # Collapsible by default
+                    model_repo = gr.Dropdown(
+                        model_repos,
+                        value=VIT_MODEL_DSV3_REPO,
+                        label="Select Model",
+                    )
+                    general_thresh = gr.Slider(
+                        0, 1, step=args.score_slider_step, value=args.score_general_threshold, label="General Tags Threshold"
+                    )
+                    character_thresh = gr.Slider(
+                        0, 1, step=args.score_slider_step, value=args.score_character_threshold, label="Character Tags Threshold"
+                    )
+                    filter_tags = gr.Textbox(
+                        value=", ".join(predefined_tags),
+                        label="Filter Tags (comma-separated)",
+                        placeholder="Add tags to filter out (e.g., winter, red, from above)",
+                        lines=9
+                    )
+
+
+            with gr.Column():
+                output = gr.Textbox(label="Output", lines=10)
 
         def process_images(files, model_repo, general_thresh, character_thresh, filter_tags):
             images = [Image.open(file.name) for file in files]