# !/usr/bin/env python # -*- coding: UTF-8 -*- import os import time import numpy as np import torch from .Sapiens_Pytorch import SapiensPredictor, SapiensConfig from .Sapiens_Pytorch.classes_and_palettes import GOLIATH_CLASSES_FIX from .sapiens.predictor import Sapiens2Predictor from .utils import get_models_path, tensor2cv, tensor2pil from comfy_api.latest import io import folder_paths from .Sapiens_Pytorch.detector import Detector node_cr_path = os.path.dirname(os.path.abspath(__file__)) device = torch.device( "cuda:0") if torch.cuda.is_available() else torch.device( "mps") if torch.backends.mps.is_available() else torch.device( "cpu") # add checkpoints dir weigths_sapiens_current_path = os.path.join(folder_paths.models_dir, "sapiens") if not os.path.exists(weigths_sapiens_current_path): os.makedirs(weigths_sapiens_current_path) folder_paths.add_model_folder_path("sapiens", weigths_sapiens_current_path) # sapiens dir class SapiensLoader(io.ComfyNode): @classmethod def define_schema(cls): return io.Schema( node_id="SapiensLoader", display_name="SapiensLoader", category="Sapiens", inputs=[ io.Combo.Input("seg",options= ["none"] + [i for i in folder_paths.get_filename_list("sapiens") if "seg" in i or "mIoU" in i ] ), io.Combo.Input("depth",options= ["none"] + [i for i in folder_paths.get_filename_list("sapiens") if "depth" in i or ("render_people" in i and "normal" not in i) ] ), io.Combo.Input("normal",options= ["none"] + [i for i in folder_paths.get_filename_list("sapiens") if "normal" in i ] ), io.Combo.Input("pose",options= ["none"] + [i for i in folder_paths.get_filename_list("sapiens") if "pose" in i or "goliath_AP" in i ] ), io.Combo.Input("pointmap",options= ["none"] + [i for i in folder_paths.get_filename_list("sapiens") if "pointmap" in i ] ), io.Combo.Input("albedo",options= ["none"] + [i for i in folder_paths.get_filename_list("sapiens") if "albedo" in i ] ), io.Combo.Input("matting",options= ["none"] + [i for i in folder_paths.get_filename_list("sapiens") if "matting" in i ] ), io.Combo.Input("dtype",options= [ "float32", "bfloat16",]), io.Float.Input("mini_person_h", default=0.5, min=0.0, max=1.0,step=0.05,display_mode=io.NumberDisplay.number), io.Boolean.Input("show_pose_object", default=False), io.Boolean.Input("convert_torchscript", default=False), io.Boolean.Input("V2", default=True), ], outputs=[ io.Model.Output(display_name="model"), ], ) @classmethod def execute(cls, seg,depth,normal,pose,pointmap,albedo,matting,dtype,mini_person_h,show_pose_object,convert_torchscript,V2) -> io.NodeOutput: if dtype == "bfloat16": infer_dtype = torch.bfloat16 else: infer_dtype = torch.float32 if not V2 : # v1 version config = SapiensConfig() config.model_dir = weigths_sapiens_current_path config.pt_type = dtype config.show_pose_object = show_pose_object config.dtype = infer_dtype config.minimum_person_height = mini_person_h config = get_models_path(seg, depth, normal, pose, config, dtype) # currently only for TorchScript models, convert selected FP32 TorchScript Sapiens models to BF16, save them and use them if convert_torchscript: print("converting TorchScript Sapiens models to BF16...") for model_path_attr in ("local_seg_path", "local_depth_path", "local_normal_path", "local_pose_path"): model_path = getattr(config, model_path_attr) if len(model_path) and model_path.endswith("torchscript.pt2"): model_split_path = os.path.splitext(model_path) converted_model_path = model_split_path[0] + "_bf16" + model_split_path[1] print(f'converting "{model_path}" to BF16...') if os.path.exists(converted_model_path): print(f'"{converted_model_path}" already exists, not converting...') else: model = torch.jit.load(model_path) model.eval().to("cuda").to(torch.bfloat16) torch.jit.save(model, converted_model_path) print(f'"{model_path}" converted to BF16 "{converted_model_path}"') setattr(config, model_path_attr, converted_model_path) config.dtype = torch.bfloat16 model = SapiensPredictor(config) else: seg_path = folder_paths.get_full_path("sapiens",seg) if seg != "none" else None pointmap_path = folder_paths.get_full_path("sapiens",pointmap) if pointmap != "none" else None normal_path = folder_paths.get_full_path("sapiens",normal) if normal != "none" else None albedo_path = folder_paths.get_full_path("sapiens",albedo) if albedo != "none" else None pose_path = folder_paths.get_full_path("sapiens",pose) if pose != "none" else None pose_detector=Detector() if pose_path is not None else None matting_path = folder_paths.get_full_path("sapiens",matting) if matting != "none" else None if pointmap_path is not None and seg_path is None: raise("pointmap need segment model to get mask,please select a segment model") model=Sapiens2Predictor(seg_path,pointmap_path,albedo_path,normal_path,pose_path,matting_path,pose_detector,node_cr_path,device,infer_dtype) model.load_model(torch.device("cpu")) # keep model on cpu model.remove_bg=True model.show_pose_object=show_pose_object return io.NodeOutput(model) class SapiensSampler(io.ComfyNode): @classmethod def define_schema(cls): return io.Schema( node_id="SapiensSampler", display_name="SapiensSampler", category="Sapiens", inputs=[ io.Model.Input("model"), io.Image.Input("image"), io.Boolean.Input("save_pose", default=False), io.Int.Input("BG_R", default=255, min=0, max=255,step=1,display_mode=io.NumberDisplay.number), io.Int.Input("BG_G", default=255, min=0, max=255,step=1,display_mode=io.NumberDisplay.number), io.Int.Input("BG_B", default=255, min=0, max=255,step=1,display_mode=io.NumberDisplay.number), io.Conditioning.Input("cond",optional=True), ], outputs=[ io.Image.Output(display_name="seg_img"), io.Image.Output(display_name="pose_img"), io.Image.Output(display_name="depth_img"), io.Image.Output(display_name="normal_img"), io.Image.Output(display_name="pointmap_img"), io.Image.Output(display_name="albedo_img"), io.Image.Output(display_name="matting_img"), io.Mask.Output(display_name="mask"), ], ) @classmethod def execute(cls, model, image,save_pose, BG_R, BG_G, BG_B, cond=None,) -> io.NodeOutput: start = time.perf_counter() zero_tensor = torch.zeros_like(image, dtype=torch.float32, device="cpu") image_list=[image] if image.size()[0] == 1 else torch.chunk(image, chunks=image.size()[0]) albedo_list,pointmap_list,depth_list,matting_list=[],[],[],[] RGB_BG = [BG_R, BG_G, BG_B] cond=[] if not cond else cond print("start predict...") if isinstance(model, SapiensPredictor): if not torch.backends.mps.is_available(): if torch.cuda.is_available(): model.move_to_cuda() img_in=[tensor2pil(i) for i in image_list] seg_list, depth_list, normal_list, pose_list, mask_list = model(img_in, cond, RGB_BG) mask = torch.cat(mask_list, dim=0) if mask_list else torch.zeros((64, 64), dtype=torch.float32, device="cpu") if not torch.backends.mps.is_available(): if torch.cuda.is_available(): model.enable_model_cpu_offload() else: img_in = [tensor2cv(i.squeeze()) for i in image_list] seg_list,normal_list,pointmap_list,albedo_list,pose_list,matting_list,mask_list,matting_mask=model.predict(img_in,cond, RGB_BG) mask_list=matting_mask if len(matting_mask)>0 else mask_list mask= torch.zeros((64, 64), dtype=torch.float32, device="cpu") if not mask_list else torch.stack([torch.from_numpy(i).float() for i in mask_list], dim=0) depth_img = zero_tensor if not depth_list else torch.cat(depth_list, dim=0) pose_img=zero_tensor if not pose_list else torch.cat(pose_list, dim=0) seg_img=zero_tensor if not seg_list else torch.cat(seg_list, dim=0) normal_img=zero_tensor if not normal_list else torch.cat(normal_list, dim=0) pointmap_img=zero_tensor if not pointmap_list else torch.cat(pointmap_list, dim=0) albedo_img=zero_tensor if not albedo_list else torch.cat(albedo_list, dim=0) matting_image=zero_tensor if not matting_list else torch.cat(matting_list, dim=0) if pose_list and save_pose: print(f"pose counts is {len(pose_list)},Save pose as *.npy files in comfyUI output....") for i,img in enumerate(pose_list): np.save(os.path.join(folder_paths.get_output_directory(),f"{i}"),tensor2cv(img)) print(f"ALL inference took: {time.perf_counter() - start:.4f} seconds") return io.NodeOutput(seg_img,pose_img, depth_img, normal_img,pointmap_img,albedo_img,matting_image, mask ) class SapiensSplit(io.ComfyNode): @classmethod def define_schema(cls): return io.Schema( node_id="SapiensSplit", display_name="SapiensSplit", category="Sapiens", inputs=[ io.Boolean.Input("keep_all", default=False), io.Boolean.Input("Background", default=False), io.Boolean.Input("Apparel", default=False), io.Boolean.Input("Eyeglass", default=False), io.Boolean.Input("Face_Neck", default=False), io.Boolean.Input("Hair", default=False), io.Boolean.Input("Left_Foot", default=False), io.Boolean.Input("Left_Hand", default=False), io.Boolean.Input("Left_Lower_Arm", default=False), io.Boolean.Input("Left_Lower_Leg", default=False), io.Boolean.Input("Left_Shoe", default=False), io.Boolean.Input("Left_Sock", default=False), io.Boolean.Input("Left_Upper_Arm", default=False), io.Boolean.Input("Left_Upper_Leg", default=False), io.Boolean.Input("Lower_Clothing", default=False), io.Boolean.Input("Right_Foot", default=False), io.Boolean.Input("Right_Hand", default=False), io.Boolean.Input("Right_Lower_Arm", default=False), io.Boolean.Input("Right_Lower_Leg", default=False), io.Boolean.Input("Right_Shoe", default=False), io.Boolean.Input("Right_Sock", default=False), io.Boolean.Input("Right_Upper_Arm", default=False), io.Boolean.Input("Right_Upper_Leg", default=False), io.Boolean.Input("Torso", default=False), io.Boolean.Input("Upper_Clothing", default=False), io.Boolean.Input("Lower_Lip", default=False), io.Boolean.Input("Upper_Lip", default=False), io.Boolean.Input("Lower_Teeth", default=False), io.Boolean.Input("Upper_Teeth", default=False), io.Boolean.Input("Tongue", default=False), ], outputs=[ io.Conditioning.Output(display_name="conds"), ], ) @classmethod def execute(cls, keep_all,**kwargs) -> io.NodeOutput: option_names = [ "Background", "Apparel", "Eyeglass","Face_Neck", "Hair", "Left_Foot", "Left_Hand", "Left_Lower_Arm", "Left_Lower_Leg", "Left_Shoe", "Left_Sock", "Left_Upper_Arm", "Left_Upper_Leg", "Lower_Clothing", "Right_Foot", "Right_Hand", "Right_Lower_Arm", "Right_Lower_Leg", "Right_Shoe", "Right_Sock", "Right_Upper_Arm", "Right_Upper_Leg", "Torso", "Upper_Clothing", "Lower_Lip", "Upper_Lip", "Lower_Teeth", "Upper_Teeth", "Tongue" ] selected_indices = [] if not keep_all: for i, name in enumerate(option_names): if kwargs.get(name, False): selected_indices.append(i) # select body index seg_select_all = [list(GOLIATH_CLASSES_FIX)[i] for i in selected_indices] if selected_indices else "seg default hunman map." print(f"Select seg part of {seg_select_all},index is {selected_indices} ") return io.NodeOutput(selected_indices)