Update app.py

app.py

@@ -17,12 +17,21 @@ 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"
 
 # Files to download from the repos
 MODEL_FILENAME = "model.onnx"
@@ -115,14 +124,58 @@ def main():
 
     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 = ["monochrome",
+    predefined_tags = ["loli",
+                       "oppai_loli",
+                       "onee-shota",
+                       "incest",
+                       "furry",
+                       "furry_female",
+                       "shota",
+                       "male_focus",
+                       "signature",
+                       "lolita_hairband",
+                       "otoko_no_ko",
+                       "minigirl",
+                       "patreon_username",
+                       "babydoll",
+                       "monochrome",
                        "happy_birthday",
+                       "happy_new_year", 
+                       "dated", 
+                       "thought_bubble",
+                       "greyscale",
+                       "speech_bubble",
+                       "english_text",
+                       "copyright_name",
+                       "twitter_username",
+                       "patreon username",
+                       "patreon logo",
+                       "cover", 
+                       "content_rating"
+                       "cover_page", 
+                       "doujin_cover", 
+                       "sex",
+                       "artist_name",
+                       "watermark",
+                       "censored",
+                       "bar_censor",
+                       "blank_censor",
+                       "blur_censor",
                        "light_censor",
                        "mosaic_censoring"]
 
@@ -163,43 +216,82 @@ def main():
             with gr.Column():
                 output = gr.Textbox(label="Output", lines=10)
 
-        def process_images(files, model_repo, general_thresh, character_thresh, filter_tags, expected_tags_input, default_tag_input):
+        def parse_replacement_rules(rules_text):
+            """Parse user-defined tag replacement rules into a dictionary."""
+            rules = {}
+            for line in rules_text.strip().split("\n"):
+                if "->" in line:
+                    old_tags, new_tags = map(str.strip, line.split("->"))
+                    old_tags_list = tuple(map(str.strip, old_tags.lower().split(",")))
+                    new_tags_list = [tag.strip() for tag in new_tags.split(",")]
+                    rules[old_tags_list] = new_tags_list
+            return rules
+
+        def apply_replacements(tags, replacement_rules):
+            """Apply replacement rules to a set of tags."""
+            tags_set = set(tags)
+        
+            for old_tags, new_tags in replacement_rules.items():
+                if set(old_tags).issubset(tags_set):  # If all old tags exist in the set
+                    tags_set.difference_update(old_tags)  # Remove old tags
+                    tags_set.update(new_tags)  # Add new ones
+        
+            return list(tags_set)
+
+        def process_images(files, model_repo, general_thresh, character_thresh, filter_tags, replacement_rules_text):
             images = [Image.open(file.name) for file in files]
             results = predictor.predict(images, model_repo, general_thresh, character_thresh)
-
+        
             # Parse filter tags
             filter_set = set(tag.strip().lower() for tag in filter_tags.split(","))
-            expected_tags = set(tag.strip().lower() for tag in expected_tags_input.split(","))  # Define expected tags dynamically
-            default_tag = default_tag_input.strip()  # Define default tag dynamically
-
+        
+            # Parse user-defined replacements
+            replacement_rules = parse_replacement_rules(replacement_rules_text)
+        
             # Generate formatted output
             prompts = []
-            for i, (general_tags, character_tags) in enumerate(results):
-                # Replace underscores with spaces for both character and general tags
-                character_part = [tag.replace('_', ' ') for tag in character_tags if tag.lower() not in filter_set]
-                general_part = [tag.replace('_', ' ') for tag in general_tags if tag.lower() not in filter_set]
-                
-                # Check if expected tags are present
-                if not any(tag in expected_tags for tag in character_part + general_part):
-                    general_part.append(default_tag)
-                
-                # Construct the prompt based on the presence of character_part
-                prompt = ", ".join(character_part + general_part)
-                prompts.append(prompt)
-
-            # Join all prompts with blank lines
+            for general_tags, character_tags in results:
+                # Apply replacements
+                general_tags = apply_replacements(general_tags, replacement_rules)
+                character_tags = apply_replacements(character_tags, replacement_rules)
+        
+                # Remove filtered tags and format
+                general_tags = [tag.replace('_', ' ') for tag in general_tags if tag.lower() not in filter_set]
+                character_tags = [tag.replace('_', ' ') for tag in character_tags if tag.lower() not in filter_set]
+        
+                # Construct final prompt
+                if character_tags:
+                    prompts.append(f"{', '.join(character_tags)}, {', '.join(general_tags)}")
+                else:
+                    prompts.append(", ".join(general_tags))
+        
             return "\n\n".join(prompts)
 
-
-
-        submit.click(
-            process_images,
-            inputs=[image_files, model_repo, general_thresh, character_thresh, filter_tags],
-            outputs=output
-        )
-
-    demo.queue(max_size=10)
-    demo.launch()
-
-if __name__ == "__main__":
-    main()
+        # Modify UI to include replacement rules input
+        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():
+                    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")
+                        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)", lines=3)
+        
+                    submit = gr.Button(value="Process Images", variant="primary")
+        
+                with gr.Column():
+                    output = gr.Textbox(label="Output", lines=10)
+        
+            # Separate input for tag replacements
+            with gr.Accordion("Tag Replacements", open=False):
+                replacement_rules_text = gr.Textbox(label="Enter replacement rules (one per line)", placeholder="e.g.,\n1boy -> 1girl\nwinter, indoors, living room -> summer, outdoors", lines=5)
+        
+            submit.click(process_images, inputs=[image_files, model_repo, general_thresh, character_thresh, filter_tags, replacement_rules_text], outputs=output)
+
+demo.queue(max_size=10)
+demo.launch()