From f16192bfa786ee3046097f5536b104e39cbf6f2b Mon Sep 17 00:00:00 2001 From: chflame163 Date: Fri, 30 Jan 2026 14:31:25 +0800 Subject: [PATCH] Ultra nodes support vitmatte-base-composition-1k --- py/ben_ultra.py | 4 ++-- py/birefnet_ultra.py | 4 ++-- py/birefnet_ultra_v2.py | 4 ++-- py/evf_sam_ultra.py | 4 ++-- py/florence2_ultra.py | 4 ++-- py/human_parts_ultra.py | 4 ++-- py/imagefunc.py | 23 ++++++++++++++++++++--- py/person_mask_ultra_v2.py | 4 ++-- py/sam_2_ultra.py | 12 ++++++------ py/segment_anything_ultra_v2.py | 4 ++-- py/segment_anything_ultra_v3.py | 4 ++-- py/transparent_background_ultra.py | 4 ++-- pyproject.toml | 2 +- 13 files changed, 47 insertions(+), 30 deletions(-) diff --git a/py/ben_ultra.py b/py/ben_ultra.py index e5222af..c55b3c7 100644 --- a/py/ben_ultra.py +++ b/py/ben_ultra.py @@ -991,7 +991,7 @@ class LS_BenUltra: @classmethod def INPUT_TYPES(cls): - method_list = ['VITMatte', 'VITMatte(local)', 'PyMatting', 'GuidedFilter', ] + method_list = ['VITMatte', 'VITMatte(local)', 'vitmatte-base-composition-1k', 'PyMatting', 'GuidedFilter', ] return { "required": { "ben_model": ("BEN_MODEL",), @@ -1044,7 +1044,7 @@ class LS_BenUltra: _mask = tensor2pil(mask_edge_detail(i, _mask, detail_range // 8 + 1, black_point, white_point)) else: _trimap = generate_VITMatte_trimap(_mask, detail_erode, detail_dilate) - _mask = generate_VITMatte(orig_image, _trimap, local_files_only=local_files_only, device=device, max_megapixels=max_megapixels) + _mask = generate_VITMatte(orig_image, _trimap, local_files_only=local_files_only, device=device, max_megapixels=max_megapixels, method=detail_method) _mask = tensor2pil(histogram_remap(pil2tensor(_mask), black_point, white_point)) else: _mask = tensor2pil(_mask) diff --git a/py/birefnet_ultra.py b/py/birefnet_ultra.py index 8fbcfc3..190ecd7 100644 --- a/py/birefnet_ultra.py +++ b/py/birefnet_ultra.py @@ -9,7 +9,7 @@ class BiRefNetUltra: @classmethod def INPUT_TYPES(cls): - method_list = ['VITMatte', 'VITMatte(local)', 'PyMatting', 'GuidedFilter', ] + method_list = ['VITMatte', 'VITMatte(local)', 'vitmatte-base-composition-1k', 'PyMatting', 'GuidedFilter', ] device_list = ['cuda','cpu'] return { "required": { @@ -62,7 +62,7 @@ class BiRefNetUltra: _mask = tensor2pil(mask_edge_detail(i, _mask, detail_range // 8 + 1, black_point, white_point)) else: _trimap = generate_VITMatte_trimap(_mask, detail_erode, detail_dilate) - _mask = generate_VITMatte(orig_image, _trimap, local_files_only=local_files_only, device=device, max_megapixels=max_megapixels) + _mask = generate_VITMatte(orig_image, _trimap, local_files_only=local_files_only, device=device, max_megapixels=max_megapixels, method=detail_method) _mask = tensor2pil(histogram_remap(pil2tensor(_mask), black_point, white_point)) else: _mask = tensor2pil(_mask) diff --git a/py/birefnet_ultra_v2.py b/py/birefnet_ultra_v2.py index 596ebc4..5614e01 100644 --- a/py/birefnet_ultra_v2.py +++ b/py/birefnet_ultra_v2.py @@ -112,7 +112,7 @@ class LS_BiRefNetUltraV2: @classmethod def INPUT_TYPES(cls): - method_list = ['VITMatte', 'VITMatte(local)', 'PyMatting', 'GuidedFilter', ] + method_list = ['VITMatte', 'VITMatte(local)', 'vitmatte-base-composition-1k', 'PyMatting', 'GuidedFilter', ] device_list = ['cuda', 'cpu'] return { "required": { @@ -187,7 +187,7 @@ class LS_BiRefNetUltraV2: _mask = tensor2pil(mask_edge_detail(i, _mask, detail_range // 8 + 1, black_point, white_point)) else: _trimap = generate_VITMatte_trimap(_mask, detail_erode, detail_dilate) - _mask = generate_VITMatte(orig_image, _trimap, local_files_only=local_files_only, device=device, max_megapixels=max_megapixels) + _mask = generate_VITMatte(orig_image, _trimap, local_files_only=local_files_only, device=device, max_megapixels=max_megapixels, method=detail_method) _mask = tensor2pil(histogram_remap(pil2tensor(_mask), black_point, white_point)) else: _mask = tensor2pil(_mask) diff --git a/py/evf_sam_ultra.py b/py/evf_sam_ultra.py index 72e1da6..1c4354d 100644 --- a/py/evf_sam_ultra.py +++ b/py/evf_sam_ultra.py @@ -20,7 +20,7 @@ class EVF_SAM_Ultra: model_list = ["evf-sam2", "evf-sam"] precision_list = ["fp16", "bf16", "fp32"] load_in_bit_list = ["full", "8", "4"] - method_list = ['VITMatte', 'VITMatte(local)', 'PyMatting', 'GuidedFilter', ] + method_list = ['VITMatte', 'VITMatte(local)', 'vitmatte-base-composition-1k', 'PyMatting', 'GuidedFilter', ] device_list = ['cuda', 'cpu'] return {"required": { @@ -98,7 +98,7 @@ class EVF_SAM_Ultra: else: _trimap = generate_VITMatte_trimap(_mask, detail_erode, detail_dilate) _mask = generate_VITMatte(orig_image, _trimap, local_files_only=local_files_only, device=device, - max_megapixels=max_megapixels) + max_megapixels=max_megapixels, method=detail_method) _mask = tensor2pil(histogram_remap(pil2tensor(_mask), black_point, white_point)) else: _mask = mask2image(_mask) diff --git a/py/florence2_ultra.py b/py/florence2_ultra.py index e42fbfd..a6374b3 100644 --- a/py/florence2_ultra.py +++ b/py/florence2_ultra.py @@ -457,7 +457,7 @@ class Florence2Ultra: "referring expression segmentation", "open vocabulary detection", ] - method_list = ['VITMatte', 'VITMatte(local)', 'PyMatting', 'GuidedFilter', ] + method_list = ['VITMatte', 'VITMatte(local)', 'vitmatte-base-composition-1k', 'PyMatting', 'GuidedFilter', ] device_list = ['cuda','cpu'] return { "required": { @@ -522,7 +522,7 @@ class Florence2Ultra: _mask = tensor2pil(mask_edge_detail(i, _mask, detail_range // 8 + 1, black_point, white_point)) else: _trimap = generate_VITMatte_trimap(_mask, detail_erode, detail_dilate) - _mask = generate_VITMatte(img, _trimap, local_files_only=local_files_only, device=device, max_megapixels=max_megapixels) + _mask = generate_VITMatte(img, _trimap, local_files_only=local_files_only, device=device, max_megapixels=max_megapixels, method=detail_method) _mask = tensor2pil(histogram_remap(pil2tensor(_mask), black_point, white_point)) else: _mask = tensor2pil(_mask) diff --git a/py/human_parts_ultra.py b/py/human_parts_ultra.py index 6bc7f0f..cbe433c 100644 --- a/py/human_parts_ultra.py +++ b/py/human_parts_ultra.py @@ -34,7 +34,7 @@ class LS_HumanPartsUltra: @classmethod def INPUT_TYPES(cls): - method_list = ['VITMatte', 'VITMatte(local)', 'PyMatting', 'GuidedFilter', ] + method_list = ['VITMatte', 'VITMatte(local)', 'vitmatte-base-composition-1k', 'PyMatting', 'GuidedFilter', ] device_list = ['cuda', 'cpu'] return { "required": { @@ -104,7 +104,7 @@ class LS_HumanPartsUltra: else: _trimap = generate_VITMatte_trimap(_mask, detail_erode, detail_dilate) _mask = generate_VITMatte(orig_image, _trimap, local_files_only=local_files_only, device=device, - max_megapixels=max_megapixels) + max_megapixels=max_megapixels, method=detail_method) _mask = tensor2pil(histogram_remap(pil2tensor(_mask), black_point, white_point)) else: _mask = mask2image(_mask) diff --git a/py/imagefunc.py b/py/imagefunc.py index a55178e..5209ec7 100644 --- a/py/imagefunc.py +++ b/py/imagefunc.py @@ -1569,7 +1569,19 @@ def load_VITMatte_model(model_name:str, local_files_only:bool=False) -> object: vitmatte = VITMatteModel(model, processor) return vitmatte -def generate_VITMatte(image:Image, trimap:Image, local_files_only:bool=False, device:str="cpu", max_megapixels:float=2.0) -> Image: + +def load_VITMatte_base_model(model_name:str, local_files_only:bool=False) -> object: + model_name = "vitmatte-base-composition-1k" + model_repo = "hustvl/vitmatte-base-composition-1k" + model_path = check_and_download_model(model_name, model_repo) + from transformers import VitMatteImageProcessor, VitMatteForImageMatting + model = VitMatteForImageMatting.from_pretrained(model_path, local_files_only=local_files_only) + processor = VitMatteImageProcessor.from_pretrained(model_path, local_files_only=local_files_only) + vitmatte = VITMatteModel(model, processor) + return vitmatte + +def generate_VITMatte(image:Image, trimap:Image, local_files_only:bool=False, device:str="cpu", + max_megapixels:float=2.0, method:str="VITMatte") -> Image: if image.mode != 'RGB': image = image.convert('RGB') if trimap.mode != 'L': @@ -1585,7 +1597,7 @@ def generate_VITMatte(image:Image, trimap:Image, local_files_only:bool=False, de image = image.resize((target_width, target_height), Image.BILINEAR) trimap = trimap.resize((target_width, target_height), Image.BILINEAR) # log(f"vitmatte image size {width}x{height} too large, resize to {target_width}x{target_height} for processing.") - model_name = "hustvl/vitmatte-small-composition-1k" + if device=="cpu": device = torch.device('cpu') else: @@ -1594,7 +1606,12 @@ def generate_VITMatte(image:Image, trimap:Image, local_files_only:bool=False, de else: log("vitmatte device is set to cuda, but not available, using cpu instead.") device = torch.device('cpu') - vit_matte_model = load_VITMatte_model(model_name=model_name, local_files_only=local_files_only) + if method == "vitmatte-base-composition-1k": + model_name = "hustvl/vitmatte-base-composition-1k" + vit_matte_model = load_VITMatte_base_model(model_name=model_name, local_files_only=local_files_only) + else: + model_name = "hustvl/vitmatte-small-composition-1k" + vit_matte_model = load_VITMatte_model(model_name=model_name, local_files_only=local_files_only) vit_matte_model.model.to(device) # log(f"vitmatte processing, image size = {image.width}x{image.height}, device = {device}.") inputs = vit_matte_model.processor(images=image, trimaps=trimap, return_tensors="pt") diff --git a/py/person_mask_ultra_v2.py b/py/person_mask_ultra_v2.py index a725842..bc60d62 100644 --- a/py/person_mask_ultra_v2.py +++ b/py/person_mask_ultra_v2.py @@ -19,7 +19,7 @@ class PersonMaskUltraV2: @classmethod def INPUT_TYPES(self): - method_list = ['VITMatte', 'VITMatte(local)', 'PyMatting', 'GuidedFilter', ] + method_list = ['VITMatte', 'VITMatte(local)', 'vitmatte-base-composition-1k', 'PyMatting', 'GuidedFilter', ] device_list = ['cuda','cpu'] return { "required": @@ -162,7 +162,7 @@ class PersonMaskUltraV2: detail_range // 8 + 1, black_point, white_point)) else: _trimap = generate_VITMatte_trimap(_mask, detail_erode, detail_dilate) - _mask = generate_VITMatte(orig_image, _trimap, local_files_only=local_files_only, device=device, max_megapixels=max_megapixels) + _mask = generate_VITMatte(orig_image, _trimap, local_files_only=local_files_only, device=device, max_megapixels=max_megapixels, method=detail_method) _mask = tensor2pil(histogram_remap(pil2tensor(_mask), black_point, white_point)) else: _mask = mask2image(_mask) diff --git a/py/sam_2_ultra.py b/py/sam_2_ultra.py index cad9da3..7060127 100644 --- a/py/sam_2_ultra.py +++ b/py/sam_2_ultra.py @@ -230,7 +230,7 @@ class LS_SAM2_ULTRA: ] model_precision_list = [ 'fp16','bf16','fp32'] select_list = ["all", "first", "by_index"] - method_list = ['VITMatte', 'VITMatte(local)', 'PyMatting', 'GuidedFilter', ] + method_list = ['VITMatte', 'VITMatte(local)', 'vitmatte-base-composition-1k', 'PyMatting', 'GuidedFilter', ] device_list = ['cuda','cpu'] return { "required": { @@ -419,7 +419,7 @@ class LS_SAM2_ULTRA: else: _trimap = generate_VITMatte_trimap(_mask, detail_erode, detail_dilate) _mask = generate_VITMatte(orig_image, _trimap, local_files_only=local_files_only, device=device, - max_megapixels=max_megapixels) + max_megapixels=max_megapixels, method=detail_method) _mask = tensor2pil(histogram_remap(pil2tensor(_mask), black_point, white_point)) else: _mask = tensor2pil(_mask) @@ -552,7 +552,7 @@ class LS_SAM2_ULTRA_V2: model_precision_list = [ 'fp16','bf16','fp32'] select_list = ["all", "first", "by_index"] - method_list = ['VITMatte', 'VITMatte(local)', 'PyMatting', 'GuidedFilter', ] + method_list = ['VITMatte', 'VITMatte(local)', 'vitmatte-base-composition-1k', 'PyMatting', 'GuidedFilter', ] device_list = ['cuda','cpu'] return { "required": { @@ -689,7 +689,7 @@ class LS_SAM2_ULTRA_V2: else: _trimap = generate_VITMatte_trimap(_mask, detail_erode, detail_dilate) _mask = generate_VITMatte(orig_image, _trimap, local_files_only=local_files_only, device=device, - max_megapixels=max_megapixels) + max_megapixels=max_megapixels, method=detail_method) _mask = tensor2pil(histogram_remap(pil2tensor(_mask), black_point, white_point)) else: _mask = tensor2pil(_mask) @@ -751,7 +751,7 @@ class LS_SAM2_VIDEO_ULTRA: 'sam2.1_hiera_tiny.safetensors', ] model_precision_list = ['fp16','bf16'] - method_list = ['VITMatte'] + method_list = ['VITMatte', 'vitmatte-base-composition-1k', ] device_list = ['cuda'] return { "required": { @@ -1029,7 +1029,7 @@ class LS_SAM2_VIDEO_ULTRA: if process_detail: _trimap = generate_VITMatte_trimap(pil2tensor(_mask), detail_erode, detail_dilate) _mask = generate_VITMatte(orig_image, _trimap, local_files_only=local_files_only, device=device, - max_megapixels=max_megapixels) + max_megapixels=max_megapixels, method=detail_method) _mask = tensor2pil(histogram_remap(pil2tensor(_mask), black_point, white_point)) color_image = Image.new("RGB", orig_image.size,color=mask_preview_color) diff --git a/py/segment_anything_ultra_v2.py b/py/segment_anything_ultra_v2.py index 9494f72..8adf0a7 100644 --- a/py/segment_anything_ultra_v2.py +++ b/py/segment_anything_ultra_v2.py @@ -19,7 +19,7 @@ class SegmentAnythingUltraV2: @classmethod def INPUT_TYPES(cls): - method_list = ['VITMatte', 'VITMatte(local)', 'PyMatting', 'GuidedFilter', ] + method_list = ['VITMatte', 'VITMatte(local)', 'vitmatte-base-composition-1k', 'PyMatting', 'GuidedFilter', ] device_list = ['cuda','cpu'] return { "required": { @@ -88,7 +88,7 @@ class SegmentAnythingUltraV2: _mask = tensor2pil(mask_edge_detail(i, _mask, detail_range // 8 + 1, black_point, white_point)) else: _trimap = generate_VITMatte_trimap(_mask, detail_erode, detail_dilate) - _mask = generate_VITMatte(_image, _trimap, local_files_only=local_files_only, device=device, max_megapixels=max_megapixels) + _mask = generate_VITMatte(_image, _trimap, local_files_only=local_files_only, device=device, max_megapixels=max_megapixels, method=detail_method) _mask = tensor2pil(histogram_remap(pil2tensor(_mask), black_point, white_point)) else: _mask = mask2image(_mask) diff --git a/py/segment_anything_ultra_v3.py b/py/segment_anything_ultra_v3.py index 8d7bbb8..d2fdcd1 100644 --- a/py/segment_anything_ultra_v3.py +++ b/py/segment_anything_ultra_v3.py @@ -39,7 +39,7 @@ class LS_SegmentAnythingUltraV3: @classmethod def INPUT_TYPES(cls): - method_list = ['VITMatte', 'VITMatte(local)', 'PyMatting', 'GuidedFilter', ] + method_list = ['VITMatte', 'VITMatte(local)', 'vitmatte-base-composition-1k', 'PyMatting', 'GuidedFilter', ] device_list = ['cuda','cpu'] return { "required": { @@ -102,7 +102,7 @@ class LS_SegmentAnythingUltraV3: _mask = tensor2pil(mask_edge_detail(i, _mask, detail_range // 8 + 1, black_point, white_point)) else: _trimap = generate_VITMatte_trimap(_mask, detail_erode, detail_dilate) - _mask = generate_VITMatte(_image, _trimap, local_files_only=local_files_only, device=device, max_megapixels=max_megapixels) + _mask = generate_VITMatte(_image, _trimap, local_files_only=local_files_only, device=device, max_megapixels=max_megapixels, method=detail_method) _mask = tensor2pil(histogram_remap(pil2tensor(_mask), black_point, white_point)) else: _mask = mask2image(_mask) diff --git a/py/transparent_background_ultra.py b/py/transparent_background_ultra.py index b850bd0..ea14025 100644 --- a/py/transparent_background_ultra.py +++ b/py/transparent_background_ultra.py @@ -19,7 +19,7 @@ class TransparentBackgroundUltra: @classmethod def INPUT_TYPES(cls): - method_list = ['VITMatte', 'VITMatte(local)', 'PyMatting', 'GuidedFilter', ] + method_list = ['VITMatte', 'VITMatte(local)', 'vitmatte-base-composition-1k', 'PyMatting', 'GuidedFilter', ] device_list = ['cuda','cpu'] return { @@ -78,7 +78,7 @@ class TransparentBackgroundUltra: _mask = tensor2pil(mask_edge_detail(i, _mask, detail_range // 8 + 1, black_point, white_point)) else: _trimap = generate_VITMatte_trimap(_mask, detail_erode, detail_dilate) - _mask = generate_VITMatte(orig_image, _trimap, local_files_only=local_files_only, device=device, max_megapixels=max_megapixels) + _mask = generate_VITMatte(orig_image, _trimap, local_files_only=local_files_only, device=device, max_megapixels=max_megapixels, method=detail_method) _mask = tensor2pil(histogram_remap(pil2tensor(_mask), black_point, white_point)) ret_image = RGB2RGBA(orig_image, _mask.convert('L')) diff --git a/pyproject.toml b/pyproject.toml index fad37d9..aba3865 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "ComfyUI_LayerStyle_Advance" description = "The nodes detached from ComfyUI Layer Style are mainly those with complex requirements for dependency packages." -version = "2.0.35" +version = "2.0.36" license = { text = "MIT License" } dependencies = ["numpy", "matplotlib", "scikit_image", "scikit_learn", "opencv-contrib-python", "pymatting", "timm", "blend_modes", "transformers", "diffusers", "loguru", "colour-science", "huggingface_hub", "segment_anything", "addict", "omegaconf", "yapf", "wget", "iopath", "mediapipe", "typer_config", "fastapi", "rich", "google-generativeai", "ultralytics", "transparent-background", "accelerate", "onnxruntime", "bitsandbytes", "peft", "protobuf", "hydra-core", "blind-watermark", "qrcode", "pyzbar", "psd-tools", "wandb", "zhipuai", "openai","google-genai", "fastapi","typer-config"]