From 78981d4eb7cb89ebda75fa8d7242e4d55d9f6d8c Mon Sep 17 00:00:00 2001 From: xingren23 Date: Tue, 20 Feb 2024 16:00:36 +0800 Subject: [PATCH 1/5] support extra_model_paths config --- WAS_Node_Suite.py | 46 ++++++++++++++++++++++++++++++++++++++-------- 1 file changed, 38 insertions(+), 8 deletions(-) diff --git a/WAS_Node_Suite.py b/WAS_Node_Suite.py index 9e644b7..5aec58b 100644 --- a/WAS_Node_Suite.py +++ b/WAS_Node_Suite.py @@ -155,7 +155,13 @@ def print(self, **kwargs): ALLOWED_EXT = ('.jpeg', '.jpg', '.png', '.tiff', '.gif', '.bmp', '.webp') - +#! Get model path with extra_model_paths +def get_model_path_with_extra(model_name): + if model_name in comfy_paths.folder_names_and_paths.keys(): + return comfy_paths.folder_names_and_paths[model_name][0][0] + else: + return os.path.join(MODELS_DIR, model_name) + #! INSTALLATION CLEANUP # Delete legacy nodes @@ -5840,7 +5846,7 @@ def image_rembg( from rembg import remove, new_session - os.environ['U2NET_HOME'] = os.path.join(MODELS_DIR, 'rembg') + os.environ['U2NET_HOME'] = get_model_path_with_extra('rembg') os.makedirs(os.environ['U2NET_HOME'], exist_ok=True) # Set bgcolor @@ -8491,7 +8497,7 @@ def inject_noise(self, samples, noise_std): class MiDaS_Model_Loader: def __init__(self): - self.midas_dir = os.path.join(MODELS_DIR, 'midas') + self.midas_dir = get_model_path_with_extra('midas') @classmethod def INPUT_TYPES(cls): @@ -8551,7 +8557,7 @@ def install_midas(self): class MiDaS_Depth_Approx: def __init__(self): - self.midas_dir = os.path.join(MODELS_DIR, 'midas') + self.midas_dir =get_model_path_with_extra('midas') @classmethod def INPUT_TYPES(cls): @@ -8668,7 +8674,7 @@ def install_midas(self): class MiDaS_Background_Foreground_Removal: def __init__(self): - self.midas_dir = os.path.join(MODELS_DIR, 'midas') + self.midas_dir = get_model_path_with_extra('midas') @classmethod def INPUT_TYPES(cls): @@ -10981,7 +10987,11 @@ def blip_model(self, blip_model): from .modules.BLIP.blip_module import blip_decoder +<<<<<<< HEAD blip_dir = os.path.join(MODELS_DIR, 'blip') +======= + blip_dir = get_model_path_with_extra('blip') +>>>>>>> f92bc24 (support extra_model_paths config) if not os.path.exists(blip_dir): os.makedirs(blip_dir, exist_ok=True) @@ -10999,8 +11009,13 @@ def blip_model(self, blip_model): elif blip_model == 'interrogate': from .modules.BLIP.blip_module import blip_vqa +<<<<<<< HEAD blip_dir = os.path.join(MODELS_DIR, 'blip') +======= + + blip_dir = get_model_path_with_extra('blip') +>>>>>>> f92bc24 (support extra_model_paths config) if not os.path.exists(blip_dir): os.makedirs(blip_dir, exist_ok=True) @@ -11077,8 +11092,13 @@ def transformImage(input_image, image_size, device): model = blip_model[0].to(device) else: from .modules.BLIP.blip_module import blip_decoder +<<<<<<< HEAD blip_dir = os.path.join(MODELS_DIR, 'blip') +======= + + blip_dir = get_model_path_with_extra('blip') +>>>>>>> f92bc24 (support extra_model_paths config) if not os.path.exists(blip_dir): os.makedirs(blip_dir, exist_ok=True) @@ -11106,8 +11126,13 @@ def transformImage(input_image, image_size, device): model = blip_model[0].to(device) else: from .modules.BLIP.blip_module import blip_vqa +<<<<<<< HEAD blip_dir = os.path.join(MODELS_DIR, 'blip') +======= + + blip_dir = get_model_path_with_extra('blip') +>>>>>>> f92bc24 (support extra_model_paths config) if not os.path.exists(blip_dir): os.makedirs(blip_dir, exist_ok=True) @@ -11155,7 +11180,7 @@ def INPUT_TYPES(cls): def clipseg_model(self, model): from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation - cache = os.path.join(MODELS_DIR, 'clipseg') + cache = get_model_path_with_extra('clipseg') inputs = CLIPSegProcessor.from_pretrained(model, cache_dir=cache) model = CLIPSegForImageSegmentation.from_pretrained(model, cache_dir=cache) @@ -11190,7 +11215,7 @@ def CLIPSeg_image(self, image, text=None, clipseg_model=None): from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation image = tensor2pil(image) - cache = os.path.join(MODELS_DIR, 'clipseg') + cache = get_model_path_with_extra('clipseg') if clipseg_model: inputs = clipseg_model[0] @@ -11284,7 +11309,7 @@ def CLIPSeg_images(self, image_a, image_b, text_a, text_b, image_c=None, image_d if text_f: prompts.append(text_f) - cache = os.path.join(MODELS_DIR, 'clipseg') + cache = get_model_path_with_extra('clipseg') inputs = CLIPSegProcessor.from_pretrained("CIDAS/clipseg-rd64-refined", cache_dir=cache) model = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined", cache_dir=cache) @@ -11359,8 +11384,13 @@ def sam_load_model(self, model_size): Repo.clone_from('https://github.com/facebookresearch/segment-anything', os.path.join(WAS_SUITE_ROOT, 'repos'+os.sep+'SAM')) sys.path.append(os.path.join(WAS_SUITE_ROOT, 'repos'+os.sep+'SAM')) +<<<<<<< HEAD sam_dir = os.path.join(MODELS_DIR, 'sam') +======= + + sam_dir = get_model_path_with_extra('sam') +>>>>>>> f92bc24 (support extra_model_paths config) if not os.path.exists(sam_dir): os.makedirs(sam_dir, exist_ok=True) From 39c8f099cdfed6ccad9a39ee54bd8094bdc54bfd Mon Sep 17 00:00:00 2001 From: xingren23 Date: Tue, 20 Feb 2024 16:07:33 +0800 Subject: [PATCH 2/5] merge --- WAS_Node_Suite.py | 24 ------------------------ 1 file changed, 24 deletions(-) diff --git a/WAS_Node_Suite.py b/WAS_Node_Suite.py index 5aec58b..53186a5 100644 --- a/WAS_Node_Suite.py +++ b/WAS_Node_Suite.py @@ -10987,11 +10987,7 @@ def blip_model(self, blip_model): from .modules.BLIP.blip_module import blip_decoder -<<<<<<< HEAD - blip_dir = os.path.join(MODELS_DIR, 'blip') -======= blip_dir = get_model_path_with_extra('blip') ->>>>>>> f92bc24 (support extra_model_paths config) if not os.path.exists(blip_dir): os.makedirs(blip_dir, exist_ok=True) @@ -11009,13 +11005,8 @@ def blip_model(self, blip_model): elif blip_model == 'interrogate': from .modules.BLIP.blip_module import blip_vqa -<<<<<<< HEAD - - blip_dir = os.path.join(MODELS_DIR, 'blip') -======= blip_dir = get_model_path_with_extra('blip') ->>>>>>> f92bc24 (support extra_model_paths config) if not os.path.exists(blip_dir): os.makedirs(blip_dir, exist_ok=True) @@ -11092,13 +11083,8 @@ def transformImage(input_image, image_size, device): model = blip_model[0].to(device) else: from .modules.BLIP.blip_module import blip_decoder -<<<<<<< HEAD - - blip_dir = os.path.join(MODELS_DIR, 'blip') -======= blip_dir = get_model_path_with_extra('blip') ->>>>>>> f92bc24 (support extra_model_paths config) if not os.path.exists(blip_dir): os.makedirs(blip_dir, exist_ok=True) @@ -11126,13 +11112,8 @@ def transformImage(input_image, image_size, device): model = blip_model[0].to(device) else: from .modules.BLIP.blip_module import blip_vqa -<<<<<<< HEAD - blip_dir = os.path.join(MODELS_DIR, 'blip') -======= - blip_dir = get_model_path_with_extra('blip') ->>>>>>> f92bc24 (support extra_model_paths config) if not os.path.exists(blip_dir): os.makedirs(blip_dir, exist_ok=True) @@ -11384,13 +11365,8 @@ def sam_load_model(self, model_size): Repo.clone_from('https://github.com/facebookresearch/segment-anything', os.path.join(WAS_SUITE_ROOT, 'repos'+os.sep+'SAM')) sys.path.append(os.path.join(WAS_SUITE_ROOT, 'repos'+os.sep+'SAM')) -<<<<<<< HEAD - - sam_dir = os.path.join(MODELS_DIR, 'sam') -======= sam_dir = get_model_path_with_extra('sam') ->>>>>>> f92bc24 (support extra_model_paths config) if not os.path.exists(sam_dir): os.makedirs(sam_dir, exist_ok=True) From f54332cb1f772bdaaf99d8eea449043d480d7553 Mon Sep 17 00:00:00 2001 From: xingren23 Date: Wed, 6 Mar 2024 21:19:53 +0800 Subject: [PATCH 3/5] remove duplicate method WAS_Checkpoint_Loader --- WAS_Node_Suite.py | 17 ----------------- 1 file changed, 17 deletions(-) diff --git a/WAS_Node_Suite.py b/WAS_Node_Suite.py index 53186a5..ac3bbd2 100644 --- a/WAS_Node_Suite.py +++ b/WAS_Node_Suite.py @@ -13183,23 +13183,6 @@ def IS_CHANGED(cls, **kwargs): # CUSTOM COMFYUI NODES -class WAS_Checkpoint_Loader: - @classmethod - def INPUT_TYPES(s): - return {"required": { "config_name": (comfy_paths.get_filename_list("configs"), ), - "ckpt_name": (comfy_paths.get_filename_list("checkpoints"), )}} - RETURN_TYPES = ("MODEL", "CLIP", "VAE", TEXT_TYPE) - RETURN_NAMES = ("MODEL", "CLIP", "VAE", "NAME_STRING") - FUNCTION = "load_checkpoint" - - CATEGORY = "WAS Suite/Loaders/Advanced" - - def load_checkpoint(self, config_name, ckpt_name, output_vae=True, output_clip=True): - config_path = comfy_paths.get_full_path("configs", config_name) - ckpt_path = comfy_paths.get_full_path("checkpoints", ckpt_name) - out = comfy.sd.load_checkpoint(config_path, ckpt_path, output_vae=True, output_clip=True, embedding_directory=comfy_paths.get_folder_paths("embeddings")) - return (out[0], out[1], out[2], os.path.splitext(os.path.basename(ckpt_name))[0]) - class WAS_Checkpoint_Loader: @classmethod def INPUT_TYPES(s): From 2894fa887d96d9ff52ebe230d3931dca4c1fc03b Mon Sep 17 00:00:00 2001 From: xingren23 Date: Fri, 8 Mar 2024 19:31:53 +0800 Subject: [PATCH 4/5] cached_package speed up load time --- WAS_Node_Suite.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/WAS_Node_Suite.py b/WAS_Node_Suite.py index fc0bfdf..6efd134 100644 --- a/WAS_Node_Suite.py +++ b/WAS_Node_Suite.py @@ -236,7 +236,6 @@ def getSuiteConfig(): cstr(f"Unable to load conf file at `{WAS_CONFIG_FILE}`. Using internal config template.").error.print() return was_conf_template return was_config - return was_config def updateSuiteConfig(conf): try: @@ -336,12 +335,18 @@ def updateSuiteConfig(conf): #! SUITE SPECIFIC CLASSES & FUNCTIONS # Freeze PIP modules +cached_packages = None def packages(versions=False): - import sys - import subprocess - return [( r.decode().split('==')[0] if not versions else r.decode() ) for r in subprocess.check_output([sys.executable, '-s', '-m', 'pip', 'freeze']).split()] - + global cached_packages + if cached_packages is None: + cached_packages = subprocess.check_output([sys.executable, '-m', 'pip', 'freeze']).split() + if versions: + return [r.decode() for r in cached_packages] + else: + return [r.decode().split('==')[0] for r in cached_packages] + def install_package(package, uninstall_first: Union[List[str], str] = None): + print(f"Installing {package}...") if os.getenv("WAS_BLOCK_AUTO_INSTALL", 'False').lower() in ('true', '1', 't'): cstr(f"Preventing package install of '{package}' due to WAS_BLOCK_INSTALL env").msg.print() else: From 4f49c274fced97793e7f05e857e4ccc2b4237a04 Mon Sep 17 00:00:00 2001 From: xingren23 Date: Mon, 11 Mar 2024 10:13:36 +0800 Subject: [PATCH 5/5] don't check_deps everytime --- WAS_Node_Suite.py | 97 ++++++++++++++++++++++++----------------------- 1 file changed, 50 insertions(+), 47 deletions(-) diff --git a/WAS_Node_Suite.py b/WAS_Node_Suite.py index 6efd134..142e55d 100644 --- a/WAS_Node_Suite.py +++ b/WAS_Node_Suite.py @@ -14093,61 +14093,64 @@ def encode(self, clip, text, token_normalization, weight_interpretation, seed=0, if NODE_CLASS_MAPPINGS.__contains__("CLIPTextEncode (BlenderNeko Advanced + NSP)"): cstr('`CLIPTextEncode (BlenderNeko Advanced + NSP)` node enabled under `WAS Suite/Conditioning` menu.').msg.print() -# opencv-python-headless handling -if 'opencv-python' in packages() or 'opencv-python-headless' in packages(): - try: - import cv2 - build_info = ' '.join(cv2.getBuildInformation().split()) - if "FFMPEG: YES" in build_info: - if was_config.__contains__('show_startup_junk'): - if was_config['show_startup_junk']: - cstr("OpenCV Python FFMPEG support is enabled").msg.print() - if was_config.__contains__('ffmpeg_bin_path'): - if was_config['ffmpeg_bin_path'] == "/path/to/ffmpeg": - cstr(f"`ffmpeg_bin_path` is not set in `{WAS_CONFIG_FILE}` config file. Will attempt to use system ffmpeg binaries if available.").warning.print() - else: - if was_config.__contains__('show_startup_junk'): - if was_config['show_startup_junk']: - cstr(f"`ffmpeg_bin_path` is set to: {was_config['ffmpeg_bin_path']}").msg.print() - else: - cstr(f"OpenCV Python FFMPEG support is not enabled\033[0m. OpenCV Python FFMPEG support, and FFMPEG binaries is required for video writing.").warning.print() - except ImportError: - cstr("OpenCV Python module cannot be found. Attempting install...").warning.print() - install_package( - package='opencv-python-headless[ffmpeg]', - uninstall_first=['opencv-python', 'opencv-python-headless[ffmpeg]'] - ) +def check_deps(): + # opencv-python-headless handling + if 'opencv-python' in packages() or 'opencv-python-headless' in packages(): + try: + import cv2 + build_info = ' '.join(cv2.getBuildInformation().split()) + if "FFMPEG: YES" in build_info: + if was_config.__contains__('show_startup_junk'): + if was_config['show_startup_junk']: + cstr("OpenCV Python FFMPEG support is enabled").msg.print() + if was_config.__contains__('ffmpeg_bin_path'): + if was_config['ffmpeg_bin_path'] == "/path/to/ffmpeg": + cstr(f"`ffmpeg_bin_path` is not set in `{WAS_CONFIG_FILE}` config file. Will attempt to use system ffmpeg binaries if available.").warning.print() + else: + if was_config.__contains__('show_startup_junk'): + if was_config['show_startup_junk']: + cstr(f"`ffmpeg_bin_path` is set to: {was_config['ffmpeg_bin_path']}").msg.print() + else: + cstr(f"OpenCV Python FFMPEG support is not enabled\033[0m. OpenCV Python FFMPEG support, and FFMPEG binaries is required for video writing.").warning.print() + except ImportError: + cstr("OpenCV Python module cannot be found. Attempting install...").warning.print() + install_package( + package='opencv-python-headless[ffmpeg]', + uninstall_first=['opencv-python', 'opencv-python-headless[ffmpeg]'] + ) + try: + import cv2 + cstr("OpenCV Python installed.").msg.print() + except ImportError: + cstr("OpenCV Python module still cannot be imported. There is a system conflict.").error.print() + else: + install_package('opencv-python-headless[ffmpeg]') try: import cv2 cstr("OpenCV Python installed.").msg.print() except ImportError: cstr("OpenCV Python module still cannot be imported. There is a system conflict.").error.print() -else: - install_package('opencv-python-headless[ffmpeg]') - try: - import cv2 - cstr("OpenCV Python installed.").msg.print() - except ImportError: - cstr("OpenCV Python module still cannot be imported. There is a system conflict.").error.print() -# scipy handling -if 'scipy' not in packages(): - install_package('scipy') + # scipy handling + if 'scipy' not in packages(): + install_package('scipy') + try: + import scipy + except ImportError as e: + cstr("Unable to import tools for certain masking procedures.").msg.print() + print(e) + + # scikit-image handling try: - import scipy + import skimage except ImportError as e: - cstr("Unable to import tools for certain masking procedures.").msg.print() - print(e) - -# scikit-image handling -try: - import skimage -except ImportError as e: - install_package( - package='scikit-image', - uninstall_first=['scikit-image'] - ) - import skimage + install_package( + package='scikit-image', + uninstall_first=['scikit-image'] + ) + import skimage +# Check for dependencies +# check_deps() was_conf = getSuiteConfig()