diff --git a/__init__.py b/__init__.py index ef16bd3..d8e82d9 100644 --- a/__init__.py +++ b/__init__.py @@ -36,7 +36,9 @@ init() node_list = [ "llm.llm_LP", - "vlm.llava_LP", + "vlm.autotagger_LP", + "vlm.llava_LP", + "vlm.ram_LP", ] NODE_CLASS_MAPPINGS = {} diff --git a/nodes/vlm/autotagger_LP.json b/nodes/vlm/autotagger_LP.json new file mode 100644 index 0000000..b3ac373 --- /dev/null +++ b/nodes/vlm/autotagger_LP.json @@ -0,0 +1,26 @@ +{ + "name": "Autotagger", + "logging": false, + "settings": { + "model": "wd-eva02-large-tagger-v3", + "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", + "Z3D-E621-Convnext": "{HF_ENDPOINT}/silveroxides/Z3D-E621-Convnext" + } +} diff --git a/nodes/vlm/autotagger_LP.py b/nodes/vlm/autotagger_LP.py new file mode 100644 index 0000000..0a4e767 --- /dev/null +++ b/nodes/vlm/autotagger_LP.py @@ -0,0 +1,429 @@ +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 +import json +import shutil +import inspect +from tqdm import tqdm + +sys.path.insert(0, os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy")) + +supported_autotaggers_extensions = set(['.onnx']) + +try: + folder_paths.folder_names_and_paths["autotaggers"] = (folder_paths.folder_names_and_paths["autotaggers"][0], supported_autotaggers_extensions) +except: + if not os.path.isdir(os.path.join(folder_paths.models_dir, "autotaggers")): + os.mkdir(os.path.join(folder_paths.models_dir, "autotaggers")) + + folder_paths.folder_names_and_paths["autotaggers"] = ([os.path.join(folder_paths.models_dir, "autotaggers")], supported_autotaggers_extensions) + +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"(LevelPixel Advanced:{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/levelpixeladvanced") + 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("autotagger_LP.user.json") + if not os.path.exists(config_path): + config_path = get_ext_dir("autotagger_LP.json") + + if not os.path.exists(config_path): + log("Missing autotagger_LP.json and autotagger_LP.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, loop=None): + 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("levelpixeladvanced/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("levelpixeladvanced/update_status", { + "node": node, + "progress": progress, + "text": text + }, client_id) + +config_autotagger = get_extension_config() + +defaults = { + "model": "wd-eva02-large-tagger-v3", + "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_autotagger.get("settings", {})) + +models_dir = folder_paths.get_folder_paths("autotaggers")[0] +if not os.path.exists(models_dir): + os.makedirs(models_dir) + +known_models = list(config_autotagger["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_autotagger["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("/levelpixeladvanced/autotagger/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 Autotagger: + @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 = False + + CATEGORY = "LevelPixel/VLM" + + 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 = { + "Autotagger|LP": Autotagger, +} +NODE_DISPLAY_NAME_MAPPINGS = { + "Autotagger|LP": "Autotagger [LP]", +} diff --git a/nodes/vlm/ram_LP.py b/nodes/vlm/ram_LP.py new file mode 100644 index 0000000..5eb7674 --- /dev/null +++ b/nodes/vlm/ram_LP.py @@ -0,0 +1,113 @@ +from PIL import Image +import torch +import gc +import numpy as np +from ram.models import ram +from ram.models import ram_plus +from ram.models import tag2text +from ram import inference_ram # as inference +from ram import inference_tag2text # as inference +from ram import get_transform +from folder_paths import models_dir, add_model_folder_path, get_folder_paths, get_filename_list, get_full_path +import numpy + +rams_dir = f"{models_dir}/rams" + +try: + if rams_dir not in get_folder_paths("rams"): + raise KeyError +except KeyError: + add_model_folder_path("rams", rams_dir) + +print(f"[Recognize Anything] RAM models found: {', '.join(get_filename_list('rams'))}") + + +class RecognizeAnything: + MODEL_NAMES = ["ram_swin_large_14m.pth", "ram_plus_swin_large_14m.pth", "tag2text_swin_14m.pth"] # other/newer models can be added here + DEVICES = ["cpu", "gpu"] if torch.cuda.is_available() else ["cpu"] + IMAGE_SIZES = [384, 448] + + def __init__(self): + self.model = None + self.modelname = "" + self.image_size = 384 + # self.inference = None + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "model": (s.MODEL_NAMES, {"default": s.MODEL_NAMES[1]},), + "device": (s.DEVICES, {"default": s.DEVICES[0]},), + "spec_tag2text": ("STRING", {"default": ""},), + # "tag2text_threshold": ("INT", {"min": 0, "max": 100, "default": 68},), + } + } + + RETURN_TYPES = ("STRING", "STRING", "STRING",) + RETURN_NAMES = ("tags", "spec_tags", "caption") + FUNCTION = "interrogate" + OUTPUT_NODE = False + CATEGORY = "LevelPixel/VLM" + + def interrogate(self, image:torch.Tensor, model:str, device:str, spec_tag2text:str): + dev = torch.device("cuda" if device.lower() == "gpu" else "cpu") + model_path = get_full_path("rams", model) + if not model_path: + raise NameError(f"Model '{model}' not found. Make sure it is in the '/models/rams' folder or add the path in 'extra_model_paths.yaml'") + self.transform = get_transform(image_size=RecognizeAnything.IMAGE_SIZES[0]) + + # delete some tags that may disturb captioning + # 127: "quarter"; 2961: "back", 3351: "two"; 3265: "three"; 3338: "four"; 3355: "five"; 3359: "one" + delete_tag_index = [127, 2961, 3351, 3265, 3338, 3355, 3359] + + if (self.model == None) or (self.modelname != model) or (device != self.device): + del self.model + gc.collect() + if (device == "cpu") and torch.cuda.is_available(): + torch.cuda.empty_cache() + self.model = None + print(f"[Recognize Anything]: loading model {model_path}, please stand by....") + + if model == RecognizeAnything.MODEL_NAMES[0]: + self.model = ram(pretrained=model_path, image_size=self.image_size, vit='swin_l') + elif model == RecognizeAnything.MODEL_NAMES[1]: + self.model = ram_plus(pretrained=model_path, image_size=self.image_size, vit='swin_l') + elif model == RecognizeAnything.MODEL_NAMES[2]: + self.model = tag2text(pretrained=model_path, image_size=self.image_size, vit='swin_b', delete_tag_index=delete_tag_index) + # self.model.threshold = tag2text_threshold # threshold for tagging + else: + raise ValueError('No valid model was selected', model) + + self.model.eval() + self.model = self.model.to(dev) + + self.modelname = model + self.device = device + + tag0 = tag1 = tag2 = '' + + for im in image: + i = 255. * im.cpu().numpy() + img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) + image = self.transform(img).unsqueeze(0).to(dev) + if model == RecognizeAnything.MODEL_NAMES[2]: + res = inference_tag2text(image, self.model, input_tag=spec_tag2text) + else: + res = inference_ram(image, self.model) + + # remove duplicates and reformat the string: + tag0 = ", ".join(list(numpy.unique(res[0].split(' | ')))) + '\n' + tag1 += ", ".join(list(numpy.unique(res[1].split(' | ')))) + '\n' if (len(res) > 1) and res[1] else '\n' + tag2 += res[2] + '\n' if (len(res) > 2) and res[2] else '\n' + + return(tag0, tag1, tag2,) + +NODE_CLASS_MAPPINGS = { + "Recognize Anything|LP": RecognizeAnything, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "Recognize Anything|LP": "Recognize Anything [LP]", +} \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 4cfee58..e2141b6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,7 @@ name = "comfyui-level-pixel-advanced" description = "Advanced nodes of the Level Pixel company. Includes convenient advanced nodes for working with LLM и VLM models (LLaVa) with GGUF format. Qwen2.5-VL and Qwen2.5 supported. Also included is a node for the RAM model. Nodes have the ability to automatically unload models from VRAM." version = "1.2.0" license = { file = "LICENSE" } -dependencies = ["torch>=2.0.1", "torchvision>=0.15.2", "transformers>=4.46", "pillow>=9.4.0", "numpy>=1.26.4", "matplotlib", "scikit-build-core>=0.10.7"] +dependencies = ["torch>=2.0.1", "torchvision>=0.15.2", "transformers>=4.46", "pillow>=9.4.0", "numpy>=1.26.4", "matplotlib", "scikit-build-core>=0.10.7", "onnxruntime-gpu>=1.20.0", "onnxruntime>=1.20.0"] [project.urls] Repository = "https://github.com/LevelPixel/ComfyUI-LevelPixel-Advanced" diff --git a/requirements.txt b/requirements.txt index a7a07ba..887f05d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,3 +4,6 @@ pillow>=9.4.0 numpy>=1.26.4 matplotlib scikit-build-core>=0.10.7 +onnxruntime-gpu>=1.20.0 +onnxruntime>=1.20.0 +ram @ git+https://github.com/xinyu1205/recognize-anything.git