Update app.py

app.py

@@ -15,6 +15,25 @@ 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"
+
+# Files to download from the repos
 MODEL_FILENAME = "model.onnx"
 LABEL_FILENAME = "selected_tags.csv"
 
@@ -25,39 +44,11 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--score-character-threshold", type=float, default=1.0)
     return parser.parse_args()
 
-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 load_labels(dataframe) -> list[str]:
+    tag_names = dataframe["name"].tolist()
+    general_indexes = list(np.where(dataframe["category"] == 0)[0])
+    character_indexes = list(np.where(dataframe["category"] == 4)[0])
+    return tag_names, general_indexes, character_indexes
 
 class Predictor:
     def __init__(self):
@@ -75,7 +66,7 @@ class Predictor:
 
         csv_path, model_path = self.download_model(model_repo)
         tags_df = pd.read_csv(csv_path)
-        self.tag_names, self.general_indexes, self.character_indexes = tags_df["name"].tolist(), list(np.where(tags_df["category"] == 0)[0]), list(np.where(tags_df["category"] == 4)[0])
+        self.tag_names, self.general_indexes, self.character_indexes = load_labels(tags_df)
 
         model = rt.InferenceSession(model_path)
         _, height, width, _ = model.get_inputs()[0].shape
@@ -83,70 +74,183 @@ class Predictor:
         self.last_loaded_repo = model_repo
         self.model = model
 
-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)
-
-    # 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)
-        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]
-
-        prompts.append(f"{', '.join(character_tags)}, {', '.join(general_tags)}" if character_tags else ", ".join(general_tags))
-
-    return "\n\n".join(prompts)
-
-args = parse_args()
-predictor = Predictor()
-
-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>")
-    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="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")
-
-            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", lines=5, value="1boy -> 1girl")
-
-    with gr.Accordion("Fallback Rules", open=False):
-        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)
-
-demo.queue()
-demo.launch()
+    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
+        pad_top = (max_dim - image.height) // 2
+        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)
+
+
+    def predict(self, images, model_repo, general_thresh, character_thresh):
+        self.load_model(model_repo)
+        results = []
+
+        for image in images:
+            image = self.prepare_image(image)
+            input_name = self.model.get_inputs()[0].name
+            label_name = self.model.get_outputs()[0].name
+            preds = self.model.run([label_name], {input_name: image})[0]
+
+            labels = list(zip(self.tag_names, preds[0].astype(float)))
+            general_res = [x[0] for i, x in enumerate(labels) if i in self.general_indexes and x[1] > general_thresh]
+            character_res = [x[0] for i, x in enumerate(labels) if i in self.character_indexes and x[1] > character_thresh]
+            results.append((general_res, character_res))
+
+        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()