diff --git a/comfyui-impact-pack.py b/comfyui-impact-pack.py index edf917e..7089d4f 100644 --- a/comfyui-impact-pack.py +++ b/comfyui-impact-pack.py @@ -63,7 +63,10 @@ from mmdet.apis import (inference_detector, init_detector) from PIL import Image +import model_management + def load_mmdet(model_path): + print(model_management.vram_state) model_config = os.path.splitext(model_path)[0] + ".py" model = init_detector(model_config, model_path, device="cpu") return model @@ -154,13 +157,19 @@ folder_paths.folder_names_and_paths["mmdets_bbox"] = ([os.path.join(model_path, folder_paths.folder_names_and_paths["mmdets_segm"] = ([os.path.join(model_path, "mmdets", "segm")], folder_paths.supported_pt_extensions) folder_paths.folder_names_and_paths["mmdets"] = ([os.path.join(model_path, "mmdets")], folder_paths.supported_pt_extensions) +class NO_BBOX_MODEL: + ERROR = "" + +class NO_SEGM_MODEL: + ERROR = "" + class MMDetLoader: @classmethod def INPUT_TYPES(s): bboxs = [ "bbox/"+x for x in folder_paths.get_filename_list("mmdets_bbox") ] segms = [ "segm/"+x for x in folder_paths.get_filename_list("mmdets_segm") ] return {"required": { "model_name": (bboxs + segms, )}} - RETURN_TYPES = ("MODEL", ) + RETURN_TYPES = ("BBOX_MODEL", "SEGM_MODEL") FUNCTION = "load_mmdet" CATEGORY = "ImpactPack" @@ -168,7 +177,11 @@ class MMDetLoader: def load_mmdet(self, model_name): mmdet_path = folder_paths.get_full_path("mmdets", model_name) model = load_mmdet(mmdet_path) - return (model, ) + + if model_name.startswith("bbox"): + return (model, NO_SEGM_MODEL()) + else: + return (NO_BBOX_MODEL(), model) class SegmDetector: input_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "input") @@ -176,7 +189,7 @@ class SegmDetector: def INPUT_TYPES(s): return {"required": { - "model": ("MODEL", ), + "segm_model": ("SEGM_MODEL", ), "image": ("IMAGE", ), "threshold": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), } @@ -187,8 +200,8 @@ class SegmDetector: CATEGORY = "ImpactPack" - def doit(self, model, image, threshold): - mmdet_results = inference_segm(model, image, threshold) + def doit(self, segm_model, image, threshold): + mmdet_results = inference_segm(segm_model, image, threshold) segmasks = create_segmasks(mmdet_results) mask = combine_masks(segmasks) @@ -196,10 +209,22 @@ class SegmDetector: return (image,) class BboxDetector(SegmDetector): - def doit(self, model, image, threshold): - mmdet_results = inference_bbox(model, image, threshold) + @classmethod + def INPUT_TYPES(s): + return {"required": + { + "bbox_model": ("BBOX_MODEL", ), + "image": ("IMAGE", ), + "threshold": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), + } + } + + + def doit(self, bbox_model, image, threshold): + mmdet_results = inference_bbox(bbox_model, image, threshold) segmasks = create_segmasks(mmdet_results) mask = combine_masks(segmasks) + image = pil2tensor(mask) return (image,) diff --git a/misc/example.json b/misc/example.json index be5a9d3..7d5c745 100644 --- a/misc/example.json +++ b/misc/example.json @@ -1,153 +1,7 @@ { - "last_node_id": 9, - "last_link_id": 10, + "last_node_id": 21, + "last_link_id": 23, "nodes": [ - { - "id": 7, - "type": "MMDetLoader", - "pos": [ - 515, - 771 - ], - "size": { - "0": 443, - "1": 58 - }, - "flags": {}, - "order": 0, - "mode": 0, - "outputs": [ - { - "name": "MODEL", - "type": "MODEL", - "links": [ - 7 - ], - "slot_index": 0 - } - ], - "properties": { - "Node name for S&R": "MMDetLoader" - }, - "widgets_values": [ - "segm/mmdet_dd-person_mask2former.pth" - ] - }, - { - "id": 2, - "type": "MMDetLoader", - "pos": [ - 515, - 180 - ], - "size": { - "0": 439, - "1": 62 - }, - "flags": { - "collapsed": false - }, - "order": 1, - "mode": 0, - "outputs": [ - { - "name": "MODEL", - "type": "MODEL", - "links": [ - 4 - ], - "slot_index": 0 - } - ], - "properties": { - "Node name for S&R": "MMDetLoader" - }, - "widgets_values": [ - "bbox/mmdet_anime-face_yolov3.pth" - ], - "shape": 2 - }, - { - "id": 6, - "type": "SegmDetector", - "pos": [ - 178, - 765 - ], - "size": { - "0": 210, - "1": 78 - }, - "flags": {}, - "order": 4, - "mode": 0, - "inputs": [ - { - "name": "model", - "type": "MODEL", - "link": 7 - }, - { - "name": "image", - "type": "IMAGE", - "link": 8 - } - ], - "outputs": [ - { - "name": "IMAGE", - "type": "IMAGE", - "links": [ - 10 - ], - "slot_index": 0 - } - ], - "properties": { - "Node name for S&R": "SegmDetector" - }, - "widgets_values": [ - 0.49999999999999956 - ] - }, - { - "id": 1, - "type": "LoadImage", - "pos": [ - 510, - 314 - ], - "size": { - "0": 448, - "1": 387 - }, - "flags": {}, - "order": 2, - "mode": 0, - "outputs": [ - { - "name": "IMAGE", - "type": "IMAGE", - "links": [ - 6, - 8 - ], - "slot_index": 0 - }, - { - "name": "MASK", - "type": "MASK", - "links": null - } - ], - "properties": { - "Node name for S&R": "LoadImage" - }, - "widgets_values": [ - "snh.jpg", - "image" - ] - }, { "id": 4, "type": "PreviewImage", @@ -166,7 +20,7 @@ { "name": "images", "type": "IMAGE", - "link": 5 + "link": 20 } ], "title": "PreviewImage - Faces", @@ -174,49 +28,6 @@ "Node name for S&R": "PreviewImage" } }, - { - "id": 5, - "type": "BboxDetector", - "pos": [ - 1048, - 180 - ], - "size": { - "0": 210, - "1": 78 - }, - "flags": {}, - "order": 3, - "mode": 0, - "inputs": [ - { - "name": "model", - "type": "MODEL", - "link": 4 - }, - { - "name": "image", - "type": "IMAGE", - "link": 6 - } - ], - "outputs": [ - { - "name": "IMAGE", - "type": "IMAGE", - "links": [ - 5 - ], - "slot_index": 0 - } - ], - "properties": { - "Node name for S&R": "BboxDetector" - }, - "widgets_values": [ - 0.49999999999999956 - ] - }, { "id": 9, "type": "PreviewImage", @@ -235,59 +46,256 @@ { "name": "images", "type": "IMAGE", - "link": 10 + "link": 23 } ], "title": "PreviewImage - Persons", "properties": { "Node name for S&R": "PreviewImage" } + }, + { + "id": 17, + "type": "BboxDetector", + "pos": [ + 1017, + 136 + ], + "size": { + "0": 315, + "1": 78 + }, + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "bbox_model", + "type": "BBOX_MODEL", + "link": 18 + }, + { + "name": "image", + "type": "IMAGE", + "link": 21 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 20 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "BboxDetector" + }, + "widgets_values": [ + 0.5 + ] + }, + { + "id": 1, + "type": "LoadImage", + "pos": [ + 510, + 314 + ], + "size": { + "0": 448, + "1": 387 + }, + "flags": {}, + "order": 0, + "mode": 0, + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 21, + 22 + ], + "slot_index": 0 + }, + { + "name": "MASK", + "type": "MASK", + "links": null + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "snh.jpg", + "image" + ] + }, + { + "id": 21, + "type": "SegmDetector", + "pos": [ + 122, + 810 + ], + "size": { + "0": 315, + "1": 78 + }, + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "segm_model", + "type": "SEGM_MODEL", + "link": 19 + }, + { + "name": "image", + "type": "IMAGE", + "link": 22 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 23 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "SegmDetector" + }, + "widgets_values": [ + 0.5 + ] + }, + { + "id": 18, + "type": "MMDetLoader", + "pos": [ + 511, + 135 + ], + "size": { + "0": 437, + "1": 87 + }, + "flags": {}, + "order": 1, + "mode": 0, + "outputs": [ + { + "name": "BBOX_MODEL", + "type": "BBOX_MODEL", + "links": [ + 18 + ], + "slot_index": 0 + }, + { + "name": "SEGM_MODEL", + "type": "SEGM_MODEL", + "links": null + } + ], + "properties": { + "Node name for S&R": "MMDetLoader" + }, + "widgets_values": [ + "bbox/mmdet_anime-face_yolov3.pth" + ] + }, + { + "id": 20, + "type": "MMDetLoader", + "pos": [ + 523, + 801 + ], + "size": { + "0": 437, + "1": 87 + }, + "flags": {}, + "order": 2, + "mode": 0, + "outputs": [ + { + "name": "BBOX_MODEL", + "type": "BBOX_MODEL", + "links": [], + "slot_index": 0 + }, + { + "name": "SEGM_MODEL", + "type": "SEGM_MODEL", + "links": [ + 19 + ], + "slot_index": 1 + } + ], + "properties": { + "Node name for S&R": "MMDetLoader" + }, + "widgets_values": [ + "segm/mmdet_dd-person_mask2former.pth" + ] } ], "links": [ [ - 4, - 2, + 18, + 18, 0, - 5, + 17, 0, - "MODEL" + "BBOX_MODEL" ], [ - 5, - 5, + 19, + 20, + 1, + 21, + 0, + "SEGM_MODEL" + ], + [ + 20, + 17, 0, 4, 0, "IMAGE" ], [ - 6, + 21, 1, 0, - 5, + 17, 1, "IMAGE" ], [ - 7, - 7, - 0, - 6, - 0, - "MODEL" - ], - [ - 8, + 22, 1, 0, - 6, + 21, 1, "IMAGE" ], [ - 10, - 6, + 23, + 21, 0, 9, 0, diff --git a/misc/example.png b/misc/example.png index 37c5175..f8bf3ca 100644 Binary files a/misc/example.png and b/misc/example.png differ