rembgNode update

"briarmbg","u2net","u2netp","u2net_human_seg","u2net_cloth_seg","silueta","isnet-general-use","isnet-anime"

README.md

@@ -2,6 +2,14 @@
 
 > [discord](https://discord.gg/cXs9vZSqeK)
 
+#### 
+[comfyui-ultralytics-yolo](https://github.com/shadowcz007/comfyui-ultralytics-yolo)
+
+[comfyui-moondream](https://github.com/shadowcz007/comfyui-moondream)
+
+[comfyui-CLIPSeg](https://github.com/shadowcz007/comfyui-CLIPSeg)
+
+
 ## 🚀🚗🚚🏃 Workflow-to-APP 
 - 新增AppInfo节点，可以通过简单的配置，把workflow转变为一个Web APP。
 - 支持多个web app 切换
@@ -149,6 +157,12 @@ Add edges to an image.
 from [simple-lama-inpainting](https://github.com/enesmsahin/simple-lama-inpainting)
 
 
+> rembgNode
+
+"briarmbg","u2net","u2netp","u2net_human_seg","u2net_cloth_seg","silueta","isnet-general-use","isnet-anime"
+
+
+
 ### Improvement 
 
 - Add "help" option to the context menu for each node.
@@ -165,8 +179,6 @@ An improvement has been made to directly redirect to GitHub to search for missin
 
 [Download rembg Models](https://github.com/danielgatis/rembg/tree/main#Models),move to:models/rembg
 
-[Download CLIPSeg](https://huggingface.co/CIDAS/clipseg-rd64-refined/tree/main), move to : models/clipseg
-
 [Download lama](https://github.com/enesmsahin/simple-lama-inpainting/releases/download/v0.1.0/big-lama.pt), move to : models/lama
 
 [Download Salesforce/blip-image-captioning-base](https://huggingface.co/Salesforce/blip-image-captioning-base), move to : models/clip_interrogator/Salesforce/blip-image-captioning-base
@@ -204,23 +216,18 @@ If you are using a venv, make sure you have it activated before installation and
 pip3 install -r requirements.txt
 ```
 
-#### 
-[comfyui-ultralytics-yolo](https://github.com/shadowcz007/comfyui-ultralytics-yolo)
-[comfyui-moondream](https://github.com/shadowcz007/comfyui-moondream)
+
+
 
 
 #### Chinese community
 访问 [www.mixcomfy.com](https://www.mixcomfy.com)，获得更多内测功能，关注微信公众号：Mixlab无界社区
 
 
-#### Thanks:
-[ComfyUI-CLIPSeg](https://github.com/biegert/ComfyUI-CLIPSeg/tree/main)
-
 #### discussions:
 [discussions](https://github.com/shadowcz007/comfyui-mixlab-nodes/discussions)
 
 
-
 <picture>
   <source
     media="(prefers-color-scheme: dark)"

__init__.py

@@ -534,7 +534,7 @@ from .nodes.PromptNode import EmbeddingPrompt,RandomPrompt,PromptSlide,PromptSim
 from .nodes.ImageNode import SaveImageToLocal,SplitImage,GridOutput,GetImageSize_,MirroredImage,ImageColorTransfer,NoiseImage,TransparentImage,GradientImage,LoadImagesFromPath,LoadImagesFromURL,ResizeImage,TextImage,SvgImage,Image3D,ShowLayer,NewLayer,MergeLayers,CenterImage,AreaToMask,SmoothMask,SplitLongMask,ImageCropByAlpha,EnhanceImage,FaceToMask
 # from .nodes.Vae import VAELoader,VAEDecode
 from .nodes.ScreenShareNode import ScreenShareNode,FloatingVideo
-from .nodes.Clipseg import CLIPSeg,CombineMasks
+
 from .nodes.ChatGPT import ChatGPTNode,ShowTextForGPT,CharacterInText
 from .nodes.Audio import GamePal,SpeechRecognition,SpeechSynthesis
 from .nodes.Utils import CreateLoraNames,CreateSampler_names,CreateCkptNames,CreateSeedNode,TESTNODE_,AppInfo,IntNumber,FloatSlider,TextInput,ColorInput,FontInput,TextToNumber,DynamicDelayProcessor,LimitNumber,SwitchByIndex,MultiplicationNode
@@ -580,8 +580,6 @@ NODE_CLASS_MAPPINGS = {
     # "VAEDecodeConsistencyDecoder":VAEDecode,
     "ScreenShare":ScreenShareNode,
     "FloatingVideo":FloatingVideo,
-    "CLIPSeg_":CLIPSeg,
-    "CombineMasks_":CombineMasks,
     "ChatGPTOpenAI":ChatGPTNode,
     "ShowTextForGPT":ShowTextForGPT,
     "CharacterInText":CharacterInText,

data/extension-node-map.json

@@ -4775,7 +4775,6 @@
             "PromptImage",
             "SaveImageToLocal",
             "AreaToMask",
-            "CLIPSeg_",
             "CharacterInText",
             "ChatGPTOpenAI",
             "Color",
@@ -4783,7 +4782,6 @@
             "CkptNames_",
             "SamplerNames_",
             "LoraNames_",
-            "CombineMasks_",
             "EnhanceImage",
             "GradientImage",
             "FaceToMask",

nodes/ChatGPT.py

@@ -14,7 +14,13 @@ def generate_random_string(length):
     letters = string.ascii_letters + string.digits
     return ''.join(random.choice(letters) for _ in range(length))
 
+class AnyType(str):
+  """A special class that is always equal in not equal comparisons. Credit to pythongosssss"""
 
+  def __ne__(self, __value: object) -> bool:
+    return False
+
+any_type = AnyType("*")
 
 #  判断是否是azure服务
 def is_azure_url(url):
@@ -196,6 +202,15 @@ class ShowTextForGPT:
     CATEGORY = "♾️Mixlab/GPT"
 
     def run(self, text,output_dir=[""]):
+        
+        # 类型纠正
+        texts=[]
+        for t in text:
+            if not isinstance(t, str):
+                t = str(t)
+            texts.append(t)
+
+        text=texts
 
         if len(output_dir)==1 and (output_dir[0]=='' or os.path.dirname(output_dir[0])==''):
             t='\n'.join(text)
@@ -272,7 +287,7 @@ class CharacterInText:
 
     def run(self, text,character,start_index):
         # print(text,character,start_index)
-        b=1 if character in text else 0
+        b=1 if character.lower() in text.lower() else 0
         
         return (b+start_index,)
 

nodes/Clipseg.py

@@ -1,272 +0,0 @@
-#### Thanks:
-# [ComfyUI-CLIPSeg](https://github.com/biegert/ComfyUI-CLIPSeg/tree/main)
-
-from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation
-
-from PIL import Image
-import torch
-import torchvision.transforms as T
-import numpy as np
-
-from torchvision.transforms.functional import to_pil_image
-import matplotlib.pyplot as plt
-import matplotlib.cm as cm
-
-
-import cv2
-
-from scipy.ndimage import gaussian_filter
-
-from typing import Optional, Tuple
-
-import warnings,os
-warnings.filterwarnings("ignore", category=UserWarning, module="torch")
-warnings.filterwarnings("ignore", category=UserWarning, module="safetensors")
-
-import folder_paths
-import logging
-
-logger = logging.getLogger('CLIPSeg nodes')
-
-clipseg_model_dir = os.path.join(folder_paths.models_dir, "clipseg")
-
-if not os.path.exists(clipseg_model_dir):
-    print(f"## clipseg model not found: {clipseg_model_dir},pls download from https://huggingface.co/CIDAS/clipseg-rd64-refined/tree/main")
-    clipseg_model_dir='CIDAS/clipseg-rd64-refined'
-
-"""Helper methods for CLIPSeg nodes"""
-
-# Tensor to PIL
-def tensor2pil(image):
-    return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
-
-# Convert PIL to Tensor
-def pil2tensor(image):
-    return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
-
-
-def tensor_to_numpy(tensor: torch.Tensor) -> np.ndarray:
-    """Convert a tensor to a numpy array and scale its values to 0-255."""
-    array = tensor.numpy().squeeze()
-    return (array * 255).astype(np.uint8)
-
-def numpy_to_tensor(array: np.ndarray) -> torch.Tensor:
-    """Convert a numpy array to a tensor and scale its values from 0-255 to 0-1."""
-    array = array.astype(np.float32) / 255.0
-    return torch.from_numpy(array)[None,]
-
-def apply_colormap(mask: torch.Tensor, colormap) -> np.ndarray:
-    """Apply a colormap to a tensor and convert it to a numpy array."""
-    colored_mask = colormap(mask.numpy())[:, :, :3]
-    return (colored_mask * 255).astype(np.uint8)
-
-def resize_image(image: np.ndarray, dimensions: Tuple[int, int]) -> np.ndarray:
-    """Resize an image to the given dimensions using linear interpolation."""
-    return cv2.resize(image, dimensions, interpolation=cv2.INTER_LINEAR)
-
-def overlay_image(background: np.ndarray, foreground: np.ndarray, alpha: float) -> np.ndarray:
-    """Overlay the foreground image onto the background with a given opacity (alpha)."""
-    return cv2.addWeighted(background, 1 - alpha, foreground, alpha, 0)
-
-def dilate_mask(mask: torch.Tensor, dilation_factor: float) -> torch.Tensor:
-    """Dilate a mask using a square kernel with a given dilation factor."""
-    kernel_size = int(dilation_factor * 2) + 1
-    kernel = np.ones((kernel_size, kernel_size), np.uint8)
-    mask_dilated = cv2.dilate(mask.numpy(), kernel, iterations=1)
-    return torch.from_numpy(mask_dilated)
-
-
-
-class CLIPSeg:
-
-    def __init__(self):
-        pass
-    
-    @classmethod
-    def INPUT_TYPES(s):
-        """
-            Return a dictionary which contains config for all input fields.
-            Some types (string): "MODEL", "VAE", "CLIP", "CONDITIONING", "LATENT", "IMAGE", "INT", "STRING", "FLOAT".
-            Input types "INT", "STRING" or "FLOAT" are special values for fields on the node.
-            The type can be a list for selection.
-
-            Returns: `dict`:
-                - Key input_fields_group (`string`): Can be either required, hidden or optional. A node class must have property `required`
-                - Value input_fields (`dict`): Contains input fields config:
-                    * Key field_name (`string`): Name of a entry-point method's argument
-                    * Value field_config (`tuple`):
-                        + First value is a string indicate the type of field or a list for selection.
-                        + Secound value is a config for type "INT", "STRING" or "FLOAT".
-        """
-        return {"required":
-                    {
-                        "image": ("IMAGE",),
-                        "text": ("STRING", {"multiline": False,"dynamicPrompts": False}),
-                        
-                     },
-                "optional":
-                    {
-                        "blur": ("FLOAT", {"min": 0, "max": 15, "step": 0.1, "default": 3}),
-                        "threshold": ("FLOAT", {"min": 0, "max": 1, "step": 0.05, "default": 0.3}),
-                        "dilation_factor": ("INT", {"min": 0, "max": 10, "step": 1, "default": 4}),
-                    }
-                }
-
-    CATEGORY = "♾️Mixlab/Mask"
-    RETURN_TYPES = ("MASK", "IMAGE", "IMAGE",)
-    RETURN_NAMES = ("Mask","Heatmap Mask", "BW Mask")
-
-    # INPUT_IS_LIST = True
-    OUTPUT_IS_LIST = (False,False,False,)
-
-    FUNCTION = "segment_image"
-    def segment_image(self, image: torch.Tensor, text: str, blur: float, threshold: float, dilation_factor: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """Create a segmentation mask from an image and a text prompt using CLIPSeg.
-
-        Args:
-            image (torch.Tensor): The image to segment.
-            text (str): The text prompt to use for segmentation.
-            blur (float): How much to blur the segmentation mask.
-            threshold (float): The threshold to use for binarizing the segmentation mask.
-            dilation_factor (int): How much to dilate the segmentation mask.
-
-        Returns:
-            Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: The segmentation mask, the heatmap mask, and the binarized mask.
-        """
-            
-        # Convert the Tensor to a PIL image
-        image_np = image.numpy().squeeze()  # Remove the first dimension (batch size of 1)
-        # Convert the numpy array back to the original range (0-255) and data type (uint8)
-        image_np = (image_np * 255).astype(np.uint8)
-        # Create a PIL image from the numpy array
-        i = Image.fromarray(image_np, mode="RGB")
-
-        processor = CLIPSegProcessor.from_pretrained(clipseg_model_dir)
-        model = CLIPSegForImageSegmentation.from_pretrained(clipseg_model_dir)
-        
-        prompt = text
-        
-        input_prc = processor(text=prompt, images=i, padding="max_length", return_tensors="pt")
-        
-        # Predict the segemntation mask
-        with torch.no_grad():
-            outputs = model(**input_prc)
-        
-        tensor = torch.sigmoid(outputs[0]) # get the mask
-        
-        # Apply a threshold to the original tensor to cut off low values
-        thresh = threshold
-        tensor_thresholded = torch.where(tensor > thresh, tensor, torch.tensor(0, dtype=torch.float))
-
-        # Apply Gaussian blur to the thresholded tensor
-        sigma = blur
-        tensor_smoothed = gaussian_filter(tensor_thresholded.numpy(), sigma=sigma)
-        tensor_smoothed = torch.from_numpy(tensor_smoothed)
-
-        # Normalize the smoothed tensor to [0, 1]
-        mask_normalized = (tensor_smoothed - tensor_smoothed.min()) / (tensor_smoothed.max() - tensor_smoothed.min())
-
-        # Dilate the normalized mask
-        mask_dilated = dilate_mask(mask_normalized, dilation_factor)
-
-        # Convert the mask to a heatmap and a binary mask
-        heatmap = apply_colormap(mask_dilated, cm.viridis)
-        binary_mask = apply_colormap(mask_dilated, cm.Greys_r)
-
-        # Overlay the heatmap and binary mask on the original image
-        dimensions = (image_np.shape[1], image_np.shape[0])
-        heatmap_resized = resize_image(heatmap, dimensions)
-        binary_mask_resized = resize_image(binary_mask, dimensions)
-
-        alpha_heatmap, alpha_binary = 0.5, 1
-        overlay_heatmap = overlay_image(image_np, heatmap_resized, alpha_heatmap)
-        overlay_binary = overlay_image(image_np, binary_mask_resized, alpha_binary)
-
-        # Convert the numpy arrays to tensors
-        image_out_heatmap = numpy_to_tensor(overlay_heatmap)
-        image_out_binary = numpy_to_tensor(overlay_binary)
-
-        # Save or display the resulting binary mask
-        binary_mask_image = Image.fromarray(binary_mask_resized[..., 0])
-
-        # convert PIL image to numpy array
-        tensor_bw = binary_mask_image.convert("L")
-        tensor_bw=pil2tensor(tensor_bw)
-        # tensor_bw = np.array(tensor_bw).astype(np.float32) / 255.0
-        # tensor_bw = torch.from_numpy(tensor_bw)[None,]
-        # tensor_bw = tensor_bw.squeeze(0)[..., 0]
-
-        return (tensor_bw, image_out_heatmap, image_out_binary,)
-
-    #OUTPUT_NODE = False
-
-class CombineMasks:
-    def __init__(self):
-        pass
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required":
-                    {
-                        "input_image": ("IMAGE", ),
-                        "mask_1": ("MASK", ), 
-                        "mask_2": ("MASK", ),
-                    },
-                "optional": 
-                    {
-                        "mask_3": ("MASK",), 
-                    },
-                }
-        
-    CATEGORY = "♾️Mixlab/Mask"
-    RETURN_TYPES = ("MASK", "IMAGE", "IMAGE",)
-    RETURN_NAMES = ("Combined Mask","Heatmap Mask", "BW Mask")
-
-    FUNCTION = "combine_masks"
-            
-    def combine_masks(self, input_image: torch.Tensor, mask_1: torch.Tensor, mask_2: torch.Tensor, mask_3: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """A method that combines two or three masks into one mask. Takes in tensors and returns the mask as a tensor, as well as the heatmap and binary mask as tensors."""
-        
-        # Combine masks
-        if mask_1 is not None:
-            mask_1 = mask_1.squeeze()
-        if mask_2 is not None:
-            mask_2 = mask_2.squeeze()
-        if mask_3 is not None:
-            mask_3 = mask_3.squeeze()
-        
-        print(mask_1.shape,mask_2.shape , mask_3.shape)
-        combined_mask = mask_1 + mask_2 + mask_3 if mask_3 is not None else mask_1 + mask_2
-        # print(combined_mask)
-
-        # Convert image and masks to numpy arrays
-        image_np = tensor_to_numpy(input_image)
-        heatmap = apply_colormap(combined_mask, cm.viridis)
-        binary_mask = apply_colormap(combined_mask, cm.Greys_r)
-        
-        # Resize heatmap and binary mask to match the original image dimensions
-        dimensions = (image_np.shape[1], image_np.shape[0])
-        # print('heatmap',heatmap)
-        if dimensions is None or dimensions[0] == 0 or dimensions[1] == 0:
-            raise ValueError("Invalid dimensions")
-        
-        heatmap_resized = resize_image(heatmap, dimensions)
-        binary_mask_resized = resize_image(binary_mask, dimensions)
-
-        # Overlay the heatmap and binary mask onto the original image
-        alpha_heatmap, alpha_binary = 0.5, 1
-        overlay_heatmap = overlay_image(image_np, heatmap_resized, alpha_heatmap)
-        overlay_binary = overlay_image(image_np, binary_mask_resized, alpha_binary)
-
-        # Convert overlays to tensors
-        image_out_heatmap = numpy_to_tensor(overlay_heatmap)
-        image_out_binary = numpy_to_tensor(overlay_binary)
-
-        return combined_mask, image_out_heatmap, image_out_binary
-
-# A dictionary that contains all nodes you want to export with their names
-# NOTE: names should be globally unique
-# NODE_CLASS_MAPPINGS = {
-#     "CLIPSeg": CLIPSeg,
-#     "CombineSegMasks": CombineMasks,
-# }

nodes/RembgNode.py

@@ -8,6 +8,467 @@ import comfy.utils
 import numpy as np
 import torch
 
+from huggingface_hub import hf_hub_download
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision.transforms.functional import normalize
+# BRIA-RMBG-1.4 / briarmbg.py
+class REBNCONV(nn.Module):
+    def __init__(self,in_ch=3,out_ch=3,dirate=1,stride=1):
+        super(REBNCONV,self).__init__()
+
+        self.conv_s1 = nn.Conv2d(in_ch,out_ch,3,padding=1*dirate,dilation=1*dirate,stride=stride)
+        self.bn_s1 = nn.BatchNorm2d(out_ch)
+        self.relu_s1 = nn.ReLU(inplace=True)
+
+    def forward(self,x):
+
+        hx = x
+        xout = self.relu_s1(self.bn_s1(self.conv_s1(hx)))
+
+        return xout
+
+## upsample tensor 'src' to have the same spatial size with tensor 'tar'
+def _upsample_like(src,tar):
+
+    src = F.interpolate(src,size=tar.shape[2:],mode='bilinear')
+
+    return src
+
+
+### RSU-7 ###
+class RSU7(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3, img_size=512):
+        super(RSU7,self).__init__()
+
+        self.in_ch = in_ch
+        self.mid_ch = mid_ch
+        self.out_ch = out_ch
+
+        self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1) ## 1 -> 1/2
+
+        self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
+        self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool3 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool4 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool5 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv6 = REBNCONV(mid_ch,mid_ch,dirate=1)
+
+        self.rebnconv7 = REBNCONV(mid_ch,mid_ch,dirate=2)
+
+        self.rebnconv6d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv5d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
+
+    def forward(self,x):
+        b, c, h, w = x.shape
+
+        hx = x
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+        hx = self.pool3(hx3)
+
+        hx4 = self.rebnconv4(hx)
+        hx = self.pool4(hx4)
+
+        hx5 = self.rebnconv5(hx)
+        hx = self.pool5(hx5)
+
+        hx6 = self.rebnconv6(hx)
+
+        hx7 = self.rebnconv7(hx6)
+
+        hx6d =  self.rebnconv6d(torch.cat((hx7,hx6),1))
+        hx6dup = _upsample_like(hx6d,hx5)
+
+        hx5d =  self.rebnconv5d(torch.cat((hx6dup,hx5),1))
+        hx5dup = _upsample_like(hx5d,hx4)
+
+        hx4d = self.rebnconv4d(torch.cat((hx5dup,hx4),1))
+        hx4dup = _upsample_like(hx4d,hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4dup,hx3),1))
+        hx3dup = _upsample_like(hx3d,hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
+        hx2dup = _upsample_like(hx2d,hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
+
+        return hx1d + hxin
+
+
+### RSU-6 ###
+class RSU6(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU6,self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
+        self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool3 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool4 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=1)
+
+        self.rebnconv6 = REBNCONV(mid_ch,mid_ch,dirate=2)
+
+        self.rebnconv5d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
+
+    def forward(self,x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+        hx = self.pool3(hx3)
+
+        hx4 = self.rebnconv4(hx)
+        hx = self.pool4(hx4)
+
+        hx5 = self.rebnconv5(hx)
+
+        hx6 = self.rebnconv6(hx5)
+
+
+        hx5d =  self.rebnconv5d(torch.cat((hx6,hx5),1))
+        hx5dup = _upsample_like(hx5d,hx4)
+
+        hx4d = self.rebnconv4d(torch.cat((hx5dup,hx4),1))
+        hx4dup = _upsample_like(hx4d,hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4dup,hx3),1))
+        hx3dup = _upsample_like(hx3d,hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
+        hx2dup = _upsample_like(hx2d,hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
+
+        return hx1d + hxin
+
+### RSU-5 ###
+class RSU5(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU5,self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
+        self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool3 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
+
+        self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=2)
+
+        self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
+
+    def forward(self,x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+        hx = self.pool3(hx3)
+
+        hx4 = self.rebnconv4(hx)
+
+        hx5 = self.rebnconv5(hx4)
+
+        hx4d = self.rebnconv4d(torch.cat((hx5,hx4),1))
+        hx4dup = _upsample_like(hx4d,hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4dup,hx3),1))
+        hx3dup = _upsample_like(hx3d,hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
+        hx2dup = _upsample_like(hx2d,hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
+
+        return hx1d + hxin
+
+### RSU-4 ###
+class RSU4(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU4,self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
+        self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
+
+        self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=2)
+
+        self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
+
+    def forward(self,x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+
+        hx4 = self.rebnconv4(hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4,hx3),1))
+        hx3dup = _upsample_like(hx3d,hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
+        hx2dup = _upsample_like(hx2d,hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
+
+        return hx1d + hxin
+
+### RSU-4F ###
+class RSU4F(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU4F,self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
+        self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=2)
+        self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=4)
+
+        self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=8)
+
+        self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=4)
+        self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=2)
+        self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
+
+    def forward(self,x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx2 = self.rebnconv2(hx1)
+        hx3 = self.rebnconv3(hx2)
+
+        hx4 = self.rebnconv4(hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4,hx3),1))
+        hx2d = self.rebnconv2d(torch.cat((hx3d,hx2),1))
+        hx1d = self.rebnconv1d(torch.cat((hx2d,hx1),1))
+
+        return hx1d + hxin
+
+
+class myrebnconv(nn.Module):
+    def __init__(self, in_ch=3,
+                       out_ch=1,
+                       kernel_size=3,
+                       stride=1,
+                       padding=1,
+                       dilation=1,
+                       groups=1):
+        super(myrebnconv,self).__init__()
+
+        self.conv = nn.Conv2d(in_ch,
+                              out_ch,
+                              kernel_size=kernel_size,
+                              stride=stride,
+                              padding=padding,
+                              dilation=dilation,
+                              groups=groups)
+        self.bn = nn.BatchNorm2d(out_ch)
+        self.rl = nn.ReLU(inplace=True)
+
+    def forward(self,x):
+        return self.rl(self.bn(self.conv(x)))
+
+
+class BriaRMBG(nn.Module):
+
+    def __init__(self,in_ch=3,out_ch=1):
+        super(BriaRMBG,self).__init__()
+
+        self.conv_in = nn.Conv2d(in_ch,64,3,stride=2,padding=1)
+        self.pool_in = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage1 = RSU7(64,32,64)
+        self.pool12 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage2 = RSU6(64,32,128)
+        self.pool23 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage3 = RSU5(128,64,256)
+        self.pool34 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage4 = RSU4(256,128,512)
+        self.pool45 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage5 = RSU4F(512,256,512)
+        self.pool56 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage6 = RSU4F(512,256,512)
+
+        # decoder
+        self.stage5d = RSU4F(1024,256,512)
+        self.stage4d = RSU4(1024,128,256)
+        self.stage3d = RSU5(512,64,128)
+        self.stage2d = RSU6(256,32,64)
+        self.stage1d = RSU7(128,16,64)
+
+        self.side1 = nn.Conv2d(64,out_ch,3,padding=1)
+        self.side2 = nn.Conv2d(64,out_ch,3,padding=1)
+        self.side3 = nn.Conv2d(128,out_ch,3,padding=1)
+        self.side4 = nn.Conv2d(256,out_ch,3,padding=1)
+        self.side5 = nn.Conv2d(512,out_ch,3,padding=1)
+        self.side6 = nn.Conv2d(512,out_ch,3,padding=1)
+
+        # self.outconv = nn.Conv2d(6*out_ch,out_ch,1)
+
+    def forward(self,x):
+
+        hx = x
+
+        hxin = self.conv_in(hx)
+        #hx = self.pool_in(hxin)
+
+        #stage 1
+        hx1 = self.stage1(hxin)
+        hx = self.pool12(hx1)
+
+        #stage 2
+        hx2 = self.stage2(hx)
+        hx = self.pool23(hx2)
+
+        #stage 3
+        hx3 = self.stage3(hx)
+        hx = self.pool34(hx3)
+
+        #stage 4
+        hx4 = self.stage4(hx)
+        hx = self.pool45(hx4)
+
+        #stage 5
+        hx5 = self.stage5(hx)
+        hx = self.pool56(hx5)
+
+        #stage 6
+        hx6 = self.stage6(hx)
+        hx6up = _upsample_like(hx6,hx5)
+
+        #-------------------- decoder --------------------
+        hx5d = self.stage5d(torch.cat((hx6up,hx5),1))
+        hx5dup = _upsample_like(hx5d,hx4)
+
+        hx4d = self.stage4d(torch.cat((hx5dup,hx4),1))
+        hx4dup = _upsample_like(hx4d,hx3)
+
+        hx3d = self.stage3d(torch.cat((hx4dup,hx3),1))
+        hx3dup = _upsample_like(hx3d,hx2)
+
+        hx2d = self.stage2d(torch.cat((hx3dup,hx2),1))
+        hx2dup = _upsample_like(hx2d,hx1)
+
+        hx1d = self.stage1d(torch.cat((hx2dup,hx1),1))
+
+
+        #side output
+        d1 = self.side1(hx1d)
+        d1 = _upsample_like(d1,x)
+
+        d2 = self.side2(hx2d)
+        d2 = _upsample_like(d2,x)
+
+        d3 = self.side3(hx3d)
+        d3 = _upsample_like(d3,x)
+
+        d4 = self.side4(hx4d)
+        d4 = _upsample_like(d4,x)
+
+        d5 = self.side5(hx5d)
+        d5 = _upsample_like(d5,x)
+
+        d6 = self.side6(hx6)
+        d6 = _upsample_like(d6,x)
+
+        return [F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)],[hx1d,hx2d,hx3d,hx4d,hx5d,hx6]
+    
+
+
+
+
+
 
 U2NET_HOME=os.path.join(folder_paths.models_dir, "rembg")
 os.environ["U2NET_HOME"] = U2NET_HOME
@@ -48,6 +509,70 @@ except:
     _available=False
 
 
+def briarmbg_run(images=[]):
+    mroot=os.path.join(folder_paths.models_dir, "rembg")
+    m=os.path.join(mroot,'briarmbg.pth')
+    if os.path.exists(m)==False:
+        # 下载
+        m1=hf_hub_download("briaai/RMBG-1.4",
+                       local_dir=mroot,
+                       filename='model.pth',
+                       local_dir_use_symlinks=False,
+                       endpoint='https://hf-mirror.com')
+        os.rename(m1, m)
+    
+    net=BriaRMBG()
+    if torch.cuda.is_available():
+        net.load_state_dict(torch.load(m))
+        net=net.cuda()
+    else:
+        net.load_state_dict(torch.load(m,map_location="cpu"))
+    net.eval() 
+
+    masks=[]
+    rgba_images=[]
+    rgb_images=[]
+    for orig_image in images:
+
+        w,h = orig_im_size = orig_image.size
+
+        image = orig_image.convert('RGB')
+        model_input_size = (1024, 1024)
+        image = image.resize(model_input_size, Image.BILINEAR)
+
+        im_np = np.array(image)
+        im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2,0,1)
+        im_tensor = torch.unsqueeze(im_tensor,0)
+        im_tensor = torch.divide(im_tensor,255.0)
+        im_tensor = normalize(im_tensor,[0.5,0.5,0.5],[1.0,1.0,1.0])
+        if torch.cuda.is_available():
+            im_tensor=im_tensor.cuda()
+
+        result=net(im_tensor)
+        result = torch.squeeze(F.interpolate(result[0][0], size=(h,w), mode='bilinear') ,0)
+        ma = torch.max(result)
+        mi = torch.min(result)
+        result = (result-mi)/(ma-mi)
+        im_array = (result*255).cpu().data.numpy().astype(np.uint8)
+        mask = Image.fromarray(np.squeeze(im_array))
+        # mask.save('test.png')
+        # mask=tensor2pil(result)
+        mask=mask.convert('L') 
+
+        masks.append(mask)
+
+        # rgba图
+        image_rgba =orig_image.convert("RGBA")
+        image_rgba.putalpha(mask)
+        rgba_images.append(image_rgba)
+
+        #rgb
+        rgb_image = Image.new("RGB", image_rgba.size, (0, 0, 0))
+        rgb_image.paste(image_rgba, mask=image_rgba.split()[3])
+        rgb_images.append(rgb_image)
+    return (masks,rgba_images,rgb_images)   
+
+
 def run_bg(model_name= "unet",images=[]):
     # model_name = "unet" # "isnet-general-use"
     rembg_session = new_session(model_name)
@@ -118,14 +643,16 @@ class RembgNode_:
     def INPUT_TYPES(s):
         return {"required": {
             "image": ("IMAGE",),
-            "model_name": (["u2net",
+            "model_name": ([
+                            "briarmbg",
+                            "u2net",
                             "u2netp",
                             "u2net_human_seg",
                             "u2net_cloth_seg",
                             "silueta",
                             "isnet-general-use",
                             "isnet-anime",
-                            # "sam"
+                            
                             ],),
             
                              },
@@ -153,7 +680,10 @@ class RembgNode_:
                 im=tensor2pil(im)
                 images.append(im)
 
-        masks,rgba_images,rgb_images=run_bg(model_name,images)
+        if model_name=='briarmbg':
+            masks,rgba_images,rgb_images=briarmbg_run(images)
+        else:
+            masks,rgba_images,rgb_images=run_bg(model_name,images)
 
         masks=[pil2tensor(m) for m in masks]
 

web/javascript/checkVersion_mixlab.js

@@ -3,7 +3,7 @@ import { app } from '../../../scripts/app.js'
 const repoOwner = 'shadowcz007' // 替换为仓库的所有者
 const repoName = 'comfyui-mixlab-nodes' // 替换为仓库的名称
 
-const version = 'v0.15.1'
+const version = 'v0.16.0'
 
 fetch(`https://api.github.com/repos/${repoOwner}/${repoName}/releases/latest`)
   .then(response => response.json())