diff --git a/ComfyUI-to-Python-Extension b/ComfyUI-to-Python-Extension index d3b82726a..2b96fca8b 160000 --- a/ComfyUI-to-Python-Extension +++ b/ComfyUI-to-Python-Extension @@ -1 +1 @@ -Subproject commit d3b82726a7f185f25f3d0194702ea17df47ee167 +Subproject commit 2b96fca8b52b1ede96c655a57ff9d71eedb5cef5 diff --git a/custom_functions/AI-portrait-wf-py.py b/custom_functions/AI-portrait-wf-py.py new file mode 100644 index 000000000..a86ece651 --- /dev/null +++ b/custom_functions/AI-portrait-wf-py.py @@ -0,0 +1,312 @@ +import os +import random +import sys +from typing import Sequence, Mapping, Any, Union +import torch + + +def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any: + """Returns the value at the given index of a sequence or mapping. + + If the object is a sequence (like list or string), returns the value at the given index. + If the object is a mapping (like a dictionary), returns the value at the index-th key. + + Some return a dictionary, in these cases, we look for the "results" key + + Args: + obj (Union[Sequence, Mapping]): The object to retrieve the value from. + index (int): The index of the value to retrieve. + + Returns: + Any: The value at the given index. + + Raises: + IndexError: If the index is out of bounds for the object and the object is not a mapping. + """ + try: + return obj[index] + except KeyError: + return obj["result"][index] + + +def find_path(name: str, path: str = None) -> str: + """ + Recursively looks at parent folders starting from the given path until it finds the given name. + Returns the path as a Path object if found, or None otherwise. + """ + # If no path is given, use the current working directory + if path is None: + path = os.getcwd() + + # Check if the current directory contains the name + if name in os.listdir(path): + path_name = os.path.join(path, name) + print(f"{name} found: {path_name}") + return path_name + + # Get the parent directory + parent_directory = os.path.dirname(path) + + # If the parent directory is the same as the current directory, we've reached the root and stop the search + if parent_directory == path: + return None + + # Recursively call the function with the parent directory + return find_path(name, parent_directory) + + +def add_comfyui_directory_to_sys_path() -> None: + """ + Add 'ComfyUI' to the sys.path + """ + comfyui_path = find_path("ComfyUI") + if comfyui_path is not None and os.path.isdir(comfyui_path): + sys.path.append(comfyui_path) + print(f"'{comfyui_path}' added to sys.path") + + +def add_extra_model_paths() -> None: + """ + Parse the optional extra_model_paths.yaml file and add the parsed paths to the sys.path. + """ + try: + from main import load_extra_path_config + except ImportError: + print( + "Could not import load_extra_path_config from main.py. Looking in utils.extra_config instead." + ) + from utils.extra_config import load_extra_path_config + + extra_model_paths = find_path("extra_model_paths.yaml") + + if extra_model_paths is not None: + load_extra_path_config(extra_model_paths) + else: + print("Could not find the extra_model_paths config file.") + + +add_comfyui_directory_to_sys_path() +add_extra_model_paths() + + +def import_custom_nodes() -> None: + """Find all custom nodes in the custom_nodes folder and add those node objects to NODE_CLASS_MAPPINGS + + This function sets up a new asyncio event loop, initializes the PromptServer, + creates a PromptQueue, and initializes the custom nodes. + """ + import asyncio + import execution + from nodes import init_extra_nodes + import server + + # Creating a new event loop and setting it as the default loop + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + + # Creating an instance of PromptServer with the loop + server_instance = server.PromptServer(loop) + execution.PromptQueue(server_instance) + + # Initializing custom nodes + init_extra_nodes() + + +from nodes import LoadImage, NODE_CLASS_MAPPINGS + + +def ai_portrait_infer(image_path): + import_custom_nodes() + with torch.inference_mode(): + preprocbuildpipe = NODE_CLASS_MAPPINGS["PreprocBuildPipe"]() + preprocbuildpipe_22 = preprocbuildpipe.load_models( + prompt_cgpath="prompt", + templates_cgpath="templates", + wd14_cgpath="wd14_tagger", + lut_path="lut", + ) + + loadimage = LoadImage() + loadimage_24 = loadimage.load_image(image=image_path) + + facewarppipebuilder = NODE_CLASS_MAPPINGS["FaceWarpPipeBuilder"]() + facewarppipebuilder_31 = facewarppipebuilder.load_models( + detect_model_path="facedetect/scrfd_10g_bnkps_shape640x640.onnx", + deca_dir="deca", + gpu_choose="cuda:0", + ) + + pulidmodelsloader = NODE_CLASS_MAPPINGS["PulidModelsLoader"]() + pulidmodelsloader_34 = pulidmodelsloader.load_models( + base_model_path="checkpoints/LEOSAM_SDXL_v5.0.safetensors", + canny_model_dir="controlnet/Canny", + pulid_model_dir="pulid", + eva_clip_path="eva_clip/EVA02_CLIP_L_336_psz14_s6B.pt", + insightface_dir="insightface", + facexlib_dir="facexlib", + ) + + ipadapterpipebuilder = NODE_CLASS_MAPPINGS["IpadapterPipeBuilder"]() + ipadapterpipebuilder_43 = ipadapterpipebuilder.load_models( + base_model_path="checkpoints/LEOSAM_SDXL_v5.0.safetensors", + ipadapt_model_dir="controlnet/IPAdapter", + buffalo_model_path="insightface", + inference_steps=20, + guidance_scale=7, + ip_adapter_scale=0.9, + dtype="float16", + device="cuda:0", + ) + + enhancebuildpipe = NODE_CLASS_MAPPINGS["EnhanceBuildPipe"]() + enhancebuildpipe_46 = enhancebuildpipe.load_models( + lut_path="lut", gpu_choose="cuda:2", sr_type="ESRGAN", half=True + ) + + faceswappipebuilder = NODE_CLASS_MAPPINGS["FaceSwapPipeBuilder"]() + faceswappipebuilder_51 = faceswappipebuilder.load_models( + swap_own_model="faceswap/swapper_own.pth", + arcface_model="faceswap/arcface_checkpoint.tar", + facealign_config_dir="face_align", + phase1_model="facealign/p1.pt", + phase2_model="facealign/p2.pt", + device="cuda:0", + ) + + preprocgetconds = NODE_CLASS_MAPPINGS["PreprocGetConds"]() + preprocsplitconds = NODE_CLASS_MAPPINGS["PreprocSplitConds"]() + facewarpdetectfacesmethod = NODE_CLASS_MAPPINGS["FaceWarpDetectFacesMethod"]() + facewarpgetfaces3dinfomethod = NODE_CLASS_MAPPINGS[ + "FaceWarpGetFaces3DinfoMethod" + ]() + facewarpwarp3dfaceimgmaskmethod = NODE_CLASS_MAPPINGS[ + "FaceWarpWarp3DfaceImgMaskMethod" + ]() + pulidinferclass = NODE_CLASS_MAPPINGS["PulidInferClass"]() + ipadapterinferclass = NODE_CLASS_MAPPINGS["IpadapterInferClass"]() + faceswapdetectpts = NODE_CLASS_MAPPINGS["FaceSwapDetectPts"]() + faceswapmethod = NODE_CLASS_MAPPINGS["FaceSwapMethod"]() + enhancemainprocess = NODE_CLASS_MAPPINGS["EnhanceMainProcess"]() + + # for q in range(10): + preprocgetconds_23 = preprocgetconds.prepare_conditions( + style="FORMAL", + gender="MALE", + is_child=False, + model=get_value_at_index(preprocbuildpipe_22, 0), + src_img=get_value_at_index(loadimage_24, 0), + ) + + preprocsplitconds_27 = preprocsplitconds.split_conditions( + pipe_conditions=get_value_at_index(preprocgetconds_23, 0) + ) + + facewarpdetectfacesmethod_33 = facewarpdetectfacesmethod.detect_faces( + model=get_value_at_index(facewarppipebuilder_31, 0), + image=get_value_at_index(preprocsplitconds_27, 0), + ) + + facewarpgetfaces3dinfomethod_32 = ( + facewarpgetfaces3dinfomethod.get_faces_3dinfo( + model=get_value_at_index(facewarppipebuilder_31, 0), + image=get_value_at_index(preprocsplitconds_27, 0), + faces=get_value_at_index(facewarpdetectfacesmethod_33, 0), + ) + ) + + facewarpwarp3dfaceimgmaskmethod_36 = ( + facewarpwarp3dfaceimgmaskmethod.warp_3d_face( + model=get_value_at_index(facewarppipebuilder_31, 0), + user_dict=get_value_at_index(facewarpgetfaces3dinfomethod_32, 0), + template_image=get_value_at_index(preprocsplitconds_27, 3), + template_mask_img=get_value_at_index(preprocsplitconds_27, 11), + ) + ) + + pulidinferclass_35 = pulidinferclass.pulid_infer( + prompt=get_value_at_index(preprocsplitconds_27, 1), + negative_prompt=get_value_at_index(preprocsplitconds_27, 2), + strength=0.7, + model=get_value_at_index(pulidmodelsloader_34, 0), + template_image=get_value_at_index( + facewarpwarp3dfaceimgmaskmethod_36, 0 + ), + canny_control=get_value_at_index(preprocsplitconds_27, 4), + user_image=get_value_at_index(preprocsplitconds_27, 0), + mask=get_value_at_index(facewarpwarp3dfaceimgmaskmethod_36, 1), + ) + + facewarpwarp3dfaceimgmaskmethod_41 = ( + facewarpwarp3dfaceimgmaskmethod.warp_3d_face( + model=get_value_at_index(facewarppipebuilder_31, 0), + user_dict=get_value_at_index(facewarpgetfaces3dinfomethod_32, 0), + template_image=get_value_at_index(pulidinferclass_35, 0), + template_mask_img=get_value_at_index( + facewarpwarp3dfaceimgmaskmethod_36, 1 + ), + ) + ) + + ipadapterinferclass_44 = ipadapterinferclass.ipadapt_process( + prompt=get_value_at_index(preprocsplitconds_27, 1), + negative_prompt=get_value_at_index(preprocsplitconds_27, 2), + seed=random.randint(1, 2**64), + strength=0.3, + model=get_value_at_index(ipadapterpipebuilder_43, 0), + template=get_value_at_index(facewarpwarp3dfaceimgmaskmethod_41, 0), + mask=get_value_at_index(facewarpwarp3dfaceimgmaskmethod_41, 1), + croped_face=get_value_at_index(pulidinferclass_35, 1), + ) + + faceswapdetectpts_56 = faceswapdetectpts.detect_face_pts( + ptstype="5", + model=get_value_at_index(faceswappipebuilder_51, 0), + src_image=get_value_at_index(preprocsplitconds_27, 0), + src_faces=get_value_at_index(facewarpdetectfacesmethod_33, 0), + ) + + facewarpdetectfacesmethod_55 = facewarpdetectfacesmethod.detect_faces( + model=get_value_at_index(facewarppipebuilder_31, 0), + image=get_value_at_index(ipadapterinferclass_44, 0), + ) + + faceswapdetectpts_54 = faceswapdetectpts.detect_face_pts( + ptstype="5", + model=get_value_at_index(faceswappipebuilder_51, 0), + src_image=get_value_at_index(ipadapterinferclass_44, 0), + src_faces=get_value_at_index(facewarpdetectfacesmethod_55, 0), + ) + + faceswapdetectpts_52 = faceswapdetectpts.detect_face_pts( + ptstype="256", + model=get_value_at_index(faceswappipebuilder_51, 0), + src_image=get_value_at_index(ipadapterinferclass_44, 0), + src_faces=get_value_at_index(facewarpdetectfacesmethod_55, 0), + ) + + faceswapmethod_53 = faceswapmethod.swap_face( + model=get_value_at_index(faceswappipebuilder_51, 0), + src_image=get_value_at_index(preprocsplitconds_27, 0), + two_stage_image=get_value_at_index(ipadapterinferclass_44, 0), + source_5pts=get_value_at_index(faceswapdetectpts_56, 0), + target_5pts=get_value_at_index(faceswapdetectpts_54, 0), + target_256pts=get_value_at_index(faceswapdetectpts_52, 0), + ) + + _, _, final_img = enhancemainprocess.enhance_process( + style="FORMAL", + is_front="False", + model=get_value_at_index(enhancebuildpipe_46, 0), + conditions=get_value_at_index(preprocgetconds_23, 0), + src_img=get_value_at_index(faceswapmethod_53, 0), + ) + + return final_img + + +import cv2 +if __name__ == "__main__": + img_path = '/home/user/works/projs/WebServer/distributed-server-node/submodules/VisualForge/ComfyUI/input/portraitInput.jpg' + + final_img = ai_portrait_infer(img_path) + cv2.imwrite('fuckaigc.png', final_img)