通义大模型修改Comfyui节点代码

Comfyui已经使用了一段时间,但是因为其开源和众多节点依赖不同导致其使用非常繁杂,加上Github和huggingface的访问问题,任何升级或者安装节点都有可能导致工作流出错,结果就是Comfyui的使用大部分时间都是在解决依赖和节点以及模型的安装上,但这个问题应该暂时是无解的,想使用Comfyui就只能花时间去解决这些问题。这里记录一下通过修改节点代码解决问题的过程,Comfyui官方包几乎每周都有升级,为的是加快速度,更新的技术使用等等,但很多节点只是个人维护,出现不兼容,运行不了是很正常的事情。
DZ-FaceDetailer是一个处理脸部的节点,作者在24年6月之后就不再更新了,升级Comfyui后运行工作流出现各种问题,这次是因为使用comfyui_facetools这个节点出了问题,发现comfyui_facetools处理脸部节点用的是DZ-FaceDetailer节点,解决依赖后问题还再,决定修改代码。
把运行工作流出现的错误日志发给通义,它会告诉你问题所在,以及解决的几种办法,然后一段段处理代码。最后修改节点下的DZFaceDetailer.py文件


import torch
from mediapipe import solutions
import cv2
import numpy as np
from PIL import Image, ImageFilter
from ultralytics import YOLO
from ultralytics.nn.tasks import DetectionModel  # 确保导入DetectionModel
from ultralytics.nn.modules.conv import Conv, Concat  # 导入Conv和Concat模块
from ultralytics.nn.modules.block import C2f, Bottleneck, SPPF, DFL  # 导入C2f, Bottleneck, SPPF和DFL模块
from ultralytics.nn.modules.head import Detect  # 导入Detect模块
from torch.nn import Sequential, Conv2d, BatchNorm2d, SiLU, ModuleList, MaxPool2d, Upsample  # 直接从torch.nn导入Sequential, Conv2d, BatchNorm2d, SiLU, ModuleList, MaxPool2d和Upsample
import os
import comfy
import nodes
from folder_paths import base_path

# 在加载模型前,明确添加DetectionModel到安全全局变量中
torch.serialization.add_safe_globals([DetectionModel, Sequential, Conv, Conv2d, BatchNorm2d, SiLU, C2f, Bottleneck, SPPF, DFL, ModuleList, MaxPool2d, Upsample, Concat, Detect])

face_model_path = os.path.join(base_path, "models/dz_facedetailer/yolo/face_yolov8n.pt")
MASK_CONTROL = ["dilate", "erode", "disabled"]
MASK_TYPE = ["face", "box"]

# 定义一个全局变量来存储模型实例,避免重复加载
global_face_model = None

class FaceDetailer:
    @classmethod
    def INPUT_TYPES(s):
        return {"required":
                {
                    "model": ("MODEL",),
                    "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
                    "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
                    "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}),
                    "sampler_name": (comfy.samplers.KSampler.SAMPLERS, ),
                    "scheduler": (comfy.samplers.KSampler.SCHEDULERS, ),
                    "positive": ("CONDITIONING", ),
                    "negative": ("CONDITIONING", ),
                    "latent_image": ("LATENT", ),
                    "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
                    "latent_image": ("LATENT", ),
                    "vae": ("VAE",),
                    "mask_blur": ("INT", {"default": 0, "min": 0, "max": 100}),
                    "mask_type": (MASK_TYPE, ),
                    "mask_control": (MASK_CONTROL, ),
                    "dilate_mask_value": ("INT", {"default": 3, "min": 0, "max": 100}),
                    "erode_mask_value": ("INT", {"default": 3, "min": 0, "max": 100}),
                }
                }

    RETURN_TYPES = ("LATENT", "MASK",)

    FUNCTION = "detailer"

    CATEGORY = "face_detailer"

    def detailer(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise, vae, mask_blur, mask_type, mask_control, dilate_mask_value, erode_mask_value):

        global global_face_model

        # 如果模型尚未加载,则加载之
        if global_face_model is None:
            try:
                with torch.serialization.safe_globals():
                    global_face_model = YOLO(face_model_path)  # 尝试加载YOLO模型
            except Exception as e:
                print(f"Failed to load with safe globals: {e}")
                
                # 如果所有方法都失败了,并且你完全信任该模型文件的来源,可以尝试以下方法:
                # 注意:这可能会带来安全风险,请谨慎使用
                global_face_model = torch.load(face_model_path, map_location='cpu', weights_only=False)
                
        # input latent decoded to tensor image for processing
        tensor_img = vae.decode(latent_image["samples"])

        batch_size = tensor_img.shape[0]

        mask = Detection().detect_faces(tensor_img, batch_size, mask_type, mask_control, mask_blur, dilate_mask_value, erode_mask_value)

        latent_mask = set_mask(latent_image, mask)

        latent = nodes.common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_mask, denoise=denoise)

        return (latent[0], latent[0]["noise_mask"],)

class Detection:
    def __init__(self):
        pass

    def detect_faces(self, tensor_img, batch_size, mask_type, mask_control, mask_blur, mask_dilate, mask_erode):
        mask_imgs = []
        for i in range(0, batch_size):
            # print(input_tensor_img[i, :,:,:].shape)
            # convert input latent to numpy array for yolo model
            img = image2nparray(tensor_img[i], False)
            # Process the face mesh or make the face box for masking
            if mask_type == "box":
                final_mask = facebox_mask(img)
            else:
                final_mask = facemesh_mask(img)

            final_mask = self.mask_control(final_mask, mask_control, mask_blur, mask_dilate, mask_erode)

            final_mask = np.array(Image.fromarray(final_mask).getchannel('A')).astype(np.float32) / 255.0
            # Convert mask to tensor and assign the mask to the input tensor
            final_mask = torch.from_numpy(final_mask)

            mask_imgs.append(final_mask)

        final_mask = torch.stack(mask_imgs)

        return final_mask

    def mask_control(self, numpy_img, mask_control, mask_blur, mask_dilate, mask_erode):
        numpy_image = numpy_img.copy();
        # Erode/Dilate mask
        if mask_control == "dilate":
            if mask_dilate > 0:
                numpy_image = self.dilate_mask(numpy_image, mask_dilate)
        elif mask_control == "erode":
            if mask_erode > 0:
                numpy_image = self.erode_mask(numpy_image, mask_erode)
        if mask_blur > 0:
            final_mask_image = Image.fromarray(numpy_image)
            blurred_mask_image = final_mask_image.filter(
                ImageFilter.GaussianBlur(radius=mask_blur))
            numpy_image = np.array(blurred_mask_image)

        return numpy_image

    def erode_mask(self, mask, dilate):
        # I use erode function because the mask is inverted
        # later I will fix it
        kernel = np.ones((int(dilate), int(dilate)), np.uint8)
        dilated_mask = cv2.dilate(mask, kernel, iterations=1)
        return dilated_mask

    def dilate_mask(self, mask, erode):
        # I use dilate function because the mask is inverted like the other function
        # later I will fix it
        kernel = np.ones((int(erode), int(erode)), np.uint8)
        eroded_mask = cv2.erode(mask, kernel, iterations=1)
        return eroded_mask

def facebox_mask(image):
    # Create an empty image with alpha
    mask = np.zeros((image.shape[0], image.shape[1], 4), dtype=np.uint8)

    # setup yolov8n face detection model
    face_model = YOLO(face_model_path)
    face_bbox = face_model(image)
    boxes = face_bbox[0].boxes
    # box = boxes[0].xyxy
    for box in boxes.xyxy:
        x_min, y_min, x_max, y_max = box.tolist()
        # Calculate the center of the bounding box
        center_x = (x_min + x_max) / 2
        center_y = (y_min + y_max) / 2

        # Calcule the maximum width and height
        width = x_max - x_min
        height = y_max - y_min
        max_size = max(width, height)

        # Get the new WxH for a ratio of 1:1
        new_width = max_size
        new_height = max_size

        # Calculate the new coordinates
        new_x_min = int(center_x - new_width / 2)
        new_y_min = int(center_y - new_height / 2)
        new_x_max = int(center_x + new_width / 2)
        new_y_max = int(center_y + new_height / 2)

        # print((new_x_min, new_y_min), (new_x_max, new_y_max))
        # set the square in the face location
        cv2.rectangle(mask, (new_x_min, new_y_min), (new_x_max, new_y_max), (0, 0, 0, 255), -1)

    # mask[:, :, 3] = ~mask[:, :, 3]  # invert the mask

    return mask


def facemesh_mask(image):

    faces_mask = []

    # Empty image with the same shape as input
    mask = np.zeros(
        (image.shape[0], image.shape[1], 4), dtype=np.uint8)
    
    # setup yolov8n face detection model
    face_model = YOLO(face_model_path)
    face_bbox = face_model(image)
    boxes = face_bbox[0].boxes
    # box = boxes[0].xyxy
    for box in boxes.xyxy:
        x_min, y_min, x_max, y_max = box.tolist()
        # Calculate the center of the bounding box
        center_x = (x_min + x_max) / 2
        center_y = (y_min + y_max) / 2

        # Calcule the maximum width and height
        width = x_max - x_min
        height = y_max - y_min
        max_size = max(width, height)

        # Get the new WxH for a ratio of 1:1
        new_width = max_size
        new_height = max_size

        # Calculate the new coordinates
        new_x_min = int(center_x - new_width / 2)
        new_y_min = int(center_y - new_height / 2)
        new_x_max = int(center_x + new_width / 2)
        new_y_max = int(center_y + new_height / 2)

        # print((new_x_min, new_y_min), (new_x_max, new_y_max))
        # set the square in the face location
        face = image[new_y_min:new_y_max, new_x_min:new_x_max, :]

        mp_face_mesh = solutions.face_mesh
        face_mesh = mp_face_mesh.FaceMesh(static_image_mode=True, max_num_faces=1, min_detection_confidence=0.5)
        results = face_mesh.process(cv2.cvtColor(face, cv2.COLOR_BGR2RGB))
        if results.multi_face_landmarks:
            for face_landmarks in results.multi_face_landmarks:
                # List of detected face points
                points = []
                for landmark in face_landmarks.landmark:
                    cx, cy = int(
                        landmark.x * face.shape[1]), int(landmark.y * face.shape[0])
                    points.append([cx, cy])

                face_mask = np.zeros((face.shape[0], face.shape[1], 4), dtype=np.uint8)

                # Obtain the countour of the face
                convex_hull = cv2.convexHull(np.array(points))
                
                # Fill the contour and store it in alpha for the mask
                cv2.fillConvexPoly(face_mask, convex_hull, (0, 0, 0, 255))

                faces_mask.append([face_mask, [new_x_min, new_x_max, new_y_min, new_y_max]])
            
    for face_mask in faces_mask:
        paste_numpy_images(mask, face_mask[0], face_mask[1][0], face_mask[1][1], face_mask[1][2], face_mask[1][3])

    # print(f"{len(faces_mask)} faces detected")
    # mask[:, :, 3] = ~mask[:, :, 3]
    return mask


def paste_numpy_images(target_image, source_image, x_min, x_max, y_min, y_max):
    # Paste the source image into the target image at the specified coordinates
    target_image[y_min:y_max, x_min:x_max, :] = source_image

    return target_image



def image2nparray(image, BGR):
    """
    convert tensor image to numpy array

    Args:
        image (Tensor): Tensor image

    Returns:
        returns: Numpy array.

    """
    narray = np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)

    if BGR:
        return narray
    else:
        return narray[:, :, ::-1]


def set_mask(samples, mask):
    s = samples.copy()
    s["noise_mask"] = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1]))
    return s

暂时解决了问题。

发表回复

您的邮箱地址不会被公开。 必填项已用 * 标注