nodes/rdancer-ComfyUI_Florence2SAM2
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

allow to select device

Browse Source 
this should fix the problem of not being able to run the node on Windows
-- just select device: "cpu" in the dropdown
This commit is contained in:
Jan Minar
2024-10-04 01:27:36 +01:00
parent 2c89b43499
commit 7b191f63cf
4 changed files with 219 additions and 156 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+5 -2
View File
@@ -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))
app.py
+10 -4
View File
@@ -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
workflows/workflow.json
+204 -150
View File
@@ -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
workflows/workflow.png
BIN
View File
Binary file not shown.
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 1.3 MiB

After

Width:  |  Height:  |  Size: 954 KiB