From 9ab68757de6038093364feb7749794fa5ccb1104 Mon Sep 17 00:00:00 2001 From: chflame163 Date: Thu, 18 Jul 2024 15:33:33 +0800 Subject: [PATCH] fix vitmatte load models from local --- README.MD | 3 +++ README_CN.MD | 5 +++-- py/imagefunc.py | 2 ++ pyproject.toml | 2 +- 4 files changed, 9 insertions(+), 3 deletions(-) diff --git a/README.MD b/README.MD index f922877..b96e786 100644 --- a/README.MD +++ b/README.MD @@ -80,6 +80,7 @@ When this error has occurred, please check the network environment. ## Update **If the dependency package error after updating, please reinstall the relevant dependency packages.
+* Change the VitMatte model of the [Ultra](#Ultra) node to a local call. Please download [all files of vitmatte model](https://huggingface.co/hustvl/vitmatte-small-composition-1k/tree/main) to the ```ComfyUI/models/vitmatte``` folder. * [GetColorToneV2](#GetColorToneV2) node add the ```mask``` method to the color selection option, which can accurately obtain the main color and average color within the mask. * [ImageScaleByAspectRatioV2](#ImageScaleByAspectRatioV2) node add the "background_color" option. * [LUT Apply](#LUT) Add the "strength" option. @@ -1390,6 +1391,8 @@ Note: When running for the first time, you need to download the vitmate model fi After successfully downloading the model, you can use ```VITMatte(local)``` without accessing the network. * VitMatte's options: ```device``` set whether to use CUDA for vitimate operations, which is about 5 times faster than CPU. ```max_megapixels``` set the maximum image size for vitmate operation, and oversized images will be reduced in size. For 16G VRAM, it is recommended to set it to 3. +Download the vitmatte model files [all files of vitmatte model](https://huggingface.co/hustvl/vitmatte-small-composition-1k/tree/main) to the ```ComfyUI/models/vitmatte``` folder. + The following figure is an example of the difference in output between three methods. ![image](image/mask_edge_ultra_detail_v2_example.jpg) diff --git a/README_CN.MD b/README_CN.MD index fb9bc9a..01641fe 100644 --- a/README_CN.MD +++ b/README_CN.MD @@ -80,6 +80,7 @@ git clone https://github.com/chflame163/ComfyUI_LayerStyle.git ## 更新说明 **如果本插件更新后出现依赖包错误，请重新安装相关依赖包。 +* [Ultra](#Ultra) 节点的VitMatte模型改为本地调用，请下载[所有的vitmatte模型文件](https://huggingface.co/hustvl/vitmatte-small-composition-1k/tree/main)到```ComfyUI/models/vitmatte```文件夹。 * [GetColorToneV2](#GetColorToneV2) 节点的取色选项增加```mask```方法，可精确获取遮罩内的主色和平均色。 * [ImageScaleByAspectRatioV2](#ImageScaleByAspectRatioV2) 节点增加background_color选项。 * [LUT Apply](#LUT) 节点增加strenght选项。 @@ -1370,10 +1371,10 @@ mask为可选输入项，如果这里输入遮罩，将作用于输出结果。 * ```GuideFilter``` 使用 opencv guidedfilter 根据颜色相似度对边缘进行羽化，对于边缘具有很强的颜色分离时效果最佳。 以上两种方法的代码来着spacepxl的[ComfyUI-Image-Filters](https://github.com/spacepxl/ComfyUI-Image-Filters)的Alpha Matte节点，感谢原作者。 * ```VitMatte``` 使用transfromer vit模型进行高质量的边缘处理，保留边缘细节，甚至可以生成半透明遮罩。 -注：首次运行时需要下载vitmatte模型文件，等待自动下载完成即可。如果无法完成下载，可运行命令```huggingface-cli download hustvl/vitmatte-small-composition-1k```手动下载模型。 -模型成功下载之后可以使用VITMatte(local)无需访问网络。 * VitMatte的选项:```device``` 设置是否使用cuda进行vitmatte运算，cuda运算速度比cpu快5倍左右。```max_megapixels```设置vitmatte运算的最大图片尺寸，超大的图片将缩小处理。对于16G显存建议设置为3。 +请下载[所有的 vitmatte 模型文件](https://huggingface.co/hustvl/vitmatte-small-composition-1k/tree/main)到```ComfyUI/models/vitmatte```文件夹。 + 下图为三种方法输出区别的示例。 ![image](image/mask_edge_ultra_detail_v2_example.jpg) diff --git a/py/imagefunc.py b/py/imagefunc.py index 1fa46e7..61522fa 100644 --- a/py/imagefunc.py +++ b/py/imagefunc.py @@ -22,6 +22,7 @@ import scipy.ndimage import cv2 import random import time +from pathlib import Path from tqdm import tqdm from functools import lru_cache from typing import Union, List @@ -1490,6 +1491,7 @@ class VITMatteModel: self.processor = processor def load_VITMatte_model(model_name:str, local_files_only:bool=False) -> object: + model_name = Path(os.path.join(folder_paths.models_dir, "vitmatte")) from transformers import VitMatteImageProcessor, VitMatteForImageMatting model = VitMatteForImageMatting.from_pretrained(model_name, local_files_only=local_files_only) processor = VitMatteImageProcessor.from_pretrained(model_name, local_files_only=local_files_only) diff --git a/pyproject.toml b/pyproject.toml index 195042c..34e8a57 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui_layerstyle" description = "A set of nodes for ComfyUI it generate image like Adobe Photoshop's Layer Style. the Drop Shadow is first completed node, and follow-up work is in progress." -version = "1.0.14" +version = "1.0.15" license = "MIT" dependencies = ["numpy", "pillow", "torch", "matplotlib", "Scipy", "scikit_image", "opencv-contrib-python", "pymatting", "segment_anything", "timm", "addict", "yapf", "colour-science", "wget", "mediapipe", "loguru", "typer_config", "fastapi", "rich", "google-generativeai", "diffusers", "omegaconf", "tqdm", "transformers", "kornia", "image-reward", "ultralytics", "blend_modes", "blind-watermark", "qrcode", "pyzbar", "psd-tools"]