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, )