Spaces:
Runtime error
Runtime error
| #!/usr/bin/env python | |
| from __future__ import annotations | |
| import argparse | |
| import functools | |
| import html | |
| import os | |
| import gradio as gr | |
| import huggingface_hub | |
| import numpy as np | |
| import onnxruntime as rt | |
| import pandas as pd | |
| import piexif | |
| import piexif.helper | |
| import PIL.Image | |
| from Utils import dbimutils | |
| TITLE = "WaifuDiffusion v1.4 Tags" | |
| DESCRIPTION = """ | |
| Demo for [SmilingWolf/wd-v1-4-vit-tagger](https://huggingface.co/SmilingWolf/wd-v1-4-vit-tagger) and [SmilingWolf/wd-v1-4-convnext-tagger](https://huggingface.co/SmilingWolf/wd-v1-4-convnext-tagger) with "ready to copy" prompt and a prompt analyzer. | |
| Modified from [NoCrypt/DeepDanbooru_string](https://huggingface.co/spaces/NoCrypt/DeepDanbooru_string) | |
| Modified from [hysts/DeepDanbooru](https://huggingface.co/spaces/hysts/DeepDanbooru) | |
| PNG Info code forked from [AUTOMATIC1111/stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui) | |
| Example image by [ほし☆☆☆](https://www.pixiv.net/en/users/43565085) | |
| """ | |
| HF_TOKEN = os.environ["HF_TOKEN"] | |
| VIT_MODEL_REPO = "SmilingWolf/wd-v1-4-vit-tagger" | |
| CONV_MODEL_REPO = "SmilingWolf/wd-v1-4-convnext-tagger" | |
| MODEL_FILENAME = "model.onnx" | |
| LABEL_FILENAME = "selected_tags.csv" | |
| def parse_args() -> argparse.Namespace: | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--score-slider-step", type=float, default=0.05) | |
| parser.add_argument("--score-threshold", type=float, default=0.35) | |
| parser.add_argument("--share", action="store_true") | |
| return parser.parse_args() | |
| def load_model(model_repo: str, model_filename: str) -> rt.InferenceSession: | |
| path = huggingface_hub.hf_hub_download( | |
| model_repo, model_filename, use_auth_token=HF_TOKEN | |
| ) | |
| model = rt.InferenceSession(path) | |
| return model | |
| def load_labels() -> list[str]: | |
| path = huggingface_hub.hf_hub_download( | |
| VIT_MODEL_REPO, LABEL_FILENAME, use_auth_token=HF_TOKEN | |
| ) | |
| df = pd.read_csv(path)["name"].tolist() | |
| return df | |
| def plaintext_to_html(text): | |
| text = ( | |
| "<p>" + "<br>\n".join([f"{html.escape(x)}" for x in text.split("\n")]) + "</p>" | |
| ) | |
| return text | |
| def predict( | |
| image: PIL.Image.Image, | |
| selected_model: str, | |
| score_threshold: float, | |
| models: dict, | |
| labels: list[str], | |
| ): | |
| rawimage = image | |
| model = models[selected_model] | |
| _, height, width, _ = model.get_inputs()[0].shape | |
| # Alpha to white | |
| image = image.convert("RGBA") | |
| new_image = PIL.Image.new("RGBA", image.size, "WHITE") | |
| new_image.paste(image, mask=image) | |
| image = new_image.convert("RGB") | |
| image = np.asarray(image) | |
| # PIL RGB to OpenCV BGR | |
| image = image[:, :, ::-1] | |
| image = dbimutils.make_square(image, height) | |
| image = dbimutils.smart_resize(image, height) | |
| image = image.astype(np.float32) | |
| image = np.expand_dims(image, 0) | |
| input_name = model.get_inputs()[0].name | |
| label_name = model.get_outputs()[0].name | |
| probs = model.run([label_name], {input_name: image})[0] | |
| labels = list(zip(labels, probs[0].astype(float))) | |
| # First 4 labels are actually ratings: pick one with argmax | |
| ratings_names = labels[:4] | |
| rating = dict(ratings_names) | |
| # Everything else is tags: pick any where prediction confidence > threshold | |
| tags_names = labels[4:] | |
| res = [x for x in tags_names if x[1] > score_threshold] | |
| res = dict(res) | |
| b = dict(sorted(res.items(), key=lambda item: item[1], reverse=True)) | |
| a = ( | |
| ", ".join(list(b.keys())) | |
| .replace("_", " ") | |
| .replace("(", "\(") | |
| .replace(")", "\)") | |
| ) | |
| c = ", ".join(list(b.keys())) | |
| items = rawimage.info | |
| geninfo = "" | |
| if "exif" in rawimage.info: | |
| exif = piexif.load(rawimage.info["exif"]) | |
| exif_comment = (exif or {}).get("Exif", {}).get(piexif.ExifIFD.UserComment, b"") | |
| try: | |
| exif_comment = piexif.helper.UserComment.load(exif_comment) | |
| except ValueError: | |
| exif_comment = exif_comment.decode("utf8", errors="ignore") | |
| items["exif comment"] = exif_comment | |
| geninfo = exif_comment | |
| for field in [ | |
| "jfif", | |
| "jfif_version", | |
| "jfif_unit", | |
| "jfif_density", | |
| "dpi", | |
| "exif", | |
| "loop", | |
| "background", | |
| "timestamp", | |
| "duration", | |
| ]: | |
| items.pop(field, None) | |
| geninfo = items.get("parameters", geninfo) | |
| info = f""" | |
| <p><h4>PNG Info</h4></p> | |
| """ | |
| for key, text in items.items(): | |
| info += ( | |
| f""" | |
| <div> | |
| <p><b>{plaintext_to_html(str(key))}</b></p> | |
| <p>{plaintext_to_html(str(text))}</p> | |
| </div> | |
| """.strip() | |
| + "\n" | |
| ) | |
| if len(info) == 0: | |
| message = "Nothing found in the image." | |
| info = f"<div><p>{message}<p></div>" | |
| return (a, c, rating, res, info) | |
| def main(): | |
| args = parse_args() | |
| vit_model = load_model(VIT_MODEL_REPO, MODEL_FILENAME) | |
| conv_model = load_model(CONV_MODEL_REPO, MODEL_FILENAME) | |
| labels = load_labels() | |
| models = {"ViT": vit_model, "ConvNext": conv_model} | |
| func = functools.partial(predict, models=models, labels=labels) | |
| gr.Interface( | |
| fn=func, | |
| inputs=[ | |
| gr.Image(type="pil", label="Input"), | |
| gr.Radio(["ViT", "ConvNext"], value="ViT", label="Model"), | |
| gr.Slider( | |
| 0, | |
| 1, | |
| step=args.score_slider_step, | |
| value=args.score_threshold, | |
| label="Score Threshold", | |
| ), | |
| ], | |
| outputs=[ | |
| gr.Textbox(label="Output (string)"), | |
| gr.Textbox(label="Output (raw string)"), | |
| gr.Label(label="Rating"), | |
| gr.Label(label="Output (label)"), | |
| gr.HTML(), | |
| ], | |
| examples=[["power.jpg", "ViT", 0.5]], | |
| title=TITLE, | |
| description=DESCRIPTION, | |
| allow_flagging="never", | |
| ).launch( | |
| enable_queue=True, | |
| share=args.share, | |
| ) | |
| if __name__ == "__main__": | |
| main() | |