From d94bef1ee5ed3656f5ff1bb2830a4ffd94f40935 Mon Sep 17 00:00:00 2001 From: chflame163 Date: Fri, 30 Jan 2026 14:31:45 +0800 Subject: [PATCH] Ultra nodes support vitmatte-base-composition-1k --- py/imagefunc.py | 23 ++++++++++++++++++++--- py/mask_edge_ultra_detail_v2.py | 4 ++-- py/mask_edge_ultra_detail_v3.py | 6 +++--- py/rmbg_ultra_v2.py | 4 ++-- py/segformer_ultra.py | 12 ++++++------ pyproject.toml | 2 +- 6 files changed, 34 insertions(+), 17 deletions(-) diff --git a/py/imagefunc.py b/py/imagefunc.py index 8f40d79..9a4b2c4 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/mask_edge_ultra_detail_v2.py b/py/mask_edge_ultra_detail_v2.py index f289d6d..f78bb38 100644 --- a/py/mask_edge_ultra_detail_v2.py +++ b/py/mask_edge_ultra_detail_v2.py @@ -12,7 +12,7 @@ class MaskEdgeUltraDetailV2: @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": { @@ -77,7 +77,7 @@ class MaskEdgeUltraDetailV2: _mask = tensor2pil(mask_edge_detail(_image, _mask, detail_range//8, black_point, white_point)) else: _trimap = generate_VITMatte_trimap(_mask, edge_erode, edte_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=method) _mask = tensor2pil(histogram_remap(pil2tensor(_mask), black_point, white_point)) ret_image = RGB2RGBA(orig_image, _mask.convert('L')) diff --git a/py/mask_edge_ultra_detail_v3.py b/py/mask_edge_ultra_detail_v3.py index 0e217e0..82747ab 100644 --- a/py/mask_edge_ultra_detail_v3.py +++ b/py/mask_edge_ultra_detail_v3.py @@ -12,7 +12,7 @@ class MaskEdgeUltraDetailV3: @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": { @@ -100,11 +100,11 @@ class MaskEdgeUltraDetailV3: processed_trimap = tensor2pil(mask_edge_detail(_image, _mask, trimap_detail_range//8, black_point, white_point)) else: _trimap = generate_VITMatte_trimap(_mask, mask_edge_erode, mask_edge_dilate) - processed_mask = generate_VITMatte(orig_image, _trimap, local_files_only=local_files_only, device=device, max_megapixels=max_megapixels) + processed_mask = generate_VITMatte(orig_image, _trimap, local_files_only=local_files_only, device=device, max_megapixels=max_megapixels, method=method ) processed_mask = tensor2pil(histogram_remap(pil2tensor(processed_mask), black_point, white_point)) if transparent_trimap is not None: _trimap = generate_VITMatte_trimap(_mask, transparent_trimap_erode, transparent_trimap_dilate) - processed_trimap = generate_VITMatte(orig_image, _trimap, local_files_only=local_files_only, device=device, max_megapixels=max_megapixels) + processed_trimap = generate_VITMatte(orig_image, _trimap, local_files_only=local_files_only, device=device, max_megapixels=max_megapixels, method=method) processed_trimap = tensor2pil(histogram_remap(pil2tensor(processed_trimap), black_point, white_point)) if transparent_trimap is not None: diff --git a/py/rmbg_ultra_v2.py b/py/rmbg_ultra_v2.py index 794831c..2b41f02 100644 --- a/py/rmbg_ultra_v2.py +++ b/py/rmbg_ultra_v2.py @@ -12,7 +12,7 @@ class RmBgUltraV2: @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": { @@ -61,7 +61,7 @@ class RmBgUltraV2: _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 = mask2image(_mask) diff --git a/py/segformer_ultra.py b/py/segformer_ultra.py index 62f7d5d..2740225 100644 --- a/py/segformer_ultra.py +++ b/py/segformer_ultra.py @@ -69,7 +69,7 @@ class Segformer_B2_Clothes: @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": { @@ -180,7 +180,7 @@ class Segformer_B2_Clothes: 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) @@ -461,7 +461,7 @@ class SegformerUltraV2: @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": { @@ -522,7 +522,7 @@ class SegformerUltraV2: 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) @@ -838,7 +838,7 @@ class LS_SegformerUltraV3: @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": { "image": ("IMAGE",), @@ -906,7 +906,7 @@ class LS_SegformerUltraV3: 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/pyproject.toml b/pyproject.toml index 7b9c746..a8806da 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 = "2.0.37" +version = "2.0.38" license = {text = "MIT License"} dependencies = ["numpy", "pillow", "torch", "matplotlib", "Scipy", "scikit_image", "scikit_learn", "opencv-contrib-python", "pymatting", "timm", "colour-science", "transformers", "blend_modes", "huggingface_hub", "loguru"]