Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
1258746886 | ||
|
|
ed128b0ad6 | ||
|
|
f0db08acd6 | ||
|
|
7c655e3080 | ||
|
|
1b9871c3df | ||
|
|
7568aaf243 | ||
|
|
a76be8450d | ||
|
|
5564ee1246 | ||
|
|
a6e9251521 | ||
|
|
0bee093916 | ||
|
|
13a9878823 | ||
|
|
acd35d50f8 | ||
|
|
5c0d99e72d |
@@ -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 切换
|
||||
@@ -31,12 +39,14 @@ APP-JSON:
|
||||
- [image-to-image](./example/Image-to-Image_2.json)
|
||||
- text-to-text
|
||||
|
||||
> 暂时支持8种节点作为界面上的输入节点:Load Image、CLIPTextEncode、PromptSlide、TextInput_、Color、FloatSlider、IntNumber、CheckpointLoaderSimple、LoraLoader
|
||||
> 暂时支持 9 种节点作为界面上的输入节点:Load Image、VHS_LoadVideo、CLIPTextEncode、PromptSlide、TextInput_、Color、FloatSlider、IntNumber、CheckpointLoaderSimple、LoraLoader
|
||||
|
||||
> 输出节点:PreviewImage 、SaveImage、ShowTextForGPT、VHS_VideoCombine、PromptImage
|
||||
|
||||
> seed统一输入控件,支持:SamplerCustom、KSampler
|
||||
|
||||
> [ps插件](https://github.com/shadowcz007/comfyui-ps-plugin)
|
||||
|
||||
|
||||
## 🏃🚗🚚🚀 Real-time Design
|
||||
> ScreenShareNode & FloatingVideoNode. Now comfyui supports capturing screen pixel streams from any software and can be used for LCM-Lora integration. Let's get started with implementation and design! 💻🌐
|
||||
@@ -149,6 +159,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 +181,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 +218,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)"
|
||||
|
||||
+3
-4
@@ -534,10 +534,10 @@ 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
|
||||
from .nodes.Utils import CreateLoraNames,CreateSampler_names,CreateCkptNames,CreateSeedNode,TESTNODE_,TESTNODE_TOKEN,AppInfo,IntNumber,FloatSlider,TextInput,ColorInput,FontInput,TextToNumber,DynamicDelayProcessor,LimitNumber,SwitchByIndex,MultiplicationNode
|
||||
from .nodes.Mask import OutlineMask,FeatheredMask
|
||||
|
||||
|
||||
@@ -546,6 +546,7 @@ from .nodes.Mask import OutlineMask,FeatheredMask
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"AppInfo":AppInfo,
|
||||
"TESTNODE_":TESTNODE_,
|
||||
"TESTNODE_TOKEN":TESTNODE_TOKEN,
|
||||
"RandomPrompt":RandomPrompt,
|
||||
"EmbeddingPrompt":EmbeddingPrompt,
|
||||
"PromptSlide":PromptSlide,
|
||||
@@ -580,8 +581,6 @@ NODE_CLASS_MAPPINGS = {
|
||||
# "VAEDecodeConsistencyDecoder":VAEDecode,
|
||||
"ScreenShare":ScreenShareNode,
|
||||
"FloatingVideo":FloatingVideo,
|
||||
"CLIPSeg_":CLIPSeg,
|
||||
"CombineMasks_":CombineMasks,
|
||||
"ChatGPTOpenAI":ChatGPTNode,
|
||||
"ShowTextForGPT":ShowTextForGPT,
|
||||
"CharacterInText":CharacterInText,
|
||||
|
||||
@@ -4775,7 +4775,6 @@
|
||||
"PromptImage",
|
||||
"SaveImageToLocal",
|
||||
"AreaToMask",
|
||||
"CLIPSeg_",
|
||||
"CharacterInText",
|
||||
"ChatGPTOpenAI",
|
||||
"Color",
|
||||
@@ -4783,7 +4782,6 @@
|
||||
"CkptNames_",
|
||||
"SamplerNames_",
|
||||
"LoraNames_",
|
||||
"CombineMasks_",
|
||||
"EnhanceImage",
|
||||
"GradientImage",
|
||||
"FaceToMask",
|
||||
|
||||
+16
-1
@@ -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,)
|
||||
|
||||
|
||||
@@ -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,
|
||||
# }
|
||||
+533
-3
@@ -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]
|
||||
|
||||
|
||||
+37
-3
@@ -381,7 +381,7 @@ class TextInput:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {
|
||||
"text": ("STRING",{"multiline": True,"default": ""}),
|
||||
"text": ("STRING",{"multiline": True,"default": ""})
|
||||
},
|
||||
}
|
||||
|
||||
@@ -695,7 +695,9 @@ class ListStatistics:
|
||||
class TESTNODE_:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": { "ANY":(any_type,), },
|
||||
return {"required": {
|
||||
"ANY":(any_type,),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = (any_type,)
|
||||
@@ -715,10 +717,42 @@ class TESTNODE_:
|
||||
|
||||
# 调用count_types方法进行统计
|
||||
result = list_stats.count_types(ANY)
|
||||
|
||||
|
||||
|
||||
return {"ui": {"data": result,"type":[str(type(ANY[0]))]}, "result": (ANY,)}
|
||||
|
||||
|
||||
class TESTNODE_TOKEN:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {
|
||||
"text":("STRING", {"forceInput": True,}),
|
||||
"clip": ("CLIP", )
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
|
||||
FUNCTION = "run"
|
||||
|
||||
CATEGORY = "♾️Mixlab/__TEST"
|
||||
|
||||
OUTPUT_NODE = True
|
||||
INPUT_IS_LIST = False
|
||||
OUTPUT_IS_LIST = (False,)
|
||||
|
||||
def run(self,text,clip=None):
|
||||
# print(text)
|
||||
|
||||
tokens = clip.tokenize(text)
|
||||
|
||||
tokens=[v for v in tokens.values()][0][0]
|
||||
|
||||
tokens=json.dumps(tokens)
|
||||
|
||||
return (tokens,)
|
||||
|
||||
|
||||
|
||||
class CreateSeedNode:
|
||||
def __init__(self):
|
||||
|
||||
+95
-13
@@ -435,8 +435,9 @@
|
||||
body
|
||||
})
|
||||
|
||||
// console.log(resp)
|
||||
|
||||
let data = await resp.json()
|
||||
// console.log(data)
|
||||
let { name, subfolder } = data
|
||||
let src = `${url}/view?filename=${encodeURIComponent(
|
||||
name
|
||||
@@ -888,6 +889,52 @@
|
||||
}
|
||||
|
||||
|
||||
function generateRainbowVideo() {
|
||||
// 创建一个canvas元素
|
||||
const canvas = document.createElement('canvas');
|
||||
canvas.width = 640; // 设置canvas宽度
|
||||
canvas.height = 480; // 设置canvas高度
|
||||
const context = canvas.getContext('2d');
|
||||
|
||||
// 绘制第一帧彩虹
|
||||
context.fillStyle = 'red';
|
||||
context.fillRect(0, 0, canvas.width / 2, canvas.height);
|
||||
context.fillStyle = 'orange';
|
||||
context.fillRect(canvas.width / 2, 0, canvas.width / 2, canvas.height);
|
||||
|
||||
// 绘制第二帧彩虹
|
||||
context.fillStyle = 'yellow';
|
||||
context.fillRect(0, 0, canvas.width / 2, canvas.height);
|
||||
context.fillStyle = 'green';
|
||||
context.fillRect(canvas.width / 2, 0, canvas.width / 2, canvas.height);
|
||||
|
||||
// 创建一个新的video元素
|
||||
const video = document.createElement('video');
|
||||
video.setAttribute('controls', true)
|
||||
// 将canvas绘制到video中
|
||||
const stream = canvas.captureStream();
|
||||
video.srcObject = stream;
|
||||
|
||||
// 导出视频
|
||||
const mediaRecorder = new MediaRecorder(stream);
|
||||
const chunks = [];
|
||||
mediaRecorder.ondataavailable = function (event) {
|
||||
chunks.push(event.data);
|
||||
};
|
||||
mediaRecorder.onstop = function () {
|
||||
const blob = new Blob(chunks, { type: 'video/mp4' });
|
||||
const url = URL.createObjectURL(blob);
|
||||
video.src = url;
|
||||
};
|
||||
mediaRecorder.start();
|
||||
setTimeout(function () {
|
||||
mediaRecorder.stop();
|
||||
}, 1000); // 设置录制时长,这里设置为1秒
|
||||
|
||||
return video;
|
||||
}
|
||||
|
||||
|
||||
async function calculateImageHash(blob) {
|
||||
const buffer = await blob.arrayBuffer();
|
||||
const hashBuffer = await crypto.subtle.digest('SHA-256', buffer);
|
||||
@@ -1081,16 +1128,19 @@
|
||||
inputData = inputData.filter(inp => inp);
|
||||
// console.log('inputData',inputData)
|
||||
inputData.forEach(data => {
|
||||
// console.log(data)
|
||||
console.log(data)
|
||||
// Check if the class_type is "LoadImage"
|
||||
if (data.class_type === "LoadImage") {
|
||||
if (data.class_type === "LoadImage" || data.class_type === "VHS_LoadVideo") {
|
||||
|
||||
let isVideoUpload = data.class_type === "VHS_LoadVideo";
|
||||
|
||||
// Create a container for the upload control
|
||||
const uploadContainer = document.createElement("div");
|
||||
uploadContainer.className = 'card';
|
||||
|
||||
// Create a label for the upload control
|
||||
const nameLabel = document.createElement("label");
|
||||
nameLabel.textContent = data.title || "LoadImage: ";
|
||||
nameLabel.textContent = data.title || (isVideoUpload ? "LoadVideo: " : "LoadImage: ");
|
||||
nameLabel.style.marginBottom = '12px'
|
||||
uploadContainer.appendChild(nameLabel);
|
||||
|
||||
@@ -1109,25 +1159,48 @@
|
||||
const btnFromClipboard = document.createElement("button");
|
||||
btnFromClipboard.style = `width: 156px; margin-left: 18px;`
|
||||
btnFromClipboard.innerText = 'paste from clipboard'
|
||||
actionDiv.appendChild(btnFromClipboard);
|
||||
if (!isVideoUpload) actionDiv.appendChild(btnFromClipboard);
|
||||
|
||||
const btnForImageEdit = document.createElement("button");
|
||||
btnForImageEdit.style = ` width: 32px; background: none;margin-left: 18px;`
|
||||
btnForImageEdit.innerHTML = '<?xml version="1.0" ?><svg version="1.1" style="width: 24px;" viewBox="0 0 50 50" xml:space="preserve" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink"><g id="Layer_1_1_"><path d="M18.293,31.707h6.414l24-24l-6.414-6.414l-24,24V31.707z M45.879,7.707l-3.586,3.586l-3.586-3.586l3.586-3.586 L45.879,7.707z M20.293,26.121l17-17l3.586,3.586l-17,17h-3.586V26.121z"/><polygon points="43.293,19.707 41.293,19.707 41.293,46.707 3.293,46.707 3.293,8.707 31.293,8.707 31.293,6.707 1.293,6.707 1.293,48.707 43.293,48.707 "/></g></svg>'
|
||||
actionDiv.appendChild(btnForImageEdit);
|
||||
if (!isVideoUpload) actionDiv.appendChild(btnForImageEdit);
|
||||
|
||||
|
||||
uploadContainer.appendChild(actionDiv)
|
||||
|
||||
// Create an image element to display the uploaded image
|
||||
const imageElement = document.createElement("img");
|
||||
imageElement.src = base64Df
|
||||
let imageElement = document.createElement("img");
|
||||
if (isVideoUpload) {
|
||||
imageElement = document.createElement('video');
|
||||
imageElement.setAttribute('controls', true)
|
||||
let [subfolder, name] = data.inputs.video.split('/');
|
||||
// console.log(subfolder,name)
|
||||
if (!name) {
|
||||
subfolder = "";
|
||||
name = data.inputs.video;
|
||||
}
|
||||
let url = `${get_url()}/view?filename=${encodeURIComponent(name)}&type=input&subfolder=${subfolder}&rand=${Math.random()}`
|
||||
imageElement.src = url;
|
||||
} else {
|
||||
// 图片
|
||||
let [subfolder, name] = data.inputs.image.split('/');
|
||||
if (!name) {
|
||||
subfolder = "";
|
||||
name = data.inputs.image;
|
||||
}
|
||||
// imageElement.src = base64Df
|
||||
let url = `${get_url()}/view?filename=${encodeURIComponent(name)}&type=input&subfolder=${subfolder}&rand=${Math.random()}`
|
||||
imageElement.src = url;
|
||||
}
|
||||
|
||||
imageElement.style.maxWidth = '200px';
|
||||
|
||||
|
||||
btnFromClipboard.addEventListener('click', (event) => handleClipboardImage(imageElement, data));
|
||||
|
||||
btnForImageEdit.addEventListener('click', e => editImage(imageElement, data))
|
||||
if (!isVideoUpload) btnFromClipboard.addEventListener('click', (event) => handleClipboardImage(imageElement, data));
|
||||
|
||||
if (!isVideoUpload) btnForImageEdit.addEventListener('click', e => editImage(imageElement, data))
|
||||
|
||||
|
||||
uploadImageInput.addEventListener('click', (event) => {
|
||||
@@ -1145,15 +1218,24 @@
|
||||
reader.onloadend = async function () {
|
||||
// 获取读取的文件内容,即 Blob 对象
|
||||
const fileBlob = new Blob([reader.result], { type: file.type });
|
||||
|
||||
// console.log( file.type.split('/')[1])
|
||||
let hashId = await calculateImageHash(fileBlob)
|
||||
|
||||
if (hashId == window._appData.data[data.id].hashId) return
|
||||
|
||||
let { url, name } = await uploadImage(fileBlob)
|
||||
let { url, name } = await uploadImage(fileBlob, '.' + file.type.split('/')[1])
|
||||
if (isVideoUpload) {
|
||||
imageElement.srcObject = null;
|
||||
}
|
||||
// 在这里可以对 Blob 对象进行进一步处理
|
||||
imageElement.src = url;
|
||||
window._appData.data[data.id].inputs.image = name;
|
||||
|
||||
if (isVideoUpload) {
|
||||
window._appData.data[data.id].inputs.video = name;
|
||||
} else {
|
||||
window._appData.data[data.id].inputs.image = name;
|
||||
}
|
||||
|
||||
window._appData.data[data.id].hashId = hashId;
|
||||
|
||||
console.log("上传的文件:", url, data.id, name);
|
||||
|
||||
@@ -277,7 +277,7 @@ async function save (json, download = false, showInfo = true) {
|
||||
|
||||
function getInputsAndOutputs () {
|
||||
const inputs =
|
||||
`LoadImage CLIPTextEncode PromptSlide TextInput_ Color FloatSlider IntNumber CheckpointLoaderSimple LoraLoader`.split(
|
||||
`LoadImage VHS_LoadVideo CLIPTextEncode PromptSlide TextInput_ Color FloatSlider IntNumber CheckpointLoaderSimple LoraLoader`.split(
|
||||
' '
|
||||
),
|
||||
outputs = `PreviewImage SaveImage ShowTextForGPT VHS_VideoCombine`.split(
|
||||
|
||||
@@ -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.17.0'
|
||||
|
||||
fetch(`https://api.github.com/repos/${repoOwner}/${repoName}/releases/latest`)
|
||||
.then(response => response.json())
|
||||
|
||||
+147
-40
@@ -9,33 +9,31 @@ import {
|
||||
|
||||
import { smart_init, addSmartMenu } from './smart_connect.js'
|
||||
|
||||
let isScriptLoaded = {};
|
||||
let isScriptLoaded = {}
|
||||
|
||||
function loadExternalScript(url) {
|
||||
function loadExternalScript (url) {
|
||||
return new Promise((resolve, reject) => {
|
||||
if (isScriptLoaded[url]) {
|
||||
resolve();
|
||||
return;
|
||||
resolve()
|
||||
return
|
||||
}
|
||||
|
||||
const script = document.createElement('script');
|
||||
script.src = url;
|
||||
const script = document.createElement('script')
|
||||
script.src = url
|
||||
script.onload = () => {
|
||||
isScriptLoaded[url]= true;
|
||||
resolve();
|
||||
};
|
||||
script.onerror = reject;
|
||||
document.head.appendChild(script);
|
||||
});
|
||||
isScriptLoaded[url] = true
|
||||
resolve()
|
||||
}
|
||||
script.onerror = reject
|
||||
document.head.appendChild(script)
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
//
|
||||
|
||||
function createChart(chartDom,nodes){
|
||||
|
||||
var myChart = echarts.init(chartDom);
|
||||
var option;
|
||||
function createChart (chartDom, nodes) {
|
||||
var myChart = echarts.init(chartDom)
|
||||
var option
|
||||
|
||||
console.log(nodes)
|
||||
option = {
|
||||
@@ -46,24 +44,25 @@ function createChart(chartDom,nodes){
|
||||
{
|
||||
name: 'nodeA',
|
||||
value: 10,
|
||||
children: Array.from(nodes,n=>{
|
||||
children: Array.from(nodes, n => {
|
||||
return {
|
||||
name:n.type,
|
||||
value:n.count
|
||||
name: n.type,
|
||||
value: n.count
|
||||
}
|
||||
})
|
||||
},
|
||||
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
};
|
||||
|
||||
option && myChart.setOption(option);
|
||||
}
|
||||
|
||||
option && myChart.setOption(option)
|
||||
}
|
||||
|
||||
async function createNodesCharts () {
|
||||
await loadExternalScript('/extensions/comfyui-mixlab-nodes/lib/echarts.min.js')
|
||||
await loadExternalScript(
|
||||
'/extensions/comfyui-mixlab-nodes/lib/echarts.min.js'
|
||||
)
|
||||
const templates = await loadTemplate()
|
||||
var nodes = {}
|
||||
Array.from(templates, t => {
|
||||
@@ -74,7 +73,6 @@ async function createNodesCharts () {
|
||||
}
|
||||
})
|
||||
nodes = Object.values(nodes).sort((a, b) => b.count - a.count)
|
||||
|
||||
|
||||
const menu = document.querySelector('.comfy-menu')
|
||||
const separator = document.createElement('div')
|
||||
@@ -86,11 +84,9 @@ async function createNodesCharts () {
|
||||
menu.append(separator)
|
||||
|
||||
const appsButton = document.createElement('button')
|
||||
appsButton.textContent = 'Nodes';
|
||||
appsButton.textContent = 'Nodes'
|
||||
|
||||
|
||||
|
||||
appsButton.onclick = () => {
|
||||
appsButton.onclick = () => {
|
||||
let div = document.querySelector('#mixlab_apps')
|
||||
if (!div) {
|
||||
div = document.createElement('div')
|
||||
@@ -150,10 +146,9 @@ async function createNodesCharts () {
|
||||
div.appendChild(btn)
|
||||
|
||||
let chartDom = document.createElement('div')
|
||||
chartDom.style=`height:80vh;width:450px`
|
||||
chartDom.className='chart'
|
||||
chartDom.style = `height:80vh;width:450px`
|
||||
chartDom.className = 'chart'
|
||||
div.appendChild(chartDom)
|
||||
|
||||
}
|
||||
if (div.style.display == 'flex') {
|
||||
div.style.display = 'none'
|
||||
@@ -173,14 +168,11 @@ async function createNodesCharts () {
|
||||
background-color: var(--comfy-menu-bg);
|
||||
padding: 10px;
|
||||
border: 1px solid black;z-index: 999999999;padding-top: 0;`
|
||||
};
|
||||
|
||||
|
||||
createChart(div.querySelector('.chart'),nodes)
|
||||
}
|
||||
|
||||
createChart(div.querySelector('.chart'), nodes)
|
||||
}
|
||||
menu.append(appsButton)
|
||||
|
||||
}
|
||||
|
||||
function copyNodeValues (src, dest) {
|
||||
@@ -366,7 +358,14 @@ injectCSS(`::-webkit-scrollbar {
|
||||
animation-name: loading_mixlab;
|
||||
animation-duration: 2s;
|
||||
animation-iteration-count: infinite;
|
||||
}`)
|
||||
}
|
||||
|
||||
.dynamic_prompt{
|
||||
border-left: 2px solid var(--input-text);
|
||||
}
|
||||
|
||||
|
||||
`)
|
||||
|
||||
async function getCustomnodeMappings (mode = 'url') {
|
||||
// mode = "local";
|
||||
@@ -824,6 +823,73 @@ const loadTemplate = async () => {
|
||||
return templates ?? []
|
||||
}
|
||||
|
||||
function drawBadge (node, orig, restArgs) {
|
||||
let ctx = restArgs[0]
|
||||
const r = orig?.apply?.(node, restArgs)
|
||||
|
||||
if (
|
||||
!node.flags.collapsed &&
|
||||
node.constructor.title_mode != LiteGraph.NO_TITLE
|
||||
) {
|
||||
let text = `#${node.id} `
|
||||
|
||||
let nick = node.getNickname()
|
||||
if (nick) {
|
||||
if (nick == 'ComfyUI') {
|
||||
nick = '🦊'
|
||||
}
|
||||
|
||||
if (nick.length > 25) {
|
||||
text += nick.substring(0, 23) + '..'
|
||||
} else {
|
||||
text += nick
|
||||
}
|
||||
}
|
||||
|
||||
if (text != '') {
|
||||
let fgColor = 'white'
|
||||
let bgColor = '#0F1F0F'
|
||||
let visible = true
|
||||
|
||||
ctx.save()
|
||||
ctx.font = '12px sans-serif'
|
||||
const sz = ctx.measureText(text)
|
||||
ctx.fillStyle = bgColor
|
||||
ctx.beginPath()
|
||||
ctx.roundRect(
|
||||
node.size[0] - sz.width - 12,
|
||||
-LiteGraph.NODE_TITLE_HEIGHT - 20,
|
||||
sz.width + 12,
|
||||
20,
|
||||
5
|
||||
)
|
||||
ctx.fill()
|
||||
|
||||
ctx.fillStyle = fgColor
|
||||
ctx.fillText(
|
||||
text,
|
||||
node.size[0] - sz.width - 6,
|
||||
-LiteGraph.NODE_TITLE_HEIGHT - 6
|
||||
)
|
||||
ctx.restore()
|
||||
|
||||
if (node.has_errors) {
|
||||
ctx.save()
|
||||
ctx.font = 'bold 14px sans-serif'
|
||||
const sz2 = ctx.measureText(node.type)
|
||||
ctx.fillStyle = 'white'
|
||||
ctx.fillText(
|
||||
node.type,
|
||||
node.size[0] / 2 - sz2.width / 2,
|
||||
node.size[1] / 2
|
||||
)
|
||||
ctx.restore()
|
||||
}
|
||||
}
|
||||
}
|
||||
return r
|
||||
}
|
||||
|
||||
app.registerExtension({
|
||||
name: 'Comfy.Mixlab.ui',
|
||||
init () {
|
||||
@@ -1397,6 +1463,47 @@ app.registerExtension({
|
||||
|
||||
// createNodesCharts()
|
||||
},
|
||||
nodeCreated (node) {
|
||||
if (node.widgets) {
|
||||
// Locate dynamic prompt text widgets
|
||||
// Include any widgets with dynamicPrompts set to true, and customtext
|
||||
|
||||
for (let index = 0; index < node.widgets.length; index++) {
|
||||
const widget = node.widgets[index]
|
||||
if (
|
||||
(widget.type === 'customtext' && widget.dynamicPrompts !== false) ||
|
||||
widget.dynamicPrompts
|
||||
) {
|
||||
widget.element.classList.add('dynamic_prompt')
|
||||
|
||||
widget.element.addEventListener('mouseover', e => {
|
||||
// console.log(node.widgets_values[index])
|
||||
if (node.widgets_values[index])
|
||||
widget.element.setAttribute('title', node.widgets_values[index])
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 右上角的badge是否已经绘制
|
||||
if (!node.badge_enabled) {
|
||||
if (!node.getNickname) {
|
||||
node.getNickname = function () {
|
||||
if (node.nickname) {
|
||||
return node.nickname
|
||||
}
|
||||
return
|
||||
// return getNickname(node, node.comfyClass.trim())
|
||||
}
|
||||
}
|
||||
|
||||
const orig = node.__proto__.onDrawForeground
|
||||
node.onDrawForeground = function (ctx) {
|
||||
drawBadge(node, orig, arguments)
|
||||
}
|
||||
node.badge_enabled = true
|
||||
}
|
||||
},
|
||||
async loadedGraphNode (node, app) {
|
||||
// console.log(
|
||||
// '#ui init',
|
||||
|
||||
Reference in New Issue
Block a user