app.py

from os import getenv
from pathlib import Path
from typing import Optional

import gradio as gr
import numpy as np
import onnxruntime as rt
from PIL import Image

from tagger.common import LabelData, load_labels_hf, preprocess_image
from tagger.model import create_session

TITLE = "WaifuDiffusion Tagger"
DESCRIPTION = """
Tag images with the WaifuDiffusion Tagger models!

Primarily used as a backend for a Discord bot.
"""
HF_TOKEN = getenv("HF_TOKEN", None)

MODEL_VARIANTS: dict[str, str] = {
    "v3": {
        "SwinV2": "SmilingWolf/wd-swinv2-tagger-v3",
        "ConvNeXT": "SmilingWolf/wd-convnext-tagger-v3",
        "ViT": "SmilingWolf/wd-vit-tagger-v3",
    },
    "v2": {
        "MOAT": "SmilingWolf/wd-v1-4-moat-tagger-v2",
        "SwinV2": "SmilingWolf/wd-v1-4-swinv2-tagger-v2",
        "ConvNeXT": "SmilingWolf/wd-v1-4-convnext-tagger-v2",
        "ConvNeXTv2": "SmilingWolf/wd-v1-4-convnextv2-tagger-v2",
        "ViT": "SmilingWolf/wd-v1-4-vit-tagger-v2",
    },
}
# prepopulate cache keys in model cache
cache_keys = ["-".join([x, y]) for x in MODEL_VARIANTS.keys() for y in MODEL_VARIANTS[x].keys()]
loaded_models: dict[str, Optional[rt.InferenceSession]] = {k: None for k in cache_keys}

# get the repo root (or the current working directory if running in ipython)
WORK_DIR = Path(__file__).parent.resolve() if "__file__" in globals() else Path().resolve()
# allowed extensions
IMAGE_EXTENSIONS = [".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp", ".tiff", ".tif"]

# get the example images
example_images = sorted(
    [
        str(x.relative_to(WORK_DIR))
        for x in WORK_DIR.joinpath("examples").iterdir()
        if x.is_file() and x.suffix.lower() in IMAGE_EXTENSIONS
    ]
)


def load_model(version: str, variant: str) -> rt.InferenceSession:
    global loaded_models

    # resolve the repo name
    model_repo = MODEL_VARIANTS.get(version, {}).get(variant, None)
    if model_repo is None:
        raise ValueError(f"Unknown model variant: {version}-{variant}")

    cache_key = f"{version}-{variant}"
    if loaded_models.get(cache_key, None) is None:
        # save model to cache
        loaded_models[cache_key] = create_session(model_repo, token=HF_TOKEN)

    return loaded_models[cache_key]


def mcut_threshold(probs: np.ndarray) -> float:
    """
    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).
    """
    probs = probs[probs.argsort()[::-1]]
    diffs = probs[:-1] - probs[1:]
    idx = diffs.argmax()
    thresh = (probs[idx] + probs[idx + 1]) / 2
    return float(thresh)


def predict(
    image: Image.Image,
    version: str,
    variant: str,
    gen_threshold: float = 0.35,
    gen_use_mcut: bool = False,
    char_threshold: float = 0.85,
    char_use_mcut: bool = False,
):
    # join variant for cache key
    model: rt.InferenceSession = load_model(version, variant)
    # load labels
    labels: LabelData = load_labels_hf(MODEL_VARIANTS[version][variant])

    # get input size and name
    _, h, w, _ = model.get_inputs()[0].shape
    input_name = model.get_inputs()[0].name
    output_name = model.get_outputs()[0].name

    # preprocess image
    image = preprocess_image(image, (h, w))

    # turn into BGR24 numpy array of N,H,W,C since thats what these want
    inputs = image.convert("RGB").convert("BGR;24")
    inputs = np.array(inputs).astype(np.float32)
    inputs = np.expand_dims(inputs, axis=0)

    # Run the ONNX model
    probs = model.run([output_name], {input_name: inputs})

    # Convert indices+probs to labels
    probs = list(zip(labels.names, probs[0][0].astype(float)))

    # First 4 labels are actually ratings
    rating_labels = dict([probs[i] for i in labels.rating])

    # General labels, pick any where prediction confidence > threshold
    if gen_use_mcut:
        gen_array = np.array([probs[i][1] for i in labels.general])
        gen_threshold = mcut_threshold(gen_array)

    gen_labels = [probs[i] for i in labels.general]
    gen_labels = dict([x for x in gen_labels if x[1] > gen_threshold])
    gen_labels = dict(sorted(gen_labels.items(), key=lambda item: item[1], reverse=True))

    # Character labels, pick any where prediction confidence > threshold
    if char_use_mcut:
        char_array = np.array([probs[i][1] for i in labels.character])
        char_threshold = round(mcut_threshold(char_array), 2)

    char_labels = [probs[i] for i in labels.character]
    char_labels = dict([x for x in char_labels if x[1] > char_threshold])
    char_labels = dict(sorted(char_labels.items(), key=lambda item: item[1], reverse=True))

    # Combine general and character labels, sort by confidence
    combined_names = [x for x in gen_labels]
    combined_names.extend([x for x in char_labels])

    # Convert to a string suitable for use as a training caption
    caption = ", ".join(combined_names)
    booru = caption.replace("_", " ").replace("(", "\(").replace(")", "\)")

    return image, caption, booru, rating_labels, char_labels, char_threshold, gen_labels, gen_threshold


css = """
#gen_mcut, #char_mcut {
    padding-top: var(--scale-3);
}
#gen_threshold.dimmed, #char_threshold.dimmed {
    filter: brightness(75%);
}
"""

with gr.Blocks(theme="NoCrypt/miku", analytics_enabled=False, title=TITLE, css=css) as demo:
    with gr.Row(equal_height=False):
        with gr.Column(min_width=720):
            with gr.Group():
                img_input = gr.Image(
                    label="Input",
                    type="pil",
                    image_mode="RGB",
                    sources=["upload", "clipboard"],
                )
                show_processed = gr.Checkbox(label="Show Preprocessed Image", value=False)
            with gr.Row():
                version = gr.Radio(
                    choices=list(MODEL_VARIANTS.keys()),
                    label="Model Version",
                    value="v3",
                    min_width=160,
                    scale=1,
                )  # gen_threshold > div.wrap.hide
                variant = gr.Radio(
                    choices=list(MODEL_VARIANTS[version.value].keys()),
                    label="Model Variant",
                    value="SwinV2",
                    min_width=560,
                )
            with gr.Group():
                with gr.Row():
                    gen_threshold = gr.Slider(
                        minimum=0.0,
                        maximum=1.0,
                        value=0.35,
                        step=0.01,
                        label="General Tag Threshold",
                        scale=5,
                        elem_id="gen_threshold",
                    )
                    gen_mcut = gr.Checkbox(label="Use Max-Cut", value=False, scale=1, elem_id="gen_mcut")
                with gr.Row():
                    char_threshold = gr.Slider(
                        minimum=0.0,
                        maximum=1.0,
                        value=0.85,
                        step=0.01,
                        label="Character Tag Threshold",
                        scale=5,
                        elem_id="char_threshold",
                    )
                    char_mcut = gr.Checkbox(label="Use Max-Cut", value=False, scale=1, elem_id="char_mcut")
            with gr.Row():
                clear = gr.ClearButton(
                    components=[],
                    variant="secondary",
                    size="lg",
                )
                submit = gr.Button(value="Submit", variant="primary", size="lg")

        with gr.Column(min_width=720):
            img_output = gr.Image(
                label="Preprocessed Image", type="pil", image_mode="RGB", scale=1, visible=False
            )
            with gr.Group():
                caption = gr.Textbox(label="Caption", show_copy_button=True)
                tags = gr.Textbox(label="Tags", show_copy_button=True)
            with gr.Group():
                rating = gr.Label(label="Rating")
            with gr.Group():
                char_mcut_out = gr.Number(label="Max-Cut Threshold", precision=2, visible=False)
                character = gr.Label(label="Character")
            with gr.Group():
                gen_mcut_out = gr.Number(label="Max-Cut Threshold", precision=2, visible=False)
                general = gr.Label(label="General")

    with gr.Row():
        examples = [[imgpath, 0.35, mc, 0.85, mc] for mc in [False, True] for imgpath in example_images]

        examples = gr.Examples(
            examples=examples,
            inputs=[img_input, gen_threshold, gen_mcut, char_threshold, char_mcut],
        )

    # tell clear button which components to clear
    clear.add([img_input, img_output, caption, rating, character, general])

    def on_select_variant(evt: gr.SelectData, variant: str):
        if evt.selected:
            choices = list(MODEL_VARIANTS[variant])
            return gr.update(choices=choices, value=choices[0])
        return gr.update()

    version.select(on_select_variant, inputs=[version], outputs=[variant])

    # show/hide processed image
    def on_change_show(val: gr.Checkbox):
        return gr.update(visible=val)

    show_processed.select(on_change_show, inputs=[show_processed], outputs=[img_output])

    # handle mcut thresholding (auto-calculate threshold from probs, disable slider)
    def on_change_mcut(val: gr.Checkbox):
        return (
            gr.update(interactive=not val, elem_classes=["dimmed"] if val else []),
            gr.update(visible=val),
        )

    gen_mcut.change(on_change_mcut, inputs=[gen_mcut], outputs=[gen_threshold, gen_mcut_out])
    char_mcut.change(on_change_mcut, inputs=[char_mcut], outputs=[char_threshold, char_mcut_out])

    submit.click(
        predict,
        inputs=[img_input, version, variant, gen_threshold, gen_mcut, char_threshold, char_mcut],
        outputs=[img_output, caption, tags, rating, character, char_threshold, general, gen_threshold],
        api_name="predict",
    )

if __name__ == "__main__":
    demo.queue(max_size=10)
    if getenv("SPACE_ID", None) is not None:
        demo.launch()
    else:
        demo.launch(
            server_name="0.0.0.0",
            server_port=7871,
        )