From 1f8b19fe5effa1b7ab9811fa4afa76c474feac5d Mon Sep 17 00:00:00 2001
From: yolain <yolain@163.com>
Date: Sat, 20 Jul 2024 21:01:56 +0800
Subject: [PATCH] Add inspyrenet to easy imageRembg

---
 README.en.md                      |  1 +
 README.md                         |  1 +
 py/image.py                       | 41 ++++++++++++++++++++++---------
 web/js/easy/easyDynamicWidgets.js | 11 +++++++++
 4 files changed, 43 insertions(+), 11 deletions(-)

diff --git a/README.en.md b/README.en.md
index 01c3151..2d49250 100644
--- a/README.en.md
+++ b/README.en.md
@@ -56,6 +56,7 @@ Double-click install.bat to install the required dependencies
 
 **v1.2.1**
 
+- Added **inspyrenet** to `easy imageRemBg` 
 - Added `easy controlnetLoader++`
 - Added **PLUS (kolors genernal)** preset to `easy ipadapterApply` and `easy ipadapterApplyADV` (Supported kolors ipadapter）
 - Added `easy kolorsLoader` - Code based on [MinusZoneAI](https://github.com/MinusZoneAI/ComfyUI-Kolors-MZ)'s and [kijai](https://github.com/kijai/ComfyUI-KwaiKolorsWrapper)'s repo, thanks for their contribution.
diff --git a/README.md b/README.md
index 135d92c..557d2a2 100644
--- a/README.md
+++ b/README.md
@@ -64,6 +64,7 @@ git clone https://github.com/yolain/ComfyUI-Easy-Use
 
 **v1.2.1**
 
+- `easy imageRemBg` 增加 **inspyrenet** 选项
 - 增加 `easy controlnetLoader++`
 - `easy ipadapterApply` 和 `easy ipadapterApplyADV` 增加 **PLUS (kolors genernal)** 预置项 (支持kolors IPA，可能会改名）
 - 去除 `easy positive` `easy negative` 等prompt节点的自动将中文翻译功能，自动翻译仅在 `easy a1111Loader` 等不支持中文TE的加载器中生效
diff --git a/py/image.py b/py/image.py
index 407000e..3a5ef11 100644
--- a/py/image.py
+++ b/py/image.py
@@ -674,9 +674,13 @@ class imageRemBg:
     return {
       "required": {
         "images": ("IMAGE",),
-        "rem_mode": (("RMBG-1.4",),),
+        "rem_mode": (("RMBG-1.4","Inspyrenet"),),
         "image_output": (["Hide", "Preview", "Save", "Hide/Save"], {"default": "Preview"}),
         "save_prefix": ("STRING", {"default": "ComfyUI"}),
+
+      },
+      "optional":{
+        "torchscript_jit": ("BOOLEAN", {"default": False}),
       },
       "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
     }
@@ -688,7 +692,9 @@ class imageRemBg:
 
   CATEGORY = "EasyUse/Image"
 
-  def remove(self, rem_mode, images, image_output, save_prefix, prompt=None, extra_pnginfo=None):
+  def remove(self, rem_mode, images, image_output, save_prefix, torchscript_jit=False, prompt=None, extra_pnginfo=None):
+    new_images = list()
+    masks = list()
     if rem_mode == "RMBG-1.4":
       # load model
       model_url = REMBG_MODELS[rem_mode]['model_url']
@@ -702,8 +708,6 @@ class imageRemBg:
       net.eval()
       # prepare input
       model_input_size = [1024, 1024]
-      new_images = list()
-      masks = list()
       for image in images:
         orig_im = tensor2pil(image)
         w, h = orig_im.size
@@ -721,18 +725,33 @@ class imageRemBg:
       new_images = torch.cat(new_images, dim=0)
       masks = torch.cat(masks, dim=0)
 
+    elif rem_mode == "Inspyrenet":
+      from tqdm import tqdm
+      try:
+        from transparent_background import Remover
+      except:
+          install_package("transparent_background")
+          from transparent_background import Remover
 
-      results = easySave(new_images, save_prefix, image_output, prompt, extra_pnginfo)
+      remover = Remover(jit=torchscript_jit)
 
-      if image_output in ("Hide", "Hide/Save"):
-        return {"ui": {},
-                "result": (new_images, masks)}
+      for img in tqdm(images, "Inspyrenet Rembg"):
+        mid = remover.process(tensor2pil(img), type='rgba')
+        out = pil2tensor(mid)
+        new_images.append(out)
+        mask = out[:, :, :, 3]
+        masks.append(mask)
+      new_images = torch.cat(new_images, dim=0)
+      masks = torch.cat(masks, dim=0)
 
-      return {"ui": {"images": results},
+    results = easySave(new_images, save_prefix, image_output, prompt, extra_pnginfo)
+
+    if image_output in ("Hide", "Hide/Save"):
+      return {"ui": {},
               "result": (new_images, masks)}
 
-    else:
-      return (None, None)
+    return {"ui": {"images": results},
+            "result": (new_images, masks)}
 
 # 图像选择器
 class imageChooser(PreviewImage):
diff --git a/web/js/easy/easyDynamicWidgets.js b/web/js/easy/easyDynamicWidgets.js
index ebd875c..c75f7d5 100644
--- a/web/js/easy/easyDynamicWidgets.js
+++ b/web/js/easy/easyDynamicWidgets.js
@@ -442,6 +442,17 @@ function widgetLogic(node, widget) {
 		}
 		updateNodeHeight(node)
 	}
+
+	if(widget.name == 'rem_mode'){
+		switch (widget.value){
+			case 'Inspyrenet':
+				toggleWidget(node, findWidgetByName(node, 'torchscript_jit'), true)
+				break
+			default:
+				toggleWidget(node, findWidgetByName(node, 'torchscript_jit'), false)
+				break
+		}
+	}
 }
 
 function widgetLogic2(node, widget) {