功能修改

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

.gitignore

@@ -1,21 +0,0 @@
-# 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,20 +8,166 @@ 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_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
+##### 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
 
-Then run the install.bat file, and if there are no errors, it's ready.
+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.
 
-1. Model addresses and storage paths:
-   - Lama model:
-     https://huggingface.co/lllyasviel/Annotators/resolve/main/ControlNetLama.pth  ..\ComfyUI\models\lama\ControlNetLama.pth
+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
 
-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
+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 
 
-3. 执行修改脚本
-windows  执行“修改文件.bat”文件，未报错后就可以了
-linux  执行“修改文件.sh”文件，未报错后就可以了
+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)
+    '''优化输出结束'''
+``` 

README.md

@@ -8,21 +8,170 @@
 
 ![Alt text](解压存放路径及名称.png)
 
-##### 特别说明该版本执行最新版“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
+##### 特别说明该版本执行最新版“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
 
+1. 根据图片解压到指定插件目录及名称，
+将insightface.rar解压到..\ComfyUI_windows_portable\python_embeded\Lib\site-packages目录
+将venv.rar解压到..\ComfyUI_windows_portable\python_embeded\Lib\site-packages目录
 然后运行install.bat文件，未报错后就可以了，
 
-1. 模型地址，及存放路径：
+2. 模型地址，及存放路径：
 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
 
-2. image-face-fusion模型换脸模型
+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模型换脸模型
 链接：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 目录
 
-3. 执行修改脚本
-windows  执行“修改文件.bat”文件，未报错后就可以了
-linux  执行“修改文件.sh”文件，未报错后就可以了
+
+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)
+    '''优化输出结束'''
+``` 

__init__.py

@@ -7,7 +7,6 @@ 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 = {}
 
@@ -22,5 +21,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

@@ -1,854 +0,0 @@
-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

@@ -1,655 +0,0 @@
-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

@@ -1,216 +0,0 @@
-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

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

config/weChat.yaml

@@ -0,0 +1,116 @@
+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,11 +1,3 @@
-#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": 2070780056835,
+      "seed": 611578276912800,
       "steps": 8,
       "cfg": 1,
       "sampler_name": "euler_ancestral",
@@ -86,6 +86,19 @@
       "title": "VAE解码"
     }
   },
+  "9": {
+    "inputs": {
+      "filename_prefix": "ComfyUI",
+      "images": [
+        "8",
+        0
+      ]
+    },
+    "class_type": "SaveImage",
+    "_meta": {
+      "title": "保存图像"
+    }
+  },
   "10": {
     "inputs": {
       "lora_name": "lcm-lora-sdv1-5.safetensors",
@@ -175,364 +188,5 @@
     "_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": "GetImageSize",
+    "class_type": "GetImage_(Width&Height) _O",
     "_meta": {
-      "title": "GetImageSize"
+      "title": "GetImage_(Width&Height) _O"
     }
   },
   "17": {

js/LamAllInOne.js

@@ -0,0 +1,158 @@
+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,24 +10,18 @@ app.registerExtension({
         if (showNames.indexOf(nodeData.name)>=0) {
             const onDrawForeground = nodeType.prototype.onDrawForeground;
             nodeType.prototype.onDrawForeground = function (ctx) {
-                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)
+                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 
-
+                return r;
             };
         }
         var names=["MultiTextConcatenate",'MultiTextSelelct',"MultiIntFormula","MultiParamFormula"]
@@ -52,7 +46,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"){
@@ -62,21 +56,11 @@ app.registerExtension({
                     widget.draw = function (ctx,parentNode, widgetWidth, y, widgetHeight) {
                         draw?.apply(this, arguments);
                         if (this.inputEl.hidden) return;
-                        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) ;
+                        const transform = ctx.getTransform();
                         let maxIoputSize=Math.max(...[parentNode.inputs.length,parentNode.outputs.length])
                         Object.assign(this.inputEl.style, {
-                            left: `${transform.a * margin + transform.e}px`,
-                            top: `${transform.d + transform.f}px`,
-                            width: `${w}px`,
+                            top: `${ maxIoputSize*transform.d*20+transform.d*5+transform.f}px`,
+                            height: `${32.0 * transform.d}px`,
                         });
                     }
                 });

js/litegraphCore.js

@@ -1,27 +1,57 @@
 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) {','LiteGraph.pointerListenerAdd(dialog,"leave", function(e) {','timeout_close = setTimeout(function() {']
+    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);']
     for(var i=0;i<repStrs.length;i++){ //如果对应几个字符串存在缺失，就放弃操作
         if(funstr.indexOf(repStrs[i])==-1){
             return;
         }
     }
     funstr=funstr.replace('function(event, options) {','');
-    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.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.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,7 +2,8 @@ import cv2
 import os
 import torch
 import numpy as np
-from .face_fusion.image_face_fusion import ImageFaceFusion
+from .src.face_fusion.image_face_fusion import ImageFaceFusion
+from .third_part.GFPGAN.gfpgan import GFPGANer
 from PIL import Image
 import folder_paths
 
@@ -12,6 +13,7 @@ class FaceFusion:
         return {"required":
                     {"template_image": ("IMAGE", ),
                     "user_image": ("IMAGE", ),
+                    "facerestore": ([ True, False ], ),
                     },
                 }
 
@@ -21,8 +23,12 @@ class FaceFusion:
     RETURN_NAMES = ("图片",)
 
     FUNCTION = "face_fusion"
-    def face_fusion(self,template_image,user_image):
+    def face_fusion(self,template_image,user_image,facerestore):
         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
@@ -38,11 +44,15 @@ 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,9 +3,6 @@ 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):
@@ -18,13 +15,9 @@ 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"
@@ -32,34 +25,31 @@ class ForEnd:
 
     CATEGORY = "lam"
 
-    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)+"论结束........."
-
+    def for_end_fun(self,total,i,port,obj):
         if i>=total:
-            result = result+ '\n整个循环结束'
-        else:
-            r = promptFun(PromptServer.instance,pdata)
+            return { "ui": { "text":'循环结束' } }
 
-        if StatusInfo.endswith('循环结束')==False and len(StatusInfo)>0:
-            result=StatusInfo+'\n'+result
+        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
 
-        pdata['extra_data']['extra_pnginfo']['StatusInfo']=result
-
-        return { "ui": { "text":result} }
+        return { "ui": { "text":"第"+str(index+1)+"次循环结果："+result} }
 
 NODE_CLASS_MAPPINGS = {
     "ForEnd": ForEnd

py/GLM3Prompt.py

@@ -1,47 +0,0 @@
-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

@@ -0,0 +1,56 @@
+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

@@ -0,0 +1,48 @@
+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,6 +3,7 @@ import cv2
 import os
 import torch
 import numpy as np
+from ffmpy import FFmpeg
 
 class LamLoadVideo:
     def __init__(self):
@@ -66,9 +67,8 @@ 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)
-            #提取音频
-            cmd = r"ffmpeg -y -hide_banner -loglevel error -i %s %s"%(videoPath, filePathName)
-            os.system(cmd)  
+            ff = FFmpeg(inputs={videoPath: None},outputs={filePathName: '-f {} -vn'.format('mp3')})
+            ff.run()
         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,10 +65,14 @@ 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":"翻译结果："+result['result']['trans_result'][0]['dst']},"result": (result['result']['trans_result'][0]['dst'],)}
+            return { "ui": { "text":"翻译结果："+texts},"result": (texts,)}
 
     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

@@ -0,0 +1,224 @@
+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

@@ -0,0 +1,45 @@
+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

@@ -0,0 +1,71 @@
+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

@@ -0,0 +1,56 @@
+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,6 +3,7 @@ import cv2
 import os
 import torch
 import numpy as np
+from ffmpy import FFmpeg
 
 class VideoAddAudio:
     def __init__(self):
@@ -39,9 +40,10 @@ 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)
-        #插入音频
-        cmd = r'ffmpeg -y -hide_banner -loglevel error -i "%s" -i "%s" -vcodec copy "%s"' % (videoPath, audioPath, result)
-        os.system(cmd)  
+        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()
         return {"ui": {"text": "视频插入音频成功，保存路径："+result,
         'videos':[{'filename':file,'type':'output','subfolder':'video'}]}}
 

py/VideoFaceFusion.py

@@ -1,108 +1,17 @@
 import cv2
 import os
-import sys
 import torch
-import model_management
 import numpy as np
-from .facelib.utils.face_restoration_helper import FaceRestoreHelper
-from .facelib.detection.retinaface import retinaface
-from .face_fusion.image_face_fusion import ImageFaceFusion
+from .src.face_fusion.image_face_fusion import ImageFaceFusion
+from .third_part.GFPGAN.gfpgan import GFPGANer
 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)
@@ -114,12 +23,7 @@ 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"
 
@@ -128,13 +32,10 @@ class VideoFaceFusion:
     OUTPUT_NODE = True
 
     FUNCTION = "face_fusion"
-    def face_fusion(self,templateVideoPath,user_image,filename_prefix,facedetection,upscale_model=None):
+    def face_fusion(self,templateVideoPath,user_image,filename_prefix):
         image_face_fusion = ImageFaceFusion(model_dir=folder_paths.models_dir+'/image-face-fusion')
-        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)
+        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
 
@@ -168,58 +69,19 @@ class VideoFaceFusion:
         with imageio.get_writer(videoPathTemp, fps=fps) as video:
             for i in tqdm(range(len(image_list)), '写入进度:'):
                 image=image_list[i]
-                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)# 写入视频
-
+                _,_,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)
-        #提取音频
-        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)  
+        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,

py/VideoRoopFaceSwap.py

@@ -0,0 +1,104 @@
+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,147 +14,8 @@ 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':
@@ -231,36 +92,6 @@ def send_sync(self, event, data, sid=None,port=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
@@ -273,18 +104,6 @@ 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]
@@ -303,18 +122,6 @@ 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'])
@@ -339,6 +146,8 @@ 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, '', '')
@@ -349,117 +158,6 @@ 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 '您还没有排队，请先发送您的指令'
@@ -478,12 +176,8 @@ 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,page_number=page_number)
+        datas=db.get_many_data(openId)
         db.close_con()
         rdataList=[]
         if "type" in request.rel_url.query:
@@ -797,8 +491,6 @@ 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/facelib/detection/__init__.py

@@ -1,100 +0,0 @@
-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

@@ -1,219 +0,0 @@
-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

@@ -1,317 +0,0 @@
-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

@@ -1,370 +0,0 @@
-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

@@ -1,196 +0,0 @@
-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/yolov5face/face_detector.py

@@ -1,142 +0,0 @@
-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

@@ -1,299 +0,0 @@
-# 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

@@ -1,45 +0,0 @@
-# # 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

@@ -1,235 +0,0 @@
-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

@@ -1,47 +0,0 @@
-# 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

@@ -1,45 +0,0 @@
-# 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

@@ -1,12 +0,0 @@
-# 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

@@ -1,35 +0,0 @@
-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

@@ -1,5 +0,0 @@
-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

@@ -1,271 +0,0 @@
-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

@@ -1,40 +0,0 @@
-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

@@ -1,23 +0,0 @@
-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

@@ -1,140 +0,0 @@
-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/utils/__init__.py

@@ -1,7 +0,0 @@
-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

@@ -1,455 +0,0 @@
-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

@@ -1,248 +0,0 @@
-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

@@ -1,143 +0,0 @@
-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,7 +7,6 @@ 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,14 +7,8 @@ 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,10 +4,8 @@ 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 .directory import DirectoryVisualizer
-from .noop import NoopVisualizer
+from ..visualizers.directory import DirectoryVisualizer
+from ..visualizers.noop import NoopVisualizer
 
 
 def make_visualizer(kind, **kwargs):

py/scripts/swapper.py

@@ -0,0 +1,60 @@
+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

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

py/src/audio2exp_models/audio2exp.py

@@ -0,0 +1,41 @@
+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

@@ -0,0 +1,74 @@
+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

@@ -0,0 +1,94 @@
+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

@@ -0,0 +1,94 @@
+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

py/src/audio2pose_models/audio_encoder.py

@@ -0,0 +1,64 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+class Conv2d(nn.Module):
+    def __init__(self, cin, cout, kernel_size, stride, padding, residual=False, *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
+
+    def forward(self, x):
+        out = self.conv_block(x)
+        if self.residual:
+            out += x
+        return self.act(out)
+
+class AudioEncoder(nn.Module):
+    def __init__(self, wav2lip_checkpoint, device):
+        super(AudioEncoder, self).__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, we do not need to load wav2lip model here.
+        # wav2lip_state_dict = torch.load(wav2lip_checkpoint, map_location=torch.device(device))['state_dict']
+        # state_dict = self.audio_encoder.state_dict()
+
+        # for k,v in wav2lip_state_dict.items():
+        #     if 'audio_encoder' in k:
+        #         state_dict[k.replace('module.audio_encoder.', '')] = v
+        # self.audio_encoder.load_state_dict(state_dict)
+
+
+    def forward(self, audio_sequences):
+        # audio_sequences = (B, T, 1, 80, 16)
+        B = audio_sequences.size(0)
+
+        audio_sequences = torch.cat([audio_sequences[:, i] for i in range(audio_sequences.size(1))], dim=0)
+
+        audio_embedding = self.audio_encoder(audio_sequences) # B, 512, 1, 1
+        dim = audio_embedding.shape[1]
+        audio_embedding = audio_embedding.reshape((B, -1, dim, 1, 1))
+
+        return audio_embedding.squeeze(-1).squeeze(-1) #B seq_len+1 512 

py/src/audio2pose_models/audio_encoder_two.py

@@ -0,0 +1,64 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+class Conv2d(nn.Module):
+    def __init__(self, cin, cout, kernel_size, stride, padding, residual=False, *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
+
+    def forward(self, x):
+        out = self.conv_block(x)
+        if self.residual:
+            out += x
+        return self.act(out)
+
+class AudioEncoder(nn.Module):
+    def __init__(self, wav2lip_checkpoint, device):
+        super(AudioEncoder, self).__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
+        wav2lip_state_dict = torch.load(wav2lip_checkpoint, map_location=torch.device(device))['state_dict']
+        state_dict = self.audio_encoder.state_dict()
+
+        for k,v in wav2lip_state_dict.items():
+            if 'audio_encoder' in k:
+                state_dict[k.replace('module.audio_encoder.', '')] = v
+        self.audio_encoder.load_state_dict(state_dict)
+
+
+    def forward(self, audio_sequences):
+        # audio_sequences = (B, T, 1, 80, 16)
+        B = audio_sequences.size(0)
+
+        audio_sequences = torch.cat([audio_sequences[:, i] for i in range(audio_sequences.size(1))], dim=0)
+
+        audio_embedding = self.audio_encoder(audio_sequences) # B, 512, 1, 1
+        dim = audio_embedding.shape[1]
+        audio_embedding = audio_embedding.reshape((B, -1, dim, 1, 1))
+
+        return audio_embedding.squeeze(-1).squeeze(-1) #B seq_len+1 512 

py/src/audio2pose_models/cvae.py

@@ -0,0 +1,149 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+from .res_unet import ResUnet
+
+def class2onehot(idx, class_num):
+
+    assert torch.max(idx).item() < class_num
+    onehot = torch.zeros(idx.size(0), class_num).to(idx.device)
+    onehot.scatter_(1, idx, 1)
+    return onehot
+
+class CVAE(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        encoder_layer_sizes = cfg.MODEL.CVAE.ENCODER_LAYER_SIZES
+        decoder_layer_sizes = cfg.MODEL.CVAE.DECODER_LAYER_SIZES
+        latent_size = cfg.MODEL.CVAE.LATENT_SIZE
+        num_classes = cfg.DATASET.NUM_CLASSES
+        audio_emb_in_size = cfg.MODEL.CVAE.AUDIO_EMB_IN_SIZE
+        audio_emb_out_size = cfg.MODEL.CVAE.AUDIO_EMB_OUT_SIZE
+        seq_len = cfg.MODEL.CVAE.SEQ_LEN
+
+        self.latent_size = latent_size
+
+        self.encoder = ENCODER(encoder_layer_sizes, latent_size, num_classes,
+                                audio_emb_in_size, audio_emb_out_size, seq_len)
+        self.decoder = DECODER(decoder_layer_sizes, latent_size, num_classes,
+                                audio_emb_in_size, audio_emb_out_size, seq_len)
+    def reparameterize(self, mu, logvar):
+        std = torch.exp(0.5 * logvar)
+        eps = torch.randn_like(std)
+        return mu + eps * std
+
+    def forward(self, batch):
+        batch = self.encoder(batch)
+        mu = batch['mu']
+        logvar = batch['logvar']
+        z = self.reparameterize(mu, logvar)
+        batch['z'] = z
+        return self.decoder(batch)
+
+    def test(self, batch):
+        '''
+        class_id = batch['class']
+        z = torch.randn([class_id.size(0), self.latent_size]).to(class_id.device)
+        batch['z'] = z
+        '''
+        return self.decoder(batch)
+
+class ENCODER(nn.Module):
+    def __init__(self, layer_sizes, latent_size, num_classes, 
+                audio_emb_in_size, audio_emb_out_size, seq_len):
+        super().__init__()
+
+        self.resunet = ResUnet()
+        self.num_classes = num_classes
+        self.seq_len = seq_len
+
+        self.MLP = nn.Sequential()
+        layer_sizes[0] += latent_size + seq_len*audio_emb_out_size + 6
+        for i, (in_size, out_size) in enumerate(zip(layer_sizes[:-1], layer_sizes[1:])):
+            self.MLP.add_module(
+                name="L{:d}".format(i), module=nn.Linear(in_size, out_size))
+            self.MLP.add_module(name="A{:d}".format(i), module=nn.ReLU())
+
+        self.linear_means = nn.Linear(layer_sizes[-1], latent_size)
+        self.linear_logvar = nn.Linear(layer_sizes[-1], latent_size)
+        self.linear_audio = nn.Linear(audio_emb_in_size, audio_emb_out_size)
+
+        self.classbias = nn.Parameter(torch.randn(self.num_classes, latent_size))
+
+    def forward(self, batch):
+        class_id = batch['class']
+        pose_motion_gt = batch['pose_motion_gt']                             #bs seq_len 6
+        ref = batch['ref']                             #bs 6
+        bs = pose_motion_gt.shape[0]
+        audio_in = batch['audio_emb']                          # bs seq_len audio_emb_in_size
+
+        #pose encode
+        pose_emb = self.resunet(pose_motion_gt.unsqueeze(1))          #bs 1 seq_len 6 
+        pose_emb = pose_emb.reshape(bs, -1)                    #bs seq_len*6
+
+        #audio mapping
+        print(audio_in.shape)
+        audio_out = self.linear_audio(audio_in)                # bs seq_len audio_emb_out_size
+        audio_out = audio_out.reshape(bs, -1)
+
+        class_bias = self.classbias[class_id]                  #bs latent_size
+        x_in = torch.cat([ref, pose_emb, audio_out, class_bias], dim=-1) #bs seq_len*(audio_emb_out_size+6)+latent_size
+        x_out = self.MLP(x_in)
+
+        mu = self.linear_means(x_out)
+        logvar = self.linear_means(x_out)                      #bs latent_size 
+
+        batch.update({'mu':mu, 'logvar':logvar})
+        return batch
+
+class DECODER(nn.Module):
+    def __init__(self, layer_sizes, latent_size, num_classes, 
+                audio_emb_in_size, audio_emb_out_size, seq_len):
+        super().__init__()
+
+        self.resunet = ResUnet()
+        self.num_classes = num_classes
+        self.seq_len = seq_len
+
+        self.MLP = nn.Sequential()
+        input_size = latent_size + seq_len*audio_emb_out_size + 6
+        for i, (in_size, out_size) in enumerate(zip([input_size]+layer_sizes[:-1], layer_sizes)):
+            self.MLP.add_module(
+                name="L{:d}".format(i), module=nn.Linear(in_size, out_size))
+            if i+1 < len(layer_sizes):
+                self.MLP.add_module(name="A{:d}".format(i), module=nn.ReLU())
+            else:
+                self.MLP.add_module(name="sigmoid", module=nn.Sigmoid())
+        
+        self.pose_linear = nn.Linear(6, 6)
+        self.linear_audio = nn.Linear(audio_emb_in_size, audio_emb_out_size)
+
+        self.classbias = nn.Parameter(torch.randn(self.num_classes, latent_size))
+
+    def forward(self, batch):
+
+        z = batch['z']                                          #bs latent_size
+        bs = z.shape[0]
+        class_id = batch['class']
+        ref = batch['ref']                             #bs 6
+        audio_in = batch['audio_emb']                           # bs seq_len audio_emb_in_size
+        #print('audio_in: ', audio_in[:, :, :10])
+
+        audio_out = self.linear_audio(audio_in)                 # bs seq_len audio_emb_out_size
+        #print('audio_out: ', audio_out[:, :, :10])
+        audio_out = audio_out.reshape([bs, -1])                 # bs seq_len*audio_emb_out_size
+        class_bias = self.classbias[class_id]                   #bs latent_size
+
+        z = z + class_bias
+        x_in = torch.cat([ref, z, audio_out], dim=-1)
+        x_out = self.MLP(x_in)                                  # bs layer_sizes[-1]
+        x_out = x_out.reshape((bs, self.seq_len, -1))
+
+        #print('x_out: ', x_out)
+
+        pose_emb = self.resunet(x_out.unsqueeze(1))             #bs 1 seq_len 6
+
+        pose_motion_pred = self.pose_linear(pose_emb.squeeze(1))       #bs seq_len 6
+
+        batch.update({'pose_motion_pred':pose_motion_pred})
+        return batch

py/src/audio2pose_models/discriminator.py

@@ -0,0 +1,76 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+class ConvNormRelu(nn.Module):
+    def __init__(self, conv_type='1d', in_channels=3, out_channels=64, downsample=False,
+                 kernel_size=None, stride=None, padding=None, norm='BN', leaky=False):
+        super().__init__()
+        if kernel_size is None:
+            if downsample:
+                kernel_size, stride, padding = 4, 2, 1
+            else:
+                kernel_size, stride, padding = 3, 1, 1
+
+        if conv_type == '2d':
+            self.conv = nn.Conv2d(
+                in_channels,
+                out_channels,
+                kernel_size,
+                stride,
+                padding,
+                bias=False,
+            )
+            if norm == 'BN':
+                self.norm = nn.BatchNorm2d(out_channels)
+            elif norm == 'IN':
+                self.norm = nn.InstanceNorm2d(out_channels)
+            else:
+                raise NotImplementedError
+        elif conv_type == '1d':
+            self.conv = nn.Conv1d(
+                in_channels,
+                out_channels,
+                kernel_size,
+                stride,
+                padding,
+                bias=False,
+            )
+            if norm == 'BN':
+                self.norm = nn.BatchNorm1d(out_channels)
+            elif norm == 'IN':
+                self.norm = nn.InstanceNorm1d(out_channels)
+            else:
+                raise NotImplementedError
+        nn.init.kaiming_normal_(self.conv.weight)
+
+        self.act = nn.LeakyReLU(negative_slope=0.2, inplace=False) if leaky else nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        x = self.conv(x)
+        if isinstance(self.norm, nn.InstanceNorm1d):
+            x = self.norm(x.permute((0, 2, 1))).permute((0, 2, 1))  # normalize on [C]
+        else:
+            x = self.norm(x)
+        x = self.act(x)
+        return x
+
+
+class PoseSequenceDiscriminator(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        leaky = self.cfg.MODEL.DISCRIMINATOR.LEAKY_RELU
+
+        self.seq = nn.Sequential(
+            ConvNormRelu('1d', cfg.MODEL.DISCRIMINATOR.INPUT_CHANNELS, 256, downsample=True, leaky=leaky),  # B, 256, 64
+            ConvNormRelu('1d', 256, 512, downsample=True, leaky=leaky),  # B, 512, 32
+            ConvNormRelu('1d', 512, 1024, kernel_size=3, stride=1, padding=1, leaky=leaky),  # B, 1024, 16
+            nn.Conv1d(1024, 1, kernel_size=3, stride=1, padding=1, bias=True)  # B, 1, 16
+        )
+
+    def forward(self, x):
+        x = x.reshape(x.size(0), x.size(1), -1).transpose(1, 2)
+        x = self.seq(x)
+        x = x.squeeze(1)
+        return x

py/src/audio2pose_models/networks.py

@@ -0,0 +1,140 @@
+import torch.nn as nn
+import torch
+
+
+class ResidualConv(nn.Module):
+    def __init__(self, input_dim, output_dim, stride, padding):
+        super(ResidualConv, self).__init__()
+
+        self.conv_block = nn.Sequential(
+            nn.BatchNorm2d(input_dim),
+            nn.ReLU(),
+            nn.Conv2d(
+                input_dim, output_dim, kernel_size=3, stride=stride, padding=padding
+            ),
+            nn.BatchNorm2d(output_dim),
+            nn.ReLU(),
+            nn.Conv2d(output_dim, output_dim, kernel_size=3, padding=1),
+        )
+        self.conv_skip = nn.Sequential(
+            nn.Conv2d(input_dim, output_dim, kernel_size=3, stride=stride, padding=1),
+            nn.BatchNorm2d(output_dim),
+        )
+
+    def forward(self, x):
+
+        return self.conv_block(x) + self.conv_skip(x)
+
+
+class Upsample(nn.Module):
+    def __init__(self, input_dim, output_dim, kernel, stride):
+        super(Upsample, self).__init__()
+
+        self.upsample = nn.ConvTranspose2d(
+            input_dim, output_dim, kernel_size=kernel, stride=stride
+        )
+
+    def forward(self, x):
+        return self.upsample(x)
+
+
+class Squeeze_Excite_Block(nn.Module):
+    def __init__(self, channel, reduction=16):
+        super(Squeeze_Excite_Block, self).__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.fc = nn.Sequential(
+            nn.Linear(channel, channel // reduction, bias=False),
+            nn.ReLU(inplace=True),
+            nn.Linear(channel // reduction, channel, bias=False),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, x):
+        b, c, _, _ = x.size()
+        y = self.avg_pool(x).view(b, c)
+        y = self.fc(y).view(b, c, 1, 1)
+        return x * y.expand_as(x)
+
+
+class ASPP(nn.Module):
+    def __init__(self, in_dims, out_dims, rate=[6, 12, 18]):
+        super(ASPP, self).__init__()
+
+        self.aspp_block1 = nn.Sequential(
+            nn.Conv2d(
+                in_dims, out_dims, 3, stride=1, padding=rate[0], dilation=rate[0]
+            ),
+            nn.ReLU(inplace=True),
+            nn.BatchNorm2d(out_dims),
+        )
+        self.aspp_block2 = nn.Sequential(
+            nn.Conv2d(
+                in_dims, out_dims, 3, stride=1, padding=rate[1], dilation=rate[1]
+            ),
+            nn.ReLU(inplace=True),
+            nn.BatchNorm2d(out_dims),
+        )
+        self.aspp_block3 = nn.Sequential(
+            nn.Conv2d(
+                in_dims, out_dims, 3, stride=1, padding=rate[2], dilation=rate[2]
+            ),
+            nn.ReLU(inplace=True),
+            nn.BatchNorm2d(out_dims),
+        )
+
+        self.output = nn.Conv2d(len(rate) * out_dims, out_dims, 1)
+        self._init_weights()
+
+    def forward(self, x):
+        x1 = self.aspp_block1(x)
+        x2 = self.aspp_block2(x)
+        x3 = self.aspp_block3(x)
+        out = torch.cat([x1, x2, x3], dim=1)
+        return self.output(out)
+
+    def _init_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight)
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+
+class Upsample_(nn.Module):
+    def __init__(self, scale=2):
+        super(Upsample_, self).__init__()
+
+        self.upsample = nn.Upsample(mode="bilinear", scale_factor=scale)
+
+    def forward(self, x):
+        return self.upsample(x)
+
+
+class AttentionBlock(nn.Module):
+    def __init__(self, input_encoder, input_decoder, output_dim):
+        super(AttentionBlock, self).__init__()
+
+        self.conv_encoder = nn.Sequential(
+            nn.BatchNorm2d(input_encoder),
+            nn.ReLU(),
+            nn.Conv2d(input_encoder, output_dim, 3, padding=1),
+            nn.MaxPool2d(2, 2),
+        )
+
+        self.conv_decoder = nn.Sequential(
+            nn.BatchNorm2d(input_decoder),
+            nn.ReLU(),
+            nn.Conv2d(input_decoder, output_dim, 3, padding=1),
+        )
+
+        self.conv_attn = nn.Sequential(
+            nn.BatchNorm2d(output_dim),
+            nn.ReLU(),
+            nn.Conv2d(output_dim, 1, 1),
+        )
+
+    def forward(self, x1, x2):
+        out = self.conv_encoder(x1) + self.conv_decoder(x2)
+        out = self.conv_attn(out)
+        return out * x2

py/src/audio2pose_models/res_unet.py

@@ -0,0 +1,65 @@
+import torch
+import torch.nn as nn
+from .networks import ResidualConv, Upsample
+
+
+class ResUnet(nn.Module):
+    def __init__(self, channel=1, filters=[32, 64, 128, 256]):
+        super(ResUnet, self).__init__()
+
+        self.input_layer = nn.Sequential(
+            nn.Conv2d(channel, filters[0], kernel_size=3, padding=1),
+            nn.BatchNorm2d(filters[0]),
+            nn.ReLU(),
+            nn.Conv2d(filters[0], filters[0], kernel_size=3, padding=1),
+        )
+        self.input_skip = nn.Sequential(
+            nn.Conv2d(channel, filters[0], kernel_size=3, padding=1)
+        )
+
+        self.residual_conv_1 = ResidualConv(filters[0], filters[1], stride=(2,1), padding=1)
+        self.residual_conv_2 = ResidualConv(filters[1], filters[2], stride=(2,1), padding=1)
+
+        self.bridge = ResidualConv(filters[2], filters[3], stride=(2,1), padding=1)
+
+        self.upsample_1 = Upsample(filters[3], filters[3], kernel=(2,1), stride=(2,1))
+        self.up_residual_conv1 = ResidualConv(filters[3] + filters[2], filters[2], stride=1, padding=1)
+
+        self.upsample_2 = Upsample(filters[2], filters[2], kernel=(2,1), stride=(2,1))
+        self.up_residual_conv2 = ResidualConv(filters[2] + filters[1], filters[1], stride=1, padding=1)
+
+        self.upsample_3 = Upsample(filters[1], filters[1], kernel=(2,1), stride=(2,1))
+        self.up_residual_conv3 = ResidualConv(filters[1] + filters[0], filters[0], stride=1, padding=1)
+
+        self.output_layer = nn.Sequential(
+            nn.Conv2d(filters[0], 1, 1, 1),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, x):
+        # Encode
+        x1 = self.input_layer(x) + self.input_skip(x)
+        x2 = self.residual_conv_1(x1)
+        x3 = self.residual_conv_2(x2)
+        # Bridge
+        x4 = self.bridge(x3)
+
+        # Decode
+        x4 = self.upsample_1(x4)
+        x5 = torch.cat([x4, x3], dim=1)
+
+        x6 = self.up_residual_conv1(x5)
+
+        x6 = self.upsample_2(x6)
+        x7 = torch.cat([x6, x2], dim=1)
+
+        x8 = self.up_residual_conv2(x7)
+
+        x8 = self.upsample_3(x8)
+        x9 = torch.cat([x8, x1], dim=1)
+
+        x10 = self.up_residual_conv3(x9)
+
+        output = self.output_layer(x10)
+
+        return output

py/src/config/auido2exp.yaml

@@ -0,0 +1,58 @@
+DATASET:
+  TRAIN_FILE_LIST: /apdcephfs_cq2/share_1290939/wenxuazhang/code/file_list/train.txt
+  EVAL_FILE_LIST: /apdcephfs_cq2/share_1290939/wenxuazhang/code/file_list/val.txt
+  TRAIN_BATCH_SIZE: 32
+  EVAL_BATCH_SIZE: 32
+  EXP: True
+  EXP_DIM: 64
+  FRAME_LEN: 32
+  COEFF_LEN: 73
+  NUM_CLASSES: 46
+  AUDIO_ROOT_PATH: /apdcephfs_cq2/share_1290939/wenxuazhang/voxceleb1/wav
+  COEFF_ROOT_PATH: /apdcephfs_cq2/share_1290939/wenxuazhang/voxceleb1/wav2lip_3dmm
+  LMDB_PATH: /apdcephfs_cq2/share_1290939/shadowcun/datasets/VoxCeleb/v1/imdb
+  DEBUG: True
+  NUM_REPEATS: 2
+  T: 40
+  
+
+MODEL:
+  FRAMEWORK: V2
+  AUDIOENCODER:
+    LEAKY_RELU: True
+    NORM: 'IN'
+  DISCRIMINATOR:
+    LEAKY_RELU: False
+    INPUT_CHANNELS: 6
+  CVAE:
+    AUDIO_EMB_IN_SIZE: 512
+    AUDIO_EMB_OUT_SIZE: 128
+    SEQ_LEN: 32
+    LATENT_SIZE: 256
+    ENCODER_LAYER_SIZES: [192, 1024]
+    DECODER_LAYER_SIZES: [1024, 192]
+    
+
+TRAIN:
+  MAX_EPOCH: 300
+  GENERATOR:
+    LR: 2.0e-5
+  DISCRIMINATOR:
+    LR: 1.0e-5
+  LOSS:
+    W_FEAT: 0
+    W_COEFF_EXP: 2
+    W_LM: 1.0e-2
+    W_LM_MOUTH: 0
+    W_REG: 0
+    W_SYNC: 0
+    W_COLOR: 0
+    W_EXPRESSION: 0
+    W_LIPREADING: 0.01
+    W_LIPREADING_VV: 0
+    W_EYE_BLINK: 4
+
+TAG:
+  NAME:  small_dataset
+
+

py/src/config/auido2pose.yaml

@@ -0,0 +1,49 @@
+DATASET:
+  TRAIN_FILE_LIST: /apdcephfs_cq2/share_1290939/wenxuazhang/code/audio2pose_unet_noAudio/dataset/train_33.txt
+  EVAL_FILE_LIST: /apdcephfs_cq2/share_1290939/wenxuazhang/code/audio2pose_unet_noAudio/dataset/val.txt
+  TRAIN_BATCH_SIZE: 64
+  EVAL_BATCH_SIZE: 1
+  EXP: True
+  EXP_DIM: 64
+  FRAME_LEN: 32
+  COEFF_LEN: 73
+  NUM_CLASSES: 46
+  AUDIO_ROOT_PATH: /apdcephfs_cq2/share_1290939/wenxuazhang/voxceleb1/wav
+  COEFF_ROOT_PATH: /apdcephfs_cq2/share_1290939/shadowcun/datasets/VoxCeleb/v1/imdb
+  DEBUG: True
+  
+
+MODEL:
+  AUDIOENCODER:
+    LEAKY_RELU: True
+    NORM: 'IN'
+  DISCRIMINATOR:
+    LEAKY_RELU: False
+    INPUT_CHANNELS: 6
+  CVAE:
+    AUDIO_EMB_IN_SIZE: 512
+    AUDIO_EMB_OUT_SIZE: 6
+    SEQ_LEN: 32
+    LATENT_SIZE: 64
+    ENCODER_LAYER_SIZES: [192, 128]
+    DECODER_LAYER_SIZES: [128, 192]
+    
+
+TRAIN:
+  MAX_EPOCH: 150
+  GENERATOR:
+    LR: 1.0e-4
+  DISCRIMINATOR:
+    LR: 1.0e-4
+  LOSS:
+    LAMBDA_REG: 1
+    LAMBDA_LANDMARKS: 0
+    LAMBDA_VERTICES: 0
+    LAMBDA_GAN_MOTION: 0.7
+    LAMBDA_GAN_COEFF: 0
+    LAMBDA_KL: 1
+
+TAG:
+  NAME: cvae_UNET_useAudio_usewav2lipAudioEncoder
+
+

py/src/config/facerender.yaml

@@ -0,0 +1,45 @@
+model_params:
+  common_params:
+    num_kp: 15 
+    image_channel: 3                    
+    feature_channel: 32
+    estimate_jacobian: False   # True
+  kp_detector_params:
+     temperature: 0.1
+     block_expansion: 32            
+     max_features: 1024
+     scale_factor: 0.25         # 0.25
+     num_blocks: 5
+     reshape_channel: 16384  # 16384 = 1024 * 16
+     reshape_depth: 16
+  he_estimator_params:
+     block_expansion: 64            
+     max_features: 2048
+     num_bins: 66
+  generator_params:
+    block_expansion: 64
+    max_features: 512
+    num_down_blocks: 2
+    reshape_channel: 32
+    reshape_depth: 16         # 512 = 32 * 16
+    num_resblocks: 6
+    estimate_occlusion_map: True
+    dense_motion_params:
+      block_expansion: 32
+      max_features: 1024
+      num_blocks: 5
+      reshape_depth: 16
+      compress: 4
+  discriminator_params:
+    scales: [1]
+    block_expansion: 32                 
+    max_features: 512
+    num_blocks: 4
+    sn: True
+  mapping_params:
+      coeff_nc: 70
+      descriptor_nc: 1024
+      layer: 3
+      num_kp: 15
+      num_bins: 66
+

py/src/config/facerender_still.yaml

@@ -0,0 +1,45 @@
+model_params:
+  common_params:
+    num_kp: 15 
+    image_channel: 3                    
+    feature_channel: 32
+    estimate_jacobian: False   # True
+  kp_detector_params:
+     temperature: 0.1
+     block_expansion: 32            
+     max_features: 1024
+     scale_factor: 0.25         # 0.25
+     num_blocks: 5
+     reshape_channel: 16384  # 16384 = 1024 * 16
+     reshape_depth: 16
+  he_estimator_params:
+     block_expansion: 64            
+     max_features: 2048
+     num_bins: 66
+  generator_params:
+    block_expansion: 64
+    max_features: 512
+    num_down_blocks: 2
+    reshape_channel: 32
+    reshape_depth: 16         # 512 = 32 * 16
+    num_resblocks: 6
+    estimate_occlusion_map: True
+    dense_motion_params:
+      block_expansion: 32
+      max_features: 1024
+      num_blocks: 5
+      reshape_depth: 16
+      compress: 4
+  discriminator_params:
+    scales: [1]
+    block_expansion: 32                 
+    max_features: 512
+    num_blocks: 4
+    sn: True
+  mapping_params:
+      coeff_nc: 73
+      descriptor_nc: 1024
+      layer: 3
+      num_kp: 15
+      num_bins: 66
+

py/src/dain_model/base_predictor.py

@@ -0,0 +1,66 @@
+import os
+import numpy as np
+from PIL import Image
+import paddle
+
+
+class BasePredictor(object):
+    def __init__(self):
+        pass
+
+    def build_inference_model(self):
+        if paddle.in_dynamic_mode():
+            # todo self.model = build_model(self.cfg)
+            pass
+        else:
+            place = paddle.get_device()
+            self.exe = paddle.static.Executor(place)
+            file_names = os.listdir(self.weight_path)
+            for file_name in file_names:
+                if file_name.find('model') > -1:
+                    model_file = file_name
+                elif file_name.find('param') > -1:
+                    param_file = file_name
+
+            self.program, self.feed_names, self.fetch_targets = paddle.static.load_inference_model(
+                self.weight_path,
+                executor=self.exe,
+                model_filename=model_file,
+                params_filename=param_file)
+
+    def base_forward(self, inputs):
+        if paddle.in_dynamic_mode():
+            out = self.model(inputs)
+        else:
+            feed_dict = {}
+            if isinstance(inputs, dict):
+                feed_dict = inputs
+            elif isinstance(inputs, (list, tuple)):
+                for i, feed_name in enumerate(self.feed_names):
+                    feed_dict[feed_name] = inputs[i]
+            else:
+                feed_dict[self.feed_names[0]] = inputs
+
+            out = self.exe.run(self.program,
+                               fetch_list=self.fetch_targets,
+                               feed=feed_dict)
+
+        return out
+
+    def is_image(self, input):
+        try:
+            if isinstance(input, (np.ndarray, Image.Image)):
+                return True
+            elif isinstance(input, str):
+                if not os.path.isfile(input):
+                    raise ValueError('input must be a file')
+                img = Image.open(input)
+                _ = img.size
+                return True
+            else:
+                return False
+        except:
+            return False
+
+    def run(self):
+        raise NotImplementedError

py/src/dain_model/dain_predictor.py

@@ -0,0 +1,312 @@
+import os, sys
+import cv2
+import glob
+import shutil
+import numpy as np
+from tqdm import tqdm
+from imageio import imread, imsave
+
+from .base_predictor import BasePredictor
+
+DAIN_WEIGHT_URL = 'https://paddlegan.bj.bcebos.com/applications/DAIN_weight.tar'
+
+
+def video2frames(video_path, outpath, **kargs):
+    def _dict2str(kargs):
+        cmd_str = ''
+        for k, v in kargs.items():
+            cmd_str += (' ' + str(k) + ' ' + str(v))
+        return cmd_str
+
+    ffmpeg = ['ffmpeg ', ' -y -loglevel ', ' error ']
+    vid_name = os.path.basename(video_path).split('.')[0]
+    out_full_path = os.path.join(outpath, vid_name)
+
+    if not os.path.exists(out_full_path):
+        os.makedirs(out_full_path)
+
+    # video file name
+    outformat = os.path.join(out_full_path, '%08d.png')
+
+    cmd = ffmpeg + [' -i ', video_path, ' -start_number ', ' 0 ', outformat]
+
+    cmd = ''.join(cmd) + _dict2str(kargs)
+
+    if os.system(cmd) != 0:
+        raise RuntimeError('ffmpeg process video: {} error'.format(vid_name))
+
+    sys.stdout.flush()
+    return out_full_path
+
+
+def frames2video(frame_path, video_path, r, w, h):
+    out = cv2.VideoWriter(video_path, cv2.VideoWriter_fourcc(*'DIVX'), int(r), (w, h))
+    filelist = sorted(glob.glob(os.path.join(frame_path, '*.png')))
+    for img_path in filelist:
+        img = cv2.imread(img_path)
+        out.write(img)
+    out.release()
+
+
+class DAINPredictor(BasePredictor):
+    def __init__(self,
+                 output='output',
+                 weight_path=None,
+                 time_step=None,
+                 use_gpu=True,
+                 remove_duplicates=False):
+        self.output_path = os.path.join(output, 'DAIN')
+        self.weight_path = weight_path
+        self.time_step = time_step
+        self.key_frame_thread = 0
+        self.remove_duplicates = remove_duplicates
+
+        self.build_inference_model()
+
+    def run(self, video_path):
+        cmd='rm -rf {}/*'.format(self.output_path)
+        os.system(cmd)
+        frame_path_input = os.path.join(self.output_path, 'frames-input')
+        frame_path_interpolated = os.path.join(self.output_path,
+                                               'frames-interpolated')
+        frame_path_combined = os.path.join(self.output_path, 'frames-combined')
+        video_path_output = os.path.join(self.output_path, 'videos-output')
+
+        if not os.path.exists(self.output_path):
+            os.makedirs(self.output_path)
+        if not os.path.exists(frame_path_input):
+            os.makedirs(frame_path_input)
+        if not os.path.exists(frame_path_interpolated):
+            os.makedirs(frame_path_interpolated)
+        if not os.path.exists(frame_path_combined):
+            os.makedirs(frame_path_combined)
+        if not os.path.exists(video_path_output):
+            os.makedirs(video_path_output)
+
+        timestep = self.time_step
+        num_frames = int(1.0 / timestep) - 1
+
+        cap = cv2.VideoCapture(video_path)
+        fps = cap.get(cv2.CAP_PROP_FPS)
+        print("Old fps (frame rate): ", fps)
+
+        times_interp = int(1.0 / timestep)
+        r2 = str(int(fps) * times_interp)
+        print("New fps (frame rate): ", r2)
+
+        out_path = video2frames(video_path, frame_path_input)
+
+        vidname = os.path.basename(video_path).split('.')[0]
+
+        frames = sorted(glob.glob(os.path.join(out_path, '*.png')))
+
+        if self.remove_duplicates:
+            frames = self.remove_duplicate_frames(out_path)
+
+        img = imread(frames[0])
+
+        int_width = img.shape[1]
+        int_height = img.shape[0]
+        channel = img.shape[2]
+        if not channel == 3:
+            return
+
+        if int_width != ((int_width >> 7) << 7):
+            int_width_pad = (((int_width >> 7) + 1) << 7)  # more than necessary
+            padding_left = int((int_width_pad - int_width) / 2)
+            padding_right = int_width_pad - int_width - padding_left
+        else:
+            padding_left = 32
+            padding_right = 32
+
+        if int_height != ((int_height >> 7) << 7):
+            int_height_pad = (
+                    ((int_height >> 7) + 1) << 7)  # more than necessary
+            padding_top = int((int_height_pad - int_height) / 2)
+            padding_bottom = int_height_pad - int_height - padding_top
+        else:
+            padding_top = 32
+            padding_bottom = 32
+
+        frame_num = len(frames)
+
+        if not os.path.exists(os.path.join(frame_path_interpolated, vidname)):
+            os.makedirs(os.path.join(frame_path_interpolated, vidname))
+        if not os.path.exists(os.path.join(frame_path_combined, vidname)):
+            os.makedirs(os.path.join(frame_path_combined, vidname))
+
+        for i in tqdm(range(frame_num - 1)):
+            first = frames[i]
+            second = frames[i + 1]
+            first_index = int(first.split(os.sep)[-1].split('.')[-2])
+            second_index = int(second.split(os.sep)[-1].split('.')[-2])
+
+            img_first = imread(first)
+            img_second = imread(second)
+
+            X0 = img_first.astype('float32').transpose((2, 0, 1)) / 255
+            X1 = img_second.astype('float32').transpose((2, 0, 1)) / 255
+
+            assert (X0.shape[1] == X1.shape[1])
+            assert (X0.shape[2] == X1.shape[2])
+
+            X0 = np.pad(X0, ((0, 0), (padding_top, padding_bottom), \
+                             (padding_left, padding_right)), mode='edge')
+            X1 = np.pad(X1, ((0, 0), (padding_top, padding_bottom), \
+                             (padding_left, padding_right)), mode='edge')
+
+            X0 = np.expand_dims(X0, axis=0)
+            X1 = np.expand_dims(X1, axis=0)
+
+            X0 = np.expand_dims(X0, axis=0)
+            X1 = np.expand_dims(X1, axis=0)
+
+            X = np.concatenate((X0, X1), axis=0)
+
+            o = self.base_forward(X)
+
+            y_ = o[0]
+
+            y_ = [
+                np.transpose(
+                    255.0 * item.clip(
+                        0, 1.0)[0, :, padding_top:padding_top + int_height,
+                            padding_left:padding_left + int_width],
+                    (1, 2, 0)) for item in y_
+            ]
+            if self.remove_duplicates:
+                num_frames = times_interp * (second_index - first_index) - 1
+                time_offsets = [
+                    kk * timestep for kk in range(1, 1 + num_frames, 1)
+                ]
+                start = times_interp * first_index + 1
+                for item, time_offset in zip(y_, time_offsets):
+                    out_dir = os.path.join(frame_path_interpolated, vidname,
+                                           "{:08d}.png".format(start))
+                    imsave(out_dir, np.round(item).astype(np.uint8))
+                    start = start + 1
+
+            else:
+                time_offsets = [
+                    kk * timestep for kk in range(1, 1 + num_frames, 1)
+                ]
+
+                count = 1
+                for item, time_offset in zip(y_, time_offsets):
+                    out_dir = os.path.join(
+                        frame_path_interpolated, vidname,
+                        "{:0>6d}_{:0>4d}.png".format(i, count))
+                    count = count + 1
+                    imsave(out_dir, np.round(item).astype(np.uint8))
+
+        input_dir = os.path.join(frame_path_input, vidname)
+        interpolated_dir = os.path.join(frame_path_interpolated, vidname)
+        combined_dir = os.path.join(frame_path_combined, vidname)
+
+        if self.remove_duplicates:
+            self.combine_frames_with_rm(input_dir, interpolated_dir,
+                                        combined_dir, times_interp)
+
+        else:
+            num_frames = int(1.0 / timestep) - 1
+            self.combine_frames(input_dir, interpolated_dir, combined_dir,
+                                num_frames)
+
+        frame_pattern_combined = os.path.join(frame_path_combined, vidname)
+        # video_pattern_output = os.path.join(video_path_output, vidname + '.mp4')
+        video_pattern_output = os.path.join(video_path_output, vidname + '.avi')
+        if os.path.exists(video_pattern_output):
+            os.remove(video_pattern_output)
+        frames2video(frame_pattern_combined, video_pattern_output, r2, int_width, int_height)
+
+        return frame_pattern_combined, video_pattern_output
+
+    def combine_frames(self, input, interpolated, combined, num_frames):
+        frames1 = sorted(glob.glob(os.path.join(input, '*.png')))
+        frames2 = sorted(glob.glob(os.path.join(interpolated, '*.png')))
+        num1 = len(frames1)
+        num2 = len(frames2)
+
+        for i in range(num1):
+            src = frames1[i]
+            imgname = int(src.split(os.sep)[-1].split('.')[-2])
+            assert i == imgname
+            dst = os.path.join(combined,
+                               '{:08d}.png'.format(i * (num_frames + 1)))
+            shutil.copy2(src, dst)
+            if i < num1 - 1:
+                try:
+                    for k in range(num_frames):
+                        src = frames2[i * num_frames + k]
+                        dst = os.path.join(
+                            combined,
+                            '{:08d}.png'.format(i * (num_frames + 1) + k + 1))
+                        shutil.copy2(src, dst)
+                except Exception as e:
+                    print(e)
+
+    def combine_frames_with_rm(self, input, interpolated, combined,
+                               times_interp):
+        frames1 = sorted(glob.glob(os.path.join(input, '*.png')))
+        frames2 = sorted(glob.glob(os.path.join(interpolated, '*.png')))
+        num1 = len(frames1)
+        num2 = len(frames2)
+
+        for i in range(num1):
+            src = frames1[i]
+            index = int(src.split(os.sep)[-1].split('.')[-2])
+            dst = os.path.join(combined,
+                               '{:08d}.png'.format(times_interp * index))
+            shutil.copy2(src, dst)
+
+        for i in range(num2):
+            src = frames2[i]
+            imgname = src.split(os.sep)[-1]
+            dst = os.path.join(combined, imgname)
+            shutil.copy2(src, dst)
+
+    def remove_duplicate_frames(self, paths):
+        def dhash(image, hash_size=8):
+            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+            resized = cv2.resize(gray, (hash_size + 1, hash_size))
+            diff = resized[:, 1:] > resized[:, :-1]
+            return sum([2 ** i for (i, v) in enumerate(diff.flatten()) if v])
+
+        hashes = {}
+        max_interp = 9
+        image_paths = sorted(glob.glob(os.path.join(paths, '*.png')))
+        for image_path in image_paths:
+            image = cv2.imread(image_path)
+            h = dhash(image)
+            p = hashes.get(h, [])
+            p.append(image_path)
+            hashes[h] = p
+
+        for (h, hashed_paths) in hashes.items():
+            if len(hashed_paths) > 1:
+                first_index = int(hashed_paths[0].split(
+                    os.sep)[-1].split('.')[-2])
+                last_index = int(hashed_paths[-1].split(
+                    os.sep)[-1].split('.')[-2]) + 1
+                gap = 2 * (last_index - first_index) - 1
+                if gap > 2 * max_interp:
+                    cut1 = len(hashed_paths) // 3
+                    cut2 = cut1 * 2
+                    for p in hashed_paths[1:cut1 - 1]:
+                        os.remove(p)
+                    for p in hashed_paths[cut1 + 1:cut2]:
+                        os.remove(p)
+                    for p in hashed_paths[cut2 + 1:]:
+                        os.remove(p)
+                if gap > max_interp:
+                    mid = len(hashed_paths) // 2
+                    for p in hashed_paths[1:mid - 1]:
+                        os.remove(p)
+                    for p in hashed_paths[mid + 1:]:
+                        os.remove(p)
+                else:
+                    for p in hashed_paths[1:]:
+                        os.remove(p)
+
+        frames = sorted(glob.glob(os.path.join(paths, '*.png')))
+        return frames

py/src/face3d/data/__init__.py

@@ -0,0 +1,116 @@
+"""This package includes all the modules related to data loading and preprocessing
+
+ To add a custom dataset class called 'dummy', you need to add a file called 'dummy_dataset.py' and define a subclass 'DummyDataset' inherited from BaseDataset.
+ You need to implement four functions:
+    -- <__init__>:                      initialize the class, first call BaseDataset.__init__(self, opt).
+    -- <__len__>:                       return the size of dataset.
+    -- <__getitem__>:                   get a data point from data loader.
+    -- <modify_commandline_options>:    (optionally) add dataset-specific options and set default options.
+
+Now you can use the dataset class by specifying flag '--dataset_mode dummy'.
+See our template dataset class 'template_dataset.py' for more details.
+"""
+import numpy as np
+import importlib
+import torch.utils.data
+from face3d.data.base_dataset import BaseDataset
+
+
+def find_dataset_using_name(dataset_name):
+    """Import the module "data/[dataset_name]_dataset.py".
+
+    In the file, the class called DatasetNameDataset() will
+    be instantiated. It has to be a subclass of BaseDataset,
+    and it is case-insensitive.
+    """
+    dataset_filename = "data." + dataset_name + "_dataset"
+    datasetlib = importlib.import_module(dataset_filename)
+
+    dataset = None
+    target_dataset_name = dataset_name.replace('_', '') + 'dataset'
+    for name, cls in datasetlib.__dict__.items():
+        if name.lower() == target_dataset_name.lower() \
+           and issubclass(cls, BaseDataset):
+            dataset = cls
+
+    if dataset is None:
+        raise NotImplementedError("In %s.py, there should be a subclass of BaseDataset with class name that matches %s in lowercase." % (dataset_filename, target_dataset_name))
+
+    return dataset
+
+
+def get_option_setter(dataset_name):
+    """Return the static method <modify_commandline_options> of the dataset class."""
+    dataset_class = find_dataset_using_name(dataset_name)
+    return dataset_class.modify_commandline_options
+
+
+def create_dataset(opt, rank=0):
+    """Create a dataset given the option.
+
+    This function wraps the class CustomDatasetDataLoader.
+        This is the main interface between this package and 'train.py'/'test.py'
+
+    Example:
+        >>> from data import create_dataset
+        >>> dataset = create_dataset(opt)
+    """
+    data_loader = CustomDatasetDataLoader(opt, rank=rank)
+    dataset = data_loader.load_data()
+    return dataset
+
+class CustomDatasetDataLoader():
+    """Wrapper class of Dataset class that performs multi-threaded data loading"""
+
+    def __init__(self, opt, rank=0):
+        """Initialize this class
+
+        Step 1: create a dataset instance given the name [dataset_mode]
+        Step 2: create a multi-threaded data loader.
+        """
+        self.opt = opt
+        dataset_class = find_dataset_using_name(opt.dataset_mode)
+        self.dataset = dataset_class(opt)
+        self.sampler = None
+        print("rank %d %s dataset [%s] was created" % (rank, self.dataset.name, type(self.dataset).__name__))
+        if opt.use_ddp and opt.isTrain:
+            world_size = opt.world_size
+            self.sampler = torch.utils.data.distributed.DistributedSampler(
+                    self.dataset,
+                    num_replicas=world_size,
+                    rank=rank,
+                    shuffle=not opt.serial_batches
+                )
+            self.dataloader = torch.utils.data.DataLoader(
+                        self.dataset,
+                        sampler=self.sampler,
+                        num_workers=int(opt.num_threads / world_size), 
+                        batch_size=int(opt.batch_size / world_size), 
+                        drop_last=True)
+        else:
+            self.dataloader = torch.utils.data.DataLoader(
+                self.dataset,
+                batch_size=opt.batch_size,
+                shuffle=(not opt.serial_batches) and opt.isTrain,
+                num_workers=int(opt.num_threads),
+                drop_last=True
+            )
+
+    def set_epoch(self, epoch):
+        self.dataset.current_epoch = epoch
+        if self.sampler is not None:
+            self.sampler.set_epoch(epoch)
+
+    def load_data(self):
+        return self
+
+    def __len__(self):
+        """Return the number of data in the dataset"""
+        return min(len(self.dataset), self.opt.max_dataset_size)
+
+    def __iter__(self):
+        """Return a batch of data"""
+        for i, data in enumerate(self.dataloader):
+            if i * self.opt.batch_size >= self.opt.max_dataset_size:
+                break
+            yield data

py/src/face3d/data/base_dataset.py

@@ -0,0 +1,125 @@
+"""This module implements an abstract base class (ABC) 'BaseDataset' for datasets.
+
+It also includes common transformation functions (e.g., get_transform, __scale_width), which can be later used in subclasses.
+"""
+import random
+import numpy as np
+import torch.utils.data as data
+from PIL import Image
+import torchvision.transforms as transforms
+from abc import ABC, abstractmethod
+
+
+class BaseDataset(data.Dataset, ABC):
+    """This class is an abstract base class (ABC) for datasets.
+
+    To create a subclass, you need to implement the following four functions:
+    -- <__init__>:                      initialize the class, first call BaseDataset.__init__(self, opt).
+    -- <__len__>:                       return the size of dataset.
+    -- <__getitem__>:                   get a data point.
+    -- <modify_commandline_options>:    (optionally) add dataset-specific options and set default options.
+    """
+
+    def __init__(self, opt):
+        """Initialize the class; save the options in the class
+
+        Parameters:
+            opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
+        """
+        self.opt = opt
+        # self.root = opt.dataroot
+        self.current_epoch = 0
+
+    @staticmethod
+    def modify_commandline_options(parser, is_train):
+        """Add new dataset-specific options, and rewrite default values for existing options.
+
+        Parameters:
+            parser          -- original option parser
+            is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
+
+        Returns:
+            the modified parser.
+        """
+        return parser
+
+    @abstractmethod
+    def __len__(self):
+        """Return the total number of images in the dataset."""
+        return 0
+
+    @abstractmethod
+    def __getitem__(self, index):
+        """Return a data point and its metadata information.
+
+        Parameters:
+            index - - a random integer for data indexing
+
+        Returns:
+            a dictionary of data with their names. It ususally contains the data itself and its metadata information.
+        """
+        pass
+
+
+def get_transform(grayscale=False):
+    transform_list = []
+    if grayscale:
+        transform_list.append(transforms.Grayscale(1))
+    transform_list += [transforms.ToTensor()]
+    return transforms.Compose(transform_list)
+
+def get_affine_mat(opt, size):
+    shift_x, shift_y, scale, rot_angle, flip = 0., 0., 1., 0., False
+    w, h = size
+
+    if 'shift' in opt.preprocess:
+        shift_pixs = int(opt.shift_pixs)
+        shift_x = random.randint(-shift_pixs, shift_pixs)
+        shift_y = random.randint(-shift_pixs, shift_pixs)
+    if 'scale' in opt.preprocess:
+        scale = 1 + opt.scale_delta * (2 * random.random() - 1)
+    if 'rot' in opt.preprocess:
+        rot_angle = opt.rot_angle * (2 * random.random() - 1)
+        rot_rad = -rot_angle * np.pi/180
+    if 'flip' in opt.preprocess:
+        flip = random.random() > 0.5
+
+    shift_to_origin = np.array([1, 0, -w//2, 0, 1, -h//2, 0, 0, 1]).reshape([3, 3])
+    flip_mat = np.array([-1 if flip else 1, 0, 0, 0, 1, 0, 0, 0, 1]).reshape([3, 3])
+    shift_mat = np.array([1, 0, shift_x, 0, 1, shift_y, 0, 0, 1]).reshape([3, 3])
+    rot_mat = np.array([np.cos(rot_rad), np.sin(rot_rad), 0, -np.sin(rot_rad), np.cos(rot_rad), 0, 0, 0, 1]).reshape([3, 3])
+    scale_mat = np.array([scale, 0, 0, 0, scale, 0, 0, 0, 1]).reshape([3, 3])
+    shift_to_center = np.array([1, 0, w//2, 0, 1, h//2, 0, 0, 1]).reshape([3, 3])
+    
+    affine = shift_to_center @ scale_mat @ rot_mat @ shift_mat @ flip_mat @ shift_to_origin    
+    affine_inv = np.linalg.inv(affine)
+    return affine, affine_inv, flip
+
+def apply_img_affine(img, affine_inv, method=Image.BICUBIC):
+    return img.transform(img.size, Image.AFFINE, data=affine_inv.flatten()[:6], resample=Image.BICUBIC)
+
+def apply_lm_affine(landmark, affine, flip, size):
+    _, h = size
+    lm = landmark.copy()
+    lm[:, 1] = h - 1 - lm[:, 1]
+    lm = np.concatenate((lm, np.ones([lm.shape[0], 1])), -1)
+    lm = lm @ np.transpose(affine)
+    lm[:, :2] = lm[:, :2] / lm[:, 2:]
+    lm = lm[:, :2]
+    lm[:, 1] = h - 1 - lm[:, 1]
+    if flip:
+        lm_ = lm.copy()
+        lm_[:17] = lm[16::-1]
+        lm_[17:22] = lm[26:21:-1]
+        lm_[22:27] = lm[21:16:-1]
+        lm_[31:36] = lm[35:30:-1]
+        lm_[36:40] = lm[45:41:-1]
+        lm_[40:42] = lm[47:45:-1]
+        lm_[42:46] = lm[39:35:-1]
+        lm_[46:48] = lm[41:39:-1]
+        lm_[48:55] = lm[54:47:-1]
+        lm_[55:60] = lm[59:54:-1]
+        lm_[60:65] = lm[64:59:-1]
+        lm_[65:68] = lm[67:64:-1]
+        lm = lm_
+    return lm

py/src/face3d/data/flist_dataset.py

@@ -0,0 +1,125 @@
+"""This script defines the custom dataset for Deep3DFaceRecon_pytorch
+"""
+
+import os.path
+from data.base_dataset import BaseDataset, get_transform, get_affine_mat, apply_img_affine, apply_lm_affine
+from data.image_folder import make_dataset
+from PIL import Image
+import random
+import util.util as util
+import numpy as np
+import json
+import torch
+from scipy.io import loadmat, savemat
+import pickle
+from util.preprocess import align_img, estimate_norm
+from util.load_mats import load_lm3d
+
+
+def default_flist_reader(flist):
+    """
+    flist format: impath label\nimpath label\n ...(same to caffe's filelist)
+    """
+    imlist = []
+    with open(flist, 'r') as rf:
+        for line in rf.readlines():
+            impath = line.strip()
+            imlist.append(impath)
+
+    return imlist
+
+def jason_flist_reader(flist):
+    with open(flist, 'r') as fp:
+        info = json.load(fp)
+    return info
+
+def parse_label(label):
+    return torch.tensor(np.array(label).astype(np.float32))
+
+
+class FlistDataset(BaseDataset):
+    """
+    It requires one directories to host training images '/path/to/data/train'
+    You can train the model with the dataset flag '--dataroot /path/to/data'.
+    """
+
+    def __init__(self, opt):
+        """Initialize this dataset class.
+
+        Parameters:
+            opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
+        """
+        BaseDataset.__init__(self, opt)
+        
+        self.lm3d_std = load_lm3d(opt.bfm_folder)
+        
+        msk_names = default_flist_reader(opt.flist)
+        self.msk_paths = [os.path.join(opt.data_root, i) for i in msk_names]
+
+        self.size = len(self.msk_paths) 
+        self.opt = opt
+        
+        self.name = 'train' if opt.isTrain else 'val'
+        if '_' in opt.flist:
+            self.name += '_' + opt.flist.split(os.sep)[-1].split('_')[0]
+        
+
+    def __getitem__(self, index):
+        """Return a data point and its metadata information.
+
+        Parameters:
+            index (int)      -- a random integer for data indexing
+
+        Returns a dictionary that contains A, B, A_paths and B_paths
+            img (tensor)       -- an image in the input domain
+            msk (tensor)       -- its corresponding attention mask
+            lm  (tensor)       -- its corresponding 3d landmarks
+            im_paths (str)     -- image paths
+            aug_flag (bool)    -- a flag used to tell whether its raw or augmented
+        """
+        msk_path = self.msk_paths[index % self.size]  # make sure index is within then range
+        img_path = msk_path.replace('mask/', '')
+        lm_path = '.'.join(msk_path.replace('mask', 'landmarks').split('.')[:-1]) + '.txt'
+
+        raw_img = Image.open(img_path).convert('RGB')
+        raw_msk = Image.open(msk_path).convert('RGB')
+        raw_lm = np.loadtxt(lm_path).astype(np.float32)
+
+        _, img, lm, msk = align_img(raw_img, raw_lm, self.lm3d_std, raw_msk)
+        
+        aug_flag = self.opt.use_aug and self.opt.isTrain
+        if aug_flag:
+            img, lm, msk = self._augmentation(img, lm, self.opt, msk)
+        
+        _, H = img.size
+        M = estimate_norm(lm, H)
+        transform = get_transform()
+        img_tensor = transform(img)
+        msk_tensor = transform(msk)[:1, ...]
+        lm_tensor = parse_label(lm)
+        M_tensor = parse_label(M)
+
+
+        return {'imgs': img_tensor, 
+                'lms': lm_tensor, 
+                'msks': msk_tensor, 
+                'M': M_tensor,
+                'im_paths': img_path, 
+                'aug_flag': aug_flag,
+                'dataset': self.name}
+
+    def _augmentation(self, img, lm, opt, msk=None):
+        affine, affine_inv, flip = get_affine_mat(opt, img.size)
+        img = apply_img_affine(img, affine_inv)
+        lm = apply_lm_affine(lm, affine, flip, img.size)
+        if msk is not None:
+            msk = apply_img_affine(msk, affine_inv, method=Image.BILINEAR)
+        return img, lm, msk
+    
+
+
+
+    def __len__(self):
+        """Return the total number of images in the dataset.
+        """
+        return self.size

py/src/face3d/data/image_folder.py

@@ -0,0 +1,66 @@
+"""A modified image folder class
+
+We modify the official PyTorch image folder (https://github.com/pytorch/vision/blob/master/torchvision/datasets/folder.py)
+so that this class can load images from both current directory and its subdirectories.
+"""
+import numpy as np
+import torch.utils.data as data
+
+from PIL import Image
+import os
+import os.path
+
+IMG_EXTENSIONS = [
+    '.jpg', '.JPG', '.jpeg', '.JPEG',
+    '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP',
+    '.tif', '.TIF', '.tiff', '.TIFF',
+]
+
+
+def is_image_file(filename):
+    return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
+
+
+def make_dataset(dir, max_dataset_size=float("inf")):
+    images = []
+    assert os.path.isdir(dir) or os.path.islink(dir), '%s is not a valid directory' % dir
+
+    for root, _, fnames in sorted(os.walk(dir, followlinks=True)):
+        for fname in fnames:
+            if is_image_file(fname):
+                path = os.path.join(root, fname)
+                images.append(path)
+    return images[:min(max_dataset_size, len(images))]
+
+
+def default_loader(path):
+    return Image.open(path).convert('RGB')
+
+
+class ImageFolder(data.Dataset):
+
+    def __init__(self, root, transform=None, return_paths=False,
+                 loader=default_loader):
+        imgs = make_dataset(root)
+        if len(imgs) == 0:
+            raise(RuntimeError("Found 0 images in: " + root + "\n"
+                               "Supported image extensions are: " + ",".join(IMG_EXTENSIONS)))
+
+        self.root = root
+        self.imgs = imgs
+        self.transform = transform
+        self.return_paths = return_paths
+        self.loader = loader
+
+    def __getitem__(self, index):
+        path = self.imgs[index]
+        img = self.loader(path)
+        if self.transform is not None:
+            img = self.transform(img)
+        if self.return_paths:
+            return img, path
+        else:
+            return img
+
+    def __len__(self):
+        return len(self.imgs)

py/src/face3d/data/template_dataset.py

@@ -0,0 +1,75 @@
+"""Dataset class template
+
+This module provides a template for users to implement custom datasets.
+You can specify '--dataset_mode template' to use this dataset.
+The class name should be consistent with both the filename and its dataset_mode option.
+The filename should be <dataset_mode>_dataset.py
+The class name should be <Dataset_mode>Dataset.py
+You need to implement the following functions:
+    -- <modify_commandline_options>:　Add dataset-specific options and rewrite default values for existing options.
+    -- <__init__>: Initialize this dataset class.
+    -- <__getitem__>: Return a data point and its metadata information.
+    -- <__len__>: Return the number of images.
+"""
+from data.base_dataset import BaseDataset, get_transform
+# from data.image_folder import make_dataset
+# from PIL import Image
+
+
+class TemplateDataset(BaseDataset):
+    """A template dataset class for you to implement custom datasets."""
+    @staticmethod
+    def modify_commandline_options(parser, is_train):
+        """Add new dataset-specific options, and rewrite default values for existing options.
+
+        Parameters:
+            parser          -- original option parser
+            is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
+
+        Returns:
+            the modified parser.
+        """
+        parser.add_argument('--new_dataset_option', type=float, default=1.0, help='new dataset option')
+        parser.set_defaults(max_dataset_size=10, new_dataset_option=2.0)  # specify dataset-specific default values
+        return parser
+
+    def __init__(self, opt):
+        """Initialize this dataset class.
+
+        Parameters:
+            opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
+
+        A few things can be done here.
+        - save the options (have been done in BaseDataset)
+        - get image paths and meta information of the dataset.
+        - define the image transformation.
+        """
+        # save the option and dataset root
+        BaseDataset.__init__(self, opt)
+        # get the image paths of your dataset;
+        self.image_paths = []  # You can call sorted(make_dataset(self.root, opt.max_dataset_size)) to get all the image paths under the directory self.root
+        # define the default transform function. You can use <base_dataset.get_transform>; You can also define your custom transform function
+        self.transform = get_transform(opt)
+
+    def __getitem__(self, index):
+        """Return a data point and its metadata information.
+
+        Parameters:
+            index -- a random integer for data indexing
+
+        Returns:
+            a dictionary of data with their names. It usually contains the data itself and its metadata information.
+
+        Step 1: get a random image path: e.g., path = self.image_paths[index]
+        Step 2: load your data from the disk: e.g., image = Image.open(path).convert('RGB').
+        Step 3: convert your data to a PyTorch tensor. You can use helpder functions such as self.transform. e.g., data = self.transform(image)
+        Step 4: return a data point as a dictionary.
+        """
+        path = 'temp'    # needs to be a string
+        data_A = None    # needs to be a tensor
+        data_B = None    # needs to be a tensor
+        return {'data_A': data_A, 'data_B': data_B, 'path': path}
+
+    def __len__(self):
+        """Return the total number of images."""
+        return len(self.image_paths)

py/src/face3d/extract_kp_videos.py

@@ -0,0 +1,109 @@
+import os
+import cv2
+import time
+import glob
+import argparse
+import face_alignment
+import numpy as np
+from PIL import Image
+from tqdm import tqdm
+from itertools import cycle
+
+from torch.multiprocessing import Pool, Process, set_start_method
+
+class KeypointExtractor():
+    def __init__(self, device):
+        self.detector = face_alignment.FaceAlignment(face_alignment.LandmarksType._2D, 
+                                                     device=device)   
+
+    def extract_keypoint(self, images, name=None, info=True):
+        if isinstance(images, list):
+            keypoints = []
+            if info:
+                i_range = tqdm(images,desc='landmark Det:')
+            else:
+                i_range = images
+
+            for image in i_range:
+                current_kp = self.extract_keypoint(image)
+                if np.mean(current_kp) == -1 and keypoints:
+                    keypoints.append(keypoints[-1])
+                else:
+                    keypoints.append(current_kp[None])
+
+            keypoints = np.concatenate(keypoints, 0)
+            if name is not None:
+                np.savetxt(os.path.splitext(name)[0]+'.txt', keypoints.reshape(-1))
+            return keypoints
+        else:
+            while True:
+                try:
+                    keypoints = self.detector.get_landmarks_from_image(np.array(images))[0]
+                    break
+                except RuntimeError as e:
+                    if str(e).startswith('CUDA'):
+                        print("Warning: out of memory, sleep for 1s")
+                        time.sleep(1)
+                    else:
+                        print(e)
+                        break    
+                except TypeError:
+                    print('No face detected in this image')
+                    shape = [68, 2]
+                    keypoints = -1. * np.ones(shape)                    
+                    break
+            if name is not None:
+                np.savetxt(os.path.splitext(name)[0]+'.txt', keypoints.reshape(-1))
+            return keypoints
+
+def read_video(filename):
+    frames = []
+    cap = cv2.VideoCapture(filename)
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if ret:
+            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            frame = Image.fromarray(frame)
+            frames.append(frame)
+        else:
+            break
+    cap.release()
+    return frames
+
+def run(data):
+    filename, opt, device = data
+    os.environ['CUDA_VISIBLE_DEVICES'] = device
+    kp_extractor = KeypointExtractor()
+    images = read_video(filename)
+    name = filename.split('/')[-2:]
+    os.makedirs(os.path.join(opt.output_dir, name[-2]), exist_ok=True)
+    kp_extractor.extract_keypoint(
+        images, 
+        name=os.path.join(opt.output_dir, name[-2], name[-1])
+    )
+
+if __name__ == '__main__':
+    set_start_method('spawn')
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument('--input_dir', type=str, help='the folder of the input files')
+    parser.add_argument('--output_dir', type=str, help='the folder of the output files')
+    parser.add_argument('--device_ids', type=str, default='0,1')
+    parser.add_argument('--workers', type=int, default=4)
+
+    opt = parser.parse_args()
+    filenames = list()
+    VIDEO_EXTENSIONS_LOWERCASE = {'mp4'}
+    VIDEO_EXTENSIONS = VIDEO_EXTENSIONS_LOWERCASE.union({f.upper() for f in VIDEO_EXTENSIONS_LOWERCASE})
+    extensions = VIDEO_EXTENSIONS
+    
+    for ext in extensions:
+        os.listdir(f'{opt.input_dir}')
+        print(f'{opt.input_dir}/*.{ext}')
+        filenames = sorted(glob.glob(f'{opt.input_dir}/*.{ext}'))
+    print('Total number of videos:', len(filenames))
+    pool = Pool(opt.workers)
+    args_list = cycle([opt])
+    device_ids = opt.device_ids.split(",")
+    device_ids = cycle(device_ids)
+    for data in tqdm(pool.imap_unordered(run, zip(filenames, args_list, device_ids))):
+        None

py/src/face3d/extract_kp_videos_safe.py

@@ -0,0 +1,126 @@
+import os
+import cv2
+import time
+import glob
+import argparse
+import numpy as np
+from PIL import Image
+import torch
+from tqdm import tqdm
+from itertools import cycle
+from facexlib.alignment import init_alignment_model, landmark_98_to_68
+from facexlib.detection import init_detection_model
+from torch.multiprocessing import Pool, Process, set_start_method
+import folder_paths
+
+
+class KeypointExtractor():
+    def __init__(self, device='cuda'):
+        root_path = 'gfpgan/weights'
+        self.detector = init_alignment_model('awing_fan',device=device, model_rootpath=root_path)   
+        self.det_net = init_detection_model('retinaface_resnet50', half=False,device=device, model_rootpath=root_path)
+
+    def extract_keypoint(self, images, name=None, info=True):
+        if isinstance(images, list):
+            keypoints = []
+            if info:
+                i_range = tqdm(images,desc='landmark Det:')
+            else:
+                i_range = images
+
+            for image in i_range:
+                current_kp = self.extract_keypoint(image)
+                # current_kp = self.detector.get_landmarks(np.array(image))
+                if np.mean(current_kp) == -1 and keypoints:
+                    keypoints.append(keypoints[-1])
+                else:
+                    keypoints.append(current_kp[None])
+
+            keypoints = np.concatenate(keypoints, 0)
+            np.savetxt(os.path.splitext(name)[0]+'.txt', keypoints.reshape(-1))
+            return keypoints
+        else:
+            while True:
+                try:
+                    with torch.no_grad():
+                        # face detection -> face alignment.
+                        img = np.array(images)
+                        bboxes = self.det_net.detect_faces(images, 0.97)
+                        
+                        bboxes = bboxes[0]
+                        img = img[int(bboxes[1]):int(bboxes[3]), int(bboxes[0]):int(bboxes[2]), :]
+
+                        keypoints = landmark_98_to_68(self.detector.get_landmarks(img)) # [0]
+
+                        #### keypoints to the original location
+                        keypoints[:,0] += int(bboxes[0])
+                        keypoints[:,1] += int(bboxes[1])
+
+                        break
+                except RuntimeError as e:
+                    if str(e).startswith('CUDA'):
+                        print("Warning: out of memory, sleep for 1s")
+                        time.sleep(1)
+                    else:
+                        print(e)
+                        break    
+                except TypeError:
+                    print('No face detected in this image')
+                    shape = [68, 2]
+                    keypoints = -1. * np.ones(shape)                    
+                    break
+            if name is not None:
+                np.savetxt(os.path.splitext(name)[0]+'.txt', keypoints.reshape(-1))
+            return keypoints
+
+def read_video(filename):
+    frames = []
+    cap = cv2.VideoCapture(filename)
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if ret:
+            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            frame = Image.fromarray(frame)
+            frames.append(frame)
+        else:
+            break
+    cap.release()
+    return frames
+
+def run(data):
+    filename, opt, device = data
+    os.environ['CUDA_VISIBLE_DEVICES'] = device
+    kp_extractor = KeypointExtractor()
+    images = read_video(filename)
+    name = filename.split('/')[-2:]
+    os.makedirs(os.path.join(opt.output_dir, name[-2]), exist_ok=True)
+    kp_extractor.extract_keypoint(
+        images, 
+        name=os.path.join(opt.output_dir, name[-2], name[-1])
+    )
+
+if __name__ == '__main__':
+    set_start_method('spawn')
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument('--input_dir', type=str, help='the folder of the input files')
+    parser.add_argument('--output_dir', type=str, help='the folder of the output files')
+    parser.add_argument('--device_ids', type=str, default='0,1')
+    parser.add_argument('--workers', type=int, default=4)
+
+    opt = parser.parse_args()
+    filenames = list()
+    VIDEO_EXTENSIONS_LOWERCASE = {'mp4'}
+    VIDEO_EXTENSIONS = VIDEO_EXTENSIONS_LOWERCASE.union({f.upper() for f in VIDEO_EXTENSIONS_LOWERCASE})
+    extensions = VIDEO_EXTENSIONS
+    
+    for ext in extensions:
+        os.listdir(f'{opt.input_dir}')
+        print(f'{opt.input_dir}/*.{ext}')
+        filenames = sorted(glob.glob(f'{opt.input_dir}/*.{ext}'))
+    print('Total number of videos:', len(filenames))
+    pool = Pool(opt.workers)
+    args_list = cycle([opt])
+    device_ids = opt.device_ids.split(",")
+    device_ids = cycle(device_ids)
+    for data in tqdm(pool.imap_unordered(run, zip(filenames, args_list, device_ids))):
+        None

py/src/face3d/models/__init__.py

@@ -0,0 +1,67 @@
+"""This package contains modules related to objective functions, optimizations, and network architectures.
+
+To add a custom model class called 'dummy', you need to add a file called 'dummy_model.py' and define a subclass DummyModel inherited from BaseModel.
+You need to implement the following five functions:
+    -- <__init__>:                      initialize the class; first call BaseModel.__init__(self, opt).
+    -- <set_input>:                     unpack data from dataset and apply preprocessing.
+    -- <forward>:                       produce intermediate results.
+    -- <optimize_parameters>:           calculate loss, gradients, and update network weights.
+    -- <modify_commandline_options>:    (optionally) add model-specific options and set default options.
+
+In the function <__init__>, you need to define four lists:
+    -- self.loss_names (str list):          specify the training losses that you want to plot and save.
+    -- self.model_names (str list):         define networks used in our training.
+    -- self.visual_names (str list):        specify the images that you want to display and save.
+    -- self.optimizers (optimizer list):    define and initialize optimizers. You can define one optimizer for each network. If two networks are updated at the same time, you can use itertools.chain to group them. See cycle_gan_model.py for an usage.
+
+Now you can use the model class by specifying flag '--model dummy'.
+See our template model class 'template_model.py' for more details.
+"""
+
+import importlib
+from .base_model import BaseModel
+
+
+def find_model_using_name(model_name):
+    """Import the module "models/[model_name]_model.py".
+
+    In the file, the class called DatasetNameModel() will
+    be instantiated. It has to be a subclass of BaseModel,
+    and it is case-insensitive.
+    """
+    model_filename = "face3d.models." + model_name + "_model"
+    modellib = importlib.import_module(model_filename)
+    model = None
+    target_model_name = model_name.replace('_', '') + 'model'
+    for name, cls in modellib.__dict__.items():
+        if name.lower() == target_model_name.lower() \
+           and issubclass(cls, BaseModel):
+            model = cls
+
+    if model is None:
+        print("In %s.py, there should be a subclass of BaseModel with class name that matches %s in lowercase." % (model_filename, target_model_name))
+        exit(0)
+
+    return model
+
+
+def get_option_setter(model_name):
+    """Return the static method <modify_commandline_options> of the model class."""
+    model_class = find_model_using_name(model_name)
+    return model_class.modify_commandline_options
+
+
+def create_model(opt):
+    """Create a model given the option.
+
+    This function warps the class CustomDatasetDataLoader.
+    This is the main interface between this package and 'train.py'/'test.py'
+
+    Example:
+        >>> from models import create_model
+        >>> model = create_model(opt)
+    """
+    model = find_model_using_name(opt.model)
+    instance = model(opt)
+    print("model [%s] was created" % type(instance).__name__)
+    return instance

py/src/face3d/models/arcface_torch/README.md

@@ -0,0 +1,164 @@
+# Distributed Arcface Training in Pytorch
+
+This is a deep learning library that makes face recognition efficient, and effective, which can train tens of millions
+identity on a single server.
+
+## Requirements
+
+- Install [pytorch](http://pytorch.org) (torch>=1.6.0), our doc for [install.md](docs/install.md).
+- `pip install -r requirements.txt`.
+- Download the dataset
+  from [https://github.com/deepinsight/insightface/tree/master/recognition/_datasets_](https://github.com/deepinsight/insightface/tree/master/recognition/_datasets_)
+  .
+
+## How to Training
+
+To train a model, run `train.py` with the path to the configs:
+
+### 1. Single node, 8 GPUs:
+
+```shell
+python -m torch.distributed.launch --nproc_per_node=8 --nnodes=1 --node_rank=0 --master_addr="127.0.0.1" --master_port=1234 train.py configs/ms1mv3_r50
+```
+
+### 2. Multiple nodes, each node 8 GPUs:
+
+Node 0:
+
+```shell
+python -m torch.distributed.launch --nproc_per_node=8 --nnodes=2 --node_rank=0 --master_addr="ip1" --master_port=1234 train.py train.py configs/ms1mv3_r50
+```
+
+Node 1:
+
+```shell
+python -m torch.distributed.launch --nproc_per_node=8 --nnodes=2 --node_rank=1 --master_addr="ip1" --master_port=1234 train.py train.py configs/ms1mv3_r50
+```
+
+### 3.Training resnet2060 with 8 GPUs:
+
+```shell
+python -m torch.distributed.launch --nproc_per_node=8 --nnodes=1 --node_rank=0 --master_addr="127.0.0.1" --master_port=1234 train.py configs/ms1mv3_r2060.py
+```
+
+## Model Zoo
+
+- The models are available for non-commercial research purposes only.  
+- All models can be found in here.  
+- [Baidu Yun Pan](https://pan.baidu.com/s/1CL-l4zWqsI1oDuEEYVhj-g):   e8pw  
+- [onedrive](https://1drv.ms/u/s!AswpsDO2toNKq0lWY69vN58GR6mw?e=p9Ov5d)
+
+### Performance on [**ICCV2021-MFR**](http://iccv21-mfr.com/)
+
+ICCV2021-MFR testset consists of non-celebrities so we can ensure that it has very few overlap with public available face 
+recognition training set, such as MS1M and CASIA as they mostly collected from online celebrities. 
+As the result, we can evaluate the FAIR performance for different algorithms.  
+
+For **ICCV2021-MFR-ALL** set, TAR is measured on all-to-all 1:1 protocal, with FAR less than 0.000001(e-6). The 
+globalised multi-racial testset contains 242,143 identities and 1,624,305 images. 
+
+For **ICCV2021-MFR-MASK** set, TAR is measured on mask-to-nonmask 1:1 protocal, with FAR less than 0.0001(e-4). 
+Mask testset contains 6,964 identities, 6,964 masked images and 13,928 non-masked images. 
+There are totally 13,928 positive pairs and 96,983,824 negative pairs.
+
+| Datasets | backbone  | Training throughout | Size / MB  | **ICCV2021-MFR-MASK** | **ICCV2021-MFR-ALL** |
+| :---:    | :---      | :---                | :---       |:---                   |:---                  |     
+| MS1MV3    | r18  | -              | 91   | **47.85** | **68.33** |
+| Glint360k | r18  | 8536           | 91   | **53.32** | **72.07** |
+| MS1MV3    | r34  | -              | 130  | **58.72** | **77.36** |
+| Glint360k | r34  | 6344           | 130  | **65.10** | **83.02** |
+| MS1MV3    | r50  | 5500           | 166  | **63.85** | **80.53** |
+| Glint360k | r50  | 5136           | 166  | **70.23** | **87.08** |
+| MS1MV3    | r100 | -              | 248  | **69.09** | **84.31** |
+| Glint360k | r100 | 3332           | 248  | **75.57** | **90.66** |
+| MS1MV3    | mobilefacenet | 12185 | 7.8  | **41.52** | **65.26** |        
+| Glint360k | mobilefacenet | 11197 | 7.8  | **44.52** | **66.48** |  
+
+### Performance on IJB-C and Verification Datasets
+
+|   Datasets | backbone      | IJBC(1e-05) | IJBC(1e-04) | agedb30 | cfp_fp | lfw  |  log    |
+| :---:      |    :---       | :---          | :---  | :---  |:---   |:---    |:---     |  
+| MS1MV3     | r18      | 92.07 | 94.66 | 97.77 | 97.73 | 99.77 |[log](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/ms1mv3_arcface_r18_fp16/training.log)|         
+| MS1MV3     | r34      | 94.10 | 95.90 | 98.10 | 98.67 | 99.80 |[log](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/ms1mv3_arcface_r34_fp16/training.log)|        
+| MS1MV3     | r50      | 94.79 | 96.46 | 98.35 | 98.96 | 99.83 |[log](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/ms1mv3_arcface_r50_fp16/training.log)|         
+| MS1MV3     | r100     | 95.31 | 96.81 | 98.48 | 99.06 | 99.85 |[log](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/ms1mv3_arcface_r100_fp16/training.log)|        
+| MS1MV3     | **r2060**| 95.34 | 97.11 | 98.67 | 99.24 | 99.87 |[log](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/ms1mv3_arcface_r2060_fp16/training.log)|
+| Glint360k  |r18-0.1   | 93.16 | 95.33 | 97.72 | 97.73 | 99.77 |[log](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/glint360k_cosface_r18_fp16_0.1/training.log)| 
+| Glint360k  |r34-0.1   | 95.16 | 96.56 | 98.33 | 98.78 | 99.82 |[log](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/glint360k_cosface_r34_fp16_0.1/training.log)| 
+| Glint360k  |r50-0.1   | 95.61 | 96.97 | 98.38 | 99.20 | 99.83 |[log](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/glint360k_cosface_r50_fp16_0.1/training.log)| 
+| Glint360k  |r100-0.1  | 95.88 | 97.32 | 98.48 | 99.29 | 99.82 |[log](https://raw.githubusercontent.com/anxiangsir/insightface_arcface_log/master/glint360k_cosface_r100_fp16_0.1/training.log)|
+
+[comment]: <> (More details see [model.md]&#40;docs/modelzoo.md&#41; in docs.)
+
+
+## [Speed Benchmark](docs/speed_benchmark.md)
+
+**Arcface Torch** can train large-scale face recognition training set efficiently and quickly. When the number of
+classes in training sets is greater than 300K and the training is sufficient, partial fc sampling strategy will get same
+accuracy with several times faster training performance and smaller GPU memory. 
+Partial FC is a sparse variant of the model parallel architecture for large sacle  face recognition. Partial FC use a 
+sparse softmax, where each batch dynamicly sample a subset of class centers for training. In each iteration, only a 
+sparse part of the parameters will be updated, which can reduce a lot of GPU memory and calculations. With Partial FC, 
+we can scale trainset of 29 millions identities, the largest to date. Partial FC also supports multi-machine distributed 
+training and mixed precision training.
+
+![Image text](https://github.com/anxiangsir/insightface_arcface_log/blob/master/partial_fc_v2.png)
+
+More details see 
+[speed_benchmark.md](docs/speed_benchmark.md) in docs.
+
+### 1. Training speed of different parallel methods (samples / second), Tesla V100 32GB * 8. (Larger is better)
+
+`-` means training failed because of gpu memory limitations.
+
+| Number of Identities in Dataset | Data Parallel | Model Parallel | Partial FC 0.1 |
+| :---    | :--- | :--- | :--- |
+|125000   | 4681         | 4824          | 5004     |
+|1400000  | **1672**     | 3043          | 4738     |
+|5500000  | **-**        | **1389**      | 3975     |
+|8000000  | **-**        | **-**         | 3565     |
+|16000000 | **-**        | **-**         | 2679     |
+|29000000 | **-**        | **-**         | **1855** |
+
+### 2. GPU memory cost of different parallel methods (MB per GPU), Tesla V100 32GB * 8. (Smaller is better)
+
+| Number of Identities in Dataset | Data Parallel | Model Parallel | Partial FC 0.1 |
+| :---    | :---      | :---      | :---  |
+|125000   | 7358      | 5306      | 4868  |
+|1400000  | 32252     | 11178     | 6056  |
+|5500000  | **-**     | 32188     | 9854  |
+|8000000  | **-**     | **-**     | 12310 |
+|16000000 | **-**     | **-**     | 19950 |
+|29000000 | **-**     | **-**     | 32324 |
+
+## Evaluation ICCV2021-MFR and IJB-C
+
+More details see [eval.md](docs/eval.md) in docs.
+
+## Test
+
+We tested many versions of PyTorch. Please create an issue if you are having trouble.  
+
+- [x] torch 1.6.0
+- [x] torch 1.7.1
+- [x] torch 1.8.0
+- [x] torch 1.9.0
+
+## Citation
+
+```
+@inproceedings{deng2019arcface,
+  title={Arcface: Additive angular margin loss for deep face recognition},
+  author={Deng, Jiankang and Guo, Jia and Xue, Niannan and Zafeiriou, Stefanos},
+  booktitle={Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition},
+  pages={4690--4699},
+  year={2019}
+}
+@inproceedings{an2020partical_fc,
+  title={Partial FC: Training 10 Million Identities on a Single Machine},
+  author={An, Xiang and Zhu, Xuhan and Xiao, Yang and Wu, Lan and Zhang, Ming and Gao, Yuan and Qin, Bin and
+  Zhang, Debing and Fu Ying},
+  booktitle={Arxiv 2010.05222},
+  year={2020}
+}
+```

py/src/face3d/models/arcface_torch/backbones/__init__.py

@@ -0,0 +1,25 @@
+from .iresnet import iresnet18, iresnet34, iresnet50, iresnet100, iresnet200
+from .mobilefacenet import get_mbf
+
+
+def get_model(name, **kwargs):
+    # resnet
+    if name == "r18":
+        return iresnet18(False, **kwargs)
+    elif name == "r34":
+        return iresnet34(False, **kwargs)
+    elif name == "r50":
+        return iresnet50(False, **kwargs)
+    elif name == "r100":
+        return iresnet100(False, **kwargs)
+    elif name == "r200":
+        return iresnet200(False, **kwargs)
+    elif name == "r2060":
+        from .iresnet2060 import iresnet2060
+        return iresnet2060(False, **kwargs)
+    elif name == "mbf":
+        fp16 = kwargs.get("fp16", False)
+        num_features = kwargs.get("num_features", 512)
+        return get_mbf(fp16=fp16, num_features=num_features)
+    else:
+        raise ValueError()

py/src/face3d/models/arcface_torch/backbones/iresnet.py

@@ -0,0 +1,187 @@
+import torch
+from torch import nn
+
+__all__ = ['iresnet18', 'iresnet34', 'iresnet50', 'iresnet100', 'iresnet200']
+
+
+def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes,
+                     out_planes,
+                     kernel_size=3,
+                     stride=stride,
+                     padding=dilation,
+                     groups=groups,
+                     bias=False,
+                     dilation=dilation)
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution"""
+    return nn.Conv2d(in_planes,
+                     out_planes,
+                     kernel_size=1,
+                     stride=stride,
+                     bias=False)
+
+
+class IBasicBlock(nn.Module):
+    expansion = 1
+    def __init__(self, inplanes, planes, stride=1, downsample=None,
+                 groups=1, base_width=64, dilation=1):
+        super(IBasicBlock, self).__init__()
+        if groups != 1 or base_width != 64:
+            raise ValueError('BasicBlock only supports groups=1 and base_width=64')
+        if dilation > 1:
+            raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
+        self.bn1 = nn.BatchNorm2d(inplanes, eps=1e-05,)
+        self.conv1 = conv3x3(inplanes, planes)
+        self.bn2 = nn.BatchNorm2d(planes, eps=1e-05,)
+        self.prelu = nn.PReLU(planes)
+        self.conv2 = conv3x3(planes, planes, stride)
+        self.bn3 = nn.BatchNorm2d(planes, eps=1e-05,)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+        out = self.bn1(x)
+        out = self.conv1(out)
+        out = self.bn2(out)
+        out = self.prelu(out)
+        out = self.conv2(out)
+        out = self.bn3(out)
+        if self.downsample is not None:
+            identity = self.downsample(x)
+        out += identity
+        return out
+
+
+class IResNet(nn.Module):
+    fc_scale = 7 * 7
+    def __init__(self,
+                 block, layers, dropout=0, num_features=512, zero_init_residual=False,
+                 groups=1, width_per_group=64, replace_stride_with_dilation=None, fp16=False):
+        super(IResNet, self).__init__()
+        self.fp16 = fp16
+        self.inplanes = 64
+        self.dilation = 1
+        if replace_stride_with_dilation is None:
+            replace_stride_with_dilation = [False, False, False]
+        if len(replace_stride_with_dilation) != 3:
+            raise ValueError("replace_stride_with_dilation should be None "
+                             "or a 3-element tuple, got {}".format(replace_stride_with_dilation))
+        self.groups = groups
+        self.base_width = width_per_group
+        self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(self.inplanes, eps=1e-05)
+        self.prelu = nn.PReLU(self.inplanes)
+        self.layer1 = self._make_layer(block, 64, layers[0], stride=2)
+        self.layer2 = self._make_layer(block,
+                                       128,
+                                       layers[1],
+                                       stride=2,
+                                       dilate=replace_stride_with_dilation[0])
+        self.layer3 = self._make_layer(block,
+                                       256,
+                                       layers[2],
+                                       stride=2,
+                                       dilate=replace_stride_with_dilation[1])
+        self.layer4 = self._make_layer(block,
+                                       512,
+                                       layers[3],
+                                       stride=2,
+                                       dilate=replace_stride_with_dilation[2])
+        self.bn2 = nn.BatchNorm2d(512 * block.expansion, eps=1e-05,)
+        self.dropout = nn.Dropout(p=dropout, inplace=True)
+        self.fc = nn.Linear(512 * block.expansion * self.fc_scale, num_features)
+        self.features = nn.BatchNorm1d(num_features, eps=1e-05)
+        nn.init.constant_(self.features.weight, 1.0)
+        self.features.weight.requires_grad = False
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.normal_(m.weight, 0, 0.1)
+            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, IBasicBlock):
+                    nn.init.constant_(m.bn2.weight, 0)
+
+    def _make_layer(self, block, planes, blocks, stride=1, dilate=False):
+        downsample = None
+        previous_dilation = self.dilation
+        if dilate:
+            self.dilation *= stride
+            stride = 1
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride),
+                nn.BatchNorm2d(planes * block.expansion, eps=1e-05, ),
+            )
+        layers = []
+        layers.append(
+            block(self.inplanes, planes, stride, downsample, self.groups,
+                  self.base_width, previous_dilation))
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(
+                block(self.inplanes,
+                      planes,
+                      groups=self.groups,
+                      base_width=self.base_width,
+                      dilation=self.dilation))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        with torch.cuda.amp.autocast(self.fp16):
+            x = self.conv1(x)
+            x = self.bn1(x)
+            x = self.prelu(x)
+            x = self.layer1(x)
+            x = self.layer2(x)
+            x = self.layer3(x)
+            x = self.layer4(x)
+            x = self.bn2(x)
+            x = torch.flatten(x, 1)
+            x = self.dropout(x)
+        x = self.fc(x.float() if self.fp16 else x)
+        x = self.features(x)
+        return x
+
+
+def _iresnet(arch, block, layers, pretrained, progress, **kwargs):
+    model = IResNet(block, layers, **kwargs)
+    if pretrained:
+        raise ValueError()
+    return model
+
+
+def iresnet18(pretrained=False, progress=True, **kwargs):
+    return _iresnet('iresnet18', IBasicBlock, [2, 2, 2, 2], pretrained,
+                    progress, **kwargs)
+
+
+def iresnet34(pretrained=False, progress=True, **kwargs):
+    return _iresnet('iresnet34', IBasicBlock, [3, 4, 6, 3], pretrained,
+                    progress, **kwargs)
+
+
+def iresnet50(pretrained=False, progress=True, **kwargs):
+    return _iresnet('iresnet50', IBasicBlock, [3, 4, 14, 3], pretrained,
+                    progress, **kwargs)
+
+
+def iresnet100(pretrained=False, progress=True, **kwargs):
+    return _iresnet('iresnet100', IBasicBlock, [3, 13, 30, 3], pretrained,
+                    progress, **kwargs)
+
+
+def iresnet200(pretrained=False, progress=True, **kwargs):
+    return _iresnet('iresnet200', IBasicBlock, [6, 26, 60, 6], pretrained,
+                    progress, **kwargs)
+

py/src/face3d/models/arcface_torch/backbones/iresnet2060.py

@@ -0,0 +1,176 @@
+import torch
+from torch import nn
+
+assert torch.__version__ >= "1.8.1"
+from torch.utils.checkpoint import checkpoint_sequential
+
+__all__ = ['iresnet2060']
+
+
+def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes,
+                     out_planes,
+                     kernel_size=3,
+                     stride=stride,
+                     padding=dilation,
+                     groups=groups,
+                     bias=False,
+                     dilation=dilation)
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution"""
+    return nn.Conv2d(in_planes,
+                     out_planes,
+                     kernel_size=1,
+                     stride=stride,
+                     bias=False)
+
+
+class IBasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None,
+                 groups=1, base_width=64, dilation=1):
+        super(IBasicBlock, self).__init__()
+        if groups != 1 or base_width != 64:
+            raise ValueError('BasicBlock only supports groups=1 and base_width=64')
+        if dilation > 1:
+            raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
+        self.bn1 = nn.BatchNorm2d(inplanes, eps=1e-05, )
+        self.conv1 = conv3x3(inplanes, planes)
+        self.bn2 = nn.BatchNorm2d(planes, eps=1e-05, )
+        self.prelu = nn.PReLU(planes)
+        self.conv2 = conv3x3(planes, planes, stride)
+        self.bn3 = nn.BatchNorm2d(planes, eps=1e-05, )
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+        out = self.bn1(x)
+        out = self.conv1(out)
+        out = self.bn2(out)
+        out = self.prelu(out)
+        out = self.conv2(out)
+        out = self.bn3(out)
+        if self.downsample is not None:
+            identity = self.downsample(x)
+        out += identity
+        return out
+
+
+class IResNet(nn.Module):
+    fc_scale = 7 * 7
+
+    def __init__(self,
+                 block, layers, dropout=0, num_features=512, zero_init_residual=False,
+                 groups=1, width_per_group=64, replace_stride_with_dilation=None, fp16=False):
+        super(IResNet, self).__init__()
+        self.fp16 = fp16
+        self.inplanes = 64
+        self.dilation = 1
+        if replace_stride_with_dilation is None:
+            replace_stride_with_dilation = [False, False, False]
+        if len(replace_stride_with_dilation) != 3:
+            raise ValueError("replace_stride_with_dilation should be None "
+                             "or a 3-element tuple, got {}".format(replace_stride_with_dilation))
+        self.groups = groups
+        self.base_width = width_per_group
+        self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(self.inplanes, eps=1e-05)
+        self.prelu = nn.PReLU(self.inplanes)
+        self.layer1 = self._make_layer(block, 64, layers[0], stride=2)
+        self.layer2 = self._make_layer(block,
+                                       128,
+                                       layers[1],
+                                       stride=2,
+                                       dilate=replace_stride_with_dilation[0])
+        self.layer3 = self._make_layer(block,
+                                       256,
+                                       layers[2],
+                                       stride=2,
+                                       dilate=replace_stride_with_dilation[1])
+        self.layer4 = self._make_layer(block,
+                                       512,
+                                       layers[3],
+                                       stride=2,
+                                       dilate=replace_stride_with_dilation[2])
+        self.bn2 = nn.BatchNorm2d(512 * block.expansion, eps=1e-05, )
+        self.dropout = nn.Dropout(p=dropout, inplace=True)
+        self.fc = nn.Linear(512 * block.expansion * self.fc_scale, num_features)
+        self.features = nn.BatchNorm1d(num_features, eps=1e-05)
+        nn.init.constant_(self.features.weight, 1.0)
+        self.features.weight.requires_grad = False
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.normal_(m.weight, 0, 0.1)
+            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, IBasicBlock):
+                    nn.init.constant_(m.bn2.weight, 0)
+
+    def _make_layer(self, block, planes, blocks, stride=1, dilate=False):
+        downsample = None
+        previous_dilation = self.dilation
+        if dilate:
+            self.dilation *= stride
+            stride = 1
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride),
+                nn.BatchNorm2d(planes * block.expansion, eps=1e-05, ),
+            )
+        layers = []
+        layers.append(
+            block(self.inplanes, planes, stride, downsample, self.groups,
+                  self.base_width, previous_dilation))
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(
+                block(self.inplanes,
+                      planes,
+                      groups=self.groups,
+                      base_width=self.base_width,
+                      dilation=self.dilation))
+
+        return nn.Sequential(*layers)
+
+    def checkpoint(self, func, num_seg, x):
+        if self.training:
+            return checkpoint_sequential(func, num_seg, x)
+        else:
+            return func(x)
+
+    def forward(self, x):
+        with torch.cuda.amp.autocast(self.fp16):
+            x = self.conv1(x)
+            x = self.bn1(x)
+            x = self.prelu(x)
+            x = self.layer1(x)
+            x = self.checkpoint(self.layer2, 20, x)
+            x = self.checkpoint(self.layer3, 100, x)
+            x = self.layer4(x)
+            x = self.bn2(x)
+            x = torch.flatten(x, 1)
+            x = self.dropout(x)
+        x = self.fc(x.float() if self.fp16 else x)
+        x = self.features(x)
+        return x
+
+
+def _iresnet(arch, block, layers, pretrained, progress, **kwargs):
+    model = IResNet(block, layers, **kwargs)
+    if pretrained:
+        raise ValueError()
+    return model
+
+
+def iresnet2060(pretrained=False, progress=True, **kwargs):
+    return _iresnet('iresnet2060', IBasicBlock, [3, 128, 1024 - 128, 3], pretrained, progress, **kwargs)

py/src/face3d/models/arcface_torch/backbones/mobilefacenet.py

@@ -0,0 +1,130 @@
+'''
+Adapted from https://github.com/cavalleria/cavaface.pytorch/blob/master/backbone/mobilefacenet.py
+Original author cavalleria
+'''
+
+import torch.nn as nn
+from torch.nn import Linear, Conv2d, BatchNorm1d, BatchNorm2d, PReLU, Sequential, Module
+import torch
+
+
+class Flatten(Module):
+    def forward(self, x):
+        return x.view(x.size(0), -1)
+
+
+class ConvBlock(Module):
+    def __init__(self, in_c, out_c, kernel=(1, 1), stride=(1, 1), padding=(0, 0), groups=1):
+        super(ConvBlock, self).__init__()
+        self.layers = nn.Sequential(
+            Conv2d(in_c, out_c, kernel, groups=groups, stride=stride, padding=padding, bias=False),
+            BatchNorm2d(num_features=out_c),
+            PReLU(num_parameters=out_c)
+        )
+
+    def forward(self, x):
+        return self.layers(x)
+
+
+class LinearBlock(Module):
+    def __init__(self, in_c, out_c, kernel=(1, 1), stride=(1, 1), padding=(0, 0), groups=1):
+        super(LinearBlock, self).__init__()
+        self.layers = nn.Sequential(
+            Conv2d(in_c, out_c, kernel, stride, padding, groups=groups, bias=False),
+            BatchNorm2d(num_features=out_c)
+        )
+
+    def forward(self, x):
+        return self.layers(x)
+
+
+class DepthWise(Module):
+    def __init__(self, in_c, out_c, residual=False, kernel=(3, 3), stride=(2, 2), padding=(1, 1), groups=1):
+        super(DepthWise, self).__init__()
+        self.residual = residual
+        self.layers = nn.Sequential(
+            ConvBlock(in_c, out_c=groups, kernel=(1, 1), padding=(0, 0), stride=(1, 1)),
+            ConvBlock(groups, groups, groups=groups, kernel=kernel, padding=padding, stride=stride),
+            LinearBlock(groups, out_c, kernel=(1, 1), padding=(0, 0), stride=(1, 1))
+        )
+
+    def forward(self, x):
+        short_cut = None
+        if self.residual:
+            short_cut = x
+        x = self.layers(x)
+        if self.residual:
+            output = short_cut + x
+        else:
+            output = x
+        return output
+
+
+class Residual(Module):
+    def __init__(self, c, num_block, groups, kernel=(3, 3), stride=(1, 1), padding=(1, 1)):
+        super(Residual, self).__init__()
+        modules = []
+        for _ in range(num_block):
+            modules.append(DepthWise(c, c, True, kernel, stride, padding, groups))
+        self.layers = Sequential(*modules)
+
+    def forward(self, x):
+        return self.layers(x)
+
+
+class GDC(Module):
+    def __init__(self, embedding_size):
+        super(GDC, self).__init__()
+        self.layers = nn.Sequential(
+            LinearBlock(512, 512, groups=512, kernel=(7, 7), stride=(1, 1), padding=(0, 0)),
+            Flatten(),
+            Linear(512, embedding_size, bias=False),
+            BatchNorm1d(embedding_size))
+
+    def forward(self, x):
+        return self.layers(x)
+
+
+class MobileFaceNet(Module):
+    def __init__(self, fp16=False, num_features=512):
+        super(MobileFaceNet, self).__init__()
+        scale = 2
+        self.fp16 = fp16
+        self.layers = nn.Sequential(
+            ConvBlock(3, 64 * scale, kernel=(3, 3), stride=(2, 2), padding=(1, 1)),
+            ConvBlock(64 * scale, 64 * scale, kernel=(3, 3), stride=(1, 1), padding=(1, 1), groups=64),
+            DepthWise(64 * scale, 64 * scale, kernel=(3, 3), stride=(2, 2), padding=(1, 1), groups=128),
+            Residual(64 * scale, num_block=4, groups=128, kernel=(3, 3), stride=(1, 1), padding=(1, 1)),
+            DepthWise(64 * scale, 128 * scale, kernel=(3, 3), stride=(2, 2), padding=(1, 1), groups=256),
+            Residual(128 * scale, num_block=6, groups=256, kernel=(3, 3), stride=(1, 1), padding=(1, 1)),
+            DepthWise(128 * scale, 128 * scale, kernel=(3, 3), stride=(2, 2), padding=(1, 1), groups=512),
+            Residual(128 * scale, num_block=2, groups=256, kernel=(3, 3), stride=(1, 1), padding=(1, 1)),
+        )
+        self.conv_sep = ConvBlock(128 * scale, 512, kernel=(1, 1), stride=(1, 1), padding=(0, 0))
+        self.features = GDC(num_features)
+        self._initialize_weights()
+
+    def _initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+                if m.bias is not None:
+                    m.bias.data.zero_()
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.Linear):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+                if m.bias is not None:
+                    m.bias.data.zero_()
+
+    def forward(self, x):
+        with torch.cuda.amp.autocast(self.fp16):
+            x = self.layers(x)
+        x = self.conv_sep(x.float() if self.fp16 else x)
+        x = self.features(x)
+        return x
+
+
+def get_mbf(fp16, num_features):
+    return MobileFaceNet(fp16, num_features)

py/src/face3d/models/arcface_torch/configs/3millions.py

@@ -0,0 +1,23 @@
+from easydict import EasyDict as edict
+
+# configs for test speed
+
+config = edict()
+config.loss = "arcface"
+config.network = "r50"
+config.resume = False
+config.output = None
+config.embedding_size = 512
+config.sample_rate = 1.0
+config.fp16 = True
+config.momentum = 0.9
+config.weight_decay = 5e-4
+config.batch_size = 128
+config.lr = 0.1  # batch size is 512
+
+config.rec = "synthetic"
+config.num_classes = 300 * 10000
+config.num_epoch = 30
+config.warmup_epoch = -1
+config.decay_epoch = [10, 16, 22]
+config.val_targets = []

py/src/face3d/models/arcface_torch/configs/3millions_pfc.py

@@ -0,0 +1,23 @@
+from easydict import EasyDict as edict
+
+# configs for test speed
+
+config = edict()
+config.loss = "arcface"
+config.network = "r50"
+config.resume = False
+config.output = None
+config.embedding_size = 512
+config.sample_rate = 0.1
+config.fp16 = True
+config.momentum = 0.9
+config.weight_decay = 5e-4
+config.batch_size = 128
+config.lr = 0.1  # batch size is 512
+
+config.rec = "synthetic"
+config.num_classes = 300 * 10000
+config.num_epoch = 30
+config.warmup_epoch = -1
+config.decay_epoch = [10, 16, 22]
+config.val_targets = []

py/src/face3d/models/arcface_torch/configs/base.py

@@ -0,0 +1,56 @@
+from easydict import EasyDict as edict
+
+# make training faster
+# our RAM is 256G
+# mount -t tmpfs -o size=140G  tmpfs /train_tmp
+
+config = edict()
+config.loss = "arcface"
+config.network = "r50"
+config.resume = False
+config.output = "ms1mv3_arcface_r50"
+
+config.dataset = "ms1m-retinaface-t1"
+config.embedding_size = 512
+config.sample_rate = 1
+config.fp16 = False
+config.momentum = 0.9
+config.weight_decay = 5e-4
+config.batch_size = 128
+config.lr = 0.1  # batch size is 512
+
+if config.dataset == "emore":
+    config.rec = "/train_tmp/faces_emore"
+    config.num_classes = 85742
+    config.num_image = 5822653
+    config.num_epoch = 16
+    config.warmup_epoch = -1
+    config.decay_epoch = [8, 14, ]
+    config.val_targets = ["lfw", ]
+
+elif config.dataset == "ms1m-retinaface-t1":
+    config.rec = "/train_tmp/ms1m-retinaface-t1"
+    config.num_classes = 93431
+    config.num_image = 5179510
+    config.num_epoch = 25
+    config.warmup_epoch = -1
+    config.decay_epoch = [11, 17, 22]
+    config.val_targets = ["lfw", "cfp_fp", "agedb_30"]
+
+elif config.dataset == "glint360k":
+    config.rec = "/train_tmp/glint360k"
+    config.num_classes = 360232
+    config.num_image = 17091657
+    config.num_epoch = 20
+    config.warmup_epoch = -1
+    config.decay_epoch = [8, 12, 15, 18]
+    config.val_targets = ["lfw", "cfp_fp", "agedb_30"]
+
+elif config.dataset == "webface":
+    config.rec = "/train_tmp/faces_webface_112x112"
+    config.num_classes = 10572
+    config.num_image = "forget"
+    config.num_epoch = 34
+    config.warmup_epoch = -1
+    config.decay_epoch = [20, 28, 32]
+    config.val_targets = ["lfw", "cfp_fp", "agedb_30"]

py/src/face3d/models/arcface_torch/configs/glint360k_mbf.py

@@ -0,0 +1,26 @@
+from easydict import EasyDict as edict
+
+# make training faster
+# our RAM is 256G
+# mount -t tmpfs -o size=140G  tmpfs /train_tmp
+
+config = edict()
+config.loss = "cosface"
+config.network = "mbf"
+config.resume = False
+config.output = None
+config.embedding_size = 512
+config.sample_rate = 0.1
+config.fp16 = True
+config.momentum = 0.9
+config.weight_decay = 2e-4
+config.batch_size = 128
+config.lr = 0.1  # batch size is 512
+
+config.rec = "/train_tmp/glint360k"
+config.num_classes = 360232
+config.num_image = 17091657
+config.num_epoch = 20
+config.warmup_epoch = -1
+config.decay_epoch = [8, 12, 15, 18]
+config.val_targets = ["lfw", "cfp_fp", "agedb_30"]

py/src/face3d/models/arcface_torch/configs/glint360k_r100.py

@@ -0,0 +1,26 @@
+from easydict import EasyDict as edict
+
+# make training faster
+# our RAM is 256G
+# mount -t tmpfs -o size=140G  tmpfs /train_tmp
+
+config = edict()
+config.loss = "cosface"
+config.network = "r100"
+config.resume = False
+config.output = None
+config.embedding_size = 512
+config.sample_rate = 1.0
+config.fp16 = True
+config.momentum = 0.9
+config.weight_decay = 5e-4
+config.batch_size = 128
+config.lr = 0.1  # batch size is 512
+
+config.rec = "/train_tmp/glint360k"
+config.num_classes = 360232
+config.num_image = 17091657
+config.num_epoch = 20
+config.warmup_epoch = -1
+config.decay_epoch = [8, 12, 15, 18]
+config.val_targets = ["lfw", "cfp_fp", "agedb_30"]

py/src/face3d/models/arcface_torch/configs/glint360k_r18.py

@@ -0,0 +1,26 @@
+from easydict import EasyDict as edict
+
+# make training faster
+# our RAM is 256G
+# mount -t tmpfs -o size=140G  tmpfs /train_tmp
+
+config = edict()
+config.loss = "cosface"
+config.network = "r18"
+config.resume = False
+config.output = None
+config.embedding_size = 512
+config.sample_rate = 1.0
+config.fp16 = True
+config.momentum = 0.9
+config.weight_decay = 5e-4
+config.batch_size = 128
+config.lr = 0.1  # batch size is 512
+
+config.rec = "/train_tmp/glint360k"
+config.num_classes = 360232
+config.num_image = 17091657
+config.num_epoch = 20
+config.warmup_epoch = -1
+config.decay_epoch = [8, 12, 15, 18]
+config.val_targets = ["lfw", "cfp_fp", "agedb_30"]