From 9d4881316473fa5a5d6d5e04db4a2f3dc2a3fcae Mon Sep 17 00:00:00 2001 From: AI Lab <129358391+1038lab@users.noreply.github.com> Date: Thu, 1 Jan 2026 19:26:59 -0800 Subject: [PATCH] Delete AILab_FaceSegment.py --- AILab_FaceSegment.py | 293 ------------------------------------------- 1 file changed, 293 deletions(-) delete mode 100644 AILab_FaceSegment.py diff --git a/AILab_FaceSegment.py b/AILab_FaceSegment.py deleted file mode 100644 index bebd1d2..0000000 --- a/AILab_FaceSegment.py +++ /dev/null @@ -1,293 +0,0 @@ -# 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": ("COLORCODE", {"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) - 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)" -} \ No newline at end of file