diff --git a/AILab_BiRefNet.py b/AILab_BiRefNet.py index a299a08..6465fbe 100644 --- a/AILab_BiRefNet.py +++ b/AILab_BiRefNet.py @@ -350,8 +350,8 @@ class BiRefNetRMBG: } } - RETURN_TYPES = ("IMAGE", "MASK") - RETURN_NAMES = ("IMAGE", "MASK") + RETURN_TYPES = ("IMAGE", "MASK", "IMAGE") + RETURN_NAMES = ("IMAGE", "MASK", "MASK_IMAGE") FUNCTION = "process_image" CATEGORY = "🧪AILab/🧽RMBG" @@ -447,7 +447,16 @@ class BiRefNetRMBG: processed_masks.append(pil2tensor(mask)) - return (torch.cat(processed_images, dim=0), torch.cat(processed_masks, dim=0)) + # Create mask image for visualization + mask_images = [] + for mask_tensor in processed_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) + + return (torch.cat(processed_images, dim=0), torch.cat(processed_masks, dim=0), mask_image_output) except Exception as e: handle_model_error(f"Error in image processing: {str(e)}") diff --git a/AILab_BodySegment.py b/AILab_BodySegment.py index 5ae34db..29dedb1 100644 --- a/AILab_BodySegment.py +++ b/AILab_BodySegment.py @@ -85,8 +85,8 @@ class BodySegment: }, } - RETURN_TYPES = ("IMAGE", "MASK") - RETURN_NAMES = ("IMAGE", "MASK") + RETURN_TYPES = ("IMAGE", "MASK", "IMAGE") + RETURN_NAMES = ("IMAGE", "MASK", "MASK_IMAGE") FUNCTION = "segment_body" CATEGORY = "🧪AILab/🧽RMBG" @@ -217,11 +217,20 @@ class BodySegment: 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) + return (batch_tensor, batch_masks, mask_image_output) except Exception as e: self.clear_model() diff --git a/AILab_ClothSegment.py b/AILab_ClothSegment.py index 93a4b8f..d311f37 100644 --- a/AILab_ClothSegment.py +++ b/AILab_ClothSegment.py @@ -85,8 +85,8 @@ class ClothesSegment: }, } - RETURN_TYPES = ("IMAGE", "MASK") - RETURN_NAMES = ("IMAGE", "MASK") + RETURN_TYPES = ("IMAGE", "MASK", "IMAGE") + RETURN_NAMES = ("IMAGE", "MASK", "MASK_IMAGE") FUNCTION = "segment_clothes" CATEGORY = "🧪AILab/🧽RMBG" @@ -253,11 +253,20 @@ class ClothesSegment: 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) + return (batch_tensor, batch_masks, mask_image_output) except Exception as e: self.clear_model() diff --git a/AILab_FaceSegment.py b/AILab_FaceSegment.py index 2acf66f..acd15d5 100644 --- a/AILab_FaceSegment.py +++ b/AILab_FaceSegment.py @@ -1,281 +1,295 @@ -# 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 = [ - # "Background", # Not a facial feature - "Skin", "Nose", "Eyeglasses", "Left-eye", "Right-eye", - "Left-eyebrow", "Right-eyebrow", "Left-ear", "Right-ear", "Mouth", - "Upper-lip", "Lower-lip", "Hair", "Earring", "Neck", - # "Hat", # Not a facial feature - # "Necklace", # Not a facial feature - # "Clothing" # Not a facial feature - ] - - tooltips = { - "process_res": "Processing resolution (higher = more VRAM)", - "mask_blur": "Blur amount for mask edges", - "mask_offset": "Expand/Shrink mask boundary", - "background_color": "Choose background color (Alpha = transparent)", - "invert_output": "Invert both image and mask output", - } - - 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"]}), - "background_color": (["Alpha", "black", "white", "gray", "green", "blue", "red"], {"default": "Alpha", "tooltip": tooltips["background_color"]}), - "invert_output": ("BOOLEAN", {"default": False, "tooltip": tooltips["invert_output"]}), - }, - } - - RETURN_TYPES = ("IMAGE", "MASK") - RETURN_NAMES = ("IMAGE", "MASK") - 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_color="Alpha", 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", "Eyes", "Mouth"] - - # 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_color == "Alpha": - rgba_image = RGB2RGBA(orig_image, mask_image) - result_image = pil2tensor(rgba_image) - else: - bg_colors = { - "black": (0, 0, 0), - "white": (255, 255, 255), - "gray": (128, 128, 128), - "green": (0, 255, 0), - "blue": (0, 0, 255), - "red": (255, 0, 0) - } - - rgba_image = RGB2RGBA(orig_image, mask_image) - bg_image = Image.new('RGBA', orig_image.size, (*bg_colors[background_color], 255)) - 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)) - - # Prepare final output - batch_tensor = torch.cat(batch_tensor, dim=0) - batch_masks = torch.cat(batch_masks, dim=0) - - return (batch_tensor, batch_masks) - - except Exception as e: - self.clear_model() - raise RuntimeError(f"Error in Face Parsing processing: {str(e)}") - finally: - if not self.model.training: - self.clear_model() - -NODE_CLASS_MAPPINGS = { - "FaceSegment": FaceSegment -} - -NODE_DISPLAY_NAME_MAPPINGS = { - "FaceSegment": "Face Segment (RMBG)" +# 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 = [ + # "Background", # Not a facial feature + "Skin", "Nose", "Eyeglasses", "Left-eye", "Right-eye", + "Left-eyebrow", "Right-eyebrow", "Left-ear", "Right-ear", "Mouth", + "Upper-lip", "Lower-lip", "Hair", "Earring", "Neck", + # "Hat", # Not a facial feature + # "Necklace", # Not a facial feature + # "Clothing" # Not a facial feature + ] + + tooltips = { + "process_res": "Processing resolution (higher = more VRAM)", + "mask_blur": "Blur amount for mask edges", + "mask_offset": "Expand/Shrink mask boundary", + "background_color": "Choose background color (Alpha = transparent)", + "invert_output": "Invert both image and mask output", + } + + 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"]}), + "background_color": (["Alpha", "black", "white", "gray", "green", "blue", "red"], {"default": "Alpha", "tooltip": tooltips["background_color"]}), + "invert_output": ("BOOLEAN", {"default": False, "tooltip": tooltips["invert_output"]}), + }, + } + + 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_color="Alpha", 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_color == "Alpha": + rgba_image = RGB2RGBA(orig_image, mask_image) + result_image = pil2tensor(rgba_image) + else: + bg_colors = { + "black": (0, 0, 0), + "white": (255, 255, 255), + "gray": (128, 128, 128), + "green": (0, 255, 0), + "blue": (0, 0, 255), + "red": (255, 0, 0) + } + + rgba_image = RGB2RGBA(orig_image, mask_image) + bg_image = Image.new('RGBA', orig_image.size, (*bg_colors[background_color], 255)) + 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 diff --git a/AILab_FashionSegment.py b/AILab_FashionSegment.py index 6125751..d37f524 100644 --- a/AILab_FashionSegment.py +++ b/AILab_FashionSegment.py @@ -156,9 +156,9 @@ class FashionSegmentClothing: return { "required": { "images": ("IMAGE",), - "accessories_options": ("ACCESSORIES_OPTIONS",), }, "optional": { + "accessories_options": ("ACCESSORIES_OPTIONS",), **{cls_name: ("BOOLEAN", {"default": False}) for cls_name in clothing_classes}, "process_res": ("INT", {"default": 512, "min": 128, "max": 2048, "step": 32}), @@ -169,8 +169,8 @@ class FashionSegmentClothing: }, } - RETURN_TYPES = ("IMAGE", "MASK") - RETURN_NAMES = ("images", "mask") + RETURN_TYPES = ("IMAGE", "MASK", "IMAGE") + RETURN_NAMES = ("IMAGE", "MASK", "MASK_IMAGE") FUNCTION = "segment_fashion" CATEGORY = "🧪AILab/🧽RMBG" @@ -328,10 +328,19 @@ class FashionSegmentClothing: 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) + batch_tensor = torch.cat(batch_tensor, dim=0) batch_masks = torch.cat(batch_masks, dim=0) - return (batch_tensor, batch_masks) + return (batch_tensor, batch_masks, mask_image_output) except Exception as e: self.clear_model() diff --git a/AILab_RMBG.py b/AILab_RMBG.py index be84348..121dd08 100644 --- a/AILab_RMBG.py +++ b/AILab_RMBG.py @@ -129,8 +129,7 @@ class BaseModelLoader: hf_hub_download( repo_id=model_info["repo_id"], filename=filename, - local_dir=cache_dir, - local_dir_use_symlinks=False + local_dir=cache_dir ) return True, "Model files downloaded successfully" @@ -229,7 +228,7 @@ class RMBGModel(BaseModelLoader): raise RuntimeError("Could not find suitable model class") except Exception as custom_e: - handle_model_error(f"Failed to load model in compatibility mode: {str(custom_e)}\nConsider downgrading transformers to version 4.48.3: pip install transformers==4.48.3") + handle_model_error(f"Failed to load model in compatibility mode: {str(custom_e)}") else: raise ae except Exception as e: