diff --git a/__init__.py b/__init__.py index 63becd7..9a1513b 100644 --- a/__init__.py +++ b/__init__.py @@ -310,6 +310,8 @@ NODE_CLASS_MAPPINGS = { NODE_DISPLAY_NAME_MAPPINGS = { + "SAMLoader": "SAMLoader (Impact)", + "BboxDetectorSEGS": "BBOX Detector (SEGS)", "SegmDetectorSEGS": "SEGM Detector (SEGS)", "ONNXDetectorSEGS": "ONNX Detector (SEGS/legacy) - use BBOXDetector", diff --git a/modules/impact/config.py b/modules/impact/config.py index 6a7aeeb..3989298 100644 --- a/modules/impact/config.py +++ b/modules/impact/config.py @@ -2,7 +2,7 @@ import configparser import os -version = "V4.53" +version = "V4.53.1" dependency_version = 19 diff --git a/modules/impact/core.py b/modules/impact/core.py index 524eba2..d0942df 100644 --- a/modules/impact/core.py +++ b/modules/impact/core.py @@ -610,6 +610,7 @@ def make_sam_mask(sam_model, segs, image, detection_hint, dilation, else: mask = torch.zeros((8, 8), dtype=torch.float32, device="cpu") # empty mask + mask = utils.make_3d_mask(mask) return mask diff --git a/modules/impact/detectors.py b/modules/impact/detectors.py index fc1a569..82cdf84 100644 --- a/modules/impact/detectors.py +++ b/modules/impact/detectors.py @@ -203,8 +203,8 @@ class SimpleDetectorForEach: @staticmethod def detect(bbox_detector, image, bbox_threshold, bbox_dilation, crop_factor, drop_size, - sub_threshold, sub_dilation, sub_bbox_expansion, - sam_mask_hint_threshold, post_dilation=0, sam_model_opt=None, segm_detector_opt=None): + sub_threshold, sub_dilation, sub_bbox_expansion, + sam_mask_hint_threshold, post_dilation=0, sam_model_opt=None, segm_detector_opt=None): if len(image) > 1: raise Exception('[Impact Pack] ERROR: SimpleDetectorForEach does not allow image batches.\nPlease refer to https://github.com/ltdrdata/ComfyUI-extension-tutorials/blob/Main/ComfyUI-Impact-Pack/tutorial/batching-detailer.md for more information.') diff --git a/modules/impact/impact_pack.py b/modules/impact/impact_pack.py index 3eda053..fc0c321 100644 --- a/modules/impact/impact_pack.py +++ b/modules/impact/impact_pack.py @@ -81,9 +81,10 @@ class CLIPSegDetectorProvider: class SAMLoader: @classmethod def INPUT_TYPES(cls): + models = [x for x in folder_paths.get_filename_list("sams") if 'hq' not in x] return { "required": { - "model_name": (folder_paths.get_filename_list("sams"), ), + "model_name": (models, ), "device_mode": (["AUTO", "Prefer GPU", "CPU"],), } } @@ -1559,7 +1560,7 @@ class MasksToMaskList: def doit(self, masks): if masks is None: - empty_mask = torch.zeros((64,64), dtype=torch.float32, device="cpu") + empty_mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu") return ([empty_mask], ) res = [] @@ -1569,6 +1570,8 @@ class MasksToMaskList: print(f"mask len: {len(res)}") + res = [make_3d_mask(x) for x in res] + return (res, ) @@ -1589,23 +1592,20 @@ class MaskListToMaskBatch: def doit(self, mask): if len(mask) == 1: - if len(mask[0].shape) == 2: - mask = mask[0].unsqueeze(0) + mask = make_3d_mask(mask[0]) return (mask,) elif len(mask) > 1: - mask1 = mask[0] - if len(mask1.shape) == 2: - mask1 = mask1.unsqueeze(0) + mask1 = make_3d_mask(mask[0]) for mask2 in mask[1:]: - if len(mask2.shape) == 2: - mask2 = mask2.unsqueeze(0) + mask2 = make_3d_mask(mask2) if mask1.shape[1:] != mask2.shape[1:]: mask2 = comfy.utils.common_upscale(mask2.movedim(-1, 1), mask1.shape[2], mask1.shape[1], "lanczos", "center").movedim(1, -1) mask1 = torch.cat((mask1, mask2), dim=0) + return (mask1,) else: - empty_mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu").unsqueeze(0) + empty_mask = torch.zeros((1, 64, 64), dtype=torch.float32, device="cpu").unsqueeze(0) return (empty_mask,) diff --git a/modules/impact/segs_nodes.py b/modules/impact/segs_nodes.py index d21b76f..b3e26f7 100644 --- a/modules/impact/segs_nodes.py +++ b/modules/impact/segs_nodes.py @@ -643,6 +643,7 @@ class SEGSToMaskList: if len(masks) == 0: empty_mask = torch.zeros(segs[0], dtype=torch.float32, device="cpu") masks = [empty_mask] + masks = [utils.make_3d_mask(mask) for mask in masks] return (masks,) @@ -655,14 +656,14 @@ class SEGSToMaskBatch: } RETURN_TYPES = ("MASK",) - OUTPUT_IS_LIST = (True,) FUNCTION = "doit" CATEGORY = "ImpactPack/Util" def doit(self, segs): masks = core.segs_to_masklist(segs) - mask_batch = torch.stack(masks, dim=0) + masks = [utils.make_3d_mask(mask) for mask in masks] + mask_batch = torch.concat(masks) return (mask_batch,) @@ -826,7 +827,9 @@ class DilateMask: def doit(self, mask, dilation): mask = core.dilate_mask(mask.numpy(), dilation) - return (torch.from_numpy(mask), ) + mask = torch.from_numpy(mask) + mask = utils.make_3d_mask(mask) + return (mask, ) class GaussianBlurMask: @@ -987,7 +990,7 @@ class EmptySEGS: class SegsToCombinedMask: @classmethod def INPUT_TYPES(s): - return {"required": { "segs": ("SEGS",), } } + return {"required": {"segs": ("SEGS",), }} RETURN_TYPES = ("MASK",) FUNCTION = "doit" @@ -995,7 +998,9 @@ class SegsToCombinedMask: CATEGORY = "ImpactPack/Operation" def doit(self, segs): - return (core.segs_to_combined_mask(segs),) + mask = core.segs_to_combined_mask(segs) + mask = utils.make_3d_mask(mask) + return (mask,) class MediaPipeFaceMeshToSEGS: