循环嵌套优化

中文搜索闪退问题修改

Signed-off-by: yanlang0123 <yanlang0123@gmail.com>

.gitignore

@@ -0,0 +1,21 @@
+# Mac
+.DS_Store
+**/.DS_Store
+
+# vim/vi
+*.swp
+
+# JavaScript
+node_modules/
+.node_modules/
+.eslintcache
+unpackage/dist/build/
+unpackage/dist/dev/
+
+# python
+*.pyc
+config/baidu.json
+ckpts/*
+config/lamWeChat.db
+config/weChat.yaml
+config/settings_nodes.json

README.en.md

@@ -8,166 +8,20 @@ Download and place in the plugin directory of comfyUI, as shown below:
 
 ![Alt text](解压存放路径及名称.png)
 
-##### Special Note for this version, execute the latest version "ComfyUI_windows_portable_nvidia_cu118_or_cpu.7z" with cu118 inside, the address is as follows:
-https://github.com/comfyanonymous/ComfyUI/releases/download/latest/ComfyUI_windows_portable_nvidia_cu118_or_cpu.7z
+##### Special Note: For this version, execute the latest version "ComfyUI_windows_portable_nvidia_cu121_or_cpu.7z" with cu121 inside. The version address is as follows:
+https://github.com/comfyanonymous/ComfyUI/releases/download/latest/ComfyUI_windows_portable_nvidia_cu121_or_cpu.7z
 
-1. Unzip according to the image to the specified plugin directory and name. Extract insightface.rar to the directory ..\ComfyUI_windows_portable\python_embeded\Lib\site-packages. Extract venv.rar to the directory ..\ComfyUI_windows_portable\python_embeded\Lib\site-packages. Then run the install.bat file, and if there are no errors, it's ready.
+Then run the install.bat file, and if there are no errors, it's ready.
 
-2. Model addresses and storage paths:
-   Lama model:
-   https://huggingface.co/lllyasviel/Annotators/resolve/main/ControlNetLama.pth  ..\ComfyUI\models\lama\ControlNetLama.pth
-   SadTalker models:
-   https://github.com/OpenTalker/SadTalker/releases/download/v0.0.2-rc/mapping_00109-model.pth.tar ..\ComfyUI\models\SadTalker\mapping_00109-model.pth.tar
-   https://github.com/OpenTalker/SadTalker/releases/download/v0.0.2-rc/mapping_00229-model.pth.tar ..\ComfyUI\models\SadTalker\mapping_00229-model.pth.tar
-   https://github.com/OpenTalker/SadTalker/releases/download/v0.0.2-rc/SadTalker_V0.0.2_256.safetensors ..\ComfyUI\models\SadTalker\SadTalker_V0.0.2_256.safetensors
-   https://github.com/OpenTalker/SadTalker/releases/download/v0.0.2-rc/SadTalker_V0.0.2_512.safetensors ..\ComfyUI\models\SadTalker\SadTalker_V0.0.2_512.safetensors
+1. Model addresses and storage paths:
+   - Lama model:
+     https://huggingface.co/lllyasviel/Annotators/resolve/main/ControlNetLama.pth  ..\ComfyUI\models\lama\ControlNetLama.pth
 
-3. Gfpgan is based on the startup directory, follow the prompts accordingly.
-   https://github.com/xinntao/facexlib/releases/download/v0.1.0/alignment_WFLW_4HG.pth ..\ComfyUI_windows_portable\gfpgan\weights\alignment_WFLW_4HG.pth 
-   https://github.com/xinntao/facexlib/releases/download/v0.1.0/detection_Resnet50_Final.pth ..\ComfyUI_windows_portable\gfpgan\weights\detection_Resnet50_Final.pth 
-   https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth ..\ComfyUI_windows_portable\gfpgan\weights\GFPGANv1.4.pth 
-   https://github.com/xinntao/facexlib/releases/download/v0.2.2/parsing_parsenet.pth ..\ComfyUI_windows_portable\gfpgan\weights\parsing_parsenet.pth 
+2. Image-face-fusion model (Face swapping model)
+   - Link: https://pan.baidu.com/s/19DOgJQ_RHNAjfNrzSr2uTQ?pwd=gf0p 
+   - Extraction Code: gf0p
+   - Extract to the directory ..\ComfyUI\models\image-face-fusion
 
-4. Image-face-fusion model (Face swapping model)
-   Link: https://pan.baidu.com/s/19DOgJQ_RHNAjfNrzSr2uTQ?pwd=gf0p 
-   Extraction Code: gf0p
-   Extract to the directory ..\ComfyUI\models\image-face-fusion
-
-5. Roop-face-swap model (Face swapping model)
-   Link: https://pan.baidu.com/s/1cJeRtqgdeNW21Hljwuv-tA?pwd=b4yy 
-   Extraction Code: b4yy
-   Extract to the directory ..\ComfyUI\models\roop-face-swap
-
-6. Modify the execution.py file in the comfyUI root directory. 
-搜索 “def recursive_execute” 
-位置参考图如下：
-![Alt text](修改位置.png)
-新增内容：
-```python
-#循环添加代码----------开始----------
-def get_del_keys(key, prompt,uniqueIds):
-    keys=[]
-    for k,v in prompt.items():
-        if k in uniqueIds :
-            continue
-        for k1,v1 in v['inputs'].items():
-            if type(v1)==list and v1[0]==key:
-                keys.append(k)
-                keys=keys+get_del_keys(k,prompt,uniqueIds)
-    return keys
-#循环添加代码-----------结束---------
-```
-```python
-#循环添加代码----------开始----------
-    startNum=None
-    startData={}
-    delKeys=[]
-    backhaul={}
-    clTypes=['ForInnerEnd','IfInnerExecute','DoWhileEnd']
-    oldPrompt=None
-    if class_type in clTypes:
-        oldPrompt=copy.deepcopy(prompt)
-
-    if class_type=='ForInnerEnd':
-        startNum=prompt[unique_id]['inputs']['total'][0]
-        inputNum=prompt[unique_id]['inputs']['obj'][0]
-        maxKeyStr=sorted(list(prompt.keys()), key=lambda x: int(x.split(':')[0]))[-1]
-        maxKey = int(maxKeyStr.split(':')[0])
-        delKeys=list(set(get_del_keys(startNum,prompt,[startNum,unique_id]))) 
-        delKeys.append(startNum)
-        delKeys = list(filter(lambda x: x != inputNum, delKeys))
-        for key in delKeys:
-            outputs.pop(key, None)
-        
-        startInput=prompt[startNum]['inputs']
-        if isinstance(startInput['total'],list) or isinstance(startInput['stop'],list) or isinstance(startInput['i'],list):
-            result = recursive_execute(server, prompt, outputs, startNum, extra_data, executed, prompt_id, outputs_ui, object_storage)
-            if result[0] is not True:
-                    return result
-            
-        for x in startInput:
-            input_data = startInput[x]
-            if isinstance(input_data, list):
-                startData[x]=outputs[input_data[0]][input_data[1]][0]
-            else:
-                startData[x]=input_data
-
-        for i in range(startData['i']+1,startData['total'],startData['stop']):
-            prompt[str(maxKey+i)]=prompt[inputNum]
-            prompt[unique_id]['inputs']['obj'+str(i)]=[str(maxKey+i),prompt[unique_id]['inputs']['obj'][-1]]
-            if i==startInput['i']+1:
-                backhaul['obj'+str(i)]=prompt[unique_id]['inputs']['obj']
-            else:
-                backhaul['obj'+str(i)]=[str(maxKey+i-startInput['stop']),prompt[unique_id]['inputs']['obj'][-1]]
-    elif class_type=='DoWhileEnd':
-        startNum=prompt[unique_id]['inputs']['start'][0]
-        inputNum=prompt[unique_id]['inputs']['ANY'][0]
-        delKeys=list(set(get_del_keys(startNum,prompt,[startNum,unique_id]))) 
-        delKeys.append(startNum)
-        delKeys.append(inputNum)
-
-    #循环添加代码-----------结束---------
-```
-```python
-                #循环添加代码----------开始----------
-                if class_type=='ForInnerEnd' and x !='obj' and x.startswith('obj'):
-                    if startNum!=None:
-                        if isinstance(prompt[startNum]['inputs']['i'],list):
-                            prompt[startNum]['inputs']['i']=startData['i']+startData['stop']
-                        else:
-                            prompt[startNum]['inputs']['i']=prompt[startNum]['inputs']['i']+startData['stop']
-                        prompt[startNum]['inputs']['obj']=backhaul[x]
-                        for key in delKeys:
-                            outputs.pop(key, None)
-                #循环添加代码-----------结束---------
-```
-```python
-                #判断选择添加代码-----------开始---------
-                if class_type=='DoWhileEnd' and x == 'ANY':
-                    any=outputs[inputNum][prompt[unique_id]['inputs']['ANY'][1]][0]
-                    i=0
-                    while any:
-                        for key in delKeys:
-                            outputs.pop(key, None)
-                        i=i+1
-                        outputs['i']=[[i]]
-                        prompt[startNum]['inputs']['i']=['i',0]
-                        result = recursive_execute(server, prompt, outputs, input_unique_id, extra_data, executed, prompt_id, outputs_ui, object_storage)
-                        if result[0] is not True:
-                            return result
-                        any=outputs[inputNum][prompt[unique_id]['inputs']['ANY'][1]][0]
-            
-            if class_type=='IfInnerExecute' and x=='ANY':
-                if outputs[input_unique_id][output_index][0]:
-                    inputs['IF_FALSE']=oldPrompt[unique_id]['inputs']['IF_TRUE']
-                else:
-                    inputs['IF_TRUE']=oldPrompt[unique_id]['inputs']['IF_FALSE']
-            #判断选择添加代码-----------结束---------
-```
-```python
-        #循环添加代码----------开始----------
-        if class_type=='ForInnerEnd' or class_type=='DoWhileEnd':
-            prompt[startNum]['inputs']=oldPrompt[startNum]['inputs']
-            prompt[unique_id]['inputs']=oldPrompt[unique_id]['inputs']
-            outputs.pop(startNum, None)
-            result = recursive_execute(server, prompt, outputs, startNum, extra_data, executed, prompt_id, outputs_ui, object_storage)
-            if result[0] is not True:
-                   return result
-        if class_type=='IfInnerExecute':
-            prompt[unique_id]['inputs']=oldPrompt[unique_id]['inputs']
-        #循环添加代码-----------结束---------
-```
-搜索 “def validate_prompt
-```python
-'''优化输出开始'''
-    inputKeys=[]
-    for k,v in prompt.items():
-        for k1,v1 in v['inputs'].items():
-            if type(v1)==list and len(v1)==2:
-                inputKeys.append(v1[0])
-    for x in prompt:
-        class_ = nodes.NODE_CLASS_MAPPINGS[prompt[x]['class_type']]
-        if hasattr(class_, 'OUTPUT_NODE') and class_.OUTPUT_NODE == True and x not in inputKeys:
-            outputs.add(x)
-    '''优化输出结束'''
-``` 
+3. 执行修改脚本
+windows  执行“修改文件.bat”文件，未报错后就可以了
+linux  执行“修改文件.sh”文件，未报错后就可以了

README.md

@@ -8,170 +8,21 @@
 
 ![Alt text](解压存放路径及名称.png)
 
-##### 特别说明该版本执行最新版“ComfyUI_windows_portable_nvidia_cu118_or_cpu.7z” 里面带cu118的版本地址如下：
-https://github.com/comfyanonymous/ComfyUI/releases/download/latest/ComfyUI_windows_portable_nvidia_cu118_or_cpu.7z
+##### 特别说明该版本执行最新版“ComfyUI_windows_portable_nvidia_cu121_or_cpu.7z” 里面带cu121的版本地址如下：
+https://github.com/comfyanonymous/ComfyUI/releases/download/latest/ComfyUI_windows_portable_nvidia_cu121_or_cpu.7z
 
-1. 根据图片解压到指定插件目录及名称，
-将insightface.rar解压到..\ComfyUI_windows_portable\python_embeded\Lib\site-packages目录
-将venv.rar解压到..\ComfyUI_windows_portable\python_embeded\Lib\site-packages目录
 然后运行install.bat文件，未报错后就可以了，
 
-2. 模型地址，及存放路径：
+1. 模型地址，及存放路径：
 lama模型：
 https://huggingface.co/lllyasviel/Annotators/resolve/main/ControlNetLama.pth  ..\ComfyUI\models\lama\ControlNetLama.pth
-SadTalker模型：
-https://github.com/OpenTalker/SadTalker/releases/download/v0.0.2-rc/mapping_00109-model.pth.tar ..\ComfyUI\models\SadTalker\mapping_00109-model.pth.tar
-https://github.com/OpenTalker/SadTalker/releases/download/v0.0.2-rc/mapping_00229-model.pth.tar ..\ComfyUI\models\SadTalker\mapping_00229-model.pth.tar
-https://github.com/OpenTalker/SadTalker/releases/download/v0.0.2-rc/SadTalker_V0.0.2_256.safetensors ..\ComfyUI\models\SadTalker\SadTalker_V0.0.2_256.safetensors
-https://github.com/OpenTalker/SadTalker/releases/download/v0.0.2-rc/SadTalker_V0.0.2_512.safetensors ..\ComfyUI\models\SadTalker\SadTalker_V0.0.2_512.safetensors
 
-3. gfpgan是根据启动目录来的，根据提示对应
-https://github.com/xinntao/facexlib/releases/download/v0.1.0/alignment_WFLW_4HG.pth ..\ComfyUI_windows_portable\gfpgan\weights\alignment_WFLW_4HG.pth 
-https://github.com/xinntao/facexlib/releases/download/v0.1.0/detection_Resnet50_Final.pth ..\ComfyUI_windows_portable\gfpgan\weights\detection_Resnet50_Final.pth 
-https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth ..\ComfyUI_windows_portable\gfpgan\weights\GFPGANv1.4.pth 
-https://github.com/xinntao/facexlib/releases/download/v0.2.2/parsing_parsenet.pth ..\ComfyUI_windows_portable\gfpgan\weights\parsing_parsenet.pth 
-
-4. image-face-fusion模型换脸模型
+2. image-face-fusion模型换脸模型
 链接：https://pan.baidu.com/s/19DOgJQ_RHNAjfNrzSr2uTQ?pwd=gf0p 
 提取码：gf0p
 解压到 ..\ComfyUI\models\image-face-fusion 目录
 
-5. roop-face-swap模型换脸模型
-链接：https://pan.baidu.com/s/1cJeRtqgdeNW21Hljwuv-tA?pwd=b4yy 
-提取码：b4yy
-解压到 ..\ComfyUI\models\roop-face-swap 目录
 
-
-6. 修改comfyUI根目录下的execution.py文件，修改内容
-搜索 “def recursive_execute” 
-位置参考图如下：
-![Alt text](修改位置.png)
-新增内容：
-```python
-#循环添加代码----------开始----------
-def get_del_keys(key, prompt,uniqueIds):
-    keys=[]
-    for k,v in prompt.items():
-        if k in uniqueIds :
-            continue
-        for k1,v1 in v['inputs'].items():
-            if type(v1)==list and v1[0]==key:
-                keys.append(k)
-                keys=keys+get_del_keys(k,prompt,uniqueIds)
-    return keys
-#循环添加代码-----------结束---------
-```
-```python
-#循环添加代码----------开始----------
-    startNum=None
-    startData={}
-    delKeys=[]
-    backhaul={}
-    clTypes=['ForInnerEnd','IfInnerExecute','DoWhileEnd']
-    oldPrompt=None
-    if class_type in clTypes:
-        oldPrompt=copy.deepcopy(prompt)
-
-    if class_type=='ForInnerEnd':
-        startNum=prompt[unique_id]['inputs']['total'][0]
-        inputNum=prompt[unique_id]['inputs']['obj'][0]
-        maxKeyStr=sorted(list(prompt.keys()), key=lambda x: int(x.split(':')[0]))[-1]
-        maxKey = int(maxKeyStr.split(':')[0])
-        delKeys=list(set(get_del_keys(startNum,prompt,[startNum,unique_id]))) 
-        delKeys.append(startNum)
-        delKeys = list(filter(lambda x: x != inputNum, delKeys))
-        for key in delKeys:
-            outputs.pop(key, None)
-        
-        startInput=prompt[startNum]['inputs']
-        if isinstance(startInput['total'],list) or isinstance(startInput['stop'],list) or isinstance(startInput['i'],list):
-            result = recursive_execute(server, prompt, outputs, startNum, extra_data, executed, prompt_id, outputs_ui, object_storage)
-            if result[0] is not True:
-                    return result
-            
-        for x in startInput:
-            input_data = startInput[x]
-            if isinstance(input_data, list):
-                startData[x]=outputs[input_data[0]][input_data[1]][0]
-            else:
-                startData[x]=input_data
-
-        for i in range(startData['i']+1,startData['total'],startData['stop']):
-            prompt[str(maxKey+i)]=prompt[inputNum]
-            prompt[unique_id]['inputs']['obj'+str(i)]=[str(maxKey+i),prompt[unique_id]['inputs']['obj'][-1]]
-            if i==startInput['i']+1:
-                backhaul['obj'+str(i)]=prompt[unique_id]['inputs']['obj']
-            else:
-                backhaul['obj'+str(i)]=[str(maxKey+i-startInput['stop']),prompt[unique_id]['inputs']['obj'][-1]]
-    elif class_type=='DoWhileEnd':
-        startNum=prompt[unique_id]['inputs']['start'][0]
-        inputNum=prompt[unique_id]['inputs']['ANY'][0]
-        delKeys=list(set(get_del_keys(startNum,prompt,[startNum,unique_id]))) 
-        delKeys.append(startNum)
-        delKeys.append(inputNum)
-
-    #循环添加代码-----------结束---------
-```
-```python
-                #循环添加代码----------开始----------
-                if class_type=='ForInnerEnd' and x !='obj' and x.startswith('obj'):
-                    if startNum!=None:
-                        if isinstance(prompt[startNum]['inputs']['i'],list):
-                            prompt[startNum]['inputs']['i']=startData['i']+startData['stop']
-                        else:
-                            prompt[startNum]['inputs']['i']=prompt[startNum]['inputs']['i']+startData['stop']
-                        prompt[startNum]['inputs']['obj']=backhaul[x]
-                        for key in delKeys:
-                            outputs.pop(key, None)
-                #循环添加代码-----------结束---------
-```
-```python
-                #判断选择添加代码-----------开始---------
-                if class_type=='DoWhileEnd' and x == 'ANY':
-                    any=outputs[inputNum][prompt[unique_id]['inputs']['ANY'][1]][0]
-                    i=0
-                    while any:
-                        for key in delKeys:
-                            outputs.pop(key, None)
-                        i=i+1
-                        outputs['i']=[[i]]
-                        prompt[startNum]['inputs']['i']=['i',0]
-                        result = recursive_execute(server, prompt, outputs, input_unique_id, extra_data, executed, prompt_id, outputs_ui, object_storage)
-                        if result[0] is not True:
-                            return result
-                        any=outputs[inputNum][prompt[unique_id]['inputs']['ANY'][1]][0]
-            
-            if class_type=='IfInnerExecute' and x=='ANY':
-                if outputs[input_unique_id][output_index][0]:
-                    inputs['IF_FALSE']=oldPrompt[unique_id]['inputs']['IF_TRUE']
-                else:
-                    inputs['IF_TRUE']=oldPrompt[unique_id]['inputs']['IF_FALSE']
-            #判断选择添加代码-----------结束---------
-```
-```python
-        #循环添加代码----------开始----------
-        if class_type=='ForInnerEnd' or class_type=='DoWhileEnd':
-            prompt[startNum]['inputs']=oldPrompt[startNum]['inputs']
-            prompt[unique_id]['inputs']=oldPrompt[unique_id]['inputs']
-            outputs.pop(startNum, None)
-            result = recursive_execute(server, prompt, outputs, startNum, extra_data, executed, prompt_id, outputs_ui, object_storage)
-            if result[0] is not True:
-                   return result
-        if class_type=='IfInnerExecute':
-            prompt[unique_id]['inputs']=oldPrompt[unique_id]['inputs']
-        #循环添加代码-----------结束---------
-```
-搜索 “def validate_prompt
-```python
-'''优化输出开始'''
-    inputKeys=[]
-    for k,v in prompt.items():
-        for k1,v1 in v['inputs'].items():
-            if type(v1)==list and len(v1)==2:
-                inputKeys.append(v1[0])
-    for x in prompt:
-        class_ = nodes.NODE_CLASS_MAPPINGS[prompt[x]['class_type']]
-        if hasattr(class_, 'OUTPUT_NODE') and class_.OUTPUT_NODE == True and x not in inputKeys:
-            outputs.add(x)
-    '''优化输出结束'''
-``` 
+3. 执行修改脚本
+windows  执行“修改文件.bat”文件，未报错后就可以了
+linux  执行“修改文件.sh”文件，未报错后就可以了

__init__.py

@@ -7,6 +7,7 @@ from .lam import init, get_ext_dir
 repo_dir = os.path.dirname(os.path.realpath(__file__))
 sys.path.insert(0, repo_dir)
 original_modules = sys.modules.copy()
+
 NODE_CLASS_MAPPINGS = {}
 NODE_DISPLAY_NAME_MAPPINGS = {}
 
@@ -21,5 +22,5 @@ if init():
         NODE_CLASS_MAPPINGS = {**NODE_CLASS_MAPPINGS, **imported_module.NODE_CLASS_MAPPINGS}
         NODE_DISPLAY_NAME_MAPPINGS = {**NODE_DISPLAY_NAME_MAPPINGS, **imported_module.NODE_DISPLAY_NAME_MAPPINGS}
 
-WEB_DIRECTORY = "./js"
+WEB_DIRECTORY = "./js" 
 __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS","WEB_DIRECTORY"]

backup/execution.py

@@ -0,0 +1,854 @@
+import sys
+import copy
+import logging
+import threading
+import heapq
+import traceback
+import inspect
+from typing import List, Literal, NamedTuple, Optional
+
+import torch
+import nodes
+
+import comfy.model_management
+
+def get_input_data(inputs, class_def, unique_id, outputs={}, prompt={}, extra_data={}):
+    valid_inputs = class_def.INPUT_TYPES()
+    input_data_all = {}
+    for x in inputs:
+        input_data = inputs[x]
+        if isinstance(input_data, list):
+            input_unique_id = input_data[0]
+            output_index = input_data[1]
+            if input_unique_id not in outputs:
+                input_data_all[x] = (None,)
+                continue
+            obj = outputs[input_unique_id][output_index]
+            input_data_all[x] = obj
+        else:
+            if ("required" in valid_inputs and x in valid_inputs["required"]) or ("optional" in valid_inputs and x in valid_inputs["optional"]):
+                input_data_all[x] = [input_data]
+
+    if "hidden" in valid_inputs:
+        h = valid_inputs["hidden"]
+        for x in h:
+            if h[x] == "PROMPT":
+                input_data_all[x] = [prompt]
+            if h[x] == "EXTRA_PNGINFO":
+                input_data_all[x] = [extra_data.get('extra_pnginfo', None)]
+            if h[x] == "UNIQUE_ID":
+                input_data_all[x] = [unique_id]
+    return input_data_all
+
+def map_node_over_list(obj, input_data_all, func, allow_interrupt=False):
+    # check if node wants the lists
+    input_is_list = False
+    if hasattr(obj, "INPUT_IS_LIST"):
+        input_is_list = obj.INPUT_IS_LIST
+
+    if len(input_data_all) == 0:
+        max_len_input = 0
+    else:
+        max_len_input = max([len(x) for x in input_data_all.values()])
+     
+    # get a slice of inputs, repeat last input when list isn't long enough
+    def slice_dict(d, i):
+        d_new = dict()
+        for k,v in d.items():
+            d_new[k] = v[i if len(v) > i else -1]
+        return d_new
+    
+    results = []
+    if input_is_list:
+        if allow_interrupt:
+            nodes.before_node_execution()
+        results.append(getattr(obj, func)(**input_data_all))
+    elif max_len_input == 0:
+        if allow_interrupt:
+            nodes.before_node_execution()
+        results.append(getattr(obj, func)())
+    else:
+        for i in range(max_len_input):
+            if allow_interrupt:
+                nodes.before_node_execution()
+            results.append(getattr(obj, func)(**slice_dict(input_data_all, i)))
+    return results
+
+def get_output_data(obj, input_data_all):
+    
+    results = []
+    uis = []
+    return_values = map_node_over_list(obj, input_data_all, obj.FUNCTION, allow_interrupt=True)
+
+    for r in return_values:
+        if isinstance(r, dict):
+            if 'ui' in r:
+                uis.append(r['ui'])
+            if 'result' in r:
+                results.append(r['result'])
+        else:
+            results.append(r)
+    
+    output = []
+    if len(results) > 0:
+        # check which outputs need concatenating
+        output_is_list = [False] * len(results[0])
+        if hasattr(obj, "OUTPUT_IS_LIST"):
+            output_is_list = obj.OUTPUT_IS_LIST
+
+        # merge node execution results
+        for i, is_list in zip(range(len(results[0])), output_is_list):
+            if is_list:
+                output.append([x for o in results for x in o[i]])
+            else:
+                output.append([o[i] for o in results])
+
+    ui = dict()    
+    if len(uis) > 0:
+        ui = {k: [y for x in uis for y in x[k]] for k in uis[0].keys()}
+    return output, ui
+
+def format_value(x):
+    if x is None:
+        return None
+    elif isinstance(x, (int, float, bool, str)):
+        return x
+    else:
+        return str(x)
+
+def recursive_execute(server, prompt, outputs, current_item, extra_data, executed, prompt_id, outputs_ui, object_storage):
+    unique_id = current_item
+    inputs = prompt[unique_id]['inputs']
+    class_type = prompt[unique_id]['class_type']
+    class_def = nodes.NODE_CLASS_MAPPINGS[class_type]
+    if unique_id in outputs:
+        return (True, None, None)
+
+    for x in inputs:
+        input_data = inputs[x]
+
+        if isinstance(input_data, list):
+            input_unique_id = input_data[0]
+            output_index = input_data[1]
+            if input_unique_id not in outputs:
+                result = recursive_execute(server, prompt, outputs, input_unique_id, extra_data, executed, prompt_id, outputs_ui, object_storage)
+                if result[0] is not True:
+                    # Another node failed further upstream
+                    return result
+
+    input_data_all = None
+    try:
+        input_data_all = get_input_data(inputs, class_def, unique_id, outputs, prompt, extra_data)
+        if server.client_id is not None:
+            server.last_node_id = unique_id
+            server.send_sync("executing", { "node": unique_id, "prompt_id": prompt_id }, server.client_id)
+
+        obj = object_storage.get((unique_id, class_type), None)
+        if obj is None:
+            obj = class_def()
+            object_storage[(unique_id, class_type)] = obj
+
+        output_data, output_ui = get_output_data(obj, input_data_all)
+        outputs[unique_id] = output_data
+        if len(output_ui) > 0:
+            outputs_ui[unique_id] = output_ui
+            if server.client_id is not None:
+                server.send_sync("executed", { "node": unique_id, "output": output_ui, "prompt_id": prompt_id }, server.client_id)
+    except comfy.model_management.InterruptProcessingException as iex:
+        logging.info("Processing interrupted")
+
+        # skip formatting inputs/outputs
+        error_details = {
+            "node_id": unique_id,
+        }
+
+        return (False, error_details, iex)
+    except Exception as ex:
+        typ, _, tb = sys.exc_info()
+        exception_type = full_type_name(typ)
+        input_data_formatted = {}
+        if input_data_all is not None:
+            input_data_formatted = {}
+            for name, inputs in input_data_all.items():
+                input_data_formatted[name] = [format_value(x) for x in inputs]
+
+        output_data_formatted = {}
+        for node_id, node_outputs in outputs.items():
+            output_data_formatted[node_id] = [[format_value(x) for x in l] for l in node_outputs]
+
+        logging.error(f"!!! Exception during processing!!! {ex}")
+        logging.error(traceback.format_exc())
+
+        error_details = {
+            "node_id": unique_id,
+            "exception_message": str(ex),
+            "exception_type": exception_type,
+            "traceback": traceback.format_tb(tb),
+            "current_inputs": input_data_formatted,
+            "current_outputs": output_data_formatted
+        }
+        return (False, error_details, ex)
+
+    executed.add(unique_id)
+
+    return (True, None, None)
+
+def recursive_will_execute(prompt, outputs, current_item, memo={}):
+    unique_id = current_item
+
+    if unique_id in memo:
+        return memo[unique_id]
+
+    inputs = prompt[unique_id]['inputs']
+    will_execute = []
+    if unique_id in outputs:
+        return []
+
+    for x in inputs:
+        input_data = inputs[x]
+        if isinstance(input_data, list):
+            input_unique_id = input_data[0]
+            output_index = input_data[1]
+            if input_unique_id not in outputs:
+                will_execute += recursive_will_execute(prompt, outputs, input_unique_id, memo)
+
+    memo[unique_id] = will_execute + [unique_id]
+    return memo[unique_id]
+
+def recursive_output_delete_if_changed(prompt, old_prompt, outputs, current_item):
+    unique_id = current_item
+    inputs = prompt[unique_id]['inputs']
+    class_type = prompt[unique_id]['class_type']
+    class_def = nodes.NODE_CLASS_MAPPINGS[class_type]
+
+    is_changed_old = ''
+    is_changed = ''
+    to_delete = False
+    if hasattr(class_def, 'IS_CHANGED'):
+        if unique_id in old_prompt and 'is_changed' in old_prompt[unique_id]:
+            is_changed_old = old_prompt[unique_id]['is_changed']
+        if 'is_changed' not in prompt[unique_id]:
+            input_data_all = get_input_data(inputs, class_def, unique_id, outputs)
+            if input_data_all is not None:
+                try:
+                    #is_changed = class_def.IS_CHANGED(**input_data_all)
+                    is_changed = map_node_over_list(class_def, input_data_all, "IS_CHANGED")
+                    prompt[unique_id]['is_changed'] = is_changed
+                except:
+                    to_delete = True
+        else:
+            is_changed = prompt[unique_id]['is_changed']
+
+    if unique_id not in outputs:
+        return True
+
+    if not to_delete:
+        if is_changed != is_changed_old:
+            to_delete = True
+        elif unique_id not in old_prompt:
+            to_delete = True
+        elif inputs == old_prompt[unique_id]['inputs']:
+            for x in inputs:
+                input_data = inputs[x]
+
+                if isinstance(input_data, list):
+                    input_unique_id = input_data[0]
+                    output_index = input_data[1]
+                    if input_unique_id in outputs:
+                        to_delete = recursive_output_delete_if_changed(prompt, old_prompt, outputs, input_unique_id)
+                    else:
+                        to_delete = True
+                    if to_delete:
+                        break
+        else:
+            to_delete = True
+
+    if to_delete:
+        d = outputs.pop(unique_id)
+        del d
+    return to_delete
+
+class PromptExecutor:
+    def __init__(self, server):
+        self.server = server
+        self.reset()
+
+    def reset(self):
+        self.outputs = {}
+        self.object_storage = {}
+        self.outputs_ui = {}
+        self.status_messages = []
+        self.success = True
+        self.old_prompt = {}
+
+    def add_message(self, event, data, broadcast: bool):
+        self.status_messages.append((event, data))
+        if self.server.client_id is not None or broadcast:
+            self.server.send_sync(event, data, self.server.client_id)
+
+    def handle_execution_error(self, prompt_id, prompt, current_outputs, executed, error, ex):
+        node_id = error["node_id"]
+        class_type = prompt[node_id]["class_type"]
+
+        # First, send back the status to the frontend depending
+        # on the exception type
+        if isinstance(ex, comfy.model_management.InterruptProcessingException):
+            mes = {
+                "prompt_id": prompt_id,
+                "node_id": node_id,
+                "node_type": class_type,
+                "executed": list(executed),
+            }
+            self.add_message("execution_interrupted", mes, broadcast=True)
+        else:
+            mes = {
+                "prompt_id": prompt_id,
+                "node_id": node_id,
+                "node_type": class_type,
+                "executed": list(executed),
+
+                "exception_message": error["exception_message"],
+                "exception_type": error["exception_type"],
+                "traceback": error["traceback"],
+                "current_inputs": error["current_inputs"],
+                "current_outputs": error["current_outputs"],
+            }
+            self.add_message("execution_error", mes, broadcast=False)
+        
+        # Next, remove the subsequent outputs since they will not be executed
+        to_delete = []
+        for o in self.outputs:
+            if (o not in current_outputs) and (o not in executed):
+                to_delete += [o]
+                if o in self.old_prompt:
+                    d = self.old_prompt.pop(o)
+                    del d
+        for o in to_delete:
+            d = self.outputs.pop(o)
+            del d
+
+    def execute(self, prompt, prompt_id, extra_data={}, execute_outputs=[]):
+        nodes.interrupt_processing(False)
+
+        if "client_id" in extra_data:
+            self.server.client_id = extra_data["client_id"]
+        else:
+            self.server.client_id = None
+
+        self.status_messages = []
+        self.add_message("execution_start", { "prompt_id": prompt_id}, broadcast=False)
+
+        with torch.inference_mode():
+            #delete cached outputs if nodes don't exist for them
+            to_delete = []
+            for o in self.outputs:
+                if o not in prompt:
+                    to_delete += [o]
+            for o in to_delete:
+                d = self.outputs.pop(o)
+                del d
+            to_delete = []
+            for o in self.object_storage:
+                if o[0] not in prompt:
+                    to_delete += [o]
+                else:
+                    p = prompt[o[0]]
+                    if o[1] != p['class_type']:
+                        to_delete += [o]
+            for o in to_delete:
+                d = self.object_storage.pop(o)
+                del d
+
+            for x in prompt:
+                recursive_output_delete_if_changed(prompt, self.old_prompt, self.outputs, x)
+
+            current_outputs = set(self.outputs.keys())
+            for x in list(self.outputs_ui.keys()):
+                if x not in current_outputs:
+                    d = self.outputs_ui.pop(x)
+                    del d
+
+            comfy.model_management.cleanup_models(keep_clone_weights_loaded=True)
+            self.add_message("execution_cached",
+                          { "nodes": list(current_outputs) , "prompt_id": prompt_id},
+                          broadcast=False)
+            executed = set()
+            output_node_id = None
+            to_execute = []
+
+            for node_id in list(execute_outputs):
+                to_execute += [(0, node_id)]
+
+            while len(to_execute) > 0:
+                #always execute the output that depends on the least amount of unexecuted nodes first
+                memo = {}
+                to_execute = sorted(list(map(lambda a: (len(recursive_will_execute(prompt, self.outputs, a[-1], memo)), a[-1]), to_execute)))
+                output_node_id = to_execute.pop(0)[-1]
+
+                # This call shouldn't raise anything if there's an error deep in
+                # the actual SD code, instead it will report the node where the
+                # error was raised
+                self.success, error, ex = recursive_execute(self.server, prompt, self.outputs, output_node_id, extra_data, executed, prompt_id, self.outputs_ui, self.object_storage)
+                if self.success is not True:
+                    self.handle_execution_error(prompt_id, prompt, current_outputs, executed, error, ex)
+                    break
+
+            for x in executed:
+                self.old_prompt[x] = copy.deepcopy(prompt[x])
+            self.server.last_node_id = None
+            if comfy.model_management.DISABLE_SMART_MEMORY:
+                comfy.model_management.unload_all_models()
+
+
+
+def validate_inputs(prompt, item, validated):
+    unique_id = item
+    if unique_id in validated:
+        return validated[unique_id]
+
+    inputs = prompt[unique_id]['inputs']
+    class_type = prompt[unique_id]['class_type']
+    obj_class = nodes.NODE_CLASS_MAPPINGS[class_type]
+
+    class_inputs = obj_class.INPUT_TYPES()
+    required_inputs = class_inputs['required']
+
+    errors = []
+    valid = True
+
+    validate_function_inputs = []
+    if hasattr(obj_class, "VALIDATE_INPUTS"):
+        validate_function_inputs = inspect.getfullargspec(obj_class.VALIDATE_INPUTS).args
+
+    for x in required_inputs:
+        if x not in inputs:
+            error = {
+                "type": "required_input_missing",
+                "message": "Required input is missing",
+                "details": f"{x}",
+                "extra_info": {
+                    "input_name": x
+                }
+            }
+            errors.append(error)
+            continue
+
+        val = inputs[x]
+        info = required_inputs[x]
+        type_input = info[0]
+        if isinstance(val, list):
+            if len(val) != 2:
+                error = {
+                    "type": "bad_linked_input",
+                    "message": "Bad linked input, must be a length-2 list of [node_id, slot_index]",
+                    "details": f"{x}",
+                    "extra_info": {
+                        "input_name": x,
+                        "input_config": info,
+                        "received_value": val
+                    }
+                }
+                errors.append(error)
+                continue
+
+            o_id = val[0]
+            o_class_type = prompt[o_id]['class_type']
+            r = nodes.NODE_CLASS_MAPPINGS[o_class_type].RETURN_TYPES
+            if r[val[1]] != type_input:
+                received_type = r[val[1]]
+                details = f"{x}, {received_type} != {type_input}"
+                error = {
+                    "type": "return_type_mismatch",
+                    "message": "Return type mismatch between linked nodes",
+                    "details": details,
+                    "extra_info": {
+                        "input_name": x,
+                        "input_config": info,
+                        "received_type": received_type,
+                        "linked_node": val
+                    }
+                }
+                errors.append(error)
+                continue
+            try:
+                r = validate_inputs(prompt, o_id, validated)
+                if r[0] is False:
+                    # `r` will be set in `validated[o_id]` already
+                    valid = False
+                    continue
+            except Exception as ex:
+                typ, _, tb = sys.exc_info()
+                valid = False
+                exception_type = full_type_name(typ)
+                reasons = [{
+                    "type": "exception_during_inner_validation",
+                    "message": "Exception when validating inner node",
+                    "details": str(ex),
+                    "extra_info": {
+                        "input_name": x,
+                        "input_config": info,
+                        "exception_message": str(ex),
+                        "exception_type": exception_type,
+                        "traceback": traceback.format_tb(tb),
+                        "linked_node": val
+                    }
+                }]
+                validated[o_id] = (False, reasons, o_id)
+                continue
+        else:
+            try:
+                if type_input == "INT":
+                    val = int(val)
+                    inputs[x] = val
+                if type_input == "FLOAT":
+                    val = float(val)
+                    inputs[x] = val
+                if type_input == "STRING":
+                    val = str(val)
+                    inputs[x] = val
+            except Exception as ex:
+                error = {
+                    "type": "invalid_input_type",
+                    "message": f"Failed to convert an input value to a {type_input} value",
+                    "details": f"{x}, {val}, {ex}",
+                    "extra_info": {
+                        "input_name": x,
+                        "input_config": info,
+                        "received_value": val,
+                        "exception_message": str(ex)
+                    }
+                }
+                errors.append(error)
+                continue
+
+            if len(info) > 1:
+                if "min" in info[1] and val < info[1]["min"]:
+                    error = {
+                        "type": "value_smaller_than_min",
+                        "message": "Value {} smaller than min of {}".format(val, info[1]["min"]),
+                        "details": f"{x}",
+                        "extra_info": {
+                            "input_name": x,
+                            "input_config": info,
+                            "received_value": val,
+                        }
+                    }
+                    errors.append(error)
+                    continue
+                if "max" in info[1] and val > info[1]["max"]:
+                    error = {
+                        "type": "value_bigger_than_max",
+                        "message": "Value {} bigger than max of {}".format(val, info[1]["max"]),
+                        "details": f"{x}",
+                        "extra_info": {
+                            "input_name": x,
+                            "input_config": info,
+                            "received_value": val,
+                        }
+                    }
+                    errors.append(error)
+                    continue
+
+            if x not in validate_function_inputs:
+                if isinstance(type_input, list):
+                    if val not in type_input:
+                        input_config = info
+                        list_info = ""
+
+                        # Don't send back gigantic lists like if they're lots of
+                        # scanned model filepaths
+                        if len(type_input) > 20:
+                            list_info = f"(list of length {len(type_input)})"
+                            input_config = None
+                        else:
+                            list_info = str(type_input)
+
+                        error = {
+                            "type": "value_not_in_list",
+                            "message": "Value not in list",
+                            "details": f"{x}: '{val}' not in {list_info}",
+                            "extra_info": {
+                                "input_name": x,
+                                "input_config": input_config,
+                                "received_value": val,
+                            }
+                        }
+                        errors.append(error)
+                        continue
+
+    if len(validate_function_inputs) > 0:
+        input_data_all = get_input_data(inputs, obj_class, unique_id)
+        input_filtered = {}
+        for x in input_data_all:
+            if x in validate_function_inputs:
+                input_filtered[x] = input_data_all[x]
+
+        #ret = obj_class.VALIDATE_INPUTS(**input_filtered)
+        ret = map_node_over_list(obj_class, input_filtered, "VALIDATE_INPUTS")
+        for x in input_filtered:
+            for i, r in enumerate(ret):
+                if r is not True:
+                    details = f"{x}"
+                    if r is not False:
+                        details += f" - {str(r)}"
+
+                    error = {
+                        "type": "custom_validation_failed",
+                        "message": "Custom validation failed for node",
+                        "details": details,
+                        "extra_info": {
+                            "input_name": x,
+                            "input_config": info,
+                            "received_value": val,
+                        }
+                    }
+                    errors.append(error)
+                    continue
+
+    if len(errors) > 0 or valid is not True:
+        ret = (False, errors, unique_id)
+    else:
+        ret = (True, [], unique_id)
+
+    validated[unique_id] = ret
+    return ret
+
+def full_type_name(klass):
+    module = klass.__module__
+    if module == 'builtins':
+        return klass.__qualname__
+    return module + '.' + klass.__qualname__
+
+def validate_prompt(prompt):
+    outputs = set()
+    for x in prompt:
+        if 'class_type' not in prompt[x]:
+            error = {
+                "type": "invalid_prompt",
+                "message": f"Cannot execute because a node is missing the class_type property.",
+                "details": f"Node ID '#{x}'",
+                "extra_info": {}
+            }
+            return (False, error, [], [])
+
+        class_type = prompt[x]['class_type']
+        class_ = nodes.NODE_CLASS_MAPPINGS.get(class_type, None)
+        if class_ is None:
+            error = {
+                "type": "invalid_prompt",
+                "message": f"Cannot execute because node {class_type} does not exist.",
+                "details": f"Node ID '#{x}'",
+                "extra_info": {}
+            }
+            return (False, error, [], [])
+
+        if hasattr(class_, 'OUTPUT_NODE') and class_.OUTPUT_NODE is True:
+            outputs.add(x)
+
+    if len(outputs) == 0:
+        error = {
+            "type": "prompt_no_outputs",
+            "message": "Prompt has no outputs",
+            "details": "",
+            "extra_info": {}
+        }
+        return (False, error, [], [])
+
+    good_outputs = set()
+    errors = []
+    node_errors = {}
+    validated = {}
+    for o in outputs:
+        valid = False
+        reasons = []
+        try:
+            m = validate_inputs(prompt, o, validated)
+            valid = m[0]
+            reasons = m[1]
+        except Exception as ex:
+            typ, _, tb = sys.exc_info()
+            valid = False
+            exception_type = full_type_name(typ)
+            reasons = [{
+                "type": "exception_during_validation",
+                "message": "Exception when validating node",
+                "details": str(ex),
+                "extra_info": {
+                    "exception_type": exception_type,
+                    "traceback": traceback.format_tb(tb)
+                }
+            }]
+            validated[o] = (False, reasons, o)
+
+        if valid is True:
+            good_outputs.add(o)
+        else:
+            logging.error(f"Failed to validate prompt for output {o}:")
+            if len(reasons) > 0:
+                logging.error("* (prompt):")
+                for reason in reasons:
+                    logging.error(f"  - {reason['message']}: {reason['details']}")
+            errors += [(o, reasons)]
+            for node_id, result in validated.items():
+                valid = result[0]
+                reasons = result[1]
+                # If a node upstream has errors, the nodes downstream will also
+                # be reported as invalid, but there will be no errors attached.
+                # So don't return those nodes as having errors in the response.
+                if valid is not True and len(reasons) > 0:
+                    if node_id not in node_errors:
+                        class_type = prompt[node_id]['class_type']
+                        node_errors[node_id] = {
+                            "errors": reasons,
+                            "dependent_outputs": [],
+                            "class_type": class_type
+                        }
+                        logging.error(f"* {class_type} {node_id}:")
+                        for reason in reasons:
+                            logging.error(f"  - {reason['message']}: {reason['details']}")
+                    node_errors[node_id]["dependent_outputs"].append(o)
+            logging.error("Output will be ignored")
+
+    if len(good_outputs) == 0:
+        errors_list = []
+        for o, errors in errors:
+            for error in errors:
+                errors_list.append(f"{error['message']}: {error['details']}")
+        errors_list = "\n".join(errors_list)
+
+        error = {
+            "type": "prompt_outputs_failed_validation",
+            "message": "Prompt outputs failed validation",
+            "details": errors_list,
+            "extra_info": {}
+        }
+
+        return (False, error, list(good_outputs), node_errors)
+
+    return (True, None, list(good_outputs), node_errors)
+
+MAXIMUM_HISTORY_SIZE = 10000
+
+class PromptQueue:
+    def __init__(self, server):
+        self.server = server
+        self.mutex = threading.RLock()
+        self.not_empty = threading.Condition(self.mutex)
+        self.task_counter = 0
+        self.queue = []
+        self.currently_running = {}
+        self.history = {}
+        self.flags = {}
+        server.prompt_queue = self
+
+    def put(self, item):
+        with self.mutex:
+            heapq.heappush(self.queue, item)
+            self.server.queue_updated()
+            self.not_empty.notify()
+
+    def get(self, timeout=None):
+        with self.not_empty:
+            while len(self.queue) == 0:
+                self.not_empty.wait(timeout=timeout)
+                if timeout is not None and len(self.queue) == 0:
+                    return None
+            item = heapq.heappop(self.queue)
+            i = self.task_counter
+            self.currently_running[i] = copy.deepcopy(item)
+            self.task_counter += 1
+            self.server.queue_updated()
+            return (item, i)
+
+    class ExecutionStatus(NamedTuple):
+        status_str: Literal['success', 'error']
+        completed: bool
+        messages: List[str]
+
+    def task_done(self, item_id, outputs,
+                  status: Optional['PromptQueue.ExecutionStatus']):
+        with self.mutex:
+            prompt = self.currently_running.pop(item_id)
+            if len(self.history) > MAXIMUM_HISTORY_SIZE:
+                self.history.pop(next(iter(self.history)))
+
+            status_dict: Optional[dict] = None
+            if status is not None:
+                status_dict = copy.deepcopy(status._asdict())
+
+            self.history[prompt[1]] = {
+                "prompt": prompt,
+                "outputs": copy.deepcopy(outputs),
+                'status': status_dict,
+            }
+            self.server.queue_updated()
+
+    def get_current_queue(self):
+        with self.mutex:
+            out = []
+            for x in self.currently_running.values():
+                out += [x]
+            return (out, copy.deepcopy(self.queue))
+
+    def get_tasks_remaining(self):
+        with self.mutex:
+            return len(self.queue) + len(self.currently_running)
+
+    def wipe_queue(self):
+        with self.mutex:
+            self.queue = []
+            self.server.queue_updated()
+
+    def delete_queue_item(self, function):
+        with self.mutex:
+            for x in range(len(self.queue)):
+                if function(self.queue[x]):
+                    if len(self.queue) == 1:
+                        self.wipe_queue()
+                    else:
+                        self.queue.pop(x)
+                        heapq.heapify(self.queue)
+                    self.server.queue_updated()
+                    return True
+        return False
+
+    def get_history(self, prompt_id=None, max_items=None, offset=-1):
+        with self.mutex:
+            if prompt_id is None:
+                out = {}
+                i = 0
+                if offset < 0 and max_items is not None:
+                    offset = len(self.history) - max_items
+                for k in self.history:
+                    if i >= offset:
+                        out[k] = self.history[k]
+                        if max_items is not None and len(out) >= max_items:
+                            break
+                    i += 1
+                return out
+            elif prompt_id in self.history:
+                return {prompt_id: copy.deepcopy(self.history[prompt_id])}
+            else:
+                return {}
+
+    def wipe_history(self):
+        with self.mutex:
+            self.history = {}
+
+    def delete_history_item(self, id_to_delete):
+        with self.mutex:
+            self.history.pop(id_to_delete, None)
+
+    def set_flag(self, name, data):
+        with self.mutex:
+            self.flags[name] = data
+            self.not_empty.notify()
+
+    def get_flags(self, reset=True):
+        with self.mutex:
+            if reset:
+                ret = self.flags
+                self.flags = {}
+                return ret
+            else:
+                return self.flags.copy()

backup/server.py

@@ -0,0 +1,655 @@
+import os
+import sys
+import asyncio
+import traceback
+
+import nodes
+import folder_paths
+import execution
+import uuid
+import urllib
+import json
+import glob
+import struct
+import ssl
+from PIL import Image, ImageOps
+from PIL.PngImagePlugin import PngInfo
+from io import BytesIO
+
+import aiohttp
+from aiohttp import web
+import logging
+
+import mimetypes
+from comfy.cli_args import args
+import comfy.utils
+import comfy.model_management
+
+from app.user_manager import UserManager
+
+class BinaryEventTypes:
+    PREVIEW_IMAGE = 1
+    UNENCODED_PREVIEW_IMAGE = 2
+
+async def send_socket_catch_exception(function, message):
+    try:
+        await function(message)
+    except (aiohttp.ClientError, aiohttp.ClientPayloadError, ConnectionResetError) as err:
+        logging.warning("send error: {}".format(err))
+
+@web.middleware
+async def cache_control(request: web.Request, handler):
+    response: web.Response = await handler(request)
+    if request.path.endswith('.js') or request.path.endswith('.css'):
+        response.headers.setdefault('Cache-Control', 'no-cache')
+    return response
+
+def create_cors_middleware(allowed_origin: str):
+    @web.middleware
+    async def cors_middleware(request: web.Request, handler):
+        if request.method == "OPTIONS":
+            # Pre-flight request. Reply successfully:
+            response = web.Response()
+        else:
+            response = await handler(request)
+
+        response.headers['Access-Control-Allow-Origin'] = allowed_origin
+        response.headers['Access-Control-Allow-Methods'] = 'POST, GET, DELETE, PUT, OPTIONS'
+        response.headers['Access-Control-Allow-Headers'] = 'Content-Type, Authorization'
+        response.headers['Access-Control-Allow-Credentials'] = 'true'
+        return response
+
+    return cors_middleware
+
+class PromptServer():
+    def __init__(self, loop):
+        PromptServer.instance = self
+
+        mimetypes.init()
+        mimetypes.types_map['.js'] = 'application/javascript; charset=utf-8'
+
+        self.user_manager = UserManager()
+        self.supports = ["custom_nodes_from_web"]
+        self.prompt_queue = None
+        self.loop = loop
+        self.messages = asyncio.Queue()
+        self.number = 0
+
+        middlewares = [cache_control]
+        if args.enable_cors_header:
+            middlewares.append(create_cors_middleware(args.enable_cors_header))
+
+        max_upload_size = round(args.max_upload_size * 1024 * 1024)
+        self.app = web.Application(client_max_size=max_upload_size, middlewares=middlewares)
+        self.sockets = dict()
+        self.web_root = os.path.join(os.path.dirname(
+            os.path.realpath(__file__)), "web")
+        routes = web.RouteTableDef()
+        self.routes = routes
+        self.last_node_id = None
+        self.client_id = None
+
+        self.on_prompt_handlers = []
+
+        @routes.get('/ws')
+        async def websocket_handler(request):
+            ws = web.WebSocketResponse()
+            await ws.prepare(request)
+            sid = request.rel_url.query.get('clientId', '')
+            if sid:
+                # Reusing existing session, remove old
+                self.sockets.pop(sid, None)
+            else:
+                sid = uuid.uuid4().hex
+
+            self.sockets[sid] = ws
+
+            try:
+                # Send initial state to the new client
+                await self.send("status", { "status": self.get_queue_info(), 'sid': sid }, sid)
+                # On reconnect if we are the currently executing client send the current node
+                if self.client_id == sid and self.last_node_id is not None:
+                    await self.send("executing", { "node": self.last_node_id }, sid)
+                    
+                async for msg in ws:
+                    if msg.type == aiohttp.WSMsgType.ERROR:
+                        logging.warning('ws connection closed with exception %s' % ws.exception())
+            finally:
+                self.sockets.pop(sid, None)
+            return ws
+
+        @routes.get("/")
+        async def get_root(request):
+            return web.FileResponse(os.path.join(self.web_root, "index.html"))
+
+        @routes.get("/embeddings")
+        def get_embeddings(self):
+            embeddings = folder_paths.get_filename_list("embeddings")
+            return web.json_response(list(map(lambda a: os.path.splitext(a)[0], embeddings)))
+
+        @routes.get("/extensions")
+        async def get_extensions(request):
+            files = glob.glob(os.path.join(
+                glob.escape(self.web_root), 'extensions/**/*.js'), recursive=True)
+            
+            extensions = list(map(lambda f: "/" + os.path.relpath(f, self.web_root).replace("\\", "/"), files))
+            
+            for name, dir in nodes.EXTENSION_WEB_DIRS.items():
+                files = glob.glob(os.path.join(glob.escape(dir), '**/*.js'), recursive=True)
+                extensions.extend(list(map(lambda f: "/extensions/" + urllib.parse.quote(
+                    name) + "/" + os.path.relpath(f, dir).replace("\\", "/"), files)))
+
+            return web.json_response(extensions)
+
+        def get_dir_by_type(dir_type):
+            if dir_type is None:
+                dir_type = "input"
+
+            if dir_type == "input":
+                type_dir = folder_paths.get_input_directory()
+            elif dir_type == "temp":
+                type_dir = folder_paths.get_temp_directory()
+            elif dir_type == "output":
+                type_dir = folder_paths.get_output_directory()
+
+            return type_dir, dir_type
+
+        def image_upload(post, image_save_function=None):
+            image = post.get("image")
+            overwrite = post.get("overwrite")
+
+            image_upload_type = post.get("type")
+            upload_dir, image_upload_type = get_dir_by_type(image_upload_type)
+
+            if image and image.file:
+                filename = image.filename
+                if not filename:
+                    return web.Response(status=400)
+
+                subfolder = post.get("subfolder", "")
+                full_output_folder = os.path.join(upload_dir, os.path.normpath(subfolder))
+                filepath = os.path.abspath(os.path.join(full_output_folder, filename))
+
+                if os.path.commonpath((upload_dir, filepath)) != upload_dir:
+                    return web.Response(status=400)
+
+                if not os.path.exists(full_output_folder):
+                    os.makedirs(full_output_folder)
+
+                split = os.path.splitext(filename)
+
+                if overwrite is not None and (overwrite == "true" or overwrite == "1"):
+                    pass
+                else:
+                    i = 1
+                    while os.path.exists(filepath):
+                        filename = f"{split[0]} ({i}){split[1]}"
+                        filepath = os.path.join(full_output_folder, filename)
+                        i += 1
+
+                if image_save_function is not None:
+                    image_save_function(image, post, filepath)
+                else:
+                    with open(filepath, "wb") as f:
+                        f.write(image.file.read())
+
+                return web.json_response({"name" : filename, "subfolder": subfolder, "type": image_upload_type})
+            else:
+                return web.Response(status=400)
+
+        @routes.post("/upload/image")
+        async def upload_image(request):
+            post = await request.post()
+            return image_upload(post)
+
+
+        @routes.post("/upload/mask")
+        async def upload_mask(request):
+            post = await request.post()
+
+            def image_save_function(image, post, filepath):
+                original_ref = json.loads(post.get("original_ref"))
+                filename, output_dir = folder_paths.annotated_filepath(original_ref['filename'])
+
+                # validation for security: prevent accessing arbitrary path
+                if filename[0] == '/' or '..' in filename:
+                    return web.Response(status=400)
+
+                if output_dir is None:
+                    type = original_ref.get("type", "output")
+                    output_dir = folder_paths.get_directory_by_type(type)
+
+                if output_dir is None:
+                    return web.Response(status=400)
+
+                if original_ref.get("subfolder", "") != "":
+                    full_output_dir = os.path.join(output_dir, original_ref["subfolder"])
+                    if os.path.commonpath((os.path.abspath(full_output_dir), output_dir)) != output_dir:
+                        return web.Response(status=403)
+                    output_dir = full_output_dir
+
+                file = os.path.join(output_dir, filename)
+
+                if os.path.isfile(file):
+                    with Image.open(file) as original_pil:
+                        metadata = PngInfo()
+                        if hasattr(original_pil,'text'):
+                            for key in original_pil.text:
+                                metadata.add_text(key, original_pil.text[key])
+                        original_pil = original_pil.convert('RGBA')
+                        mask_pil = Image.open(image.file).convert('RGBA')
+
+                        # alpha copy
+                        new_alpha = mask_pil.getchannel('A')
+                        original_pil.putalpha(new_alpha)
+                        original_pil.save(filepath, compress_level=4, pnginfo=metadata)
+
+            return image_upload(post, image_save_function)
+
+        @routes.get("/view")
+        async def view_image(request):
+            if "filename" in request.rel_url.query:
+                filename = request.rel_url.query["filename"]
+                filename,output_dir = folder_paths.annotated_filepath(filename)
+
+                # validation for security: prevent accessing arbitrary path
+                if filename[0] == '/' or '..' in filename:
+                    return web.Response(status=400)
+
+                if output_dir is None:
+                    type = request.rel_url.query.get("type", "output")
+                    output_dir = folder_paths.get_directory_by_type(type)
+
+                if output_dir is None:
+                    return web.Response(status=400)
+
+                if "subfolder" in request.rel_url.query:
+                    full_output_dir = os.path.join(output_dir, request.rel_url.query["subfolder"])
+                    if os.path.commonpath((os.path.abspath(full_output_dir), output_dir)) != output_dir:
+                        return web.Response(status=403)
+                    output_dir = full_output_dir
+
+                filename = os.path.basename(filename)
+                file = os.path.join(output_dir, filename)
+
+                if os.path.isfile(file):
+                    if 'preview' in request.rel_url.query:
+                        with Image.open(file) as img:
+                            preview_info = request.rel_url.query['preview'].split(';')
+                            image_format = preview_info[0]
+                            if image_format not in ['webp', 'jpeg'] or 'a' in request.rel_url.query.get('channel', ''):
+                                image_format = 'webp'
+
+                            quality = 90
+                            if preview_info[-1].isdigit():
+                                quality = int(preview_info[-1])
+
+                            buffer = BytesIO()
+                            if image_format in ['jpeg'] or request.rel_url.query.get('channel', '') == 'rgb':
+                                img = img.convert("RGB")
+                            img.save(buffer, format=image_format, quality=quality)
+                            buffer.seek(0)
+
+                            return web.Response(body=buffer.read(), content_type=f'image/{image_format}',
+                                                headers={"Content-Disposition": f"filename=\"{filename}\""})
+
+                    if 'channel' not in request.rel_url.query:
+                        channel = 'rgba'
+                    else:
+                        channel = request.rel_url.query["channel"]
+
+                    if channel == 'rgb':
+                        with Image.open(file) as img:
+                            if img.mode == "RGBA":
+                                r, g, b, a = img.split()
+                                new_img = Image.merge('RGB', (r, g, b))
+                            else:
+                                new_img = img.convert("RGB")
+
+                            buffer = BytesIO()
+                            new_img.save(buffer, format='PNG')
+                            buffer.seek(0)
+
+                            return web.Response(body=buffer.read(), content_type='image/png',
+                                                headers={"Content-Disposition": f"filename=\"{filename}\""})
+
+                    elif channel == 'a':
+                        with Image.open(file) as img:
+                            if img.mode == "RGBA":
+                                _, _, _, a = img.split()
+                            else:
+                                a = Image.new('L', img.size, 255)
+
+                            # alpha img
+                            alpha_img = Image.new('RGBA', img.size)
+                            alpha_img.putalpha(a)
+                            alpha_buffer = BytesIO()
+                            alpha_img.save(alpha_buffer, format='PNG')
+                            alpha_buffer.seek(0)
+
+                            return web.Response(body=alpha_buffer.read(), content_type='image/png',
+                                                headers={"Content-Disposition": f"filename=\"{filename}\""})
+                    else:
+                        return web.FileResponse(file, headers={"Content-Disposition": f"filename=\"{filename}\""})
+
+            return web.Response(status=404)
+
+        @routes.get("/view_metadata/{folder_name}")
+        async def view_metadata(request):
+            folder_name = request.match_info.get("folder_name", None)
+            if folder_name is None:
+                return web.Response(status=404)
+            if not "filename" in request.rel_url.query:
+                return web.Response(status=404)
+
+            filename = request.rel_url.query["filename"]
+            if not filename.endswith(".safetensors"):
+                return web.Response(status=404)
+
+            safetensors_path = folder_paths.get_full_path(folder_name, filename)
+            if safetensors_path is None:
+                return web.Response(status=404)
+            out = comfy.utils.safetensors_header(safetensors_path, max_size=1024*1024)
+            if out is None:
+                return web.Response(status=404)
+            dt = json.loads(out)
+            if not "__metadata__" in dt:
+                return web.Response(status=404)
+            return web.json_response(dt["__metadata__"])
+
+        @routes.get("/system_stats")
+        async def get_queue(request):
+            device = comfy.model_management.get_torch_device()
+            device_name = comfy.model_management.get_torch_device_name(device)
+            vram_total, torch_vram_total = comfy.model_management.get_total_memory(device, torch_total_too=True)
+            vram_free, torch_vram_free = comfy.model_management.get_free_memory(device, torch_free_too=True)
+            system_stats = {
+                "system": {
+                    "os": os.name,
+                    "python_version": sys.version,
+                    "embedded_python": os.path.split(os.path.split(sys.executable)[0])[1] == "python_embeded"
+                },
+                "devices": [
+                    {
+                        "name": device_name,
+                        "type": device.type,
+                        "index": device.index,
+                        "vram_total": vram_total,
+                        "vram_free": vram_free,
+                        "torch_vram_total": torch_vram_total,
+                        "torch_vram_free": torch_vram_free,
+                    }
+                ]
+            }
+            return web.json_response(system_stats)
+
+        @routes.get("/prompt")
+        async def get_prompt(request):
+            return web.json_response(self.get_queue_info())
+
+        def node_info(node_class):
+            obj_class = nodes.NODE_CLASS_MAPPINGS[node_class]
+            info = {}
+            info['input'] = obj_class.INPUT_TYPES()
+            info['output'] = obj_class.RETURN_TYPES
+            info['output_is_list'] = obj_class.OUTPUT_IS_LIST if hasattr(obj_class, 'OUTPUT_IS_LIST') else [False] * len(obj_class.RETURN_TYPES)
+            info['output_name'] = obj_class.RETURN_NAMES if hasattr(obj_class, 'RETURN_NAMES') else info['output']
+            info['name'] = node_class
+            info['display_name'] = nodes.NODE_DISPLAY_NAME_MAPPINGS[node_class] if node_class in nodes.NODE_DISPLAY_NAME_MAPPINGS.keys() else node_class
+            info['description'] = obj_class.DESCRIPTION if hasattr(obj_class,'DESCRIPTION') else ''
+            info['category'] = 'sd'
+            if hasattr(obj_class, 'OUTPUT_NODE') and obj_class.OUTPUT_NODE == True:
+                info['output_node'] = True
+            else:
+                info['output_node'] = False
+
+            if hasattr(obj_class, 'CATEGORY'):
+                info['category'] = obj_class.CATEGORY
+            return info
+
+        @routes.get("/object_info")
+        async def get_object_info(request):
+            out = {}
+            for x in nodes.NODE_CLASS_MAPPINGS:
+                try:
+                    out[x] = node_info(x)
+                except Exception as e:
+                    logging.error(f"[ERROR] An error occurred while retrieving information for the '{x}' node.")
+                    logging.error(traceback.format_exc())
+            return web.json_response(out)
+
+        @routes.get("/object_info/{node_class}")
+        async def get_object_info_node(request):
+            node_class = request.match_info.get("node_class", None)
+            out = {}
+            if (node_class is not None) and (node_class in nodes.NODE_CLASS_MAPPINGS):
+                out[node_class] = node_info(node_class)
+            return web.json_response(out)
+
+        @routes.get("/history")
+        async def get_history(request):
+            max_items = request.rel_url.query.get("max_items", None)
+            if max_items is not None:
+                max_items = int(max_items)
+            return web.json_response(self.prompt_queue.get_history(max_items=max_items))
+
+        @routes.get("/history/{prompt_id}")
+        async def get_history(request):
+            prompt_id = request.match_info.get("prompt_id", None)
+            return web.json_response(self.prompt_queue.get_history(prompt_id=prompt_id))
+
+        @routes.get("/queue")
+        async def get_queue(request):
+            queue_info = {}
+            current_queue = self.prompt_queue.get_current_queue()
+            queue_info['queue_running'] = current_queue[0]
+            queue_info['queue_pending'] = current_queue[1]
+            return web.json_response(queue_info)
+
+        @routes.post("/prompt")
+        async def post_prompt(request):
+            logging.info("got prompt")
+            resp_code = 200
+            out_string = ""
+            json_data =  await request.json()
+            json_data = self.trigger_on_prompt(json_data)
+
+            if "number" in json_data:
+                number = float(json_data['number'])
+            else:
+                number = self.number
+                if "front" in json_data:
+                    if json_data['front']:
+                        number = -number
+
+                self.number += 1
+
+            if "prompt" in json_data:
+                prompt = json_data["prompt"]
+                valid = execution.validate_prompt(prompt)
+                extra_data = {}
+                if "extra_data" in json_data:
+                    extra_data = json_data["extra_data"]
+
+                if "client_id" in json_data:
+                    extra_data["client_id"] = json_data["client_id"]
+                if valid[0]:
+                    prompt_id = str(uuid.uuid4())
+                    outputs_to_execute = valid[2]
+                    self.prompt_queue.put((number, prompt_id, prompt, extra_data, outputs_to_execute))
+                    response = {"prompt_id": prompt_id, "number": number, "node_errors": valid[3]}
+                    return web.json_response(response)
+                else:
+                    logging.warning("invalid prompt: {}".format(valid[1]))
+                    return web.json_response({"error": valid[1], "node_errors": valid[3]}, status=400)
+            else:
+                return web.json_response({"error": "no prompt", "node_errors": []}, status=400)
+
+        @routes.post("/queue")
+        async def post_queue(request):
+            json_data =  await request.json()
+            if "clear" in json_data:
+                if json_data["clear"]:
+                    self.prompt_queue.wipe_queue()
+            if "delete" in json_data:
+                to_delete = json_data['delete']
+                for id_to_delete in to_delete:
+                    delete_func = lambda a: a[1] == id_to_delete
+                    self.prompt_queue.delete_queue_item(delete_func)
+
+            return web.Response(status=200)
+
+        @routes.post("/interrupt")
+        async def post_interrupt(request):
+            nodes.interrupt_processing()
+            return web.Response(status=200)
+
+        @routes.post("/free")
+        async def post_free(request):
+            json_data = await request.json()
+            unload_models = json_data.get("unload_models", False)
+            free_memory = json_data.get("free_memory", False)
+            if unload_models:
+                self.prompt_queue.set_flag("unload_models", unload_models)
+            if free_memory:
+                self.prompt_queue.set_flag("free_memory", free_memory)
+            return web.Response(status=200)
+
+        @routes.post("/history")
+        async def post_history(request):
+            json_data =  await request.json()
+            if "clear" in json_data:
+                if json_data["clear"]:
+                    self.prompt_queue.wipe_history()
+            if "delete" in json_data:
+                to_delete = json_data['delete']
+                for id_to_delete in to_delete:
+                    self.prompt_queue.delete_history_item(id_to_delete)
+
+            return web.Response(status=200)
+        
+    def add_routes(self):
+        self.user_manager.add_routes(self.routes)
+        self.app.add_routes(self.routes)
+
+        for name, dir in nodes.EXTENSION_WEB_DIRS.items():
+            self.app.add_routes([
+                web.static('/extensions/' + urllib.parse.quote(name), dir),
+            ])
+
+        self.app.add_routes([
+            web.static('/', self.web_root),
+        ])
+
+    def get_queue_info(self):
+        prompt_info = {}
+        exec_info = {}
+        exec_info['queue_remaining'] = self.prompt_queue.get_tasks_remaining()
+        prompt_info['exec_info'] = exec_info
+        return prompt_info
+
+    async def send(self, event, data, sid=None):
+        if event == BinaryEventTypes.UNENCODED_PREVIEW_IMAGE:
+            await self.send_image(data, sid=sid)
+        elif isinstance(data, (bytes, bytearray)):
+            await self.send_bytes(event, data, sid)
+        else:
+            await self.send_json(event, data, sid)
+
+    def encode_bytes(self, event, data):
+        if not isinstance(event, int):
+            raise RuntimeError(f"Binary event types must be integers, got {event}")
+
+        packed = struct.pack(">I", event)
+        message = bytearray(packed)
+        message.extend(data)
+        return message
+
+    async def send_image(self, image_data, sid=None):
+        image_type = image_data[0]
+        image = image_data[1]
+        max_size = image_data[2]
+        if max_size is not None:
+            if hasattr(Image, 'Resampling'):
+                resampling = Image.Resampling.BILINEAR
+            else:
+                resampling = Image.ANTIALIAS
+
+            image = ImageOps.contain(image, (max_size, max_size), resampling)
+        type_num = 1
+        if image_type == "JPEG":
+            type_num = 1
+        elif image_type == "PNG":
+            type_num = 2
+
+        bytesIO = BytesIO()
+        header = struct.pack(">I", type_num)
+        bytesIO.write(header)
+        image.save(bytesIO, format=image_type, quality=95, compress_level=1)
+        preview_bytes = bytesIO.getvalue()
+        await self.send_bytes(BinaryEventTypes.PREVIEW_IMAGE, preview_bytes, sid=sid)
+
+    async def send_bytes(self, event, data, sid=None):
+        message = self.encode_bytes(event, data)
+
+        if sid is None:
+            sockets = list(self.sockets.values())
+            for ws in sockets:
+                await send_socket_catch_exception(ws.send_bytes, message)
+        elif sid in self.sockets:
+            await send_socket_catch_exception(self.sockets[sid].send_bytes, message)
+
+    async def send_json(self, event, data, sid=None):
+        message = {"type": event, "data": data}
+
+        if sid is None:
+            sockets = list(self.sockets.values())
+            for ws in sockets:
+                await send_socket_catch_exception(ws.send_json, message)
+        elif sid in self.sockets:
+            await send_socket_catch_exception(self.sockets[sid].send_json, message)
+
+    def send_sync(self, event, data, sid=None):
+        self.loop.call_soon_threadsafe(
+            self.messages.put_nowait, (event, data, sid))
+
+    def queue_updated(self):
+        self.send_sync("status", { "status": self.get_queue_info() })
+
+    async def publish_loop(self):
+        while True:
+            msg = await self.messages.get()
+            await self.send(*msg)
+
+    async def start(self, address, port, verbose=True, call_on_start=None):
+        runner = web.AppRunner(self.app, access_log=None)
+        await runner.setup()
+        ssl_ctx = None
+        scheme = "http"
+        if args.tls_keyfile and args.tls_certfile:
+                ssl_ctx = ssl.SSLContext(protocol=ssl.PROTOCOL_TLS_SERVER, verify_mode=ssl.CERT_NONE)
+                ssl_ctx.load_cert_chain(certfile=args.tls_certfile,
+                                keyfile=args.tls_keyfile)
+                scheme = "https"
+
+        site = web.TCPSite(runner, address, port, ssl_context=ssl_ctx)
+        await site.start()
+
+        if verbose:
+            logging.info("Starting server\n")
+            logging.info("To see the GUI go to: {}://{}:{}".format(scheme, address, port))
+        if call_on_start is not None:
+            call_on_start(scheme, address, port)
+
+    def add_on_prompt_handler(self, handler):
+        self.on_prompt_handlers.append(handler)
+
+    def trigger_on_prompt(self, json_data):
+        for handler in self.on_prompt_handlers:
+            try:
+                json_data = handler(json_data)
+            except Exception as e:
+                logging.warning(f"[ERROR] An error occurred during the on_prompt_handler processing")
+                logging.warning(traceback.format_exc())
+
+        return json_data

changeRestore.py

@@ -0,0 +1,216 @@
+import os,sys
+#读取py文件为text字符串
+editlist=[{
+	"name": "server.py",
+	"isEditStr":"elif 'prompt_id' in json_data:",
+	"changs":[
+		{
+			"primit":'''else:
+                return web.json_response({"error": "no prompt", "node_errors": []}, status=400)''',
+			"edit":'''elif 'prompt_id' in json_data:
+                return web.json_response(json_data)
+            else:
+                return web.json_response({"error": "no prompt", "node_errors": []}, status=400)'''
+		}
+	]
+},{
+	"name": "execution.py",
+	"isEditStr":"def get_del_keys(key, prompt,uniqueIds)",
+	"changs":[
+		{
+			"primit":'''def recursive_execute''',
+			"edit":'''#循环添加代码----------开始----------
+def get_del_keys(key, prompt,uniqueIds):
+    keys=[]
+    for k,v in prompt.items():
+        if k in uniqueIds :
+            continue
+        for k1,v1 in v['inputs'].items():
+            if type(v1)==list and v1[0]==key:
+                keys.append(k)
+                keys=keys+get_del_keys(k,prompt,uniqueIds+[k])
+    return keys
+#循环添加代码-----------结束---------
+def recursive_execute'''
+		},{
+			"primit":'''class_def = nodes.NODE_CLASS_MAPPINGS[class_type]
+    if unique_id in outputs:
+        return (True, None, None)''',
+			"edit":'''class_def = nodes.NODE_CLASS_MAPPINGS[class_type]
+    if unique_id in outputs:
+        return (True, None, None)
+
+    #循环添加代码----------开始----------
+    startNum=None
+    startData={}
+    delKeys=[]
+    backhaul={}
+    clTypes=['ForInnerEnd','IfInnerExecute','DoWhileEnd']
+    oldPrompt=None
+    if class_type in clTypes:
+        oldPrompt=copy.deepcopy(prompt)
+
+    if class_type=='ForInnerEnd':
+        startNum=prompt[unique_id]['inputs']['total'][0]
+        inputNum=prompt[unique_id]['inputs']['obj'][0]
+        maxKeyStr=sorted(list(prompt.keys()), key=lambda x: int(x.split(':')[0]))[-1]
+        maxKey = int(maxKeyStr.split(':')[0])
+        delKeys=list(set(get_del_keys(startNum,prompt,[startNum,unique_id]))) 
+        delKeys.append(startNum)
+        delKeys = list(filter(lambda x: x != inputNum, delKeys))
+        for key in delKeys:
+            outputs.pop(key, None)
+        
+        startInput=prompt[startNum]['inputs']
+        if isinstance(startInput['total'],list) or isinstance(startInput['stop'],list) or isinstance(startInput['i'],list):
+            result = recursive_execute(server, prompt, outputs, startNum, extra_data, executed, prompt_id, outputs_ui, object_storage)
+            if result[0] is not True:
+                    return result
+            
+        for x in startInput:
+            input_data = startInput[x]
+            if isinstance(input_data, list):
+                startData[x]=outputs[input_data[0]][input_data[1]][0]
+            else:
+                startData[x]=input_data
+
+        for i in range(startData['i']+1,startData['total'],startData['stop']):
+            prompt[str(maxKey+i)]=prompt[inputNum]
+            prompt[unique_id]['inputs']['obj'+str(i)]=[str(maxKey+i),prompt[unique_id]['inputs']['obj'][-1]]
+            if i==startData['i']+1:
+                backhaul['obj'+str(i)]=prompt[unique_id]['inputs']['obj']
+            else:
+                backhaul['obj'+str(i)]=[str(maxKey+i-startData['stop']),prompt[unique_id]['inputs']['obj'][-1]]
+    elif class_type=='DoWhileEnd':
+        startNum=prompt[unique_id]['inputs']['start'][0]
+        inputNum=prompt[unique_id]['inputs']['ANY'][0]
+        delKeys=list(set(get_del_keys(startNum,prompt,[startNum,unique_id]))) 
+        delKeys.append(startNum)
+        delKeys.append(inputNum)
+
+    #循环添加代码-----------结束---------'''
+		}
+		,{
+			"primit":'''result = recursive_execute(server, prompt, outputs, input_unique_id, extra_data, executed, prompt_id, outputs_ui, object_storage)
+                if result[0] is not True:
+                    # Another node failed further upstream
+                    return result''',
+			"edit":'''#循环添加代码----------开始----------
+                if class_type=='ForInnerEnd' and x !='obj' and x.startswith('obj'):
+                    if startNum!=None:
+                        if isinstance(prompt[startNum]['inputs']['i'],list):
+                            prompt[startNum]['inputs']['i']=startData['i']+startData['stop']
+                        else:
+                            prompt[startNum]['inputs']['i']=prompt[startNum]['inputs']['i']+startData['stop']
+                        prompt[startNum]['inputs']['obj']=backhaul[x]
+                        for key in delKeys:
+                            outputs.pop(key, None)
+                #循环添加代码-----------结束---------
+                result = recursive_execute(server, prompt, outputs, input_unique_id, extra_data, executed, prompt_id, outputs_ui, object_storage)
+                if result[0] is not True:
+                    # Another node failed further upstream
+                    return result
+                
+                #判断选择添加代码-----------开始---------
+                if class_type=='DoWhileEnd' and x == 'ANY':
+                    any=outputs[inputNum][prompt[unique_id]['inputs']['ANY'][1]][0]
+                    i=0
+                    while any:
+                        for key in delKeys:
+                            outputs.pop(key, None)
+                        i=i+1
+                        outputs['i']=[[i]]
+                        prompt[startNum]['inputs']['i']=['i',0]
+                        result = recursive_execute(server, prompt, outputs, input_unique_id, extra_data, executed, prompt_id, outputs_ui, object_storage)
+                        if result[0] is not True:
+                            return result
+                        any=outputs[inputNum][prompt[unique_id]['inputs']['ANY'][1]][0]
+            
+            if class_type=='IfInnerExecute' and x=='ANY':
+                if outputs[input_unique_id][output_index][0]:
+                    inputs['IF_FALSE']=oldPrompt[unique_id]['inputs']['IF_TRUE']
+                else:
+                    inputs['IF_TRUE']=oldPrompt[unique_id]['inputs']['IF_FALSE']
+            #判断选择添加代码-----------结束---------'''
+		}
+		,{
+			"primit":'''output_data, output_ui = get_output_data(obj, input_data_all)
+        outputs[unique_id] = output_data''',
+			"edit":'''output_data, output_ui = get_output_data(obj, input_data_all)
+        outputs[unique_id] = output_data
+        #循环添加代码----------开始----------
+        if class_type=='ForInnerEnd' or class_type=='DoWhileEnd':
+            prompt[startNum]['inputs']=oldPrompt[startNum]['inputs']
+            prompt[unique_id]['inputs']=oldPrompt[unique_id]['inputs']
+            outputs.pop(startNum, None)
+            result = recursive_execute(server, prompt, outputs, startNum, extra_data, executed, prompt_id, outputs_ui, object_storage)
+            if result[0] is not True:
+                   return result
+        if class_type=='IfInnerExecute':
+            prompt[unique_id]['inputs']=oldPrompt[unique_id]['inputs']
+        #循环添加代码-----------结束---------'''
+		}
+		,{
+			"primit":'''def validate_prompt(prompt):
+    outputs = set()''',
+			"edit":'''def validate_prompt(prompt):
+    outputs = set()
+    inputKeys=[]
+    for k,v in prompt.items():
+        for k1,v1 in v['inputs'].items():
+            if type(v1)==list and len(v1)==2:
+                inputKeys.append(v1[0])'''
+		}
+		,{
+			"primit":'''if hasattr(class_, 'OUTPUT_NODE') and class_.OUTPUT_NODE is True:''',
+			"edit":'''if hasattr(class_, 'OUTPUT_NODE') and class_.OUTPUT_NODE is True and x not in inputKeys: #添加and x not in inputKeys'''
+		}
+	]
+}
+]
+filePath=os.path.dirname(os.path.abspath(__file__))
+print('当前目录：',filePath)
+base_path=filePath.split('custom_nodes')[0]
+backupPath=os.path.join(filePath,'backup')
+def read_py_file(file_path):
+	with open(file_path, 'r',encoding='utf-8') as file:
+		text = file.read()
+	return text
+
+def write_py_file(file_path, text):
+	with open(file_path, 'w',encoding='utf-8') as file:
+		file.write(text)
+
+
+if __name__ == '__main__':
+	isRestore='--res' in sys.argv
+	print('-------开始还原-------' if isRestore else '--------开始修改--------')
+	for i in editlist:
+		file_path=os.path.join(base_path,i['name'])
+		if not os.path.exists(file_path):
+			print(file_path+'文件不存在')
+			continue
+		if not os.path.exists(backupPath):
+			os.makedirs(backupPath)
+		
+		if isRestore:
+			if not os.path.exists(os.path.join(backupPath,i['name'])):
+				print('备份文件'+os.path.join(backupPath,i['name'])+'不存在无法还原')
+				continue
+			text=read_py_file(os.path.join(backupPath,i['name']))
+			write_py_file(file_path, text)
+		else:
+			text=read_py_file(file_path)
+			if text.find(i['isEditStr'])!=-1:
+				continue
+			write_py_file(os.path.join(backupPath,i['name']), text)
+			textOld=text+''
+			for j in i['changs']:
+				if text.find(j['primit'])==-1:
+					print(i['name']+'未找到修改位置，修改失败')
+					text=textOld
+					break
+				text=text.replace(j['primit'],j['edit'])
+			write_py_file(file_path, text)
+	print('-------还原完成-------' if isRestore else '--------修改完成--------')
+

config/settings_nodes.json

@@ -1 +0,0 @@
-{}

config/weChat.yaml

@@ -1,116 +0,0 @@
-wechat:
-  appid: wx1c2b19f7af1d6241                        #公众号appid 
-  secret: 915d587b5ae4e148e100429e11ca4d60         #公众号秘钥
-  serverAddress: http://139.9.53.208/ #服务器地址 微信公众配置地址的前缀
-  freeSize: 3 #免费使用次数
-  isEnterprise: false #是否为企业微信
-  adminNo: b0pOVFM2dnRsZkdLeWl2Zlk2bG9MTFNjUTNGUQ== #管理员微信编号启动会后可发送我的编号获取
-  adminWeChat: yanlang123456                        #管理员微信号
-  access_token: '' #不用填系统用于保存临时值
-  access_token_expires_at: 0 #默认0
-  commands:                                        #功能列表
-    图生视频:
-      filename: img2video.json                     #功能流程图json文件
-      params:                                      #参数列表
-        image:                                     #参数名称(几个固定参数名称必须如下：图片：image  随机种子：seed 正向提示词：prompt 反向提示词：negative) 
-          isRequired: true                         #是否必填
-          keys:                                    #参数对应json中对应的key
-          - '38'
-          - inputs
-          - image_path
-          zhName: 图片                             #参数中文名称
-        motion_bucket:
-          default: 127
-          isRequired: true
-          keys:
-          - '12'
-          - inputs
-          - motion_bucket_id
-          max: 200
-          min: 0
-          type: number
-          zhName: 运动量
-        seed:
-          default: -1
-          isRequired: true
-          keys:
-          - '3'
-          - inputs
-          - seed
-          type: text
-          zhName: 随机种子
-      replyText: 请发送您的图片发送“ok”开始任务，如：“ok”     #指令说明回复
-    文生图:
-      filename: text2img.json
-      params:
-        ckpt:
-          isRequired: false
-          keys:
-          - '4'
-          - inputs
-          - ckpt_name
-          options:          #选项参数
-            冷却混合料SD1.5: chilloutmix_NiPrunedFp32Fix.safetensors #中文名称:英文名称
-            卡通动画片SD1.5: yamer_Cartoon_xenoArcadiaCD.safetensors
-            无法控制的SDXL: juggernautXL_version6Rundiffusion.safetensors
-            梦想塑造者SD1.5: dreamshaper_8.safetensors
-          zhName: 基础模型
-        height:
-          default: 512     #默认值
-          isRequired: false
-          keys:
-          - '14'
-          - inputs
-          - height
-          max: 1024         #最大值
-          min: 512          #最小值
-          type: number      #类型(输入型参数目前只支持 number/text)
-          zhName: 高度
-        lora:
-          isRequired: false
-          keys:
-          - '10'
-          - inputs
-          - lora_name
-          options:
-            SD1.5-LCM加速: lcm-lora-sdv1-5.safetensors
-            SDXL-LCM加速: lcm-lora-sdxl.safetensors
-          zhName: LoRA
-        negative:
-          isRequired: false
-          keys:
-          - '7'
-          - inputs
-          - text
-          zhName: 反向提示词
-        prompt:
-          isRequired: true
-          keys:
-          - '20'
-          - inputs
-          - text_trans
-          zhName: 正向提示词
-        seed:
-          default: -1
-          isRequired: true
-          keys:
-          - '3'
-          - inputs
-          - seed
-          type: text
-          zhName: 随机种子
-        width:
-          default: 512
-          isRequired: false
-          keys:
-          - '14'
-          - inputs
-          - width
-          max: 1024
-          min: 512
-          type: number
-          zhName: 宽度
-      replyText: 请发送您的中文或英文提示词用“ok”结束，如：“一只小狗ok”
-  query_commands: #查询指令（固定）
-  - 查询排队情况
-  - 我的编号

config/weChat.yaml.template

@@ -1,3 +1,11 @@
+#redis:
+#  host: 127.0.0.1
+#  port: 6379
+#  password: 
+#  basePath: 192.168.1.8:8188 #comfyui地址
+#  isMain: true       #是否主comfyUI
+#  modelPriority: false #是否模型优先
+
 wechat:
   appid: wxa50ed345b253964d                        #公众号appid 
   secret: 327ac9ba06d214d2fcfe2e9ed0520684         #公众号秘钥

config/workflow/color2painting.json

@@ -1,7 +1,7 @@
 {
   "3": {
     "inputs": {
-      "seed": 611578276912800,
+      "seed": 2070780056835,
       "steps": 8,
       "cfg": 1,
       "sampler_name": "euler_ancestral",
@@ -86,19 +86,6 @@
       "title": "VAE解码"
     }
   },
-  "9": {
-    "inputs": {
-      "filename_prefix": "ComfyUI",
-      "images": [
-        "8",
-        0
-      ]
-    },
-    "class_type": "SaveImage",
-    "_meta": {
-      "title": "保存图像"
-    }
-  },
   "10": {
     "inputs": {
       "lora_name": "lcm-lora-sdv1-5.safetensors",
@@ -188,5 +175,364 @@
     "_meta": {
       "title": "反向提示词"
     }
+  },
+  "29": {
+    "inputs": {
+      "seed": 19,
+      "steps": 25,
+      "cfg": 6,
+      "sampler_name": "dpmpp_2m",
+      "scheduler": "karras",
+      "denoise": 0.9,
+      "model": [
+        "36",
+        0
+      ],
+      "positive": [
+        "48",
+        0
+      ],
+      "negative": [
+        "48",
+        1
+      ],
+      "latent_image": [
+        "48",
+        2
+      ]
+    },
+    "class_type": "KSampler",
+    "_meta": {
+      "title": "K采样器"
+    }
+  },
+  "31": {
+    "inputs": {
+      "text": [
+        "52",
+        0
+      ],
+      "clip": [
+        "4",
+        1
+      ]
+    },
+    "class_type": "CLIPTextEncode",
+    "_meta": {
+      "title": "CLIP文本编码器"
+    }
+  },
+  "32": {
+    "inputs": {
+      "text": "blurry, noisy, messy, glitch, distorted, malformed, ill, horror, naked, nipples",
+      "clip": [
+        "4",
+        1
+      ]
+    },
+    "class_type": "CLIPTextEncode",
+    "_meta": {
+      "title": "CLIP文本编码器"
+    }
+  },
+  "33": {
+    "inputs": {
+      "samples": [
+        "29",
+        0
+      ],
+      "vae": [
+        "4",
+        2
+      ]
+    },
+    "class_type": "VAEDecode",
+    "_meta": {
+      "title": "VAE解码"
+    }
+  },
+  "34": {
+    "inputs": {
+      "images": [
+        "56",
+        0
+      ]
+    },
+    "class_type": "PreviewImage",
+    "_meta": {
+      "title": "预览图像"
+    }
+  },
+  "36": {
+    "inputs": {
+      "weight_style": 1,
+      "weight_composition": 1,
+      "expand_style": false,
+      "combine_embeds": "concat",
+      "start_at": 0,
+      "end_at": 1,
+      "embeds_scaling": "V only",
+      "model": [
+        "37",
+        0
+      ],
+      "ipadapter": [
+        "37",
+        1
+      ],
+      "image_style": [
+        "39",
+        0
+      ],
+      "image_composition": [
+        "47",
+        0
+      ]
+    },
+    "class_type": "IPAdapterStyleComposition",
+    "_meta": {
+      "title": "IPAdapter风格合成SDXL"
+    }
+  },
+  "37": {
+    "inputs": {
+      "preset": "STANDARD (medium strength)",
+      "model": [
+        "4",
+        0
+      ]
+    },
+    "class_type": "IPAdapterUnifiedLoader",
+    "_meta": {
+      "title": "IPAdapter加载器"
+    }
+  },
+  "39": {
+    "inputs": {
+      "interpolation": "LANCZOS",
+      "crop_position": "pad",
+      "sharpening": 0,
+      "image": [
+        "8",
+        0
+      ]
+    },
+    "class_type": "PrepImageForClipVision",
+    "_meta": {
+      "title": "CLIP视觉图像处理"
+    }
+  },
+  "40": {
+    "inputs": {
+      "strength": 0.55,
+      "start_percent": 0,
+      "end_percent": 0.75,
+      "positive": [
+        "31",
+        0
+      ],
+      "negative": [
+        "32",
+        0
+      ],
+      "control_net": [
+        "41",
+        0
+      ],
+      "image": [
+        "49",
+        0
+      ]
+    },
+    "class_type": "ControlNetApplyAdvanced",
+    "_meta": {
+      "title": "ControlNet应用(高级)"
+    }
+  },
+  "41": {
+    "inputs": {
+      "control_net_name": "t2i-adapter-lineart-sdxl-1.0.safetensors"
+    },
+    "class_type": "ControlNetLoader",
+    "_meta": {
+      "title": "ControlNet加载器"
+    }
+  },
+  "43": {
+    "inputs": {
+      "strength": 1,
+      "start_percent": 0,
+      "end_percent": 0.3,
+      "positive": [
+        "40",
+        0
+      ],
+      "negative": [
+        "40",
+        1
+      ],
+      "control_net": [
+        "44",
+        0
+      ],
+      "image": [
+        "45",
+        0
+      ]
+    },
+    "class_type": "ControlNetApplyAdvanced",
+    "_meta": {
+      "title": "ControlNet应用(高级)"
+    }
+  },
+  "44": {
+    "inputs": {
+      "control_net_name": "diffusers_xl_depth_mid.safetensors"
+    },
+    "class_type": "ControlNetLoader",
+    "_meta": {
+      "title": "ControlNet加载器"
+    }
+  },
+  "45": {
+    "inputs": {
+      "ckpt_name": "depth_anything_vits14.pth",
+      "resolution": 960,
+      "image": [
+        "55",
+        0
+      ]
+    },
+    "class_type": "DepthAnythingPreprocessor",
+    "_meta": {
+      "title": "DA深度预处理器"
+    }
+  },
+  "47": {
+    "inputs": {
+      "interpolation": "LANCZOS",
+      "crop_position": "pad",
+      "sharpening": 0,
+      "image": [
+        "55",
+        0
+      ]
+    },
+    "class_type": "PrepImageForClipVision",
+    "_meta": {
+      "title": "CLIP视觉图像处理"
+    }
+  },
+  "48": {
+    "inputs": {
+      "positive": [
+        "43",
+        0
+      ],
+      "negative": [
+        "43",
+        1
+      ],
+      "vae": [
+        "4",
+        2
+      ],
+      "pixels": [
+        "55",
+        0
+      ]
+    },
+    "class_type": "InstructPixToPixConditioning",
+    "_meta": {
+      "title": "InstructPixToPix条件"
+    }
+  },
+  "49": {
+    "inputs": {
+      "coarse": "disable",
+      "resolution": 512,
+      "image": [
+        "55",
+        0
+      ]
+    },
+    "class_type": "LineArtPreprocessor",
+    "_meta": {
+      "title": "LineArt艺术线预处理器"
+    }
+  },
+  "52": {
+    "inputs": {
+      "model": "wd-v1-4-moat-tagger-v2",
+      "threshold": 0.35,
+      "character_threshold": 0.85,
+      "replace_underscore": "",
+      "trailing_comma": "1girl, solo, long_hair, looking_at_viewer, shirt, black_hair, upper_body, black_eyes, lips, realistic",
+      "exclude_tags": "",
+      "image": [
+        "55",
+        0
+      ]
+    },
+    "class_type": "WD14Tagger|pysssss",
+    "_meta": {
+      "title": "WD14反推提示词"
+    }
+  },
+  "55": {
+    "inputs": {
+      "image_path": [
+        "59",
+        0
+      ],
+      "RGBA": "false",
+      "filename_text_extension": "true"
+    },
+    "class_type": "LamLoadPathImage",
+    "_meta": {
+      "title": "加载网络图片或本地图片"
+    }
+  },
+  "56": {
+    "inputs": {
+      "ANY": [
+        "60",
+        0
+      ],
+      "IF_TRUE": [
+        "33",
+        0
+      ],
+      "IF_FALSE": [
+        "8",
+        0
+      ]
+    },
+    "class_type": "IfInnerExecute",
+    "_meta": {
+      "title": "判断选择"
+    }
+  },
+  "59": {
+    "inputs": {
+      "text": ""
+    },
+    "class_type": "TextInput_",
+    "_meta": {
+      "title": "文本"
+    }
+  },
+  "60": {
+    "inputs": {
+      "expression": "len(p0)>0",
+      "p0": [
+        "59",
+        0
+      ]
+    },
+    "class_type": "MultiParamFormula",
+    "_meta": {
+      "title": "多参代码表达式"
+    }
   }
 }

config/workflow/line2drawing.json

@@ -163,9 +163,9 @@
         0
       ]
     },
-    "class_type": "GetImage_(Width&Height) _O",
+    "class_type": "GetImageSize",
     "_meta": {
-      "title": "GetImage_(Width&Height) _O"
+      "title": "GetImageSize"
     }
   },
   "17": {

js/LamAllInOne.js

@@ -1,158 +0,0 @@
-import { app } from "/scripts/app.js";
-import { $el } from "../../../scripts/ui.js";
-import { ComfyWidgets } from "/scripts/widgets.js";
-import { api } from "../../../scripts/api.js";
-//background-image: url("http://alexfeds.com/wp-content/uploads/2018/04/save_icon.svg");
-$el("style", {
-	textContent: `
-    .lam_btn {
-        border: none;
-        color: grey;
-        padding: 12px 12px;
-        background-color: #494848;
-        font-size: 16px;
-        cursor: pointer;
-        background-repeat: repeat-y;
-    }
-`,
-parent: document.body,
-})
-
-let pb_cache = {};
-async function getPrompt(name) {
-    if(pb_cache[name])
-		return pb_cache[name];
-	else {
-        const resp = await api.fetchApi(`/lam/getPrompt?name=${name}`);
-        if (resp.status === 200) {
-            let data = await resp.json();
-            pb_cache[name] = data;
-            return data;
-        }
-        return undefined;
-    }
-    
-}
-function getTagList(tags,cat) {
-    let rlist=[]
-    Object.keys(tags).forEach((k) => {
-        if (typeof tags[k] === "string") {
-            let t=[k, tags[k]]
-            rlist.push($el(
-                "li.lam-model-tag",
-                {
-                    dataset: {
-                        tag: t[1],
-                        name: t[0],
-                        cat:cat,
-                        weight: 1
-                    },
-                    $: (el) => {
-                        el.onclick = () => {
-                            el.classList.toggle("lam-model-tag--selected");
-                        };
-                    },
-                },
-                [
-                    $el("p", {
-                        textContent: t[0],
-                    }),
-                    $el("span.lable", {
-                        textContent: t[1],
-                    }),
-                ]
-            ))
-        }else{
-            rlist.push(...getTagList(tags[k],cat))
-        }
-    });
-    return rlist
-}
-function getSelList(tags) {
-    let rlist=[]
-    Object.keys(tags).forEach((k) => {
-        rlist.push($el(
-            "li.lam-model-tag",
-            {
-                dataset: {
-                    name:  tags[k]['name'],
-                    tag: tags[k]['tag'],
-                    cat: tags[k]['cat'],
-                    weight: tags[k]['weight'],
-                },
-                // $: (el) => {
-                //     el.onclick = () => {
-                //         el.classList.add("lam-model-tag--del");
-                //     };
-                // },
-            },
-            [
-                $el("span.btn", {
-                    textContent: "-",
-                    $: (el) => {
-                        el.onclick = () => {
-                            el.parentElement.dataset.weight = (((Number(el.parentElement.dataset.weight)*1000)-(0.05*1000))/1000).toFixed(2);
-                        };
-                    },
-                }),
-                $el("p", {
-                    textContent: tags[k]['name'],
-                }),
-                $el("span.lable", {
-                    textContent: tags[k]['weight'],
-                }),
-                $el("span.lable", {
-                    textContent: tags[k]['cat'],
-                }),
-                $el("span.btn", {
-                    textContent: "+",
-                    $: (el) => {
-                        el.onclick = () => {
-                            el.parentElement.dataset.weight = (((Number(el.parentElement.dataset.weight)*1000)+(0.05*1000))/1000).toFixed(2);
-                        };
-                    },
-                }),
-                $el("span.btn", {
-                    textContent: "X",
-                    $: (el) => {
-                        el.onclick = () => {
-                            el.parentElement.classList.add("lam-model-tag--del");
-                        };
-                    },
-                }),
-            ]
-        ))
-    })
-    return rlist;
-}
-// Displays input text on a node
-app.registerExtension({
-    name: "LamAllInOne",
-    async beforeRegisterNodeDef(nodeType, nodeData, app) {
-		var names=['LamAllInOne']
-        if (names.indexOf(nodeData.name)>=0) {
-            const onNodeCreated = nodeType.prototype.onNodeCreated;
-            nodeType.prototype.onNodeCreated = function() {
-                const r = onNodeCreated?.apply(this, arguments);
-                this.widgets.splice(this.widgets.findIndex(obj => obj.name === 'prompt'),1);
-                this.widgets.splice(this.widgets.findIndex(obj => obj.name === 'negative'),1);
-                const promptText = $el("textarea.comfy-multiline-input.dynamic_prompt",{
-                    placeholder:"正向提示词",
-                    style:{
-                        width: '100%',
-                        height: '50%'  
-                    }
-                });
-                const btn1=$el('i.lam_btn',{
-                    style:{'mask-image':'url(/lam/getImage?type=icons&&name=close.svg)'}
-                })
-                const prompt =this.addDOMWidget('prompt',"list",$el("div",[promptText,btn1]))
-                const negative =this.addDOMWidget('negative',"list",$el("div",[]))
-
-                this.setSize([700, 800]);
-                return r;
-            }
-
-        }
-    },
-});

js/MultiTextConcatenate.js

@@ -10,18 +10,24 @@ app.registerExtension({
         if (showNames.indexOf(nodeData.name)>=0) {
             const onDrawForeground = nodeType.prototype.onDrawForeground;
             nodeType.prototype.onDrawForeground = function (ctx) {
-                const r = onDrawForeground?.apply?.(this, arguments);
-                const v = app.nodeOutputs?.[this.id + ""];
-                if (!this.flags.collapsed && v) {
-                    const text = v.value[0] + "";
-                    ctx.save();
-                    ctx.font = "bold 12px sans-serif";
-                    ctx.fillStyle = "dodgerblue";
-                    const sz = ctx.measureText(text);
-                    ctx.fillText(text, this.size[0]/2 - (sz.width + 5)/2, LiteGraph.NODE_SLOT_HEIGHT * 1);
-                    ctx.restore();
+                try {
+                    const r = onDrawForeground?.apply?.(this, arguments);
+                    const v = app.nodeOutputs?.[this.id + ""];
+                    if (!this.flags.collapsed && v) {
+                        const text = v.value[0] + "";
+                        ctx.save();
+                        ctx.font = "bold 12px sans-serif";
+                        ctx.fillStyle = "dodgerblue";
+                        const sz = ctx.measureText(text);
+                        ctx.fillText(text, this.size[0]/2 - (sz.width + 5)/2, LiteGraph.NODE_SLOT_HEIGHT * 1);
+                        ctx.restore();
+                    }
+                    return r;
+                } catch(e) {
+                    console.log(e)
                 }
-                return r;
+                return 
+
             };
         }
         var names=["MultiTextConcatenate",'MultiTextSelelct',"MultiIntFormula","MultiParamFormula"]
@@ -46,7 +52,7 @@ app.registerExtension({
                     this.outputPrefix="r"
                 }
                 function changCustomtext(){
-                    this.setSize( this.computeSize() );
+                    //this.setSize( this.computeSize() );
                 }
                 this.widgets.forEach(function(widget) {
                     if(widget.type!="customtext"){
@@ -56,11 +62,21 @@ app.registerExtension({
                     widget.draw = function (ctx,parentNode, widgetWidth, y, widgetHeight) {
                         draw?.apply(this, arguments);
                         if (this.inputEl.hidden) return;
-                        const transform = ctx.getTransform();
+                        const  margin = 20,
+                        clientRectBound = ctx.canvas.getBoundingClientRect(),
+                        transform = new DOMMatrix()
+                        .scaleSelf(
+                            clientRectBound.width / ctx.canvas.width,
+                            clientRectBound.height / ctx.canvas.height
+                        )
+                        .multiplySelf(ctx.getTransform())
+                        .translateSelf(margin, margin + y),
+                        w = (widgetWidth - (margin * 2+10) * transform.a) ;
                         let maxIoputSize=Math.max(...[parentNode.inputs.length,parentNode.outputs.length])
                         Object.assign(this.inputEl.style, {
-                            top: `${ maxIoputSize*transform.d*20+transform.d*5+transform.f}px`,
-                            height: `${32.0 * transform.d}px`,
+                            left: `${transform.a * margin + transform.e}px`,
+                            top: `${transform.d + transform.f}px`,
+                            width: `${w}px`,
                         });
                     }
                 });

js/litegraphCore.js

@@ -1,57 +1,27 @@
 import { app } from "/scripts/app.js";
-import { $el } from "../../../scripts/ui.js";
-$el("style", {
-	textContent: `
-    .node-zh-cn-name {
-        background-color: rgb(128, 213, 247);
-        color: #000;
-        display: flex;
-        align-items: center;
-        gap: 5px;
-        padding: 2px 5px;
-        cursor: pointer;
-    }
-    
-    .node-zh-cn-name:hover {
-        outline: 2px solid dodgerblue;
-    }
-    .node-zh-cn-name p {
-        margin: 0;
-    }
-    .node-zh-cn-name span {
-        text-align: center;
-        border-radius: 5px;
-        background-color: dodgerblue;
-        color: #fff;
-        padding: 2px;
-        position: relative;
-        min-width: 20px;
-        overflow: hidden;
-    }
-`});
 function addConvertToGroupOptions() {
     //中文搜索修改
     const showSearchBox = LGraphCanvas.prototype.showSearchBox;
     var funstr=showSearchBox+'';
-    var repStrs=['function(event, options) {','var filtered = keys.filter( inner_test_filter );'
-    ,'if ((!options.show_all_if_empty || str) && type.toLowerCase().indexOf(str) === -1)','help.dataset["type"] = escape(type);']
+    var repStrs=['function(event, options) {','LiteGraph.pointerListenerAdd(dialog,"leave", function(e) {','timeout_close = setTimeout(function() {']
     for(var i=0;i<repStrs.length;i++){ //如果对应几个字符串存在缺失，就放弃操作
         if(funstr.indexOf(repStrs[i])==-1){
             return;
         }
     }
     funstr=funstr.replace('function(event, options) {','');
-    funstr=funstr.replace('var filtered = keys.filter( inner_test_filter );', `var filtered = keys.filter( inner_test_filter );
-    filtered=filtered.map(function(e){ 
-        if(/.*[\u4e00-\u9fa5]+.*$/.test(LiteGraph.registered_node_types[e]['title'])) {
-            return [e,LiteGraph.registered_node_types[e]['title']];
-        }else{
-            return e;
-        }
-    });`);
-    funstr=funstr.replace('if ((!options.show_all_if_empty || str) && type.toLowerCase().indexOf(str) === -1)', `if ((!options.show_all_if_empty || str) && (type.toLowerCase().indexOf(str) === -1 && ctor.title.indexOf(str) === -1 ))`);
-    //funstr=funstr.replace('help.innerText = type;', `help.innerHTML = type instanceof Array?'<div class="node-zh-cn-name"><p>'+type[0]+'</p><span>'+type[1]+'</span></div>':'<div class="node-zh-cn-name"><p>'+type+'</p></div>'`);
-    funstr=funstr.replace('help.dataset["type"] = escape(type);',`help.dataset["type"]=type instanceof Array?escape(type[0]):escape(type);`);
+    funstr=funstr.replace('LiteGraph.pointerListenerAdd(dialog,"leave", function(e) {', `let isComposing=false;
+            LiteGraph.pointerListenerAdd(dialog,"compositionstart", function(e) {
+                isComposing=true;
+            });
+            LiteGraph.pointerListenerAdd(dialog,"compositionend", function(e) {
+                isComposing=false;
+            });
+            LiteGraph.pointerListenerAdd(dialog,"leave", function(e) {`);
+    funstr=funstr.replace('timeout_close = setTimeout(function() {', `if (isComposing){
+                    return;
+                }
+                timeout_close = setTimeout(function() {`);
     funstr=funstr.substring(0,funstr.length-2);
     LGraphCanvas.prototype.showSearchBox = new Function('event','options',funstr);
 }

js/openPoseEditorPlus.js

@@ -723,7 +723,7 @@ class OpenPose {
             img.style.height = "40px";
             img.addEventListener("dragend", (e) =>  {
               let name=e.target.alt;
-              let x=e.x;
+              let x=e.x;  
               let y=e.y;
               let width=thi.canvas.width;
               let height=thi.canvas.height;

py/FaceFusion.py

@@ -2,8 +2,7 @@ import cv2
 import os
 import torch
 import numpy as np
-from .src.face_fusion.image_face_fusion import ImageFaceFusion
-from .third_part.GFPGAN.gfpgan import GFPGANer
+from .face_fusion.image_face_fusion import ImageFaceFusion
 from PIL import Image
 import folder_paths
 
@@ -13,7 +12,6 @@ class FaceFusion:
         return {"required":
                     {"template_image": ("IMAGE", ),
                     "user_image": ("IMAGE", ),
-                    "facerestore": ([ True, False ], ),
                     },
                 }
 
@@ -23,12 +21,8 @@ class FaceFusion:
     RETURN_NAMES = ("图片",)
 
     FUNCTION = "face_fusion"
-    def face_fusion(self,template_image,user_image,facerestore):
+    def face_fusion(self,template_image,user_image):
         image_face_fusion = ImageFaceFusion(model_dir=folder_paths.models_dir+'/image-face-fusion')
-        if facerestore:
-            restorer_model = GFPGANer(model_path=os.path.join(folder_paths.models_dir+'/upscale_models','GFPGANv1.4.pth'), upscale=1, arch='clean',
-                                channel_multiplier=2, bg_upsampler=None)
-
         imaget_np=template_image.numpy()
         imageu_np=user_image.numpy()
         imgt_shape=imaget_np.shape
@@ -44,15 +38,11 @@ class FaceFusion:
             imaget=np.uint8(imaget)
             image = image_face_fusion.inference(imaget,imageu)
             image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-            if facerestore:
-                _,_,rimage=restorer_model.enhance(image)
-                if rimage is not None:
-                    image=rimage
+            
             image = np.array(image).astype(np.float32) / 255.0
             image = torch.from_numpy(image)[None,]
             image = image.squeeze()
             sample_frames.append(image)
-
         return (torch.stack(sample_frames),)
 
 NODE_CLASS_MAPPINGS = {

py/ForEnd.py

@@ -3,6 +3,9 @@ import numpy as np
 import requests
 import json
 from .src.utils.uitls import AlwaysEqualProxy
+from server import PromptServer
+from .WeChatAuth import prompt as promptFun
+from copy import deepcopy
 
 class ForEnd:
     def __init__(self):
@@ -15,9 +18,13 @@ class ForEnd:
             "required": {
                 "total": ("INT", {"forceInput": True}),
                 "i": ("INT",{"forceInput": True}),
-                "port": ("INT", {"default": 8188, "min": 1, "max": 99999}),
                 "obj": (AlwaysEqualProxy("*"),),
             },
+            "hidden": {
+                "unique_id": "UNIQUE_ID",           #节点编号
+                "prompt": "PROMPT",                 #流程节点信息
+                "extra_pnginfo": "EXTRA_PNGINFO"    #前端流程图信息
+            },
         }
     RETURN_TYPES = ()
     FUNCTION = "for_end_fun"
@@ -25,31 +32,34 @@ class ForEnd:
 
     CATEGORY = "lam"
 
-    def for_end_fun(self,total,i,port,obj):
+    def for_end_fun(self,total,i,obj,unique_id,prompt,extra_pnginfo):
+        pdata=json.loads('{}')
+        pdata['client_id']=PromptServer.instance.client_id
+        pdata['extra_data']={'extra_pnginfo':deepcopy(extra_pnginfo)}
+        pdata['prompt']=deepcopy(prompt)
+        start_id=prompt[unique_id]['inputs']['total'][0]
+        StatusInfo=pdata['extra_data']['extra_pnginfo']['StatusInfo'] if 'StatusInfo' in pdata['extra_data']['extra_pnginfo'] else ''
+        index=pdata['extra_data']['extra_pnginfo']['index'] if 'index' in pdata['extra_data']['extra_pnginfo'] else 0
+        if index==0:
+            StatusInfo=''
+
+        i=i+pdata['prompt'][start_id]['inputs']['stop']
+        pdata['prompt'][start_id]['inputs']['i']=i
+        pdata['extra_data']['extra_pnginfo']['index']=index+1
+        
+        result = "第"+str(index+1)+"论结束........."
+
         if i>=total:
-            return { "ui": { "text":'循环结束' } }
+            result = result+ '\n整个循环结束'
+        else:
+            r = promptFun(PromptServer.instance,pdata)
 
-        r = requests.get("http://127.0.0.1:"+str(port)+"/queue")
-        result = r.text
-        #print(result)
-        data = json.loads(result)
-        #rdata=data[list(data.keys())[0]]
-        rdata=data['queue_running'][0] if len(data['queue_running'])>0 else []
-        index=0
-        if len(rdata)>0:
-            pdata=json.loads('{}')
-            pdata['client_id']=rdata[3]['client_id']
-            pdata['extra_data']={'extra_pnginfo':rdata[3]['extra_pnginfo']}
-            pdata['prompt']=rdata[2]
-            for key in list(pdata['prompt'].keys()):
-                if pdata['prompt'][key]['class_type']=='ForStart':
-                    pdata['prompt'][key]['inputs']['i']=i
-                    index=i//pdata['prompt'][key]['inputs']['stop']
-                    break
-            r = requests.post("http://127.0.0.1:"+str(port)+"/prompt",json=pdata)
-            result = r.text
+        if StatusInfo.endswith('循环结束')==False and len(StatusInfo)>0:
+            result=StatusInfo+'\n'+result
 
-        return { "ui": { "text":"第"+str(index+1)+"次循环结果："+result} }
+        pdata['extra_data']['extra_pnginfo']['StatusInfo']=result
+
+        return { "ui": { "text":result} }
 
 NODE_CLASS_MAPPINGS = {
     "ForEnd": ForEnd

py/GLM3Prompt.py

@@ -0,0 +1,47 @@
+import json
+import openai
+
+class GLM3Prompt:
+    """
+    ChatGLM3接口调用
+    """
+    def __init__(self):
+        pass
+    
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "glm3_url": ("STRING",{"default": "http://localhost:8000/v1"}),
+                "system_prompt": ("STRING", {"multiline": True,"default":"你是名绘图提示词专家，你会根据客户的描述内容，精确返回对应图片提示，注意提示词必须是英文"}),
+                "text": ("STRING", {"multiline": True})
+            },
+        }
+
+    RETURN_TYPES = ("STRING",)
+    RETURN_NAMES = ("提示词",)
+
+    FUNCTION = "translate"
+
+    #OUTPUT_NODE = False
+
+    CATEGORY = "lam"
+
+    def translate(self, glm3_url,system_prompt,text):
+        openai.api_base = glm3_url
+        openai.api_key = "xxx"
+        params = dict(model="chatglm3", messages=[{"role": "system", "content": system_prompt},{"role": "user", "content": text}], stream=False)
+        response = openai.ChatCompletion.create(**params)
+        reply = response.choices[0].message.content
+        return (reply,)
+
+# A dictionary that contains all nodes you want to export with their names
+# NOTE: names should be globally unique
+NODE_CLASS_MAPPINGS = {
+    "GLM3Prompt": GLM3Prompt
+}
+
+# A dictionary that contains the friendly/humanly readable titles for the nodes
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "GLM3Prompt": "ChatGLM3提示词生成工具"
+}

py/Image2TalkingFace.py

@@ -1,56 +0,0 @@
-import folder_paths
-import cv2
-import os, shutil
-import torch
-import numpy as np
-from .src.gradio_demo import SadTalker  
-
-class Image2TalkingFace:
-    def __init__(self):
-        self.output_dir = os.path.join(folder_paths.get_output_directory(), 'video')
-        self.checkpoint_path=os.path.join(folder_paths.models_dir, 'SadTalker')
-        if not os.path.exists(self.output_dir):
-            os .makedirs(self.output_dir)
-        if not os.path.exists(self.checkpoint_path):
-            os .makedirs(self.checkpoint_path)
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required":
-                    {
-                        "images": ("IMAGE", ),
-                        "audioPath": ("STRING", {"forceInput": True}),
-                        "pose_style": ("INT", {"default": 0, "min": 0, "max": 64}),
-                        "size_of_image": ([256, 512], ),
-                        "preprocess_type":(['crop', 'resize','full', 'extcrop', 'extfull'], ),
-                        "is_still_mode": ([ False, True ], ),
-                        "batch_size": ("INT", {"default": 1, "min": 1, "max": 10}),
-                        "enhancer": ([ False, True ], ),
-                        "filename_prefix": ("STRING", {"default": "comfyUI"}),
-                    },
-                }
-
-    CATEGORY = "lam"
-
-    RETURN_TYPES = ("STRING",)
-    RETURN_NAMES = ("视频文件名",)
-    FUNCTION = "sad_talker"
-    OUTPUT_NODE = True
-
-    def sad_talker(self, images,audioPath,pose_style,size_of_image,preprocess_type,is_still_mode,batch_size,enhancer,filename_prefix):
-        sad_talker = SadTalker(self.checkpoint_path, "custom_nodes/ComfyUI_Lam/py/src/config", lazy_load=True)
-        autio_path=sad_talker.run_image_2_video(images, audioPath,preprocess_type,is_still_mode,enhancer,batch_size,size_of_image,pose_style,result_dir=self.output_dir)
-        full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir)
-        file = f"{filename}_{counter:05}_.mp4"
-        shutil.move(autio_path, os.path.join(full_output_folder, file))
-        autio_path=os.path.join(full_output_folder, file)
-        return {"ui": {"text": "视频合成成功，保存路径："+autio_path,
-        'videos':[{'filename':file,'type':'output','subfolder':'video'}]}, "result": (autio_path,)}
-
-NODE_CLASS_MAPPINGS = {
-    "Image2TalkingFace": Image2TalkingFace
-}
-
-# A dictionary that contains the friendly/humanly readable titles for the nodes
-NODE_DISPLAY_NAME_MAPPINGS = {
-    "Image2TalkingFace": "图像转说话的人脸动画"
-}

py/LamAllInOne.py

@@ -1,48 +0,0 @@
-from server import PromptServer
-from aiohttp import web
-import json
-import os
-import yaml
-
-
-class LamAllInOne:
-    """
-    提示词选择工具
-    """
-    def __init__(self):
-        pass
-    
-    @classmethod
-    def INPUT_TYPES(s):
-        
-        return {
-            "required": {
-                "prompt": ("STRING", {"multiline": True, "default": ""}),
-                "negative": ("STRING", {"multiline": True, "default": ""}),
-            },
-            "hidden": {"unique_id": "UNIQUE_ID","wprompt":"PROMPT"},
-        }
-
-    RETURN_TYPES = ("STRING","STRING",)
-    RETURN_NAMES = ("提示词","反向词",)
-
-    FUNCTION = "all_in_one_prompt"
-
-    #OUTPUT_NODE = False
-
-    CATEGORY = "lam"
-
-    def all_in_one_prompt(self,prompt,negative):
-        
-        return (prompt,negative,)
-
-# A dictionary that contains all nodes you want to export with their names
-# NOTE: names should be globally unique
-NODE_CLASS_MAPPINGS = {
-    "LamAllInOne": LamAllInOne
-}
-
-# A dictionary that contains the friendly/humanly readable titles for the nodes
-NODE_DISPLAY_NAME_MAPPINGS = {
-    "LamAllInOne": "提示词工具"
-}

py/LoadVideo.py

@@ -3,7 +3,6 @@ import cv2
 import os
 import torch
 import numpy as np
-from ffmpy import FFmpeg
 
 class LamLoadVideo:
     def __init__(self):
@@ -67,8 +66,9 @@ class LamLoadVideo:
             full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir)
             file = f"{filename}_{counter:05}_.mp3"
             filePathName=os.path.join(full_output_folder, file)
-            ff = FFmpeg(inputs={videoPath: None},outputs={filePathName: '-f {} -vn'.format('mp3')})
-            ff.run()
+            #提取音频
+            cmd = r"ffmpeg -y -hide_banner -loglevel error -i %s %s"%(videoPath, filePathName)
+            os.system(cmd)  
         return {"ui": {"text": "音频提取成功，保存路径："+filePathName if extract_audio else '需要提取音频请设置extract_audio为True'}, "result": (torch.stack(sample_frames),fps,frames,filePathName,)}
 
 NODE_CLASS_MAPPINGS = {

py/PromptTranslator.py

@@ -65,14 +65,10 @@ class PromptTranslator:
         # Send request
         r = requests.post(url, params=payload, headers=headers)
         result = r.json()
-        texts=''
-        for result in result['result']['trans_result']:
-            if result['dst']:
-                texts+=result['dst']+'\n'
         if 'error_code' in result:
             return { "ui": { "text":"翻译失败:"+result['error_msg']},"result": (text) }
         else:
-            return { "ui": { "text":"翻译结果："+texts},"result": (texts,)}
+            return { "ui": { "text":"翻译结果："+result['result']['trans_result'][0]['dst']},"result": (result['result']['trans_result'][0]['dst'],)}
 
     def getAccessToken(self,apiKey,secretKey):
         url = f"https://aip.baidubce.com/oauth/2.0/token?grant_type=client_credentials&client_id={apiKey}&client_secret={secretKey}"

py/ReferenceOnly.py

@@ -1,224 +0,0 @@
-import torch
-#说明
-#1、 T2i 为原版的reference only 只能用于文生图，如果要图生图使用 使用Latent 复合的话 会有相互污染出现
-#2、个人改写了代码，增加的了三个节点，主要是下面三个：
-#   ①、用于图生图的 ReferenceOnlySimple_Img2Img
-#   ②、用于文生图的使用遮罩：ReferenceOnlySimple_Masks
-#   ③、用于图生图的 双reference only 参考的  ReferenceOnly_TwoReference_Img2Img  Ps:大概率是个玩具
-#使用方法：从网盘下载python 文件，放到Comfyui的 custom_nodes 目录下即可，
-# 示列：“E:\ComfyUI\ComfyUI\custom_nodes”
-#PS:本插件主要用于对于人物一致性的探索，代码为个人改写的，仅仅只是尝试探索，个人B站账号：zhairy ,QQ：366310700 
-#本插件开源，请勿用于倒卖，也勿用于其他违规用途！
-#ReferenceOnly文生图节点
-class ReferenceOnlySimple_T2i:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "model": ("MODEL",),
-                              "reference": ("LATENT",),
-                              "batch_size": ("INT", {"default": 1, "min": 1, "max": 64})
-                              }}
-
-    RETURN_TYPES = ("MODEL", "LATENT")
-    FUNCTION = "reference_only_T2i"
-    CATEGORY = "custom_node_experiments"
-
-    def reference_only_T2i(self, model, reference, batch_size):
-        model_reference = model.clone()
-        size_latent = list(reference["samples"].shape)
-        size_latent[0] = batch_size
-        latent = {}
-        latent["samples"] = torch.zeros(size_latent)
-  
-        batch = latent["samples"].shape[0] + reference["samples"].shape[0]
-        # batch = reference["samples"].shape[0]
-        def reference_apply(q, k, v, extra_options):
-            k = k.clone().repeat(1, 2, 1)
-            offset = 0
-            if q.shape[0] > batch:
-                offset = batch
-
-            for o in range(0, q.shape[0], batch):
-                for x in range(1, batch):
-                    k[x + o, q.shape[1]:] = q[o,:]
-
-            return q, k, k
-
-        model_reference.set_model_attn1_patch(reference_apply)
-        out_latent = torch.cat((reference["samples"], latent["samples"]))
-        if "noise_mask" in latent:
-            mask = latent["noise_mask"]
-        else:
-            mask = torch.ones((64,64), dtype=torch.float32, device="cpu")
-
-        if len(mask.shape) < 3:
-            mask = mask.unsqueeze(0)
-        if mask.shape[0] < latent["samples"].shape[0]:
-            print(latent["samples"].shape, mask.shape)
-            mask = mask.repeat(latent["samples"].shape[0], 1, 1)
-
-        out_mask = torch.zeros((1,mask.shape[1],mask.shape[2]), dtype=torch.float32, device="cpu")
-        return (model_reference, {"samples": out_latent, "noise_mask": torch.cat((out_mask, mask))})
-    
-#ReferenceOnly图生图节点
-class ReferenceOnlySimple_Img2Img:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "model": ("MODEL",),
-                              "reference": ("LATENT",),
-                              "latent": ("LATENT",),
-                              "batch_size": ("INT", {"default": 1, "min": 1, "max": 64})
-                              }}
-
-    RETURN_TYPES = ("MODEL", "LATENT")
-    FUNCTION = "reference_only_img2img"
-    CATEGORY = "custom_node_experiments"
-
-    def reference_only_img2img(self, model, reference, latent, batch_size):
-        model_reference = model.clone()
-        size_latent = list(reference["samples"].shape)
-        size_latent[0] = batch_size
-
-        batch = latent["samples"].shape[0] + reference["samples"].shape[0]
-        def reference_apply(q, k, v, extra_options):
-            k = k.clone().repeat(1, 2, 1)
-            offset = 0
-            if q.shape[0] > batch:
-                offset = batch
-
-            for o in range(0, q.shape[0], batch):
-                for x in range(1, batch):
-                    k[x + o, q.shape[1]:] = q[o,:]
-
-            return q, k, k
-
-        model_reference.set_model_attn1_patch(reference_apply)
-        out_latent = torch.cat((reference["samples"], latent["samples"]))
-        if "noise_mask" in latent:
-            mask = latent["noise_mask"]
-        else:
-            mask = torch.ones((64,64), dtype=torch.float32, device="cpu")
-
-        if len(mask.shape) < 3:
-            mask = mask.unsqueeze(0)
-        if mask.shape[0] < latent["samples"].shape[0]:
-            print(latent["samples"].shape, mask.shape)
-            mask = mask.repeat(latent["samples"].shape[0], 1, 1)
-
-        out_mask = torch.zeros((1,mask.shape[1],mask.shape[2]), dtype=torch.float32, device="cpu")
-        return (model_reference, {"samples": out_latent, "noise_mask": torch.cat((out_mask, mask))})
-
-
-#ReferenceOnly遮罩节点
-class ReferenceOnlySimple_Masks:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "model": ("MODEL",),
-                              "reference": ("LATENT",),
-                              "Mask_latent": ("LATENT",),
-                              "batch_size": ("INT", {"default": 1, "min": 1, "max": 64})
-                              }}
-
-    RETURN_TYPES = ("MODEL", "LATENT")
-    FUNCTION = "reference_only_mask"
-    CATEGORY = "custom_node_experiments"
-
-    def reference_only_mask(self, model, reference, Mask_latent, batch_size):
-        model_reference = model.clone()
-        size_latent = list(reference["samples"].shape)
-        size_latent[0] = batch_size
-
-        batch = Mask_latent["samples"].shape[0] + reference["samples"].shape[0]
-        def reference_apply(q, k, v, extra_options):
-            k = k.clone().repeat(1, 2, 1)
-            offset = 0
-            if q.shape[0] > batch:
-                offset = batch
-            for o in range(0, q.shape[0], batch):
-                for x in range(1, batch):
-                    k[x + o, q.shape[1]:] = q[o,:]
-            return q, k, k
-
-        model_reference.set_model_attn1_patch(reference_apply)
-        out_latent = torch.cat((reference["samples"], Mask_latent["samples"]))
-        if "noise_mask" in Mask_latent:
-            mask = Mask_latent["noise_mask"]
-        else:
-            mask = torch.ones((64,64), dtype=torch.float32, device="cpu")
-
-        if len(mask.shape) < 3:
-            mask = mask.unsqueeze(0)
-        if mask.shape[0] < Mask_latent["samples"].shape[0]:
-            print(Mask_latent["samples"].shape, mask.shape)
-            mask = mask.repeat(Mask_latent["samples"].shape[0], 1, 1)
-
-        if len(mask.shape) < 4:
-            out_mask = torch.zeros((1,mask.shape[1],mask.shape[2]), dtype=torch.float32, device="cpu")
-        else:
-            out_mask = torch.zeros((1,1,mask.shape[2],mask.shape[3]), dtype=torch.float32, device="cpu")
-
-        return (model_reference, {"samples": out_latent, "noise_mask": torch.cat((out_mask, mask))})  
-
-#ReferenceOnly_双参考图
-class ReferenceOnly_TwoReference_Img2Img:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "model": ("MODEL",),
-                              "reference": ("LATENT",),
-                              "reference2": ("LATENT",),
-                              "latent": ("LATENT",),
-                              "batch_size": ("INT", {"default": 1, "min": 1, "max": 64})
-                              }}
-
-    RETURN_TYPES = ("MODEL", "LATENT")
-    FUNCTION = "reference_only_tworeference_img2img"
-    CATEGORY = "custom_node_experiments"
-
-    def reference_only_tworeference_img2img(self, model, reference, reference2,latent,batch_size):
-        model_reference = model.clone()
-        size_latent = list(reference["samples"].shape)
-        size_latent[0] = batch_size
-
-        batch = latent["samples"].shape[0] + reference["samples"].shape[0] + reference2["samples"].shape[0]
-
-        def reference_apply(q, k, v, extra_options):
-            k = k.clone().repeat(1, 2, 1)
-            offset = 0
-            if q.shape[0] > batch:
-                offset = batch
-
-            re = extra_options["transformer_index"] % 2
-            for o in range(0, q.shape[0], batch):
-                for x in range(1, batch):
-                    k[x + o, q.shape[1]:] = q[o + re,:]
-
-            return q, k, k
-
-        model_reference.set_model_attn1_patch(reference_apply)
-        out_latent = torch.cat((reference["samples"], reference2["samples"], latent["samples"]))
-        if "noise_mask" in latent:
-            mask = latent["noise_mask"]
-        else:
-            mask = torch.ones((64,64), dtype=torch.float32, device="cpu")
-
-        if len(mask.shape) < 3:
-            mask = mask.unsqueeze(0)
-        if mask.shape[0] < latent["samples"].shape[0]:
-            print(latent["samples"].shape, mask.shape)
-            mask = mask.repeat(latent["samples"].shape[0], 1, 1)
-
-        out_mask = torch.zeros((1,mask.shape[1],mask.shape[2]), dtype=torch.float32, device="cpu")
-        return (model_reference, {"samples": out_latent, "noise_mask": torch.cat((out_mask,out_mask, mask))})
-
-NODE_CLASS_MAPPINGS = {
-    "ReferenceOnlySimple_T2i": ReferenceOnlySimple_T2i,
-    "ReferenceOnlySimple_Img2Img": ReferenceOnlySimple_Img2Img,
-    "ReferenceOnlySimple_Masks": ReferenceOnlySimple_Masks,
-    "ReferenceOnly_TwoReference_Img2Img":ReferenceOnly_TwoReference_Img2Img,
-}
-
-NODE_DISPLAY_NAME_MAPPINGS = {
-    "ReferenceOnlySimple_T2i": "仅参考文生图节点",
-    "ReferenceOnlySimple_Img2Img": "仅参考图生图",
-    "ReferenceOnlySimple_Masks": "仅参考遮罩节点",
-    "ReferenceOnly_TwoReference_Img2Img": "仅参考双参考图"
-}

py/RoleManagement.py

@@ -1,45 +0,0 @@
-from PIL import Image
-import numpy as np
-import torch
-import os
-import folder_paths
-from .src.utils.uitls import AlwaysEqualProxy,AlwaysTupleZero
-
-class RoleManagement:
-    def __init__(self):
-        self.input_role_dir = os.path.join(folder_paths.get_input_directory(), 'role')
-        if not os.path.exists(self.input_role_dir):
-            os .makedirs(self.input_role_dir)
-
-    @classmethod
-    def INPUT_TYPES(cls):
-        input_role_dir = os.path.join(folder_paths.get_input_directory(), 'role')
-        if not os.path.exists(input_role_dir):
-            os .makedirs(input_role_dir)
-        files = [f for f in os.listdir(input_role_dir) if os.path.isdir(os.path.join(input_role_dir, f))]
-        return {
-            "required": {
-                "role": (sorted(files), ),
-            }
-        }
-    RETURN_TYPES = AlwaysTupleZero(AlwaysEqualProxy("*"),) #"IMAGE","IMAGE","STRING",
-    RETURN_NAMES = ("调试参数",) #"脸部图","姿态图","表情",
-    FUNCTION = "get_role"
-
-    CATEGORY = "lam"
-
-    def get_role(self, **kwargs):
-        params = {}
-        for arg in kwargs:
-            params[arg] = kwargs[arg]
-        #video_path = folder_paths.get_annotated_filepath(video,self.input_role_dir)
-        return (params, )
-
-NODE_CLASS_MAPPINGS = {
-    "RoleManagement": RoleManagement
-}
-
-# A dictionary that contains the friendly/humanly readable titles for the nodes
-NODE_DISPLAY_NAME_MAPPINGS = {
-    "RoleManagement": "角色管理"
-}

py/RoopFaceSwap.py

@@ -1,71 +0,0 @@
-import cv2
-import os
-import torch
-import numpy as np
-from .scripts.swapper import getFaceSwapModel,get_face_single
-from .third_part.GFPGAN.gfpgan import GFPGANer
-from PIL import Image
-import folder_paths
-
-class RoopFaceSwap:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required":
-                    {"template_image": ("IMAGE", ),
-                    "user_image": ("IMAGE", ),
-                    "facerestore": ([ True, False ], ),
-                    },
-                }
-
-    CATEGORY = "lam"
-
-    RETURN_TYPES = ("IMAGE",)
-    RETURN_NAMES = ("图片",)
-
-    FUNCTION = "swap_face"
-
-    def swap_face(self,template_image,user_image,facerestore):
-        face_swapper = getFaceSwapModel(folder_paths.models_dir+'/roop-face-swap/inswapper_128.onnx')
-        if facerestore:
-            restorer_model = GFPGANer(model_path=os.path.join(folder_paths.models_dir+'/upscale_models','GFPGANv1.4.pth'), upscale=1, arch='clean',
-                                channel_multiplier=2, bg_upsampler=None)
-
-        imaget_np=template_image.numpy()
-        imageu_np=user_image.numpy()
-        imgt_shape=imaget_np.shape
-        imau_shape=imageu_np.shape
-        if imgt_shape[0]<=0 or imau_shape[0]<=0:
-            return (template_image,)
-
-        sample_frames = []
-        imageu=imageu_np[0]*255
-        imageu=np.uint8(imageu)
-        source_face = get_face_single(imageu, face_index=0)
-        for i in range(imgt_shape[0]):
-            imaget=imaget_np[i]*255
-            imaget=np.uint8(imaget)
-            target_face = get_face_single(imaget, face_index=0)
-            if target_face is not None and source_face is not None:
-                image = face_swapper.get(imaget, target_face, source_face)
-            else:
-                image = imaget
-            #image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-            if facerestore:
-                _,_,rimage=restorer_model.enhance(image)
-                if rimage is not None:
-                    image=rimage
-            image = np.array(image).astype(np.float32) / 255.0
-            image = torch.from_numpy(image)[None,]
-            image = image.squeeze()
-            sample_frames.append(image)
-
-        return (torch.stack(sample_frames),)
-
-NODE_CLASS_MAPPINGS = {
-    "RoopFaceSwap": RoopFaceSwap
-}
-
-# A dictionary that contains the friendly/humanly readable titles for the nodes
-NODE_DISPLAY_NAME_MAPPINGS = {
-    "RoopFaceSwap": "人脸替换"
-}

py/Video2TalkingFace.py

@@ -1,56 +0,0 @@
-import folder_paths
-import cv2
-import os,shutil
-import torch
-import numpy as np
-from .src.gradio_demo import SadTalker  
-
-class Video2TalkingFace:
-    def __init__(self):
-        self.output_dir = os.path.join(folder_paths.get_output_directory(), 'video')
-        self.checkpoint_path=os.path.join(folder_paths.models_dir, 'SadTalker')
-        if not os.path.exists(self.output_dir):
-            os .makedirs(self.output_dir)
-        if not os.path.exists(self.checkpoint_path):
-            os .makedirs(self.checkpoint_path)
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required":
-                    {
-                        "videoPath": ("STRING", {"forceInput": True}),
-                        "audioPath": ("STRING", {"forceInput": True}),
-                        "batch_size": ("INT", {"default": 1, "min": 1, "max": 10}),
-                        "enhancer": ([ 'none','lip','face' ], ),
-                        "filename_prefix": ("STRING", {"default": "comfyUI"}),
-                    },
-                }
-
-    CATEGORY = "lam"
-
-    RETURN_TYPES = ("STRING",)
-    RETURN_NAMES = ("视频文件名",)
-    FUNCTION = "sad_talker"
-    OUTPUT_NODE = True
-
-    def sad_talker(self, videoPath,audioPath,batch_size,enhancer,filename_prefix):
-        if os.path.exists(videoPath) == False:
-            raise Exception('视频文件不存在')
-        if os.path.exists(audioPath) == False:
-            raise Exception('音频文件不存在')
-        sad_talker = SadTalker(self.checkpoint_path, "custom_nodes/ComfyUI_Lam/py/src/config", lazy_load=True)
-        autio_path=sad_talker.run_video_2_video(videoPath,audioPath,enhancer,batch_size,result_dir=self.output_dir)
-        full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir)
-        file = f"{filename}_{counter:05}_.mp4"
-        shutil.move(autio_path, os.path.join(full_output_folder, file))
-        autio_path=os.path.join(full_output_folder, file)
-        return {"ui": {"text": "视频合成成功，保存路径："+autio_path,
-        "filename_prefix": ("STRING", {"default": "comfyUI"}),}, "result": (autio_path,)}
-
-NODE_CLASS_MAPPINGS = {
-    "Video2TalkingFace": Video2TalkingFace
-}
-
-# A dictionary that contains the friendly/humanly readable titles for the nodes
-NODE_DISPLAY_NAME_MAPPINGS = {
-    "Video2TalkingFace": "视频转说话的人脸动画"
-}

py/VideoAddAudio.py

@@ -3,7 +3,6 @@ import cv2
 import os
 import torch
 import numpy as np
-from ffmpy import FFmpeg
 
 class VideoAddAudio:
     def __init__(self):
@@ -40,10 +39,9 @@ class VideoAddAudio:
         full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir)
         file = f"{filename}_{counter:05}_.mp4"
         result=os.path.join(full_output_folder, file)
-        ff = FFmpeg(
-            inputs={videoPath: None, audioPath: None},
-            outputs={result: '-map 0:v -map 1:a -c:v copy -c:a {} -shortest'.format(_codec)})
-        ff.run()
+        #插入音频
+        cmd = r'ffmpeg -y -hide_banner -loglevel error -i "%s" -i "%s" -vcodec copy "%s"' % (videoPath, audioPath, result)
+        os.system(cmd)  
         return {"ui": {"text": "视频插入音频成功，保存路径："+result,
         'videos':[{'filename':file,'type':'output','subfolder':'video'}]}}
 

py/VideoFaceFusion.py

@@ -1,17 +1,108 @@
 import cv2
 import os
+import sys
 import torch
+import model_management
 import numpy as np
-from .src.face_fusion.image_face_fusion import ImageFaceFusion
-from .third_part.GFPGAN.gfpgan import GFPGANer
+from .facelib.utils.face_restoration_helper import FaceRestoreHelper
+from .facelib.detection.retinaface import retinaface
+from .face_fusion.image_face_fusion import ImageFaceFusion
 from PIL import Image
 import folder_paths
-from ffmpy import FFmpeg
 import imageio
 from tqdm import tqdm 
+from torchvision.transforms.functional import normalize
+
+def img2tensor(imgs, bgr2rgb=True, float32=True):
+    """Numpy array to tensor.
+
+    Args:
+        imgs (list[ndarray] | ndarray): Input images.
+        bgr2rgb (bool): Whether to change bgr to rgb.
+        float32 (bool): Whether to change to float32.
+
+    Returns:
+        list[tensor] | tensor: Tensor images. If returned results only have
+            one element, just return tensor.
+    """
+
+    def _totensor(img, bgr2rgb, float32):
+        if img.shape[2] == 3 and bgr2rgb:
+            if img.dtype == 'float64':
+                img = img.astype('float32')
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        img = torch.from_numpy(img.transpose(2, 0, 1))
+        if float32:
+            img = img.float()
+        return img
+
+    if isinstance(imgs, list):
+        return [_totensor(img, bgr2rgb, float32) for img in imgs]
+    else:
+        return _totensor(imgs, bgr2rgb, float32)
+
+
+def tensor2img(tensor, rgb2bgr=True, out_type=np.uint8, min_max=(0, 1)):
+    """Convert torch Tensors into image numpy arrays.
+
+    After clamping to [min, max], values will be normalized to [0, 1].
+
+    Args:
+        tensor (Tensor or list[Tensor]): Accept shapes:
+            1) 4D mini-batch Tensor of shape (B x 3/1 x H x W);
+            2) 3D Tensor of shape (3/1 x H x W);
+            3) 2D Tensor of shape (H x W).
+            Tensor channel should be in RGB order.
+        rgb2bgr (bool): Whether to change rgb to bgr.
+        out_type (numpy type): output types. If ``np.uint8``, transform outputs
+            to uint8 type with range [0, 255]; otherwise, float type with
+            range [0, 1]. Default: ``np.uint8``.
+        min_max (tuple[int]): min and max values for clamp.
+
+    Returns:
+        (Tensor or list): 3D ndarray of shape (H x W x C) OR 2D ndarray of
+        shape (H x W). The channel order is BGR.
+    """
+    if not (torch.is_tensor(tensor) or (isinstance(tensor, list) and all(torch.is_tensor(t) for t in tensor))):
+        raise TypeError(f'tensor or list of tensors expected, got {type(tensor)}')
+
+    if torch.is_tensor(tensor):
+        tensor = [tensor]
+    result = []
+    for _tensor in tensor:
+        _tensor = _tensor.squeeze(0).float().detach().cpu().clamp_(*min_max)
+        _tensor = (_tensor - min_max[0]) / (min_max[1] - min_max[0])
+
+        n_dim = _tensor.dim()
+        if n_dim == 4:
+            img_np = make_grid(_tensor, nrow=int(math.sqrt(_tensor.size(0))), normalize=False).numpy()
+            img_np = img_np.transpose(1, 2, 0)
+            if rgb2bgr:
+                img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
+        elif n_dim == 3:
+            img_np = _tensor.numpy()
+            img_np = img_np.transpose(1, 2, 0)
+            if img_np.shape[2] == 1:  # gray image
+                img_np = np.squeeze(img_np, axis=2)
+            else:
+                if rgb2bgr:
+                    img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
+        elif n_dim == 2:
+            img_np = _tensor.numpy()
+        else:
+            raise TypeError('Only support 4D, 3D or 2D tensor. ' f'But received with dimension: {n_dim}')
+        if out_type == np.uint8:
+            # Unlike MATLAB, numpy.unit8() WILL NOT round by default.
+            img_np = (img_np * 255.0).round()
+        img_np = img_np.astype(out_type)
+        result.append(img_np)
+    if len(result) == 1:
+        result = result[0]
+    return result
 
 class VideoFaceFusion:
     def __init__(self):
+        self.face_helper = None
         self.output_dir = os.path.join(folder_paths.get_output_directory(), 'video')
         if not os.path.exists(self.output_dir):
             os .makedirs(self.output_dir)
@@ -23,7 +114,12 @@ class VideoFaceFusion:
                     "user_image": ("IMAGE", ),
                     "filename_prefix": ("STRING", {"default": "comfyUI"}),
                     },
+                "optional": {
+                    "facedetection": (["retinaface_resnet50", "retinaface_mobile0.25", "YOLOv5l", "YOLOv5n"],),
+                    "upscale_model": ("UPSCALE_MODEL",),
+                    }
                 }
+                
 
     CATEGORY = "lam"
 
@@ -32,10 +128,13 @@ class VideoFaceFusion:
     OUTPUT_NODE = True
 
     FUNCTION = "face_fusion"
-    def face_fusion(self,templateVideoPath,user_image,filename_prefix):
+    def face_fusion(self,templateVideoPath,user_image,filename_prefix,facedetection,upscale_model=None):
         image_face_fusion = ImageFaceFusion(model_dir=folder_paths.models_dir+'/image-face-fusion')
-        restorer_model = GFPGANer(model_path=os.path.join(folder_paths.models_dir+'/upscale_models','GFPGANv1.4.pth'), upscale=1, arch='clean',
-                                channel_multiplier=2, bg_upsampler=None)
+        if upscale_model:
+            device = model_management.get_torch_device()
+            upscale_model.to(device)
+            if self.face_helper is None:
+                self.face_helper = FaceRestoreHelper(1, face_size=512, crop_ratio=(1, 1), det_model=facedetection, save_ext='png', use_parse=True, device=device)
         imageu_np=user_image.numpy()
         imau_shape=imageu_np.shape
 
@@ -69,19 +168,58 @@ class VideoFaceFusion:
         with imageio.get_writer(videoPathTemp, fps=fps) as video:
             for i in tqdm(range(len(image_list)), '写入进度:'):
                 image=image_list[i]
-                _,_,rimage=restorer_model.enhance(image)
-                if rimage is not None:
-                    image=rimage
-                video.append_data(image)# 写入视频
+                if upscale_model is None or self.face_helper is None:
+                    video.append_data(image)# 写入视频
+                    continue
+
+                image = image.squeeze()
+                image = image[:, :, ::-1]
+                original_resolution = image.shape[0:2]
+
+                self.face_helper.clean_all()
+                self.face_helper.read_image(image)
+                self.face_helper.get_face_landmarks_5(only_center_face=False, resize=640, eye_dist_threshold=5)
+                self.face_helper.align_warp_face()
+
+                restored_face = None
+                for idx, cropped_face in enumerate(self.face_helper.cropped_faces):
+                    cropped_face_t = img2tensor(cropped_face / 255., bgr2rgb=True, float32=True)
+                    normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
+                    cropped_face_t = cropped_face_t.unsqueeze(0).to(device)
+
+                    try:
+                        with torch.no_grad():
+                            #output = upscale_model(cropped_face_t, w=strength, adain=True)[0]
+                            output = upscale_model(cropped_face_t)[0]
+                            restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))
+                        del output
+                        torch.cuda.empty_cache()
+                    except Exception as error:
+                        print(f'\tFailed inference for CodeFormer: {error}', file=sys.stderr)
+                        restored_face = tensor2img(cropped_face_t, rgb2bgr=True, min_max=(-1, 1))
+
+                    restored_face = restored_face.astype('uint8')
+                    self.face_helper.add_restored_face(restored_face)
+
+                self.face_helper.get_inverse_affine(None)
+
+                restored_img = self.face_helper.paste_faces_to_input_image()
+                restored_img = restored_img[:, :, ::-1]
+
+                if original_resolution != restored_img.shape[0:2]:
+                    restored_img = cv2.resize(restored_img, (0, 0), fx=original_resolution[1]/restored_img.shape[1], fy=original_resolution[0]/restored_img.shape[0], interpolation=cv2.INTER_LINEAR)
+
+                self.face_helper.clean_all()
+                video.append_data(restored_img)# 写入视频
+
         fileAudio = f"{filename}_{counter:05}_temp.mp3"
         audioFilePathName=os.path.join(full_output_folder, fileAudio)
-        aff = FFmpeg(inputs={templateVideoPath: None},outputs={audioFilePathName: '-f {} -vn'.format('mp3')})
-        aff.run()
-        _codec = 'copy'
-        vff = FFmpeg(
-            inputs={videoPathTemp: None, audioFilePathName: None},
-            outputs={videoPath: '-map 0:v -map 1:a -c:v copy -c:a {} -shortest'.format(_codec)})
-        vff.run()
+        #提取音频
+        cmd = r"ffmpeg -y -hide_banner -loglevel error -i %s %s"%(templateVideoPath, audioFilePathName)
+        os.system(cmd)  
+        #插入音频
+        cmd = r'ffmpeg -y -hide_banner -loglevel error -i "%s" -i "%s" -vcodec copy "%s"' % (videoPathTemp, audioFilePathName, videoPath)
+        os.system(cmd)  
         os.remove(videoPathTemp)
         os.remove(audioFilePathName)
         return {"ui": {"text": "视频保存成功文件地址："+videoPath,

py/VideoRoopFaceSwap.py

@@ -1,104 +0,0 @@
-import cv2
-import os
-import torch
-import numpy as np
-from .scripts.swapper import getFaceSwapModel,get_face_single
-from .third_part.GFPGAN.gfpgan import GFPGANer
-from PIL import Image
-import folder_paths
-from ffmpy import FFmpeg
-import imageio
-from tqdm import tqdm 
-
-class VideoRoopFaceSwap:
-    def __init__(self):
-        self.output_dir = os.path.join(folder_paths.get_output_directory(), 'video')
-        if not os.path.exists(self.output_dir):
-            os .makedirs(self.output_dir)
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required":
-                    {"templateVideoPath": ("STRING", {"forceInput": True}),
-                    "user_image": ("IMAGE", ),
-                    "filename_prefix": ("STRING", {"default": "comfyUI"})
-                    },
-                }
-
-    CATEGORY = "lam"
-
-    RETURN_TYPES = ("STRING",)
-    RETURN_NAMES = ("视频文件名",)
-    OUTPUT_NODE = True
-
-    FUNCTION = "face_fusion"
-    def face_fusion(self,templateVideoPath,user_image,filename_prefix):
-        face_swapper = getFaceSwapModel(folder_paths.models_dir+'/roop-face-swap/inswapper_128.onnx')
-        restorer_model = GFPGANer(model_path=os.path.join(folder_paths.models_dir+'/upscale_models','GFPGANv1.4.pth'), upscale=1, arch='clean',
-                                channel_multiplier=2, bg_upsampler=None)
-        imageu_np=user_image.numpy()
-        imau_shape=imageu_np.shape
-
-        if imau_shape[0]<=0:
-            return {"ui": {"text": "替换人脸照片不能为空"}, "result": ("",)}
-
-        imageu=imageu_np[0]*255
-        imageu=np.uint8(imageu)
-        cap = cv2.VideoCapture(templateVideoPath)
-        fps = cap.get(cv2.CAP_PROP_FPS)
-        length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-        flag = cap.isOpened()
-        if not flag:
-            return {"ui": {"text": "视频打开失败"}, "result": ("",)}
-        image_list=[]
-        source_face = get_face_single(imageu, face_index=0)
-        target_face=None
-        for frame_idx in tqdm(range(length), '换脸进度:'):
-        #while True:
-            flag, frame_c = cap.read()
-            if not flag:  # 如果已经读取到最后一帧则退出
-                break
-            frame = cv2.cvtColor(frame_c, cv2.COLOR_BGR2RGB) #bgr转rgb
-            if target_face is None:
-                target_face = get_face_single(frame, face_index=0)
-            if target_face is not None and source_face is not None:
-                image = face_swapper.get(frame, target_face, source_face)
-            else:
-                image = frame
-            #image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) #bgr转rgb
-            image_list.append(image)
-            
-        full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir, imau_shape[1], imau_shape[0])
-        fileTemp = f"{filename}_{counter:05}_temp.mp4"
-        videoPathTemp = os.path.join(full_output_folder, fileTemp)
-        file = f"{filename}_{counter:05}_.mp4"
-        videoPath = os.path.join(full_output_folder, file)
-        with imageio.get_writer(videoPathTemp, fps=fps) as video:
-            for i in tqdm(range(len(image_list)), '写入进度:'):
-                image=image_list[i]
-                _,_,rimage=restorer_model.enhance(image)
-                if rimage is not None:
-                    image=rimage
-                video.append_data(image)# 写入视频
-        fileAudio = f"{filename}_{counter:05}_temp.mp3"
-        audioFilePathName=os.path.join(full_output_folder, fileAudio)
-        aff = FFmpeg(inputs={templateVideoPath: None},outputs={audioFilePathName: '-f {} -vn'.format('mp3')})
-        aff.run()
-        _codec = 'copy'
-        vff = FFmpeg(
-            inputs={videoPathTemp: None, audioFilePathName: None},
-            outputs={videoPath: '-map 0:v -map 1:a -c:v copy -c:a {} -shortest'.format(_codec)})
-        vff.run()
-        os.remove(videoPathTemp)
-        os.remove(audioFilePathName)
-        return {"ui": {"text": "视频保存成功文件地址："+videoPath,
-        'videos':[{'filename':file,'type':'output','subfolder':'video'}]}, "result": (videoPath,)}
-
-NODE_CLASS_MAPPINGS = {
-    "VideoRoopFaceSwap": VideoRoopFaceSwap
-}
-
-# A dictionary that contains the friendly/humanly readable titles for the nodes
-NODE_DISPLAY_NAME_MAPPINGS = {
-    "VideoRoopFaceSwap": "视频脸部替换"
-}

py/WeChatAuth.py

@@ -14,8 +14,147 @@ import execution
 import uuid
 import folder_paths
 import traceback
+from comfy.cli_args import args
+from .src.wechat.redisSub import PubSub,run_with_reconnect,r
+from threading import Thread, current_thread
+from typing import List, Literal, NamedTuple, Optional
+import copy
+import asyncio
+ 
+@run_with_reconnect
+def subscribe(p):
+    while True:
+        msg = p.parse_response()
+        if msg[0]=='message' and len(msg)==3:
+            try:
+                message=json.loads(msg[2])
+                if message['event']=='addTask':
+                    ckptSetCount(message)
+                    prompt(PromptServer.instance,message['data'])
+                elif message['event']=='taskDone':
+                    task_done(PromptServer.instance.prompt_queue,message['item_id'],message['data'])
+                elif message['event']=='sendImage':
+                    base64_to_file(message['data'],message['filename'],message['type'],message['subfolder'])
+                elif message['event']=='dowFile':
+                    filedata=r.get(message['filename'])
+                    base64_to_file(filedata,message['filename'],message['type'],message['subfolder'])
+                    r.delete(message['filename'])
+                elif str(message['event'])=='2':
+                    filedata = Image.open(BytesIO(base64_to_b64decode(message['data'][1])))
+                    data=tuple([message['data'][0],filedata,message['data'][2]])
+                    send_sync(PromptServer.instance,message['event'],data,sid=message['sid'],port=message['port'])
+                else:
+                    send_sync(PromptServer.instance,message['event'],message['data'],sid=message['sid'],port=message['port'])
+            except Exception as e:
+                print('订阅消息处理异常',e)
+    
+@run_with_reconnect
+def ckptSetCount(message):
+    if 'ckptName' in message and message['ckptName']:
+        val=r.get('ckpt:'+Config().redis['basePath']+':'+message['ckptName'])
+        if val==None:
+            val=3
+        else:
+            val=int(val)+1
+        r.set('ckpt:'+Config().redis['basePath']+':'+message['ckptName'],val)
+        keys=r.keys('ckpt:'+Config().redis['basePath']+':*')
+        for key in keys:
+            val=int(r.get(key))
+            if val<=1:
+                r.delete(key)
+            else:
+                r.set(key,val-1)
+                
+@run_with_reconnect
+def sendPublish(channel,data):
+    if Config().redis['basePath'] == channel:
+        jsondata=json.loads(data)
+        if jsondata['event']=='addTask':
+            ckptSetCount(jsondata)
+            prompt(PromptServer.instance,jsondata['data'])
+            return 
+    r.publish(channel,data)
+
+@run_with_reconnect
+def refresh_heartbeat():
+    print('----心跳线程-----')
+    while True:
+        val=r.get('heartbeat:'+Config().redis['basePath'])
+        if val==None:
+            val=0
+        r.setex('heartbeat:'+Config().redis['basePath'], 3, val)
+        time.sleep(2)
+
+
+@run_with_reconnect
+def send_sync(self, event, data, sid=None,port=None): #继承父类的send_sync方法
+    if hasattr(self,"pub")==False and r: #添加订阅消息
+        setattr(self,"pub",PubSub().subscribe(Config().redis['basePath']))
+        if Config().redis['isMain']:
+            r.set('mainPath',Config().redis['basePath'])
+        keys=r.keys('ckpt:'+Config().redis['basePath']+':*')
+        for key in keys:
+            r.delete(key)
+        Thread(target=refresh_heartbeat,daemon=True, args=()).start()
+        Thread(target=subscribe,daemon=True, args=(self.pub,)).start()
+        if self.prompt_queue:
+            self.prompt_queue.task_done=types.MethodType(task_done,self.prompt_queue)
+    
+
+    if r :
+        if Config().redis['isMain']==False and event not in ['crystools.monitor']:
+            mainPath=r.get('mainPath')
+            if mainPath:
+                if event=='status':
+                    val=data['status']['exec_info']['queue_remaining']
+                    r.setex('heartbeat:'+Config().redis['basePath'], 3, val)
+                if event=='executed':
+                    if 'images' in data['output']:
+                        for img in data['output']['images']:
+                            # imgStr=image_to_base64(img['filename'],img['type'],img['subfolder'])
+                            # filename=base64_encode(Config().redis['basePath'])+img['filename']
+                            # msg={'event':'sendImage','port':Config().redis['basePath'],'filename':filename,'data':imgStr
+                            #     ,'type':img['type'],'subfolder':img['subfolder']}
+                            # sendPublish(mainPath, json.dumps(msg))
+                            # img['filename']=filename
+                            filename=base64_encode(Config().redis['basePath'])+img['filename']
+                            filedata=file_to_base64(img['filename'],img['type'],img['subfolder'])
+                            r.set(filename,filedata)
+                            msg={'event':'dowFile','port':Config().redis['basePath'],'filename':filename,'type':img['type'],'subfolder':img['subfolder']}
+                            sendPublish(mainPath, json.dumps(msg))
+                            img['filename']=filename
+                    elif 'gifs' in data['output']:
+                        for img in data['output']['gifs']:
+                            filename=base64_encode(Config().redis['basePath'])+img['filename']
+                            filedata=file_to_base64(img['filename'],img['type'],img['subfolder'])
+                            r.set(filename,filedata)
+                            msg={'event':'dowFile','port':Config().redis['basePath'],'filename':filename,'type':img['type'],'subfolder':img['subfolder']}
+                            sendPublish(mainPath, json.dumps(msg))
+                            img['filename']=filename
+                elif str(event)=='2':
+                    # 创建一个BytesIO对象，用于临时存储图像数据
+                    image_data = io.BytesIO()
+                    # 将图像保存到BytesIO对象中，格式为JPEG
+                    data[1].save(image_data, format='JPEG')
+                    filedata=base64_to_b64encode(image_data.getvalue())
+                    datalist=[data[0],filedata,data[2]]
+                    data=datalist
+                msg={'event':event,'port':Config().redis['basePath'],'data':data,'sid':sid}
+                sendPublish(mainPath, json.dumps(msg))
+                return 
+        elif event=='status':
+            if port==None:
+                val=data['status']['exec_info']['queue_remaining']
+                r.setex('heartbeat:'+Config().redis['basePath'], 3, val)
+            keys = r.keys('heartbeat:*')
+            queue_remaining=0
+            for key in keys:
+                val=r.get(key)
+                if val:
+                    queue_remaining+=int(val)
+
+            data['status']['exec_info']['queue_remaining']=queue_remaining
 
-def send_sync(self, event, data, sid=None): #继承父类的send_sync方法
     if hasattr(self,"clientObjPromptId")==False:
         setattr(self,"clientObjPromptId",{})
     if event=='execution_start':
@@ -92,6 +231,36 @@ def send_sync(self, event, data, sid=None): #继承父类的send_sync方法
     self.loop.call_soon_threadsafe(
         self.messages.put_nowait, (event, data, sid))
     
+MAXIMUM_HISTORY_SIZE = 10000
+def task_done(self, item_id, outputs={},status: Optional['PromptQueue.ExecutionStatus']=None):
+    print('----------task_done--------')
+    if status==None:
+        self.history[item_id] = outputs
+        self.server.send_sync("executing", { "node": None, "prompt_id": item_id }, self.server.client_id)
+        return
+    
+    with self.mutex:
+        prompt = self.currently_running.pop(item_id)
+        if len(self.history) > MAXIMUM_HISTORY_SIZE:
+            self.history.pop(next(iter(self.history)))
+
+        status_dict: Optional[dict] = None
+        if status is not None:
+            status_dict = copy.deepcopy(status._asdict())
+
+        data={
+            "prompt": prompt,
+            "outputs": copy.deepcopy(outputs),
+            'status': status_dict,
+        }
+        if r and Config().redis['isMain']==False :
+            mainPath=r.get('mainPath')
+            if mainPath:
+                sendPublish(mainPath, json.dumps({'event':'taskDone','data':data,'item_id':prompt[1]}))
+        
+        self.history[prompt[1]] = data
+        self.server.queue_updated()
+    
 def update_dict(dictionary, keys, value):
     if len(keys) == 1:
         dictionary[keys[0]] = value
@@ -104,6 +273,18 @@ def update_dict(dictionary, keys, value):
                 update_dict(item, keys[:], value)
         else:
             raise ValueError("Invalid data structure")
+        
+def getCkptName(prompt):
+    name=None
+    for k,v in prompt.items():
+        for k1,v1 in v['inputs'].items():
+            if type(v1)==str and k1.startswith('ckpt_name'):
+                name=v1
+                break
+        if name!=None:
+            break
+    return name
+
 def setPost(self,FromUserName):
     self.user_command[FromUserName]['status']='waiting' #prepare:准备 waiting:待执行  wcomplete完成
     params=self.user_command[FromUserName]
@@ -122,6 +303,18 @@ def setPost(self,FromUserName):
             update_dict(json_data,keys,params[key])
     
     json_data={"prompt":json_data,"client_id":params['openId']}
+    now = time.localtime()
+    start_time = time.strftime("%Y-%m-%d %H:%M:%S", now)
+    if r and Config().redis['isMain']:
+        prompt_id=str(uuid.uuid4())
+        self.user_command[FromUserName]['prompt_id']=prompt_id 
+        db=DataBaseUtil()
+        db.insert_data( params['openId'], json.dumps(params), prompt_id,'waiting',start_time, '', '')
+        db.close_con()
+        data=selServer(json_data,prompt_id)
+        if data:
+            return data
+                    
     json_data = self.trigger_on_prompt(json_data)
     if "number" in json_data:
         number = float(json_data['number'])
@@ -146,8 +339,6 @@ def setPost(self,FromUserName):
             prompt_id = str(uuid.uuid4())
             outputs_to_execute = valid[2]
             self.prompt_queue.put((number, prompt_id, prompt, extra_data, outputs_to_execute))
-            now = time.localtime()
-            start_time = time.strftime("%Y-%m-%d %H:%M:%S", now)
             self.user_command[FromUserName]['prompt_id']=prompt_id 
             db=DataBaseUtil()
             db.insert_data( params['openId'], json.dumps(params), prompt_id,'waiting',start_time, '', '')
@@ -158,6 +349,117 @@ def setPost(self,FromUserName):
             logging.warning("invalid prompt: {}".format(valid[1]))
             return None
         
+@run_with_reconnect
+def selServer(json_data,prompt_id):
+    json_data['prompt_id']=prompt_id
+    name=None
+    if Config().redis['modelPriority']==True :
+        name=getCkptName(json_data['prompt'])
+        if name:
+            ckkeys=r.keys('ckpt:*:'+name)
+            nport=None
+            for ckkey in ckkeys:
+                ns=ckkey.split(":")
+                nport=':'.join(ns[1:3])
+                break
+            if nport:
+                nval = r.get('heartbeat:'+nport)
+                if nval!=None:
+                    sendPublish(nport, json.dumps({'event':'addTask','data':json_data,'ckptName':name}))
+                    return {"prompt_id": prompt_id, "number": 1, "node_errors": []}
+                
+    keys=r.keys('heartbeat:*')
+    if len(keys)>1:
+        nameSize={}
+        for key in keys:
+            val=r.get(key)
+            if val!=None :
+                ns=key.split(":")
+                print(':'.join(ns[1:]))
+                ckkeys=r.keys('ckpt:'+':'.join(ns[1:])+':*')
+                if int(val)==0 and len(ckkeys)==0:
+                    sendPublish(':'.join(ns[1:]), json.dumps({'event':'addTask','data':json_data,'ckptName':name}))
+                    return {"prompt_id": prompt_id, "number": 1, "node_errors": []}
+                else:
+                    nameSize[':'.join(ns[1:])]=int(val)+len(ckkeys)
+        print('nameSize:',nameSize)
+        minKey=min(key for key, value in nameSize.items() if value == min(nameSize.values()))
+        sendPublish(minKey, json.dumps({'event':'addTask','data':json_data,'ckptName':name}))
+        return {"prompt_id": prompt_id, "number": 1, "node_errors": []}
+    return None
+def trigger_on_prompt(self,json_data,isRun=True):
+    if isRun and r and Config().redis['isMain']:
+        prompt_id=str(uuid.uuid4())
+        data=selServer(json_data,prompt_id)
+        if data:
+            return data
+        
+    for handler in self.on_prompt_handlers:
+        try:
+            json_data = handler(json_data)
+        except Exception as e:
+            logging.warning(f"[ERROR] An error occurred during the on_prompt_handler processing")
+            logging.warning(traceback.format_exc())
+
+    prompt=json_data['prompt']
+    maxKeyStr=sorted(list(prompt.keys()), key=lambda x: int(x.split(':')[0]))[-1]
+    maxKey = int(maxKeyStr.split(':')[0])
+    endNodeKeys=[x for x in prompt.keys() if prompt[x]['class_type']=='ForInnerEnd']
+    for unique_id in endNodeKeys:
+        inputNum=prompt[unique_id]['inputs']['obj']
+        maxKey=maxKey+1
+        json_data['prompt'][str(maxKey)]={ "inputs": { "expression": "p0",  "p0": inputNum },"class_type": "MultiParamFormula"}
+        json_data['prompt'][unique_id]['inputs']['obj']=[str(maxKey),0]
+
+    return json_data
+
+def prompt(self,json_data):
+    json_data = self.trigger_on_prompt(json_data,isRun=False)
+    if 'prompt_id' in json_data:
+        prompt_id=json_data['prompt_id']
+    else:
+        prompt_id = str(uuid.uuid4())
+    try:
+        if "number" in json_data:
+            number = float(json_data['number'])
+        else:
+            number = self.number
+            if "front" in json_data:
+                if json_data['front']:
+                    number = -number
+
+            self.number += 1
+
+        if "prompt" in json_data:
+            prompt = json_data["prompt"]
+            valid = execution.validate_prompt(prompt)
+            extra_data = {}
+            if "extra_data" in json_data:
+                extra_data = json_data["extra_data"]
+
+            if "client_id" in json_data:
+                extra_data["client_id"] = json_data["client_id"]
+            if valid[0]:
+                if 'prompt_id' in json_data:
+                    prompt_id=json_data['prompt_id']
+                else:
+                    prompt_id = str(uuid.uuid4())
+                outputs_to_execute = valid[2]
+                self.prompt_queue.put((number, prompt_id, prompt, extra_data, outputs_to_execute))
+                response = {"prompt_id": prompt_id, "number": number, "node_errors": valid[3]}
+                return response
+            else:
+                logging.warning("invalid prompt: {}".format(valid[1]))
+                return {"error": valid[1], "node_errors": valid[3]}
+
+        else:
+            return {"error": "no prompt", "node_errors": []}
+    except Exception as e:
+        print('prompt处理异常',e)
+        node_errors= ['prompt处理异常']
+        send_sync(self,'execution_error', {"prompt_id": prompt_id,'node_errors':node_errors}, sid=self.client_id)
+    
+        
 def getTaskRanking(self,FromUserName):
     if 'prompt_id' not in self.user_command[FromUserName]:
         return '您还没有排队，请先发送您的指令'
@@ -176,8 +478,12 @@ def getTaskRanking(self,FromUserName):
 async def getHistorys(request):
     if "openId" in request.rel_url.query:
         openId=request.rel_url.query['openId']
+        page_number=1
+        if 'page_number' in request.rel_url.query:
+            page_number=int(request.rel_url.query['page_number'])
+
         db=DataBaseUtil()
-        datas=db.get_many_data(openId)
+        datas=db.get_many_data(openId,page_number=page_number)
         db.close_con()
         rdataList=[]
         if "type" in request.rel_url.query:
@@ -491,6 +797,8 @@ async def handleMessagePost(request):
     return web.Response(status=200)
 
 PromptServer.instance.send_sync=types.MethodType(send_sync,PromptServer.instance)
+PromptServer.instance.trigger_on_prompt=types.MethodType(trigger_on_prompt,PromptServer.instance)
+
 
 NODE_CLASS_MAPPINGS = {}
 NODE_DISPLAY_NAME_MAPPINGS = {}

py/face_fusion/peppa_pig_face/face_landmark.py

@@ -82,7 +82,7 @@ class FaceLandmark:
         bbox[2] = center[0] + one_edge // 2
         bbox[3] = center[1] + one_edge // 2
 
-        bbox = bbox.astype(np.int)
+        bbox = bbox.astype(np.int32)
         crop_image = bimg[bbox[1]:bbox[3], bbox[0]:bbox[2], :]
         h, w, _ = crop_image.shape
         crop_image = cv2.resize(crop_image,

py/facelib/detection/__init__.py

@@ -0,0 +1,100 @@
+import os
+import torch
+from torch import nn
+from copy import deepcopy
+
+from ..utils import load_file_from_url
+from ..utils import download_pretrained_models
+from .yolov5face.models.common import Conv
+
+from .retinaface.retinaface import RetinaFace
+from .yolov5face.face_detector import YoloDetector
+
+
+def init_detection_model(model_name, half=False, device='cuda'):
+    if 'retinaface' in model_name:
+        model = init_retinaface_model(model_name, half, device)
+    elif 'YOLOv5' in model_name:
+        model = init_yolov5face_model(model_name, device)
+    else:
+        raise NotImplementedError(f'{model_name} is not implemented.')
+
+    return model
+
+
+def init_retinaface_model(model_name, half=False, device='cuda'):
+    if model_name == 'retinaface_resnet50':
+        model = RetinaFace(network_name='resnet50', half=half)
+        model_url = 'https://github.com/xinntao/facexlib/releases/download/v0.1.0/detection_Resnet50_Final.pth'
+    elif model_name == 'retinaface_mobile0.25':
+        model = RetinaFace(network_name='mobile0.25', half=half)
+        model_url = 'https://github.com/xinntao/facexlib/releases/download/v0.1.0/detection_mobilenet0.25_Final.pth'
+    else:
+        raise NotImplementedError(f'{model_name} is not implemented.')
+
+    model_path = load_file_from_url(url=model_url, model_dir='../../models/facedetection', progress=True, file_name=None)
+    load_net = torch.load(model_path, map_location=lambda storage, loc: storage)
+    # remove unnecessary 'module.'
+    for k, v in deepcopy(load_net).items():
+        if k.startswith('module.'):
+            load_net[k[7:]] = v
+            load_net.pop(k)
+    model.load_state_dict(load_net, strict=True)
+    model.eval()
+    model = model.to(device)
+
+    return model
+
+
+def init_yolov5face_model(model_name, device='cuda'):
+    if model_name == 'YOLOv5l':
+        model = YoloDetector(config_name='custom_nodes/facerestore/facelib/detection/yolov5face/models/yolov5l.yaml', device=device)
+        model_url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/yolov5l-face.pth'
+    elif model_name == 'YOLOv5n':
+        model = YoloDetector(config_name='custom_nodes/facerestore/facelib/detection/yolov5face/models/yolov5n.yaml', device=device)
+        model_url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/yolov5n-face.pth'
+    else:
+        raise NotImplementedError(f'{model_name} is not implemented.')
+    
+    model_path = load_file_from_url(url=model_url, model_dir='../../models/facedetection', progress=True, file_name=None)
+    load_net = torch.load(model_path, map_location=lambda storage, loc: storage)
+    model.detector.load_state_dict(load_net, strict=True)
+    model.detector.eval()
+    model.detector = model.detector.to(device).float()
+
+    for m in model.detector.modules():
+        if type(m) in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU]:
+            m.inplace = True  # pytorch 1.7.0 compatibility
+        elif isinstance(m, Conv):
+            m._non_persistent_buffers_set = set()  # pytorch 1.6.0 compatibility
+
+    return model
+
+
+# Download from Google Drive
+# def init_yolov5face_model(model_name, device='cuda'):
+#     if model_name == 'YOLOv5l':
+#         model = YoloDetector(config_name='facelib/detection/yolov5face/models/yolov5l.yaml', device=device)
+#         f_id = {'yolov5l-face.pth': '131578zMA6B2x8VQHyHfa6GEPtulMCNzV'}
+#     elif model_name == 'YOLOv5n':
+#         model = YoloDetector(config_name='facelib/detection/yolov5face/models/yolov5n.yaml', device=device)
+#         f_id = {'yolov5n-face.pth': '1fhcpFvWZqghpGXjYPIne2sw1Fy4yhw6o'}
+#     else:
+#         raise NotImplementedError(f'{model_name} is not implemented.')
+
+#     model_path = os.path.join('../../models/facedetection', list(f_id.keys())[0])
+#     if not os.path.exists(model_path):
+#         download_pretrained_models(file_ids=f_id, save_path_root='../../models/facedetection')
+
+#     load_net = torch.load(model_path, map_location=lambda storage, loc: storage)
+#     model.detector.load_state_dict(load_net, strict=True)
+#     model.detector.eval()
+#     model.detector = model.detector.to(device).float()
+
+#     for m in model.detector.modules():
+#         if type(m) in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU]:
+#             m.inplace = True  # pytorch 1.7.0 compatibility
+#         elif isinstance(m, Conv):
+#             m._non_persistent_buffers_set = set()  # pytorch 1.6.0 compatibility
+
+#     return model

py/facelib/detection/align_trans.py

@@ -0,0 +1,219 @@
+import cv2
+import numpy as np
+
+from .matlab_cp2tform import get_similarity_transform_for_cv2
+
+# reference facial points, a list of coordinates (x,y)
+REFERENCE_FACIAL_POINTS = [[30.29459953, 51.69630051], [65.53179932, 51.50139999], [48.02519989, 71.73660278],
+                           [33.54930115, 92.3655014], [62.72990036, 92.20410156]]
+
+DEFAULT_CROP_SIZE = (96, 112)
+
+
+class FaceWarpException(Exception):
+
+    def __str__(self):
+        return 'In File {}:{}'.format(__file__, super.__str__(self))
+
+
+def get_reference_facial_points(output_size=None, inner_padding_factor=0.0, outer_padding=(0, 0), default_square=False):
+    """
+    Function:
+    ----------
+        get reference 5 key points according to crop settings:
+        0. Set default crop_size:
+            if default_square:
+                crop_size = (112, 112)
+            else:
+                crop_size = (96, 112)
+        1. Pad the crop_size by inner_padding_factor in each side;
+        2. Resize crop_size into (output_size - outer_padding*2),
+            pad into output_size with outer_padding;
+        3. Output reference_5point;
+    Parameters:
+    ----------
+        @output_size: (w, h) or None
+            size of aligned face image
+        @inner_padding_factor: (w_factor, h_factor)
+            padding factor for inner (w, h)
+        @outer_padding: (w_pad, h_pad)
+            each row is a pair of coordinates (x, y)
+        @default_square: True or False
+            if True:
+                default crop_size = (112, 112)
+            else:
+                default crop_size = (96, 112);
+        !!! make sure, if output_size is not None:
+                (output_size - outer_padding)
+                = some_scale * (default crop_size * (1.0 +
+                inner_padding_factor))
+    Returns:
+    ----------
+        @reference_5point: 5x2 np.array
+            each row is a pair of transformed coordinates (x, y)
+    """
+
+    tmp_5pts = np.array(REFERENCE_FACIAL_POINTS)
+    tmp_crop_size = np.array(DEFAULT_CROP_SIZE)
+
+    # 0) make the inner region a square
+    if default_square:
+        size_diff = max(tmp_crop_size) - tmp_crop_size
+        tmp_5pts += size_diff / 2
+        tmp_crop_size += size_diff
+
+    if (output_size and output_size[0] == tmp_crop_size[0] and output_size[1] == tmp_crop_size[1]):
+
+        return tmp_5pts
+
+    if (inner_padding_factor == 0 and outer_padding == (0, 0)):
+        if output_size is None:
+            return tmp_5pts
+        else:
+            raise FaceWarpException('No paddings to do, output_size must be None or {}'.format(tmp_crop_size))
+
+    # check output size
+    if not (0 <= inner_padding_factor <= 1.0):
+        raise FaceWarpException('Not (0 <= inner_padding_factor <= 1.0)')
+
+    if ((inner_padding_factor > 0 or outer_padding[0] > 0 or outer_padding[1] > 0) and output_size is None):
+        output_size = tmp_crop_size * \
+            (1 + inner_padding_factor * 2).astype(np.int32)
+        output_size += np.array(outer_padding)
+    if not (outer_padding[0] < output_size[0] and outer_padding[1] < output_size[1]):
+        raise FaceWarpException('Not (outer_padding[0] < output_size[0] and outer_padding[1] < output_size[1])')
+
+    # 1) pad the inner region according inner_padding_factor
+    if inner_padding_factor > 0:
+        size_diff = tmp_crop_size * inner_padding_factor * 2
+        tmp_5pts += size_diff / 2
+        tmp_crop_size += np.round(size_diff).astype(np.int32)
+
+    # 2) resize the padded inner region
+    size_bf_outer_pad = np.array(output_size) - np.array(outer_padding) * 2
+
+    if size_bf_outer_pad[0] * tmp_crop_size[1] != size_bf_outer_pad[1] * tmp_crop_size[0]:
+        raise FaceWarpException('Must have (output_size - outer_padding)'
+                                '= some_scale * (crop_size * (1.0 + inner_padding_factor)')
+
+    scale_factor = size_bf_outer_pad[0].astype(np.float32) / tmp_crop_size[0]
+    tmp_5pts = tmp_5pts * scale_factor
+    #    size_diff = tmp_crop_size * (scale_factor - min(scale_factor))
+    #    tmp_5pts = tmp_5pts + size_diff / 2
+    tmp_crop_size = size_bf_outer_pad
+
+    # 3) add outer_padding to make output_size
+    reference_5point = tmp_5pts + np.array(outer_padding)
+    tmp_crop_size = output_size
+
+    return reference_5point
+
+
+def get_affine_transform_matrix(src_pts, dst_pts):
+    """
+    Function:
+    ----------
+        get affine transform matrix 'tfm' from src_pts to dst_pts
+    Parameters:
+    ----------
+        @src_pts: Kx2 np.array
+            source points matrix, each row is a pair of coordinates (x, y)
+        @dst_pts: Kx2 np.array
+            destination points matrix, each row is a pair of coordinates (x, y)
+    Returns:
+    ----------
+        @tfm: 2x3 np.array
+            transform matrix from src_pts to dst_pts
+    """
+
+    tfm = np.float32([[1, 0, 0], [0, 1, 0]])
+    n_pts = src_pts.shape[0]
+    ones = np.ones((n_pts, 1), src_pts.dtype)
+    src_pts_ = np.hstack([src_pts, ones])
+    dst_pts_ = np.hstack([dst_pts, ones])
+
+    A, res, rank, s = np.linalg.lstsq(src_pts_, dst_pts_)
+
+    if rank == 3:
+        tfm = np.float32([[A[0, 0], A[1, 0], A[2, 0]], [A[0, 1], A[1, 1], A[2, 1]]])
+    elif rank == 2:
+        tfm = np.float32([[A[0, 0], A[1, 0], 0], [A[0, 1], A[1, 1], 0]])
+
+    return tfm
+
+
+def warp_and_crop_face(src_img, facial_pts, reference_pts=None, crop_size=(96, 112), align_type='smilarity'):
+    """
+    Function:
+    ----------
+        apply affine transform 'trans' to uv
+    Parameters:
+    ----------
+        @src_img: 3x3 np.array
+            input image
+        @facial_pts: could be
+            1)a list of K coordinates (x,y)
+        or
+            2) Kx2 or 2xK np.array
+            each row or col is a pair of coordinates (x, y)
+        @reference_pts: could be
+            1) a list of K coordinates (x,y)
+        or
+            2) Kx2 or 2xK np.array
+            each row or col is a pair of coordinates (x, y)
+        or
+            3) None
+            if None, use default reference facial points
+        @crop_size: (w, h)
+            output face image size
+        @align_type: transform type, could be one of
+            1) 'similarity': use similarity transform
+            2) 'cv2_affine': use the first 3 points to do affine transform,
+                    by calling cv2.getAffineTransform()
+            3) 'affine': use all points to do affine transform
+    Returns:
+    ----------
+        @face_img: output face image with size (w, h) = @crop_size
+    """
+
+    if reference_pts is None:
+        if crop_size[0] == 96 and crop_size[1] == 112:
+            reference_pts = REFERENCE_FACIAL_POINTS
+        else:
+            default_square = False
+            inner_padding_factor = 0
+            outer_padding = (0, 0)
+            output_size = crop_size
+
+            reference_pts = get_reference_facial_points(output_size, inner_padding_factor, outer_padding,
+                                                        default_square)
+
+    ref_pts = np.float32(reference_pts)
+    ref_pts_shp = ref_pts.shape
+    if max(ref_pts_shp) < 3 or min(ref_pts_shp) != 2:
+        raise FaceWarpException('reference_pts.shape must be (K,2) or (2,K) and K>2')
+
+    if ref_pts_shp[0] == 2:
+        ref_pts = ref_pts.T
+
+    src_pts = np.float32(facial_pts)
+    src_pts_shp = src_pts.shape
+    if max(src_pts_shp) < 3 or min(src_pts_shp) != 2:
+        raise FaceWarpException('facial_pts.shape must be (K,2) or (2,K) and K>2')
+
+    if src_pts_shp[0] == 2:
+        src_pts = src_pts.T
+
+    if src_pts.shape != ref_pts.shape:
+        raise FaceWarpException('facial_pts and reference_pts must have the same shape')
+
+    if align_type == 'cv2_affine':
+        tfm = cv2.getAffineTransform(src_pts[0:3], ref_pts[0:3])
+    elif align_type == 'affine':
+        tfm = get_affine_transform_matrix(src_pts, ref_pts)
+    else:
+        tfm = get_similarity_transform_for_cv2(src_pts, ref_pts)
+
+    face_img = cv2.warpAffine(src_img, tfm, (crop_size[0], crop_size[1]))
+
+    return face_img

py/facelib/detection/matlab_cp2tform.py

@@ -0,0 +1,317 @@
+import numpy as np
+from numpy.linalg import inv, lstsq
+from numpy.linalg import matrix_rank as rank
+from numpy.linalg import norm
+
+
+class MatlabCp2tormException(Exception):
+
+    def __str__(self):
+        return 'In File {}:{}'.format(__file__, super.__str__(self))
+
+
+def tformfwd(trans, uv):
+    """
+    Function:
+    ----------
+        apply affine transform 'trans' to uv
+
+    Parameters:
+    ----------
+        @trans: 3x3 np.array
+            transform matrix
+        @uv: Kx2 np.array
+            each row is a pair of coordinates (x, y)
+
+    Returns:
+    ----------
+        @xy: Kx2 np.array
+            each row is a pair of transformed coordinates (x, y)
+    """
+    uv = np.hstack((uv, np.ones((uv.shape[0], 1))))
+    xy = np.dot(uv, trans)
+    xy = xy[:, 0:-1]
+    return xy
+
+
+def tforminv(trans, uv):
+    """
+    Function:
+    ----------
+        apply the inverse of affine transform 'trans' to uv
+
+    Parameters:
+    ----------
+        @trans: 3x3 np.array
+            transform matrix
+        @uv: Kx2 np.array
+            each row is a pair of coordinates (x, y)
+
+    Returns:
+    ----------
+        @xy: Kx2 np.array
+            each row is a pair of inverse-transformed coordinates (x, y)
+    """
+    Tinv = inv(trans)
+    xy = tformfwd(Tinv, uv)
+    return xy
+
+
+def findNonreflectiveSimilarity(uv, xy, options=None):
+    options = {'K': 2}
+
+    K = options['K']
+    M = xy.shape[0]
+    x = xy[:, 0].reshape((-1, 1))  # use reshape to keep a column vector
+    y = xy[:, 1].reshape((-1, 1))  # use reshape to keep a column vector
+
+    tmp1 = np.hstack((x, y, np.ones((M, 1)), np.zeros((M, 1))))
+    tmp2 = np.hstack((y, -x, np.zeros((M, 1)), np.ones((M, 1))))
+    X = np.vstack((tmp1, tmp2))
+
+    u = uv[:, 0].reshape((-1, 1))  # use reshape to keep a column vector
+    v = uv[:, 1].reshape((-1, 1))  # use reshape to keep a column vector
+    U = np.vstack((u, v))
+
+    # We know that X * r = U
+    if rank(X) >= 2 * K:
+        r, _, _, _ = lstsq(X, U, rcond=-1)
+        r = np.squeeze(r)
+    else:
+        raise Exception('cp2tform:twoUniquePointsReq')
+    sc = r[0]
+    ss = r[1]
+    tx = r[2]
+    ty = r[3]
+
+    Tinv = np.array([[sc, -ss, 0], [ss, sc, 0], [tx, ty, 1]])
+    T = inv(Tinv)
+    T[:, 2] = np.array([0, 0, 1])
+
+    return T, Tinv
+
+
+def findSimilarity(uv, xy, options=None):
+    options = {'K': 2}
+
+    #    uv = np.array(uv)
+    #    xy = np.array(xy)
+
+    # Solve for trans1
+    trans1, trans1_inv = findNonreflectiveSimilarity(uv, xy, options)
+
+    # Solve for trans2
+
+    # manually reflect the xy data across the Y-axis
+    xyR = xy
+    xyR[:, 0] = -1 * xyR[:, 0]
+
+    trans2r, trans2r_inv = findNonreflectiveSimilarity(uv, xyR, options)
+
+    # manually reflect the tform to undo the reflection done on xyR
+    TreflectY = np.array([[-1, 0, 0], [0, 1, 0], [0, 0, 1]])
+
+    trans2 = np.dot(trans2r, TreflectY)
+
+    # Figure out if trans1 or trans2 is better
+    xy1 = tformfwd(trans1, uv)
+    norm1 = norm(xy1 - xy)
+
+    xy2 = tformfwd(trans2, uv)
+    norm2 = norm(xy2 - xy)
+
+    if norm1 <= norm2:
+        return trans1, trans1_inv
+    else:
+        trans2_inv = inv(trans2)
+        return trans2, trans2_inv
+
+
+def get_similarity_transform(src_pts, dst_pts, reflective=True):
+    """
+    Function:
+    ----------
+        Find Similarity Transform Matrix 'trans':
+            u = src_pts[:, 0]
+            v = src_pts[:, 1]
+            x = dst_pts[:, 0]
+            y = dst_pts[:, 1]
+            [x, y, 1] = [u, v, 1] * trans
+
+    Parameters:
+    ----------
+        @src_pts: Kx2 np.array
+            source points, each row is a pair of coordinates (x, y)
+        @dst_pts: Kx2 np.array
+            destination points, each row is a pair of transformed
+            coordinates (x, y)
+        @reflective: True or False
+            if True:
+                use reflective similarity transform
+            else:
+                use non-reflective similarity transform
+
+    Returns:
+    ----------
+       @trans: 3x3 np.array
+            transform matrix from uv to xy
+        trans_inv: 3x3 np.array
+            inverse of trans, transform matrix from xy to uv
+    """
+
+    if reflective:
+        trans, trans_inv = findSimilarity(src_pts, dst_pts)
+    else:
+        trans, trans_inv = findNonreflectiveSimilarity(src_pts, dst_pts)
+
+    return trans, trans_inv
+
+
+def cvt_tform_mat_for_cv2(trans):
+    """
+    Function:
+    ----------
+        Convert Transform Matrix 'trans' into 'cv2_trans' which could be
+        directly used by cv2.warpAffine():
+            u = src_pts[:, 0]
+            v = src_pts[:, 1]
+            x = dst_pts[:, 0]
+            y = dst_pts[:, 1]
+            [x, y].T = cv_trans * [u, v, 1].T
+
+    Parameters:
+    ----------
+        @trans: 3x3 np.array
+            transform matrix from uv to xy
+
+    Returns:
+    ----------
+        @cv2_trans: 2x3 np.array
+            transform matrix from src_pts to dst_pts, could be directly used
+            for cv2.warpAffine()
+    """
+    cv2_trans = trans[:, 0:2].T
+
+    return cv2_trans
+
+
+def get_similarity_transform_for_cv2(src_pts, dst_pts, reflective=True):
+    """
+    Function:
+    ----------
+        Find Similarity Transform Matrix 'cv2_trans' which could be
+        directly used by cv2.warpAffine():
+            u = src_pts[:, 0]
+            v = src_pts[:, 1]
+            x = dst_pts[:, 0]
+            y = dst_pts[:, 1]
+            [x, y].T = cv_trans * [u, v, 1].T
+
+    Parameters:
+    ----------
+        @src_pts: Kx2 np.array
+            source points, each row is a pair of coordinates (x, y)
+        @dst_pts: Kx2 np.array
+            destination points, each row is a pair of transformed
+            coordinates (x, y)
+        reflective: True or False
+            if True:
+                use reflective similarity transform
+            else:
+                use non-reflective similarity transform
+
+    Returns:
+    ----------
+        @cv2_trans: 2x3 np.array
+            transform matrix from src_pts to dst_pts, could be directly used
+            for cv2.warpAffine()
+    """
+    trans, trans_inv = get_similarity_transform(src_pts, dst_pts, reflective)
+    cv2_trans = cvt_tform_mat_for_cv2(trans)
+
+    return cv2_trans
+
+
+if __name__ == '__main__':
+    """
+    u = [0, 6, -2]
+    v = [0, 3, 5]
+    x = [-1, 0, 4]
+    y = [-1, -10, 4]
+
+    # In Matlab, run:
+    #
+    #   uv = [u'; v'];
+    #   xy = [x'; y'];
+    #   tform_sim=cp2tform(uv,xy,'similarity');
+    #
+    #   trans = tform_sim.tdata.T
+    #   ans =
+    #       -0.0764   -1.6190         0
+    #        1.6190   -0.0764         0
+    #       -3.2156    0.0290    1.0000
+    #   trans_inv = tform_sim.tdata.Tinv
+    #    ans =
+    #
+    #       -0.0291    0.6163         0
+    #       -0.6163   -0.0291         0
+    #       -0.0756    1.9826    1.0000
+    #    xy_m=tformfwd(tform_sim, u,v)
+    #
+    #    xy_m =
+    #
+    #       -3.2156    0.0290
+    #        1.1833   -9.9143
+    #        5.0323    2.8853
+    #    uv_m=tforminv(tform_sim, x,y)
+    #
+    #    uv_m =
+    #
+    #        0.5698    1.3953
+    #        6.0872    2.2733
+    #       -2.6570    4.3314
+    """
+    u = [0, 6, -2]
+    v = [0, 3, 5]
+    x = [-1, 0, 4]
+    y = [-1, -10, 4]
+
+    uv = np.array((u, v)).T
+    xy = np.array((x, y)).T
+
+    print('\n--->uv:')
+    print(uv)
+    print('\n--->xy:')
+    print(xy)
+
+    trans, trans_inv = get_similarity_transform(uv, xy)
+
+    print('\n--->trans matrix:')
+    print(trans)
+
+    print('\n--->trans_inv matrix:')
+    print(trans_inv)
+
+    print('\n---> apply transform to uv')
+    print('\nxy_m = uv_augmented * trans')
+    uv_aug = np.hstack((uv, np.ones((uv.shape[0], 1))))
+    xy_m = np.dot(uv_aug, trans)
+    print(xy_m)
+
+    print('\nxy_m = tformfwd(trans, uv)')
+    xy_m = tformfwd(trans, uv)
+    print(xy_m)
+
+    print('\n---> apply inverse transform to xy')
+    print('\nuv_m = xy_augmented * trans_inv')
+    xy_aug = np.hstack((xy, np.ones((xy.shape[0], 1))))
+    uv_m = np.dot(xy_aug, trans_inv)
+    print(uv_m)
+
+    print('\nuv_m = tformfwd(trans_inv, xy)')
+    uv_m = tformfwd(trans_inv, xy)
+    print(uv_m)
+
+    uv_m = tforminv(trans, xy)
+    print('\nuv_m = tforminv(trans, xy)')
+    print(uv_m)

py/facelib/detection/retinaface/retinaface.py

@@ -0,0 +1,370 @@
+import cv2
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from PIL import Image
+from torchvision.models._utils import IntermediateLayerGetter as IntermediateLayerGetter
+
+from ..align_trans import get_reference_facial_points, warp_and_crop_face
+from .retinaface_net import FPN, SSH, MobileNetV1, make_bbox_head, make_class_head, make_landmark_head
+from .retinaface_utils import (PriorBox, batched_decode, batched_decode_landm, decode, decode_landm,
+                                                 py_cpu_nms)
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+
+def generate_config(network_name):
+
+    cfg_mnet = {
+        'name': 'mobilenet0.25',
+        'min_sizes': [[16, 32], [64, 128], [256, 512]],
+        'steps': [8, 16, 32],
+        'variance': [0.1, 0.2],
+        'clip': False,
+        'loc_weight': 2.0,
+        'gpu_train': True,
+        'batch_size': 32,
+        'ngpu': 1,
+        'epoch': 250,
+        'decay1': 190,
+        'decay2': 220,
+        'image_size': 640,
+        'return_layers': {
+            'stage1': 1,
+            'stage2': 2,
+            'stage3': 3
+        },
+        'in_channel': 32,
+        'out_channel': 64
+    }
+
+    cfg_re50 = {
+        'name': 'Resnet50',
+        'min_sizes': [[16, 32], [64, 128], [256, 512]],
+        'steps': [8, 16, 32],
+        'variance': [0.1, 0.2],
+        'clip': False,
+        'loc_weight': 2.0,
+        'gpu_train': True,
+        'batch_size': 24,
+        'ngpu': 4,
+        'epoch': 100,
+        'decay1': 70,
+        'decay2': 90,
+        'image_size': 840,
+        'return_layers': {
+            'layer2': 1,
+            'layer3': 2,
+            'layer4': 3
+        },
+        'in_channel': 256,
+        'out_channel': 256
+    }
+
+    if network_name == 'mobile0.25':
+        return cfg_mnet
+    elif network_name == 'resnet50':
+        return cfg_re50
+    else:
+        raise NotImplementedError(f'network_name={network_name}')
+
+
+class RetinaFace(nn.Module):
+
+    def __init__(self, network_name='resnet50', half=False, phase='test'):
+        super(RetinaFace, self).__init__()
+        self.half_inference = half
+        cfg = generate_config(network_name)
+        self.backbone = cfg['name']
+
+        self.model_name = f'retinaface_{network_name}'
+        self.cfg = cfg
+        self.phase = phase
+        self.target_size, self.max_size = 1600, 2150
+        self.resize, self.scale, self.scale1 = 1., None, None
+        self.mean_tensor = torch.tensor([[[[104.]], [[117.]], [[123.]]]]).to(device)
+        self.reference = get_reference_facial_points(default_square=True)
+        # Build network.
+        backbone = None
+        if cfg['name'] == 'mobilenet0.25':
+            backbone = MobileNetV1()
+            self.body = IntermediateLayerGetter(backbone, cfg['return_layers'])
+        elif cfg['name'] == 'Resnet50':
+            import torchvision.models as models
+            backbone = models.resnet50(pretrained=False)
+            self.body = IntermediateLayerGetter(backbone, cfg['return_layers'])
+
+        in_channels_stage2 = cfg['in_channel']
+        in_channels_list = [
+            in_channels_stage2 * 2,
+            in_channels_stage2 * 4,
+            in_channels_stage2 * 8,
+        ]
+
+        out_channels = cfg['out_channel']
+        self.fpn = FPN(in_channels_list, out_channels)
+        self.ssh1 = SSH(out_channels, out_channels)
+        self.ssh2 = SSH(out_channels, out_channels)
+        self.ssh3 = SSH(out_channels, out_channels)
+
+        self.ClassHead = make_class_head(fpn_num=3, inchannels=cfg['out_channel'])
+        self.BboxHead = make_bbox_head(fpn_num=3, inchannels=cfg['out_channel'])
+        self.LandmarkHead = make_landmark_head(fpn_num=3, inchannels=cfg['out_channel'])
+
+        self.to(device)
+        self.eval()
+        if self.half_inference:
+            self.half()
+
+    def forward(self, inputs):
+        out = self.body(inputs)
+
+        if self.backbone == 'mobilenet0.25' or self.backbone == 'Resnet50':
+            out = list(out.values())
+        # FPN
+        fpn = self.fpn(out)
+
+        # SSH
+        feature1 = self.ssh1(fpn[0])
+        feature2 = self.ssh2(fpn[1])
+        feature3 = self.ssh3(fpn[2])
+        features = [feature1, feature2, feature3]
+
+        bbox_regressions = torch.cat([self.BboxHead[i](feature) for i, feature in enumerate(features)], dim=1)
+        classifications = torch.cat([self.ClassHead[i](feature) for i, feature in enumerate(features)], dim=1)
+        tmp = [self.LandmarkHead[i](feature) for i, feature in enumerate(features)]
+        ldm_regressions = (torch.cat(tmp, dim=1))
+
+        if self.phase == 'train':
+            output = (bbox_regressions, classifications, ldm_regressions)
+        else:
+            output = (bbox_regressions, F.softmax(classifications, dim=-1), ldm_regressions)
+        return output
+
+    def __detect_faces(self, inputs):
+        # get scale
+        height, width = inputs.shape[2:]
+        self.scale = torch.tensor([width, height, width, height], dtype=torch.float32).to(device)
+        tmp = [width, height, width, height, width, height, width, height, width, height]
+        self.scale1 = torch.tensor(tmp, dtype=torch.float32).to(device)
+
+        # forawrd
+        inputs = inputs.to(device)
+        if self.half_inference:
+            inputs = inputs.half()
+        loc, conf, landmarks = self(inputs)
+
+        # get priorbox
+        priorbox = PriorBox(self.cfg, image_size=inputs.shape[2:])
+        priors = priorbox.forward().to(device)
+
+        return loc, conf, landmarks, priors
+
+    # single image detection
+    def transform(self, image, use_origin_size):
+        # convert to opencv format
+        if isinstance(image, Image.Image):
+            image = cv2.cvtColor(np.asarray(image), cv2.COLOR_RGB2BGR)
+        image = image.astype(np.float32)
+
+        # testing scale
+        im_size_min = np.min(image.shape[0:2])
+        im_size_max = np.max(image.shape[0:2])
+        resize = float(self.target_size) / float(im_size_min)
+
+        # prevent bigger axis from being more than max_size
+        if np.round(resize * im_size_max) > self.max_size:
+            resize = float(self.max_size) / float(im_size_max)
+        resize = 1 if use_origin_size else resize
+
+        # resize
+        if resize != 1:
+            image = cv2.resize(image, None, None, fx=resize, fy=resize, interpolation=cv2.INTER_LINEAR)
+
+        # convert to torch.tensor format
+        # image -= (104, 117, 123)
+        image = image.transpose(2, 0, 1)
+        image = torch.from_numpy(image).unsqueeze(0)
+
+        return image, resize
+
+    def detect_faces(
+        self,
+        image,
+        conf_threshold=0.8,
+        nms_threshold=0.4,
+        use_origin_size=True,
+    ):
+        """
+        Params:
+            imgs: BGR image
+        """
+        image, self.resize = self.transform(image, use_origin_size)
+        image = image.to(device)
+        if self.half_inference:
+            image = image.half()
+        image = image - self.mean_tensor
+
+        loc, conf, landmarks, priors = self.__detect_faces(image)
+
+        boxes = decode(loc.data.squeeze(0), priors.data, self.cfg['variance'])
+        boxes = boxes * self.scale / self.resize
+        boxes = boxes.cpu().numpy()
+
+        scores = conf.squeeze(0).data.cpu().numpy()[:, 1]
+
+        landmarks = decode_landm(landmarks.squeeze(0), priors, self.cfg['variance'])
+        landmarks = landmarks * self.scale1 / self.resize
+        landmarks = landmarks.cpu().numpy()
+
+        # ignore low scores
+        inds = np.where(scores > conf_threshold)[0]
+        boxes, landmarks, scores = boxes[inds], landmarks[inds], scores[inds]
+
+        # sort
+        order = scores.argsort()[::-1]
+        boxes, landmarks, scores = boxes[order], landmarks[order], scores[order]
+
+        # do NMS
+        bounding_boxes = np.hstack((boxes, scores[:, np.newaxis])).astype(np.float32, copy=False)
+        keep = py_cpu_nms(bounding_boxes, nms_threshold)
+        bounding_boxes, landmarks = bounding_boxes[keep, :], landmarks[keep]
+        # self.t['forward_pass'].toc()
+        # print(self.t['forward_pass'].average_time)
+        # import sys
+        # sys.stdout.flush()
+        return np.concatenate((bounding_boxes, landmarks), axis=1)
+
+    def __align_multi(self, image, boxes, landmarks, limit=None):
+
+        if len(boxes) < 1:
+            return [], []
+
+        if limit:
+            boxes = boxes[:limit]
+            landmarks = landmarks[:limit]
+
+        faces = []
+        for landmark in landmarks:
+            facial5points = [[landmark[2 * j], landmark[2 * j + 1]] for j in range(5)]
+
+            warped_face = warp_and_crop_face(np.array(image), facial5points, self.reference, crop_size=(112, 112))
+            faces.append(warped_face)
+
+        return np.concatenate((boxes, landmarks), axis=1), faces
+
+    def align_multi(self, img, conf_threshold=0.8, limit=None):
+
+        rlt = self.detect_faces(img, conf_threshold=conf_threshold)
+        boxes, landmarks = rlt[:, 0:5], rlt[:, 5:]
+
+        return self.__align_multi(img, boxes, landmarks, limit)
+
+    # batched detection
+    def batched_transform(self, frames, use_origin_size):
+        """
+        Arguments:
+            frames: a list of PIL.Image, or torch.Tensor(shape=[n, h, w, c],
+                type=np.float32, BGR format).
+            use_origin_size: whether to use origin size.
+        """
+        from_PIL = True if isinstance(frames[0], Image.Image) else False
+
+        # convert to opencv format
+        if from_PIL:
+            frames = [cv2.cvtColor(np.asarray(frame), cv2.COLOR_RGB2BGR) for frame in frames]
+            frames = np.asarray(frames, dtype=np.float32)
+
+        # testing scale
+        im_size_min = np.min(frames[0].shape[0:2])
+        im_size_max = np.max(frames[0].shape[0:2])
+        resize = float(self.target_size) / float(im_size_min)
+
+        # prevent bigger axis from being more than max_size
+        if np.round(resize * im_size_max) > self.max_size:
+            resize = float(self.max_size) / float(im_size_max)
+        resize = 1 if use_origin_size else resize
+
+        # resize
+        if resize != 1:
+            if not from_PIL:
+                frames = F.interpolate(frames, scale_factor=resize)
+            else:
+                frames = [
+                    cv2.resize(frame, None, None, fx=resize, fy=resize, interpolation=cv2.INTER_LINEAR)
+                    for frame in frames
+                ]
+
+        # convert to torch.tensor format
+        if not from_PIL:
+            frames = frames.transpose(1, 2).transpose(1, 3).contiguous()
+        else:
+            frames = frames.transpose((0, 3, 1, 2))
+            frames = torch.from_numpy(frames)
+
+        return frames, resize
+
+    def batched_detect_faces(self, frames, conf_threshold=0.8, nms_threshold=0.4, use_origin_size=True):
+        """
+        Arguments:
+            frames: a list of PIL.Image, or np.array(shape=[n, h, w, c],
+                type=np.uint8, BGR format).
+            conf_threshold: confidence threshold.
+            nms_threshold: nms threshold.
+            use_origin_size: whether to use origin size.
+        Returns:
+            final_bounding_boxes: list of np.array ([n_boxes, 5],
+                type=np.float32).
+            final_landmarks: list of np.array ([n_boxes, 10], type=np.float32).
+        """
+        # self.t['forward_pass'].tic()
+        frames, self.resize = self.batched_transform(frames, use_origin_size)
+        frames = frames.to(device)
+        frames = frames - self.mean_tensor
+
+        b_loc, b_conf, b_landmarks, priors = self.__detect_faces(frames)
+
+        final_bounding_boxes, final_landmarks = [], []
+
+        # decode
+        priors = priors.unsqueeze(0)
+        b_loc = batched_decode(b_loc, priors, self.cfg['variance']) * self.scale / self.resize
+        b_landmarks = batched_decode_landm(b_landmarks, priors, self.cfg['variance']) * self.scale1 / self.resize
+        b_conf = b_conf[:, :, 1]
+
+        # index for selection
+        b_indice = b_conf > conf_threshold
+
+        # concat
+        b_loc_and_conf = torch.cat((b_loc, b_conf.unsqueeze(-1)), dim=2).float()
+
+        for pred, landm, inds in zip(b_loc_and_conf, b_landmarks, b_indice):
+
+            # ignore low scores
+            pred, landm = pred[inds, :], landm[inds, :]
+            if pred.shape[0] == 0:
+                final_bounding_boxes.append(np.array([], dtype=np.float32))
+                final_landmarks.append(np.array([], dtype=np.float32))
+                continue
+
+            # sort
+            # order = score.argsort(descending=True)
+            # box, landm, score = box[order], landm[order], score[order]
+
+            # to CPU
+            bounding_boxes, landm = pred.cpu().numpy(), landm.cpu().numpy()
+
+            # NMS
+            keep = py_cpu_nms(bounding_boxes, nms_threshold)
+            bounding_boxes, landmarks = bounding_boxes[keep, :], landm[keep]
+
+            # append
+            final_bounding_boxes.append(bounding_boxes)
+            final_landmarks.append(landmarks)
+        # self.t['forward_pass'].toc(average=True)
+        # self.batch_time += self.t['forward_pass'].diff
+        # self.total_frame += len(frames)
+        # print(self.batch_time / self.total_frame)
+
+        return final_bounding_boxes, final_landmarks

py/facelib/detection/retinaface/retinaface_net.py

@@ -0,0 +1,196 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def conv_bn(inp, oup, stride=1, leaky=0):
+    return nn.Sequential(
+        nn.Conv2d(inp, oup, 3, stride, 1, bias=False), nn.BatchNorm2d(oup),
+        nn.LeakyReLU(negative_slope=leaky, inplace=True))
+
+
+def conv_bn_no_relu(inp, oup, stride):
+    return nn.Sequential(
+        nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
+        nn.BatchNorm2d(oup),
+    )
+
+
+def conv_bn1X1(inp, oup, stride, leaky=0):
+    return nn.Sequential(
+        nn.Conv2d(inp, oup, 1, stride, padding=0, bias=False), nn.BatchNorm2d(oup),
+        nn.LeakyReLU(negative_slope=leaky, inplace=True))
+
+
+def conv_dw(inp, oup, stride, leaky=0.1):
+    return nn.Sequential(
+        nn.Conv2d(inp, inp, 3, stride, 1, groups=inp, bias=False),
+        nn.BatchNorm2d(inp),
+        nn.LeakyReLU(negative_slope=leaky, inplace=True),
+        nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
+        nn.BatchNorm2d(oup),
+        nn.LeakyReLU(negative_slope=leaky, inplace=True),
+    )
+
+
+class SSH(nn.Module):
+
+    def __init__(self, in_channel, out_channel):
+        super(SSH, self).__init__()
+        assert out_channel % 4 == 0
+        leaky = 0
+        if (out_channel <= 64):
+            leaky = 0.1
+        self.conv3X3 = conv_bn_no_relu(in_channel, out_channel // 2, stride=1)
+
+        self.conv5X5_1 = conv_bn(in_channel, out_channel // 4, stride=1, leaky=leaky)
+        self.conv5X5_2 = conv_bn_no_relu(out_channel // 4, out_channel // 4, stride=1)
+
+        self.conv7X7_2 = conv_bn(out_channel // 4, out_channel // 4, stride=1, leaky=leaky)
+        self.conv7x7_3 = conv_bn_no_relu(out_channel // 4, out_channel // 4, stride=1)
+
+    def forward(self, input):
+        conv3X3 = self.conv3X3(input)
+
+        conv5X5_1 = self.conv5X5_1(input)
+        conv5X5 = self.conv5X5_2(conv5X5_1)
+
+        conv7X7_2 = self.conv7X7_2(conv5X5_1)
+        conv7X7 = self.conv7x7_3(conv7X7_2)
+
+        out = torch.cat([conv3X3, conv5X5, conv7X7], dim=1)
+        out = F.relu(out)
+        return out
+
+
+class FPN(nn.Module):
+
+    def __init__(self, in_channels_list, out_channels):
+        super(FPN, self).__init__()
+        leaky = 0
+        if (out_channels <= 64):
+            leaky = 0.1
+        self.output1 = conv_bn1X1(in_channels_list[0], out_channels, stride=1, leaky=leaky)
+        self.output2 = conv_bn1X1(in_channels_list[1], out_channels, stride=1, leaky=leaky)
+        self.output3 = conv_bn1X1(in_channels_list[2], out_channels, stride=1, leaky=leaky)
+
+        self.merge1 = conv_bn(out_channels, out_channels, leaky=leaky)
+        self.merge2 = conv_bn(out_channels, out_channels, leaky=leaky)
+
+    def forward(self, input):
+        # names = list(input.keys())
+        # input = list(input.values())
+
+        output1 = self.output1(input[0])
+        output2 = self.output2(input[1])
+        output3 = self.output3(input[2])
+
+        up3 = F.interpolate(output3, size=[output2.size(2), output2.size(3)], mode='nearest')
+        output2 = output2 + up3
+        output2 = self.merge2(output2)
+
+        up2 = F.interpolate(output2, size=[output1.size(2), output1.size(3)], mode='nearest')
+        output1 = output1 + up2
+        output1 = self.merge1(output1)
+
+        out = [output1, output2, output3]
+        return out
+
+
+class MobileNetV1(nn.Module):
+
+    def __init__(self):
+        super(MobileNetV1, self).__init__()
+        self.stage1 = nn.Sequential(
+            conv_bn(3, 8, 2, leaky=0.1),  # 3
+            conv_dw(8, 16, 1),  # 7
+            conv_dw(16, 32, 2),  # 11
+            conv_dw(32, 32, 1),  # 19
+            conv_dw(32, 64, 2),  # 27
+            conv_dw(64, 64, 1),  # 43
+        )
+        self.stage2 = nn.Sequential(
+            conv_dw(64, 128, 2),  # 43 + 16 = 59
+            conv_dw(128, 128, 1),  # 59 + 32 = 91
+            conv_dw(128, 128, 1),  # 91 + 32 = 123
+            conv_dw(128, 128, 1),  # 123 + 32 = 155
+            conv_dw(128, 128, 1),  # 155 + 32 = 187
+            conv_dw(128, 128, 1),  # 187 + 32 = 219
+        )
+        self.stage3 = nn.Sequential(
+            conv_dw(128, 256, 2),  # 219 +3 2 = 241
+            conv_dw(256, 256, 1),  # 241 + 64 = 301
+        )
+        self.avg = nn.AdaptiveAvgPool2d((1, 1))
+        self.fc = nn.Linear(256, 1000)
+
+    def forward(self, x):
+        x = self.stage1(x)
+        x = self.stage2(x)
+        x = self.stage3(x)
+        x = self.avg(x)
+        # x = self.model(x)
+        x = x.view(-1, 256)
+        x = self.fc(x)
+        return x
+
+
+class ClassHead(nn.Module):
+
+    def __init__(self, inchannels=512, num_anchors=3):
+        super(ClassHead, self).__init__()
+        self.num_anchors = num_anchors
+        self.conv1x1 = nn.Conv2d(inchannels, self.num_anchors * 2, kernel_size=(1, 1), stride=1, padding=0)
+
+    def forward(self, x):
+        out = self.conv1x1(x)
+        out = out.permute(0, 2, 3, 1).contiguous()
+
+        return out.view(out.shape[0], -1, 2)
+
+
+class BboxHead(nn.Module):
+
+    def __init__(self, inchannels=512, num_anchors=3):
+        super(BboxHead, self).__init__()
+        self.conv1x1 = nn.Conv2d(inchannels, num_anchors * 4, kernel_size=(1, 1), stride=1, padding=0)
+
+    def forward(self, x):
+        out = self.conv1x1(x)
+        out = out.permute(0, 2, 3, 1).contiguous()
+
+        return out.view(out.shape[0], -1, 4)
+
+
+class LandmarkHead(nn.Module):
+
+    def __init__(self, inchannels=512, num_anchors=3):
+        super(LandmarkHead, self).__init__()
+        self.conv1x1 = nn.Conv2d(inchannels, num_anchors * 10, kernel_size=(1, 1), stride=1, padding=0)
+
+    def forward(self, x):
+        out = self.conv1x1(x)
+        out = out.permute(0, 2, 3, 1).contiguous()
+
+        return out.view(out.shape[0], -1, 10)
+
+
+def make_class_head(fpn_num=3, inchannels=64, anchor_num=2):
+    classhead = nn.ModuleList()
+    for i in range(fpn_num):
+        classhead.append(ClassHead(inchannels, anchor_num))
+    return classhead
+
+
+def make_bbox_head(fpn_num=3, inchannels=64, anchor_num=2):
+    bboxhead = nn.ModuleList()
+    for i in range(fpn_num):
+        bboxhead.append(BboxHead(inchannels, anchor_num))
+    return bboxhead
+
+
+def make_landmark_head(fpn_num=3, inchannels=64, anchor_num=2):
+    landmarkhead = nn.ModuleList()
+    for i in range(fpn_num):
+        landmarkhead.append(LandmarkHead(inchannels, anchor_num))
+    return landmarkhead

py/facelib/detection/retinaface/retinaface_utils.py

@@ -1,5 +1,50 @@
-import torch
 import numpy as np
+import torch
+import torchvision
+from itertools import product as product
+from math import ceil
+
+
+class PriorBox(object):
+
+    def __init__(self, cfg, image_size=None, phase='train'):
+        super(PriorBox, self).__init__()
+        self.min_sizes = cfg['min_sizes']
+        self.steps = cfg['steps']
+        self.clip = cfg['clip']
+        self.image_size = image_size
+        self.feature_maps = [[ceil(self.image_size[0] / step), ceil(self.image_size[1] / step)] for step in self.steps]
+        self.name = 's'
+
+    def forward(self):
+        anchors = []
+        for k, f in enumerate(self.feature_maps):
+            min_sizes = self.min_sizes[k]
+            for i, j in product(range(f[0]), range(f[1])):
+                for min_size in min_sizes:
+                    s_kx = min_size / self.image_size[1]
+                    s_ky = min_size / self.image_size[0]
+                    dense_cx = [x * self.steps[k] / self.image_size[1] for x in [j + 0.5]]
+                    dense_cy = [y * self.steps[k] / self.image_size[0] for y in [i + 0.5]]
+                    for cy, cx in product(dense_cy, dense_cx):
+                        anchors += [cx, cy, s_kx, s_ky]
+
+        # back to torch land
+        output = torch.Tensor(anchors).view(-1, 4)
+        if self.clip:
+            output.clamp_(max=1, min=0)
+        return output
+
+
+def py_cpu_nms(dets, thresh):
+    """Pure Python NMS baseline."""
+    keep = torchvision.ops.nms(
+        boxes=torch.Tensor(dets[:, :4]),
+        scores=torch.Tensor(dets[:, 4]),
+        iou_threshold=thresh,
+    )
+
+    return list(keep)
 
 
 def point_form(boxes):
@@ -10,8 +55,11 @@ def point_form(boxes):
     Return:
         boxes: (tensor) Converted xmin, ymin, xmax, ymax form of boxes.
     """
-    return torch.cat((boxes[:, :2] - boxes[:, 2:]/2,     # xmin, ymin
-                     boxes[:, :2] + boxes[:, 2:]/2), 1)  # xmax, ymax
+    return torch.cat(
+        (
+            boxes[:, :2] - boxes[:, 2:] / 2,  # xmin, ymin
+            boxes[:, :2] + boxes[:, 2:] / 2),
+        1)  # xmax, ymax
 
 
 def center_size(boxes):
@@ -22,8 +70,10 @@ def center_size(boxes):
     Return:
         boxes: (tensor) Converted xmin, ymin, xmax, ymax form of boxes.
     """
-    return torch.cat((boxes[:, 2:] + boxes[:, :2])/2,  # cx, cy
-                     boxes[:, 2:] - boxes[:, :2], 1)  # w, h
+    return torch.cat(
+        (boxes[:, 2:] + boxes[:, :2]) / 2,  # cx, cy
+        boxes[:, 2:] - boxes[:, :2],
+        1)  # w, h
 
 
 def intersect(box_a, box_b):
@@ -39,10 +89,8 @@ def intersect(box_a, box_b):
     """
     A = box_a.size(0)
     B = box_b.size(0)
-    max_xy = torch.min(box_a[:, 2:].unsqueeze(1).expand(A, B, 2),
-                       box_b[:, 2:].unsqueeze(0).expand(A, B, 2))
-    min_xy = torch.max(box_a[:, :2].unsqueeze(1).expand(A, B, 2),
-                       box_b[:, :2].unsqueeze(0).expand(A, B, 2))
+    max_xy = torch.min(box_a[:, 2:].unsqueeze(1).expand(A, B, 2), box_b[:, 2:].unsqueeze(0).expand(A, B, 2))
+    min_xy = torch.max(box_a[:, :2].unsqueeze(1).expand(A, B, 2), box_b[:, :2].unsqueeze(0).expand(A, B, 2))
     inter = torch.clamp((max_xy - min_xy), min=0)
     return inter[:, :, 0] * inter[:, :, 1]
 
@@ -60,10 +108,8 @@ def jaccard(box_a, box_b):
         jaccard overlap: (tensor) Shape: [box_a.size(0), box_b.size(0)]
     """
     inter = intersect(box_a, box_b)
-    area_a = ((box_a[:, 2]-box_a[:, 0]) *
-              (box_a[:, 3]-box_a[:, 1])).unsqueeze(1).expand_as(inter)  # [A,B]
-    area_b = ((box_b[:, 2]-box_b[:, 0]) *
-              (box_b[:, 3]-box_b[:, 1])).unsqueeze(0).expand_as(inter)  # [A,B]
+    area_a = ((box_a[:, 2] - box_a[:, 0]) * (box_a[:, 3] - box_a[:, 1])).unsqueeze(1).expand_as(inter)  # [A,B]
+    area_b = ((box_b[:, 2] - box_b[:, 0]) * (box_b[:, 3] - box_b[:, 1])).unsqueeze(0).expand_as(inter)  # [A,B]
     union = area_a + area_b - inter
     return inter / union  # [A,B]
 
@@ -98,25 +144,23 @@ def match(threshold, truths, priors, variances, labels, landms, loc_t, conf_t, l
     overlap, encode the bounding boxes, then return the matched indices
     corresponding to both confidence and location preds.
     Args:
-        threshold: (float) The overlap threshold used when mathing boxes.
+        threshold: (float) The overlap threshold used when matching boxes.
         truths: (tensor) Ground truth boxes, Shape: [num_obj, 4].
         priors: (tensor) Prior boxes from priorbox layers, Shape: [n_priors,4].
         variances: (tensor) Variances corresponding to each prior coord,
             Shape: [num_priors, 4].
         labels: (tensor) All the class labels for the image, Shape: [num_obj].
         landms: (tensor) Ground truth landms, Shape [num_obj, 10].
-        loc_t: (tensor) Tensor to be filled w/ endcoded location targets.
+        loc_t: (tensor) Tensor to be filled w/ encoded location targets.
         conf_t: (tensor) Tensor to be filled w/ matched indices for conf preds.
-        landm_t: (tensor) Tensor to be filled w/ endcoded landm targets.
+        landm_t: (tensor) Tensor to be filled w/ encoded landm targets.
         idx: (int) current batch index
     Return:
-        The matched indices corresponding to 1)location 2)confidence 3)landm preds.
+        The matched indices corresponding to 1)location 2)confidence
+        3)landm preds.
     """
     # jaccard index
-    overlaps = jaccard(
-        truths,
-        point_form(priors)
-    )
+    overlaps = jaccard(truths, point_form(priors))
     # (Bipartite Matching)
     # [1,num_objects] best prior for each ground truth
     best_prior_overlap, best_prior_idx = overlaps.max(1, keepdim=True)
@@ -139,16 +183,16 @@ def match(threshold, truths, priors, variances, labels, landms, loc_t, conf_t, l
     best_truth_overlap.index_fill_(0, best_prior_idx_filter, 2)  # ensure best prior
     # TODO refactor: index  best_prior_idx with long tensor
     # ensure every gt matches with its prior of max overlap
-    for j in range(best_prior_idx.size(0)):     # 判别此anchor是预测哪一个boxes
+    for j in range(best_prior_idx.size(0)):  # 判别此anchor是预测哪一个boxes
         best_truth_idx[best_prior_idx[j]] = j
-    matches = truths[best_truth_idx]            # Shape: [num_priors,4] 此处为每一个anchor对应的bbox取出来
-    conf = labels[best_truth_idx]               # Shape: [num_priors]      此处为每一个anchor对应的label取出来
-    conf[best_truth_overlap < threshold] = 0    # label as background   overlap<0.35的全部作为负样本
+    matches = truths[best_truth_idx]  # Shape: [num_priors,4] 此处为每一个anchor对应的bbox取出来
+    conf = labels[best_truth_idx]  # Shape: [num_priors]      此处为每一个anchor对应的label取出来
+    conf[best_truth_overlap < threshold] = 0  # label as background   overlap<0.35的全部作为负样本
     loc = encode(matches, priors, variances)
 
     matches_landm = landms[best_truth_idx]
     landm = encode_landm(matches_landm, priors, variances)
-    loc_t[idx] = loc    # [num_priors,4] encoded offsets to learn
+    loc_t[idx] = loc  # [num_priors,4] encoded offsets to learn
     conf_t[idx] = conf  # [num_priors] top class label for each prior
     landm_t[idx] = landm
 
@@ -167,7 +211,7 @@ def encode(matched, priors, variances):
     """
 
     # dist b/t match center and prior's center
-    g_cxcy = (matched[:, :2] + matched[:, 2:])/2 - priors[:, :2]
+    g_cxcy = (matched[:, :2] + matched[:, 2:]) / 2 - priors[:, :2]
     # encode variance
     g_cxcy /= (variances[0] * priors[:, 2:])
     # match wh / prior wh
@@ -176,6 +220,7 @@ def encode(matched, priors, variances):
     # return target for smooth_l1_loss
     return torch.cat([g_cxcy, g_wh], 1)  # [num_priors,4]
 
+
 def encode_landm(matched, priors, variances):
     """Encode the variances from the priorbox layers into the ground truth boxes
     we have matched (based on jaccard overlap) with the prior boxes.
@@ -219,13 +264,13 @@ def decode(loc, priors, variances):
         decoded bounding box predictions
     """
 
-    boxes = torch.cat((
-        priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:],
-        priors[:, 2:] * torch.exp(loc[:, 2:] * variances[1])), 1)
+    boxes = torch.cat((priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:],
+                       priors[:, 2:] * torch.exp(loc[:, 2:] * variances[1])), 1)
     boxes[:, :2] -= boxes[:, 2:] / 2
     boxes[:, 2:] += boxes[:, :2]
     return boxes
 
+
 def decode_landm(pre, priors, variances):
     """Decode landm from predictions using priors to undo
     the encoding we did for offset regression at train time.
@@ -238,12 +283,60 @@ def decode_landm(pre, priors, variances):
     Return:
         decoded landm predictions
     """
-    landms = torch.cat((priors[:, :2] + pre[:, :2] * variances[0] * priors[:, 2:],
-                        priors[:, :2] + pre[:, 2:4] * variances[0] * priors[:, 2:],
-                        priors[:, :2] + pre[:, 4:6] * variances[0] * priors[:, 2:],
-                        priors[:, :2] + pre[:, 6:8] * variances[0] * priors[:, 2:],
-                        priors[:, :2] + pre[:, 8:10] * variances[0] * priors[:, 2:],
-                        ), dim=1)
+    tmp = (
+        priors[:, :2] + pre[:, :2] * variances[0] * priors[:, 2:],
+        priors[:, :2] + pre[:, 2:4] * variances[0] * priors[:, 2:],
+        priors[:, :2] + pre[:, 4:6] * variances[0] * priors[:, 2:],
+        priors[:, :2] + pre[:, 6:8] * variances[0] * priors[:, 2:],
+        priors[:, :2] + pre[:, 8:10] * variances[0] * priors[:, 2:],
+    )
+    landms = torch.cat(tmp, dim=1)
+    return landms
+
+
+def batched_decode(b_loc, priors, variances):
+    """Decode locations from predictions using priors to undo
+    the encoding we did for offset regression at train time.
+    Args:
+        b_loc (tensor): location predictions for loc layers,
+            Shape: [num_batches,num_priors,4]
+        priors (tensor): Prior boxes in center-offset form.
+            Shape: [1,num_priors,4].
+        variances: (list[float]) Variances of priorboxes
+    Return:
+        decoded bounding box predictions
+    """
+    boxes = (
+        priors[:, :, :2] + b_loc[:, :, :2] * variances[0] * priors[:, :, 2:],
+        priors[:, :, 2:] * torch.exp(b_loc[:, :, 2:] * variances[1]),
+    )
+    boxes = torch.cat(boxes, dim=2)
+
+    boxes[:, :, :2] -= boxes[:, :, 2:] / 2
+    boxes[:, :, 2:] += boxes[:, :, :2]
+    return boxes
+
+
+def batched_decode_landm(pre, priors, variances):
+    """Decode landm from predictions using priors to undo
+    the encoding we did for offset regression at train time.
+    Args:
+        pre (tensor): landm predictions for loc layers,
+            Shape: [num_batches,num_priors,10]
+        priors (tensor): Prior boxes in center-offset form.
+            Shape: [1,num_priors,4].
+        variances: (list[float]) Variances of priorboxes
+    Return:
+        decoded landm predictions
+    """
+    landms = (
+        priors[:, :, :2] + pre[:, :, :2] * variances[0] * priors[:, :, 2:],
+        priors[:, :, :2] + pre[:, :, 2:4] * variances[0] * priors[:, :, 2:],
+        priors[:, :, :2] + pre[:, :, 4:6] * variances[0] * priors[:, :, 2:],
+        priors[:, :, :2] + pre[:, :, 6:8] * variances[0] * priors[:, :, 2:],
+        priors[:, :, :2] + pre[:, :, 8:10] * variances[0] * priors[:, :, 2:],
+    )
+    landms = torch.cat(landms, dim=2)
     return landms
 
 
@@ -255,7 +348,7 @@ def log_sum_exp(x):
         x (Variable(tensor)): conf_preds from conf layers
     """
     x_max = x.data.max()
-    return torch.log(torch.sum(torch.exp(x-x_max), 1, keepdim=True)) + x_max
+    return torch.log(torch.sum(torch.exp(x - x_max), 1, keepdim=True)) + x_max
 
 
 # Original author: Francisco Massa:
@@ -318,13 +411,11 @@ def nms(boxes, scores, overlap=0.5, top_k=200):
         # check sizes of xx1 and xx2.. after each iteration
         w = torch.clamp(w, min=0.0)
         h = torch.clamp(h, min=0.0)
-        inter = w*h
+        inter = w * h
         # IoU = i / (area(a) + area(b) - i)
         rem_areas = torch.index_select(area, 0, idx)  # load remaining areas)
         union = (rem_areas - inter) + area[i]
-        IoU = inter/union  # store result in iou
+        IoU = inter / union  # store result in iou
         # keep only elements with an IoU <= overlap
         idx = idx[IoU.le(overlap)]
     return keep, count
-
-

py/facelib/detection/yolov5face/face_detector.py

@@ -0,0 +1,142 @@
+import copy
+import os
+from pathlib import Path
+
+import cv2
+import numpy as np
+import torch
+from torch import nn
+
+from .models.common import Conv
+from .models.yolo import Model
+from .utils.datasets import letterbox
+from .utils.general import (
+    check_img_size,
+    non_max_suppression_face,
+    scale_coords,
+    scale_coords_landmarks,
+)
+
+IS_HIGH_VERSION = tuple(map(int, torch.__version__.split('+')[0].split('.'))) >= (1, 9, 0)
+
+
+def isListempty(inList):
+    if isinstance(inList, list): # Is a list
+        return all(map(isListempty, inList))
+    return False # Not a list
+
+class YoloDetector:
+    def __init__(
+        self,
+        config_name,
+        min_face=10,
+        target_size=None,
+        device='cuda',
+    ):
+        """
+        config_name: name of .yaml config with network configuration from models/ folder.
+        min_face : minimal face size in pixels.
+        target_size : target size of smaller image axis (choose lower for faster work). e.g. 480, 720, 1080.
+                    None for original resolution.
+        """
+        self._class_path = Path(__file__).parent.absolute()
+        self.target_size = target_size
+        self.min_face = min_face
+        self.detector = Model(cfg=config_name)
+        self.device = device
+
+
+    def _preprocess(self, imgs):
+        """
+        Preprocessing image before passing through the network. Resize and conversion to torch tensor.
+        """
+        pp_imgs = []
+        for img in imgs:
+            h0, w0 = img.shape[:2]  # orig hw
+            if self.target_size:
+                r = self.target_size / min(h0, w0)  # resize image to img_size
+                if r < 1:
+                    img = cv2.resize(img, (int(w0 * r), int(h0 * r)), interpolation=cv2.INTER_LINEAR)
+
+            imgsz = check_img_size(max(img.shape[:2]), s=self.detector.stride.max())  # check img_size
+            img = letterbox(img, new_shape=imgsz)[0]
+            pp_imgs.append(img)
+        pp_imgs = np.array(pp_imgs)
+        pp_imgs = pp_imgs.transpose(0, 3, 1, 2)
+        pp_imgs = torch.from_numpy(pp_imgs).to(self.device)
+        pp_imgs = pp_imgs.float()  # uint8 to fp16/32
+        return pp_imgs / 255.0  # 0 - 255 to 0.0 - 1.0
+
+    def _postprocess(self, imgs, origimgs, pred, conf_thres, iou_thres):
+        """
+        Postprocessing of raw pytorch model output.
+        Returns:
+            bboxes: list of arrays with 4 coordinates of bounding boxes with format x1,y1,x2,y2.
+            points: list of arrays with coordinates of 5 facial keypoints (eyes, nose, lips corners).
+        """
+        bboxes = [[] for _ in range(len(origimgs))]
+        landmarks = [[] for _ in range(len(origimgs))]
+
+        pred = non_max_suppression_face(pred, conf_thres, iou_thres)
+
+        for image_id, origimg in enumerate(origimgs):
+            img_shape = origimg.shape
+            image_height, image_width = img_shape[:2]
+            gn = torch.tensor(img_shape)[[1, 0, 1, 0]]  # normalization gain whwh
+            gn_lks = torch.tensor(img_shape)[[1, 0, 1, 0, 1, 0, 1, 0, 1, 0]]  # normalization gain landmarks
+            det = pred[image_id].cpu()
+            scale_coords(imgs[image_id].shape[1:], det[:, :4], img_shape).round()
+            scale_coords_landmarks(imgs[image_id].shape[1:], det[:, 5:15], img_shape).round()
+
+            for j in range(det.size()[0]):
+                box = (det[j, :4].view(1, 4) / gn).view(-1).tolist()
+                box = list(
+                    map(int, [box[0] * image_width, box[1] * image_height, box[2] * image_width, box[3] * image_height])
+                )
+                if box[3] - box[1] < self.min_face:
+                    continue
+                lm = (det[j, 5:15].view(1, 10) / gn_lks).view(-1).tolist()
+                lm = list(map(int, [i * image_width if j % 2 == 0 else i * image_height for j, i in enumerate(lm)]))
+                lm = [lm[i : i + 2] for i in range(0, len(lm), 2)]
+                bboxes[image_id].append(box)
+                landmarks[image_id].append(lm)
+        return bboxes, landmarks
+
+    def detect_faces(self, imgs, conf_thres=0.7, iou_thres=0.5):
+        """
+        Get bbox coordinates and keypoints of faces on original image.
+        Params:
+            imgs: image or list of images to detect faces on with BGR order (convert to RGB order for inference)
+            conf_thres: confidence threshold for each prediction
+            iou_thres: threshold for NMS (filter of intersecting bboxes)
+        Returns:
+            bboxes: list of arrays with 4 coordinates of bounding boxes with format x1,y1,x2,y2.
+            points: list of arrays with coordinates of 5 facial keypoints (eyes, nose, lips corners).
+        """
+        # Pass input images through face detector
+        images = imgs if isinstance(imgs, list) else [imgs]
+        images = [cv2.cvtColor(img, cv2.COLOR_BGR2RGB) for img in images]
+        origimgs = copy.deepcopy(images)
+
+        images = self._preprocess(images)
+        
+        if IS_HIGH_VERSION:
+            with torch.inference_mode():  # for pytorch>=1.9 
+                pred = self.detector(images)[0]
+        else:
+            with torch.no_grad():  # for pytorch<1.9
+                pred = self.detector(images)[0]
+
+        bboxes, points = self._postprocess(images, origimgs, pred, conf_thres, iou_thres)
+
+        # return bboxes, points
+        if not isListempty(points):
+            bboxes = np.array(bboxes).reshape(-1,4)
+            points = np.array(points).reshape(-1,10)
+            padding = bboxes[:,0].reshape(-1,1)
+            return np.concatenate((bboxes, padding, points), axis=1)
+        else:
+            return None
+
+    def __call__(self, *args):
+        return self.predict(*args)

py/facelib/detection/yolov5face/models/common.py

@@ -0,0 +1,299 @@
+# This file contains modules common to various models
+
+import math
+
+import numpy as np
+import torch
+from torch import nn
+
+from ..utils.datasets import letterbox
+from ..utils.general import (
+    make_divisible,
+    non_max_suppression,
+    scale_coords,
+    xyxy2xywh,
+)
+
+
+def autopad(k, p=None):  # kernel, padding
+    # Pad to 'same'
+    if p is None:
+        p = k // 2 if isinstance(k, int) else [x // 2 for x in k]  # auto-pad
+    return p
+
+
+def channel_shuffle(x, groups):
+    batchsize, num_channels, height, width = x.data.size()
+    channels_per_group = torch.div(num_channels, groups, rounding_mode="trunc")
+
+    # reshape
+    x = x.view(batchsize, groups, channels_per_group, height, width)
+    x = torch.transpose(x, 1, 2).contiguous()
+
+    # flatten
+    return x.view(batchsize, -1, height, width)
+
+
+def DWConv(c1, c2, k=1, s=1, act=True):
+    # Depthwise convolution
+    return Conv(c1, c2, k, s, g=math.gcd(c1, c2), act=act)
+
+
+class Conv(nn.Module):
+    # Standard convolution
+    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):  # ch_in, ch_out, kernel, stride, padding, groups
+        super().__init__()
+        self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False)
+        self.bn = nn.BatchNorm2d(c2)
+        self.act = nn.SiLU() if act is True else (act if isinstance(act, nn.Module) else nn.Identity())
+
+    def forward(self, x):
+        return self.act(self.bn(self.conv(x)))
+
+    def fuseforward(self, x):
+        return self.act(self.conv(x))
+
+
+class StemBlock(nn.Module):
+    def __init__(self, c1, c2, k=3, s=2, p=None, g=1, act=True):
+        super().__init__()
+        self.stem_1 = Conv(c1, c2, k, s, p, g, act)
+        self.stem_2a = Conv(c2, c2 // 2, 1, 1, 0)
+        self.stem_2b = Conv(c2 // 2, c2, 3, 2, 1)
+        self.stem_2p = nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True)
+        self.stem_3 = Conv(c2 * 2, c2, 1, 1, 0)
+
+    def forward(self, x):
+        stem_1_out = self.stem_1(x)
+        stem_2a_out = self.stem_2a(stem_1_out)
+        stem_2b_out = self.stem_2b(stem_2a_out)
+        stem_2p_out = self.stem_2p(stem_1_out)
+        return self.stem_3(torch.cat((stem_2b_out, stem_2p_out), 1))
+
+
+class Bottleneck(nn.Module):
+    # Standard bottleneck
+    def __init__(self, c1, c2, shortcut=True, g=1, e=0.5):  # ch_in, ch_out, shortcut, groups, expansion
+        super().__init__()
+        c_ = int(c2 * e)  # hidden channels
+        self.cv1 = Conv(c1, c_, 1, 1)
+        self.cv2 = Conv(c_, c2, 3, 1, g=g)
+        self.add = shortcut and c1 == c2
+
+    def forward(self, x):
+        return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))
+
+
+class BottleneckCSP(nn.Module):
+    # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
+    def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):  # ch_in, ch_out, number, shortcut, groups, expansion
+        super().__init__()
+        c_ = int(c2 * e)  # hidden channels
+        self.cv1 = Conv(c1, c_, 1, 1)
+        self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False)
+        self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False)
+        self.cv4 = Conv(2 * c_, c2, 1, 1)
+        self.bn = nn.BatchNorm2d(2 * c_)  # applied to cat(cv2, cv3)
+        self.act = nn.LeakyReLU(0.1, inplace=True)
+        self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))
+
+    def forward(self, x):
+        y1 = self.cv3(self.m(self.cv1(x)))
+        y2 = self.cv2(x)
+        return self.cv4(self.act(self.bn(torch.cat((y1, y2), dim=1))))
+
+
+class C3(nn.Module):
+    # CSP Bottleneck with 3 convolutions
+    def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):  # ch_in, ch_out, number, shortcut, groups, expansion
+        super().__init__()
+        c_ = int(c2 * e)  # hidden channels
+        self.cv1 = Conv(c1, c_, 1, 1)
+        self.cv2 = Conv(c1, c_, 1, 1)
+        self.cv3 = Conv(2 * c_, c2, 1)  # act=FReLU(c2)
+        self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))
+
+    def forward(self, x):
+        return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), dim=1))
+
+
+class ShuffleV2Block(nn.Module):
+    def __init__(self, inp, oup, stride):
+        super().__init__()
+
+        if not 1 <= stride <= 3:
+            raise ValueError("illegal stride value")
+        self.stride = stride
+
+        branch_features = oup // 2
+
+        if self.stride > 1:
+            self.branch1 = nn.Sequential(
+                self.depthwise_conv(inp, inp, kernel_size=3, stride=self.stride, padding=1),
+                nn.BatchNorm2d(inp),
+                nn.Conv2d(inp, branch_features, kernel_size=1, stride=1, padding=0, bias=False),
+                nn.BatchNorm2d(branch_features),
+                nn.SiLU(),
+            )
+        else:
+            self.branch1 = nn.Sequential()
+
+        self.branch2 = nn.Sequential(
+            nn.Conv2d(
+                inp if (self.stride > 1) else branch_features,
+                branch_features,
+                kernel_size=1,
+                stride=1,
+                padding=0,
+                bias=False,
+            ),
+            nn.BatchNorm2d(branch_features),
+            nn.SiLU(),
+            self.depthwise_conv(branch_features, branch_features, kernel_size=3, stride=self.stride, padding=1),
+            nn.BatchNorm2d(branch_features),
+            nn.Conv2d(branch_features, branch_features, kernel_size=1, stride=1, padding=0, bias=False),
+            nn.BatchNorm2d(branch_features),
+            nn.SiLU(),
+        )
+
+    @staticmethod
+    def depthwise_conv(i, o, kernel_size, stride=1, padding=0, bias=False):
+        return nn.Conv2d(i, o, kernel_size, stride, padding, bias=bias, groups=i)
+
+    def forward(self, x):
+        if self.stride == 1:
+            x1, x2 = x.chunk(2, dim=1)
+            out = torch.cat((x1, self.branch2(x2)), dim=1)
+        else:
+            out = torch.cat((self.branch1(x), self.branch2(x)), dim=1)
+        out = channel_shuffle(out, 2)
+        return out
+
+
+class SPP(nn.Module):
+    # Spatial pyramid pooling layer used in YOLOv3-SPP
+    def __init__(self, c1, c2, k=(5, 9, 13)):
+        super().__init__()
+        c_ = c1 // 2  # hidden channels
+        self.cv1 = Conv(c1, c_, 1, 1)
+        self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1)
+        self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k])
+
+    def forward(self, x):
+        x = self.cv1(x)
+        return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1))
+
+
+class Focus(nn.Module):
+    # Focus wh information into c-space
+    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):  # ch_in, ch_out, kernel, stride, padding, groups
+        super().__init__()
+        self.conv = Conv(c1 * 4, c2, k, s, p, g, act)
+
+    def forward(self, x):  # x(b,c,w,h) -> y(b,4c,w/2,h/2)
+        return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1))
+
+
+class Concat(nn.Module):
+    # Concatenate a list of tensors along dimension
+    def __init__(self, dimension=1):
+        super().__init__()
+        self.d = dimension
+
+    def forward(self, x):
+        return torch.cat(x, self.d)
+
+
+class NMS(nn.Module):
+    # Non-Maximum Suppression (NMS) module
+    conf = 0.25  # confidence threshold
+    iou = 0.45  # IoU threshold
+    classes = None  # (optional list) filter by class
+
+    def forward(self, x):
+        return non_max_suppression(x[0], conf_thres=self.conf, iou_thres=self.iou, classes=self.classes)
+
+
+class AutoShape(nn.Module):
+    # input-robust model wrapper for passing cv2/np/PIL/torch inputs. Includes preprocessing, inference and NMS
+    img_size = 640  # inference size (pixels)
+    conf = 0.25  # NMS confidence threshold
+    iou = 0.45  # NMS IoU threshold
+    classes = None  # (optional list) filter by class
+
+    def __init__(self, model):
+        super().__init__()
+        self.model = model.eval()
+
+    def autoshape(self):
+        print("autoShape already enabled, skipping... ")  # model already converted to model.autoshape()
+        return self
+
+    def forward(self, imgs, size=640, augment=False, profile=False):
+        # Inference from various sources. For height=720, width=1280, RGB images example inputs are:
+        #   OpenCV:          = cv2.imread('image.jpg')[:,:,::-1]  # HWC BGR to RGB x(720,1280,3)
+        #   PIL:             = Image.open('image.jpg')  # HWC x(720,1280,3)
+        #   numpy:           = np.zeros((720,1280,3))  # HWC
+        #   torch:           = torch.zeros(16,3,720,1280)  # BCHW
+        #   multiple:        = [Image.open('image1.jpg'), Image.open('image2.jpg'), ...]  # list of images
+
+        p = next(self.model.parameters())  # for device and type
+        if isinstance(imgs, torch.Tensor):  # torch
+            return self.model(imgs.to(p.device).type_as(p), augment, profile)  # inference
+
+        # Pre-process
+        n, imgs = (len(imgs), imgs) if isinstance(imgs, list) else (1, [imgs])  # number of images, list of images
+        shape0, shape1 = [], []  # image and inference shapes
+        for i, im in enumerate(imgs):
+            im = np.array(im)  # to numpy
+            if im.shape[0] < 5:  # image in CHW
+                im = im.transpose((1, 2, 0))  # reverse dataloader .transpose(2, 0, 1)
+            im = im[:, :, :3] if im.ndim == 3 else np.tile(im[:, :, None], 3)  # enforce 3ch input
+            s = im.shape[:2]  # HWC
+            shape0.append(s)  # image shape
+            g = size / max(s)  # gain
+            shape1.append([y * g for y in s])
+            imgs[i] = im  # update
+        shape1 = [make_divisible(x, int(self.stride.max())) for x in np.stack(shape1, 0).max(0)]  # inference shape
+        x = [letterbox(im, new_shape=shape1, auto=False)[0] for im in imgs]  # pad
+        x = np.stack(x, 0) if n > 1 else x[0][None]  # stack
+        x = np.ascontiguousarray(x.transpose((0, 3, 1, 2)))  # BHWC to BCHW
+        x = torch.from_numpy(x).to(p.device).type_as(p) / 255.0  # uint8 to fp16/32
+
+        # Inference
+        with torch.no_grad():
+            y = self.model(x, augment, profile)[0]  # forward
+        y = non_max_suppression(y, conf_thres=self.conf, iou_thres=self.iou, classes=self.classes)  # NMS
+
+        # Post-process
+        for i in range(n):
+            scale_coords(shape1, y[i][:, :4], shape0[i])
+
+        return Detections(imgs, y, self.names)
+
+
+class Detections:
+    # detections class for YOLOv5 inference results
+    def __init__(self, imgs, pred, names=None):
+        super().__init__()
+        d = pred[0].device  # device
+        gn = [torch.tensor([*(im.shape[i] for i in [1, 0, 1, 0]), 1.0, 1.0], device=d) for im in imgs]  # normalizations
+        self.imgs = imgs  # list of images as numpy arrays
+        self.pred = pred  # list of tensors pred[0] = (xyxy, conf, cls)
+        self.names = names  # class names
+        self.xyxy = pred  # xyxy pixels
+        self.xywh = [xyxy2xywh(x) for x in pred]  # xywh pixels
+        self.xyxyn = [x / g for x, g in zip(self.xyxy, gn)]  # xyxy normalized
+        self.xywhn = [x / g for x, g in zip(self.xywh, gn)]  # xywh normalized
+        self.n = len(self.pred)
+
+    def __len__(self):
+        return self.n
+
+    def tolist(self):
+        # return a list of Detections objects, i.e. 'for result in results.tolist():'
+        x = [Detections([self.imgs[i]], [self.pred[i]], self.names) for i in range(self.n)]
+        for d in x:
+            for k in ["imgs", "pred", "xyxy", "xyxyn", "xywh", "xywhn"]:
+                setattr(d, k, getattr(d, k)[0])  # pop out of list
+        return x

py/facelib/detection/yolov5face/models/experimental.py

@@ -0,0 +1,45 @@
+# # This file contains experimental modules
+
+import numpy as np
+import torch
+from torch import nn
+
+from .common import Conv
+
+
+class CrossConv(nn.Module):
+    # Cross Convolution Downsample
+    def __init__(self, c1, c2, k=3, s=1, g=1, e=1.0, shortcut=False):
+        # ch_in, ch_out, kernel, stride, groups, expansion, shortcut
+        super().__init__()
+        c_ = int(c2 * e)  # hidden channels
+        self.cv1 = Conv(c1, c_, (1, k), (1, s))
+        self.cv2 = Conv(c_, c2, (k, 1), (s, 1), g=g)
+        self.add = shortcut and c1 == c2
+
+    def forward(self, x):
+        return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))
+
+
+class MixConv2d(nn.Module):
+    # Mixed Depthwise Conv https://arxiv.org/abs/1907.09595
+    def __init__(self, c1, c2, k=(1, 3), s=1, equal_ch=True):
+        super().__init__()
+        groups = len(k)
+        if equal_ch:  # equal c_ per group
+            i = torch.linspace(0, groups - 1e-6, c2).floor()  # c2 indices
+            c_ = [(i == g).sum() for g in range(groups)]  # intermediate channels
+        else:  # equal weight.numel() per group
+            b = [c2] + [0] * groups
+            a = np.eye(groups + 1, groups, k=-1)
+            a -= np.roll(a, 1, axis=1)
+            a *= np.array(k) ** 2
+            a[0] = 1
+            c_ = np.linalg.lstsq(a, b, rcond=None)[0].round()  # solve for equal weight indices, ax = b
+
+        self.m = nn.ModuleList([nn.Conv2d(c1, int(c_[g]), k[g], s, k[g] // 2, bias=False) for g in range(groups)])
+        self.bn = nn.BatchNorm2d(c2)
+        self.act = nn.LeakyReLU(0.1, inplace=True)
+
+    def forward(self, x):
+        return x + self.act(self.bn(torch.cat([m(x) for m in self.m], 1)))

py/facelib/detection/yolov5face/models/yolo.py

@@ -0,0 +1,235 @@
+import math
+from copy import deepcopy
+from pathlib import Path
+
+import torch
+import yaml  # for torch hub
+from torch import nn
+
+from .common import (
+    C3,
+    NMS,
+    SPP,
+    AutoShape,
+    Bottleneck,
+    BottleneckCSP,
+    Concat,
+    Conv,
+    DWConv,
+    Focus,
+    ShuffleV2Block,
+    StemBlock,
+)
+from .experimental import CrossConv, MixConv2d
+from ..utils.autoanchor import check_anchor_order
+from ..utils.general import make_divisible
+from ..utils.torch_utils import copy_attr, fuse_conv_and_bn
+
+
+class Detect(nn.Module):
+    stride = None  # strides computed during build
+    export = False  # onnx export
+
+    def __init__(self, nc=80, anchors=(), ch=()):  # detection layer
+        super().__init__()
+        self.nc = nc  # number of classes
+        self.no = nc + 5 + 10  # number of outputs per anchor
+
+        self.nl = len(anchors)  # number of detection layers
+        self.na = len(anchors[0]) // 2  # number of anchors
+        self.grid = [torch.zeros(1)] * self.nl  # init grid
+        a = torch.tensor(anchors).float().view(self.nl, -1, 2)
+        self.register_buffer("anchors", a)  # shape(nl,na,2)
+        self.register_buffer("anchor_grid", a.clone().view(self.nl, 1, -1, 1, 1, 2))  # shape(nl,1,na,1,1,2)
+        self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch)  # output conv
+
+    def forward(self, x):
+        z = []  # inference output
+        if self.export:
+            for i in range(self.nl):
+                x[i] = self.m[i](x[i])
+            return x
+        for i in range(self.nl):
+            x[i] = self.m[i](x[i])  # conv
+            bs, _, ny, nx = x[i].shape  # x(bs,255,20,20) to x(bs,3,20,20,85)
+            x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
+
+            if not self.training:  # inference
+                if self.grid[i].shape[2:4] != x[i].shape[2:4]:
+                    self.grid[i] = self._make_grid(nx, ny).to(x[i].device)
+
+                y = torch.full_like(x[i], 0)
+                y[..., [0, 1, 2, 3, 4, 15]] = x[i][..., [0, 1, 2, 3, 4, 15]].sigmoid()
+                y[..., 5:15] = x[i][..., 5:15]
+
+                y[..., 0:2] = (y[..., 0:2] * 2.0 - 0.5 + self.grid[i].to(x[i].device)) * self.stride[i]  # xy
+                y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh
+
+                y[..., 5:7] = (
+                    y[..., 5:7] * self.anchor_grid[i] + self.grid[i].to(x[i].device) * self.stride[i]
+                )  # landmark x1 y1
+                y[..., 7:9] = (
+                    y[..., 7:9] * self.anchor_grid[i] + self.grid[i].to(x[i].device) * self.stride[i]
+                )  # landmark x2 y2
+                y[..., 9:11] = (
+                    y[..., 9:11] * self.anchor_grid[i] + self.grid[i].to(x[i].device) * self.stride[i]
+                )  # landmark x3 y3
+                y[..., 11:13] = (
+                    y[..., 11:13] * self.anchor_grid[i] + self.grid[i].to(x[i].device) * self.stride[i]
+                )  # landmark x4 y4
+                y[..., 13:15] = (
+                    y[..., 13:15] * self.anchor_grid[i] + self.grid[i].to(x[i].device) * self.stride[i]
+                )  # landmark x5 y5
+
+                z.append(y.view(bs, -1, self.no))
+
+        return x if self.training else (torch.cat(z, 1), x)
+
+    @staticmethod
+    def _make_grid(nx=20, ny=20):
+        # yv, xv = torch.meshgrid([torch.arange(ny), torch.arange(nx)], indexing="ij") # for pytorch>=1.10
+        yv, xv = torch.meshgrid([torch.arange(ny), torch.arange(nx)])
+        return torch.stack((xv, yv), 2).view((1, 1, ny, nx, 2)).float()
+
+
+class Model(nn.Module):
+    def __init__(self, cfg="yolov5s.yaml", ch=3, nc=None):  # model, input channels, number of classes
+        super().__init__()
+        self.yaml_file = Path(cfg).name
+        with Path(cfg).open(encoding="utf8") as f:
+            self.yaml = yaml.safe_load(f)  # model dict
+
+        # Define model
+        ch = self.yaml["ch"] = self.yaml.get("ch", ch)  # input channels
+        if nc and nc != self.yaml["nc"]:
+            self.yaml["nc"] = nc  # override yaml value
+
+        self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch])  # model, savelist
+        self.names = [str(i) for i in range(self.yaml["nc"])]  # default names
+
+        # Build strides, anchors
+        m = self.model[-1]  # Detect()
+        if isinstance(m, Detect):
+            s = 128  # 2x min stride
+            m.stride = torch.tensor([s / x.shape[-2] for x in self.forward(torch.zeros(1, ch, s, s))])  # forward
+            m.anchors /= m.stride.view(-1, 1, 1)
+            check_anchor_order(m)
+            self.stride = m.stride
+            self._initialize_biases()  # only run once
+
+    def forward(self, x):
+        return self.forward_once(x)  # single-scale inference, train
+
+    def forward_once(self, x):
+        y = []  # outputs
+        for m in self.model:
+            if m.f != -1:  # if not from previous layer
+                x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f]  # from earlier layers
+
+            x = m(x)  # run
+            y.append(x if m.i in self.save else None)  # save output
+
+        return x
+
+    def _initialize_biases(self, cf=None):  # initialize biases into Detect(), cf is class frequency
+        # https://arxiv.org/abs/1708.02002 section 3.3
+        m = self.model[-1]  # Detect() module
+        for mi, s in zip(m.m, m.stride):  # from
+            b = mi.bias.view(m.na, -1)  # conv.bias(255) to (3,85)
+            b.data[:, 4] += math.log(8 / (640 / s) ** 2)  # obj (8 objects per 640 image)
+            b.data[:, 5:] += math.log(0.6 / (m.nc - 0.99)) if cf is None else torch.log(cf / cf.sum())  # cls
+            mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)
+
+    def _print_biases(self):
+        m = self.model[-1]  # Detect() module
+        for mi in m.m:  # from
+            b = mi.bias.detach().view(m.na, -1).T  # conv.bias(255) to (3,85)
+            print(("%6g Conv2d.bias:" + "%10.3g" * 6) % (mi.weight.shape[1], *b[:5].mean(1).tolist(), b[5:].mean()))
+
+    def fuse(self):  # fuse model Conv2d() + BatchNorm2d() layers
+        print("Fusing layers... ")
+        for m in self.model.modules():
+            if isinstance(m, Conv) and hasattr(m, "bn"):
+                m.conv = fuse_conv_and_bn(m.conv, m.bn)  # update conv
+                delattr(m, "bn")  # remove batchnorm
+                m.forward = m.fuseforward  # update forward
+            elif type(m) is nn.Upsample:
+                m.recompute_scale_factor = None  # torch 1.11.0 compatibility
+        return self
+
+    def nms(self, mode=True):  # add or remove NMS module
+        present = isinstance(self.model[-1], NMS)  # last layer is NMS
+        if mode and not present:
+            print("Adding NMS... ")
+            m = NMS()  # module
+            m.f = -1  # from
+            m.i = self.model[-1].i + 1  # index
+            self.model.add_module(name=str(m.i), module=m)  # add
+            self.eval()
+        elif not mode and present:
+            print("Removing NMS... ")
+            self.model = self.model[:-1]  # remove
+        return self
+
+    def autoshape(self):  # add autoShape module
+        print("Adding autoShape... ")
+        m = AutoShape(self)  # wrap model
+        copy_attr(m, self, include=("yaml", "nc", "hyp", "names", "stride"), exclude=())  # copy attributes
+        return m
+
+
+def parse_model(d, ch):  # model_dict, input_channels(3)
+    anchors, nc, gd, gw = d["anchors"], d["nc"], d["depth_multiple"], d["width_multiple"]
+    na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors  # number of anchors
+    no = na * (nc + 5)  # number of outputs = anchors * (classes + 5)
+
+    layers, save, c2 = [], [], ch[-1]  # layers, savelist, ch out
+    for i, (f, n, m, args) in enumerate(d["backbone"] + d["head"]):  # from, number, module, args
+        m = eval(m) if isinstance(m, str) else m  # eval strings
+        for j, a in enumerate(args):
+            try:
+                args[j] = eval(a) if isinstance(a, str) else a  # eval strings
+            except:
+                pass
+
+        n = max(round(n * gd), 1) if n > 1 else n  # depth gain
+        if m in [
+            Conv,
+            Bottleneck,
+            SPP,
+            DWConv,
+            MixConv2d,
+            Focus,
+            CrossConv,
+            BottleneckCSP,
+            C3,
+            ShuffleV2Block,
+            StemBlock,
+        ]:
+            c1, c2 = ch[f], args[0]
+
+            c2 = make_divisible(c2 * gw, 8) if c2 != no else c2
+
+            args = [c1, c2, *args[1:]]
+            if m in [BottleneckCSP, C3]:
+                args.insert(2, n)
+                n = 1
+        elif m is nn.BatchNorm2d:
+            args = [ch[f]]
+        elif m is Concat:
+            c2 = sum(ch[-1 if x == -1 else x + 1] for x in f)
+        elif m is Detect:
+            args.append([ch[x + 1] for x in f])
+            if isinstance(args[1], int):  # number of anchors
+                args[1] = [list(range(args[1] * 2))] * len(f)
+        else:
+            c2 = ch[f]
+
+        m_ = nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args)  # module
+        t = str(m)[8:-2].replace("__main__.", "")  # module type
+        np = sum(x.numel() for x in m_.parameters())  # number params
+        m_.i, m_.f, m_.type, m_.np = i, f, t, np  # attach index, 'from' index, type, number params
+        save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1)  # append to savelist
+        layers.append(m_)
+        ch.append(c2)
+    return nn.Sequential(*layers), sorted(save)

py/facelib/detection/yolov5face/models/yolov5l.yaml

@@ -0,0 +1,47 @@
+# parameters
+nc: 1  # number of classes
+depth_multiple: 1.0  # model depth multiple
+width_multiple: 1.0  # layer channel multiple
+
+# anchors
+anchors:
+  - [4,5,  8,10,  13,16]  # P3/8
+  - [23,29,  43,55,  73,105]  # P4/16
+  - [146,217,  231,300,  335,433]  # P5/32
+
+# YOLOv5 backbone
+backbone:
+  # [from, number, module, args]
+  [[-1, 1, StemBlock, [64, 3, 2]],  # 0-P1/2
+   [-1, 3, C3, [128]],
+   [-1, 1, Conv, [256, 3, 2]],      # 2-P3/8
+   [-1, 9, C3, [256]],
+   [-1, 1, Conv, [512, 3, 2]],      # 4-P4/16
+   [-1, 9, C3, [512]],
+   [-1, 1, Conv, [1024, 3, 2]],     # 6-P5/32
+   [-1, 1, SPP, [1024, [3,5,7]]],
+   [-1, 3, C3, [1024, False]],      # 8
+  ]
+
+# YOLOv5 head
+head:
+  [[-1, 1, Conv, [512, 1, 1]],
+   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
+   [[-1, 5], 1, Concat, [1]],  # cat backbone P4
+   [-1, 3, C3, [512, False]],  # 12
+
+   [-1, 1, Conv, [256, 1, 1]],
+   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
+   [[-1, 3], 1, Concat, [1]],  # cat backbone P3
+   [-1, 3, C3, [256, False]],  # 16 (P3/8-small)
+
+   [-1, 1, Conv, [256, 3, 2]],
+   [[-1, 13], 1, Concat, [1]],  # cat head P4
+   [-1, 3, C3, [512, False]],  # 19 (P4/16-medium)
+
+   [-1, 1, Conv, [512, 3, 2]],
+   [[-1, 9], 1, Concat, [1]],  # cat head P5
+   [-1, 3, C3, [1024, False]],  # 22 (P5/32-large)
+
+   [[16, 19, 22], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
+  ]

py/facelib/detection/yolov5face/models/yolov5n.yaml

@@ -0,0 +1,45 @@
+# parameters
+nc: 1  # number of classes
+depth_multiple: 1.0  # model depth multiple
+width_multiple: 1.0  # layer channel multiple
+
+# anchors
+anchors:
+  - [4,5,  8,10,  13,16]  # P3/8
+  - [23,29,  43,55,  73,105]  # P4/16
+  - [146,217,  231,300,  335,433]  # P5/32
+
+# YOLOv5 backbone
+backbone:
+  # [from, number, module, args]
+  [[-1, 1, StemBlock, [32, 3, 2]],    # 0-P2/4
+   [-1, 1, ShuffleV2Block, [128, 2]], # 1-P3/8
+   [-1, 3, ShuffleV2Block, [128, 1]], # 2
+   [-1, 1, ShuffleV2Block, [256, 2]], # 3-P4/16
+   [-1, 7, ShuffleV2Block, [256, 1]], # 4
+   [-1, 1, ShuffleV2Block, [512, 2]], # 5-P5/32
+   [-1, 3, ShuffleV2Block, [512, 1]], # 6
+  ]
+
+# YOLOv5 head
+head:
+  [[-1, 1, Conv, [128, 1, 1]],
+   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
+   [[-1, 4], 1, Concat, [1]],  # cat backbone P4
+   [-1, 1, C3, [128, False]],  # 10
+
+   [-1, 1, Conv, [128, 1, 1]],
+   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
+   [[-1, 2], 1, Concat, [1]],  # cat backbone P3
+   [-1, 1, C3, [128, False]],  # 14 (P3/8-small)
+
+   [-1, 1, Conv, [128, 3, 2]],
+   [[-1, 11], 1, Concat, [1]],  # cat head P4
+   [-1, 1, C3, [128, False]],  # 17 (P4/16-medium)
+
+   [-1, 1, Conv, [128, 3, 2]],
+   [[-1, 7], 1, Concat, [1]],  # cat head P5
+   [-1, 1, C3, [128, False]],  # 20 (P5/32-large)
+
+   [[14, 17, 20], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
+  ]

py/facelib/detection/yolov5face/utils/autoanchor.py

@@ -0,0 +1,12 @@
+# Auto-anchor utils
+
+
+def check_anchor_order(m):
+    # Check anchor order against stride order for YOLOv5 Detect() module m, and correct if necessary
+    a = m.anchor_grid.prod(-1).view(-1)  # anchor area
+    da = a[-1] - a[0]  # delta a
+    ds = m.stride[-1] - m.stride[0]  # delta s
+    if da.sign() != ds.sign():  # same order
+        print("Reversing anchor order")
+        m.anchors[:] = m.anchors.flip(0)
+        m.anchor_grid[:] = m.anchor_grid.flip(0)

py/facelib/detection/yolov5face/utils/datasets.py

@@ -0,0 +1,35 @@
+import cv2
+import numpy as np
+
+
+def letterbox(img, new_shape=(640, 640), color=(114, 114, 114), auto=True, scale_fill=False, scaleup=True):
+    # Resize image to a 32-pixel-multiple rectangle https://github.com/ultralytics/yolov3/issues/232
+    shape = img.shape[:2]  # current shape [height, width]
+    if isinstance(new_shape, int):
+        new_shape = (new_shape, new_shape)
+
+    # Scale ratio (new / old)
+    r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
+    if not scaleup:  # only scale down, do not scale up (for better test mAP)
+        r = min(r, 1.0)
+
+    # Compute padding
+    ratio = r, r  # width, height ratios
+    new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
+    dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh padding
+    if auto:  # minimum rectangle
+        dw, dh = np.mod(dw, 64), np.mod(dh, 64)  # wh padding
+    elif scale_fill:  # stretch
+        dw, dh = 0.0, 0.0
+        new_unpad = (new_shape[1], new_shape[0])
+        ratio = new_shape[1] / shape[1], new_shape[0] / shape[0]  # width, height ratios
+
+    dw /= 2  # divide padding into 2 sides
+    dh /= 2
+
+    if shape[::-1] != new_unpad:  # resize
+        img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
+    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
+    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
+    img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # add border
+    return img, ratio, (dw, dh)

py/facelib/detection/yolov5face/utils/extract_ckpt.py

@@ -0,0 +1,5 @@
+import torch
+import sys
+sys.path.insert(0,'./facelib/detection/yolov5face')
+model = torch.load('facelib/detection/yolov5face/yolov5n-face.pt', map_location='cpu')['model']
+torch.save(model.state_dict(),'../../models/facedetection')

py/facelib/detection/yolov5face/utils/general.py

@@ -0,0 +1,271 @@
+import math
+import time
+
+import numpy as np
+import torch
+import torchvision
+
+
+def check_img_size(img_size, s=32):
+    # Verify img_size is a multiple of stride s
+    new_size = make_divisible(img_size, int(s))  # ceil gs-multiple
+    # if new_size != img_size:
+    #     print(f"WARNING: --img-size {img_size:g} must be multiple of max stride {s:g}, updating to {new_size:g}")
+    return new_size
+
+
+def make_divisible(x, divisor):
+    # Returns x evenly divisible by divisor
+    return math.ceil(x / divisor) * divisor
+
+
+def xyxy2xywh(x):
+    # Convert nx4 boxes from [x1, y1, x2, y2] to [x, y, w, h] where xy1=top-left, xy2=bottom-right
+    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
+    y[:, 0] = (x[:, 0] + x[:, 2]) / 2  # x center
+    y[:, 1] = (x[:, 1] + x[:, 3]) / 2  # y center
+    y[:, 2] = x[:, 2] - x[:, 0]  # width
+    y[:, 3] = x[:, 3] - x[:, 1]  # height
+    return y
+
+
+def xywh2xyxy(x):
+    # Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
+    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
+    y[:, 0] = x[:, 0] - x[:, 2] / 2  # top left x
+    y[:, 1] = x[:, 1] - x[:, 3] / 2  # top left y
+    y[:, 2] = x[:, 0] + x[:, 2] / 2  # bottom right x
+    y[:, 3] = x[:, 1] + x[:, 3] / 2  # bottom right y
+    return y
+
+
+def scale_coords(img1_shape, coords, img0_shape, ratio_pad=None):
+    # Rescale coords (xyxy) from img1_shape to img0_shape
+    if ratio_pad is None:  # calculate from img0_shape
+        gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1])  # gain  = old / new
+        pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2  # wh padding
+    else:
+        gain = ratio_pad[0][0]
+        pad = ratio_pad[1]
+
+    coords[:, [0, 2]] -= pad[0]  # x padding
+    coords[:, [1, 3]] -= pad[1]  # y padding
+    coords[:, :4] /= gain
+    clip_coords(coords, img0_shape)
+    return coords
+
+
+def clip_coords(boxes, img_shape):
+    # Clip bounding xyxy bounding boxes to image shape (height, width)
+    boxes[:, 0].clamp_(0, img_shape[1])  # x1
+    boxes[:, 1].clamp_(0, img_shape[0])  # y1
+    boxes[:, 2].clamp_(0, img_shape[1])  # x2
+    boxes[:, 3].clamp_(0, img_shape[0])  # y2
+
+
+def box_iou(box1, box2):
+    # https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py
+    """
+    Return intersection-over-union (Jaccard index) of boxes.
+    Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
+    Arguments:
+        box1 (Tensor[N, 4])
+        box2 (Tensor[M, 4])
+    Returns:
+        iou (Tensor[N, M]): the NxM matrix containing the pairwise
+            IoU values for every element in boxes1 and boxes2
+    """
+
+    def box_area(box):
+        return (box[2] - box[0]) * (box[3] - box[1])
+
+    area1 = box_area(box1.T)
+    area2 = box_area(box2.T)
+
+    inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) - torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2)
+    return inter / (area1[:, None] + area2 - inter)
+
+
+def non_max_suppression_face(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, labels=()):
+    """Performs Non-Maximum Suppression (NMS) on inference results
+    Returns:
+         detections with shape: nx6 (x1, y1, x2, y2, conf, cls)
+    """
+
+    nc = prediction.shape[2] - 15  # number of classes
+    xc = prediction[..., 4] > conf_thres  # candidates
+
+    # Settings
+    # (pixels) maximum box width and height
+    max_wh = 4096
+    time_limit = 10.0  # seconds to quit after
+    redundant = True  # require redundant detections
+    multi_label = nc > 1  # multiple labels per box (adds 0.5ms/img)
+    merge = False  # use merge-NMS
+
+    t = time.time()
+    output = [torch.zeros((0, 16), device=prediction.device)] * prediction.shape[0]
+    for xi, x in enumerate(prediction):  # image index, image inference
+        # Apply constraints
+        x = x[xc[xi]]  # confidence
+
+        # Cat apriori labels if autolabelling
+        if labels and len(labels[xi]):
+            label = labels[xi]
+            v = torch.zeros((len(label), nc + 15), device=x.device)
+            v[:, :4] = label[:, 1:5]  # box
+            v[:, 4] = 1.0  # conf
+            v[range(len(label)), label[:, 0].long() + 15] = 1.0  # cls
+            x = torch.cat((x, v), 0)
+
+        # If none remain process next image
+        if not x.shape[0]:
+            continue
+
+        # Compute conf
+        x[:, 15:] *= x[:, 4:5]  # conf = obj_conf * cls_conf
+
+        # Box (center x, center y, width, height) to (x1, y1, x2, y2)
+        box = xywh2xyxy(x[:, :4])
+
+        # Detections matrix nx6 (xyxy, conf, landmarks, cls)
+        if multi_label:
+            i, j = (x[:, 15:] > conf_thres).nonzero(as_tuple=False).T
+            x = torch.cat((box[i], x[i, j + 15, None], x[:, 5:15], j[:, None].float()), 1)
+        else:  # best class only
+            conf, j = x[:, 15:].max(1, keepdim=True)
+            x = torch.cat((box, conf, x[:, 5:15], j.float()), 1)[conf.view(-1) > conf_thres]
+
+        # Filter by class
+        if classes is not None:
+            x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]
+
+        # If none remain process next image
+        n = x.shape[0]  # number of boxes
+        if not n:
+            continue
+
+        # Batched NMS
+        c = x[:, 15:16] * (0 if agnostic else max_wh)  # classes
+        boxes, scores = x[:, :4] + c, x[:, 4]  # boxes (offset by class), scores
+        i = torchvision.ops.nms(boxes, scores, iou_thres)  # NMS
+
+        if merge and (1 < n < 3e3):  # Merge NMS (boxes merged using weighted mean)
+            # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)
+            iou = box_iou(boxes[i], boxes) > iou_thres  # iou matrix
+            weights = iou * scores[None]  # box weights
+            x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True)  # merged boxes
+            if redundant:
+                i = i[iou.sum(1) > 1]  # require redundancy
+
+        output[xi] = x[i]
+        if (time.time() - t) > time_limit:
+            break  # time limit exceeded
+
+    return output
+
+
+def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, labels=()):
+    """Performs Non-Maximum Suppression (NMS) on inference results
+
+    Returns:
+         detections with shape: nx6 (x1, y1, x2, y2, conf, cls)
+    """
+
+    nc = prediction.shape[2] - 5  # number of classes
+    xc = prediction[..., 4] > conf_thres  # candidates
+
+    # Settings
+    # (pixels) maximum box width and height
+    max_wh = 4096
+    time_limit = 10.0  # seconds to quit after
+    redundant = True  # require redundant detections
+    multi_label = nc > 1  # multiple labels per box (adds 0.5ms/img)
+    merge = False  # use merge-NMS
+
+    t = time.time()
+    output = [torch.zeros((0, 6), device=prediction.device)] * prediction.shape[0]
+    for xi, x in enumerate(prediction):  # image index, image inference
+        x = x[xc[xi]]  # confidence
+
+        # Cat apriori labels if autolabelling
+        if labels and len(labels[xi]):
+            label_id = labels[xi]
+            v = torch.zeros((len(label_id), nc + 5), device=x.device)
+            v[:, :4] = label_id[:, 1:5]  # box
+            v[:, 4] = 1.0  # conf
+            v[range(len(label_id)), label_id[:, 0].long() + 5] = 1.0  # cls
+            x = torch.cat((x, v), 0)
+
+        # If none remain process next image
+        if not x.shape[0]:
+            continue
+
+        # Compute conf
+        x[:, 5:] *= x[:, 4:5]  # conf = obj_conf * cls_conf
+
+        # Box (center x, center y, width, height) to (x1, y1, x2, y2)
+        box = xywh2xyxy(x[:, :4])
+
+        # Detections matrix nx6 (xyxy, conf, cls)
+        if multi_label:
+            i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).T
+            x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1)
+        else:  # best class only
+            conf, j = x[:, 5:].max(1, keepdim=True)
+            x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres]
+
+        # Filter by class
+        if classes is not None:
+            x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]
+
+        # Check shape
+        n = x.shape[0]  # number of boxes
+        if not n:  # no boxes
+            continue
+
+        x = x[x[:, 4].argsort(descending=True)]  # sort by confidence
+
+        # Batched NMS
+        c = x[:, 5:6] * (0 if agnostic else max_wh)  # classes
+        boxes, scores = x[:, :4] + c, x[:, 4]  # boxes (offset by class), scores
+        i = torchvision.ops.nms(boxes, scores, iou_thres)  # NMS
+        if merge and (1 < n < 3e3):  # Merge NMS (boxes merged using weighted mean)
+            # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)
+            iou = box_iou(boxes[i], boxes) > iou_thres  # iou matrix
+            weights = iou * scores[None]  # box weights
+            x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True)  # merged boxes
+            if redundant:
+                i = i[iou.sum(1) > 1]  # require redundancy
+
+        output[xi] = x[i]
+        if (time.time() - t) > time_limit:
+            print(f"WARNING: NMS time limit {time_limit}s exceeded")
+            break  # time limit exceeded
+
+    return output
+
+
+def scale_coords_landmarks(img1_shape, coords, img0_shape, ratio_pad=None):
+    # Rescale coords (xyxy) from img1_shape to img0_shape
+    if ratio_pad is None:  # calculate from img0_shape
+        gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1])  # gain  = old / new
+        pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2  # wh padding
+    else:
+        gain = ratio_pad[0][0]
+        pad = ratio_pad[1]
+
+    coords[:, [0, 2, 4, 6, 8]] -= pad[0]  # x padding
+    coords[:, [1, 3, 5, 7, 9]] -= pad[1]  # y padding
+    coords[:, :10] /= gain
+    coords[:, 0].clamp_(0, img0_shape[1])  # x1
+    coords[:, 1].clamp_(0, img0_shape[0])  # y1
+    coords[:, 2].clamp_(0, img0_shape[1])  # x2
+    coords[:, 3].clamp_(0, img0_shape[0])  # y2
+    coords[:, 4].clamp_(0, img0_shape[1])  # x3
+    coords[:, 5].clamp_(0, img0_shape[0])  # y3
+    coords[:, 6].clamp_(0, img0_shape[1])  # x4
+    coords[:, 7].clamp_(0, img0_shape[0])  # y4
+    coords[:, 8].clamp_(0, img0_shape[1])  # x5
+    coords[:, 9].clamp_(0, img0_shape[0])  # y5
+    return coords

py/facelib/detection/yolov5face/utils/torch_utils.py

@@ -0,0 +1,40 @@
+import torch
+from torch import nn
+
+
+def fuse_conv_and_bn(conv, bn):
+    # Fuse convolution and batchnorm layers https://tehnokv.com/posts/fusing-batchnorm-and-conv/
+    fusedconv = (
+        nn.Conv2d(
+            conv.in_channels,
+            conv.out_channels,
+            kernel_size=conv.kernel_size,
+            stride=conv.stride,
+            padding=conv.padding,
+            groups=conv.groups,
+            bias=True,
+        )
+        .requires_grad_(False)
+        .to(conv.weight.device)
+    )
+
+    # prepare filters
+    w_conv = conv.weight.clone().view(conv.out_channels, -1)
+    w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var)))
+    fusedconv.weight.copy_(torch.mm(w_bn, w_conv).view(fusedconv.weight.size()))
+
+    # prepare spatial bias
+    b_conv = torch.zeros(conv.weight.size(0), device=conv.weight.device) if conv.bias is None else conv.bias
+    b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(torch.sqrt(bn.running_var + bn.eps))
+    fusedconv.bias.copy_(torch.mm(w_bn, b_conv.reshape(-1, 1)).reshape(-1) + b_bn)
+
+    return fusedconv
+
+
+def copy_attr(a, b, include=(), exclude=()):
+    # Copy attributes from b to a, options to only include [...] and to exclude [...]
+    for k, v in b.__dict__.items():
+        if (include and k not in include) or k.startswith("_") or k in exclude:
+            continue
+
+        setattr(a, k, v)

py/facelib/parsing/__init__.py

@@ -0,0 +1,23 @@
+import torch
+
+from ..utils import load_file_from_url
+from .bisenet import BiSeNet
+from .parsenet import ParseNet
+
+
+def init_parsing_model(model_name='bisenet', half=False, device='cuda'):
+    if model_name == 'bisenet':
+        model = BiSeNet(num_class=19)
+        model_url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/parsing_bisenet.pth'
+    elif model_name == 'parsenet':
+        model = ParseNet(in_size=512, out_size=512, parsing_ch=19)
+        model_url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/parsing_parsenet.pth'
+    else:
+        raise NotImplementedError(f'{model_name} is not implemented.')
+
+    model_path = load_file_from_url(url=model_url, model_dir='../../models/facedetection', progress=True, file_name=None)
+    load_net = torch.load(model_path, map_location=lambda storage, loc: storage)
+    model.load_state_dict(load_net, strict=True)
+    model.eval()
+    model = model.to(device)
+    return model

py/facelib/parsing/bisenet.py

@@ -0,0 +1,140 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .resnet import ResNet18
+
+
+class ConvBNReLU(nn.Module):
+
+    def __init__(self, in_chan, out_chan, ks=3, stride=1, padding=1):
+        super(ConvBNReLU, self).__init__()
+        self.conv = nn.Conv2d(in_chan, out_chan, kernel_size=ks, stride=stride, padding=padding, bias=False)
+        self.bn = nn.BatchNorm2d(out_chan)
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = F.relu(self.bn(x))
+        return x
+
+
+class BiSeNetOutput(nn.Module):
+
+    def __init__(self, in_chan, mid_chan, num_class):
+        super(BiSeNetOutput, self).__init__()
+        self.conv = ConvBNReLU(in_chan, mid_chan, ks=3, stride=1, padding=1)
+        self.conv_out = nn.Conv2d(mid_chan, num_class, kernel_size=1, bias=False)
+
+    def forward(self, x):
+        feat = self.conv(x)
+        out = self.conv_out(feat)
+        return out, feat
+
+
+class AttentionRefinementModule(nn.Module):
+
+    def __init__(self, in_chan, out_chan):
+        super(AttentionRefinementModule, self).__init__()
+        self.conv = ConvBNReLU(in_chan, out_chan, ks=3, stride=1, padding=1)
+        self.conv_atten = nn.Conv2d(out_chan, out_chan, kernel_size=1, bias=False)
+        self.bn_atten = nn.BatchNorm2d(out_chan)
+        self.sigmoid_atten = nn.Sigmoid()
+
+    def forward(self, x):
+        feat = self.conv(x)
+        atten = F.avg_pool2d(feat, feat.size()[2:])
+        atten = self.conv_atten(atten)
+        atten = self.bn_atten(atten)
+        atten = self.sigmoid_atten(atten)
+        out = torch.mul(feat, atten)
+        return out
+
+
+class ContextPath(nn.Module):
+
+    def __init__(self):
+        super(ContextPath, self).__init__()
+        self.resnet = ResNet18()
+        self.arm16 = AttentionRefinementModule(256, 128)
+        self.arm32 = AttentionRefinementModule(512, 128)
+        self.conv_head32 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1)
+        self.conv_head16 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1)
+        self.conv_avg = ConvBNReLU(512, 128, ks=1, stride=1, padding=0)
+
+    def forward(self, x):
+        feat8, feat16, feat32 = self.resnet(x)
+        h8, w8 = feat8.size()[2:]
+        h16, w16 = feat16.size()[2:]
+        h32, w32 = feat32.size()[2:]
+
+        avg = F.avg_pool2d(feat32, feat32.size()[2:])
+        avg = self.conv_avg(avg)
+        avg_up = F.interpolate(avg, (h32, w32), mode='nearest')
+
+        feat32_arm = self.arm32(feat32)
+        feat32_sum = feat32_arm + avg_up
+        feat32_up = F.interpolate(feat32_sum, (h16, w16), mode='nearest')
+        feat32_up = self.conv_head32(feat32_up)
+
+        feat16_arm = self.arm16(feat16)
+        feat16_sum = feat16_arm + feat32_up
+        feat16_up = F.interpolate(feat16_sum, (h8, w8), mode='nearest')
+        feat16_up = self.conv_head16(feat16_up)
+
+        return feat8, feat16_up, feat32_up  # x8, x8, x16
+
+
+class FeatureFusionModule(nn.Module):
+
+    def __init__(self, in_chan, out_chan):
+        super(FeatureFusionModule, self).__init__()
+        self.convblk = ConvBNReLU(in_chan, out_chan, ks=1, stride=1, padding=0)
+        self.conv1 = nn.Conv2d(out_chan, out_chan // 4, kernel_size=1, stride=1, padding=0, bias=False)
+        self.conv2 = nn.Conv2d(out_chan // 4, out_chan, kernel_size=1, stride=1, padding=0, bias=False)
+        self.relu = nn.ReLU(inplace=True)
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, fsp, fcp):
+        fcat = torch.cat([fsp, fcp], dim=1)
+        feat = self.convblk(fcat)
+        atten = F.avg_pool2d(feat, feat.size()[2:])
+        atten = self.conv1(atten)
+        atten = self.relu(atten)
+        atten = self.conv2(atten)
+        atten = self.sigmoid(atten)
+        feat_atten = torch.mul(feat, atten)
+        feat_out = feat_atten + feat
+        return feat_out
+
+
+class BiSeNet(nn.Module):
+
+    def __init__(self, num_class):
+        super(BiSeNet, self).__init__()
+        self.cp = ContextPath()
+        self.ffm = FeatureFusionModule(256, 256)
+        self.conv_out = BiSeNetOutput(256, 256, num_class)
+        self.conv_out16 = BiSeNetOutput(128, 64, num_class)
+        self.conv_out32 = BiSeNetOutput(128, 64, num_class)
+
+    def forward(self, x, return_feat=False):
+        h, w = x.size()[2:]
+        feat_res8, feat_cp8, feat_cp16 = self.cp(x)  # return res3b1 feature
+        feat_sp = feat_res8  # replace spatial path feature with res3b1 feature
+        feat_fuse = self.ffm(feat_sp, feat_cp8)
+
+        out, feat = self.conv_out(feat_fuse)
+        out16, feat16 = self.conv_out16(feat_cp8)
+        out32, feat32 = self.conv_out32(feat_cp16)
+
+        out = F.interpolate(out, (h, w), mode='bilinear', align_corners=True)
+        out16 = F.interpolate(out16, (h, w), mode='bilinear', align_corners=True)
+        out32 = F.interpolate(out32, (h, w), mode='bilinear', align_corners=True)
+
+        if return_feat:
+            feat = F.interpolate(feat, (h, w), mode='bilinear', align_corners=True)
+            feat16 = F.interpolate(feat16, (h, w), mode='bilinear', align_corners=True)
+            feat32 = F.interpolate(feat32, (h, w), mode='bilinear', align_corners=True)
+            return out, out16, out32, feat, feat16, feat32
+        else:
+            return out, out16, out32

py/facelib/parsing/parsenet.py

@@ -1,18 +1,19 @@
-# -*- coding: utf-8 -*-
-import torch
-import torch.nn as nn
-from torch.nn.parameter import Parameter
-from torch.nn import functional as F
+"""Modified from https://github.com/chaofengc/PSFRGAN
+"""
 import numpy as np
+import torch.nn as nn
+from torch.nn import functional as F
+
 
 class NormLayer(nn.Module):
     """Normalization Layers.
-    ------------
-    # Arguments
-        - channels: input channels, for batch norm and instance norm.
-        - input_size: input shape without batch size, for layer norm.
+
+    Args:
+        channels: input channels, for batch norm and instance norm.
+        input_size: input shape without batch size, for layer norm.
     """
-    def __init__(self, channels, normalize_shape=None, norm_type='bn', ref_channels=None):
+
+    def __init__(self, channels, normalize_shape=None, norm_type='bn'):
         super(NormLayer, self).__init__()
         norm_type = norm_type.lower()
         self.norm_type = norm_type
@@ -27,9 +28,9 @@ class NormLayer(nn.Module):
         elif norm_type == 'layer':
             self.norm = nn.LayerNorm(normalize_shape)
         elif norm_type == 'none':
-            self.norm = lambda x: x*1.0
+            self.norm = lambda x: x * 1.0
         else:
-            assert 1==0, 'Norm type {} not support.'.format(norm_type)
+            assert 1 == 0, f'Norm type {norm_type} not support.'
 
     def forward(self, x, ref=None):
         if self.norm_type == 'spade':
@@ -40,15 +41,16 @@ class NormLayer(nn.Module):
 
 class ReluLayer(nn.Module):
     """Relu Layer.
-    ------------
-    # Arguments
-        - relu type: type of relu layer, candidates are
+
+    Args:
+        relu type: type of relu layer, candidates are
             - ReLU
             - LeakyReLU: default relu slope 0.2
-            - PRelu 
+            - PRelu
             - SELU
             - none: direct pass
     """
+
     def __init__(self, channels, relu_type='relu'):
         super(ReluLayer, self).__init__()
         relu_type = relu_type.lower()
@@ -61,29 +63,38 @@ class ReluLayer(nn.Module):
         elif relu_type == 'selu':
             self.func = nn.SELU(True)
         elif relu_type == 'none':
-            self.func = lambda x: x*1.0
+            self.func = lambda x: x * 1.0
         else:
-            assert 1==0, 'Relu type {} not support.'.format(relu_type)
+            assert 1 == 0, f'Relu type {relu_type} not support.'
 
     def forward(self, x):
         return self.func(x)
 
 
 class ConvLayer(nn.Module):
-    def __init__(self, in_channels, out_channels, kernel_size=3, scale='none', norm_type='none', relu_type='none', use_pad=True, bias=True):
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size=3,
+                 scale='none',
+                 norm_type='none',
+                 relu_type='none',
+                 use_pad=True,
+                 bias=True):
         super(ConvLayer, self).__init__()
         self.use_pad = use_pad
         self.norm_type = norm_type
         if norm_type in ['bn']:
             bias = False
-        
+
         stride = 2 if scale == 'down' else 1
 
         self.scale_func = lambda x: x
         if scale == 'up':
             self.scale_func = lambda x: nn.functional.interpolate(x, scale_factor=2, mode='nearest')
 
-        self.reflection_pad = nn.ReflectionPad2d(int(np.ceil((kernel_size - 1.)/2))) 
+        self.reflection_pad = nn.ReflectionPad2d(int(np.ceil((kernel_size - 1.) / 2)))
         self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride, bias=bias)
 
         self.relu = ReluLayer(out_channels, relu_type)
@@ -103,6 +114,7 @@ class ResidualBlock(nn.Module):
     """
     Residual block recommended in: http://torch.ch/blog/2016/02/04/resnets.html
     """
+
     def __init__(self, c_in, c_out, relu_type='prelu', norm_type='bn', scale='none'):
         super(ResidualBlock, self).__init__()
 
@@ -110,18 +122,73 @@ class ResidualBlock(nn.Module):
             self.shortcut_func = lambda x: x
         else:
             self.shortcut_func = ConvLayer(c_in, c_out, 3, scale)
-        
+
         scale_config_dict = {'down': ['none', 'down'], 'up': ['up', 'none'], 'none': ['none', 'none']}
         scale_conf = scale_config_dict[scale]
 
-        self.conv1 = ConvLayer(c_in, c_out, 3, scale_conf[0], norm_type=norm_type, relu_type=relu_type) 
+        self.conv1 = ConvLayer(c_in, c_out, 3, scale_conf[0], norm_type=norm_type, relu_type=relu_type)
         self.conv2 = ConvLayer(c_out, c_out, 3, scale_conf[1], norm_type=norm_type, relu_type='none')
-  
+
     def forward(self, x):
         identity = self.shortcut_func(x)
 
         res = self.conv1(x)
         res = self.conv2(res)
         return identity + res
-        
 
+
+class ParseNet(nn.Module):
+
+    def __init__(self,
+                 in_size=128,
+                 out_size=128,
+                 min_feat_size=32,
+                 base_ch=64,
+                 parsing_ch=19,
+                 res_depth=10,
+                 relu_type='LeakyReLU',
+                 norm_type='bn',
+                 ch_range=[32, 256]):
+        super().__init__()
+        self.res_depth = res_depth
+        act_args = {'norm_type': norm_type, 'relu_type': relu_type}
+        min_ch, max_ch = ch_range
+
+        ch_clip = lambda x: max(min_ch, min(x, max_ch))  # noqa: E731
+        min_feat_size = min(in_size, min_feat_size)
+
+        down_steps = int(np.log2(in_size // min_feat_size))
+        up_steps = int(np.log2(out_size // min_feat_size))
+
+        # =============== define encoder-body-decoder ====================
+        self.encoder = []
+        self.encoder.append(ConvLayer(3, base_ch, 3, 1))
+        head_ch = base_ch
+        for i in range(down_steps):
+            cin, cout = ch_clip(head_ch), ch_clip(head_ch * 2)
+            self.encoder.append(ResidualBlock(cin, cout, scale='down', **act_args))
+            head_ch = head_ch * 2
+
+        self.body = []
+        for i in range(res_depth):
+            self.body.append(ResidualBlock(ch_clip(head_ch), ch_clip(head_ch), **act_args))
+
+        self.decoder = []
+        for i in range(up_steps):
+            cin, cout = ch_clip(head_ch), ch_clip(head_ch // 2)
+            self.decoder.append(ResidualBlock(cin, cout, scale='up', **act_args))
+            head_ch = head_ch // 2
+
+        self.encoder = nn.Sequential(*self.encoder)
+        self.body = nn.Sequential(*self.body)
+        self.decoder = nn.Sequential(*self.decoder)
+        self.out_img_conv = ConvLayer(ch_clip(head_ch), 3)
+        self.out_mask_conv = ConvLayer(ch_clip(head_ch), parsing_ch)
+
+    def forward(self, x):
+        feat = self.encoder(x)
+        x = feat + self.body(feat)
+        x = self.decoder(x)
+        out_img = self.out_img_conv(x)
+        out_mask = self.out_mask_conv(x)
+        return out_mask, out_img

py/facelib/parsing/resnet.py

@@ -1,23 +1,14 @@
-#!/usr/bin/python
-# -*- encoding: utf-8 -*-
-
-import torch
 import torch.nn as nn
 import torch.nn.functional as F
-import torch.utils.model_zoo as modelzoo
-
-# from modules.bn import InPlaceABNSync as BatchNorm2d
-
-resnet18_url = 'https://download.pytorch.org/models/resnet18-5c106cde.pth'
 
 
 def conv3x3(in_planes, out_planes, stride=1):
     """3x3 convolution with padding"""
-    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
-                     padding=1, bias=False)
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False)
 
 
 class BasicBlock(nn.Module):
+
     def __init__(self, in_chan, out_chan, stride=1):
         super(BasicBlock, self).__init__()
         self.conv1 = conv3x3(in_chan, out_chan, stride)
@@ -28,10 +19,9 @@ class BasicBlock(nn.Module):
         self.downsample = None
         if in_chan != out_chan or stride != 1:
             self.downsample = nn.Sequential(
-                nn.Conv2d(in_chan, out_chan,
-                          kernel_size=1, stride=stride, bias=False),
+                nn.Conv2d(in_chan, out_chan, kernel_size=1, stride=stride, bias=False),
                 nn.BatchNorm2d(out_chan),
-                )
+            )
 
     def forward(self, x):
         residual = self.conv1(x)
@@ -50,23 +40,22 @@ class BasicBlock(nn.Module):
 
 def create_layer_basic(in_chan, out_chan, bnum, stride=1):
     layers = [BasicBlock(in_chan, out_chan, stride=stride)]
-    for i in range(bnum-1):
+    for i in range(bnum - 1):
         layers.append(BasicBlock(out_chan, out_chan, stride=1))
     return nn.Sequential(*layers)
 
 
-class Resnet18(nn.Module):
+class ResNet18(nn.Module):
+
     def __init__(self):
-        super(Resnet18, self).__init__()
-        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
-                               bias=False)
+        super(ResNet18, self).__init__()
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
         self.bn1 = nn.BatchNorm2d(64)
         self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
         self.layer1 = create_layer_basic(64, 64, bnum=2, stride=1)
         self.layer2 = create_layer_basic(64, 128, bnum=2, stride=2)
         self.layer3 = create_layer_basic(128, 256, bnum=2, stride=2)
         self.layer4 = create_layer_basic(256, 512, bnum=2, stride=2)
-        self.init_weight()
 
     def forward(self, x):
         x = self.conv1(x)
@@ -74,36 +63,7 @@ class Resnet18(nn.Module):
         x = self.maxpool(x)
 
         x = self.layer1(x)
-        feat8 = self.layer2(x) # 1/8
-        feat16 = self.layer3(feat8) # 1/16
-        feat32 = self.layer4(feat16) # 1/32
+        feat8 = self.layer2(x)  # 1/8
+        feat16 = self.layer3(feat8)  # 1/16
+        feat32 = self.layer4(feat16)  # 1/32
         return feat8, feat16, feat32
-
-    def init_weight(self):
-        state_dict = modelzoo.load_url(resnet18_url)
-        self_state_dict = self.state_dict()
-        for k, v in state_dict.items():
-            if 'fc' in k: continue
-            self_state_dict.update({k: v})
-        self.load_state_dict(self_state_dict)
-
-    def get_params(self):
-        wd_params, nowd_params = [], []
-        for name, module in self.named_modules():
-            if isinstance(module, (nn.Linear, nn.Conv2d)):
-                wd_params.append(module.weight)
-                if not module.bias is None:
-                    nowd_params.append(module.bias)
-            elif isinstance(module,  nn.BatchNorm2d):
-                nowd_params += list(module.parameters())
-        return wd_params, nowd_params
-
-
-if __name__ == "__main__":
-    net = Resnet18()
-    x = torch.randn(16, 3, 224, 224)
-    out = net(x)
-    print(out[0].size())
-    print(out[1].size())
-    print(out[2].size())
-    net.get_params()

py/facelib/utils/__init__.py

@@ -0,0 +1,7 @@
+from .face_utils import align_crop_face_landmarks, compute_increased_bbox, get_valid_bboxes, paste_face_back
+from .misc import img2tensor, load_file_from_url, download_pretrained_models, scandir
+
+__all__ = [
+    'align_crop_face_landmarks', 'compute_increased_bbox', 'get_valid_bboxes', 'load_file_from_url', 
+    'download_pretrained_models', 'paste_face_back', 'img2tensor', 'scandir'
+]

py/facelib/utils/face_restoration_helper.py

@@ -0,0 +1,455 @@
+import cv2
+import numpy as np
+import os
+import torch
+from torchvision.transforms.functional import normalize
+
+from ..detection import init_detection_model
+from ..parsing import init_parsing_model
+from .misc import img2tensor, imwrite
+
+
+def get_largest_face(det_faces, h, w):
+
+    def get_location(val, length):
+        if val < 0:
+            return 0
+        elif val > length:
+            return length
+        else:
+            return val
+
+    face_areas = []
+    for det_face in det_faces:
+        left = get_location(det_face[0], w)
+        right = get_location(det_face[2], w)
+        top = get_location(det_face[1], h)
+        bottom = get_location(det_face[3], h)
+        face_area = (right - left) * (bottom - top)
+        face_areas.append(face_area)
+    largest_idx = face_areas.index(max(face_areas))
+    return det_faces[largest_idx], largest_idx
+
+
+def get_center_face(det_faces, h=0, w=0, center=None):
+    if center is not None:
+        center = np.array(center)
+    else:
+        center = np.array([w / 2, h / 2])
+    center_dist = []
+    for det_face in det_faces:
+        face_center = np.array([(det_face[0] + det_face[2]) / 2, (det_face[1] + det_face[3]) / 2])
+        dist = np.linalg.norm(face_center - center)
+        center_dist.append(dist)
+    center_idx = center_dist.index(min(center_dist))
+    return det_faces[center_idx], center_idx
+
+
+class FaceRestoreHelper(object):
+    """Helper for the face restoration pipeline (base class)."""
+
+    def __init__(self,
+                 upscale_factor,
+                 face_size=512,
+                 crop_ratio=(1, 1),
+                 det_model='retinaface_resnet50',
+                 save_ext='png',
+                 template_3points=False,
+                 pad_blur=False,
+                 use_parse=False,
+                 device=None):
+        self.template_3points = template_3points  # improve robustness
+        self.upscale_factor = upscale_factor
+        # the cropped face ratio based on the square face
+        self.crop_ratio = crop_ratio  # (h, w)
+        assert (self.crop_ratio[0] >= 1 and self.crop_ratio[1] >= 1), 'crop ration only supports >=1'
+        self.face_size = (int(face_size * self.crop_ratio[1]), int(face_size * self.crop_ratio[0]))
+
+        if self.template_3points:
+            self.face_template = np.array([[192, 240], [319, 240], [257, 371]])
+        else:
+            # standard 5 landmarks for FFHQ faces with 512 x 512 
+            # facexlib
+            self.face_template = np.array([[192.98138, 239.94708], [318.90277, 240.1936], [256.63416, 314.01935],
+                                           [201.26117, 371.41043], [313.08905, 371.15118]])
+
+            # dlib: left_eye: 36:41  right_eye: 42:47  nose: 30,32,33,34  left mouth corner: 48  right mouth corner: 54
+            # self.face_template = np.array([[193.65928, 242.98541], [318.32558, 243.06108], [255.67984, 328.82894],
+            #                                 [198.22603, 372.82502], [313.91018, 372.75659]])
+
+
+        self.face_template = self.face_template * (face_size / 512.0)
+        if self.crop_ratio[0] > 1:
+            self.face_template[:, 1] += face_size * (self.crop_ratio[0] - 1) / 2
+        if self.crop_ratio[1] > 1:
+            self.face_template[:, 0] += face_size * (self.crop_ratio[1] - 1) / 2
+        self.save_ext = save_ext
+        self.pad_blur = pad_blur
+        if self.pad_blur is True:
+            self.template_3points = False
+
+        self.all_landmarks_5 = []
+        self.det_faces = []
+        self.affine_matrices = []
+        self.inverse_affine_matrices = []
+        self.cropped_faces = []
+        self.restored_faces = []
+        self.pad_input_imgs = []
+
+        if device is None:
+            self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        else:
+            self.device = device
+
+        # init face detection model
+        self.face_det = init_detection_model(det_model, half=False, device=self.device)
+
+        # init face parsing model
+        self.use_parse = use_parse
+        self.face_parse = init_parsing_model(model_name='parsenet', device=self.device)
+
+    def set_upscale_factor(self, upscale_factor):
+        self.upscale_factor = upscale_factor
+
+    def read_image(self, img):
+        """img can be image path or cv2 loaded image."""
+        # self.input_img is Numpy array, (h, w, c), BGR, uint8, [0, 255]
+        if isinstance(img, str):
+            img = cv2.imread(img)
+
+        if np.max(img) > 256:  # 16-bit image
+            img = img / 65535 * 255
+        if len(img.shape) == 2:  # gray image
+            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
+        elif img.shape[2] == 4:  # BGRA image with alpha channel
+            img = img[:, :, 0:3]
+
+        self.input_img = img
+
+        if min(self.input_img.shape[:2])<512:
+            f = 512.0/min(self.input_img.shape[:2])
+            self.input_img = cv2.resize(self.input_img, (0,0), fx=f, fy=f, interpolation=cv2.INTER_LINEAR)
+
+    def get_face_landmarks_5(self,
+                             only_keep_largest=False,
+                             only_center_face=False,
+                             resize=None,
+                             blur_ratio=0.01,
+                             eye_dist_threshold=None):
+        if resize is None:
+            scale = 1
+            input_img = self.input_img
+        else:
+            h, w = self.input_img.shape[0:2]
+            scale = resize / min(h, w)
+            scale = max(1, scale) # always scale up
+            h, w = int(h * scale), int(w * scale)
+            interp = cv2.INTER_AREA if scale < 1 else cv2.INTER_LINEAR
+            input_img = cv2.resize(self.input_img, (w, h), interpolation=interp)
+
+        with torch.no_grad():
+            bboxes = self.face_det.detect_faces(input_img)
+
+        if bboxes is None or bboxes.shape[0] == 0:
+            return 0
+        else:
+            bboxes = bboxes / scale
+
+        for bbox in bboxes:
+            # remove faces with too small eye distance: side faces or too small faces
+            eye_dist = np.linalg.norm([bbox[6] - bbox[8], bbox[7] - bbox[9]])
+            if eye_dist_threshold is not None and (eye_dist < eye_dist_threshold):
+                continue
+
+            if self.template_3points:
+                landmark = np.array([[bbox[i], bbox[i + 1]] for i in range(5, 11, 2)])
+            else:
+                landmark = np.array([[bbox[i], bbox[i + 1]] for i in range(5, 15, 2)])
+            self.all_landmarks_5.append(landmark)
+            self.det_faces.append(bbox[0:5])
+            
+        if len(self.det_faces) == 0:
+            return 0
+        if only_keep_largest:
+            h, w, _ = self.input_img.shape
+            self.det_faces, largest_idx = get_largest_face(self.det_faces, h, w)
+            self.all_landmarks_5 = [self.all_landmarks_5[largest_idx]]
+        elif only_center_face:
+            h, w, _ = self.input_img.shape
+            self.det_faces, center_idx = get_center_face(self.det_faces, h, w)
+            self.all_landmarks_5 = [self.all_landmarks_5[center_idx]]
+
+        # pad blurry images
+        if self.pad_blur:
+            self.pad_input_imgs = []
+            for landmarks in self.all_landmarks_5:
+                # get landmarks
+                eye_left = landmarks[0, :]
+                eye_right = landmarks[1, :]
+                eye_avg = (eye_left + eye_right) * 0.5
+                mouth_avg = (landmarks[3, :] + landmarks[4, :]) * 0.5
+                eye_to_eye = eye_right - eye_left
+                eye_to_mouth = mouth_avg - eye_avg
+
+                # Get the oriented crop rectangle
+                # x: half width of the oriented crop rectangle
+                x = eye_to_eye - np.flipud(eye_to_mouth) * [-1, 1]
+                #  - np.flipud(eye_to_mouth) * [-1, 1]: rotate 90 clockwise
+                # norm with the hypotenuse: get the direction
+                x /= np.hypot(*x)  # get the hypotenuse of a right triangle
+                rect_scale = 1.5
+                x *= max(np.hypot(*eye_to_eye) * 2.0 * rect_scale, np.hypot(*eye_to_mouth) * 1.8 * rect_scale)
+                # y: half height of the oriented crop rectangle
+                y = np.flipud(x) * [-1, 1]
+
+                # c: center
+                c = eye_avg + eye_to_mouth * 0.1
+                # quad: (left_top, left_bottom, right_bottom, right_top)
+                quad = np.stack([c - x - y, c - x + y, c + x + y, c + x - y])
+                # qsize: side length of the square
+                qsize = np.hypot(*x) * 2
+                border = max(int(np.rint(qsize * 0.1)), 3)
+
+                # get pad
+                # pad: (width_left, height_top, width_right, height_bottom)
+                pad = (int(np.floor(min(quad[:, 0]))), int(np.floor(min(quad[:, 1]))), int(np.ceil(max(quad[:, 0]))),
+                       int(np.ceil(max(quad[:, 1]))))
+                pad = [
+                    max(-pad[0] + border, 1),
+                    max(-pad[1] + border, 1),
+                    max(pad[2] - self.input_img.shape[0] + border, 1),
+                    max(pad[3] - self.input_img.shape[1] + border, 1)
+                ]
+
+                if max(pad) > 1:
+                    # pad image
+                    pad_img = np.pad(self.input_img, ((pad[1], pad[3]), (pad[0], pad[2]), (0, 0)), 'reflect')
+                    # modify landmark coords
+                    landmarks[:, 0] += pad[0]
+                    landmarks[:, 1] += pad[1]
+                    # blur pad images
+                    h, w, _ = pad_img.shape
+                    y, x, _ = np.ogrid[:h, :w, :1]
+                    mask = np.maximum(1.0 - np.minimum(np.float32(x) / pad[0],
+                                                       np.float32(w - 1 - x) / pad[2]),
+                                      1.0 - np.minimum(np.float32(y) / pad[1],
+                                                       np.float32(h - 1 - y) / pad[3]))
+                    blur = int(qsize * blur_ratio)
+                    if blur % 2 == 0:
+                        blur += 1
+                    blur_img = cv2.boxFilter(pad_img, 0, ksize=(blur, blur))
+                    # blur_img = cv2.GaussianBlur(pad_img, (blur, blur), 0)
+
+                    pad_img = pad_img.astype('float32')
+                    pad_img += (blur_img - pad_img) * np.clip(mask * 3.0 + 1.0, 0.0, 1.0)
+                    pad_img += (np.median(pad_img, axis=(0, 1)) - pad_img) * np.clip(mask, 0.0, 1.0)
+                    pad_img = np.clip(pad_img, 0, 255)  # float32, [0, 255]
+                    self.pad_input_imgs.append(pad_img)
+                else:
+                    self.pad_input_imgs.append(np.copy(self.input_img))
+
+        return len(self.all_landmarks_5)
+
+    def align_warp_face(self, save_cropped_path=None, border_mode='constant'):
+        """Align and warp faces with face template.
+        """
+        if self.pad_blur:
+            assert len(self.pad_input_imgs) == len(
+                self.all_landmarks_5), f'Mismatched samples: {len(self.pad_input_imgs)} and {len(self.all_landmarks_5)}'
+        for idx, landmark in enumerate(self.all_landmarks_5):
+            # use 5 landmarks to get affine matrix
+            # use cv2.LMEDS method for the equivalence to skimage transform
+            # ref: https://blog.csdn.net/yichxi/article/details/115827338
+            affine_matrix = cv2.estimateAffinePartial2D(landmark, self.face_template, method=cv2.LMEDS)[0]
+            self.affine_matrices.append(affine_matrix)
+            # warp and crop faces
+            if border_mode == 'constant':
+                border_mode = cv2.BORDER_CONSTANT
+            elif border_mode == 'reflect101':
+                border_mode = cv2.BORDER_REFLECT101
+            elif border_mode == 'reflect':
+                border_mode = cv2.BORDER_REFLECT
+            if self.pad_blur:
+                input_img = self.pad_input_imgs[idx]
+            else:
+                input_img = self.input_img
+            cropped_face = cv2.warpAffine(
+                input_img, affine_matrix, self.face_size, borderMode=border_mode, borderValue=(135, 133, 132))  # gray
+            self.cropped_faces.append(cropped_face)
+            # save the cropped face
+            if save_cropped_path is not None:
+                path = os.path.splitext(save_cropped_path)[0]
+                save_path = f'{path}_{idx:02d}.{self.save_ext}'
+                imwrite(cropped_face, save_path)
+
+    def get_inverse_affine(self, save_inverse_affine_path=None):
+        """Get inverse affine matrix."""
+        for idx, affine_matrix in enumerate(self.affine_matrices):
+            inverse_affine = cv2.invertAffineTransform(affine_matrix)
+            inverse_affine *= self.upscale_factor
+            self.inverse_affine_matrices.append(inverse_affine)
+            # save inverse affine matrices
+            if save_inverse_affine_path is not None:
+                path, _ = os.path.splitext(save_inverse_affine_path)
+                save_path = f'{path}_{idx:02d}.pth'
+                torch.save(inverse_affine, save_path)
+
+
+    def add_restored_face(self, face):
+        self.restored_faces.append(face)
+
+
+    def paste_faces_to_input_image(self, save_path=None, upsample_img=None, draw_box=False, face_upsampler=None):
+        h, w, _ = self.input_img.shape
+        h_up, w_up = int(h * self.upscale_factor), int(w * self.upscale_factor)
+
+        if upsample_img is None:
+            # simply resize the background
+            # upsample_img = cv2.resize(self.input_img, (w_up, h_up), interpolation=cv2.INTER_LANCZOS4)
+            upsample_img = cv2.resize(self.input_img, (w_up, h_up), interpolation=cv2.INTER_LINEAR)
+        else:
+            upsample_img = cv2.resize(upsample_img, (w_up, h_up), interpolation=cv2.INTER_LANCZOS4)
+
+        assert len(self.restored_faces) == len(
+            self.inverse_affine_matrices), ('length of restored_faces and affine_matrices are different.')
+        
+        inv_mask_borders = []
+        for restored_face, inverse_affine in zip(self.restored_faces, self.inverse_affine_matrices):
+            if face_upsampler is not None:
+                restored_face = face_upsampler.enhance(restored_face, outscale=self.upscale_factor)[0]
+                inverse_affine /= self.upscale_factor
+                inverse_affine[:, 2] *= self.upscale_factor
+                face_size = (self.face_size[0]*self.upscale_factor, self.face_size[1]*self.upscale_factor)
+            else:
+                # Add an offset to inverse affine matrix, for more precise back alignment
+                if self.upscale_factor > 1:
+                    extra_offset = 0.5 * self.upscale_factor
+                else:
+                    extra_offset = 0
+                inverse_affine[:, 2] += extra_offset
+                face_size = self.face_size
+            inv_restored = cv2.warpAffine(restored_face, inverse_affine, (w_up, h_up))
+
+            # if draw_box or not self.use_parse:  # use square parse maps
+            #     mask = np.ones(face_size, dtype=np.float32)
+            #     inv_mask = cv2.warpAffine(mask, inverse_affine, (w_up, h_up))
+            #     # remove the black borders
+            #     inv_mask_erosion = cv2.erode(
+            #         inv_mask, np.ones((int(2 * self.upscale_factor), int(2 * self.upscale_factor)), np.uint8))
+            #     pasted_face = inv_mask_erosion[:, :, None] * inv_restored
+            #     total_face_area = np.sum(inv_mask_erosion)  # // 3
+            #     # add border
+            #     if draw_box:
+            #         h, w = face_size
+            #         mask_border = np.ones((h, w, 3), dtype=np.float32)
+            #         border = int(1400/np.sqrt(total_face_area))
+            #         mask_border[border:h-border, border:w-border,:] = 0
+            #         inv_mask_border = cv2.warpAffine(mask_border, inverse_affine, (w_up, h_up))
+            #         inv_mask_borders.append(inv_mask_border)
+            #     if not self.use_parse:
+            #         # compute the fusion edge based on the area of face
+            #         w_edge = int(total_face_area**0.5) // 20
+            #         erosion_radius = w_edge * 2
+            #         inv_mask_center = cv2.erode(inv_mask_erosion, np.ones((erosion_radius, erosion_radius), np.uint8))
+            #         blur_size = w_edge * 2
+            #         inv_soft_mask = cv2.GaussianBlur(inv_mask_center, (blur_size + 1, blur_size + 1), 0)
+            #         if len(upsample_img.shape) == 2:  # upsample_img is gray image
+            #             upsample_img = upsample_img[:, :, None]
+            #         inv_soft_mask = inv_soft_mask[:, :, None]
+
+            # always use square mask
+            mask = np.ones(face_size, dtype=np.float32)
+            inv_mask = cv2.warpAffine(mask, inverse_affine, (w_up, h_up))
+            # remove the black borders
+            inv_mask_erosion = cv2.erode(
+                inv_mask, np.ones((int(2 * self.upscale_factor), int(2 * self.upscale_factor)), np.uint8))
+            pasted_face = inv_mask_erosion[:, :, None] * inv_restored
+            total_face_area = np.sum(inv_mask_erosion)  # // 3
+            # add border
+            if draw_box:
+                h, w = face_size
+                mask_border = np.ones((h, w, 3), dtype=np.float32)
+                border = int(1400/np.sqrt(total_face_area))
+                mask_border[border:h-border, border:w-border,:] = 0
+                inv_mask_border = cv2.warpAffine(mask_border, inverse_affine, (w_up, h_up))
+                inv_mask_borders.append(inv_mask_border)
+            # compute the fusion edge based on the area of face
+            w_edge = int(total_face_area**0.5) // 20
+            erosion_radius = w_edge * 2
+            inv_mask_center = cv2.erode(inv_mask_erosion, np.ones((erosion_radius, erosion_radius), np.uint8))
+            blur_size = w_edge * 2
+            inv_soft_mask = cv2.GaussianBlur(inv_mask_center, (blur_size + 1, blur_size + 1), 0)
+            if len(upsample_img.shape) == 2:  # upsample_img is gray image
+                upsample_img = upsample_img[:, :, None]
+            inv_soft_mask = inv_soft_mask[:, :, None]
+
+            # parse mask
+            if self.use_parse:
+                # inference
+                face_input = cv2.resize(restored_face, (512, 512), interpolation=cv2.INTER_LINEAR)
+                face_input = img2tensor(face_input.astype('float32') / 255., bgr2rgb=True, float32=True)
+                normalize(face_input, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
+                face_input = torch.unsqueeze(face_input, 0).to(self.device)
+                with torch.no_grad():
+                    out = self.face_parse(face_input)[0]
+                out = out.argmax(dim=1).squeeze().cpu().numpy()
+
+                parse_mask = np.zeros(out.shape)
+                MASK_COLORMAP = [0, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 0, 255, 0, 0, 0]
+                for idx, color in enumerate(MASK_COLORMAP):
+                    parse_mask[out == idx] = color
+                #  blur the mask
+                parse_mask = cv2.GaussianBlur(parse_mask, (101, 101), 11)
+                parse_mask = cv2.GaussianBlur(parse_mask, (101, 101), 11)
+                # remove the black borders
+                thres = 10
+                parse_mask[:thres, :] = 0
+                parse_mask[-thres:, :] = 0
+                parse_mask[:, :thres] = 0
+                parse_mask[:, -thres:] = 0
+                parse_mask = parse_mask / 255.
+
+                parse_mask = cv2.resize(parse_mask, face_size)
+                parse_mask = cv2.warpAffine(parse_mask, inverse_affine, (w_up, h_up), flags=3)
+                inv_soft_parse_mask = parse_mask[:, :, None]
+                # pasted_face = inv_restored
+                fuse_mask = (inv_soft_parse_mask<inv_soft_mask).astype('int')
+                inv_soft_mask = inv_soft_parse_mask*fuse_mask + inv_soft_mask*(1-fuse_mask)
+
+            if len(upsample_img.shape) == 3 and upsample_img.shape[2] == 4:  # alpha channel
+                alpha = upsample_img[:, :, 3:]
+                upsample_img = inv_soft_mask * pasted_face + (1 - inv_soft_mask) * upsample_img[:, :, 0:3]
+                upsample_img = np.concatenate((upsample_img, alpha), axis=2)
+            else:
+                upsample_img = inv_soft_mask * pasted_face + (1 - inv_soft_mask) * upsample_img
+
+        if np.max(upsample_img) > 256:  # 16-bit image
+            upsample_img = upsample_img.astype(np.uint16)
+        else:
+            upsample_img = upsample_img.astype(np.uint8)
+
+        # draw bounding box
+        if draw_box:
+            # upsample_input_img = cv2.resize(input_img, (w_up, h_up))
+            img_color = np.ones([*upsample_img.shape], dtype=np.float32)
+            img_color[:,:,0] = 0
+            img_color[:,:,1] = 255
+            img_color[:,:,2] = 0
+            for inv_mask_border in inv_mask_borders:
+                upsample_img = inv_mask_border * img_color + (1 - inv_mask_border) * upsample_img
+                # upsample_input_img = inv_mask_border * img_color + (1 - inv_mask_border) * upsample_input_img
+
+        if save_path is not None:
+            path = os.path.splitext(save_path)[0]
+            save_path = f'{path}.{self.save_ext}'
+            imwrite(upsample_img, save_path)
+        return upsample_img
+
+    def clean_all(self):
+        self.all_landmarks_5 = []
+        self.restored_faces = []
+        self.affine_matrices = []
+        self.cropped_faces = []
+        self.inverse_affine_matrices = []
+        self.det_faces = []
+        self.pad_input_imgs = []

py/facelib/utils/face_utils.py

@@ -0,0 +1,248 @@
+import cv2
+import numpy as np
+import torch
+
+
+def compute_increased_bbox(bbox, increase_area, preserve_aspect=True):
+    left, top, right, bot = bbox
+    width = right - left
+    height = bot - top
+
+    if preserve_aspect:
+        width_increase = max(increase_area, ((1 + 2 * increase_area) * height - width) / (2 * width))
+        height_increase = max(increase_area, ((1 + 2 * increase_area) * width - height) / (2 * height))
+    else:
+        width_increase = height_increase = increase_area
+    left = int(left - width_increase * width)
+    top = int(top - height_increase * height)
+    right = int(right + width_increase * width)
+    bot = int(bot + height_increase * height)
+    return (left, top, right, bot)
+
+
+def get_valid_bboxes(bboxes, h, w):
+    left = max(bboxes[0], 0)
+    top = max(bboxes[1], 0)
+    right = min(bboxes[2], w)
+    bottom = min(bboxes[3], h)
+    return (left, top, right, bottom)
+
+
+def align_crop_face_landmarks(img,
+                              landmarks,
+                              output_size,
+                              transform_size=None,
+                              enable_padding=True,
+                              return_inverse_affine=False,
+                              shrink_ratio=(1, 1)):
+    """Align and crop face with landmarks.
+
+    The output_size and transform_size are based on width. The height is
+    adjusted based on shrink_ratio_h/shring_ration_w.
+
+    Modified from:
+    https://github.com/NVlabs/ffhq-dataset/blob/master/download_ffhq.py
+
+    Args:
+        img (Numpy array): Input image.
+        landmarks (Numpy array): 5 or 68 or 98 landmarks.
+        output_size (int): Output face size.
+        transform_size (ing): Transform size. Usually the four time of
+            output_size.
+        enable_padding (float): Default: True.
+        shrink_ratio (float | tuple[float] | list[float]): Shring the whole
+            face for height and width (crop larger area). Default: (1, 1).
+
+    Returns:
+        (Numpy array): Cropped face.
+    """
+    lm_type = 'retinaface_5'  # Options: dlib_5, retinaface_5
+
+    if isinstance(shrink_ratio, (float, int)):
+        shrink_ratio = (shrink_ratio, shrink_ratio)
+    if transform_size is None:
+        transform_size = output_size * 4
+
+    # Parse landmarks
+    lm = np.array(landmarks)
+    if lm.shape[0] == 5 and lm_type == 'retinaface_5':
+        eye_left = lm[0]
+        eye_right = lm[1]
+        mouth_avg = (lm[3] + lm[4]) * 0.5
+    elif lm.shape[0] == 5 and lm_type == 'dlib_5':
+        lm_eye_left = lm[2:4]
+        lm_eye_right = lm[0:2]
+        eye_left = np.mean(lm_eye_left, axis=0)
+        eye_right = np.mean(lm_eye_right, axis=0)
+        mouth_avg = lm[4]
+    elif lm.shape[0] == 68:
+        lm_eye_left = lm[36:42]
+        lm_eye_right = lm[42:48]
+        eye_left = np.mean(lm_eye_left, axis=0)
+        eye_right = np.mean(lm_eye_right, axis=0)
+        mouth_avg = (lm[48] + lm[54]) * 0.5
+    elif lm.shape[0] == 98:
+        lm_eye_left = lm[60:68]
+        lm_eye_right = lm[68:76]
+        eye_left = np.mean(lm_eye_left, axis=0)
+        eye_right = np.mean(lm_eye_right, axis=0)
+        mouth_avg = (lm[76] + lm[82]) * 0.5
+
+    eye_avg = (eye_left + eye_right) * 0.5
+    eye_to_eye = eye_right - eye_left
+    eye_to_mouth = mouth_avg - eye_avg
+
+    # Get the oriented crop rectangle
+    # x: half width of the oriented crop rectangle
+    x = eye_to_eye - np.flipud(eye_to_mouth) * [-1, 1]
+    #  - np.flipud(eye_to_mouth) * [-1, 1]: rotate 90 clockwise
+    # norm with the hypotenuse: get the direction
+    x /= np.hypot(*x)  # get the hypotenuse of a right triangle
+    rect_scale = 1  # TODO: you can edit it to get larger rect
+    x *= max(np.hypot(*eye_to_eye) * 2.0 * rect_scale, np.hypot(*eye_to_mouth) * 1.8 * rect_scale)
+    # y: half height of the oriented crop rectangle
+    y = np.flipud(x) * [-1, 1]
+
+    x *= shrink_ratio[1]  # width
+    y *= shrink_ratio[0]  # height
+
+    # c: center
+    c = eye_avg + eye_to_mouth * 0.1
+    # quad: (left_top, left_bottom, right_bottom, right_top)
+    quad = np.stack([c - x - y, c - x + y, c + x + y, c + x - y])
+    # qsize: side length of the square
+    qsize = np.hypot(*x) * 2
+
+    quad_ori = np.copy(quad)
+    # Shrink, for large face
+    # TODO: do we really need shrink
+    shrink = int(np.floor(qsize / output_size * 0.5))
+    if shrink > 1:
+        h, w = img.shape[0:2]
+        rsize = (int(np.rint(float(w) / shrink)), int(np.rint(float(h) / shrink)))
+        img = cv2.resize(img, rsize, interpolation=cv2.INTER_AREA)
+        quad /= shrink
+        qsize /= shrink
+
+    # Crop
+    h, w = img.shape[0:2]
+    border = max(int(np.rint(qsize * 0.1)), 3)
+    crop = (int(np.floor(min(quad[:, 0]))), int(np.floor(min(quad[:, 1]))), int(np.ceil(max(quad[:, 0]))),
+            int(np.ceil(max(quad[:, 1]))))
+    crop = (max(crop[0] - border, 0), max(crop[1] - border, 0), min(crop[2] + border, w), min(crop[3] + border, h))
+    if crop[2] - crop[0] < w or crop[3] - crop[1] < h:
+        img = img[crop[1]:crop[3], crop[0]:crop[2], :]
+        quad -= crop[0:2]
+
+    # Pad
+    # pad: (width_left, height_top, width_right, height_bottom)
+    h, w = img.shape[0:2]
+    pad = (int(np.floor(min(quad[:, 0]))), int(np.floor(min(quad[:, 1]))), int(np.ceil(max(quad[:, 0]))),
+           int(np.ceil(max(quad[:, 1]))))
+    pad = (max(-pad[0] + border, 0), max(-pad[1] + border, 0), max(pad[2] - w + border, 0), max(pad[3] - h + border, 0))
+    if enable_padding and max(pad) > border - 4:
+        pad = np.maximum(pad, int(np.rint(qsize * 0.3)))
+        img = np.pad(img, ((pad[1], pad[3]), (pad[0], pad[2]), (0, 0)), 'reflect')
+        h, w = img.shape[0:2]
+        y, x, _ = np.ogrid[:h, :w, :1]
+        mask = np.maximum(1.0 - np.minimum(np.float32(x) / pad[0],
+                                           np.float32(w - 1 - x) / pad[2]),
+                          1.0 - np.minimum(np.float32(y) / pad[1],
+                                           np.float32(h - 1 - y) / pad[3]))
+        blur = int(qsize * 0.02)
+        if blur % 2 == 0:
+            blur += 1
+        blur_img = cv2.boxFilter(img, 0, ksize=(blur, blur))
+
+        img = img.astype('float32')
+        img += (blur_img - img) * np.clip(mask * 3.0 + 1.0, 0.0, 1.0)
+        img += (np.median(img, axis=(0, 1)) - img) * np.clip(mask, 0.0, 1.0)
+        img = np.clip(img, 0, 255)  # float32, [0, 255]
+        quad += pad[:2]
+
+    # Transform use cv2
+    h_ratio = shrink_ratio[0] / shrink_ratio[1]
+    dst_h, dst_w = int(transform_size * h_ratio), transform_size
+    template = np.array([[0, 0], [0, dst_h], [dst_w, dst_h], [dst_w, 0]])
+    # use cv2.LMEDS method for the equivalence to skimage transform
+    # ref: https://blog.csdn.net/yichxi/article/details/115827338
+    affine_matrix = cv2.estimateAffinePartial2D(quad, template, method=cv2.LMEDS)[0]
+    cropped_face = cv2.warpAffine(
+        img, affine_matrix, (dst_w, dst_h), borderMode=cv2.BORDER_CONSTANT, borderValue=(135, 133, 132))  # gray
+
+    if output_size < transform_size:
+        cropped_face = cv2.resize(
+            cropped_face, (output_size, int(output_size * h_ratio)), interpolation=cv2.INTER_LINEAR)
+
+    if return_inverse_affine:
+        dst_h, dst_w = int(output_size * h_ratio), output_size
+        template = np.array([[0, 0], [0, dst_h], [dst_w, dst_h], [dst_w, 0]])
+        # use cv2.LMEDS method for the equivalence to skimage transform
+        # ref: https://blog.csdn.net/yichxi/article/details/115827338
+        affine_matrix = cv2.estimateAffinePartial2D(
+            quad_ori, np.array([[0, 0], [0, output_size], [dst_w, dst_h], [dst_w, 0]]), method=cv2.LMEDS)[0]
+        inverse_affine = cv2.invertAffineTransform(affine_matrix)
+    else:
+        inverse_affine = None
+    return cropped_face, inverse_affine
+
+
+def paste_face_back(img, face, inverse_affine):
+    h, w = img.shape[0:2]
+    face_h, face_w = face.shape[0:2]
+    inv_restored = cv2.warpAffine(face, inverse_affine, (w, h))
+    mask = np.ones((face_h, face_w, 3), dtype=np.float32)
+    inv_mask = cv2.warpAffine(mask, inverse_affine, (w, h))
+    # remove the black borders
+    inv_mask_erosion = cv2.erode(inv_mask, np.ones((2, 2), np.uint8))
+    inv_restored_remove_border = inv_mask_erosion * inv_restored
+    total_face_area = np.sum(inv_mask_erosion) // 3
+    # compute the fusion edge based on the area of face
+    w_edge = int(total_face_area**0.5) // 20
+    erosion_radius = w_edge * 2
+    inv_mask_center = cv2.erode(inv_mask_erosion, np.ones((erosion_radius, erosion_radius), np.uint8))
+    blur_size = w_edge * 2
+    inv_soft_mask = cv2.GaussianBlur(inv_mask_center, (blur_size + 1, blur_size + 1), 0)
+    img = inv_soft_mask * inv_restored_remove_border + (1 - inv_soft_mask) * img
+    # float32, [0, 255]
+    return img
+
+
+if __name__ == '__main__':
+    import os
+
+    from custom_nodes.facerestore.facelib.detection import init_detection_model
+    from custom_nodes.facerestore.facelib.utils.face_restoration_helper import get_largest_face
+
+    img_path = '/home/wxt/datasets/ffhq/ffhq_wild/00009.png'
+    img_name = os.splitext(os.path.basename(img_path))[0]
+
+    # initialize model
+    det_net = init_detection_model('retinaface_resnet50', half=False)
+    img_ori = cv2.imread(img_path)
+    h, w = img_ori.shape[0:2]
+    # if larger than 800, scale it
+    scale = max(h / 800, w / 800)
+    if scale > 1:
+        img = cv2.resize(img_ori, (int(w / scale), int(h / scale)), interpolation=cv2.INTER_LINEAR)
+
+    with torch.no_grad():
+        bboxes = det_net.detect_faces(img, 0.97)
+    if scale > 1:
+        bboxes *= scale  # the score is incorrect
+    bboxes = get_largest_face(bboxes, h, w)[0]
+
+    landmarks = np.array([[bboxes[i], bboxes[i + 1]] for i in range(5, 15, 2)])
+
+    cropped_face, inverse_affine = align_crop_face_landmarks(
+        img_ori,
+        landmarks,
+        output_size=512,
+        transform_size=None,
+        enable_padding=True,
+        return_inverse_affine=True,
+        shrink_ratio=(1, 1))
+
+    cv2.imwrite(f'tmp/{img_name}_cropeed_face.png', cropped_face)
+    img = paste_face_back(img_ori, cropped_face, inverse_affine)
+    cv2.imwrite(f'tmp/{img_name}_back.png', img)

py/facelib/utils/misc.py

@@ -0,0 +1,143 @@
+import cv2
+import os
+import os.path as osp
+import torch
+from torch.hub import download_url_to_file, get_dir
+from urllib.parse import urlparse
+# from basicsr.utils.download_util import download_file_from_google_drive
+#import gdown
+
+
+ROOT_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+
+def download_pretrained_models(file_ids, save_path_root):
+    os.makedirs(save_path_root, exist_ok=True)
+
+    for file_name, file_id in file_ids.items():
+        file_url = 'https://drive.google.com/uc?id='+file_id
+        save_path = osp.abspath(osp.join(save_path_root, file_name))
+        if osp.exists(save_path):
+            user_response = input(f'{file_name} already exist. Do you want to cover it? Y/N\n')
+            if user_response.lower() == 'y':
+                print(f'Covering {file_name} to {save_path}')
+                print("skipping gdown in facelib/utils/misc.py "+file_url)
+                #gdown.download(file_url, save_path, quiet=False)
+                # download_file_from_google_drive(file_id, save_path)
+            elif user_response.lower() == 'n':
+                print(f'Skipping {file_name}')
+            else:
+                raise ValueError('Wrong input. Only accepts Y/N.')
+        else:
+            print(f'Downloading {file_name} to {save_path}')
+            print("skipping gdown in facelib/utils/misc.py "+file_url)
+            #gdown.download(file_url, save_path, quiet=False)
+            # download_file_from_google_drive(file_id, save_path)
+
+
+def imwrite(img, file_path, params=None, auto_mkdir=True):
+    """Write image to file.
+
+    Args:
+        img (ndarray): Image array to be written.
+        file_path (str): Image file path.
+        params (None or list): Same as opencv's :func:`imwrite` interface.
+        auto_mkdir (bool): If the parent folder of `file_path` does not exist,
+            whether to create it automatically.
+
+    Returns:
+        bool: Successful or not.
+    """
+    if auto_mkdir:
+        dir_name = os.path.abspath(os.path.dirname(file_path))
+        os.makedirs(dir_name, exist_ok=True)
+    return cv2.imwrite(file_path, img, params)
+
+
+def img2tensor(imgs, bgr2rgb=True, float32=True):
+    """Numpy array to tensor.
+
+    Args:
+        imgs (list[ndarray] | ndarray): Input images.
+        bgr2rgb (bool): Whether to change bgr to rgb.
+        float32 (bool): Whether to change to float32.
+
+    Returns:
+        list[tensor] | tensor: Tensor images. If returned results only have
+            one element, just return tensor.
+    """
+
+    def _totensor(img, bgr2rgb, float32):
+        if img.shape[2] == 3 and bgr2rgb:
+            if img.dtype == 'float64':
+                img = img.astype('float32')
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        img = torch.from_numpy(img.transpose(2, 0, 1))
+        if float32:
+            img = img.float()
+        return img
+
+    if isinstance(imgs, list):
+        return [_totensor(img, bgr2rgb, float32) for img in imgs]
+    else:
+        return _totensor(imgs, bgr2rgb, float32)
+
+
+def load_file_from_url(url, model_dir=None, progress=True, file_name=None):
+    """Ref:https://github.com/1adrianb/face-alignment/blob/master/face_alignment/utils.py
+    """
+    if model_dir is None:
+        hub_dir = get_dir()
+        model_dir = os.path.join(hub_dir, 'checkpoints')
+
+    os.makedirs(os.path.join(ROOT_DIR, model_dir), exist_ok=True)
+
+    parts = urlparse(url)
+    filename = os.path.basename(parts.path)
+    if file_name is not None:
+        filename = file_name
+    cached_file = os.path.abspath(os.path.join(ROOT_DIR, model_dir, filename))
+    if not os.path.exists(cached_file):
+        print(f'Downloading: "{url}" to {cached_file}\n')
+        download_url_to_file(url, cached_file, hash_prefix=None, progress=progress)
+    return cached_file
+
+
+def scandir(dir_path, suffix=None, recursive=False, full_path=False):
+    """Scan a directory to find the interested files.
+    Args:
+        dir_path (str): Path of the directory.
+        suffix (str | tuple(str), optional): File suffix that we are
+            interested in. Default: None.
+        recursive (bool, optional): If set to True, recursively scan the
+            directory. Default: False.
+        full_path (bool, optional): If set to True, include the dir_path.
+            Default: False.
+    Returns:
+        A generator for all the interested files with relative paths.
+    """
+
+    if (suffix is not None) and not isinstance(suffix, (str, tuple)):
+        raise TypeError('"suffix" must be a string or tuple of strings')
+
+    root = dir_path
+
+    def _scandir(dir_path, suffix, recursive):
+        for entry in os.scandir(dir_path):
+            if not entry.name.startswith('.') and entry.is_file():
+                if full_path:
+                    return_path = entry.path
+                else:
+                    return_path = osp.relpath(entry.path, root)
+
+                if suffix is None:
+                    yield return_path
+                elif return_path.endswith(suffix):
+                    yield return_path
+            else:
+                if recursive:
+                    yield from _scandir(entry.path, suffix=suffix, recursive=recursive)
+                else:
+                    continue
+
+    return _scandir(dir_path, suffix=suffix, recursive=recursive)

py/lama/saicinpainting/training/data/masks.py

@@ -7,6 +7,7 @@ from enum import Enum
 import cv2
 import numpy as np
 
+# from lama.saicinpainting.evaluation.masks.mask import SegmentationMask
 from ...utils import LinearRamp
 
 LOGGER = logging.getLogger(__name__)

py/lama/saicinpainting/training/trainers/base.py

@@ -7,8 +7,14 @@ import pytorch_lightning as ptl
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+# from torch.utils.data import DistributedSampler
 
+# from lama.saicinpainting.evaluation import make_evaluator
+# from lama.saicinpainting.training.data.datasets import make_default_train_dataloader, make_default_val_dataloader
+# from lama.saicinpainting.training.losses.adversarial import make_discrim_loss
+# from lama.saicinpainting.training.losses.perceptual import PerceptualLoss, ResNetPL
 from ..modules import make_generator  #, make_discriminator
+# from lama.saicinpainting.training.visualizers import make_visualizer
 from ...utils import add_prefix_to_keys, average_dicts, set_requires_grad, flatten_dict, \
     get_has_ddp_rank
 

py/lama/saicinpainting/training/trainers/default.py

@@ -4,8 +4,10 @@ import torch
 import torch.nn.functional as F
 from omegaconf import OmegaConf
 
+# from lama.saicinpainting.training.data.datasets import make_constant_area_crop_params
 from ..losses.distance_weighting import make_mask_distance_weighter
 from ..losses.feature_matching import feature_matching_loss, masked_l1_loss
+# from lama.saicinpainting.training.modules.fake_fakes import FakeFakesGenerator
 from .base import BaseInpaintingTrainingModule, make_multiscale_noise
 from ...utils import add_prefix_to_keys, get_ramp
 

py/lama/saicinpainting/training/visualizers/__init__.py

@@ -1,7 +1,7 @@
 import logging
 
-from ..visualizers.directory import DirectoryVisualizer
-from ..visualizers.noop import NoopVisualizer
+from .directory import DirectoryVisualizer
+from .noop import NoopVisualizer
 
 
 def make_visualizer(kind, **kwargs):

py/scripts/swapper.py

@@ -1,60 +0,0 @@
-import numpy as np
-import cv2
-import threading
-from typing import Any, Optional, List
-import insightface
-from insightface.app.common import Face
-try:
-    import torch.cuda as cuda
-except:
-    cuda = None
-if cuda is not None:
-    if cuda.is_available():
-        providers = ["CUDAExecutionProvider"]
-    else:
-        providers = ["CPUExecutionProvider"]
-else:
-    providers = ["CPUExecutionProvider"]
-
-THREAD_LOCK = threading.Lock()
-FACE_ANALYSER = None
-FS_MODEL = None
-CURRENT_FS_MODEL_PATH = None
-
-def getFaceSwapModel(model_path: str):
-    global FS_MODEL
-    global CURRENT_FS_MODEL_PATH
-    if CURRENT_FS_MODEL_PATH is None or CURRENT_FS_MODEL_PATH != model_path:
-        CURRENT_FS_MODEL_PATH = model_path
-        FS_MODEL = insightface.model_zoo.get_model(model_path, providers=providers)
-
-    return FS_MODEL
-
-def get_face_analyser() -> Any:
-    global FACE_ANALYSER
-
-    with THREAD_LOCK:
-        if FACE_ANALYSER is None:
-            FACE_ANALYSER = insightface.app.FaceAnalysis(name='buffalo_l', providers=providers)
-    return FACE_ANALYSER
-
-def getFaceAnalyser():
-    face_analyser=insightface.app.FaceAnalysis(name="buffalo_l", providers=providers)
-    return face_analyser
-
-def get_face_single(img_data: np.ndarray,face_index=0, det_size=(640, 640)):
-    face_analyser = get_face_analyser()
-    face_analyser.prepare(ctx_id=0, det_size=det_size)
-    face = face_analyser.get(img_data)
-
-    if len(face) == 0 and det_size[0] > 320 and det_size[1] > 320:
-        det_size_half = (det_size[0] // 2, det_size[1] // 2)
-        return get_face_single(img_data, face_index=face_index, det_size=det_size_half)
-
-    try:
-        return sorted(face, key=lambda x: x.bbox[0])[face_index]
-    except IndexError:
-        return None
-
-
-

py/settings_nodes.json

@@ -1 +0,0 @@
-{}

py/src/audio2exp_models/audio2exp.py

@@ -1,41 +0,0 @@
-from tqdm import tqdm
-import torch
-from torch import nn
-
-
-class Audio2Exp(nn.Module):
-    def __init__(self, netG, cfg, device, prepare_training_loss=False):
-        super(Audio2Exp, self).__init__()
-        self.cfg = cfg
-        self.device = device
-        self.netG = netG.to(device)
-
-    def test(self, batch):
-
-        mel_input = batch['indiv_mels']                         # bs T 1 80 16
-        bs = mel_input.shape[0]
-        T = mel_input.shape[1]
-
-        exp_coeff_pred = []
-
-        for i in tqdm(range(0, T, 10),'audio2exp:'): # every 10 frames
-            
-            current_mel_input = mel_input[:,i:i+10]
-
-            #ref = batch['ref'][:, :, :64].repeat((1,current_mel_input.shape[1],1))           #bs T 64
-            ref = batch['ref'][:, :, :64][:, i:i+10]
-            ratio = batch['ratio_gt'][:, i:i+10]                               #bs T
-
-            audiox = current_mel_input.view(-1, 1, 80, 16)                  # bs*T 1 80 16
-
-            curr_exp_coeff_pred  = self.netG(audiox, ref, ratio)         # bs T 64 
-
-            exp_coeff_pred += [curr_exp_coeff_pred]
-
-        # BS x T x 64
-        results_dict = {
-            'exp_coeff_pred': torch.cat(exp_coeff_pred, axis=1)
-            }
-        return results_dict
-
-

py/src/audio2exp_models/networks.py

@@ -1,74 +0,0 @@
-import torch
-import torch.nn.functional as F
-from torch import nn
-
-class Conv2d(nn.Module):
-    def __init__(self, cin, cout, kernel_size, stride, padding, residual=False, use_act = True, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.conv_block = nn.Sequential(
-                            nn.Conv2d(cin, cout, kernel_size, stride, padding),
-                            nn.BatchNorm2d(cout)
-                            )
-        self.act = nn.ReLU()
-        self.residual = residual
-        self.use_act = use_act
-
-    def forward(self, x):
-        out = self.conv_block(x)
-        if self.residual:
-            out += x
-        
-        if self.use_act:
-            return self.act(out)
-        else:
-            return out
-
-class SimpleWrapperV2(nn.Module):
-    def __init__(self) -> None:
-        super().__init__()
-        self.audio_encoder = nn.Sequential(
-            Conv2d(1, 32, kernel_size=3, stride=1, padding=1),
-            Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True),
-            Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True),
-
-            Conv2d(32, 64, kernel_size=3, stride=(3, 1), padding=1),
-            Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True),
-            Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True),
-
-            Conv2d(64, 128, kernel_size=3, stride=3, padding=1),
-            Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True),
-            Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True),
-
-            Conv2d(128, 256, kernel_size=3, stride=(3, 2), padding=1),
-            Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True),
-
-            Conv2d(256, 512, kernel_size=3, stride=1, padding=0),
-            Conv2d(512, 512, kernel_size=1, stride=1, padding=0),
-            )
-
-        #### load the pre-trained audio_encoder 
-        #self.audio_encoder = self.audio_encoder.to(device)  
-        '''
-        wav2lip_state_dict = torch.load('/apdcephfs_cq2/share_1290939/wenxuazhang/checkpoints/wav2lip.pth')['state_dict']
-        state_dict = self.audio_encoder.state_dict()
-
-        for k,v in wav2lip_state_dict.items():
-            if 'audio_encoder' in k:
-                print('init:', k)
-                state_dict[k.replace('module.audio_encoder.', '')] = v
-        self.audio_encoder.load_state_dict(state_dict)
-        '''
-
-        self.mapping1 = nn.Linear(512+64+1, 64)
-        #self.mapping2 = nn.Linear(30, 64)
-        #nn.init.constant_(self.mapping1.weight, 0.)
-        nn.init.constant_(self.mapping1.bias, 0.)
-
-    def forward(self, x, ref, ratio):
-        x = self.audio_encoder(x).view(x.size(0), -1)
-        ref_reshape = ref.reshape(x.size(0), -1)
-        ratio = ratio.reshape(x.size(0), -1)
-        
-        y = self.mapping1(torch.cat([x, ref_reshape, ratio], dim=1)) 
-        out = y.reshape(ref.shape[0], ref.shape[1], -1) #+ ref # resudial
-        return out

py/src/audio2pose_models/audio2pose.py

@@ -1,94 +0,0 @@
-import torch
-from torch import nn
-from .cvae import CVAE
-from .discriminator import PoseSequenceDiscriminator
-from .audio_encoder import AudioEncoder
-
-class Audio2Pose(nn.Module):
-    def __init__(self, cfg, wav2lip_checkpoint, device='cuda'):
-        super().__init__()
-        self.cfg = cfg
-        self.seq_len = cfg.MODEL.CVAE.SEQ_LEN
-        self.latent_dim = cfg.MODEL.CVAE.LATENT_SIZE
-        self.device = device
-
-        self.audio_encoder = AudioEncoder(wav2lip_checkpoint, device)
-        self.audio_encoder.eval()
-        for param in self.audio_encoder.parameters():
-            param.requires_grad = False
-
-        self.netG = CVAE(cfg)
-        self.netD_motion = PoseSequenceDiscriminator(cfg)
-        
-        
-    def forward(self, x):
-
-        batch = {}
-        coeff_gt = x['gt'].cuda().squeeze(0)           #bs frame_len+1 73
-        batch['pose_motion_gt'] = coeff_gt[:, 1:, 64:70] - coeff_gt[:, :1, 64:70] #bs frame_len 6
-        batch['ref'] = coeff_gt[:, 0, 64:70]  #bs  6
-        batch['class'] = x['class'].squeeze(0).cuda() # bs
-        indiv_mels= x['indiv_mels'].cuda().squeeze(0) # bs seq_len+1 80 16
-
-        # forward
-        audio_emb_list = []
-        audio_emb = self.audio_encoder(indiv_mels[:, 1:, :, :].unsqueeze(2)) #bs seq_len 512
-        batch['audio_emb'] = audio_emb
-        batch = self.netG(batch)
-
-        pose_motion_pred = batch['pose_motion_pred']           # bs frame_len 6
-        pose_gt = coeff_gt[:, 1:, 64:70].clone()               # bs frame_len 6
-        pose_pred = coeff_gt[:, :1, 64:70] + pose_motion_pred  # bs frame_len 6
-
-        batch['pose_pred'] = pose_pred
-        batch['pose_gt'] = pose_gt
-
-        return batch
-
-    def test(self, x):
-
-        batch = {}
-        ref = x['ref']                            #bs 1 70
-        batch['ref'] = x['ref'][:,0,-6:]  
-        batch['class'] = x['class']  
-        bs = ref.shape[0]
-        
-        indiv_mels= x['indiv_mels']               # bs T 1 80 16
-        indiv_mels_use = indiv_mels[:, 1:]        # we regard the ref as the first frame
-        num_frames = x['num_frames']
-        num_frames = int(num_frames) - 1
-
-        #  
-        div = num_frames//self.seq_len
-        re = num_frames%self.seq_len
-        audio_emb_list = []
-        pose_motion_pred_list = [torch.zeros(batch['ref'].unsqueeze(1).shape, dtype=batch['ref'].dtype, 
-                                                device=batch['ref'].device)]
-
-        for i in range(div):
-            z = torch.randn(bs, self.latent_dim).to(ref.device)
-            batch['z'] = z
-            audio_emb = self.audio_encoder(indiv_mels_use[:, i*self.seq_len:(i+1)*self.seq_len,:,:,:]) #bs seq_len 512
-            batch['audio_emb'] = audio_emb
-            batch = self.netG.test(batch)
-            pose_motion_pred_list.append(batch['pose_motion_pred'])  #list of bs seq_len 6
-        
-        if re != 0:
-            z = torch.randn(bs, self.latent_dim).to(ref.device)
-            batch['z'] = z
-            audio_emb = self.audio_encoder(indiv_mels_use[:, -1*self.seq_len:,:,:,:]) #bs seq_len  512
-            if audio_emb.shape[1] != self.seq_len:
-                pad_dim = self.seq_len-audio_emb.shape[1]
-                pad_audio_emb = audio_emb[:, :1].repeat(1, pad_dim, 1) 
-                audio_emb = torch.cat([pad_audio_emb, audio_emb], 1) 
-            batch['audio_emb'] = audio_emb
-            batch = self.netG.test(batch)
-            pose_motion_pred_list.append(batch['pose_motion_pred'][:,-1*re:,:])   
-        
-        pose_motion_pred = torch.cat(pose_motion_pred_list, dim = 1)
-        batch['pose_motion_pred'] = pose_motion_pred
-
-        pose_pred = ref[:, :1, -6:] + pose_motion_pred  # bs T 6
-
-        batch['pose_pred'] = pose_pred
-        return batch

py/src/audio2pose_models/audio2pose_two.py

@@ -1,94 +0,0 @@
-import torch
-from torch import nn
-from .cvae import CVAE
-from .discriminator import PoseSequenceDiscriminator
-from .audio_encoder import AudioEncoder
-
-class Audio2Pose(nn.Module):
-    def __init__(self, cfg, wav2lip_checkpoint, device='cuda'):
-        super().__init__()
-        self.cfg = cfg
-        self.seq_len = cfg.MODEL.CVAE.SEQ_LEN
-        self.latent_dim = cfg.MODEL.CVAE.LATENT_SIZE
-        self.device = device
-
-        self.audio_encoder = AudioEncoder(wav2lip_checkpoint, device)
-        self.audio_encoder.eval()
-        for param in self.audio_encoder.parameters():
-            param.requires_grad = False
-
-        self.netG = CVAE(cfg)
-        self.netD_motion = PoseSequenceDiscriminator(cfg)
-        
-        
-    def forward(self, x):
-
-        batch = {}
-        coeff_gt = x['gt'].cuda().squeeze(0)           #bs frame_len+1 73
-        batch['pose_motion_gt'] = coeff_gt[:, 1:, -9:-3] - coeff_gt[:, :1, -9:-3] #bs frame_len 6
-        batch['ref'] = coeff_gt[:, 0, -9:-3]  #bs  6
-        batch['class'] = x['class'].squeeze(0).cuda() # bs
-        indiv_mels= x['indiv_mels'].cuda().squeeze(0) # bs seq_len+1 80 16
-
-        # forward
-        audio_emb_list = []
-        audio_emb = self.audio_encoder(indiv_mels[:, 1:, :, :].unsqueeze(2)) #bs seq_len 512
-        batch['audio_emb'] = audio_emb
-        batch = self.netG(batch)
-
-        pose_motion_pred = batch['pose_motion_pred']           # bs frame_len 6
-        pose_gt = coeff_gt[:, 1:, -9:-3].clone()               # bs frame_len 6
-        pose_pred = coeff_gt[:, :1, -9:-3] + pose_motion_pred  # bs frame_len 6
-
-        batch['pose_pred'] = pose_pred
-        batch['pose_gt'] = pose_gt
-
-        return batch
-
-    def test(self, x):
-
-        batch = {}
-        ref = x['ref']                            #bs 1 70
-        batch['ref'] = x['ref'][:,0,-6:]  
-        batch['class'] = x['class']  
-        bs = ref.shape[0]
-        
-        indiv_mels= x['indiv_mels']               # bs T 1 80 16
-        indiv_mels_use = indiv_mels[:, 1:]        # we regard the ref as the first frame
-        num_frames = x['num_frames']
-        num_frames = int(num_frames) - 1
-
-        #  
-        div = num_frames//self.seq_len
-        re = num_frames%self.seq_len
-        audio_emb_list = []
-        pose_motion_pred_list = [torch.zeros(batch['ref'].unsqueeze(1).shape, dtype=batch['ref'].dtype, 
-                                                device=batch['ref'].device)]
-
-        for i in range(div):
-            z = torch.randn(bs, self.latent_dim).to(ref.device)
-            batch['z'] = z
-            audio_emb = self.audio_encoder(indiv_mels_use[:, i*self.seq_len:(i+1)*self.seq_len,:,:,:]) #bs seq_len 512
-            batch['audio_emb'] = audio_emb
-            batch = self.netG.test(batch)
-            pose_motion_pred_list.append(batch['pose_motion_pred'])  #list of bs seq_len 6
-        
-        if re != 0:
-            z = torch.randn(bs, self.latent_dim).to(ref.device)
-            batch['z'] = z
-            audio_emb = self.audio_encoder(indiv_mels_use[:, -1*self.seq_len:,:,:,:]) #bs seq_len  512
-            if audio_emb.shape[1] != self.seq_len:
-                pad_dim = self.seq_len-audio_emb.shape[1]
-                pad_audio_emb = audio_emb[:, :1].repeat(1, pad_dim, 1) 
-                audio_emb = torch.cat([pad_audio_emb, audio_emb], 1) 
-            batch['audio_emb'] = audio_emb
-            batch = self.netG.test(batch)
-            pose_motion_pred_list.append(batch['pose_motion_pred'][:,-1*re:,:])   
-        
-        pose_motion_pred = torch.cat(pose_motion_pred_list, dim = 1)
-        batch['pose_motion_pred'] = pose_motion_pred
-
-        pose_pred = ref[:, :1, -6:] + pose_motion_pred  # bs T 6
-
-        batch['pose_pred'] = pose_pred
-        return batch