diff --git a/__init__.py b/__init__.py index 86a6da7..ce26c4a 100644 --- a/__init__.py +++ b/__init__.py @@ -53,9 +53,11 @@ class F2S2GenerateMask: def INPUT_TYPES(cls): model_list = list(sam_model_to_config_map.keys()) model_list.sort() + device_list = ["cuda", "cpu"] return { "required": { "sam2_model": (model_list, {"default": "sam2_hiera_small.pt"}), + "device": (device_list,), "image": ("IMAGE",), "prompt": ("STRING", {"default": "subject"}), } @@ -66,14 +68,15 @@ class F2S2GenerateMask: FUNCTION = "_process_image" CATEGORY = "💃rDancer" - def _process_image(self, sam2_model: str, image: torch.Tensor, prompt: str = None): + def _process_image(self, sam2_model: str, device: str, image: torch.Tensor, prompt: str = None): + torch_device = torch.device(device) prompt = prompt.strip() if prompt else "" annotated_images, masks, masked_images = [], [], [] # Convert image from tensor to PIL # the image has an extra batch dimension, despite the variable name for img in image: img = tensor2pil(img).convert("RGB") - annotated_image, mask, masked_image = process_image(sam2_model, img, prompt) + annotated_image, mask, masked_image = process_image(torch_device, sam2_model, img, prompt) annotated_images.append(pil2tensor(annotated_image)) masks.append(pil2tensor(mask)) masked_images.append(pil2tensor(masked_image)) diff --git a/app.py b/app.py index 4ddfb1b..da41bcc 100644 --- a/app.py +++ b/app.py @@ -70,7 +70,7 @@ except ImportError: # VIDEO_TARGET_DIRECTORY = "tmp" # create_directory(directory_path=VIDEO_TARGET_DIRECTORY) -DEVICE = torch.device("cuda") +DEVICE = None #torch.device("cuda") # DEVICE = torch.device("cpu") torch.autocast(device_type="cuda", dtype=torch.bfloat16).__enter__() @@ -112,10 +112,16 @@ def annotate_image(image, detections): # gr.Textbox(visible=text == IMAGE_CAPTION_GROUNDING_MASKS_MODE), # ] -def lazy_load_models(sam_image_model: str): +def lazy_load_models(device: torch.device, sam_image_model: str): global SAM_IMAGE_MODEL global FLORENCE_MODEL global FLORENCE_PROCESSOR + global DEVICE + if device != DEVICE: + SAM_IMAGE_MODEL = None + FLORENCE_MODEL = None + FLORENCE_PROCESSOR = None + DEVICE = device if SAM_IMAGE_MODEL is None: SAM_IMAGE_MODEL = load_sam_image_model(device=DEVICE, checkpoint=sam_image_model) elif SAM_IMAGE_MODEL != sam_image_model: @@ -124,8 +130,8 @@ def lazy_load_models(sam_image_model: str): if FLORENCE_MODEL is None or FLORENCE_PROCESSOR is None: FLORENCE_MODEL, FLORENCE_PROCESSOR = load_florence_model(device=DEVICE) -def process_image(sam_image_model: str, image: Image.Image, promt: str) -> Tuple[Optional[Image.Image], Optional[Image.Image], Optional[Image.Image]]: - lazy_load_models(sam_image_model) +def process_image(device: torch.device, sam_image_model: str, image: Image.Image, promt: str) -> Tuple[Optional[Image.Image], Optional[Image.Image], Optional[Image.Image]]: + lazy_load_models(device, sam_image_model) annotated_image, mask_list = _process_image(IMAGE_OPEN_VOCABULARY_DETECTION_MODE, image, promt) if mask_list is not None and len(mask_list) > 0: mask = np.any(mask_list, axis=0) # Merge masks into a single mask diff --git a/workflows/workflow.json b/workflows/workflow.json index ca4833e..deb306c 100644 --- a/workflows/workflow.json +++ b/workflows/workflow.json @@ -1,14 +1,14 @@ { - "last_node_id": 52, - "last_link_id": 50, + "last_node_id": 58, + "last_link_id": 60, "nodes": [ { "id": 30, "type": "LoadImage", - "pos": [ - -241, - 63 - ], + "pos": { + "0": 360, + "1": 520 + }, "size": { "0": 278.3598327636719, "1": 400.1376647949219 @@ -16,6 +16,7 @@ "flags": {}, "order": 0, "mode": 0, + "inputs": [], "outputs": [ { "name": "IMAGE", @@ -23,8 +24,8 @@ "links": [ 28 ], - "shape": 3, - "slot_index": 0 + "slot_index": 0, + "shape": 3 }, { "name": "MASK", @@ -42,15 +43,136 @@ ] }, { - "id": 32, - "type": "LayerMask: MaskPreview", - "pos": [ - 553.8103742161225, - 63 - ], + "id": 36, + "type": "LoadImage", + "pos": { + "0": 650, + "1": 520 + }, "size": { - "0": 477.1796569824219, - "1": 385.8974914550781 + "0": 278.3598327636719, + "1": 400.1376647949219 + }, + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 29 + ], + "slot_index": 0, + "shape": 3 + }, + { + "name": "MASK", + "type": "MASK", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "Screenshot 2024-09-16 221234.png", + "image" + ] + }, + { + "id": 31, + "type": "PreviewImage", + "pos": { + "0": 360, + "1": 970 + }, + "size": { + "0": 471.4238586425781, + "1": 380.51617431640625 + }, + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 58 + } + ], + "outputs": [], + "properties": { + "Node name for S&R": "PreviewImage" + }, + "widgets_values": [] + }, + { + "id": 51, + "type": "PreviewImage", + "pos": { + "0": 1160, + "1": 970 + }, + "size": { + "0": 470.9900207519531, + "1": 380.91436767578125 + }, + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 57 + } + ], + "outputs": [], + "properties": { + "Node name for S&R": "PreviewImage" + }, + "widgets_values": [] + }, + { + "id": 58, + "type": "PreviewImage", + "pos": { + "0": 1160, + "1": 525 + }, + "size": { + "0": 470.9900207519531, + "1": 380.91436767578125 + }, + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 60 + } + ], + "outputs": [], + "properties": { + "Node name for S&R": "PreviewImage" + }, + "widgets_values": [] + }, + { + "id": 57, + "type": "MaskToImage", + "pos": { + "0": 860, + "1": 1320 + }, + "size": { + "0": 264.5999755859375, + "1": 26 }, "flags": {}, "order": 5, @@ -59,21 +181,32 @@ { "name": "mask", "type": "MASK", - "link": 49 + "link": 59 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 60 + ], + "slot_index": 0, + "shape": 3 } ], "properties": { - "Node name for S&R": "LayerMask: MaskPreview" + "Node name for S&R": "MaskToImage" }, - "color": "rgba(27, 80, 119, 0.7)" + "widgets_values": [] }, { "id": 37, "type": "ImageBatch", - "pos": [ - 263, - 661 - ], + "pos": { + "0": 880, + "1": 1010 + }, "size": { "0": 210, "1": 46 @@ -98,116 +231,28 @@ "name": "IMAGE", "type": "IMAGE", "links": [ - 47 + 56 ], - "shape": 3, - "slot_index": 0 - } - ], - "properties": { - "Node name for S&R": "ImageBatch" - } - }, - { - "id": 51, - "type": "PreviewImage", - "pos": [ - 560, - 510 - ], - "size": [ - 470.99003119854433, - 380.9143663249481 - ], - "flags": {}, - "order": 6, - "mode": 0, - "inputs": [ - { - "name": "images", - "type": "IMAGE", - "link": 50 - } - ], - "properties": { - "Node name for S&R": "PreviewImage" - } - }, - { - "id": 36, - "type": "LoadImage", - "pos": [ - 49, - 63 - ], - "size": { - "0": 278.3598327636719, - "1": 400.1376647949219 - }, - "flags": {}, - "order": 1, - "mode": 0, - "outputs": [ - { - "name": "IMAGE", - "type": "IMAGE", - "links": [ - 29 - ], - "shape": 3, - "slot_index": 0 - }, - { - "name": "MASK", - "type": "MASK", - "links": null, + "slot_index": 0, "shape": 3 } ], "properties": { - "Node name for S&R": "LoadImage" + "Node name for S&R": "ImageBatch" }, - "widgets_values": [ - "Screenshot 2024-09-16 221234.png", - "image" - ] + "widgets_values": [] }, { - "id": 31, - "type": "PreviewImage", - "pos": [ - -241, - 510 - ], - "size": { - "0": 471.4238586425781, - "1": 380.51617431640625 - }, - "flags": {}, - "order": 4, - "mode": 0, - "inputs": [ - { - "name": "images", - "type": "IMAGE", - "link": 48 - } - ], - "properties": { - "Node name for S&R": "PreviewImage" - } - }, - { - "id": 52, + "id": 56, "type": "RdancerFlorence2SAM2GenerateMask", - "pos": [ - 263, - 768.9143663249481 - ], - "size": [ - 285.0516880999385, - 122 - ], + "pos": { + "0": 840, + "1": 1120 + }, + "size": { + "0": 303.60308837890625, + "1": 146 + }, "flags": {}, "order": 3, "mode": 0, @@ -215,7 +260,7 @@ { "name": "image", "type": "IMAGE", - "link": 47 + "link": 56 } ], "outputs": [ @@ -223,28 +268,28 @@ "name": "annotated_image", "type": "IMAGE", "links": [ - 48 + 58 ], - "shape": 3, - "slot_index": 0 + "slot_index": 0, + "shape": 3 }, { "name": "mask", "type": "MASK", "links": [ - 49 + 59 ], - "shape": 3, - "slot_index": 1 + "slot_index": 1, + "shape": 3 }, { "name": "masked_image", "type": "IMAGE", "links": [ - 50 + 57 ], - "shape": 3, - "slot_index": 2 + "slot_index": 2, + "shape": 3 } ], "properties": { @@ -252,6 +297,7 @@ }, "widgets_values": [ "sam2_hiera_small.pt", + "cuda", "products" ] } @@ -274,34 +320,42 @@ "IMAGE" ], [ - 47, + 56, 37, 0, - 52, + 56, 0, "IMAGE" ], [ - 48, - 52, + 57, + 56, + 2, + 51, + 0, + "IMAGE" + ], + [ + 58, + 56, 0, 31, 0, "IMAGE" ], [ - 49, - 52, + 59, + 56, 1, - 32, + 57, 0, "MASK" ], [ - 50, - 52, - 2, - 51, + 60, + 57, + 0, + 58, 0, "IMAGE" ] @@ -310,11 +364,11 @@ "config": {}, "extra": { "ds": { - "scale": 1.486436280241443, - "offset": { - "0": 418.2350075396343, - "1": 89.96842480074564 - } + "scale": 1.2839025177495011, + "offset": [ + -279.5410202465874, + -342.9758340690436 + ] } }, "version": 0.4 diff --git a/workflows/workflow.png b/workflows/workflow.png index 076921d..d9568d6 100644 Binary files a/workflows/workflow.png and b/workflows/workflow.png differ