diff --git a/__init__.py b/__init__.py index d8e82d9..ae51f2f 100644 --- a/__init__.py +++ b/__init__.py @@ -35,6 +35,9 @@ check_requirements_installed(llama_cpp_agent_path) init() node_list = [ + "image.birefnet_LP", + "image.rmbg_LP", + "image.rembg_LP", "llm.llm_LP", "vlm.autotagger_LP", "vlm.llava_LP", diff --git a/install_init.py b/install_init.py index c06c824..bcbaa5e 100644 --- a/install_init.py +++ b/install_init.py @@ -154,6 +154,7 @@ def install_llama(system_info): # Build from source with appropriate acceleration try: + print("WARNING: Building the llama-cpp-python wheels can take UP TO AN HOUR due to the specific build conditions. This is not an error, just wait.") if system_info.get('metal', False): print("Building llama-cpp-python from source with Metal support") os.environ['CMAKE_ARGS'] = "-DGGML_METAL=on" @@ -327,17 +328,6 @@ def get_comfy_dir(subpath=None, mkdir=False): def should_install_js(): return not hasattr(PromptServer.instance, "supports") or "custom_nodes_from_web" not in PromptServer.instance.supports -def has_nvidia_cuda(): - try: - output = subprocess.check_output( - ['nvidia-smi', '--query-gpu=name', '--format=csv,noheader'], - stderr=subprocess.DEVNULL, - encoding='utf-8' - ).strip() - return bool(output) - except (subprocess.CalledProcessError, FileNotFoundError): - return False - def is_installed(pkg_name: str, min_version: str = '') -> bool: try: ver = metadata.version(pkg_name) @@ -354,7 +344,8 @@ def uninstall(pkg: str): subprocess.check_call([sys.executable, '-m', 'pip', 'uninstall', '-y', pkg]) def install_onnxruntime(): - gpu = has_nvidia_cuda() + import torch + gpu = torch.cuda.is_available() print(f"Found NVIDIA GPU: {gpu}") if gpu: if is_installed('onnxruntime'): diff --git a/node_list.json b/node_list.json index 5c1b51f..7308d85 100644 --- a/node_list.json +++ b/node_list.json @@ -9,5 +9,9 @@ "LLava Sampler Advanced [LP]":"Generates text using the VLM. Advanced settings.", "LLava Simple [LP]":"Generates text using the loaded VLM from LLava Loader. Simple settings.", "LLava Advanced [LP]":"Generates text using the loaded VLM from LLava Loader. Advanced settings.", - "Recognize Anything (RAM) [LP]":"Recognize Anything (RAM) - node for generating tags for an image using the Recognize Anything Plus Model" + "Recognize Anything (RAM) [LP]":"Recognize Anything (RAM) - node for generating tags for an image using the Recognize Anything Plus Model", + "Image Remove Background (rembg) [LP]":"Removes background for images with user-defined model. This node uses the rembg library. It is NOT recommended for use in normal scenarios.", + "Image Remove Background (BiRefNet) [LP]":"Removes background for images with user-defined model type of BiRefNet. Recommended for fast and high-quality background removal", + "Image Remove Background (RMBG) [LP]":"Removes background for images with user-defined model. This node uses the RMBG models. Recommended for ultra-high precision", + "Multimodal Generator Advanced [LP]":"Node for Multimodal neural models based on GGUF. Supports Qwen2.5-VL of GGUF format. Based on llama-mtmd." } \ No newline at end of file diff --git a/nodes/image/birefnet_LP.py b/nodes/image/birefnet_LP.py new file mode 100644 index 0000000..142d343 --- /dev/null +++ b/nodes/image/birefnet_LP.py @@ -0,0 +1,421 @@ +import os +import torch +from PIL import Image, ImageFilter +from torchvision import transforms +import numpy as np +import folder_paths +from huggingface_hub import hf_hub_download +import sys +import importlib.util +from safetensors.torch import load_file +import cv2 + +device = "cuda" if torch.cuda.is_available() else "cpu" + +folder_paths.add_model_folder_path("rmbg", os.path.join(folder_paths.models_dir, "RMBG")) + +MODEL_CONFIG = { + "BiRefNet-HR-matting": { + "repo_id": "1038lab/BiRefNet", + "files": { + "birefnet.py": "birefnet.py", + "BiRefNet_config.py": "BiRefNet_config.py", + "BiRefNet-HR-matting.safetensors": "BiRefNet-HR-matting.safetensors", + "config.json": "config.json" + }, + "cache_dir": "BiRefNet", + "description": "High resolution matting model", + "default_res": 2048, + "max_res": 2560, + "min_res": 1024 + }, + "BiRefNet-HR": { + "repo_id": "1038lab/BiRefNet", + "files": { + "birefnet.py": "birefnet.py", + "BiRefNet_config.py": "BiRefNet_config.py", + "BiRefNet-HR.safetensors": "BiRefNet-HR.safetensors", + "config.json": "config.json" + }, + "cache_dir": "BiRefNet", + "description": "High resolution general purpose model", + "default_res": 2048, + "max_res": 2560, + "min_res": 1024 + }, + "BiRefNet-general": { + "repo_id": "1038lab/BiRefNet", + "files": { + "birefnet.py": "birefnet.py", + "BiRefNet_config.py": "BiRefNet_config.py", + "BiRefNet-general.safetensors": "BiRefNet-general.safetensors", + "config.json": "config.json" + }, + "cache_dir": "BiRefNet", + "description": "General purpose model with balanced performance", + "default_res": 1024, + "max_res": 2048, + "min_res": 512 + }, + "BiRefNet_512x512": { + "repo_id": "1038lab/BiRefNet", + "files": { + "birefnet.py": "birefnet.py", + "BiRefNet_config.py": "BiRefNet_config.py", + "BiRefNet_512x512.safetensors": "BiRefNet_512x512.safetensors", + "config.json": "config.json" + }, + "cache_dir": "BiRefNet", + "description": "Optimized for 512x512 resolution, faster processing", + "default_res": 512, + "max_res": 1024, + "min_res": 256, + "force_res": True + }, + "BiRefNet-portrait": { + "repo_id": "1038lab/BiRefNet", + "files": { + "birefnet.py": "birefnet.py", + "BiRefNet_config.py": "BiRefNet_config.py", + "BiRefNet-portrait.safetensors": "BiRefNet-portrait.safetensors", + "config.json": "config.json" + }, + "cache_dir": "BiRefNet", + "description": "Optimized for portrait/human matting", + "default_res": 1024, + "max_res": 2048, + "min_res": 512 + }, + "BiRefNet-matting": { + "repo_id": "1038lab/BiRefNet", + "files": { + "birefnet.py": "birefnet.py", + "BiRefNet_config.py": "BiRefNet_config.py", + "BiRefNet-matting.safetensors": "BiRefNet-matting.safetensors", + "config.json": "config.json" + }, + "cache_dir": "BiRefNet", + "description": "General purpose matting model", + "default_res": 1024, + "max_res": 2048, + "min_res": 512 + }, + "BiRefNet_lite": { + "repo_id": "1038lab/BiRefNet", + "files": { + "birefnet_lite.py": "birefnet_lite.py", + "BiRefNet_config.py": "BiRefNet_config.py", + "BiRefNet_lite.safetensors": "BiRefNet_lite.safetensors", + "config.json": "config.json" + }, + "cache_dir": "BiRefNet", + "description": "Lightweight version for faster processing", + "default_res": 1024, + "max_res": 2048, + "min_res": 512 + }, + "BiRefNet_lite-2K": { + "repo_id": "1038lab/BiRefNet", + "files": { + "birefnet_lite.py": "birefnet_lite.py", + "BiRefNet_config.py": "BiRefNet_config.py", + "BiRefNet_lite-2K.safetensors": "BiRefNet_lite-2K.safetensors", + "config.json": "config.json" + }, + "cache_dir": "BiRefNet", + "description": "Lightweight version optimized for 2K resolution", + "default_res": 2048, + "max_res": 2560, + "min_res": 1024 + } +} + +def tensor2pil(image): + return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)) + +def pil2tensor(image): + return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0) + +def handle_model_error(message): + print(f"[BiRefNet ERROR] {message}") + raise RuntimeError(message) + +def refine_foreground(image_bchw, masks_b1hw): + b, c, h, w = image_bchw.shape + if b != masks_b1hw.shape[0]: + raise ValueError("images and masks must have the same batch size") + + image_np = image_bchw.cpu().numpy() + mask_np = masks_b1hw.cpu().numpy() + + refined_fg = [] + for i in range(b): + mask = mask_np[i, 0] + thresh = 0.45 + mask_binary = (mask > thresh).astype(np.float32) + + edge_blur = cv2.GaussianBlur(mask_binary, (3, 3), 0) + transition_mask = np.logical_and(mask > 0.05, mask < 0.95) + + alpha = 0.85 + mask_refined = np.where(transition_mask, + alpha * mask + (1-alpha) * edge_blur, + mask_binary) + + edge_region = np.logical_and(mask > 0.2, mask < 0.8) + mask_refined = np.where(edge_region, + mask_refined * 0.98, + mask_refined) + + result = [] + for c in range(image_np.shape[1]): + channel = image_np[i, c] + refined = channel * mask_refined + result.append(refined) + + refined_fg.append(np.stack(result)) + + return torch.from_numpy(np.stack(refined_fg)) + +class BiRefNetModel: + def __init__(self): + self.model = None + self.current_model_version = None + self.base_cache_dir = os.path.join(folder_paths.models_dir, "RMBG") + + def get_cache_dir(self, model_name): + return os.path.join(self.base_cache_dir, MODEL_CONFIG[model_name]["cache_dir"]) + + def check_model_cache(self, model_name): + cache_dir = self.get_cache_dir(model_name) + + if not os.path.exists(cache_dir): + return False, "Model directory not found" + + missing_files = [] + for filename in MODEL_CONFIG[model_name]["files"].keys(): + if not os.path.exists(os.path.join(cache_dir, filename)): + missing_files.append(filename) + + if missing_files: + return False, f"Missing model files: {', '.join(missing_files)}" + + return True, "Model cache verified" + + def download_model(self, model_name): + cache_dir = self.get_cache_dir(model_name) + + try: + os.makedirs(cache_dir, exist_ok=True) + print(f"Downloading {model_name} model files...") + + for filename in MODEL_CONFIG[model_name]["files"].keys(): + print(f"Downloading {filename}...") + hf_hub_download( + repo_id=MODEL_CONFIG[model_name]["repo_id"], + filename=filename, + local_dir=cache_dir, + local_dir_use_symlinks=False + ) + + return True, "Model files downloaded successfully" + + except Exception as e: + return False, f"Error downloading model files: {str(e)}" + + def clear_model(self): + if self.model is not None: + self.model.cpu() + del self.model + self.model = None + self.current_model_version = None + torch.cuda.empty_cache() + print("Model cleared from memory") + + def load_model(self, model_name): + if self.current_model_version != model_name: + self.clear_model() + + cache_dir = self.get_cache_dir(model_name) + model_filename = [k for k in MODEL_CONFIG[model_name]["files"].keys() if k.endswith('.py') and k != "BiRefNet_config.py"][0] + model_path = os.path.join(cache_dir, model_filename) + config_path = os.path.join(cache_dir, "BiRefNet_config.py") + weights_filename = [k for k in MODEL_CONFIG[model_name]["files"].keys() if k.endswith('.safetensors')][0] + weights_path = os.path.join(cache_dir, weights_filename) + + try: + with open(model_path, 'r', encoding='utf-8') as f: + model_content = f.read() + model_content = model_content.replace("from .BiRefNet_config", "from BiRefNet_config") + with open(model_path, 'w', encoding='utf-8') as f: + f.write(model_content) + + spec = importlib.util.spec_from_file_location("BiRefNet_config", config_path) + config_module = importlib.util.module_from_spec(spec) + sys.modules["BiRefNet_config"] = config_module + spec.loader.exec_module(config_module) + + spec = importlib.util.spec_from_file_location("birefnet", model_path) + model_module = importlib.util.module_from_spec(spec) + sys.modules["birefnet"] = model_module + spec.loader.exec_module(model_module) + + self.model = model_module.BiRefNet(config_module.BiRefNetConfig()) + + state_dict = load_file(weights_path) + self.model.load_state_dict(state_dict) + + self.model.eval() + self.model.half() + torch.set_float32_matmul_precision('high') + self.model.to(device) + self.current_model_version = model_name + + except Exception as e: + handle_model_error(f"Error loading BiRefNet model: {str(e)}") + + def process_image(self, image, params): + try: + transform_image = transforms.Compose([ + transforms.Resize((params["process_res"], params["process_res"]), + interpolation=transforms.InterpolationMode.BICUBIC), + transforms.ToTensor(), + transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) + ]) + + orig_image = tensor2pil(image) + w, h = orig_image.size + + input_tensor = transform_image(orig_image).unsqueeze(0).to(device).half() + + with torch.no_grad(): + preds = self.model(input_tensor) + pred = preds[-1].sigmoid().cpu() + + pred = pred[0].squeeze() + pred_pil = transforms.ToPILImage()(pred) + mask = pred_pil.resize((w, h), Image.BICUBIC) + + return mask + + except Exception as e: + handle_model_error(f"Error in BiRefNet processing: {str(e)}") + +class ImageRemoveBackgroundBiRefNet: + def __init__(self): + self.model = BiRefNetModel() + + @classmethod + def INPUT_TYPES(s): + tooltips = { + "image": "Input image to be processed for background removal.", + "model": "Select the BiRefNet model variant to use.", + "mask_blur": "Specify the amount of blur to apply to the mask edges (0 for no blur, higher values for more blur).", + "mask_offset": "Adjust the mask boundary (positive values expand the mask, negative values shrink it).", + "invert_output": "Enable to invert both the image and mask output (useful for certain effects).", + "refine_foreground": "Use Fast Foreground Colour Estimation to optimize transparent background", + "background": "Choose background type: Alpha (transparent) or Color (custom background color).", + "background_color": "Choose background color (Alpha = transparent)" + } + return { + "required": { + "image": ("IMAGE", {"tooltip": tooltips["image"]}), + "model": (list(MODEL_CONFIG.keys()), {"tooltip": tooltips["model"]}), + }, + "optional": { + "mask_blur": ("INT", {"default": 0, "min": 0, "max": 64, "step": 1, "tooltip": tooltips["mask_blur"]}), + "mask_offset": ("INT", {"default": 0, "min": -20, "max": 20, "step": 1, "tooltip": tooltips["mask_offset"]}), + "invert_output": ("BOOLEAN", {"default": False, "tooltip": tooltips["invert_output"]}), + "refine_foreground": ("BOOLEAN", {"default": False, "tooltip": tooltips["refine_foreground"]}), + "background": (["Alpha", "Color"], {"default": "Alpha", "tooltip": tooltips["background"]}), + "background_color": ("COLOR", {"default": "#222222", "tooltip": tooltips["background_color"]}), + } + } + + RETURN_TYPES = ("IMAGE", "MASK", "IMAGE") + RETURN_NAMES = ("IMAGE", "MASK", "MASK_IMAGE") + FUNCTION = "process_image" + CATEGORY = "LevelPixel/Image" + + def process_image(self, image, model, **params): + try: + model_config = MODEL_CONFIG[model] + process_res = model_config.get("default_res", 1024) + if model_config.get("force_res", False): + base_res = 512 + process_res = ((process_res + base_res - 1) // base_res) * base_res + else: + process_res = process_res // 32 * 32 + print(f"Using {model} model with {process_res} resolution") + params["process_res"] = process_res + processed_images = [] + processed_masks = [] + cache_status, message = self.model.check_model_cache(model) + if not cache_status: + print(f"Cache check: {message}") + print("Downloading required model files...") + download_status, download_message = self.model.download_model(model) + if not download_status: + handle_model_error(download_message) + print("Model files downloaded successfully") + self.model.load_model(model) + for img in image: + mask = self.model.process_image(img, params) + if params["mask_blur"] > 0: + mask = mask.filter(ImageFilter.GaussianBlur(radius=params["mask_blur"])) + if params["mask_offset"] != 0: + if params["mask_offset"] > 0: + for _ in range(params["mask_offset"]): + mask = mask.filter(ImageFilter.MaxFilter(3)) + else: + for _ in range(-params["mask_offset"]): + mask = mask.filter(ImageFilter.MinFilter(3)) + if params["invert_output"]: + mask = Image.fromarray(255 - np.array(mask)) + img_tensor = torch.from_numpy(np.array(tensor2pil(img))).permute(2, 0, 1).unsqueeze(0) / 255.0 + mask_tensor = torch.from_numpy(np.array(mask)).unsqueeze(0).unsqueeze(0) / 255.0 + if params.get("refine_foreground", False): + refined_fg = refine_foreground(img_tensor, mask_tensor) + refined_fg = tensor2pil(refined_fg[0].permute(1, 2, 0)) + orig_image = tensor2pil(img) + r, g, b = refined_fg.split() + foreground = Image.merge('RGBA', (r, g, b, mask)) + else: + orig_image = tensor2pil(img) + orig_rgba = orig_image.convert("RGBA") + r, g, b, _ = orig_rgba.split() + foreground = Image.merge('RGBA', (r, g, b, mask)) + if params["background"] == "Alpha": + processed_images.append(pil2tensor(foreground)) + else: + def hex_to_rgba(hex_color): + hex_color = hex_color.lstrip('#') + if len(hex_color) == 6: + r, g, b = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16) + a = 255 + elif len(hex_color) == 8: + r, g, b, a = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16), int(hex_color[6:8], 16) + else: + raise ValueError("Invalid color format") + return (r, g, b, a) + rgba = hex_to_rgba(params["background_color"]) + bg_image = Image.new('RGBA', orig_image.size, rgba) + composite_image = Image.alpha_composite(bg_image, foreground) + processed_images.append(pil2tensor(composite_image.convert("RGB"))) + processed_masks.append(pil2tensor(mask)) + mask_images = [] + for mask_tensor in processed_masks: + mask_image = mask_tensor.reshape((-1, 1, mask_tensor.shape[-2], mask_tensor.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3) + mask_images.append(mask_image) + mask_image_output = torch.cat(mask_images, dim=0) + return (torch.cat(processed_images, dim=0), torch.cat(processed_masks, dim=0), mask_image_output) + except Exception as e: + handle_model_error(f"Error in image processing: {str(e)}") + +NODE_CLASS_MAPPINGS = { + "ImageRemoveBackgroundBiRefNet|LP": ImageRemoveBackgroundBiRefNet +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ImageRemoveBackgroundBiRefNet|LP": "Image Remove Background (BiRefNet) [LP]", +} \ No newline at end of file diff --git a/nodes/image/rembg_LP.py b/nodes/image/rembg_LP.py new file mode 100644 index 0000000..c6c8be4 --- /dev/null +++ b/nodes/image/rembg_LP.py @@ -0,0 +1,61 @@ +import numpy as np +import torch +from PIL import Image + +model_list = [ +'birefnet-general', # BEST! Recommended! A pre-trained model for general use cases. Source - https://github.com/ZhengPeng7/BiRefNet +'birefnet-massive', # A pre-trained model with massive dataset Source - https://github.com/ZhengPeng7/BiRefNet +'birefnet-general-lite', # A light pre-trained model for general use cases. Source - https://github.com/ZhengPeng7/BiRefNet +'birefnet-portrait', # A pre-trained model for human portraits. Source - https://github.com/ZhengPeng7/BiRefNet +'birefnet-dis', # A pre-trained model for dichotomous image segmentation (DIS). Source - https://github.com/ZhengPeng7/BiRefNet +'birefnet-hrsod', # A pre-trained model for high-resolution salient object detection (HRSOD). Source - https://github.com/ZhengPeng7/BiRefNet +'birefnet-cod', # A pre-trained model for concealed object detection (COD). Source - https://github.com/ZhengPeng7/BiRefNet +'u2net', # A pre-trained model for general use cases. Source - https://github.com/xuebinqin/U-2-Net +'u2netp', # A lightweight version of u2net model. Source - https://github.com/xuebinqin/U-2-Net +'u2net_human_seg', # A pre-trained model for human segmentation. Source - https://github.com/xuebinqin/U-2-Net +'u2net_cloth_seg', # A pre-trained model for Cloths Parsing from human portrait. Here clothes are parsed into 3 category: Upper body, Lower body and Full body. Source - https://github.com/levindabhi/cloth-segmentation +'silueta', # Same as u2net but the size is reduced to 43Mb. Source - https://github.com/xuebinqin/U-2-Net/issues/295 +'isnet-general-use', # A new pre-trained model for general use cases. Source - https://github.com/xuebinqin/DIS +'isnet-anime', # A high-accuracy segmentation for anime character. Source - https://github.com/SkyTNT/anime-segmentation +'sam' # A pre-trained model for any use cases. Source - https://github.com/danielgatis/rembg/releases/download/v0.0.0/vit_b-decoder-quant.onnx +] + +def pil2tensor(image): + return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0) + +def tensor2pil(image): + return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)) + +def list_model(): + return model_list + +class ImageRemoveBackground: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model_name": (list_model(), ), + "image": ("IMAGE",), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "remove_background" + CATEGORY = "LevelPixel/Image" + + def remove_background(self, image, model_name): + from rembg import new_session, remove + session = new_session(model_name) + image = pil2tensor(remove(tensor2pil(image), session = session)) + return (image,) + +NODE_CLASS_MAPPINGS = { + "ImageRemoveBackground|LP": ImageRemoveBackground +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ImageRemoveBackground|LP": "Image Remove Background (rembg) [LP]", +} diff --git a/nodes/image/rmbg_LP.py b/nodes/image/rmbg_LP.py new file mode 100644 index 0000000..48f6e0d --- /dev/null +++ b/nodes/image/rmbg_LP.py @@ -0,0 +1,646 @@ +import os +import torch +from PIL import Image +from torchvision import transforms +import numpy as np +import folder_paths +from PIL import ImageFilter +import torch.nn.functional as F +from huggingface_hub import hf_hub_download +import sys +import importlib.util +from transformers import AutoModelForImageSegmentation +import cv2 +import types + +device = "cuda" if torch.cuda.is_available() else "cpu" + +folder_paths.add_model_folder_path("rmbg", os.path.join(folder_paths.models_dir, "RMBG")) + +AVAILABLE_MODELS = { + "RMBG-2.0": { + "type": "rmbg", + "repo_id": "1038lab/RMBG-2.0", + "files": { + "config.json": "config.json", + "model.safetensors": "model.safetensors", + "birefnet.py": "birefnet.py", + "BiRefNet_config.py": "BiRefNet_config.py" + }, + "cache_dir": "RMBG-2.0" + }, + "INSPYRENET": { + "type": "inspyrenet", + "repo_id": "1038lab/inspyrenet", + "files": { + "inspyrenet.safetensors": "inspyrenet.safetensors" + }, + "cache_dir": "INSPYRENET" + }, + "BEN": { + "type": "ben", + "repo_id": "1038lab/BEN", + "files": { + "model.py": "model.py", + "BEN_Base.pth": "BEN_Base.pth" + }, + "cache_dir": "BEN" + }, + "BEN2": { + "type": "ben2", + "repo_id": "1038lab/BEN2", + "files": { + "BEN2_Base.pth": "BEN2_Base.pth", + "BEN2.py": "BEN2.py" + }, + "cache_dir": "BEN2" + } +} + +def tensor2pil(image): + return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)) + +def pil2tensor(image): + return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0) + +def handle_model_error(message): + print(f"[RMBG ERROR] {message}") + raise RuntimeError(message) + +class BaseModelLoader: + def __init__(self): + self.model = None + self.current_model_version = None + self.base_cache_dir = os.path.join(folder_paths.models_dir, "RMBG") + + def get_cache_dir(self, model_name): + cache_path = os.path.join(self.base_cache_dir, AVAILABLE_MODELS[model_name]["cache_dir"]) + os.makedirs(cache_path, exist_ok=True) + return cache_path + + def check_model_cache(self, model_name): + model_info = AVAILABLE_MODELS[model_name] + cache_dir = self.get_cache_dir(model_name) + + if not os.path.exists(cache_dir): + return False, "Model directory not found" + + missing_files = [] + for filename in model_info["files"].keys(): + if not os.path.exists(os.path.join(cache_dir, model_info["files"][filename])): + missing_files.append(filename) + + if missing_files: + return False, f"Missing model files: {', '.join(missing_files)}" + + return True, "Model cache verified" + + def download_model(self, model_name): + model_info = AVAILABLE_MODELS[model_name] + cache_dir = self.get_cache_dir(model_name) + + try: + os.makedirs(cache_dir, exist_ok=True) + print(f"Downloading {model_name} model files...") + + for filename in model_info["files"].keys(): + print(f"Downloading {filename}...") + hf_hub_download( + repo_id=model_info["repo_id"], + filename=filename, + local_dir=cache_dir + ) + + return True, "Model files downloaded successfully" + + except Exception as e: + return False, f"Error downloading model files: {str(e)}" + + def clear_model(self): + if self.model is not None: + self.model.cpu() + del self.model + + import gc + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + self.model = None + self.current_model_version = None + +class RMBGModel(BaseModelLoader): + def __init__(self): + super().__init__() + + def load_model(self, model_name): + if self.current_model_version != model_name: + self.clear_model() + + cache_dir = self.get_cache_dir(model_name) + try: + try: + self.model = AutoModelForImageSegmentation.from_pretrained( + cache_dir, + trust_remote_code=True, + local_files_only=True + ) + except AttributeError as ae: + if "'Config' object has no attribute 'get_text_config'" in str(ae): + print("[RMBG WARNING] Detected newer transformers version, using compatibility mode...") + try: + from transformers import PreTrainedModel + import json + + config_path = os.path.join(cache_dir, "config.json") + with open(config_path, 'r') as f: + config = json.load(f) + + birefnet_path = os.path.join(cache_dir, "birefnet.py") + BiRefNetConfig_path = os.path.join(cache_dir, "BiRefNet_config.py") + + config_spec = importlib.util.spec_from_file_location("BiRefNetConfig", BiRefNetConfig_path) + config_module = importlib.util.module_from_spec(config_spec) + sys.modules["BiRefNetConfig"] = config_module + config_spec.loader.exec_module(config_module) + + with open(birefnet_path, 'r') as f: + birefnet_content = f.read() + + birefnet_content = birefnet_content.replace( + "from .BiRefNet_config import BiRefNetConfig", + "from BiRefNetConfig import BiRefNetConfig" + ) + + module_name = f"custom_birefnet_model_{hash(birefnet_path)}" + module = types.ModuleType(module_name) + sys.modules[module_name] = module + exec(birefnet_content, module.__dict__) + + for attr_name in dir(module): + attr = getattr(module, attr_name) + if isinstance(attr, type) and issubclass(attr, PreTrainedModel) and attr != PreTrainedModel: + BiRefNetConfig = getattr(config_module, "BiRefNetConfig") + model_config = BiRefNetConfig() + self.model = attr(model_config) + + weights_path = os.path.join(cache_dir, "model.safetensors") + try: + try: + import safetensors.torch + self.model.load_state_dict(safetensors.torch.load_file(weights_path)) + except ImportError: + from transformers.modeling_utils import load_state_dict + state_dict = load_state_dict(weights_path) + self.model.load_state_dict(state_dict) + except Exception as load_error: + pytorch_weights = os.path.join(cache_dir, "pytorch_model.bin") + if os.path.exists(pytorch_weights): + self.model.load_state_dict(torch.load(pytorch_weights, map_location="cpu")) + else: + raise RuntimeError(f"Failed to load weights: {str(load_error)}") + break + + if self.model is None: + raise RuntimeError("Could not find suitable model class") + + except Exception as custom_e: + handle_model_error(f"Failed to load model in compatibility mode: {str(custom_e)}") + else: + raise ae + except Exception as e: + handle_model_error(f"Error loading model: {str(e)}") + + self.model.eval() + for param in self.model.parameters(): + param.requires_grad = False + + torch.set_float32_matmul_precision('high') + self.model.to(device) + self.current_model_version = model_name + + def process_image(self, images, model_name, params): + try: + self.load_model(model_name) + + transform_image = transforms.Compose([ + transforms.Resize((params["process_res"], params["process_res"])), + transforms.ToTensor(), + transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) + ]) + + if isinstance(images, torch.Tensor): + if len(images.shape) == 3: + images = [images] + else: + images = [img for img in images] + + original_sizes = [tensor2pil(img).size for img in images] + + input_tensors = [transform_image(tensor2pil(img)).unsqueeze(0) for img in images] + input_batch = torch.cat(input_tensors, dim=0).to(device) + + with torch.no_grad(): + outputs = self.model(input_batch) + + if isinstance(outputs, list) and len(outputs) > 0: + results = outputs[-1].sigmoid().cpu() + elif isinstance(outputs, dict) and 'logits' in outputs: + results = outputs['logits'].sigmoid().cpu() + elif isinstance(outputs, torch.Tensor): + results = outputs.sigmoid().cpu() + else: + try: + if hasattr(outputs, 'last_hidden_state'): + results = outputs.last_hidden_state.sigmoid().cpu() + else: + for k, v in outputs.items(): + if isinstance(v, torch.Tensor): + results = v.sigmoid().cpu() + break + except: + handle_model_error("Unable to recognize model output format") + + masks = [] + + for i, (result, (orig_w, orig_h)) in enumerate(zip(results, original_sizes)): + result = result.squeeze() + result = result * (1 + (1 - params["sensitivity"])) + result = torch.clamp(result, 0, 1) + + result = F.interpolate(result.unsqueeze(0).unsqueeze(0), + size=(orig_h, orig_w), + mode='bilinear').squeeze() + + masks.append(tensor2pil(result)) + + return masks + + except Exception as e: + handle_model_error(f"Error in batch processing: {str(e)}") + +class InspyrenetModel(BaseModelLoader): + def __init__(self): + super().__init__() + + def load_model(self, model_name): + if self.current_model_version != model_name: + self.clear_model() + + try: + import transparent_background + self.model = transparent_background.Remover() + self.current_model_version = model_name + except ImportError: + try: + import pip + pip.main(['install', 'transparent_background']) + import transparent_background + self.model = transparent_background.Remover() + self.current_model_version = model_name + except Exception as e: + handle_model_error(f"Failed to install transparent_background: {str(e)}") + + def process_image(self, image, model_name, params): + try: + self.load_model(model_name) + + orig_image = tensor2pil(image) + w, h = orig_image.size + + # Resize for processing + aspect_ratio = h / w + new_w = params["process_res"] + new_h = int(params["process_res"] * aspect_ratio) + resized_image = orig_image.resize((new_w, new_h), Image.LANCZOS) + + # Process image + foreground = self.model.process(resized_image, type='rgba') + foreground = foreground.resize((w, h), Image.LANCZOS) + mask = foreground.split()[-1] + + return mask + + except Exception as e: + handle_model_error(f"Error in Inspyrenet processing: {str(e)}") + +class BENModel(BaseModelLoader): + def __init__(self): + super().__init__() + + def load_model(self, model_name): + if self.current_model_version != model_name: + self.clear_model() + + cache_dir = self.get_cache_dir(model_name) + model_path = os.path.join(cache_dir, "model.py") + module_name = f"custom_ben_model_{hash(model_path)}" + + spec = importlib.util.spec_from_file_location(module_name, model_path) + ben_module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = ben_module + spec.loader.exec_module(ben_module) + + model_weights_path = os.path.join(cache_dir, "BEN_Base.pth") + self.model = ben_module.BEN_Base() + self.model.loadcheckpoints(model_weights_path) + + self.model.eval() + for param in self.model.parameters(): + param.requires_grad = False + + torch.set_float32_matmul_precision('high') + self.model.to(device) + self.current_model_version = model_name + + def process_image(self, image, model_name, params): + try: + self.load_model(model_name) + + orig_image = tensor2pil(image) + w, h = orig_image.size + + aspect_ratio = h / w + new_w = params["process_res"] + new_h = int(params["process_res"] * aspect_ratio) + resized_image = orig_image.resize((new_w, new_h), Image.LANCZOS) + + processed_input = resized_image.convert("RGBA") + + with torch.no_grad(): + _, foreground = self.model.inference(processed_input) + + foreground = foreground.resize((w, h), Image.LANCZOS) + mask = foreground.split()[-1] + + return mask + + except Exception as e: + handle_model_error(f"Error in BEN processing: {str(e)}") + +class BEN2Model(BaseModelLoader): + def __init__(self): + super().__init__() + + def load_model(self, model_name): + if self.current_model_version != model_name: + self.clear_model() + + try: + cache_dir = self.get_cache_dir(model_name) + model_path = os.path.join(cache_dir, "BEN2.py") + module_name = f"custom_ben2_model_{hash(model_path)}" + + spec = importlib.util.spec_from_file_location(module_name, model_path) + ben2_module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = ben2_module + spec.loader.exec_module(ben2_module) + + model_weights_path = os.path.join(cache_dir, "BEN2_Base.pth") + self.model = ben2_module.BEN_Base() + self.model.loadcheckpoints(model_weights_path) + + self.model.eval() + for param in self.model.parameters(): + param.requires_grad = False + + torch.set_float32_matmul_precision('high') + self.model.to(device) + self.current_model_version = model_name + + except Exception as e: + handle_model_error(f"Error loading BEN2 model: {str(e)}") + + def process_image(self, images, model_name, params): + try: + self.load_model(model_name) + + if isinstance(images, torch.Tensor): + if len(images.shape) == 3: + images = [images] + else: + images = [img for img in images] + + batch_size = 3 + all_masks = [] + + for i in range(0, len(images), batch_size): + batch_images = images[i:i + batch_size] + batch_pil_images = [] + original_sizes = [] + + for img in batch_images: + orig_image = tensor2pil(img) + w, h = orig_image.size + original_sizes.append((w, h)) + + aspect_ratio = h / w + new_w = params["process_res"] + new_h = int(params["process_res"] * aspect_ratio) + resized_image = orig_image.resize((new_w, new_h), Image.LANCZOS) + processed_input = resized_image.convert("RGBA") + batch_pil_images.append(processed_input) + + with torch.no_grad(): + try: + foregrounds = self.model.inference(batch_pil_images) + if not isinstance(foregrounds, list): + foregrounds = [foregrounds] + except Exception as e: + handle_model_error(f"Error in BEN2 inference: {str(e)}") + + for foreground, (orig_w, orig_h) in zip(foregrounds, original_sizes): + foreground = foreground.resize((orig_w, orig_h), Image.LANCZOS) + mask = foreground.split()[-1] + all_masks.append(mask) + + if len(all_masks) == 1: + return all_masks[0] + return all_masks + + except Exception as e: + handle_model_error(f"Error in BEN2 processing: {str(e)}") + +def refine_foreground(image_bchw, masks_b1hw): + b, c, h, w = image_bchw.shape + if b != masks_b1hw.shape[0]: + raise ValueError("images and masks must have the same batch size") + + image_np = image_bchw.cpu().numpy() + mask_np = masks_b1hw.cpu().numpy() + + refined_fg = [] + for i in range(b): + mask = mask_np[i, 0] + thresh = 0.45 + mask_binary = (mask > thresh).astype(np.float32) + + edge_blur = cv2.GaussianBlur(mask_binary, (3, 3), 0) + transition_mask = np.logical_and(mask > 0.05, mask < 0.95) + + alpha = 0.85 + mask_refined = np.where(transition_mask, + alpha * mask + (1-alpha) * edge_blur, + mask_binary) + + edge_region = np.logical_and(mask > 0.2, mask < 0.8) + mask_refined = np.where(edge_region, + mask_refined * 0.98, + mask_refined) + + result = [] + for c in range(image_np.shape[1]): + channel = image_np[i, c] + refined = channel * mask_refined + result.append(refined) + + refined_fg.append(np.stack(result)) + + return torch.from_numpy(np.stack(refined_fg)) + +class ImageRemoveBackgroundRMBG: + def __init__(self): + self.models = { + "RMBG-2.0": RMBGModel(), + "INSPYRENET": InspyrenetModel(), + "BEN": BENModel(), + "BEN2": BEN2Model() + } + + @classmethod + def INPUT_TYPES(s): + tooltips = { + "image": "Input image to be processed for background removal.", + "model": "Select the background removal model to use (RMBG-2.0, INSPYRENET, BEN).", + "sensitivity": "Adjust the strength of mask detection (higher values result in more aggressive detection).", + "process_res": "Set the processing resolution (higher values require more VRAM and may increase processing time).", + "mask_blur": "Specify the amount of blur to apply to the mask edges (0 for no blur, higher values for more blur).", + "mask_offset": "Adjust the mask boundary (positive values expand the mask, negative values shrink it).", + "background": "Choose output type: Alpha (transparent) or Color (custom background color).", + "background_color": "Pick background color (supports alpha, use color picker).", + "invert_output": "Enable to invert both the image and mask output (useful for certain effects).", + "refine_foreground": "Use Fast Foreground Colour Estimation to optimize transparent background" + } + + return { + "required": { + "image": ("IMAGE", {"tooltip": tooltips["image"]}), + "model": (list(AVAILABLE_MODELS.keys()), {"tooltip": tooltips["model"]}), + }, + "optional": { + "sensitivity": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": tooltips["sensitivity"]}), + "process_res": ("INT", {"default": 1024, "min": 256, "max": 2048, "step": 8, "tooltip": tooltips["process_res"]}), + "mask_blur": ("INT", {"default": 0, "min": 0, "max": 64, "step": 1, "tooltip": tooltips["mask_blur"]}), + "mask_offset": ("INT", {"default": 0, "min": -64, "max": 64, "step": 1, "tooltip": tooltips["mask_offset"]}), + "invert_output": ("BOOLEAN", {"default": False, "tooltip": tooltips["invert_output"]}), + "refine_foreground": ("BOOLEAN", {"default": False, "tooltip": tooltips["refine_foreground"]}), + "background": (["Alpha", "Color"], {"default": "Alpha", "tooltip": tooltips["background"]}), + "background_color": ("COLOR", {"default": "#222222", "tooltip": tooltips["background_color"]}), + } + } + + RETURN_TYPES = ("IMAGE", "MASK", "IMAGE") + RETURN_NAMES = ("IMAGE", "MASK", "MASK_IMAGE") + FUNCTION = "process_image" + CATEGORY = "LevelPixel/Image" + + def process_image(self, image, model, **params): + try: + processed_images = [] + processed_masks = [] + + model_instance = self.models[model] + + cache_status, message = model_instance.check_model_cache(model) + if not cache_status: + print(f"Cache check: {message}") + print("Downloading required model files...") + download_status, download_message = model_instance.download_model(model) + if not download_status: + handle_model_error(download_message) + print("Model files downloaded successfully") + + for img in image: + mask = model_instance.process_image(img, model, params) + + if isinstance(mask, list): + masks = [m.convert("L") for m in mask if isinstance(m, Image.Image)] + mask = masks[0] if masks else None + elif isinstance(mask, Image.Image): + mask = mask.convert("L") + + mask_tensor = pil2tensor(mask) + mask_tensor = mask_tensor * (1 + (1 - params["sensitivity"])) + mask_tensor = torch.clamp(mask_tensor, 0, 1) + mask = tensor2pil(mask_tensor) + + if params["mask_blur"] > 0: + mask = mask.filter(ImageFilter.GaussianBlur(radius=params["mask_blur"])) + + if params["mask_offset"] != 0: + if params["mask_offset"] > 0: + for _ in range(params["mask_offset"]): + mask = mask.filter(ImageFilter.MaxFilter(3)) + else: + for _ in range(-params["mask_offset"]): + mask = mask.filter(ImageFilter.MinFilter(3)) + + if params["invert_output"]: + mask = Image.fromarray(255 - np.array(mask)) + + img_tensor = torch.from_numpy(np.array(tensor2pil(img))).permute(2, 0, 1).unsqueeze(0) / 255.0 + mask_tensor = torch.from_numpy(np.array(mask)).unsqueeze(0).unsqueeze(0) / 255.0 + + orig_image = tensor2pil(img) + + if params.get("refine_foreground", False): + refined_fg = refine_foreground(img_tensor, mask_tensor) + refined_fg = tensor2pil(refined_fg[0].permute(1, 2, 0)) + r, g, b = refined_fg.split() + foreground = Image.merge('RGBA', (r, g, b, mask)) + else: + orig_rgba = orig_image.convert("RGBA") + r, g, b, _ = orig_rgba.split() + foreground = Image.merge('RGBA', (r, g, b, mask)) + + if params["background"] == "Color": + def hex_to_rgba(hex_color): + hex_color = hex_color.lstrip('#') + if len(hex_color) == 6: + r, g, b = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16) + a = 255 + elif len(hex_color) == 8: + r, g, b, a = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16), int(hex_color[6:8], 16) + else: + raise ValueError("Invalid color format") + return (r, g, b, a) + rgba = hex_to_rgba(params["background_color"]) + bg_image = Image.new('RGBA', orig_image.size, rgba) + composite_image = Image.alpha_composite(bg_image, foreground) + processed_images.append(pil2tensor(composite_image.convert("RGB"))) + else: + processed_images.append(pil2tensor(foreground)) + + processed_masks.append(pil2tensor(mask)) + + mask_images = [] + for mask_tensor in processed_masks: + mask_image = mask_tensor.reshape((-1, 1, mask_tensor.shape[-2], mask_tensor.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3) + mask_images.append(mask_image) + + mask_image_output = torch.cat(mask_images, dim=0) + + return (torch.cat(processed_images, dim=0), torch.cat(processed_masks, dim=0), mask_image_output) + + except Exception as e: + handle_model_error(f"Error in image processing: {str(e)}") + empty_mask = torch.zeros((image.shape[0], image.shape[2], image.shape[3])) + empty_mask_image = empty_mask.reshape((-1, 1, empty_mask.shape[-2], empty_mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3) + return (image, empty_mask, empty_mask_image) + +NODE_CLASS_MAPPINGS = { + "ImageRemoveBackgroundRMBG|LP": ImageRemoveBackgroundRMBG +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ImageRemoveBackgroundRMBG|LP": "Image Remove Background (RMBG) [LP]", +} \ No newline at end of file diff --git a/nodes/vlm/llama-mtmd-cli/where_to_download_the_library.txt b/nodes/vlm/llama-mtmd-cli/where_to_download_the_library.txt new file mode 100644 index 0000000..96bdd20 --- /dev/null +++ b/nodes/vlm/llama-mtmd-cli/where_to_download_the_library.txt @@ -0,0 +1,14 @@ +To make this node work, you need to download the Llama.cpp libraries manually from this link: +https://github.com/ggml-org/llama.cpp/releases + +Download the archive that suits your system. Probably, for you it will be llama-b5317-bin-win--cu12.4-x64.zip + +After downloading, go to the path: +ComfyUI\custom_nodes\ComfyUI-LevelPixel-Advanced\nodes\vlm\llama-mtmd-cli\ + +And UNZIP the downloaded archive here. + +The llama-mtmd-cli folder should contain the llama-mtmd-cli.exe file and other libraries from the archive: +ComfyUI\custom_nodes\ComfyUI-LevelPixel-Advanced\nodes\vlm\llama-mtmd-cli\llama-mtmd-cli.exe + +After that, run the generation again. \ No newline at end of file diff --git a/nodes/vlm/llava_LP.py b/nodes/vlm/llava_LP.py index 7c4bd1e..4556374 100644 --- a/nodes/vlm/llava_LP.py +++ b/nodes/vlm/llava_LP.py @@ -18,7 +18,7 @@ except: os.mkdir(os.path.join(folder_paths.models_dir, "LLavacheckpoints")) folder_paths.folder_names_and_paths["LLavacheckpoints"] = ([os.path.join(folder_paths.models_dir, "LLavacheckpoints")], supported_LLava_extensions) - + class LLavaLoader: @classmethod def INPUT_TYPES(s): @@ -290,22 +290,21 @@ class LLavaAdvanced: def generate_text_full_advanced(self, ckpt_name, clip_name, max_ctx, gpu_layers, n_threads, image, system_msg, prompt, max_tokens, temperature, top_p, top_k, frequency_penalty, presence_penalty, repeat_penalty, seed, unload): - + clip_path = folder_paths.get_full_path("LLavacheckpoints", clip_name) self.clip = Llava16ChatHandler(clip_model_path=clip_path, verbose=False) - + ckpt_path = folder_paths.get_full_path("LLavacheckpoints", ckpt_name) self.llm = Llama(model_path = ckpt_path, chat_handler=self.clip, offload_kqv=True, f16_kv=True, - use_mlock=False, embedding=False, n_batch=1024, last_n_tokens_size=1024, - verbose=True, seed=42, n_ctx = max_ctx, n_gpu_layers=gpu_layers, n_threads=n_threads, - logits_all=True, echo=False) - + use_mlock=False, embedding=False, n_batch=1024, last_n_tokens_size=1024, + verbose=True, seed=42, n_ctx = max_ctx, n_gpu_layers=gpu_layers, n_threads=n_threads, + logits_all=True, echo=False) + pil_image = ToPILImage()(image[0].permute(2, 0, 1)) buffer = BytesIO() pil_image.save(buffer, format="PNG") - image_bytes = buffer.getvalue() base64_string = f"data:image/jpeg;base64,{base64.b64encode(image_bytes).decode('utf-8')}" @@ -347,6 +346,129 @@ class LLavaAdvanced: torch.cuda.empty_cache() return (f"{response['choices'][0]['message']['content']}", ) + +class MultimodalGeneratorAdvanced: + def __init__(self): + self.llm = None + self.clip = None + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "ckpt_name": (folder_paths.get_filename_list("LLavacheckpoints"), ), + "clip_name": (folder_paths.get_filename_list("LLavacheckpoints"), ), + "max_ctx": ("INT", {"default": 4096, "min": 128, "max": 128000, "step": 64}), + "gpu_layers": ("INT", {"default": 27, "min": 0, "max": 100, "step": 1}), + "n_threads": ("INT", {"default": 8, "min": 1, "max": 100, "step": 1}), + "max_tokens": ("INT", {"default": 512, "min": 1, "max": 2048, "step": 1}), + "temperature": ("FLOAT", {"default": 0.1, "min": 0.01, "max": 1.0, "step": 0.01}), + "top_p": ("FLOAT", {"default": 0.95, "min": 0.1, "max": 1.0, "step": 0.01}), + "top_k": ("INT", {"default": 40, "step": 1}), + "frequency_penalty": ("FLOAT", {"default": 0.0, "step": 0.01}), + "presence_penalty": ("FLOAT", {"default": 0.0, "step": 0.01}), + "repeat_penalty": ("FLOAT", {"default": 1.1, "step": 0.01}), + "seed": ("INT", {"default": 42, "step": 1}), + #"unload": ("BOOLEAN", {"default": False}), + "prompt": ("STRING", {"multiline": True, "default": ""}), + #"system_msg": ("STRING", {"multiline": True, "default": "You are an assistant who perfectly describes images."}), + } + } + + RETURN_TYPES = ("STRING",) + FUNCTION = "multimodal_generator_advanced" + CATEGORY = "LevelPixel/VLM" + + def multimodal_generator_advanced(self, ckpt_name, clip_name, max_ctx, gpu_layers, n_threads, image, + #system_msg, + prompt, max_tokens, temperature, top_p, top_k, frequency_penalty, + presence_penalty, repeat_penalty, seed): #, unload): + + clip_path = folder_paths.get_full_path("LLavacheckpoints", clip_name) + ckpt_path = folder_paths.get_full_path("LLavacheckpoints", ckpt_name) + + import tempfile + + pil_image = ToPILImage()(image[0].permute(2, 0, 1)) + + buffer = BytesIO() + pil_image.save(buffer, format="PNG") + + with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp: + temp_image_path = tmp.name + pil_image.save(temp_image_path) + + import subprocess + from pathlib import Path + + relative_cli_path="llama-mtmd-cli\\llama-mtmd-cli.exe" + cli = Path(__file__).parent / relative_cli_path + if not cli.is_file(): + raise FileNotFoundError(f"""Library llama-mtmd-cli.exe not founded: {cli} + + SOLUTION: + + To make this node work, you need to download the Llama.cpp libraries manually from this link: + https://github.com/ggml-org/llama.cpp/releases + + Download the archive that suits your system. Probably, for you it will be llama-b5317-bin-win--cu12.4-x64.zip + + After downloading, go to the path: + ComfyUI\\custom_nodes\\ComfyUI-LevelPixel-Advanced\\nodes\\vlm\\llama-mtmd-cli + + And unzip the downloaded archive here. + + The llama-mtmd-cli folder should contain the llama-mtmd-cli.exe file and other libraries from the archive. + + After that, run the generation again.""") + + cmd = [ + str(cli), + '-m', ckpt_path, + '--mmproj', clip_path, + '-p', prompt, + '--image', temp_image_path, + '--ctx-size', max_ctx, + '--gpu-layers', gpu_layers, + '--threads', n_threads, + '--n-predict', max_tokens, + '--temp', temperature, + '--top-p', top_p, + '--top-k', top_k, + '--frequency-penalty', frequency_penalty, + '--presence-penalty', presence_penalty, + '--repeat-penalty', repeat_penalty, + '--seed', seed, + #'-sys', system_msg + ] + + cmd = [str(x) for x in cmd] + + print(cmd) + + print("Generate by llama-mtmd-cli") + result = subprocess.run( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + shell=True + ) + + output = result.stdout.decode('utf-8').strip() + + def parse_output(output: str) -> str: + lines = output.strip().splitlines() + if len(lines) > 1: + return '\n'.join(lines[1:]).strip() + return output.strip() + + clean_output = parse_output(output) + + if os.path.exists(temp_image_path): + os.remove(temp_image_path) + + return (clean_output, ) NODE_CLASS_MAPPINGS = { "LLavaLoader|LP": LLavaLoader, @@ -355,6 +477,7 @@ NODE_CLASS_MAPPINGS = { "LLavaSamplerAdvanced|LP": LLavaSamplerAdvanced, "LLavaSimple|LP": LLavaSimple, "LLavaAdvanced|LP": LLavaAdvanced, + "MultimodalGeneratorAdvanced|LP": MultimodalGeneratorAdvanced, } NODE_DISPLAY_NAME_MAPPINGS = { @@ -364,4 +487,5 @@ NODE_DISPLAY_NAME_MAPPINGS = { "LLavaSamplerAdvanced|LP": "LLava Sampler Advanced [LP]", "LLavaSimple|LP": "LLava Simple [LP]", "LLavaAdvanced|LP": "LLava Advanced [LP]", + "MultimodalGeneratorAdvanced|LP": "Multimodal Generator Advanced [LP]", } diff --git a/nodes/vlm/ram_LP.py b/nodes/vlm/ram_LP.py index a135947..aa5d906 100644 --- a/nodes/vlm/ram_LP.py +++ b/nodes/vlm/ram_LP.py @@ -2,6 +2,9 @@ from PIL import Image import torch import gc import numpy as np +import warnings +warnings.filterwarnings("ignore", category=FutureWarning, module="timm\\.models") + from ram.models import ram from ram.models import ram_plus from ram.models import tag2text diff --git a/pyproject.toml b/pyproject.toml index 2f3489a..ca2a560 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,9 +1,22 @@ [project] name = "comfyui-levelpixel-advanced" -description = "Advanced nodes of the Level Pixel company (aka levelpixel, LP). Includes convenient advanced nodes for working with LLM и VLM models (LLaVa) with GGUF format. Qwen2.5-VL and Qwen2.5 supported. Also included is a node for the RAM model. Nodes have the ability to automatically unload models from VRAM." +description = "Advanced nodes of the [Level Pixel](https://github.com/LevelPixel/ComfyUI-LevelPixel-Advanced) company (aka levelpixel, LP). Includes convenient advanced nodes for working with LLM и VLM models (LLaVa, Multimodal Generator) with GGUF format. Qwen2.5-VL and Qwen2.5 supported, based on llama-mtmd. Also included is a node for the RAM model. Nodes have the ability to automatically unload models from VRAM. This package also includes Autotagger (WD), Image Remove Background (based on RMBG-2.0, on BiRefNet, on BiRefNet-HR-Matter, on RemBG)." version = "1.2.3" license = { file = "LICENSE" } -dependencies = ["torch>=2.0.1", "torchvision>=0.15.2", "transformers>=4.46", "pillow>=9.4.0", "numpy>=1.26.4", "matplotlib", "scikit-build-core>=0.10.7", "timm>=0.4.12", "fairscale>=0.4.4"] +dependencies = [ + "torch>=2.0.1", + "torchvision>=0.15.2", + "pillow>=9.4.0", + "numpy>=1.26.4", + "huggingface-hub>=0.19.0", + "transformers>=4.35.0", + "safetensors>=0.3.0", + "scikit-build-core>=0.10.7", + "transparent-background>=1.2.4", + "timm>=0.4.12", + "tqdm>=4.65.0", + "fairscale>=0.4.4", + "rembg>=2.0.66" ] [project.urls] Repository = "https://github.com/LevelPixel/ComfyUI-LevelPixel-Advanced" diff --git a/requirements.txt b/requirements.txt index 6b720cd..cc5218c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,8 +2,14 @@ torch>=2.0.1 torchvision>=0.15.2 pillow>=9.4.0 numpy>=1.26.4 -matplotlib +huggingface-hub>=0.19.0 +transformers>=4.35.0 +safetensors>=0.3.0 scikit-build-core>=0.10.7 +transparent-background>=1.2.4 timm>=0.4.12 +tqdm>=4.65.0 fairscale>=0.4.4 +rembg>=2.0.66 +matplotlib ram @ git+https://github.com/xinyu1205/recognize-anything.git