Update app.py

app.py

@@ -1,7 +1,319 @@
+import argparse
+import os
 import gradio as gr
+import huggingface_hub
+import numpy as np
+import onnxruntime as rt
+import pandas as pd
+from PIL import Image
+import json  # Added for loading metadata.json from the inference file
 
-def greet(name):
-    return "Hello " + name + "!!"
+TITLE = "WaifuDiffusion Tagger"
+DESCRIPTION = """
+Demo for the WaifuDiffusion tagger models
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+Example image by [ほし☆☆☆](https://www.pixiv.net/en/users/43565085)
+"""
+
+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"
+
+kaomojis = [
+    "0_0",
+    "(o)_(o)",
+    "+_+",
+    "+_-",
+    "._.",
+    "<o>_<o>",
+    "<|>_<|>",
+    "=_=",
+    ">_<",
+    "3_3",
+    "6_9",
+    ">_o",
+    "@_@",
+    "^_^",
+    "o_o",
+    "u_u",
+    "x_x",
+    "|_|",
+    "||_||",
+]
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--score-slider-step", type=float, default=0.05)
+    parser.add_argument("--score-general-threshold", type=float, default=0.35)
+    parser.add_argument("--score-character-threshold", type=float, default=0.85)
+    return parser.parse_args()
+
+def load_labels(dataframe) -> list[str]:
+    name_series = dataframe["name"]
+    name_series = name_series.map(
+        lambda x: x.replace("_", " ") if x not in kaomojis else x
+    )
+    tag_names = name_series.tolist()
+
+    rating_indexes = list(np.where(dataframe["category"] == 9)[0])
+    general_indexes = list(np.where(dataframe["category"] == 0)[0])
+    character_indexes = list(np.where(dataframe["category"] == 4)[0])
+    return tag_names, rating_indexes, general_indexes, character_indexes
+
+def mcut_threshold(probs):
+    sorted_probs = probs[probs.argsort()[::-1]]
+    difs = sorted_probs[:-1] - sorted_probs[1:]
+    t = difs.argmax()
+    thresh = (sorted_probs[t] + sorted_probs[t + 1]) / 2
+    return thresh
+
+class Predictor:
+    def __init__(self):
+        self.model_target_size = None
+        self.last_loaded_repo = None
+        # Added flag to distinguish between custom and Hugging Face models
+        self.is_custom_model = False
+
+    def download_model(self, model_repo):
+        csv_path = huggingface_hub.hf_hub_download(
+            model_repo,
+            LABEL_FILENAME,
+            use_auth_token=HF_TOKEN,
+        )
+        model_path = huggingface_hub.hf_hub_download(
+            model_repo,
+            MODEL_FILENAME,
+            use_auth_token=HF_TOKEN,
+        )
+        return csv_path, model_path
+
+    def load_model(self, model_repo, onnx_path=None, metadata_path=None):
+        # Modified to accept onnx_path and metadata_path for custom model support
+        if model_repo == "Custom Model" and onnx_path and metadata_path:
+            # Check if the custom model files have already been loaded
+            if self.last_loaded_repo == (onnx_path, metadata_path):
+                return
+            self.is_custom_model = True
+            # Load the ONNX model from the provided path (from inference file)
+            self.model = rt.InferenceSession(onnx_path)
+            # Load metadata from metadata.json (from inference file)
+            with open(metadata_path, "r", encoding="utf-8") as f:
+                metadata = json.load(f)
+            self.idx_to_tag = metadata["idx_to_tag"]
+            # Create tag_names list from idx_to_tag dictionary
+            self.tag_names = [self.idx_to_tag[str(i)] for i in range(len(self.idx_to_tag))]
+            # Set target size to 512 for custom model, as per inference file
+            self.model_target_size = 512
+            self.last_loaded_repo = (onnx_path, metadata_path)
+        else:
+            # Existing logic for Hugging Face models
+            self.is_custom_model = False
+            if self.last_loaded_repo == model_repo:
+                return
+            csv_path, model_path = self.download_model(model_repo)
+            tags_df = pd.read_csv(csv_path)
+            sep_tags = load_labels(tags_df)
+            self.tag_names = sep_tags[0]
+            self.rating_indexes = sep_tags[1]
+            self.general_indexes = sep_tags[2]
+            self.character_indexes = sep_tags[3]
+            self.model = rt.InferenceSession(model_path)
+            _, height, width, _ = self.model.get_inputs()[0].shape
+            self.model_target_size = height
+            self.last_loaded_repo = model_repo
+
+    def prepare_image(self, image):
+        if self.is_custom_model:
+            # Added preprocessing logic from inference file's preprocess_image function
+            # Adapted to take a PIL image instead of a file path
+            target_size = self.model_target_size
+            img = image.convert("RGB")
+            w, h = img.size
+            aspect = w / h
+            if aspect > 1:
+                new_w = target_size
+                new_h = int(new_w / aspect)
+            else:
+                new_h = target_size
+                new_w = int(new_h * aspect)
+            img = img.resize((new_w, new_h), Image.Resampling.LANCZOS)
+            background = Image.new("RGB", (target_size, target_size), (0, 0, 0))
+            paste_x = (target_size - new_w) // 2
+            paste_y = (target_size - new_h) // 2
+            background.paste(img, (paste_x, paste_y))
+            arr = np.array(background).astype("float32") / 255.0
+            arr = np.transpose(arr, (2, 0, 1))  # HWC to CHW as per inference file
+            arr = np.expand_dims(arr, axis=0)
+            return arr
+        else:
+            # Existing preprocessing logic for Hugging Face models
+            target_size = self.model_target_size
+            canvas = Image.new("RGBA", image.size, (255, 255, 255))
+            canvas.alpha_composite(image)
+            image = canvas.convert("RGB")
+            image_shape = image.size
+            max_dim = max(image_shape)
+            pad_left = (max_dim - image_shape[0]) // 2
+            pad_top = (max_dim - image_shape[1]) // 2
+            padded_image = Image.new("RGB", (max_dim, max_dim), (255, 255, 255))
+            padded_image.paste(image, (pad_left, pad_top))
+            if max_dim != target_size:
+                padded_image = padded_image.resize((target_size, target_size), Image.BICUBIC)
+            image_array = np.asarray(padded_image, dtype=np.float32)
+            image_array = image_array[:, :, ::-1]  # RGB to BGR
+            return np.expand_dims(image_array, axis=0)
+
+    def predict(
+        self,
+        image,
+        model_repo,
+        general_thresh,
+        general_mcut_enabled,
+        character_thresh,
+        character_mcut_enabled,
+        onnx_path=None,
+        metadata_path=None,
+    ):
+        # Modified to accept onnx_path and metadata_path for custom model
+        self.load_model(model_repo, onnx_path, metadata_path)
+        # Added check to ensure custom model files are provided
+        if self.is_custom_model and (onnx_path is None or metadata_path is None):
+            return "Please upload ONNX model and metadata JSON files.", {}, {}, {}
+        image_tensor = self.prepare_image(image)
+        input_name = self.model.get_inputs()[0].name
+        # Changed to use None for output names to get all outputs, supporting custom model
+        outputs = self.model.run(None, {input_name: image_tensor})
+        if self.is_custom_model:
+            # Added inference logic from inference file for custom model
+            # Handle case where model might output initial and refined predictions
+            refined_preds = outputs[1] if len(outputs) == 2 else outputs[0]
+            ref_logit = refined_preds[0]  # Shape (N_tags,)
+            # Apply sigmoid to convert logits to probabilities (from inference file)
+            ref_prob = 1.0 / (1.0 + np.exp(-ref_logit))
+            pred_indices = np.where(ref_prob >= general_thresh)[0]
+            predicted_tags = [self.tag_names[idx] for idx in pred_indices]
+            sorted_general_strings = ", ".join(predicted_tags)
+            # Custom model doesn't use category separation, so return empty for rating and character
+            rating = {}
+            character_res = {}
+            general_res = {self.tag_names[idx]: ref_prob[idx] for idx in pred_indices}
+        else:
+            # Existing inference logic for Hugging Face models
+            preds = outputs[0]  # Assumes single output tensor
+            labels = list(zip(self.tag_names, preds[0].astype(float)))
+            ratings_names = [labels[i] for i in self.rating_indexes]
+            rating = dict(ratings_names)
+            general_names = [labels[i] for i in self.general_indexes]
+            if general_mcut_enabled:
+                general_probs = np.array([x[1] for x in general_names])
+                general_thresh = mcut_threshold(general_probs)
+            general_res = [x for x in general_names if x[1] > general_thresh]
+            general_res = dict(general_res)
+            character_names = [labels[i] for i in self.character_indexes]
+            if character_mcut_enabled:
+                character_probs = np.array([x[1] for x in character_names])
+                character_thresh = mcut_threshold(character_probs)
+                character_thresh = max(0.15, character_thresh)
+            character_res = [x for x in character_names if x[1] > character_thresh]
+            character_res = dict(character_res)
+            sorted_general_strings = sorted(
+                general_res.items(),
+                key=lambda x: x[1],
+                reverse=True,
+            )
+            sorted_general_strings = [x[0] for x in sorted_general_strings]
+            sorted_general_strings = ", ".join(sorted_general_strings).replace("(", r"\(").replace(")", r"\)")
+        return sorted_general_strings, rating, character_res, general_res
+
+def main():
+    args = parse_args()
+    predictor = Predictor()
+    # Added "Custom Model" to the dropdown list to support local ONNX model
+    dropdown_list = [
+        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,
+        "Custom Model",
+    ]
+    with gr.Blocks(title=TITLE) as demo:
+        with gr.Column():
+            gr.Markdown(value=f"<h1 style='text-align: center; margin-bottom: 1rem'>{TITLE}</h1>")
+            gr.Markdown(value=DESCRIPTION)
+            with gr.Row():
+                with gr.Column(variant="panel"):
+                    image = gr.Image(type="pil", image_mode="RGBA", label="Input")
+                    model_repo = gr.Dropdown(dropdown_list, value=SWINV2_MODEL_DSV3_REPO, label="Model")
+                    # Added file inputs for ONNX model and metadata, hidden by default
+                    with gr.Row(visible=False) as custom_model_inputs:
+                        onnx_file = gr.File(label="ONNX Model File", file_types=[".onnx"])
+                        metadata_file = gr.File(label="Metadata JSON File", file_types=[".json"])
+                    with gr.Row():
+                        general_thresh = gr.Slider(0, 1, step=args.score_slider_step, value=args.score_general_threshold, label="General Tags Threshold", scale=3)
+                        general_mcut_enabled = gr.Checkbox(value=False, label="Use MCut threshold", scale=1)
+                    with gr.Row():
+                        character_thresh = gr.Slider(0, 1, step=args.score_slider_step, value=args.score_character_threshold, label="Character Tags Threshold", scale=3)
+                        character_mcut_enabled = gr.Checkbox(value=False, label="Use MCut threshold", scale=1)
+                    with gr.Row():
+                        # Updated clear button to include new file inputs
+                        clear = gr.ClearButton(
+                            components=[image, model_repo, general_thresh, general_mcut_enabled, character_thresh, character_mcut_enabled, onnx_file, metadata_file],
+                            variant="secondary",
+                            size="lg"
+                        )
+                        submit = gr.Button(value="Submit", variant="primary", size="lg")
+                with gr.Column(variant="panel"):
+                    sorted_general_strings = gr.Textbox(label="Output (string)")
+                    rating = gr.Label(label="Rating")
+                    character_res = gr.Label(label="Output (characters)")
+                    general_res = gr.Label(label="Output (tags)")
+                    clear.add([sorted_general_strings, rating, character_res, general_res])
+        # Added event listener to show/hide custom model inputs based on model selection
+        model_repo.change(
+            lambda x: gr.update(visible=(x == "Custom Model")),
+            inputs=model_repo,
+            outputs=custom_model_inputs,
+        )
+        # Updated submit event to pass onnx_file and metadata_file to predict
+        submit.click(
+            predictor.predict,
+            inputs=[image, model_repo, general_thresh, general_mcut_enabled, character_thresh, character_mcut_enabled, onnx_file, metadata_file],
+            outputs=[sorted_general_strings, rating, character_res, general_res],
+        )
+        gr.Examples(
+            [["power.jpg", SWINV2_MODEL_DSV3_REPO, 0.35, False, 0.85, False]],
+            inputs=[image, model_repo, general_thresh, general_mcut_enabled, character_thresh, character_mcut_enabled],
+        )
+    demo.queue(max_size=10)
+    demo.launch()
+
+if __name__ == "__main__":
+    main()