294 lines
13 KiB
Python
294 lines
13 KiB
Python
# ComfyUI-RMBG
|
|
# This custom node for ComfyUI provides functionality for face parsing using Segformer model.
|
|
#
|
|
# This integration script follows GPL-3.0 License.
|
|
# When using or modifying this code, please respect both the original model licenses
|
|
# and this integration's license terms.
|
|
#
|
|
# Source: https://github.com/AILab-AI/ComfyUI-RMBG
|
|
|
|
import os
|
|
import torch
|
|
import torch.nn as nn
|
|
import numpy as np
|
|
from typing import Tuple, Union
|
|
from PIL import Image, ImageFilter
|
|
from transformers import SegformerImageProcessor, AutoModelForSemanticSegmentation
|
|
import folder_paths
|
|
from huggingface_hub import hf_hub_download
|
|
import shutil
|
|
from torchvision import transforms
|
|
|
|
def pil2tensor(image: Image.Image) -> torch.Tensor:
|
|
return torch.from_numpy(np.array(image).astype(np.float32) / 255.0)[None,]
|
|
|
|
def tensor2pil(image: torch.Tensor) -> Image.Image:
|
|
return Image.fromarray(np.clip(255. * image.cpu().numpy(), 0, 255).astype(np.uint8))
|
|
|
|
def image2mask(image: Image.Image) -> torch.Tensor:
|
|
if isinstance(image, Image.Image):
|
|
image = pil2tensor(image)
|
|
return image.squeeze()[..., 0]
|
|
|
|
def mask2image(mask: torch.Tensor) -> Image.Image:
|
|
if len(mask.shape) == 2:
|
|
mask = mask.unsqueeze(0)
|
|
return tensor2pil(mask)
|
|
|
|
def RGB2RGBA(image: Image.Image, mask: Union[Image.Image, torch.Tensor]) -> Image.Image:
|
|
if isinstance(mask, torch.Tensor):
|
|
mask = mask2image(mask)
|
|
if mask.size != image.size:
|
|
mask = mask.resize(image.size, Image.Resampling.LANCZOS)
|
|
return Image.merge('RGBA', (*image.convert('RGB').split(), mask.convert('L')))
|
|
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
|
|
folder_paths.add_model_folder_path("rmbg", os.path.join(folder_paths.models_dir, "RMBG"))
|
|
|
|
AVAILABLE_MODELS = {
|
|
"face_parsing": "1038lab/segformer_face"
|
|
}
|
|
|
|
class FaceSegment:
|
|
def __init__(self):
|
|
self.processor = None
|
|
self.model = None
|
|
self.cache_dir = os.path.join(folder_paths.models_dir, "RMBG", "segformer_face")
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
available_classes = [
|
|
"Skin", "Nose", "Eyeglasses", "Left-eye", "Right-eye",
|
|
"Left-eyebrow", "Right-eyebrow", "Left-ear", "Right-ear", "Mouth",
|
|
"Upper-lip", "Lower-lip", "Hair", "Earring", "Neck",
|
|
]
|
|
tooltips = {
|
|
"process_res": "Processing resolution (higher = more VRAM)",
|
|
"mask_blur": "Blur amount for mask edges",
|
|
"mask_offset": "Expand/Shrink mask boundary",
|
|
"invert_output": "Invert both image and mask output",
|
|
"background": "Choose background type: Alpha (transparent) or Color (custom background color).",
|
|
"background_color": "Choose background color (Alpha = transparent)"
|
|
}
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
},
|
|
"optional": {
|
|
**{cls_name: ("BOOLEAN", {"default": False})
|
|
for cls_name in available_classes},
|
|
"process_res": ("INT", {"default": 512, "min": 128, "max": 2048, "step": 32, "tooltip": tooltips["process_res"]}),
|
|
"mask_blur": ("INT", {"default": 0, "min": 0, "max": 64, "step": 1, "tooltip": tooltips["mask_blur"]}),
|
|
"mask_offset": ("INT", {"default": 0, "min": -64, "max": 64, "step": 1, "tooltip": tooltips["mask_offset"]}),
|
|
"invert_output": ("BOOLEAN", {"default": False, "tooltip": tooltips["invert_output"]}),
|
|
"background": (["Alpha", "Color"], {"default": "Alpha", "tooltip": tooltips["background"]}),
|
|
"background_color": ("COLOR", {"default": "#222222", "tooltip": tooltips["background_color"]}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK", "IMAGE")
|
|
RETURN_NAMES = ("IMAGE", "MASK", "MASK_IMAGE")
|
|
FUNCTION = "segment_face"
|
|
CATEGORY = "🧪AILab/🧽RMBG"
|
|
|
|
def check_model_cache(self):
|
|
if not os.path.exists(self.cache_dir):
|
|
return False, "Model directory not found"
|
|
|
|
required_files = [
|
|
'config.json',
|
|
'model.safetensors',
|
|
'preprocessor_config.json'
|
|
]
|
|
|
|
missing_files = [f for f in required_files if not os.path.exists(os.path.join(self.cache_dir, f))]
|
|
if missing_files:
|
|
return False, f"Required model files missing: {', '.join(missing_files)}"
|
|
return True, "Model cache verified"
|
|
|
|
def clear_model(self):
|
|
if self.model is not None:
|
|
self.model.cpu()
|
|
del self.model
|
|
self.model = None
|
|
self.processor = None
|
|
torch.cuda.empty_cache()
|
|
|
|
def download_model_files(self):
|
|
model_id = AVAILABLE_MODELS["face_parsing"]
|
|
model_files = {
|
|
'config.json': 'config.json',
|
|
'model.safetensors': 'model.safetensors',
|
|
'preprocessor_config.json': 'preprocessor_config.json'
|
|
}
|
|
|
|
os.makedirs(self.cache_dir, exist_ok=True)
|
|
print(f"Downloading face parsing model files...")
|
|
|
|
try:
|
|
for save_name, repo_path in model_files.items():
|
|
print(f"Downloading {save_name}...")
|
|
downloaded_path = hf_hub_download(
|
|
repo_id=model_id,
|
|
filename=repo_path,
|
|
local_dir=self.cache_dir,
|
|
local_dir_use_symlinks=False
|
|
)
|
|
|
|
if os.path.dirname(downloaded_path) != self.cache_dir:
|
|
target_path = os.path.join(self.cache_dir, save_name)
|
|
shutil.move(downloaded_path, target_path)
|
|
return True, "Model files downloaded successfully"
|
|
except Exception as e:
|
|
return False, f"Error downloading model files: {str(e)}"
|
|
|
|
def segment_face(self, images, process_res=512, mask_blur=0, mask_offset=0, background="Alpha", background_color="#222222", invert_output=False, **class_selections):
|
|
try:
|
|
# Check and download model if needed
|
|
cache_status, message = self.check_model_cache()
|
|
if not cache_status:
|
|
print(f"Cache check: {message}")
|
|
download_status, download_message = self.download_model_files()
|
|
if not download_status:
|
|
raise RuntimeError(download_message)
|
|
|
|
# Load model if needed
|
|
if self.processor is None:
|
|
self.processor = SegformerImageProcessor.from_pretrained(self.cache_dir)
|
|
self.model = AutoModelForSemanticSegmentation.from_pretrained(self.cache_dir)
|
|
self.model.eval()
|
|
for param in self.model.parameters():
|
|
param.requires_grad = False
|
|
self.model.to(device)
|
|
|
|
# Class mapping for segmentation
|
|
class_map = {
|
|
"Background": 0, "Skin": 1, "Nose": 2, "Eyeglasses": 3,
|
|
"Left-eye": 4, "Right-eye": 5, "Left-eyebrow": 6, "Right-eyebrow": 7,
|
|
"Left-ear": 8, "Right-ear": 9, "Mouth": 10, "Upper-lip": 11,
|
|
"Lower-lip": 12, "Hair": 13, "Hat": 14, "Earring": 15,
|
|
"Necklace": 16, "Neck": 17, "Clothing": 18
|
|
}
|
|
|
|
# Get selected classes
|
|
selected_classes = [name for name, selected in class_selections.items() if selected]
|
|
if not selected_classes:
|
|
selected_classes = ["Skin", "Nose", "Left-eye", "Right-eye", "Mouth"]
|
|
|
|
# Validate selected classes
|
|
invalid_classes = [cls for cls in selected_classes if cls not in class_map]
|
|
if invalid_classes:
|
|
raise ValueError(f"Invalid class selections: {', '.join(invalid_classes)}. Valid classes are: {', '.join(class_map.keys())}")
|
|
|
|
# Image preprocessing
|
|
transform_image = transforms.Compose([
|
|
transforms.Resize((process_res, process_res)),
|
|
transforms.ToTensor(),
|
|
])
|
|
|
|
batch_tensor = []
|
|
batch_masks = []
|
|
|
|
for image in images:
|
|
orig_image = tensor2pil(image)
|
|
w, h = orig_image.size
|
|
|
|
input_tensor = transform_image(orig_image)
|
|
|
|
if input_tensor.shape[0] == 4:
|
|
input_tensor = input_tensor[:3]
|
|
|
|
input_tensor = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])(input_tensor)
|
|
|
|
input_tensor = input_tensor.unsqueeze(0).to(device)
|
|
|
|
with torch.no_grad():
|
|
outputs = self.model(input_tensor)
|
|
logits = outputs.logits.cpu()
|
|
upsampled_logits = nn.functional.interpolate(
|
|
logits,
|
|
size=(h, w),
|
|
mode="bilinear",
|
|
align_corners=False,
|
|
)
|
|
pred_seg = upsampled_logits.argmax(dim=1)[0]
|
|
|
|
# Combine selected class masks
|
|
combined_mask = None
|
|
for class_name in selected_classes:
|
|
mask = (pred_seg == class_map[class_name]).float()
|
|
if combined_mask is None:
|
|
combined_mask = mask
|
|
else:
|
|
combined_mask = torch.clamp(combined_mask + mask, 0, 1)
|
|
|
|
# Convert mask to PIL for processing
|
|
mask_image = Image.fromarray((combined_mask.numpy() * 255).astype(np.uint8))
|
|
|
|
if mask_blur > 0:
|
|
mask_image = mask_image.filter(ImageFilter.GaussianBlur(radius=mask_blur))
|
|
|
|
if mask_offset != 0:
|
|
if mask_offset > 0:
|
|
mask_image = mask_image.filter(ImageFilter.MaxFilter(size=mask_offset * 2 + 1))
|
|
else:
|
|
mask_image = mask_image.filter(ImageFilter.MinFilter(size=-mask_offset * 2 + 1))
|
|
|
|
if invert_output:
|
|
mask_image = Image.fromarray(255 - np.array(mask_image))
|
|
|
|
# Handle background color
|
|
if background == "Alpha":
|
|
rgba_image = RGB2RGBA(orig_image, mask_image)
|
|
result_image = pil2tensor(rgba_image)
|
|
else:
|
|
def hex_to_rgba(hex_color):
|
|
hex_color = hex_color.lstrip('#')
|
|
if len(hex_color) == 6:
|
|
r, g, b = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16)
|
|
a = 255
|
|
elif len(hex_color) == 8:
|
|
r, g, b, a = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16), int(hex_color[6:8], 16)
|
|
else:
|
|
raise ValueError("Invalid color format")
|
|
return (r, g, b, a)
|
|
rgba_image = RGB2RGBA(orig_image, mask_image)
|
|
background_color = params.get("background_color", "#222222")
|
|
rgba = hex_to_rgba(background_color)
|
|
bg_image = Image.new('RGBA', orig_image.size, rgba)
|
|
composite_image = Image.alpha_composite(bg_image, rgba_image)
|
|
result_image = pil2tensor(composite_image.convert('RGB'))
|
|
|
|
batch_tensor.append(result_image)
|
|
batch_masks.append(pil2tensor(mask_image))
|
|
|
|
# Create mask image for visualization
|
|
mask_images = []
|
|
for mask_tensor in batch_masks:
|
|
# Convert mask to RGB image format for visualization
|
|
mask_image = mask_tensor.reshape((-1, 1, mask_tensor.shape[-2], mask_tensor.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)
|
|
mask_images.append(mask_image)
|
|
|
|
mask_image_output = torch.cat(mask_images, dim=0)
|
|
|
|
# Prepare final output
|
|
batch_tensor = torch.cat(batch_tensor, dim=0)
|
|
batch_masks = torch.cat(batch_masks, dim=0)
|
|
|
|
return (batch_tensor, batch_masks, mask_image_output)
|
|
|
|
except Exception as e:
|
|
self.clear_model()
|
|
raise RuntimeError(f"Error in Face Parsing processing: {str(e)}")
|
|
finally:
|
|
if self.model is not None and not self.model.training:
|
|
self.clear_model()
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"FaceSegment": FaceSegment
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"FaceSegment": "Face Segment (RMBG)"
|
|
} |