Update app.py

app.py

@@ -15,24 +15,6 @@ 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 = "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"
-
-# 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"
-
 MODEL_FILENAME = "model.onnx"
 LABEL_FILENAME = "selected_tags.csv"
 
@@ -105,10 +87,22 @@ def process_images(files, model_repo, general_thresh, character_thresh, filter_t
     images = [Image.open(file.name) for file in files]
     results = predictor.predict(images, model_repo, general_thresh, character_thresh)
 
+    # Predefined examples
+    predefined_filter_tags = "watermark"  # This tag will be removed if detected
+    predefined_replacement_rules = "1boy -> 1girl"  # "1boy" will be replaced with "1girl"
+    predefined_fallback_rules = "sad, happy -> smile"  # If neither "sad" nor "happy" are present, add "smile"
+
+    # Combine predefined rules with user input
+    filter_tags = f"{predefined_filter_tags}, {filter_tags}".strip()
+    replacement_rules_text = f"{predefined_replacement_rules}\n{replacement_rules_text}".strip()
+    fallback_rules_text = f"{predefined_fallback_rules}\n{fallback_rules_text}".strip()
+
+    # Parse user-defined rules
     filter_set = set(tag.strip().lower() for tag in filter_tags.split(","))
     replacement_rules = parse_replacement_rules(replacement_rules_text)
     fallback_rules = parse_fallback_rules(fallback_rules_text)
 
+    # Generate formatted output
     prompts = []
     for general_tags, character_tags in results:
         general_tags = apply_replacements(general_tags, replacement_rules)
@@ -126,7 +120,7 @@ def process_images(files, model_repo, general_thresh, character_thresh, filter_t
 args = parse_args()
 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]
+model_repos = ["SmilingWolf/wd-swinv2-tagger-v3", "SmilingWolf/wd-convnext-tagger-v3", "SmilingWolf/wd-vit-tagger-v3"]
 
 with gr.Blocks(title=TITLE) as demo:
     gr.Markdown(f"<h1 style='text-align: center;'>{TITLE}</h1>")
@@ -137,7 +131,7 @@ with gr.Blocks(title=TITLE) as demo:
             image_files = gr.File(file_types=["image"], label="Upload Images", file_count="multiple")
 
             with gr.Accordion("Advanced Settings", open=False):
-                model_repo = gr.Dropdown(model_repos, value=VIT_MODEL_DSV3_REPO, label="Select Model")
+                model_repo = gr.Dropdown(model_repos, value="SmilingWolf/wd-vit-tagger-v3", 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")
 
@@ -147,10 +141,10 @@ with gr.Blocks(title=TITLE) as demo:
             output = gr.Textbox(label="Output", lines=10)
 
     with gr.Accordion("Tag Replacements", open=False):
-        replacement_rules_text = gr.Textbox(label="Replacement Rules", lines=5)
+        replacement_rules_text = gr.Textbox(label="Replacement Rules", lines=5, value="1boy -> 1girl")
 
     with gr.Accordion("Fallback Rules", open=False):
-        fallback_rules_text = gr.Textbox(label="Fallback Rules", lines=5)
+        fallback_rules_text = gr.Textbox(label="Fallback Rules", lines=5, value="sad, happy -> smile")
 
     submit.click(process_images, inputs=[image_files, model_repo, general_thresh, character_thresh, replacement_rules_text, fallback_rules_text], outputs=output)