update

AIFSH · Sep 13, 2024 · dd8cd91 · dd8cd91
1 parent c63c758
commit dd8cd91
Show file tree

Hide file tree

Showing 27 changed files with 918 additions and 205 deletions.
diff --git a/README.md b/README.md
@@ -11,7 +11,10 @@ a custom node for [HivisionIDPhotos](https://github.com/Zeyi-Lin/HivisionIDPhoto
 ## 教程
 - [Demo](https://www.bilibili.com/video/BV1iFpvegEY3/)
 - [一键包](https://pan.quark.cn/s/b8422210f61a)
-
+
+# 🤩 最近更新
+- 2024.9.13 增加**美白**功能
+
 ## weights
 存到项目的`ComfyUI/custom_nodes/HivisionIDPhotos-ComfyUI/hivision/creator/weights`目录下：
 - `modnet_photographic_portrait_matting.onnx` (24.7MB): [MODNet](https://github.com/ZHKKKe/MODNet)官方权重，[下载](https://github.com/Zeyi-Lin/HivisionIDPhotos/releases/download/pretrained-model/modnet_photographic_portrait_matting.onnx)

diff --git a/__init__.py b/__init__.py
@@ -17,7 +17,7 @@
 )
 from hivision.demo.utils import csv_to_size_list,csv_to_color_list
 from hivision.creator.choose_handler import choose_handler, HUMAN_MATTING_MODELS
-from hivision.demo.utils import add_watermark 
+from hivision.utils import add_watermark
 
 size_list_dict_CN = csv_to_size_list(os.path.join(now_dir, "hivision/demo/assets/size_list_CN.csv"))
 size_list_CN = list(size_list_dict_CN.keys())
@@ -326,7 +326,14 @@ def INPUT_TYPES(s):
                     "step":0.01,
                     "round": 0.001,
                     "display":"slider"
-                })
+                }),
+                "whitening_strength":("INT",{
+                    "default": 2,
+                    "min":0,
+                    "max":15,
+                    "step":1,
+                    "display":"slider"
+                }),
             }
         }
 
@@ -340,7 +347,7 @@ def INPUT_TYPES(s):
     CATEGORY = "AIFSH_HivisionIDPhotos"
 
     def gen_img(self,input_img,normal_params,change_bg_only,matting_model,
-                face_detect_model,head_measure_ratio,top_distance):
+                face_detect_model,head_measure_ratio,top_distance,whitening_strength):
         creator = IDCreator()
 
         # ------------------- 人像抠图模型选择 -------------------
@@ -356,7 +363,8 @@ def gen_img(self,input_img,normal_params,change_bg_only,matting_model,
             result = creator(input_image, size=size,
                                 head_measure_ratio=head_measure_ratio,
                                 head_top_range=(top_distance, top_distance-0.02),
-                                change_bg_only=change_bg_only)
+                                change_bg_only=change_bg_only,
+                                whitening_strength=whitening_strength)
         except FaceError:
             print("人脸数量不等于 1，请上传单张人脸的图像。")
         else: