Super-squash branch 'main' using huggingface_hub

Co-authored-by: p1atdev <p1atdev@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,13 @@
+---
+title: Wd Tagger Transformers
+emoji: 😻
+colorFrom: pink
+colorTo: yellow
+sdk: gradio
+sdk_version: 4.36.1
+app_file: app.py
+pinned: false
+license: apache-2.0
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,276 @@
+from PIL import Image
+
+import torch
+
+from transformers import (
+    AutoImageProcessor,
+    AutoModelForImageClassification,
+)
+
+import gradio as gr
+import spaces  # ZERO GPU
+
+MODEL_NAMES = ["p1atdev/wd-swinv2-tagger-v3-hf"]
+MODEL_NAME = MODEL_NAMES[0]
+
+model = AutoModelForImageClassification.from_pretrained(
+    MODEL_NAME,
+)
+model.to("cuda" if torch.cuda.is_available() else "cpu")
+processor = AutoImageProcessor.from_pretrained(MODEL_NAME, trust_remote_code=True)
+
+# ref: https://qiita.com/tregu148/items/fccccbbc47d966dd2fc2
+def gradio_copy_text(_text: None):
+    gr.Info("Copied!")
+
+COPY_ACTION_JS = """\
+(inputs, _outputs) => {
+  // inputs is the string value of the input_text
+  if (inputs.trim() !== "") {
+    navigator.clipboard.writeText(inputs);
+  }
+}"""
+
+def _people_tag(noun: str, minimum: int = 1, maximum: int = 5):
+    return (
+        [f"1{noun}"]
+        + [f"{num}{noun}s" for num in range(minimum + 1, maximum + 1)]
+        + [f"{maximum+1}+{noun}s"]
+    )
+
+
+PEOPLE_TAGS = (
+    _people_tag("girl") + _people_tag("boy") + _people_tag("other") + ["no humans"]
+)
+RATING_MAP = {
+    "general": "safe",
+    "sensitive": "sensitive",
+    "questionable": "nsfw",
+    "explicit": "explicit, nsfw",
+}
+RATING_MAP_E621 = {
+    "general": "rating_safe",
+    "sensitive": "rating_safe",
+    "questionable": "rating_questionable",
+    "explicit": "rating_explicit",
+}
+
+DESCRIPTION_MD = """
+# WD Tagger with 🤗 transformers
+Currently supports the following model(s):
+- [p1atdev/wd-swinv2-tagger-v3-hf](https://huggingface.co/p1atdev/wd-swinv2-tagger-v3-hf)
+
+""".strip()
+
+
+def character_list_to_series_list(character_list):
+    def get_series_dict():
+        import re
+
+        with open('characterfull.txt', 'r') as f:
+            lines = f.readlines()
+
+        series_dict = {}
+        for line in lines:
+            parts = line.strip().split(', ')
+            if len(parts) >= 3:
+                name = parts[-2].replace("\\", "")
+            if name.endswith(")"):
+                names = name.split("(")
+                character_name = "(".join(names[:-1])
+                if character_name.endswith(" "):
+                    name = character_name[:-1]
+            series = re.sub(r'\\[()]', '', parts[-1])
+            series_dict[name] = series
+
+        return series_dict
+    
+    output_series_tag = []
+    series_tag = ""
+    series_dict = get_series_dict()
+    for tag in character_list:
+        series_tag = series_dict.get(tag, "")
+        if tag.endswith(")"):
+            tags = tag.split("(")
+            character_tag = "(".join(tags[:-1])
+            if character_tag.endswith(" "):
+                character_tag = character_tag[:-1]
+            series_tag = tags[-1].replace(")", "")
+
+    if series_tag:
+        output_series_tag.append(series_tag)
+
+    return output_series_tag
+
+
+def get_e621_dict():
+    with open('danbooru_e621.csv', 'r', encoding="utf-8") as f:
+        lines = f.readlines()
+
+    e621_dict = {}
+    for line in lines:
+        parts = line.strip().split(',')
+        e621_dict[parts[0]] = parts[1]
+
+    return e621_dict
+
+
+def danbooru_to_e621(dtag, e621_dict):
+    def d_to_e(match, e621_dict):
+        dtag = match.group(0)
+        etag = e621_dict.get(dtag.strip().replace("_", " "), "")
+        if etag:
+            return etag
+        else:
+            return dtag
+    
+    import re
+    tag = re.sub(r'[\w ]+', lambda wrapper: d_to_e(wrapper, e621_dict), dtag, 2)
+
+    return tag
+
+def postprocess_results(
+    results: dict[str, float], general_threshold: float, character_threshold: float
+):
+    results = {
+        k: v for k, v in sorted(results.items(), key=lambda item: item[1], reverse=True)
+    }
+
+    rating = {}
+    character = {}
+    general = {}
+
+    for k, v in results.items():
+        if k.startswith("rating:"):
+            rating[k.replace("rating:", "")] = v
+            continue
+        elif k.startswith("character:"):
+            character[k.replace("character:", "")] = v
+            continue
+
+        general[k] = v
+
+    character = {k: v for k, v in character.items() if v >= character_threshold}
+    general = {k: v for k, v in general.items() if v >= general_threshold}
+
+    return rating, character, general
+
+
+def animagine_prompt(rating: list[str], character: list[str], general: list[str], tag_type):
+    people_tags: list[str] = []
+    other_tags: list[str] = []
+    if tag_type == "e621":
+        rating_tag = RATING_MAP_E621[rating[0]]
+    else:
+        rating_tag = RATING_MAP[rating[0]]
+
+    e621_dict = get_e621_dict()
+    for tag in general:
+        if tag_type == "e621":
+            tag = danbooru_to_e621(tag, e621_dict)
+        if tag in PEOPLE_TAGS:
+            people_tags.append(tag)
+        else:
+            other_tags.append(tag)
+
+    output_series_tag = character_list_to_series_list(character)
+
+    all_tags = people_tags + character + output_series_tag + other_tags + [rating_tag]
+
+    return ", ".join(all_tags)
+
+
+@spaces.GPU(enable_queue=True)
+def predict_tags(
+    image: Image.Image, general_threshold: float = 0.3, character_threshold: float = 0.8, tag_type = "danbooru"
+):
+    inputs = processor.preprocess(image, return_tensors="pt")
+
+    outputs = model(**inputs.to(model.device, model.dtype))
+    logits = torch.sigmoid(outputs.logits[0])  # take the first logits
+
+    # get probabilities
+    results = {
+        model.config.id2label[i]: float(logit.float()) for i, logit in enumerate(logits)
+    }
+
+    # rating, character, general
+    rating, character, general = postprocess_results(
+        results, general_threshold, character_threshold
+    )
+
+    prompt = animagine_prompt(
+        list(rating.keys()), list(character.keys()), list(general.keys()), tag_type
+    )
+    
+    return rating, character, general, prompt, gr.update(interactive=True,)
+
+
+def demo():
+    with gr.Blocks() as ui:
+        gr.Markdown(DESCRIPTION_MD)
+
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(label="Input image", type="pil")
+
+                with gr.Group():
+                    general_threshold = gr.Slider(
+                        label="Threshold",
+                        minimum=0.0,
+                        maximum=1.0,
+                        value=0.3,
+                        step=0.01,
+                        interactive=True,
+                    )
+                    character_threshold = gr.Slider(
+                        label="Character threshold",
+                        minimum=0.0,
+                        maximum=1.0,
+                        value=0.8,
+                        step=0.01,
+                        interactive=True,
+                    )
+                    tag_type = gr.Radio(
+                        label="Output tag conversion",
+                        info="danbooru for Animagine, e621 for Pony.",
+                        choices=["danbooru", "e621"],
+                        value="danbooru",
+                    )
+                
+                _model_radio = gr.Dropdown(
+                    choices=MODEL_NAMES,
+                    label="Model",
+                    value=MODEL_NAMES[0],
+                    interactive=True,
+                )
+
+                start_btn = gr.Button(value="Start", variant="primary")
+
+            with gr.Column():
+                
+                with gr.Group():
+                    prompt_text = gr.TextArea(label="Prompt", interactive=False)
+                    copy_btn = gr.Button(value="Copy to clipboard", interactive=False)
+
+                rating_tags_label = gr.Label(label="Rating tags")
+                character_tags_label = gr.Label(label="Character tags")
+                general_tags_label = gr.Label(label="General tags")
+
+        start_btn.click(
+            predict_tags,
+            inputs=[input_image, general_threshold, character_threshold, tag_type],
+            outputs=[
+                rating_tags_label,
+                character_tags_label,
+                general_tags_label,
+                prompt_text,
+                copy_btn,
+            ],
+        )
+        copy_btn.click(gradio_copy_text, inputs=[prompt_text], js=COPY_ACTION_JS)
+
+    return ui
+
+if __name__ == "__main__":
+    demo().queue().launch()

requirements.txt

@@ -0,0 +1,5 @@
+torch
+torchvision
+accelerate
+transformers
+spaces