import torch import torch.nn.functional as F import numpy as np import cv2 import folder_paths import comfy.model_management as model_management import kornia from .face_detector import ForbiddenVisionFaceDetector from .mask_processor import ForbiddenVisionMaskProcessor from .face_processor_integrated import ForbiddenVisionFaceProcessorIntegrated from .utils import check_for_interruption, ensure_model_directories, get_ordered_upscaler_model_list, clean_model_name class ForbiddenVisionFaceEditPrep: @classmethod def INPUT_TYPES(s): upscaler_models = get_ordered_upscaler_model_list() default_upscaler = "Fast 4x (Lanczos)" if "Fast 4x (Lanczos)" not in upscaler_models and upscaler_models: default_upscaler = upscaler_models[0] return { "required": { "image": ("IMAGE", {"tooltip": "Input image for face crop and mask."}), "face_selection": ("INT", {"default": 0, "min": 0, "max": 20, "step": 1, "tooltip": "0=All faces, 1=1st face, etc."}), "enable_segmentation": ("BOOLEAN", {"default": True, "tooltip": "Use AI segmentation. If disabled, creates oval masks."}), "detection_confidence": ("FLOAT", {"default": 0.8, "min": 0.1, "max": 1.0, "step": 0.01}), "processing_resolution": ("INT", {"default": 1024, "min": 512, "max": 2048, "step": 64}), "crop_padding": ("FLOAT", {"default": 1.6, "min": 1.0, "max": 3.0, "step": 0.1}), "mask_expansion": ("INT", {"default": 2, "min": 0, "max": 100, "step": 1}), # ── NEW ────────────────────────────────────────────────────────── "sampling_mask_blur_size": ("INT", {"default": 21, "min": 0, "max": 101, "step": 2, "tooltip": "Blur kernel size for the output mask. 0 or 1 = no blur."}), "sampling_mask_blur_strength": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 6.0, "step": 0.1, "tooltip": "Controls blur sigma relative to kernel size."}), # ───────────────────────────────────────────────────────────────── "enable_pre_upscale": ("BOOLEAN", {"default": True}), "upscaler_model": (upscaler_models, {"default": default_upscaler}), "isolate_face": ("BOOLEAN", {"default": False}), }, "optional": { "mask": ("MASK", {"tooltip": "Optional external mask. If provided, detection is skipped."}), } } RETURN_TYPES = ("IMAGE", "MASK", "FACE_INFO") RETURN_NAMES = ("cropped_face", "cropped_mask", "face_info") FUNCTION = "prepare_face" CATEGORY = "Forbidden Vision" def __init__(self): ensure_model_directories() self.face_detector = ForbiddenVisionFaceDetector() self.mask_processor = ForbiddenVisionMaskProcessor() self.upscaler_model = None self.upscaler_model_name = None # ── NEW HELPER ──────────────────────────────────────────────────────────── def _blur_mask(self, mask_tensor, blur_size, blur_strength): """ Apply Gaussian blur to a (B, H, W) mask tensor. Returns the mask unchanged when blur_size <= 1. Mirrors the logic in process_inpaint_mask() from the main node. """ if blur_size <= 1: return mask_tensor # kernel size must be odd if blur_size % 2 == 0: blur_size += 1 device = mask_tensor.device mask_4d = mask_tensor.unsqueeze(1) # (B, 1, H, W) base_sigma = (blur_size - 1) / 8.0 strength_t = torch.tensor(blur_strength - 1.0, device=device, dtype=torch.float32) multiplier = 1.0 + torch.tanh(strength_t) * 2.0 actual_sigma = base_sigma * multiplier.item() blurred = kornia.filters.gaussian_blur2d( mask_4d, (blur_size, blur_size), (actual_sigma, actual_sigma) ) return blurred.squeeze(1) # back to (B, H, W) # ───────────────────────────────────────────────────────────────────────── def load_upscaler_model(self, model_name): clean_model_name_val = clean_model_name(model_name) if clean_model_name_val in [ "Fast 4x (Bicubic AA)", "Fast 4x (Lanczos)", "Fast 2x (Bicubic AA)", "Fast 2x (Lanczos)" ]: self.upscaler_model = None self.upscaler_model_name = clean_model_name_val return True if self.upscaler_model is not None and self.upscaler_model_name == clean_model_name_val: return True try: model_path = folder_paths.get_full_path("upscale_models", clean_model_name_val) if model_path is None: print(f"Upscaler model '{clean_model_name_val}' not found.") return False from basicsr.utils import img2tensor, tensor2img from realesrgan import RealESRGANer device = model_management.get_torch_device() upsampler = RealESRGANer( model_path=model_path, scale=4, model=None, tile=0, tile_pad=10, pre_pad=0, half=not model_management.is_device_mps(device), device=device, ) self.upscaler_model = upsampler self.upscaler_model_name = clean_model_name_val print(f"Loaded upscaler model: {clean_model_name_val}") return True except Exception as e: print(f"Error loading upscaler model '{clean_model_name_val}': {e}") self.upscaler_model = None return False def fast_upscale_bicubic(self, image_np_uint8, scale=4): try: image_tensor = torch.from_numpy(image_np_uint8.astype(np.float32) / 255.0).unsqueeze(0) with torch.no_grad(): upscaled_tensor = F.interpolate( image_tensor.permute(0, 3, 1, 2), scale_factor=scale, mode='bicubic', align_corners=False, antialias=True ).permute(0, 2, 3, 1) upscaled_np = upscaled_tensor.squeeze(0).cpu().numpy() upscaled_np_uint8 = (np.clip(upscaled_np, 0, 1) * 255.0).round().astype(np.uint8) cleaned_np_uint8 = self.clean_interpolation_edges(upscaled_np_uint8) return cleaned_np_uint8 except Exception as e: print(f"Error in fast bicubic upscaling: {e}. Returning original image.") return image_np_uint8 def fast_upscale_lanczos(self, image_np_uint8, scale=4): try: h, w = image_np_uint8.shape[:2] new_h, new_w = h * scale, w * scale upscaled_np_uint8 = cv2.resize( image_np_uint8, (new_w, new_h), interpolation=cv2.INTER_LANCZOS4 ) cleaned_np_uint8 = self.clean_interpolation_edges(upscaled_np_uint8) return cleaned_np_uint8 except Exception as e: print(f"Error in fast Lanczos upscaling: {e}. Returning original image.") return image_np_uint8 def clean_interpolation_edges(self, image_np_uint8): try: image_tensor = torch.from_numpy(image_np_uint8.astype(np.float32) / 255.0) image_tensor = image_tensor.permute(2, 0, 1).unsqueeze(0) h, w = image_tensor.shape[2], image_tensor.shape[3] edge_size = max(8, min(h, w) // 32) mask = torch.ones_like(image_tensor[:, :1, :, :]) mask[:, :, :edge_size, :] = torch.linspace(0, 1, edge_size).view(1, 1, edge_size, 1) mask[:, :, h-edge_size:, :] = torch.linspace(1, 0, edge_size).view(1, 1, edge_size, 1) mask[:, :, :, :edge_size] = torch.minimum( mask[:, :, :, :edge_size], torch.linspace(0, 1, edge_size).view(1, 1, 1, edge_size) ) mask[:, :, :, w-edge_size:] = torch.minimum( mask[:, :, :, w-edge_size:], torch.linspace(1, 0, edge_size).view(1, 1, 1, edge_size) ) kernel_size = max(3, edge_size // 2) if kernel_size % 2 == 0: kernel_size += 1 sigma = kernel_size / 3.0 blurred = kornia.filters.gaussian_blur2d(image_tensor, (kernel_size, kernel_size), (sigma, sigma)) cleaned = image_tensor * mask + blurred * (1 - mask) cleaned_np = cleaned.squeeze(0).permute(1, 2, 0).cpu().numpy() cleaned_np_uint8 = (np.clip(cleaned_np, 0, 1) * 255.0).round().astype(np.uint8) return cleaned_np_uint8 except Exception as e: print(f"Error cleaning interpolation edges: {e}. Returning original image.") return image_np_uint8 def run_upscaler(self, image_np_uint8): if self.upscaler_model_name in ["Fast 4x (Bicubic AA)", "Fast 2x (Bicubic AA)"]: scale = 4 if "4x" in self.upscaler_model_name else 2 return self.fast_upscale_bicubic(image_np_uint8, scale=scale) if self.upscaler_model_name in ["Fast 4x (Lanczos)", "Fast 2x (Lanczos)"]: scale = 4 if "4x" in self.upscaler_model_name else 2 return self.fast_upscale_lanczos(image_np_uint8, scale=scale) if self.upscaler_model is None: return image_np_uint8 try: upsampler = self.upscaler_model output, _ = upsampler.enhance(image_np_uint8, outscale=1) return output except Exception as e: print(f"Error running upscaler '{self.upscaler_model_name}': {e}") return image_np_uint8 def prepare_face(self, image, face_selection, enable_segmentation, detection_confidence, processing_resolution, crop_padding, mask_expansion, sampling_mask_blur_size, sampling_mask_blur_strength, # ← NEW enable_pre_upscale, upscaler_model, isolate_face, mask=None): try: check_for_interruption() if image is None: h, w = 1024, 1024 device = model_management.get_torch_device() empty_face = torch.zeros((1, h, w, 3), dtype=torch.float32, device=device) empty_mask = torch.zeros((1, h, w), dtype=torch.float32, device=device) empty_info = { "original_image": np.zeros((h, w, 3), dtype=np.uint8), "original_image_size": (0, 0), "crop_coords": (0, 0, 0, 0), "face_bbox": (0, 0, 0, 0), "target_size": (h, w), "original_crop_size": (0, 0), "blend_mask": np.zeros((h, w), dtype=np.float32), "detection_angle": 0, } return (empty_face, empty_mask, empty_info) device = image.device h, w = image.shape[1], image.shape[2] if mask is not None: combined_mask = mask else: np_masks = self.face_detector.detect_faces( image_tensor=image, enable_segmentation=enable_segmentation, detection_confidence=detection_confidence, face_selection=face_selection ) if not np_masks: empty_face, empty_mask, empty_info = self.mask_processor.create_empty_outputs( image_tensor=image.cpu(), target_size=(processing_resolution, processing_resolution) ) empty_face = empty_face.to(device) empty_mask = empty_mask.to(device) return (empty_face, empty_mask, empty_info) face_masks = [torch.from_numpy(m).unsqueeze(0) for m in np_masks] combined_mask = torch.zeros((1, h, w), dtype=torch.float32, device=device) for mask_tensor in face_masks: mask_tensor = mask_tensor.to(device) combined_mask = torch.maximum(combined_mask, mask_tensor) target_resolution = (processing_resolution, processing_resolution) cropped_face, sampler_mask, restore_info = self.mask_processor.process_and_crop( image_tensor=image, mask_tensor=combined_mask, crop_padding=crop_padding, processing_resolution=target_resolution, mask_expansion=mask_expansion, enable_pre_upscale=enable_pre_upscale, upscaler_model_name=upscaler_model, upscaler_loader_callback=self.load_upscaler_model, upscaler_run_callback=self.run_upscaler ) cropped_face = cropped_face.to(device) sampler_mask = sampler_mask.to(device) # ── NEW: apply mask blur if requested ──────────────────────────── sampler_mask = self._blur_mask(sampler_mask, sampling_mask_blur_size, sampling_mask_blur_strength) # ───────────────────────────────────────────────────────────────── if isolate_face: mask_bchw = sampler_mask.unsqueeze(1) kh = max(3, (processing_resolution // 32) | 1) kw = kh sigma = max(1.0, kh / 3.0) blurred = kornia.filters.gaussian_blur2d(mask_bchw, (kh, kw), (sigma, sigma)) blurred = torch.clamp(blurred, 0.0, 1.0) mask_bhwc = blurred.permute(0, 2, 3, 1) bg = torch.zeros_like(cropped_face) cropped_face = cropped_face * mask_bhwc + bg * (1.0 - mask_bhwc) return (cropped_face, sampler_mask, restore_info) except model_management.InterruptProcessingException: raise except Exception as e: print(f"[Face Edit Prep] Error preparing face for edit: {e}") try: device = image.device if image is not None else model_management.get_torch_device() res = processing_resolution empty_face = torch.zeros((1, res, res, 3), dtype=torch.float32, device=device) empty_mask = torch.zeros((1, res, res), dtype=torch.float32, device=device) empty_info = { "original_image": np.zeros((res, res, 3), dtype=np.uint8), "original_image_size": (0, 0), "crop_coords": (0, 0, 0, 0), "face_bbox": (0, 0, 0, 0), "target_size": (res, res), "original_crop_size": (0, 0), "blend_mask": np.zeros((res, res), dtype=np.float32), "detection_angle": 0, } return (empty_face, empty_mask, empty_info) except: device = model_management.get_torch_device() empty_face = torch.zeros((1, 1024, 1024, 3), dtype=torch.float32, device=device) empty_mask = torch.zeros((1, 1024, 1024), dtype=torch.float32, device=device) empty_info = { "original_image": np.zeros((1024, 1024, 3), dtype=np.uint8), "original_image_size": (0, 0), "crop_coords": (0, 0, 0, 0), "face_bbox": (0, 0, 0, 0), "target_size": (1024, 1024), "original_crop_size": (0, 0), "blend_mask": np.zeros((1024, 1024), dtype=np.float32), "detection_angle": 0, } return (empty_face, empty_mask, empty_info) class ForbiddenVisionFaceEditMerge: @classmethod def INPUT_TYPES(s): return { "required": { "original_image": ("IMAGE", {"tooltip": "Original full image."}), "edited_face": ("IMAGE", {"tooltip": "Edited face image at the same resolution as the cropped face from prep node."}), "face_info": ("FACE_INFO", {"tooltip": "Face info output from the prep node."}), "blend_softness": ("INT", {"default": 8, "min": 0, "max": 200, "step": 1}), "enable_color_correction": ("BOOLEAN", {"default": True}), "color_correction_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.05}), } } RETURN_TYPES = ("IMAGE", "MASK") RETURN_NAMES = ("final_image", "compositing_mask") FUNCTION = "merge_face" CATEGORY = "Forbidden Vision" def __init__(self): ensure_model_directories() self.processor = ForbiddenVisionFaceProcessorIntegrated() def merge_face(self, original_image, edited_face, face_info, blend_softness, enable_color_correction, color_correction_strength): try: check_for_interruption() if original_image is None or edited_face is None or face_info is None: return (original_image, torch.zeros_like(original_image[:, :, :, 0])) if isinstance(face_info, (list, tuple)) and face_info and isinstance(face_info[0], dict): restore_info_list = [face_info[0]] elif isinstance(face_info, dict): restore_info_list = [face_info] else: return (original_image, torch.zeros_like(original_image[:, :, :, 0])) info = restore_info_list[0] crop_coords = info.get("crop_coords", (0, 0, 0, 0)) ox1, oy1, ox2, oy2 = crop_coords if ox2 <= ox1 or oy2 <= oy1: mask = torch.zeros((original_image.shape[0], original_image.shape[1], original_image.shape[2]), device=original_image.device) return (original_image, mask) processed_faces = [edited_face] final_image = self.processor.combine_all_faces_to_final_image( original_image, processed_faces, restore_info_list, blend_softness, enable_color_correction, color_correction_strength ) compositing_mask = self.processor.create_unified_mask( restore_info_list, original_image ) return (final_image, compositing_mask) except model_management.InterruptProcessingException: raise except Exception as e: print(f"[Face Edit Merge] Error merging edited face: {e}") try: mask = torch.zeros((original_image.shape[0], original_image.shape[1], original_image.shape[2]), device=original_image.device) return (original_image, mask) except: device = model_management.get_torch_device() mask = torch.zeros((1, 512, 512), dtype=torch.float32, device=device) return (None, mask)