Upload 9 files

.gitattributes

@@ -1,3 +1,4 @@
 /web/assets/** linguist-generated
 /web/** linguist-vendored
 comfy/text_encoders/t5_pile_tokenizer/tokenizer.model filter=lfs diff=lfs merge=lfs -text
+custom_nodes/comfyui-WD14-Tagger/models/wd-v1-4-moat-tagger-v2.onnx filter=lfs diff=lfs merge=lfs -text

custom_nodes/comfyui-WD14-Tagger/README.md

@@ -0,0 +1,49 @@
+# ComfyUI WD 1.4 Tagger
+
+A [ComfyUI](https://github.com/comfyanonymous/ComfyUI) extension allowing the interrogation of booru tags from images.
+
+Based on [SmilingWolf/wd-v1-4-tags](https://huggingface.co/spaces/SmilingWolf/wd-v1-4-tags) and [toriato/stable-diffusion-webui-wd14-tagger](https://github.com/toriato/stable-diffusion-webui-wd14-tagger)  
+All models created by [SmilingWolf](https://huggingface.co/SmilingWolf)
+
+## Installation
+1. `git clone https://github.com/pythongosssss/ComfyUI-WD14-Tagger` into the `custom_nodes` folder 
+    - e.g. `custom_nodes\ComfyUI-WD14-Tagger`  
+2. Open a Command Prompt/Terminal/etc
+3. Change to the `custom_nodes\ComfyUI-WD14-Tagger` folder you just created 
+    - e.g. `cd C:\ComfyUI_windows_portable\ComfyUI\custom_nodes\ComfyUI-WD14-Tagger` or wherever you have it installed
+4.  Install python packages
+      - **Windows Standalone installation** (embedded python):   
+       `../../../python_embeded/python.exe -s -m pip install -r requirements.txt`  
+      - **Manual/non-Windows installation**   
+        `pip install -r requirements.txt`
+
+## Usage
+Add the node via `image` -> `WD14Tagger|pysssss`  
+![image](https://github.com/pythongosssss/ComfyUI-WD14-Tagger/assets/125205205/ee6756ae-73f6-4e9f-a3da-eb87a056eb87)  
+Models are automatically downloaded at runtime if missing.  
+![image](https://github.com/pythongosssss/ComfyUI-WD14-Tagger/assets/125205205/cc09ae71-1a38-44da-afec-90f470a4b47d)  
+Supports tagging and outputting multiple batched inputs.  
+- **model**: The interrogation model to use. You can try them out here [WaifuDiffusion v1.4 Tags](https://huggingface.co/spaces/SmilingWolf/wd-v1-4-tags). The newest model (as of writing) is `MOAT` and the most popular is `ConvNextV2`.  
+- **threshold**: The score for the tag to be considered valid
+- **character_threshold**: The score for the character tag to be considered valid
+- **exclude_tags** A comma separated list of tags that should not be included in the results
+
+Quick interrogation of images is also available on any node that is displaying an image, e.g. a `LoadImage`, `SaveImage`, `PreviewImage` node.  
+Simply right click on the node (or if displaying multiple images, on the image you want to interrogate) and select `WD14 Tagger` from the menu  
+![image](https://github.com/pythongosssss/ComfyUI-WD14-Tagger/assets/125205205/11733899-6163-49f6-a22b-8dd86d910de6)
+
+Settings used for this are in the `settings` section of `pysssss.json`.
+
+### Offline Use
+Simplest way is to use it online, interrogate an image, and the model will be downloaded and cached, however if you want to manually download the models:
+- Create a `models` folder (in same folder as the `wd14tagger.py`)
+- Use URLs for models from the list in `pysssss.json`
+- Download `model.onnx` and name it with the model name e.g. `wd-v1-4-convnext-tagger-v2.onnx`
+- Download `selected_tags.csv` and name it with the model name e.g. `wd-v1-4-convnext-tagger-v2.csv`
+
+## Requirements
+`onnxruntime` (recommended, interrogation is still fast on CPU, included in requirements.txt)  
+or `onnxruntime-gpu` (allows use of GPU, many people have issues with this, if you try I can't provide support for this)
+
+## Changelog
+- 2023-05-14 - Moved to own repo, add downloading models, support multiple inputs

custom_nodes/comfyui-WD14-Tagger/__init__.py

@@ -0,0 +1,6 @@
+from .pysssss import init
+
+if init(check_imports=["onnxruntime"]):
+    from .wd14tagger import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+    WEB_DIRECTORY = "./web"
+    __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]

custom_nodes/comfyui-WD14-Tagger/models/wd-v1-4-moat-tagger-v2.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b8cef913be4c9e8d93f9f903e74271416502ce0b4b04df0ff1e2f00df488aa03
+size 326197340

custom_nodes/comfyui-WD14-Tagger/pyproject.toml

@@ -0,0 +1,14 @@
+[project]
+name = "comfyui-wd14-tagger"
+description = "A ComfyUI extension allowing the interrogation of booru tags from images."
+version = "1.0.1"
+license = { file = "LICENSE" }
+dependencies = ["onnxruntime"]
+
+[project.urls]
+Repository = "https://github.com/pythongosssss/ComfyUI-WD14-Tagger"
+
+[tool.comfy]
+PublisherId = "pythongosssss"
+DisplayName = "ComfyUI-WD14-Tagger"
+Icon = ""

custom_nodes/comfyui-WD14-Tagger/pysssss.json

@@ -0,0 +1,25 @@
+{
+	"name": "WD14Tagger",
+	"logging": false,
+	"settings": {
+		"model": "wd-v1-4-moat-tagger-v2",
+		"threshold": 0.35,
+		"character_threshold": 0.85,
+		"exclude_tags": "",
+		"ortProviders": ["CUDAExecutionProvider", "CPUExecutionProvider"],
+		"HF_ENDPOINT": "https://huggingface.co"
+	},
+	"models": {
+		"wd-eva02-large-tagger-v3": "{HF_ENDPOINT}/SmilingWolf/wd-eva02-large-tagger-v3",
+		"wd-vit-tagger-v3": "{HF_ENDPOINT}/SmilingWolf/wd-vit-tagger-v3",
+		"wd-swinv2-tagger-v3": "{HF_ENDPOINT}/SmilingWolf/wd-swinv2-tagger-v3",
+		"wd-convnext-tagger-v3": "{HF_ENDPOINT}/SmilingWolf/wd-convnext-tagger-v3",
+		"wd-v1-4-moat-tagger-v2": "{HF_ENDPOINT}/SmilingWolf/wd-v1-4-moat-tagger-v2",
+		"wd-v1-4-convnextv2-tagger-v2": "{HF_ENDPOINT}/SmilingWolf/wd-v1-4-convnextv2-tagger-v2",
+		"wd-v1-4-convnext-tagger-v2": "{HF_ENDPOINT}/SmilingWolf/wd-v1-4-convnext-tagger-v2",
+		"wd-v1-4-convnext-tagger": "{HF_ENDPOINT}/SmilingWolf/wd-v1-4-convnext-tagger",
+		"wd-v1-4-vit-tagger-v2": "{HF_ENDPOINT}/SmilingWolf/wd-v1-4-vit-tagger-v2",
+		"wd-v1-4-swinv2-tagger-v2": "{HF_ENDPOINT}/SmilingWolf/wd-v1-4-swinv2-tagger-v2",
+		"wd-v1-4-vit-tagger": "{HF_ENDPOINT}/SmilingWolf/wd-v1-4-vit-tagger"
+	}
+}

custom_nodes/comfyui-WD14-Tagger/pysssss.py

@@ -0,0 +1,242 @@
+import asyncio
+import os
+import json
+import shutil
+import inspect
+import aiohttp
+from server import PromptServer
+from tqdm import tqdm
+
+config = None
+
+
+def is_logging_enabled():
+    config = get_extension_config()
+    if "logging" not in config:
+        return False
+    return config["logging"]
+
+
+def log(message, type=None, always=False):
+    if not always and not is_logging_enabled():
+        return
+
+    if type is not None:
+        message = f"[{type}] {message}"
+
+    name = get_extension_config()["name"]
+
+    print(f"(pysssss:{name}) {message}")
+
+
+def get_ext_dir(subpath=None, mkdir=False):
+    dir = os.path.dirname(__file__)
+    if subpath is not None:
+        dir = os.path.join(dir, subpath)
+
+    dir = os.path.abspath(dir)
+
+    if mkdir and not os.path.exists(dir):
+        os.makedirs(dir)
+    return dir
+
+
+def get_comfy_dir(subpath=None):
+    dir = os.path.dirname(inspect.getfile(PromptServer))
+    if subpath is not None:
+        dir = os.path.join(dir, subpath)
+
+    dir = os.path.abspath(dir)
+
+    return dir
+
+
+def get_web_ext_dir():
+    config = get_extension_config()
+    name = config["name"]
+    dir = get_comfy_dir("web/extensions/pysssss")
+    if not os.path.exists(dir):
+        os.makedirs(dir)
+    dir += "/" + name
+    return dir
+
+
+def get_extension_config(reload=False):
+    global config
+    if reload == False and config is not None:
+        return config
+
+    config_path = get_ext_dir("pysssss.user.json")
+    if not os.path.exists(config_path):
+        config_path = get_ext_dir("pysssss.json")
+
+    if not os.path.exists(config_path):
+        log("Missing pysssss.json and pysssss.user.json, this extension may not work correctly. Please reinstall the extension.",
+            type="ERROR", always=True)
+        print(f"Extension path: {get_ext_dir()}")
+        return {"name": "Unknown", "version": -1}
+    with open(config_path, "r") as f:
+        config = json.loads(f.read())
+    return config
+
+def link_js(src, dst):
+    src = os.path.abspath(src)
+    dst = os.path.abspath(dst)
+    if os.name == "nt":
+        try:
+            import _winapi
+            _winapi.CreateJunction(src, dst)
+            return True
+        except:
+            pass
+    try:
+        os.symlink(src, dst)
+        return True
+    except:
+        import logging
+        logging.exception('')
+        return False
+
+
+def is_junction(path):
+    if os.name != "nt":
+        return False
+    try:
+        return bool(os.readlink(path))
+    except OSError:
+        return False
+
+def install_js():
+    src_dir = get_ext_dir("web/js")
+    if not os.path.exists(src_dir):
+        log("No JS")
+        return
+
+    should_install = should_install_js()
+    if should_install:
+        log("it looks like you're running an old version of ComfyUI that requires manual setup of web files, it is recommended you update your installation.", "warning", True)
+    dst_dir = get_web_ext_dir()
+    linked = os.path.islink(dst_dir) or is_junction(dst_dir)
+    if linked or os.path.exists(dst_dir):
+        if linked:
+            if should_install:
+                log("JS already linked")
+            else:
+                os.unlink(dst_dir)
+                log("JS unlinked, PromptServer will serve extension")
+        elif not should_install:
+            shutil.rmtree(dst_dir)
+            log("JS deleted, PromptServer will serve extension")
+        return
+    
+    if not should_install:
+        log("JS skipped, PromptServer will serve extension")
+        return
+    
+    if link_js(src_dir, dst_dir):
+        log("JS linked")
+        return
+
+    log("Copying JS files")
+    shutil.copytree(src_dir, dst_dir, dirs_exist_ok=True)
+
+
+def should_install_js():
+    return not hasattr(PromptServer.instance, "supports") or "custom_nodes_from_web" not in PromptServer.instance.supports
+
+
+def init(check_imports):
+    log("Init")
+
+    if check_imports is not None:
+        import importlib.util
+        for imp in check_imports:
+            spec = importlib.util.find_spec(imp)
+            if spec is None:
+                log(f"{imp} is required, please check requirements are installed.", type="ERROR", always=True)
+                return False
+
+    install_js()
+    return True
+
+
+async def download_to_file(url, destination, update_callback, is_ext_subpath=True, session=None):
+    close_session = False
+    if session is None:
+        close_session = True
+        loop = None
+        try:
+            loop = asyncio.get_event_loop()
+        except:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+
+        session = aiohttp.ClientSession(loop=loop)
+    if is_ext_subpath:
+        destination = get_ext_dir(destination)
+    try:
+        proxy = os.getenv("HTTP_PROXY") or os.getenv("http_proxy")
+        print("proxy:", proxy)
+        proxy_auth = None
+        if proxy:
+            proxy_auth = aiohttp.BasicAuth(os.getenv("PROXY_USER", ""), os.getenv("PROXY_PASS", ""))
+
+        async with session.get(url, proxy=proxy, proxy_auth=proxy_auth) as response:            
+            size = int(response.headers.get('content-length', 0)) or None
+
+            with tqdm(
+                unit='B', unit_scale=True, miniters=1, desc=url.split('/')[-1], total=size,
+            ) as progressbar:
+                with open(destination, mode='wb') as f:
+                    perc = 0
+                    async for chunk in response.content.iter_chunked(2048):
+                        f.write(chunk)
+                        progressbar.update(len(chunk))
+                        if update_callback is not None and progressbar.total is not None and progressbar.total != 0:
+                            last = perc
+                            perc = round(progressbar.n / progressbar.total, 2)
+                            if perc != last:
+                                last = perc
+                                await update_callback(perc)
+    finally:
+        if close_session and session is not None:
+            await session.close()
+
+def wait_for_async(async_fn):
+    try:
+        import concurrent.futures
+        # Check if we're in a running event loop
+        asyncio.get_running_loop()
+        # We're in a running loop, so run the async function in a separate thread
+        with concurrent.futures.ThreadPoolExecutor() as executor:
+            future = executor.submit(asyncio.run, async_fn())
+            return future.result()  # This blocks until complete
+    except RuntimeError:
+        # No running loop, safe to use asyncio.run()
+        return asyncio.run(async_fn())
+
+def update_node_status(client_id, node, text, progress=None):
+    if client_id is None:
+        client_id = PromptServer.instance.client_id
+
+    if client_id is None:
+        return
+
+    PromptServer.instance.send_sync("pysssss/update_status", {
+        "node": node,
+        "progress": progress,
+        "text": text
+    }, client_id)
+
+async def update_node_status_async(client_id, node, text, progress=None):
+    if client_id is None:
+        client_id = PromptServer.instance.client_id
+
+    if client_id is None:
+        return
+
+    await PromptServer.instance.send("pysssss/update_status", {
+        "node": node,
+        "progress": progress,
+        "text": text
+    }, client_id)

custom_nodes/comfyui-WD14-Tagger/requirements.txt

@@ -0,0 +1 @@
+onnxruntime

custom_nodes/comfyui-WD14-Tagger/wd14tagger.py

@@ -0,0 +1,210 @@
+# https://huggingface.co/spaces/SmilingWolf/wd-v1-4-tags
+
+import comfy.utils
+import asyncio
+import aiohttp
+import numpy as np
+import csv
+import os
+import sys
+import onnxruntime as ort
+from onnxruntime import InferenceSession
+from PIL import Image
+from server import PromptServer
+from aiohttp import web
+import folder_paths
+from .pysssss import get_ext_dir, get_comfy_dir, download_to_file, update_node_status, wait_for_async, get_extension_config, log
+sys.path.insert(0, os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy"))
+
+config = get_extension_config()
+
+defaults = {
+    "model": "wd-v1-4-moat-tagger-v2",
+    "threshold": 0.35,
+    "character_threshold": 0.85,
+    "replace_underscore": False,
+    "trailing_comma": False,
+    "exclude_tags": "",
+    "ortProviders": ["CUDAExecutionProvider", "CPUExecutionProvider"],
+    "HF_ENDPOINT": "https://huggingface.co"
+}
+defaults.update(config.get("settings", {}))
+
+if "wd14_tagger" in folder_paths.folder_names_and_paths:
+    models_dir = folder_paths.get_folder_paths("wd14_tagger")[0]
+    if not os.path.exists(models_dir):
+        os.makedirs(models_dir)
+else:
+    models_dir = get_ext_dir("models", mkdir=True)
+known_models = list(config["models"].keys())
+
+log("Available ORT providers: " + ", ".join(ort.get_available_providers()), "DEBUG", True)
+log("Using ORT providers: " + ", ".join(defaults["ortProviders"]), "DEBUG", True)
+
+def get_installed_models():
+    models = filter(lambda x: x.endswith(".onnx"), os.listdir(models_dir))
+    models = [m for m in models if os.path.exists(os.path.join(models_dir, os.path.splitext(m)[0] + ".csv"))]
+    return models
+
+
+async def tag(image, model_name, threshold=0.35, character_threshold=0.85, exclude_tags="", replace_underscore=True, trailing_comma=False, client_id=None, node=None):
+    if model_name.endswith(".onnx"):
+        model_name = model_name[0:-5]
+    installed = list(get_installed_models())
+    if not any(model_name + ".onnx" in s for s in installed):
+        await download_model(model_name, client_id, node)
+
+    name = os.path.join(models_dir, model_name + ".onnx")
+    model = InferenceSession(name, providers=defaults["ortProviders"])
+
+    input = model.get_inputs()[0]
+    height = input.shape[1]
+
+    # Reduce to max size and pad with white
+    ratio = float(height)/max(image.size)
+    new_size = tuple([int(x*ratio) for x in image.size])
+    image = image.resize(new_size, Image.LANCZOS)
+    square = Image.new("RGB", (height, height), (255, 255, 255))
+    square.paste(image, ((height-new_size[0])//2, (height-new_size[1])//2))
+
+    image = np.array(square).astype(np.float32)
+    image = image[:, :, ::-1]  # RGB -> BGR
+    image = np.expand_dims(image, 0)
+
+    # Read all tags from csv and locate start of each category
+    tags = []
+    general_index = None
+    character_index = None
+    with open(os.path.join(models_dir, model_name + ".csv")) as f:
+        reader = csv.reader(f)
+        next(reader)
+        for row in reader:
+            if general_index is None and row[2] == "0":
+                general_index = reader.line_num - 2
+            elif character_index is None and row[2] == "4":
+                character_index = reader.line_num - 2
+            if replace_underscore:
+                tags.append(row[1].replace("_", " "))
+            else:
+                tags.append(row[1])
+
+    label_name = model.get_outputs()[0].name
+    probs = model.run([label_name], {input.name: image})[0]
+
+    result = list(zip(tags, probs[0]))
+
+    # rating = max(result[:general_index], key=lambda x: x[1])
+    general = [item for item in result[general_index:character_index] if item[1] > threshold]
+    character = [item for item in result[character_index:] if item[1] > character_threshold]
+
+    all = character + general
+    remove = [s.strip() for s in exclude_tags.lower().split(",")]
+    all = [tag for tag in all if tag[0] not in remove]
+
+    res = ("" if trailing_comma else ", ").join((item[0].replace("(", "\\(").replace(")", "\\)") + (", " if trailing_comma else "") for item in all))
+
+    print(res)
+    return res
+
+
+async def download_model(model, client_id, node):
+    hf_endpoint = os.getenv("HF_ENDPOINT", defaults["HF_ENDPOINT"])
+    if not hf_endpoint.startswith("https://"):
+        hf_endpoint = f"https://{hf_endpoint}"
+    if hf_endpoint.endswith("/"):
+        hf_endpoint = hf_endpoint.rstrip("/")
+
+    url = config["models"][model]
+    url = url.replace("{HF_ENDPOINT}", hf_endpoint)
+    url = f"{url}/resolve/main/"
+    async with aiohttp.ClientSession(loop=asyncio.get_event_loop()) as session:
+        async def update_callback(perc):
+            nonlocal client_id
+            message = ""
+            if perc < 100:
+                message = f"Downloading {model}"
+            update_node_status(client_id, node, message, perc)
+
+        try:
+            await download_to_file(
+                f"{url}model.onnx", os.path.join(models_dir,f"{model}.onnx"), update_callback, session=session)
+            await download_to_file(
+                f"{url}selected_tags.csv", os.path.join(models_dir,f"{model}.csv"), update_callback, session=session)
+        except aiohttp.client_exceptions.ClientConnectorError as err:
+            log("Unable to download model. Download files manually or try using a HF mirror/proxy website by setting the environment variable HF_ENDPOINT=https://.....", "ERROR", True)
+            raise
+
+        update_node_status(client_id, node, None)
+
+    return web.Response(status=200)
+
+
+@PromptServer.instance.routes.get("/pysssss/wd14tagger/tag")
+async def get_tags(request):
+    if "filename" not in request.rel_url.query:
+        return web.Response(status=404)
+
+    type = request.query.get("type", "output")
+    if type not in ["output", "input", "temp"]:
+        return web.Response(status=400)
+
+    target_dir = get_comfy_dir(type)
+    image_path = os.path.abspath(os.path.join(
+        target_dir, request.query.get("subfolder", ""), request.query["filename"]))
+    c = os.path.commonpath((image_path, target_dir))
+    if os.path.commonpath((image_path, target_dir)) != target_dir:
+        return web.Response(status=403)
+
+    if not os.path.isfile(image_path):
+        return web.Response(status=404)
+
+    image = Image.open(image_path)
+
+    models = get_installed_models()
+    default = defaults["model"] + ".onnx"
+    model = default if default in models else models[0]
+
+    return web.json_response(await tag(image, model, client_id=request.rel_url.query.get("clientId", ""), node=request.rel_url.query.get("node", "")))
+
+
+class WD14Tagger:
+    @classmethod
+    def INPUT_TYPES(s):
+        extra = [name for name, _ in (os.path.splitext(m) for m in get_installed_models()) if name not in known_models]
+        models = known_models + extra
+        return {"required": {
+            "image": ("IMAGE", ),
+            "model": (models, { "default": defaults["model"] }),
+            "threshold": ("FLOAT", {"default": defaults["threshold"], "min": 0.0, "max": 1, "step": 0.05}),
+            "character_threshold": ("FLOAT", {"default": defaults["character_threshold"], "min": 0.0, "max": 1, "step": 0.05}),
+            "replace_underscore": ("BOOLEAN", {"default": defaults["replace_underscore"]}),
+            "trailing_comma": ("BOOLEAN", {"default": defaults["trailing_comma"]}),
+            "exclude_tags": ("STRING", {"default": defaults["exclude_tags"]}),
+        }}
+
+    RETURN_TYPES = ("STRING",)
+    OUTPUT_IS_LIST = (True,)
+    FUNCTION = "tag"
+    OUTPUT_NODE = True
+
+    CATEGORY = "image"
+
+    def tag(self, image, model, threshold, character_threshold, exclude_tags="", replace_underscore=False, trailing_comma=False):
+        tensor = image*255
+        tensor = np.array(tensor, dtype=np.uint8)
+
+        pbar = comfy.utils.ProgressBar(tensor.shape[0])
+        tags = []
+        for i in range(tensor.shape[0]):
+            image = Image.fromarray(tensor[i])
+            tags.append(wait_for_async(lambda: tag(image, model, threshold, character_threshold, exclude_tags, replace_underscore, trailing_comma)))
+            pbar.update(1)
+        return {"ui": {"tags": tags}, "result": (tags,)}
+
+
+NODE_CLASS_MAPPINGS = {
+    "WD14Tagger|pysssss": WD14Tagger,
+}
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "WD14Tagger|pysssss": "WD14 Tagger 🐍",
+}