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,11 +9,13 @@ 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
+import uvicorn
 
 TITLE = "WaifuDiffusion Tagger"
 DESCRIPTION = """
 Demo for the WaifuDiffusion tagger models
-
 Example image by [ほし☆☆☆](https://www.pixiv.net/en/users/43565085)
 """
 
@@ -29,34 +33,14 @@ 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"
 
-# Files to download from the repos
 MODEL_FILENAME = "model.onnx"
 LABEL_FILENAME = "selected_tags.csv"
 
-# https://github.com/toriato/stable-diffusion-webui-wd14-tagger/blob/a9eacb1eff904552d3012babfa28b57e1d3e295c/tagger/ui.py#L368
 kaomojis = [
-    "0_0",
-    "(o)_(o)",
-    "+_+",
-    "+_-",
-    "._.",
-    "<o>_<o>",
-    "<|>_<|>",
-    "=_=",
-    ">_<",
-    "3_3",
-    "6_9",
-    ">_o",
-    "@_@",
-    "^_^",
-    "o_o",
-    "u_u",
-    "x_x",
-    "|_|",
-    "||_||",
+    "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)
@@ -65,7 +49,6 @@ def parse_args() -> argparse.Namespace:
     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(
@@ -78,13 +61,9 @@ def load_labels(dataframe) -> list[str]:
     character_indexes = list(np.where(dataframe["category"] == 4)[0])
     return tag_names, rating_indexes, general_indexes, character_indexes
 
-
 def mcut_threshold(probs):
     """
     Maximum Cut Thresholding (MCut)
-    Largeron, C., Moulin, C., & Gery, M. (2012). MCut: A Thresholding Strategy
-     for Multi-label Classification. In 11th International Symposium, IDA 2012
-     (pp. 172-183).
     """
     sorted_probs = probs[probs.argsort()[::-1]]
     difs = sorted_probs[:-1] - sorted_probs[1:]
@@ -92,21 +71,14 @@ def mcut_threshold(probs):
     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
-
+        
     def download_model(self, model_repo):
-        csv_path = huggingface_hub.hf_hub_download(
-            model_repo,
-            LABEL_FILENAME,
-        )
-        model_path = huggingface_hub.hf_hub_download(
-            model_repo,
-            MODEL_FILENAME,
-        )
+        csv_path = huggingface_hub.hf_hub_download(model_repo, LABEL_FILENAME)
+        model_path = huggingface_hub.hf_hub_download(model_repo, MODEL_FILENAME)
         return csv_path, model_path
 
     def load_model(self, model_repo):
@@ -114,7 +86,6 @@ class Predictor:
             return
 
         csv_path, model_path = self.download_model(model_repo)
-
         tags_df = pd.read_csv(csv_path)
         sep_tags = load_labels(tags_df)
 
@@ -132,12 +103,11 @@ class Predictor:
 
     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")
 
-        # Pad image to square
         image_shape = image.size
         max_dim = max(image_shape)
         pad_left = (max_dim - image_shape[0]) // 2
@@ -146,47 +116,36 @@ class Predictor:
         padded_image = Image.new("RGB", (max_dim, max_dim), (255, 255, 255))
         padded_image.paste(image, (pad_left, pad_top))
 
-        # Resize
         if max_dim != target_size:
-            padded_image = padded_image.resize(
-                (target_size, target_size),
-                Image.BICUBIC,
-            )
+            padded_image = padded_image.resize((target_size, target_size), Image.BICUBIC)
 
-        # Convert to numpy array
         image_array = np.asarray(padded_image, dtype=np.float32)
-
-        # Convert PIL-native RGB to BGR
         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,
+        model_repo=SWINV2_MODEL_DSV3_REPO,
+        general_thresh=0.35,
+        general_mcut_enabled=False,
+        character_thresh=0.85,
+        character_mcut_enabled=False,
     ):
         self.load_model(model_repo)
-
+        
         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)))
-
-        # First 4 labels are actually ratings: pick one with argmax
+        
         ratings_names = [labels[i] for i in self.rating_indexes]
         rating = dict(ratings_names)
 
-        # Then we have general tags: pick any where prediction confidence > threshold
         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)
@@ -194,9 +153,7 @@ class Predictor:
         general_res = [x for x in general_names if x[1] > general_thresh]
         general_res = dict(general_res)
 
-        # Everything else is characters: pick any where prediction confidence > threshold
         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)
@@ -205,24 +162,31 @@ class Predictor:
         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("(", "\(").replace(")", "\)")
-        )
+        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
 
+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_strings, _, _, _ = predictor.predict(
+        pil_image, 
+        general_thresh=threshold
+    )
+    tags = sorted_general_strings.split(", ")
+    return JSONResponse(content={"tags": tags})
 
-def main():
+def create_demo():
     args = parse_args()
-
-    predictor = Predictor()
-
+    
     dropdown_list = [
         SWINV2_MODEL_DSV3_REPO,
         CONV_MODEL_DSV3_REPO,
@@ -238,10 +202,8 @@ def main():
 
     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)
+            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")
@@ -297,14 +259,7 @@ def main():
                     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,
-                        ]
-                    )
+                    clear.add([sorted_general_strings, rating, character_res, general_res])
 
         submit.click(
             predictor.predict,
@@ -318,7 +273,7 @@ def main():
             ],
             outputs=[sorted_general_strings, rating, character_res, general_res],
         )
-
+        
         gr.Examples(
             [["power.jpg", SWINV2_MODEL_DSV3_REPO, 0.35, False, 0.85, False]],
             inputs=[
@@ -330,10 +285,12 @@ def main():
                 character_mcut_enabled,
             ],
         )
+        
+        demo.queue(max_size=10)
+        return demo
 
-    demo.queue(max_size=10)
-    demo.launch()
-
+app = FastAPI()
+app = gr.mount_gradio_app(app, create_demo(), path="/")
 
 if __name__ == "__main__":
-    main()
+    uvicorn.run(app, host="0.0.0.0", port=7860)