Merge pull request #54 from frantic/patch-1

Avoid installing dependencies on every run

README.md

@@ -18,6 +18,23 @@ pip3 install -r requirements.txt
 
 The models will be automatically downloaded when used. You can also manually download them according to the table below. If the automatic download is slow, you can set the `HTTP_PROXY` and `HTTPS_PROXY` environment variables to use a proxy.
 
+### bert-base-uncased
+
+You can download the model from https://huggingface.co/bert-base-uncased/tree/main into the `models/bert-base-uncased` folder located in the root directory of ComfyUI, like this:
+
+```
+ComfyUI
+    models
+        bert-base-uncased
+            config.json
+            model.safetensors
+            tokenizer_config.json
+            tokenizer.json
+            vocab.txt
+```
+
+You can also skip this step. During the inference process, `bert-base-uncased` will be automatically downloaded through the `transformers` library, and its directory is typically `~/.cache/huggingface/hub/models--bert-base-uncased`.
+
 ### GroundingDino
 
 Please directly download the models and configuration files to the `models/grounding-dino` directory under the ComfyUI root directory, without modifying the file names.

__init__.py

@@ -1,10 +1,14 @@
 from .node import *
+from .install import *
 
 NODE_CLASS_MAPPINGS = {
     'SAMModelLoader (segment anything)': SAMModelLoader,
     'GroundingDinoModelLoader (segment anything)': GroundingDinoModelLoader,
     'GroundingDinoSAMSegment (segment anything)': GroundingDinoSAMSegment,
-    'InvertMask (segment anything)': InvertMask
+    'InvertMask (segment anything)': InvertMask,
+    "IsMaskEmpty": IsMaskEmptyNode,
 }
 
 __all__ = ['NODE_CLASS_MAPPINGS']
+
+

install.py

@@ -0,0 +1,18 @@
+import sys
+import os.path
+import subprocess
+
+custom_nodes_path = os.path.dirname(os.path.abspath(__file__))
+
+def build_pip_install_cmds(args):
+    if "python_embeded" in sys.executable or "python_embedded" in sys.executable:
+        return [sys.executable, '-s', '-m', 'pip', 'install'] + args
+    else:
+        return [sys.executable, '-m', 'pip', 'install'] + args
+
+def ensure_package():
+    cmds = build_pip_install_cmds(['-r', 'requirements.txt'])
+    subprocess.run(cmds, cwd=custom_nodes_path)
+
+if __name__ == "__main__":
+    ensure_package()

node.py

@@ -19,10 +19,12 @@ from local_groundingdino.datasets import transforms as T
 from local_groundingdino.util.utils import clean_state_dict as local_groundingdino_clean_state_dict
 from local_groundingdino.util.slconfig import SLConfig as local_groundingdino_SLConfig
 from local_groundingdino.models import build_model as local_groundingdino_build_model
+import glob
+import folder_paths
 
 logger = logging.getLogger('comfyui_segment_anything')
 
-sam_model_dir = os.path.join(folder_paths.models_dir, "sams")
+sam_model_dir_name = "sams"
 sam_model_list = {
     "sam_vit_h (2.56GB)": {
         "model_url": "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth"
@@ -47,8 +49,7 @@ sam_model_list = {
     }
 }
 
-groundingdino_model_dir = os.path.join(
-    folder_paths.models_dir, "grounding-dino")
+groundingdino_model_dir_name = "grounding-dino"
 groundingdino_model_list = {
     "GroundingDINO_SwinT_OGC (694MB)": {
         "config_url": "https://huggingface.co/ShilongLiu/GroundingDINO/resolve/main/GroundingDINO_SwinT_OGC.cfg.py",
@@ -60,6 +61,12 @@ groundingdino_model_list = {
     },
 }
 
+def get_bert_base_uncased_model_path():
+    comfy_bert_model_base = os.path.join(folder_paths.models_dir, 'bert-base-uncased')
+    if glob.glob(os.path.join(comfy_bert_model_base, '**/model.safetensors'), recursive=True):
+        print('grounding-dino is using models/bert-base-uncased')
+        return comfy_bert_model_base
+    return 'bert-base-uncased'
 
 def list_files(dirpath, extensions=[]):
     return [f for f in os.listdir(dirpath) if os.path.isfile(os.path.join(dirpath, f)) and f.split('.')[-1] in extensions]
@@ -71,7 +78,7 @@ def list_sam_model():
 
 def load_sam_model(model_name):
     sam_checkpoint_path = get_local_filepath(
-        sam_model_list[model_name]["model_url"], sam_model_dir)
+        sam_model_list[model_name]["model_url"], sam_model_dir_name)
     model_file_name = os.path.basename(sam_checkpoint_path)
     model_type = model_file_name.split('.')[0]
     if 'hq' not in model_type and 'mobile' not in model_type:
@@ -85,14 +92,22 @@ def load_sam_model(model_name):
 
 
 def get_local_filepath(url, dirname, local_file_name=None):
-    if not os.path.exists(dirname):
-        os.makedirs(dirname)
     if not local_file_name:
         parsed_url = urlparse(url)
         local_file_name = os.path.basename(parsed_url.path)
-    destination = os.path.join(dirname, local_file_name)
+
+    destination = folder_paths.get_full_path(dirname, local_file_name)
+    if destination:
+        logger.warn(f'using extra model: {destination}')
+        return destination
+
+    folder = os.path.join(folder_paths.models_dir, dirname)
+    if not os.path.exists(folder):
+        os.makedirs(folder)
+
+    destination = os.path.join(folder, local_file_name)
     if not os.path.exists(destination):
-        logging.warn(f'downloading {url} to {destination}')
+        logger.warn(f'downloading {url} to {destination}')
         download_url_to_file(url, destination)
     return destination
 
@@ -101,15 +116,18 @@ def load_groundingdino_model(model_name):
     dino_model_args = local_groundingdino_SLConfig.fromfile(
         get_local_filepath(
             groundingdino_model_list[model_name]["config_url"],
-            groundingdino_model_dir
+            groundingdino_model_dir_name
         ),
-
     )
+
+    if dino_model_args.text_encoder_type == 'bert-base-uncased':
+        dino_model_args.text_encoder_type = get_bert_base_uncased_model_path()
+    
     dino = local_groundingdino_build_model(dino_model_args)
     checkpoint = torch.load(
         get_local_filepath(
             groundingdino_model_list[model_name]["model_url"],
-            groundingdino_model_dir,
+            groundingdino_model_dir_name,
         ),
     )
     dino.load_state_dict(local_groundingdino_clean_state_dict(
@@ -334,22 +352,19 @@ class InvertMask:
         out = 1.0 - mask
         return (out,)
 
+class IsMaskEmptyNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "mask": ("MASK",),
+            },
+        }
+    RETURN_TYPES = ["NUMBER"]
+    RETURN_NAMES = ["boolean_number"]
 
-if __name__ == "__main__":
-    input_image = Image.open(
-        '/data/dev/comfyui-latest/custom_nodes/comfyui_segment_anything/human.jpg').convert('RGBA')
-    dino_model = load_groundingdino_model('GroundingDINO_SwinT_OGC (694MB)')
-    boxes = groundingdino_predict(
-        dino_model,
-        input_image,
-        'face . glasses . forehead',
-        0.3
-    )
-    sam_model = load_sam_model('sam_hq_vit_h (2.57GB)')
-    (output_images, output_masks) = sam_segment(
-        sam_model,
-        input_image,
-        boxes
-    )
-    for i in range(len(output_images)):
-        output_images[i].save(f"result_{i}.png")
+    FUNCTION = "main"
+    CATEGORY = "segment_anything"
+
+    def main(self, mask):
+        return (torch.all(mask == 0).int().item(), )