v0.28.0

add rembg api & webapp rembg
2024-05-16 11:49:45 +08:00 · 2024-05-16 11:49:14 +08:00 · 2024-05-16 09:02:13 +08:00 · 2024-05-16 09:02:08 +08:00
6 changed files with 302 additions and 75 deletions
@@ -8,6 +8,8 @@ import hashlib
 import datetime
 import folder_paths
 import logging
+import base64,io,re
+from PIL import Image
 from comfy.cli_args import args
 python = sys.executable

@@ -30,6 +32,8 @@ except:
    print("##nodes.ChatGPT ImportError")


+from .nodes.RembgNode import get_rembg_models,U2NET_HOME,run_briarmbg,run_rembg
+
 from server import PromptServer

 try:
@@ -98,6 +102,26 @@ install_openai()
 current_path = os.path.abspath(os.path.dirname(__file__))


+def remove_base64_prefix(base64_str):
+  """
+  去除 base64 字符串中的 data:image/*;base64, 前缀
+
+  Args:
+    base64_str: base64 编码的字符串
+
+  Returns:
+    去除前缀后的 base64 字符串
+  """
+
+  # 使用正则表达式匹配常见的前缀
+  pattern = r'^data:image\/(.*);base64,(.+)$'
+  match = re.match(pattern, base64_str)
+  if match:
+    # 如果匹配到常见的前缀，则去除前缀并返回
+    return match.group(2)
+  else:
+    # 如果不匹配到常见的前缀，则直接返回
+    return base64_str

 def calculate_md5(string):
    encoded_string = string.encode()
@@ -629,10 +653,53 @@ async def get_checkpoints(request):
            names=get_llama_models()
    except:
        print("llamafile none")
+
+    try:
+        if data['type']=='rembg':
+            names=get_rembg_models(U2NET_HOME)
+    except:
+        print("rembg none")
    
    return web.json_response({"names":names,"types":list(folder_paths.folder_names_and_paths.keys())})


+@routes.post('/mixlab/rembg')
+async def rembg_hander(request):
+    data = await request.json()
+    model=data['model']
+    result={}
+
+    data_base64=remove_base64_prefix(data['base64'])
+    image_data = base64.b64decode(data_base64)
+    
+    # 创建一个BytesIO对象
+    image_stream = io.BytesIO(image_data)
+
+    # 使用PIL Image模块读取图像
+    image = Image.open(image_stream)
+
+    if model=='briarmbg':
+        _,rgba_images,_=run_briarmbg([image])
+    else:
+        _,rgba_images,_=run_rembg(model,[image])
+
+    with io.BytesIO() as buf:
+        rgba_images[0].save(buf, format='PNG')
+        img_bytes = buf.getvalue()
+    img_base64 = base64.b64encode(img_bytes).decode('utf-8')
+    
+    try: 
+        result={
+            'data':img_base64,
+            'model':model,
+            'status':'success',
+            }
+    except Exception as e:
+            print(e)
+
+    return web.json_response(result)
+
+
@routes.post("/mixlab/prompt_result")
 async def post_prompt_result(request):
    data = await request.json()
@@ -476,6 +476,24 @@ os.environ["U2NET_HOME"] = U2NET_HOME
 global _available
 _available=False

+
+def get_rembg_models(path):
+    """从目录中获取文件并提取文件名
+    Args:
+        path: 目录路径
+    Returns:
+        文件名列表
+    """
+    filenames = []
+    for root, _, files in os.walk(path):
+        for filename in files:
+        # 过滤隐藏文件
+            if not filename.startswith('.'):
+                name, ext = os.path.splitext(os.path.basename(filename))
+                filenames.append(name)
+    return filenames
+
+
 def is_installed(package):
    try:
        spec = importlib.util.find_spec(package)
@@ -509,7 +527,7 @@ except:
    _available=False


-def briarmbg_run(images=[]):
+def run_briarmbg(images=[]):
    mroot=os.path.join(folder_paths.models_dir, "rembg")
    m=os.path.join(mroot,'briarmbg.pth')
    if os.path.exists(m)==False:
@@ -573,7 +591,7 @@ def briarmbg_run(images=[]):
    return (masks,rgba_images,rgb_images)   


-def run_bg(model_name= "unet",images=[]):
+def run_rembg(model_name= "unet",images=[]):
    # model_name = "unet" # "isnet-general-use"
    rembg_session = new_session(model_name)
    masks=[]
@@ -643,17 +661,7 @@ class RembgNode_:
    def INPUT_TYPES(s):
        return {"required": {
            "image": ("IMAGE",),
-            "model_name": ([
-                            "briarmbg",
-                            "u2net",
-                            "u2netp",
-                            "u2net_human_seg",
-                            "u2net_cloth_seg",
-                            "silueta",
-                            "isnet-general-use",
-                            "isnet-anime",
-                            
-                            ],),
+            "model_name": (get_rembg_models(U2NET_HOME),),
            
                             },
                }
@@ -681,9 +689,9 @@ class RembgNode_:
                images.append(im)

        if model_name=='briarmbg':
-            masks,rgba_images,rgb_images=briarmbg_run(images)
+            masks,rgba_images,rgb_images=run_briarmbg(images)
        else:
-            masks,rgba_images,rgb_images=run_bg(model_name,images)
+            masks,rgba_images,rgb_images=run_rembg(model_name,images)

        masks=[pil2tensor(m) for m in masks]

@@ -247,7 +247,7 @@

        .card textarea {
            width: 100%;
-            height:'fit-content';
+            height: 'fit-content';
            /* min-width: 300px; */
            margin-top: 12px;
            resize: none;
@@ -474,15 +474,15 @@
            <a class="link" href="https://www.mixcomfy.com" target="_blank">ComfyUI中文爱好者社区推荐</a>
        </div>

-        <a target="_blank" id="login_btn" href="https://www.mixcomfy.com/blog/" style="text-decoration: none;
+        <a id="login_btn" target="_blank" href="https://discord.gg/xbP2GZF6gn" 
+        style="text-decoration: none;
        color: black;font-size:12px">
            <svg height="32" aria-hidden="true" viewBox="0 0 16 16" version="1.1" width="32" data-view-component="true"
                class="octicon octicon-mark-github v-align-middle color-fg-default">
                <path
                    d="M8 0c4.42 0 8 3.58 8 8a8.013 8.013 0 0 1-5.45 7.59c-.4.08-.55-.17-.55-.38 0-.27.01-1.13.01-2.2 0-.75-.25-1.23-.54-1.48 1.78-.2 3.65-.88 3.65-3.95 0-.88-.31-1.59-.82-2.15.08-.2.36-1.02-.08-2.12 0 0-.67-.22-2.2.82-.64-.18-1.32-.27-2-.27-.68 0-1.36.09-2 .27-1.53-1.03-2.2-.82-2.2-.82-.44 1.1-.16 1.92-.08 2.12-.51.56-.82 1.28-.82 2.15 0 3.06 1.86 3.75 3.64 3.95-.23.2-.44.55-.51 1.07-.46.21-1.61.55-2.33-.66-.15-.24-.6-.83-1.23-.82-.67.01-.27.38.01.53.34.19.73.9.82 1.13.16.45.68 1.31 2.69.94 0 .67.01 1.3.01 1.49 0 .21-.15.45-.55.38A7.995 7.995 0 0 1 0 8c0-4.42 3.58-8 8-8Z">
                </path>
-            </svg> Community</a>
-
+            </svg> HELP/帮助</a>

    </div>
    <a id="author"></a>
@@ -680,11 +680,90 @@
        }


+        // 获取 rembg 模型
+        async function get_rembg_models() {
+            try {
+                const response = await fetch(`${get_url()}/mixlab/folder_paths`, {
+                    method: 'POST',
+                    headers: {
+                        'Content-Type': 'application/json'
+                    },
+                    body: JSON.stringify({
+                        type: 'rembg'
+                    })
+                })
+
+                const data = await response.json()
+                // console.log(data)
+                return data.names
+            } catch (error) {
+                console.error(error)
+            }
+        }
+
+        //自动抠图
+        async function run_rembg(model, base64) {
+            try {
+                const response = await fetch(`${get_url()}/mixlab/rembg`, {
+                    method: 'POST',
+                    headers: {
+                        'Content-Type': 'application/json'
+                    },
+                    body: JSON.stringify({
+                        model,
+                        base64
+                    })
+                })
+
+                const data = await response.json()
+                // console.log(data)
+                return data.data
+            } catch (error) {
+                console.error(error)
+            }
+        }
+
+        function convertImageToBlackBasedOnAlpha(image) {
+            const canvas = document.createElement('canvas');
+            const ctx = canvas.getContext('2d');
+
+            // Draw the image onto the canvas
+            canvas.width = image.width;
+            canvas.height = image.height;
+            ctx.drawImage(image, 0, 0);
+
+            // Get the image data from the canvas
+            const imageData = ctx.getImageData(0, 0, canvas.width, canvas.height);
+            const pixels = imageData.data;
+
+            // Modify the RGB values based on the alpha channel
+            for (let i = 0; i < pixels.length; i += 4) {
+                const alpha = pixels[i + 3];
+                if (alpha !== 0) {
+                    // Set non-transparent pixels to black
+                    pixels[i] = 0; // Red
+                    pixels[i + 1] = 0; // Green
+                    pixels[i + 2] = 0; // Blue
+                }
+            }
+
+            // Put the modified image data back onto the canvas
+            ctx.putImageData(imageData, 0, 0);
+
+            // Convert the modified canvas to base64 data URL
+            const base64ImageData = canvas.toDataURL('image/png'); // Replace 'png' with your desired image format
+
+            return base64ImageData;
+        }
+
+
+
        // 图像编辑
        async function editImage(image, data) {
            //判断mask是否有输出
            let isMask = data.options.hasMask;
            console.log(data)
+
            //app
            document.body.querySelector('.app').style.display = 'none'
            document.body.querySelector('#author').style.display = 'none'
@@ -693,11 +772,103 @@
            editor.style.display = 'block';
            const iframe = editor.querySelector('iframe');

+            // 复原
+            const resetLayer = () => {
+                let Layers = iframe.contentWindow.Layers;
+                for (const layer of Layers.get_layers()) {
+                    layer.visible = true;
+                }
+                Layers.refresh_gui()
+            }
+            //取image
+            const getImageBase64FromLayer = () => {
+                let Layers = iframe.contentWindow.Layers;
+                let tempCanvas = document.createElement("canvas");
+                let tempCtx = tempCanvas.getContext("2d");
+                let dim = Layers.get_dimensions();
+                tempCanvas.width = dim.width;
+                tempCanvas.height = dim.height;
+                for (const layer of Layers.get_layers()) {
+                    if (layer.name === 'Image_' + data.id) {
+                        layer.visible = true;
+                    } else {
+                        layer.visible = false;
+                    }
+                }
+                Layers.refresh_gui()
+
+                Layers.convert_layers_to_canvas(tempCtx);
+                return tempCanvas.toDataURL()
+            }
+
+            //add mask
+            const addMask = (id, name, image) => {
+                let Layers = iframe.contentWindow.Layers;
+                var new_mask_layer = {
+                    id,
+                    name,
+                    type: 'brush',
+                    data: [],
+                    render_function: ['brush', 'render'],
+                    width: image.naturalWidth || image.width,
+                    height: image.naturalHeight || image.height,
+
+                };
+                Layers.insert(new_mask_layer);
+            }
+
+            //add image
+            const addImage = (id, name, image) => {
+                let Layers = iframe.contentWindow.Layers;
+                var new_mask_layer = {
+                    id,
+                    name,
+                    type: 'image',
+                    data: image,
+                    width: image.naturalWidth || image.width,
+                    height: image.naturalHeight || image.height,
+                    width_original: image.naturalWidth || image.width,
+                    height_original: image.naturalHeight || image.height,
+                };
+                Layers.insert(new_mask_layer);
+            }
+
            //默认的画笔size设置大
            // let inputSize = (iframe.contentDocument.getElementById('size')).querySelector('input');
            // inputSize.value=50;
            // (iframe.contentDocument.getElementById('size')).querySelector('.increase_number').click()

+            //自动抠图
+            let autoMaskSelect = iframe.contentDocument.getElementById('automask_image_mixlab');
+            const select = iframe.contentDocument.getElementById('automask_models_mixlab');
+            let rembgModels = await get_rembg_models()
+            // 遍历模型列表并创建选项
+            for (const model of rembgModels) {
+                const option = document.createElement('option');
+                option.value = model;
+                option.textContent = model;
+                select.appendChild(option);
+            }
+
+            let autoMaskBtn = iframe.contentDocument.getElementById('automask_image_mixlab');
+            autoMaskBtn.addEventListener('click', async e => {
+                //api请求
+                let base64 = getImageBase64FromLayer()
+                resetLayer()
+
+                let res = await run_rembg(select.value, base64)
+                const match = res.match(/^data:image\/(\w+);base64,/);
+                if (!match) {
+                    res = 'data:image/png;base64,' + res
+                }
+
+                let image = await createImage(res)
+
+                let mb = convertImageToBlackBasedOnAlpha(image)
+                let mask = await createImage(mb)
+                addImage((Layers.get_layers()).length, 'Mask_' + data.id, mask)
+            })
+

            let cancelImageBtn = iframe.contentDocument.getElementById('cancel_image_mixlab');
            cancelImageBtn.addEventListener('click', e => {
@@ -707,6 +878,11 @@
            })
            // 获取 id 为 "mix" 的 button 元素
            let saveImageBtn = iframe.contentDocument.getElementById('save_image_mixlab');
+            saveImageBtn.style = `width: 98px;
+            height: 36px;
+            margin: 0 12px;
+            background-color: var(--background-color-active);
+            color: var(--text-color-active);`
            saveImageBtn.addEventListener('click', async e => {
                //保存，并更新图片
                e.preventDefault();
@@ -734,13 +910,7 @@
                    Layers.convert_layers_to_canvas(tempCtx);

                    // 复原
-                    for (const layer of Layers.get_layers()) {
-                        layer.visible = true;
-                    }
-                    Layers.refresh_gui()
-
-
-
+                    resetLayer()

                    // 获取图像数据
                    const imageData = tempCtx.getImageData(0, 0, dim.width, dim.height);
@@ -780,28 +950,22 @@


                    //更新图片
-
                    const canvas = document.createElement("canvas");
                    canvas.width = dim.width;
                    canvas.height = dim.height;

                    const ctx = canvas.getContext('2d');

-                    const defaultImage = new Image();
-                    defaultImage.src = data.options.defaultImage;
-                    defaultImage.onload = function () {
-                        ctx.drawImage(defaultImage, 0, 0, dim.width, dim.height);

-                        // 绘制base64图片
-                        const base64Image = base64
-                        const base64ImageObj = new Image();
-                        base64ImageObj.onload = function () {
-                            ctx.globalCompositeOperation = 'destination-in';
-                            ctx.drawImage(base64ImageObj, 0, 0, dim.width, dim.height);
-                            image.src = canvas.toDataURL();
-                        };
-                        base64ImageObj.src = base64;
-                    };
+                    const defaultImage = await createImage(data.options.defaultImage)
+                    ctx.drawImage(defaultImage, 0, 0, dim.width, dim.height);
+
+                    // 绘制base64图片
+                    // const base64Image = base64
+                    const base64ImageObj = await createImage(base64)
+                    ctx.globalCompositeOperation = 'destination-in';
+                    ctx.drawImage(base64ImageObj, 0, 0, dim.width, dim.height);
+                    image.src = canvas.toDataURL();

                }

@@ -817,36 +981,12 @@
            let layers1 = Layers.get_layers()

            if (!layers1.filter(l => l.name == 'Image_' + data.id)[0]) {
-                var new_layer = {
-                    id: (Layers.get_layers()).length,
-                    name: 'Image_' + data.id,
-                    type: 'image',
-                    data: image,
-                    width: image.naturalWidth || image.width,
-                    height: image.naturalHeight || image.height,
-                    width_original: image.naturalWidth || image.width,
-                    height_original: image.naturalHeight || image.height,
-                };
-
-                Layers.insert(new_layer);
+                addImage((Layers.get_layers()).length, 'Image_' + data.id, image)
            }

-
-
            if (isMask) {
                if (!layers1.filter(l => l.name == 'Mask_' + data.id)[0]) {
-                    var new_mask_layer = {
-                        id: (Layers.get_layers()).length,
-                        name: 'Mask_' + data.id,
-                        type: 'brush',
-                        data: [],
-                        render_function: ['brush', 'render'],
-                        width: image.naturalWidth || image.width,
-                        height: image.naturalHeight || image.height,
-
-                    };
-
-                    Layers.insert(new_mask_layer);
+                    addMask((Layers.get_layers()).length, 'Mask_' + data.id, image)
                }

            }
@@ -3,7 +3,7 @@ import { app } from '../../../scripts/app.js'
 const repoOwner = 'shadowcz007' // 替换为仓库的所有者
 const repoName = 'comfyui-mixlab-nodes' // 替换为仓库的名称

-const version = 'v0.27.1'
+const version = 'v0.28.0'

 fetch(`https://api.github.com/repos/${repoOwner}/${repoName}/releases/latest`)
  .then(response => response.json())
@@ -46,7 +46,7 @@ async function get_llamafile_models () {
    })

    const data = await response.json()
-    console.log(data)
+    // console.log(data)
    return data.names
  } catch (error) {
    console.error(error)
@@ -539,10 +539,6 @@ injectCSS(`::-webkit-scrollbar {
  border-left: 2px solid var(--input-text); 
 }
 
-.litegraph{
-  background: var(--bg-color)!important;
-}
-
 `)

 async function getCustomnodeMappings (mode = 'url') {
@@ -45,6 +45,15 @@
 		<div class="submenu">
 			<!-- <a class="logo" href="#">miniPaint</a> -->
 			<div class="block attributes" id="action_attributes"></div>
+			
+			<select id="automask_models_mixlab"></select>
+
+			<button id="automask_image_mixlab" type="button" style="width: 98px;
+			height: 36px;margin-right: 18px;
+			color: white;">
+				RemoveBg
+			</button>
+
 			<button id="cancel_image_mixlab" type="button" style="width: 98px;
 			height: 36px;
 			color: white;">
@@ -55,6 +64,13 @@
 			color: white;">
 				Save
 			</button>
+
+			<a  
+			target="_blank" href="https://discord.gg/xbP2GZF6gn" 
+			style="width: 98px;
+			color: white;
+			text-decoration: none;"> Help/帮助 </a>
+
 			<button class="undo_button" id="undo_button" type="button">
 				<span class="sr_only">Undo</span>
 			</button>
@@ -82,7 +98,7 @@
 		</div>

 		<div class="sidebar_right">
-			<div class="preview block">
+			<div class="preview block" style="display: none;">
 				<h2 class="trn toggle" data-target="toggle_preview">Preview</h2>
 				<div id="toggle_preview"></div>
 			</div>
@@ -92,7 +108,7 @@
 				<div class="content" id="toggle_colors"></div>
 			</div>

-			<div class="block" id="info_base">
+			<div class="block" id="info_base" style="display: none;">
 				<h2 class="trn toggle toggle-full" data-target="toggle_info">Information</h2>
 				<div class="content" id="toggle_info"></div>
 			</div>
Author	SHA1	Message	Date
shadowcz007	05857a92d5	v0.28.0 add rembg api & webapp rembg	2024-05-16 11:49:45 +08:00
shadowcz007	6bdc811286	add rembg api & webapp rembg	2024-05-16 11:49:14 +08:00
shadowcz007	469d50a5b8	Update index.html	2024-05-16 09:02:13 +08:00
shadowcz007	ef86904bfb	Update ui_mixlab.js	2024-05-16 09:02:08 +08:00