nodes/WeChatCV-Stand-In_Preprocessor_ComfyUI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

add face swap node

Browse Source
This commit is contained in:
KBRASK
2025-08-16 20:37:22 +08:00
parent db0d3d966d
commit e874ec8fcc
7 changed files with 2209 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+1
View File
@@ -0,0 +1 @@
__pycache__/
__init__.py
+3
View File
@@ -16,13 +16,16 @@ NODE_DISPLAY_NAME_MAPPINGS = {}
try:
from .nodes import FaceProcessorLoader, ApplyFaceProcessor
from .videoinput import VideoInputPreprocessor
# Map the class names to the classes themselves
NODE_CLASS_MAPPINGS["FaceProcessorLoader"] = FaceProcessorLoader
NODE_CLASS_MAPPINGS["ApplyFaceProcessor"] = ApplyFaceProcessor
NODE_CLASS_MAPPINGS["VideoInputPreprocessor"] = VideoInputPreprocessor
NODE_DISPLAY_NAME_MAPPINGS["FaceProcessorLoader"] = "Stand-In Processor Loader"
NODE_DISPLAY_NAME_MAPPINGS["ApplyFaceProcessor"] = "Apply Stand-In Processor"
NODE_DISPLAY_NAME_MAPPINGS["VideoInputPreprocessor"] = "Stand-In VideoInputPreprocessor"
logger.info("Successfully loaded all Stand-In nodes.")
example/face_swap.json
+2040
View File
File diff suppressed because it is too large Load Diff
wanvideo_Stand-In_reference_example.json → example/wanvideo_Stand-In_reference_example.json
View File
nodes.py
+27 -12
View File
@@ -11,10 +11,10 @@ from torchvision.transforms.functional import normalize
from typing import Union, Optional
import folder_paths
# --- GLOBALS ---
LOADED_PROCESSORS = {}
# --- DIRECTORY SETUP ---
def set_extra_config_model_path(extra_config_models_dir_key, models_dir_name: str):
"""Helper function to set up model directories within ComfyUI."""
models_dir_default = os.path.join(folder_paths.models_dir, models_dir_name)
@@ -29,10 +29,11 @@ def set_extra_config_model_path(extra_config_models_dir_key, models_dir_name: st
else:
folder_paths.add_model_folder_path(extra_config_models_dir_key, models_dir_default, is_default=True)
# Set up paths for YOLO and face parsing models
set_extra_config_model_path("yolo", "yolo")
set_extra_config_model_path("face_parsing", "face_parsing")
# --- MODEL DOWNLOADING ---
def download_yolo_model(model_name="model.pt"):
"""
Checks if a YOLO model exists and downloads it from a public source if not.
@@ -70,7 +71,7 @@ def download_yolo_model(model_name="model.pt"):
raise e
return model_path
# --- IMAGE UTILS ---
def _img2tensor(img: np.ndarray, bgr2rgb: bool = True) -> torch.Tensor:
if bgr2rgb:
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
@@ -91,17 +92,19 @@ def _pad_to_square(img: np.ndarray, pad_color: int = 255) -> np.ndarray:
return padded_img
def tensor_to_cv2_img(tensor: torch.Tensor) -> np.ndarray:
"""Converts a ComfyUI IMAGE tensor (B, H, W, C) to an OpenCV BGR image (H, W, C)."""
img_np = tensor.squeeze(0).cpu().numpy()
img_np = (img_np * 255).astype(np.uint8)
img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
return img_bgr
def cv2_img_to_tensor(img: np.ndarray) -> torch.Tensor:
"""Converts an OpenCV BGR image (H, W, C) to a ComfyUI IMAGE tensor (1, H, W, C)."""
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img_tensor = torch.from_numpy(img_rgb.astype(np.float32) / 255.0).unsqueeze(0)
return img_tensor
# --- NODES ---
class FaceProcessorLoader:
@classmethod
def INPUT_TYPES(s):
@@ -163,12 +166,12 @@ class ApplyFaceProcessor:
"border_thresh": ("INT", {"default": 10, "min": 0, "max": 100, "step": 1}),
"face_crop_scale": ("FLOAT", {"default": 1.5, "min": 1.0, "max": 3.0, "step": 0.1}),
"confidence_threshold": ("FLOAT", {"default": 0.5, "min": 0.1, "max": 1.0, "step": 0.05}),
# --- NEW INPUT ---
"with_neck": ("BOOLEAN", {"default": False, "label_on": "enabled", "label_off": "disabled"}),
}
}
RETURN_TYPES = ("IMAGE",)
RETURN_TYPES = ("IMAGE", "IMAGE")
RETURN_NAMES = ("processed_image", "face_rgba")
FUNCTION = "apply_processing"
CATEGORY = "Stand-In"
@@ -224,19 +227,29 @@ class ApplyFaceProcessor:
if with_neck:
final_mask_np = (parsing_map_np != 0).astype(np.uint8)
else:
parts_to_remove = [0, 14, 16]
parts_to_remove = [0, 14, 16] # Typically background, neck, clothes
final_mask_np = np.isin(parsing_map_np, parts_to_remove, invert=True).astype(np.uint8)
# --- Create the primary output (masked face on white background) ---
white_background = np.ones_like(image_resized, dtype=np.uint8) * 255
mask_3channel = cv2.cvtColor(final_mask_np * 255, cv2.COLOR_GRAY2BGR)
result_img_bgr = np.where(mask_3channel != 0, image_resized, white_background)
result_tensor = cv2_img_to_tensor(result_img_bgr)
return (result_tensor,)
# --- Create the secondary RGBA face output ---
# 1. Get the RGB channels from the resized BGR image
image_resized_rgb = cv2.cvtColor(image_resized, cv2.COLOR_BGR2RGB)
# 2. Create the alpha channel from the final mask
alpha_channel = final_mask_np * 255
# 3. Stack RGB and Alpha channels to create a 4-channel RGBA image
face_rgba_np = np.dstack((image_resized_rgb, alpha_channel))
# 4. Convert the RGBA numpy array to a tensor
face_rgba_tensor = torch.from_numpy(face_rgba_np.astype(np.float32) / 255.0).unsqueeze(0)
return (result_tensor, face_rgba_tensor,)
# --- MAPPINGS ---
NODE_CLASS_MAPPINGS = {
"FaceProcessorLoader": FaceProcessorLoader,
"ApplyFaceProcessor": ApplyFaceProcessor
@@ -245,4 +258,6 @@ NODE_CLASS_MAPPINGS = {
NODE_DISPLAY_NAME_MAPPINGS = {
"FaceProcessorLoader": "Stand-In Processor Loader",
"ApplyFaceProcessor": "Apply Stand-In Processor"
}
}
requirements.txt
+3 -3
View File
@@ -1,4 +1,4 @@
ultralytics
huggingface_hub==0.31.1
facexlib==0.3.0
onnxruntime-gpu==1.22.0
huggingface_hub>=0.31.1
facexlib>=0.3.0
onnxruntime-gpu>=1.22.0
videoinput.py
+135
View File
@@ -0,0 +1,135 @@
import torch
import cv2
import numpy as np
from torchvision.transforms.functional import normalize
from tqdm import tqdm
from PIL import Image
# Helper function to convert a ComfyUI IMAGE tensor to an OpenCV BGR image.
def tensor_to_cv2_img(tensor_frame: torch.Tensor) -> np.ndarray:
"""Converts a single PyTorch image tensor (H, W, C) to a CV2 image (H, W, C) in BGR format."""
img_np = (tensor_frame.cpu().numpy() * 255).astype(np.uint8)
return cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
# Helper function to convert a ComfyUI IMAGE tensor to a PIL Image.
def tensor_to_pil(tensor_frame: torch.Tensor, mode='RGB') -> Image.Image:
"""Converts a single PyTorch image tensor (H, W, C) to a PIL Image."""
return Image.fromarray((tensor_frame.cpu().numpy() * 255).astype(np.uint8), mode)
class VideoInputPreprocessor:
"""
Processes a batch of image frames to generate a face mask, then uses that mask's
bounding box to paste a provided RGBA image onto each frame.
"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"face_processor": ("FACE_PROCESSOR",),
"images": ("IMAGE",),
"face_rgba": ("IMAGE",),
"confidence_threshold": ("FLOAT", {"default": 0.5, "min": 0.1, "max": 1.0, "step": 0.05}),
"face_crop_scale": ("FLOAT", {"default": 1.5, "min": 1.0, "max": 3.0, "step": 0.1}),
"dilation_kernel_size": ("INT", {"default": 10, "min": 0, "max": 50, "step": 1}),
"with_neck": ("BOOLEAN", {"default": True, "label_on": "enabled", "label_off": "disabled"}),
}
}
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("processed_images",)
FUNCTION = "generate_mask_and_paste"
CATEGORY = "Stand-In"
def generate_mask_and_paste(self, face_processor, images: torch.Tensor, face_rgba: torch.Tensor, confidence_threshold: float, face_crop_scale: float, dilation_kernel_size: int, with_neck: bool):
detection_model, parsing_model, device = face_processor
total_frames, h, w = images.shape[0], images.shape[1], images.shape[2]
print(f"Processing {total_frames} frames ({w}x{h}) to paste new face.")
# --- Prepare the face_to_paste image ---
# Ensure it is RGBA and convert to PIL for easy pasting with transparency
if face_rgba.shape[3] != 4:
raise ValueError("The 'face_to_paste' image must be an RGBA image.")
# We only need the first image from the batch
face_to_paste_pil = tensor_to_pil(face_rgba[0], mode='RGBA')
processed_frames_tensors = []
for i in tqdm(range(total_frames), desc="Pasting face onto frames"):
frame_tensor = images[i]
frame_bgr = tensor_to_cv2_img(frame_tensor)
# --- 1. Generate Face Mask (same logic as before) ---
results = detection_model(frame_bgr, verbose=False)
confident_boxes = results[0].boxes.xyxy[results[0].boxes.conf > confidence_threshold]
full_mask_np = np.zeros((h, w), dtype=np.uint8)
if confident_boxes.shape[0] > 0:
areas = (confident_boxes[:, 2] - confident_boxes[:, 0]) * (confident_boxes[:, 3] - confident_boxes[:, 1])
x1, y1, x2, y2 = map(int, confident_boxes[torch.argmax(areas)])
center_x, center_y = (x1 + x2) // 2, (y1 + y2) // 2
side_len = int(max(x2 - x1, y2 - y1) * face_crop_scale)
half_side = side_len // 2
crop_y1, crop_x1 = max(center_y - half_side, 0), max(center_x - half_side, 0)
crop_y2, crop_x2 = min(center_y + half_side, h), min(center_x + half_side, w)
face_crop_bgr = frame_bgr[crop_y1:crop_y2, crop_x1:crop_x2]
if face_crop_bgr.size > 0:
face_resized = cv2.resize(face_crop_bgr, (512, 512), interpolation=cv2.INTER_AREA)
face_rgb = cv2.cvtColor(face_resized, cv2.COLOR_BGR2RGB)
face_tensor_in = torch.from_numpy(face_rgb.astype(np.float32) / 255.0).permute(2, 0, 1).unsqueeze(0).to(device)
with torch.no_grad():
normalized_face = normalize(face_tensor_in, [0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
parsing_map = parsing_model(normalized_face)[0].argmax(dim=1, keepdim=True)
parsing_map_np = parsing_map.squeeze().cpu().numpy().astype(np.uint8)
if with_neck:
final_mask_512 = (parsing_map_np != 0).astype(np.uint8) * 255
else:
parts_to_exclude = [0, 14, 15, 16, 18]
final_mask_512 = np.isin(parsing_map_np, parts_to_exclude, invert=True).astype(np.uint8) * 255
if dilation_kernel_size > 0:
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (dilation_kernel_size, dilation_kernel_size))
final_mask_512 = cv2.dilate(final_mask_512, kernel, iterations=1)
mask_resized_to_crop = cv2.resize(final_mask_512, (face_crop_bgr.shape[1], face_crop_bgr.shape[0]), interpolation=cv2.INTER_NEAREST)
full_mask_np[crop_y1:crop_y2, crop_x1:crop_x2] = mask_resized_to_crop
# --- 2. Paste Image onto Frame ---
# Use a copy to avoid modifying the original data
target_frame_pil = tensor_to_pil(frame_tensor).copy()
# Find the bounding box of the largest contour in the mask
contours, _ = cv2.findContours(full_mask_np, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
if contours:
largest_contour = max(contours, key=cv2.contourArea)
x, y, w_box, h_box = cv2.boundingRect(largest_contour)
# Resize the face to paste to the size of the bounding box
face_to_paste_resized = face_to_paste_pil.resize((w_box, h_box), Image.Resampling.LANCZOS)
# Paste the resized face onto the target frame using its own alpha channel as the mask
target_frame_pil.paste(face_to_paste_resized, (x, y), face_to_paste_resized)
processed_np = np.array(target_frame_pil.convert("RGB")).astype(np.float32) / 255.0
processed_tensor = torch.from_numpy(processed_np)
processed_frames_tensors.append(processed_tensor)
output_image_batch = torch.stack(processed_frames_tensors)
return (output_image_batch,)
NODE_CLASS_MAPPINGS = {
"VideoInputPreprocessor": VideoInputPreprocessor,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"VideoInputPreprocessor": "Stand-In VideoInputPreprocessor",
}