diff --git a/js/index.js b/js/index.js index 646ac88..c4acbc5 100644 --- a/js/index.js +++ b/js/index.js @@ -298,9 +298,6 @@ const ext = { node.widgets.find((x) => x.name === "embedding_id").value = id; const ckpt = node.widgets.find((x) => x.name === "ckpt").value; - const model_type = node.widgets.find( - (x) => x.name === "model_type" - ).value; api .fetchApi("/sam_model", { @@ -309,7 +306,6 @@ const ext = { image: connectedImageFileName, embedding_id: id, ckpt, - model_type, }), }) .then(() => { diff --git a/routes.py b/routes.py index 2d8ed43..3b2806d 100644 --- a/routes.py +++ b/routes.py @@ -7,6 +7,7 @@ import folder_paths import json import numpy as np import server +import re # For speeding up ONNX model, see https://github.com/facebookresearch/segment-anything/tree/main/demo#onnx-multithreading-with-sharedarraybuffer def inject_headers(original_handler): @@ -65,7 +66,7 @@ async def post_sam_model(request): if not os.path.exists(emb_filename): image = load_image(post.get("image")) ckpt = post.get("ckpt") - model_type = post.get("model_type") + model_type = re.findall(r'vit_[lbh]', ckpt)[0] ckpt = folder_paths.get_full_path("sams", ckpt) sam = sam_model_registry[model_type](checkpoint=ckpt) predictor = SamPredictor(sam) diff --git a/sam/sam.py b/sam/sam.py index 81ad70f..7c03ffd 100644 --- a/sam/sam.py +++ b/sam/sam.py @@ -2,6 +2,7 @@ import folder_paths import os import numpy as np import torch +import re from segment_anything import sam_model_registry, SamPredictor from einops import rearrange, repeat @@ -24,7 +25,6 @@ class SAM: return { "required": { "image": ("IMAGE",), - "model_type": (["vit_h", "vit_l", "vit_b"],), "ckpt": (folder_paths.get_filename_list("sams"),), "embedding_id": ( "STRING", @@ -40,13 +40,14 @@ class SAM: RETURN_TYPES = ("SAM_PROMPT",) FUNCTION = "load_image" - def load_image(self, image, model_type, ckpt, embedding_id, image_prompts_json): + def load_image(self, image, ckpt, embedding_id, image_prompts_json): import json global global_predictor if global_predictor is None: ckpt = folder_paths.get_full_path("sams", ckpt) + model_type = re.findall(r'vit_[lbh]', ckpt)[0] sam = sam_model_registry[model_type](checkpoint=ckpt) predictor = SamPredictor(sam) global_predictor = predictor diff --git a/sam/sam_prompt.py b/sam/sam_prompt.py index 7dd99db..81598aa 100644 --- a/sam/sam_prompt.py +++ b/sam/sam_prompt.py @@ -1,6 +1,7 @@ import folder_paths import os import numpy as np +import re from segment_anything import sam_model_registry, SamPredictor @@ -19,7 +20,6 @@ class SAM_Prompt_Image: return { "required": { "image": ("IMAGE",), - "model_type": (["vit_h", "vit_l", "vit_b"],), "ckpt": (folder_paths.get_filename_list("sams"),), "embedding_id": ( "STRING", @@ -35,12 +35,13 @@ class SAM_Prompt_Image: RETURN_TYPES = ("SAM_PROMPT",) FUNCTION = "load_image" - def load_image(self, image, model_type, ckpt, embedding_id, image_prompts_json): + def load_image(self, image, ckpt, embedding_id, image_prompts_json): import json emb_filename = f"{self.output_dir}/{embedding_id}.npy" if not os.path.exists(emb_filename): ckpt = folder_paths.get_full_path("sams", ckpt) + model_type = re.findall(r'vit_[lbh]', ckpt)[0] sam = sam_model_registry[model_type](checkpoint=ckpt) predictor = SamPredictor(sam)