Update app.py

app.py

@@ -1,6 +1,5 @@
 import argparse
 import os
-
 import gradio as gr
 import huggingface_hub
 import numpy as np
@@ -9,42 +8,24 @@ import pandas as pd
 from PIL import Image
 
 TITLE = "WaifuDiffusion Tagger"
-DESCRIPTION = """
-Demo for the WaifuDiffusion tagger models
-"""
+DESCRIPTION = "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"
+# Model Repositories
 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"
 LABEL_FILENAME = "selected_tags.csv"
 
-def parse_args() -> argparse.Namespace:
+def parse_args():
     parser = argparse.ArgumentParser()
     parser.add_argument("--score-slider-step", type=float, default=0.05)
     parser.add_argument("--score-general-threshold", type=float, default=0.3)
     parser.add_argument("--score-character-threshold", type=float, default=1.0)
     return parser.parse_args()
 
-def load_labels(dataframe) -> list[str]:
+def load_labels(dataframe):
     tag_names = dataframe["name"].tolist()
     general_indexes = list(np.where(dataframe["category"] == 0)[0])
     character_indexes = list(np.where(dataframe["category"] == 4)[0])
@@ -75,20 +56,6 @@ class Predictor:
         self.model = model
 
     def prepare_image(self, image):
-        # Create a white canvas with the same size as the input image
-        canvas = Image.new("RGBA", image.size, (255, 255, 255))
-    
-        # Ensure the input image has an alpha channel for compositing
-        if image.mode != "RGBA":
-            image = image.convert("RGBA")
-    
-        # Composite the input image onto the canvas
-        canvas.alpha_composite(image)
-    
-        # Convert to RGB (alpha channel is no longer needed)
-        image = canvas.convert("RGB")
-
-        # Resize the image to a square of size (model_target_size x model_target_size)
         max_dim = max(image.size)
         padded_image = Image.new("RGB", (max_dim, max_dim), (255, 255, 255))
         pad_left = (max_dim - image.width) // 2
@@ -96,10 +63,7 @@ class Predictor:
         padded_image.paste(image, (pad_left, pad_top))
         padded_image = padded_image.resize((self.model_target_size, self.model_target_size), Image.BICUBIC)
 
-        # Convert the image to a NumPy array
-        image_array = np.asarray(padded_image, dtype=np.float32)[:, :, ::-1]
-        return np.expand_dims(image_array, axis=0)
-
+        return np.expand_dims(np.asarray(padded_image, dtype=np.float32)[:, :, ::-1], axis=0)
 
     def predict(self, images, model_repo, general_thresh, character_thresh):
         self.load_model(model_repo)
@@ -118,139 +82,102 @@ class Predictor:
 
         return results
 
-def main():
-    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,
-    ]
-
-    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"]
-
-    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",
-                )
-
-                # 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=3
-                    )
-
-                submit = gr.Button(
-                    value="Process Images", variant="primary"
-                )
-
-            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]
-            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(","))
-
-            # 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 = ", ".join(
-                    tag.replace('_', ' ') for tag in character_tags if tag.lower() not in filter_set
-                )
-                general_part = ", ".join(
-                    tag.replace('_', ' ') for tag in general_tags if tag.lower() not in filter_set
-                )
-        
-                # Construct the prompt based on the presence of character_part
-                if character_part:
-                    prompts.append(f"{character_part}, {general_part}")
-                else:
-                    prompts.append(general_part)
-
-            # Join all prompts with blank lines
-            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()
+predictor = Predictor()
+
+def parse_replacement_rules(rules_text):
+    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 parse_fallback_rules(fallback_text):
+    fallback_rules = {}
+    for line in fallback_text.strip().split("\n"):
+        if "->" in line:
+            expected_tags, fallback_tag = map(str.strip, line.split("->"))
+            expected_tags_list = tuple(map(str.strip, expected_tags.lower().split(",")))
+            fallback_rules[expected_tags_list] = fallback_tag.strip()
+    return fallback_rules
+
+def apply_replacements(tags, replacement_rules):
+    tags_set = set(tags)
+
+    for old_tags, new_tags in replacement_rules.items():
+        if set(old_tags).issubset(tags_set):
+            tags_set.difference_update(old_tags)
+            tags_set.update(new_tags)
+
+    return list(tags_set)
+
+def apply_fallbacks(tags, fallback_rules):
+    tags_set = set(tags)
+
+    for expected_tags, fallback_tag in fallback_rules.items():
+        if not any(tag in tags_set for tag in expected_tags):
+            tags_set.add(fallback_tag)
+
+    return list(tags_set)
+
+def process_images(files, model_repo, general_thresh, character_thresh, filter_tags, replacement_rules_text, fallback_rules_text):
+    images = [Image.open(file.name) for file in files]
+    results = predictor.predict(images, model_repo, general_thresh, character_thresh)
+
+    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)
+
+    prompts = []
+    for general_tags, character_tags in results:
+        general_tags = apply_replacements(general_tags, replacement_rules)
+        character_tags = apply_replacements(character_tags, replacement_rules)
+
+        general_tags = apply_fallbacks(general_tags, fallback_rules)
+        character_tags = apply_fallbacks(character_tags, fallback_rules)
+
+        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]
+
+        if character_tags:
+            prompts.append(f"{', '.join(character_tags)}, {', '.join(general_tags)}")
+        else:
+            prompts.append(", ".join(general_tags))
+
+    return "\n\n".join(prompts)
+
+args = parse_args()
+
+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([VIT_MODEL_DSV3_REPO], 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(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)
+
+    with gr.Accordion("Tag Replacements", open=False):
+        replacement_rules_text = gr.Textbox(label="Replacement Rules", placeholder="e.g., 1boy -> 1girl", lines=5)
+
+    with gr.Accordion("Fallback Rules", open=False):
+        fallback_rules_text = gr.Textbox(label="Fallback Rules", placeholder="e.g., sad, happy -> smile", lines=5)
+
+    submit.click(process_images, 
+                 inputs=[image_files, model_repo, general_thresh, character_thresh, filter_tags, replacement_rules_text, fallback_rules_text], 
+                 outputs=output)
+
+demo.queue(max_size=10)
+demo.launch()