diff --git a/README.md b/README.md index 362bf3a..3fbd10b 100644 --- a/README.md +++ b/README.md @@ -782,10 +782,15 @@ On the basis of SegmentAnythingUltra, the following changes have been made: * device: Set whether the VitMatte to use cuda. * max_megapixels: Set the maximum size for VitMate operations. + ### SegmentAnythingUltraV3 Separate model loading from inference nodes to avoid duplicate model loading when using multiple SAM nodes. - +![image](image/segment_anything_ultra_v3_example.jpg) + +Node Options: ![image](image/segment_anything_ultra_v3_node.jpg) +Same as SegmentAnythingUltra, removed ```sam_comodel``` and ```ground-dino_comodel```, changed them to be obtained from node input. + ### LoadSegmentAnythingModels Load SegmentAnything models. diff --git a/README_CN.MD b/README_CN.MD index f5464be..7e6ad95 100644 --- a/README_CN.MD +++ b/README_CN.MD @@ -716,8 +716,11 @@ SegmentAnythingUltra的V2升级版，增加了VITMatte边缘处理方法。 ### SegmentAnythingUltraV3 将模型加载与推理节点分离，在使用多个SAM节点时避免重复加载模型。 - +![image](image/segment_anything_ultra_v3_example.jpg) + +节点选项说明: ![image](image/segment_anything_ultra_v3_node.jpg) +选项与SegmentAnythingUltra节点一致，去掉了sam_model和ground_dino_model改为从节点输入获取。 ### LoadSegmentAnythingModels 加载SegmentAnything模型。 diff --git a/image/segment_anything_ultra_v3_example.jpg b/image/segment_anything_ultra_v3_example.jpg new file mode 100644 index 0000000..4944e25 Binary files /dev/null and b/image/segment_anything_ultra_v3_example.jpg differ diff --git a/py/segment_anything_ultra_v3.py b/py/segment_anything_ultra_v3.py new file mode 100644 index 0000000..8d7bbb8 --- /dev/null +++ b/py/segment_anything_ultra_v3.py @@ -0,0 +1,130 @@ +# layerstyle advance + +from .imagefunc import * +from .segment_anything_func import * + + +class LS_LoadSAMModels: + def __init__(self): + self.NODE_NAME = 'SegmentAnythingUltra V3' + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "sam_model": (list_sam_model(), ), + "grounding_dino_model": (list_groundingdino_model(),), + }, + "optional": { + } + } + + RETURN_TYPES = ("LS_SAM_MODELS",) + RETURN_NAMES = ("sam_models", ) + FUNCTION = "load_sam_models" + CATEGORY = '😺dzNodes/LayerMask' + + def load_sam_models(self, sam_model, grounding_dino_model): + + + SAM_MODEL = load_sam_model(sam_model) + DINO_MODEL = load_groundingdino_model(grounding_dino_model) + + return ({"SAM_MODEL":SAM_MODEL, "DINO_MODEL":DINO_MODEL},) + + +class LS_SegmentAnythingUltraV3: + def __init__(self): + self.NODE_NAME = 'SegmentAnythingUltra V3' + + @classmethod + def INPUT_TYPES(cls): + method_list = ['VITMatte', 'VITMatte(local)', 'PyMatting', 'GuidedFilter', ] + device_list = ['cuda','cpu'] + return { + "required": { + "image": ("IMAGE",), + "sam_models": ("LS_SAM_MODELS", ), + "threshold": ("FLOAT", {"default": 0.3, "min": 0, "max": 1.0, "step": 0.01}), + "detail_method": (method_list,), + "detail_erode": ("INT", {"default": 6, "min": 1, "max": 255, "step": 1}), + "detail_dilate": ("INT", {"default": 6, "min": 1, "max": 255, "step": 1}), + "black_point": ("FLOAT", {"default": 0.15, "min": 0.01, "max": 0.98, "step": 0.01, "display": "slider"}), + "white_point": ("FLOAT", {"default": 0.99, "min": 0.02, "max": 0.99, "step": 0.01, "display": "slider"}), + "process_detail": ("BOOLEAN", {"default": True}), + "prompt": ("STRING", {"default": "subject"}), + "device": (device_list,), + "max_megapixels": ("FLOAT", {"default": 2.0, "min": 1, "max": 999, "step": 0.1}), + }, + "optional": { + } + } + + RETURN_TYPES = ("IMAGE", "MASK",) + RETURN_NAMES = ("image", "mask",) + FUNCTION = "segment_anything_ultra_v3" + CATEGORY = '😺dzNodes/LayerMask' + + def segment_anything_ultra_v3(self, image, sam_models, threshold, + detail_method, detail_erode, detail_dilate, + black_point, white_point, process_detail, prompt, + device, max_megapixels, + ): + + if detail_method == 'VITMatte(local)': + local_files_only = True + else: + local_files_only = False + + + + SAM_MODEL = sam_models["SAM_MODEL"] + DINO_MODEL = sam_models["DINO_MODEL"] + + ret_images = [] + ret_masks = [] + + for i in image: + i = torch.unsqueeze(i, 0) + i = pil2tensor(tensor2pil(i).convert('RGB')) + _image = tensor2pil(i).convert('RGBA') + boxes = groundingdino_predict(DINO_MODEL, _image, prompt, threshold) + if boxes.shape[0] == 0: + break + (_, _mask) = sam_segment(SAM_MODEL, _image, boxes) + _mask = _mask[0] + detail_range = detail_erode + detail_dilate + if process_detail: + if detail_method == 'GuidedFilter': + _mask = guided_filter_alpha(i, _mask, detail_range // 6 + 1) + _mask = tensor2pil(histogram_remap(_mask, black_point, white_point)) + elif detail_method == 'PyMatting': + _mask = tensor2pil(mask_edge_detail(i, _mask, detail_range // 8 + 1, black_point, white_point)) + else: + _trimap = generate_VITMatte_trimap(_mask, detail_erode, detail_dilate) + _mask = generate_VITMatte(_image, _trimap, local_files_only=local_files_only, device=device, max_megapixels=max_megapixels) + _mask = tensor2pil(histogram_remap(pil2tensor(_mask), black_point, white_point)) + else: + _mask = mask2image(_mask) + _image = RGB2RGBA(tensor2pil(i).convert('RGB'), _mask.convert('L')) + + ret_images.append(pil2tensor(_image)) + ret_masks.append(image2mask(_mask)) + if len(ret_masks) == 0: + _, height, width, _ = image.size() + empty_mask = torch.zeros((1, height, width), dtype=torch.uint8, device="cpu") + return (empty_mask, empty_mask) + + + log(f"{self.NODE_NAME} Processed {len(ret_masks)} image(s).", message_type='finish') + return (torch.cat(ret_images, dim=0), torch.cat(ret_masks, dim=0),) + +NODE_CLASS_MAPPINGS = { + "LayerMask: SegmentAnythingUltra V3": LS_SegmentAnythingUltraV3, + "LayerMask: LoadSegmentAnythingModels": LS_LoadSAMModels, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "LayerMask: SegmentAnythingUltra V3": "LayerMask: SegmentAnythingUltra V3(Advance)", + "LayerMask: LoadSegmentAnythingModels": "LayerMask: Load SegmentAnything Models(Advance)", +} diff --git a/py/zhipu_glm4.py b/py/zhipu_glm4.py new file mode 100644 index 0000000..9499be4 --- /dev/null +++ b/py/zhipu_glm4.py @@ -0,0 +1,124 @@ +# layerstyle advance +import os +import torch +import base64 +import requests +from io import BytesIO +from zhipuai import ZhipuAI +import folder_paths +from PIL import Image +from .imagefunc import log, tensor2pil, get_api_key + + +# apikey申请地址：https://bigmodel.cn/usercenter/proj-mgmt/apikeys +class LS_ZhipuImage: + + def __init__(self): + self.NODE_NAME = 'ZhipuGLM4V' + + @classmethod + def INPUT_TYPES(cls): + glm_model_list = ["glm-4v-flash", "glm-4v", "glm-4v-plus"] + return {"required":{ + "image": ("IMAGE",), + "model": (glm_model_list,), + "user_prompt": ("STRING", {"default": "describe this image", "multiline": True}), + }, + "optional": { + } + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("text",) + FUNCTION = "zhipu_glm4v" + CATEGORY = '😺dzNodes/LayerUtility' + + def zhipu_glm4v(self, image, model, user_prompt,): + client = ZhipuAI(api_key=get_api_key('zhipu_api_key')) # APIKey + + + img = tensor2pil(image).convert('RGB') + + img_data = BytesIO() + img.save(img_data, format="JPEG") + img_url = base64.b64encode(img_data.getvalue()).decode("utf-8") + + messages = [ + {"role": "user", + "content": [ + {"type": "text", + "text": user_prompt + }, + {"type": "image_url", + "image_url": {"url": img_url} + } + ] + } + + ] + + response = client.chat.completions.create( + model=model, # 填写需要调用的模型名称 + messages=messages + ) + ret_message = response.choices[0].message.content + log(f"{self.NODE_NAME} response is: {ret_message}") + + return (ret_message,) + +class LS_ZhipuText: + + def __init__(self): + self.NODE_NAME = 'ZhipuGLM4' + + @classmethod + def INPUT_TYPES(cls): + glm_model_list = ["GLM-4-Flash", "GLM-4-FlashX", "GLM-4-Plus", "GLM-4-Long","GLM-4-Air", "GLM-4-AirX"] + return {"required":{ + "model": (glm_model_list,), + "user_prompt": ("STRING", {"default": "where is the capital of France?", "multiline": True}), + "history_length": ("INT", {"default": 8, "min": 1, "max": 64, "step": 1}), + }, + "optional": { + "history": ("GLM4_HISTORY",), + } + } + + RETURN_TYPES = ("STRING", "GLM4_HISTORY",) + RETURN_NAMES = ("text", "history",) + FUNCTION = "zhipu_glm4" + CATEGORY = '😺dzNodes/LayerUtility' + + def zhipu_glm4(self, model, user_prompt, history_length, history=None): + + client = ZhipuAI(api_key=get_api_key('zhipu_api_key')) # APIKey + + if history is not None: + messages = history["messages"] + messages = messages[-history_length *2:] + else: + messages = [] + task = {"role": "user", "content": user_prompt} + messages.append(task) + + response = client.chat.completions.create( + model=model, # 填写需要调用的模型名称 + messages=messages + ) + ret_message = response.choices[0].message.content + messages.append({"role": "assistant", "content": ret_message}) + log(f"{self.NODE_NAME} response is: {ret_message}") + + return (ret_message, {"messages":messages},) + + +NODE_CLASS_MAPPINGS = { + "LayerUtility: ZhipuGLM4V": LS_ZhipuImage, + "LayerUtility: ZhipuGLM4": LS_ZhipuText, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "LayerUtility: ZhipuGLM4V": "LayerUtility: ZhipuGLM4V(Advance)", + "LayerUtility: ZhipuGLM4": "LayerUtility: ZhipuGLM4(Advance)", +} + diff --git a/pyproject.toml b/pyproject.toml index 166fedf..b78f78a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui_layerstyle_advance" description = "The nodes detached from ComfyUI Layer Style are mainly those with complex requirements for dependency packages." -version = "2.0.8" +version = "2.0.9" license = "MIT" dependencies = ["numpy", "matplotlib", "scikit_image", "scikit_learn", "opencv-contrib-python", "pymatting", "timm", "blend_modes", "transformers", "diffusers", "loguru", "colour-science", "huggingface_hub", "segment_anything", "addict", "omegaconf", "yapf", "wget", "iopath", "mediapipe", "typer_config", "fastapi", "rich", "google-generativeai", "ultralytics", "transparent-background", "accelerate", "onnxruntime", "bitsandbytes", "peft", "protobuf", "hydra-core", "blind-watermark", "qrcode", "pyzbar", "psd-tools", "wandb"]