Update app.py

app.py

@@ -1,5 +1,7 @@
 import argparse
 import os
+from typing import Optional
+import io
 
 import gradio as gr
 import huggingface_hub
@@ -7,21 +9,22 @@ import numpy as np
 import onnxruntime as rt
 import pandas as pd
 from PIL import Image
+from fastapi import FastAPI, File, UploadFile, Form
+from fastapi.responses import JSONResponse
+
+app = FastAPI()
 
 TITLE = "WaifuDiffusion Tagger"
-DESCRIPTION = """
-Demo for the WaifuDiffusion tagger models
-Example image by [ほし☆☆☆](https://www.pixiv.net/en/users/43565085)
-"""
+DESCRIPTION = "Demo for the WaifuDiffusion tagger models"
 
-# Dataset v3 series of models:
+# Dataset v3 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:
+# Dataset v2 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"
@@ -31,37 +34,8 @@ VIT_MODEL_DSV2_REPO = "SmilingWolf/wd-v1-4-vit-tagger-v2"
 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)
-    parser.add_argument("--share", action="store_true")
-    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
+kaomojis = ["0_0", "(o)_(o)", "+_+", "+_-", "._.", "<o>_<o>", "<|>_<|>", "=_=", ">_<",
+            "3_3", "6_9", ">_o", "@_@", "^_^", "o_o", "u_u", "x_x", "|_|", "||_||"]
 
 class Predictor:
     def __init__(self):
@@ -79,52 +53,40 @@ class Predictor:
 
         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]
+        name_series = tags_df["name"]
+        name_series = name_series.map(lambda x: x.replace("_", " ") if x not in kaomojis else x)
+        
+        self.tag_names = name_series.tolist()
+        self.rating_indexes = list(np.where(tags_df["category"] == 9)[0])
+        self.general_indexes = list(np.where(tags_df["category"] == 0)[0])
+        self.character_indexes = list(np.where(tags_df["category"] == 4)[0])
 
-        model = rt.InferenceSession(model_path)
-        _, height, width, _ = model.get_inputs()[0].shape
+        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
-        self.model = model
 
     def prepare_image(self, image):
-        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
+        max_dim = max(image.size)
+        pad_left = (max_dim - image.size[0]) // 2
+        pad_top = (max_dim - image.size[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)
+        if max_dim != self.model_target_size:
+            padded_image = padded_image.resize((self.model_target_size, self.model_target_size), Image.BICUBIC)
 
         image_array = np.asarray(padded_image, dtype=np.float32)
         image_array = image_array[:, :, ::-1]
         
         return np.expand_dims(image_array, axis=0)
 
-    def predict(
-        self,
-        image,
-        model_repo,
-        general_thresh,
-        general_mcut_enabled,
-        character_thresh,
-        character_mcut_enabled,
-    ):
+    def predict(self, image, model_repo=SWINV2_MODEL_DSV3_REPO, threshold=0.05):
         self.load_model(model_repo)
         
         image = self.prepare_image(image)
@@ -133,139 +95,94 @@ class Predictor:
         preds = self.model.run([label_name], {input_name: image})[0]
 
         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 = [x for x in general_names if x[1] > threshold]
         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 = sorted(general_res.items(), key=lambda x: x[1], reverse=True)
-        sorted_general_strings = [x[0] for x in sorted_general]
-        sorted_general_strings = ", ".join(sorted_general_strings).replace("(", "\(").replace(")", "\)")
-
-        return sorted_general_strings, rating, character_res, general_res
-
-def main():
-    args = parse_args()
-    predictor = Predictor()
-
-    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,
-    ]
-
+        return sorted_general, labels
+
+predictor = Predictor()
+
+@app.post("/tagging")
+async def tagging_endpoint(
+    image: UploadFile = File(...),
+    threshold: Optional[float] = Form(0.05)
+):
+    image_data = await image.read()
+    pil_image = Image.open(io.BytesIO(image_data)).convert("RGBA")
+    sorted_general, _ = predictor.predict(pil_image, threshold=threshold)
+    return JSONResponse(content={"tags": [x[0] for x in sorted_general]})
+
+def ui_predict(
+    image,
+    model_repo,
+    general_thresh,
+    general_mcut_enabled,
+    character_thresh,
+    character_mcut_enabled,
+):
+    sorted_general, all_labels = predictor.predict(image, model_repo, general_thresh)
+    
+    # Ratings
+    ratings = {all_labels[i][0]: all_labels[i][1] for i in predictor.rating_indexes}
+    
+    # Characters
+    character_labels = [all_labels[i] for i in predictor.character_indexes]
+    if character_mcut_enabled:
+        character_probs = np.array([x[1] for x in character_labels])
+        character_thresh = max(0.15, np.mean(character_probs))
+    character_res = {x[0]: x[1] for x in character_labels if x[1] > character_thresh}
+
+    # Format output
+    sorted_general_strings = ", ".join(x[0] for x in sorted_general).replace("(", "\(").replace(")", "\)")
+    return sorted_general_strings, ratings, character_res, dict(sorted_general)
+
+def create_demo():
     with gr.Blocks(title=TITLE) as demo:
-        with gr.Column():
-            gr.Markdown(f"<h1 style='text-align: center; margin-bottom: 1rem'>{TITLE}</h1>")
-            gr.Markdown(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",
-                    )
-                    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():
-                        clear = gr.ClearButton(
-                            components=[
-                                image,
-                                model_repo,
-                                general_thresh,
-                                general_mcut_enabled,
-                                character_thresh,
-                                character_mcut_enabled,
-                            ],
-                            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])
+        gr.Markdown(f"<h1 style='text-align: center; margin-bottom: 1rem'>{TITLE}</h1>")
+        gr.Markdown(DESCRIPTION)
+        
+        with gr.Row():
+            with gr.Column(variant="panel"):
+                image = gr.Image(type="pil", image_mode="RGBA", label="Input")
+                model_repo = gr.Dropdown(
+                    choices=[
+                        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
+                    ],
+                    value=SWINV2_MODEL_DSV3_REPO,
+                    label="Model"
+                )
+                with gr.Row():
+                    general_thresh = gr.Slider(0, 1, value=0.35, step=0.05, label="General Tags Threshold")
+                    general_mcut = gr.Checkbox(value=False, label="Use MCut threshold")
+                with gr.Row():
+                    character_thresh = gr.Slider(0, 1, value=0.85, step=0.05, label="Character Tags Threshold")
+                    character_mcut = gr.Checkbox(value=False, label="Use MCut threshold")
+                submit = gr.Button(value="Submit", variant="primary")
+
+            with gr.Column(variant="panel"):
+                text_output = gr.Textbox(label="Output (string)")
+                rating_output = gr.Label(label="Rating")
+                character_output = gr.Label(label="Characters")
+                general_output = gr.Label(label="Tags")
 
         submit.click(
-            predictor.predict,
-            inputs=[
-                image,
-                model_repo,
-                general_thresh,
-                general_mcut_enabled,
-                character_thresh,
-                character_mcut_enabled,
-            ],
-            outputs=[sorted_general_strings, rating, character_res, general_res],
+            ui_predict,
+            inputs=[image, model_repo, general_thresh, general_mcut,
+                   character_thresh, character_mcut],
+            outputs=[text_output, rating_output, character_output, general_output]
         )
-        
-        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()
+        return demo
+
+app = gr.mount_gradio_app(app, create_demo(), path="/")
 
 if __name__ == "__main__":
-    main()
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)