1194 lines
57 KiB
Python
1194 lines
57 KiB
Python
import torch
|
|
import torch.nn.functional as F
|
|
import numpy as np
|
|
import random
|
|
import nodes
|
|
from PIL import Image
|
|
import folder_paths
|
|
import comfy.model_management as model_management
|
|
import comfy.utils
|
|
import comfy.samplers
|
|
import comfy.sample
|
|
import comfy.sd
|
|
import latent_preview
|
|
import cv2
|
|
import os
|
|
from datetime import datetime
|
|
from .face_detector import ForbiddenVisionFaceDetector
|
|
from .mask_processor import ForbiddenVisionMaskProcessor
|
|
from .utils import safe_tensor_to_numpy, safe_numpy_to_tensor, check_for_interruption, get_yolo_models, get_sam_models
|
|
|
|
class ForbiddenVisionFaceProcessorIntegrated:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"model": ("MODEL",),
|
|
"vae": ("VAE",),
|
|
"positive": ("CONDITIONING",),
|
|
"negative": ("CONDITIONING",),
|
|
|
|
"steps": ("INT", {"default": 8, "min": 1, "max": 100}),
|
|
"cfg_scale": ("FLOAT", {"default": 3.0, "min": 0.0, "max": 30.0, "step": 0.1}),
|
|
"sampler": (comfy.samplers.KSampler.SAMPLERS, {"default": "euler_ancestral"}),
|
|
"scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"default": "sgm_uniform"}),
|
|
"denoise_strength": ("FLOAT", {"default": 0.4, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
|
|
"processing_resolution": ("INT", {"default": 1024, "min": 512, "max": 2048, "step": 64, "tooltip": "The resolution at which the cropped mask area will be processed."}),
|
|
"enable_pre_upscale": ("BOOLEAN", {"default": True, "label_on": "Enabled", "label_off": "Disabled", "tooltip": "Enable to upscale small faces with an AI model before processing."}),
|
|
"upscaler_model": (ForbiddenVisionFaceProcessorIntegrated.get_ordered_upscaler_models(), {"tooltip": "The AI model used for pre-upscaling small faces. Only active if 'enable_pre_upscale' is on."}),
|
|
"crop_padding": ("FLOAT", {"default": 1.2, "min": 1.0, "max": 2.0, "step": 0.1, "tooltip": "Padding added to the mask's bounding box before cropping."}),
|
|
|
|
"enable_mask_cleanup": ("BOOLEAN", {"default": True, "label_on": "Enabled", "label_off": "Disabled", "tooltip": "Enable intelligent cleanup to fill holes and remove noise from the input mask."}),
|
|
|
|
"face_positive_prompt": ("STRING", {"multiline": True, "default": "face focus"}),
|
|
"face_negative_prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
"replace_positive_prompt": ("BOOLEAN", {"default": False}),
|
|
"replace_negative_prompt": ("BOOLEAN", {"default": False}),
|
|
|
|
"blend_softness": ("INT", {"default": 8, "min": 0, "max": 200, "step": 1}),
|
|
"mask_expansion": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1}),
|
|
"sampling_mask_blur_size": ("INT", {"default": 21, "min": 1, "max": 101, "step": 2}),
|
|
"sampling_mask_blur_strength": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 6.0, "step": 0.1}),
|
|
"enable_vertical_flip": ("BOOLEAN", {"default": False}),
|
|
"enable_color_correction": ("BOOLEAN", {"default": True}),
|
|
"color_correction_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.1}),
|
|
|
|
"face_selection": ("INT", {"default": 0, "min": 0, "max": 20, "step": 1, "tooltip": "0=All faces, 1=1st face, etc. Used by both mask input and internal detector."}),
|
|
"detection_confidence": ("FLOAT", {"default": 0.5, "min": 0.1, "max": 1.0, "step": 0.01, "tooltip": "Confidence threshold for the internal BBox detector (YOLO)."}),
|
|
"sam_threshold": ("FLOAT", {"default": 0.4, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Confidence threshold for the internal SAM mask refinement."}),
|
|
"bbox_model": (get_yolo_models(), {"tooltip": "Internal BBOX detector model (YOLO). Used if 'mask' is not connected."}),
|
|
"sam_model": (get_sam_models(), {"tooltip": "Internal Segment Anything Model (SAM). Used if 'mask' is not connected."}),
|
|
},
|
|
"optional": {
|
|
"mask": ("MASK", {"tooltip": "Optional: Face mask for processing. If connected, internal detection is skipped."}),
|
|
"clip": ("CLIP", {"tooltip": "Optional: Required only if using face prompts."}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "MASK")
|
|
RETURN_NAMES = ("final_image", "processed_face", "side_by_side_comparison", "final_mask")
|
|
FUNCTION = "process_face_complete"
|
|
CATEGORY = "Forbidden Vision"
|
|
@staticmethod
|
|
def get_ordered_upscaler_models():
|
|
preferred_models = [
|
|
"4x_NMKD-Siax_200k.pth",
|
|
"4x-UltraSharp.pth",
|
|
"4x_foolhardy_Remacri.pth",
|
|
"4x-AnimeSharp.pth"
|
|
]
|
|
|
|
all_models = folder_paths.get_filename_list("upscale_models") or []
|
|
if not all_models:
|
|
return ["None Found"]
|
|
|
|
found_preferred = []
|
|
for preferred in preferred_models:
|
|
for model in all_models:
|
|
if preferred in model or model in preferred:
|
|
if model not in found_preferred:
|
|
found_preferred.append(model)
|
|
break
|
|
|
|
remaining_models = [model for model in all_models if model not in found_preferred]
|
|
|
|
ordered_models = found_preferred + sorted(remaining_models)
|
|
|
|
return ordered_models if ordered_models else ["None Found"]
|
|
|
|
|
|
def __init__(self):
|
|
self.face_detector = ForbiddenVisionFaceDetector()
|
|
self.mask_processor = ForbiddenVisionMaskProcessor()
|
|
self.sampling_cache = {}
|
|
self.total_cache = {}
|
|
self.upscaler_model = None
|
|
self.upscaler_model_name = None
|
|
|
|
def clear_cache(self):
|
|
"""Clear all cached data to ensure fresh processing"""
|
|
try:
|
|
if hasattr(self, 'total_cache'):
|
|
self.total_cache.clear()
|
|
print("Face processor cache cleared")
|
|
except Exception as e:
|
|
print(f"Error clearing cache: {e}")
|
|
|
|
def load_upscaler_model(self, model_name):
|
|
if self.upscaler_model is not None and self.upscaler_model_name == model_name:
|
|
return True
|
|
|
|
try:
|
|
model_path = folder_paths.get_full_path("upscale_models", model_name)
|
|
if model_path is None:
|
|
print(f"Upscaler model '{model_name}' not found.")
|
|
return False
|
|
|
|
UpscalerLoaderClass = nodes.NODE_CLASS_MAPPINGS['UpscaleModelLoader']
|
|
upscaler_loader = UpscalerLoaderClass()
|
|
|
|
self.upscaler_model = upscaler_loader.load_model(model_name)[0]
|
|
self.upscaler_model_name = model_name
|
|
|
|
return True
|
|
|
|
except model_management.InterruptProcessingException:
|
|
raise
|
|
except KeyError:
|
|
print("Error: Could not find 'UpscaleModelLoader' in nodes.NODE_CLASS_MAPPINGS.")
|
|
return False
|
|
except Exception as e:
|
|
print(f"An error occurred loading upscaler model {model_name}: {e}")
|
|
self.upscaler_model = None
|
|
self.upscaler_model_name = None
|
|
return False
|
|
|
|
def run_upscaler(self, image_np_uint8):
|
|
if self.upscaler_model is None:
|
|
return image_np_uint8
|
|
|
|
image_tensor = torch.from_numpy(image_np_uint8.astype(np.float32) / 255.0).unsqueeze(0)
|
|
|
|
try:
|
|
ImageUpscalerClass = nodes.NODE_CLASS_MAPPINGS['ImageUpscaleWithModel']
|
|
upscaler_node = ImageUpscalerClass()
|
|
|
|
with torch.no_grad():
|
|
upscaled_tensor = upscaler_node.upscale(
|
|
upscale_model=self.upscaler_model,
|
|
image=image_tensor
|
|
)[0]
|
|
|
|
except model_management.InterruptProcessingException:
|
|
raise
|
|
except KeyError:
|
|
print("Error: Could not find 'ImageUpscaleWithModel' in nodes.NODE_CLASS_MAPPINGS.")
|
|
return image_np_uint8
|
|
except Exception as e:
|
|
print(f"Error during upscaling process: {e}. Falling back to standard resize.")
|
|
import traceback
|
|
traceback.print_exc()
|
|
return image_np_uint8
|
|
|
|
upscaled_np = upscaled_tensor.squeeze(0).cpu().numpy()
|
|
upscaled_np_uint8 = (np.clip(upscaled_np, 0, 1) * 255.0).round().astype(np.uint8)
|
|
|
|
return upscaled_np_uint8
|
|
|
|
def combine_all_faces_to_final_image(self, original_image, all_processed_faces, all_restore_info, blend_softness, enable_color_correction=False, color_correction_strength=1.0):
|
|
"""
|
|
Combines all processed faces onto the original image using a unified mask and an efficient,
|
|
performance-optimized color correction workflow.
|
|
"""
|
|
try:
|
|
check_for_interruption()
|
|
|
|
if not all_processed_faces or not all_restore_info:
|
|
return original_image
|
|
|
|
original_np = self.fast_tensor_to_numpy(original_image.clone(), target_range=(0, 1))
|
|
h, w = original_np.shape[:2]
|
|
|
|
new_faces_canvas = original_np.copy()
|
|
unified_compositing_mask = np.zeros((h, w), dtype=np.float32)
|
|
|
|
for i, (processed_face_tensor, restore_info) in enumerate(zip(all_processed_faces, all_restore_info)):
|
|
check_for_interruption()
|
|
|
|
if isinstance(restore_info, list):
|
|
restore_info = restore_info[0]
|
|
|
|
crop_coords = restore_info["crop_coords"]
|
|
crop_x1, crop_y1, crop_x2, crop_y2 = crop_coords
|
|
actual_crop_w, actual_crop_h = crop_x2 - crop_x1, crop_y2 - crop_y1
|
|
|
|
if actual_crop_w <= 0 or actual_crop_h <= 0:
|
|
continue
|
|
|
|
processed_face_np = self.fast_tensor_to_numpy(processed_face_tensor, target_range=(0, 1))
|
|
|
|
if enable_color_correction and color_correction_strength > 0.0:
|
|
original_crop_np = original_np[crop_y1:crop_y2, crop_x1:crop_x2]
|
|
|
|
processed_face_np = self._perform_color_correction(
|
|
processed_face_np, original_crop_np, color_correction_strength
|
|
)
|
|
|
|
resized_processed_face = self.high_quality_resize(
|
|
processed_face_np, (actual_crop_w, actual_crop_h)
|
|
)
|
|
|
|
new_faces_canvas[crop_y1:crop_y2, crop_x1:crop_x2] = resized_processed_face
|
|
|
|
full_blend_mask = restore_info.get("blend_mask")
|
|
if full_blend_mask is None:
|
|
paste_mask_crop = np.ones((actual_crop_h, actual_crop_w), dtype=np.float32)
|
|
else:
|
|
paste_mask_crop = full_blend_mask[crop_y1:crop_y2, crop_x1:crop_x2]
|
|
|
|
compositing_mask_crop = self.create_compositing_blend_mask(paste_mask_crop, blend_softness)
|
|
|
|
unified_compositing_mask[crop_y1:crop_y2, crop_x1:crop_x2] = np.maximum(
|
|
unified_compositing_mask[crop_y1:crop_y2, crop_x1:crop_x2],
|
|
compositing_mask_crop
|
|
)
|
|
|
|
check_for_interruption()
|
|
|
|
unified_compositing_mask_3d = np.stack([unified_compositing_mask] * 3, axis=-1)
|
|
|
|
final_result_np = (new_faces_canvas * unified_compositing_mask_3d) + (original_np * (1.0 - unified_compositing_mask_3d))
|
|
|
|
final_tensor = self.fast_numpy_to_tensor(final_result_np, input_range=(0, 1))
|
|
|
|
return final_tensor
|
|
|
|
except model_management.InterruptProcessingException:
|
|
raise
|
|
except Exception as e:
|
|
print(f"Error combining faces to final image: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
return original_image
|
|
def process_face_complete(self, image, model, vae, positive, negative,
|
|
steps, cfg_scale, sampler, scheduler, denoise_strength, seed,
|
|
processing_resolution, enable_pre_upscale, upscaler_model, crop_padding, enable_mask_cleanup,
|
|
face_positive_prompt, face_negative_prompt, replace_positive_prompt, replace_negative_prompt,
|
|
blend_softness, mask_expansion,
|
|
sampling_mask_blur_size, sampling_mask_blur_strength,
|
|
enable_vertical_flip, enable_color_correction, color_correction_strength,
|
|
face_selection, detection_confidence, sam_threshold, bbox_model, sam_model,
|
|
mask=None, clip=None):
|
|
try:
|
|
check_for_interruption()
|
|
|
|
if image is None or model is None or vae is None:
|
|
print("ERROR: Required inputs (image, model, vae) are missing.")
|
|
return self.create_safe_fallback_outputs(image, processing_resolution)
|
|
|
|
original_image = image
|
|
face_masks = []
|
|
|
|
if mask is not None:
|
|
print("INFO: 'mask' input connected. Processing based on provided mask.")
|
|
face_masks = [mask]
|
|
else:
|
|
print("INFO: No 'mask' input. Running internal face detection with two-stage approach.")
|
|
np_masks = self.face_detector.detect_faces(
|
|
image_tensor=image,
|
|
bbox_model_name=bbox_model,
|
|
sam_model_name=sam_model,
|
|
detection_confidence=detection_confidence,
|
|
sam_threshold=sam_threshold,
|
|
face_selection=face_selection
|
|
)
|
|
if np_masks:
|
|
face_masks = [torch.from_numpy(m).unsqueeze(0) for m in np_masks]
|
|
|
|
if not face_masks:
|
|
print("ERROR: No face masks found from any source. Cannot proceed.")
|
|
return self.create_safe_fallback_outputs(image, processing_resolution)
|
|
|
|
sampling_cache_key = self.get_sampling_cache_key_only(
|
|
image, mask, model, vae, positive, negative, steps, cfg_scale, sampler, scheduler,
|
|
denoise_strength, seed, processing_resolution, enable_pre_upscale, upscaler_model, crop_padding,
|
|
enable_mask_cleanup, face_positive_prompt, face_negative_prompt, replace_positive_prompt,
|
|
replace_negative_prompt, mask_expansion, sampling_mask_blur_size, sampling_mask_blur_strength,
|
|
enable_vertical_flip, face_selection, clip, detection_confidence, sam_threshold, bbox_model, sam_model
|
|
)
|
|
compositing_cache_key = self.get_compositing_cache_key(
|
|
sampling_cache_key, blend_softness, enable_color_correction, color_correction_strength
|
|
)
|
|
|
|
if hasattr(self, 'total_cache') and compositing_cache_key and compositing_cache_key in self.total_cache:
|
|
print("DEBUG: Using cached results for entire processing (L2 Cache Hit)")
|
|
return self.total_cache[compositing_cache_key]
|
|
|
|
all_processed_faces, all_restore_info = [], []
|
|
|
|
if hasattr(self, 'sampling_cache') and sampling_cache_key and sampling_cache_key in self.sampling_cache:
|
|
print("DEBUG: Using cached sampling results (L1 Cache Hit). Re-compositing...")
|
|
all_processed_faces, all_restore_info = self.sampling_cache[sampling_cache_key]
|
|
else:
|
|
print("DEBUG: Running fresh sampling for all faces (Cache Miss)")
|
|
|
|
if enable_vertical_flip:
|
|
image = torch.flip(image, dims=[1])
|
|
if mask is not None:
|
|
face_masks = [torch.flip(m, dims=[1]) for m in face_masks]
|
|
|
|
original_base_image = image.clone()
|
|
|
|
for i, face_mask in enumerate(face_masks):
|
|
check_for_interruption()
|
|
print(f"DEBUG: Processing face {i+1}/{len(face_masks)}")
|
|
processed_result = self.process_single_face_unified(
|
|
original_base_image, face_mask, model, vae, positive, negative,
|
|
steps, cfg_scale, sampler, scheduler, denoise_strength, seed + i,
|
|
processing_resolution, enable_pre_upscale, upscaler_model, crop_padding, enable_mask_cleanup,
|
|
face_positive_prompt, face_negative_prompt, replace_positive_prompt, replace_negative_prompt,
|
|
mask_expansion, sampling_mask_blur_size, sampling_mask_blur_strength, clip
|
|
)
|
|
if processed_result:
|
|
_, processed_face, restore_info = processed_result
|
|
all_processed_faces.append(processed_face)
|
|
all_restore_info.append(restore_info)
|
|
if len(face_masks) > 1 and i < len(face_masks) - 1 and torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
if sampling_cache_key:
|
|
if not hasattr(self, 'sampling_cache'): self.sampling_cache = {}
|
|
self.sampling_cache[sampling_cache_key] = (all_processed_faces, all_restore_info)
|
|
|
|
if not all_processed_faces:
|
|
print("ERROR: All face processing failed. Returning original image.")
|
|
return self.create_safe_fallback_outputs(original_image, processing_resolution)
|
|
|
|
compositing_base_image = torch.flip(image, dims=[1]) if enable_vertical_flip else original_image.clone()
|
|
final_image = self.combine_all_faces_to_final_image(compositing_base_image, all_processed_faces, all_restore_info, blend_softness, enable_color_correction, color_correction_strength)
|
|
processed_face_output = self.create_combined_face_output(all_processed_faces, processing_resolution)
|
|
side_by_side = self.create_unified_comparison(original_image, all_processed_faces, all_restore_info, processing_resolution)
|
|
final_mask = self.create_unified_mask(all_restore_info, original_image)
|
|
|
|
if enable_vertical_flip:
|
|
processed_face_output = torch.flip(processed_face_output, dims=[1])
|
|
side_by_side = torch.flip(side_by_side, dims=[1])
|
|
final_mask = torch.flip(final_mask, dims=[1])
|
|
|
|
final_results = (final_image, processed_face_output, side_by_side, final_mask)
|
|
|
|
if compositing_cache_key:
|
|
if not hasattr(self, 'total_cache'): self.total_cache = {}
|
|
self.total_cache[compositing_cache_key] = final_results
|
|
|
|
return final_results
|
|
|
|
except model_management.InterruptProcessingException:
|
|
print("Face processing cancelled by user.")
|
|
raise
|
|
except Exception as e:
|
|
print(f"An error occurred during the main face processing workflow: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
(fallback_image, fallback_processed, fallback_comparison, fallback_mask) = self.create_safe_fallback_outputs(original_image, processing_resolution)
|
|
if enable_vertical_flip:
|
|
fallback_image = torch.flip(fallback_image, dims=[1])
|
|
fallback_processed = torch.flip(fallback_processed, dims=[1])
|
|
fallback_comparison = torch.flip(fallback_comparison, dims=[1])
|
|
fallback_mask = torch.flip(fallback_mask, dims=[1])
|
|
return (fallback_image, fallback_processed, fallback_comparison, fallback_mask)
|
|
|
|
def get_sampling_cache_key_only(self, image, mask, model, vae, positive, negative,
|
|
steps, cfg_scale, sampler, scheduler, denoise_strength, seed,
|
|
processing_resolution, enable_pre_upscale, upscaler_model, crop_padding, enable_mask_cleanup,
|
|
face_positive_prompt, face_negative_prompt, replace_positive_prompt, replace_negative_prompt,
|
|
mask_expansion, sampling_mask_blur_size, sampling_mask_blur_strength,
|
|
enable_vertical_flip, face_selection, clip,
|
|
detection_confidence, sam_threshold, bbox_model, sam_model):
|
|
try:
|
|
image_sample = image[0, ::96, ::96, :].flatten()[:50] if image.numel() > 2500 else image.flatten()[:50]
|
|
image_hash = hash(tuple(image_sample.cpu().numpy().round(3)))
|
|
|
|
mask_hash = 0
|
|
if mask is not None:
|
|
mask_sample = mask[0, ::96, ::96].flatten()[:50] if mask.numel() > 2500 else mask.flatten()[:50]
|
|
mask_hash = hash(tuple(mask_sample.cpu().numpy().round(3)))
|
|
else:
|
|
mask_hash = hash((detection_confidence, sam_threshold, bbox_model, sam_model))
|
|
|
|
model_hash = id(model)
|
|
vae_hash = id(vae)
|
|
|
|
pos_hash = 0
|
|
if positive and len(positive) > 0 and len(positive[0]) > 0:
|
|
pos_sample = positive[0][0].flatten()[:20].cpu().numpy() if positive[0][0].numel() > 20 else positive[0][0].flatten().cpu().numpy()
|
|
pos_hash = hash(tuple(pos_sample.round(3)))
|
|
|
|
neg_hash = 0
|
|
if negative and len(negative) > 0 and len(negative[0]) > 0:
|
|
neg_sample = negative[0][0].flatten()[:20].cpu().numpy() if negative[0][0].numel() > 20 else negative[0][0].flatten().cpu().numpy()
|
|
neg_hash = hash(tuple(neg_sample.round(3)))
|
|
|
|
face_pos_hash = hash(face_positive_prompt[:100]) if face_positive_prompt else 0
|
|
face_neg_hash = hash(face_negative_prompt[:100]) if face_negative_prompt else 0
|
|
clip_hash = id(clip) if clip else 0
|
|
|
|
return (
|
|
image_hash, mask_hash, tuple(image.shape),
|
|
model_hash, vae_hash, pos_hash, neg_hash,
|
|
steps, cfg_scale, sampler, scheduler, denoise_strength, seed,
|
|
processing_resolution, enable_pre_upscale, upscaler_model, crop_padding, enable_mask_cleanup,
|
|
face_pos_hash, face_neg_hash, replace_positive_prompt, replace_negative_prompt,
|
|
mask_expansion, sampling_mask_blur_size, sampling_mask_blur_strength,
|
|
enable_vertical_flip, face_selection, clip_hash
|
|
)
|
|
|
|
except Exception as e:
|
|
print(f"Error generating sampling cache key: {e}")
|
|
return None
|
|
|
|
def get_compositing_cache_key(self, sampling_cache_key, blend_softness, enable_color_correction, color_correction_strength):
|
|
"""Cache key for final compositing (includes compositing-only parameters)"""
|
|
try:
|
|
return (
|
|
sampling_cache_key,
|
|
blend_softness, enable_color_correction, color_correction_strength
|
|
)
|
|
except Exception as e:
|
|
print(f"Error generating compositing cache key: {e}")
|
|
return None
|
|
|
|
|
|
|
|
|
|
def process_single_face_unified(self, image, mask, model, vae, positive, negative,
|
|
steps, cfg_scale, sampler, scheduler, denoise_strength, seed,
|
|
processing_resolution, enable_pre_upscale, upscaler_model, crop_padding, enable_mask_cleanup,
|
|
face_positive_prompt, face_negative_prompt, replace_positive_prompt, replace_negative_prompt,
|
|
mask_expansion, sampling_mask_blur_size, sampling_mask_blur_strength, clip):
|
|
"""Unified method to process a single face"""
|
|
try:
|
|
if mask is None:
|
|
return None
|
|
|
|
cropped_face, sampler_mask_batch, restore_info = self.mask_processor.process_and_crop(
|
|
image_tensor=image,
|
|
mask_tensor=mask,
|
|
enable_cleanup=enable_mask_cleanup,
|
|
crop_padding=crop_padding,
|
|
processing_resolution=processing_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
|
|
)
|
|
|
|
if self.is_empty_detection(cropped_face, restore_info):
|
|
return None
|
|
|
|
enhanced_positive, enhanced_negative = self.enhance_conditioning_with_face_prompts(
|
|
positive, negative, face_positive_prompt, face_negative_prompt,
|
|
replace_positive_prompt, replace_negative_prompt, clip
|
|
)
|
|
|
|
processed_latent = self.run_inpaint_sampling(
|
|
cropped_face, sampler_mask_batch, model, vae, enhanced_positive, enhanced_negative,
|
|
steps, cfg_scale, sampler, scheduler, denoise_strength, seed,
|
|
sampling_mask_blur_size, sampling_mask_blur_strength
|
|
)
|
|
|
|
with torch.no_grad():
|
|
processed_face_batch = vae.decode(processed_latent["samples"])
|
|
if processed_face_batch.min() < -0.5:
|
|
processed_face_batch = (processed_face_batch + 1.0) / 2.0
|
|
processed_face_batch = torch.clamp(processed_face_batch, 0.0, 1.0)
|
|
|
|
|
|
return (image, processed_face_batch, restore_info)
|
|
|
|
except model_management.InterruptProcessingException:
|
|
raise
|
|
except Exception as e:
|
|
print(f"Error processing single face: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
return None
|
|
|
|
def create_combined_face_output(self, processed_faces, processing_resolution):
|
|
"""Create combined output from multiple processed faces"""
|
|
try:
|
|
if not processed_faces:
|
|
return torch.zeros((1, processing_resolution, processing_resolution, 3))
|
|
|
|
if len(processed_faces) == 1:
|
|
return processed_faces[0]
|
|
|
|
face_crops = []
|
|
for processed_face in processed_faces:
|
|
processed_np = self.fast_tensor_to_numpy(processed_face, target_range=(0, 255))
|
|
processed_resized = cv2.resize(processed_np, (processing_resolution, processing_resolution), interpolation=cv2.INTER_CUBIC)
|
|
face_crops.append(processed_resized)
|
|
|
|
if len(face_crops) == 2:
|
|
combined = np.concatenate(face_crops, axis=1)
|
|
else:
|
|
rows = []
|
|
for i in range(0, len(face_crops), 2):
|
|
row_faces = face_crops[i:i+2]
|
|
if len(row_faces) == 1:
|
|
row_faces.append(np.zeros_like(row_faces[0]))
|
|
row = np.concatenate(row_faces, axis=1)
|
|
rows.append(row)
|
|
combined = np.concatenate(rows, axis=0)
|
|
|
|
return self.fast_numpy_to_tensor(combined, input_range=(0, 255))
|
|
|
|
except model_management.InterruptProcessingException:
|
|
raise
|
|
except Exception as e:
|
|
print(f"Error creating combined face output: {e}")
|
|
return processed_faces[0] if processed_faces else torch.zeros((1, processing_resolution, processing_resolution, 3))
|
|
|
|
def create_unified_comparison(self, original_image, processed_faces, restore_info_list, processing_resolution):
|
|
"""Create unified side-by-side comparison"""
|
|
try:
|
|
if not processed_faces or not restore_info_list:
|
|
return original_image
|
|
|
|
original_np = self.fast_tensor_to_numpy(original_image, target_range=(0, 255))
|
|
comparisons = []
|
|
|
|
for i, (processed_face, restore_info) in enumerate(zip(processed_faces, restore_info_list)):
|
|
crop_coords = restore_info["crop_coords"]
|
|
crop_x1, crop_y1, crop_x2, crop_y2 = crop_coords
|
|
|
|
if crop_x2 <= crop_x1 or crop_y2 <= crop_y1:
|
|
continue
|
|
|
|
original_crop = original_np[crop_y1:crop_y2, crop_x1:crop_x2]
|
|
original_resized = cv2.resize(original_crop, (processing_resolution, processing_resolution), interpolation=cv2.INTER_CUBIC)
|
|
|
|
processed_np = self.fast_tensor_to_numpy(processed_face, target_range=(0, 255))
|
|
processed_resized = cv2.resize(processed_np, (processing_resolution, processing_resolution), interpolation=cv2.INTER_CUBIC)
|
|
|
|
face_comparison = np.concatenate([original_resized, processed_resized], axis=1)
|
|
comparisons.append(face_comparison)
|
|
|
|
if not comparisons:
|
|
return original_image
|
|
|
|
if len(comparisons) == 1:
|
|
final_comparison = comparisons[0]
|
|
else:
|
|
final_comparison = np.concatenate(comparisons, axis=0)
|
|
|
|
return self.fast_numpy_to_tensor(final_comparison, input_range=(0, 255))
|
|
|
|
except model_management.InterruptProcessingException:
|
|
raise
|
|
except Exception as e:
|
|
print(f"Error creating unified comparison: {e}")
|
|
return original_image
|
|
|
|
def create_unified_mask(self, restore_info_list, image):
|
|
"""Create unified mask from all processed faces"""
|
|
try:
|
|
if not restore_info_list:
|
|
image_shape = image.shape
|
|
return torch.zeros((1, image_shape[1], image_shape[2]), dtype=torch.float32)
|
|
|
|
print(f"DEBUG: Creating unified mask from {len(restore_info_list)} faces")
|
|
|
|
image_shape = image.shape
|
|
original_height, original_width = image_shape[1], image_shape[2]
|
|
combined_mask = np.zeros((original_height, original_width), dtype=np.float32)
|
|
|
|
for i, restore_info in enumerate(restore_info_list):
|
|
blend_mask = restore_info.get("blend_mask")
|
|
if blend_mask is not None:
|
|
print(f"DEBUG: Adding blend_mask from face {i} (shape: {blend_mask.shape})")
|
|
combined_mask = np.maximum(combined_mask, blend_mask)
|
|
else:
|
|
print(f"DEBUG: Face {i} has no blend_mask")
|
|
|
|
print(f"DEBUG: Final combined mask has {np.sum(combined_mask > 0)} non-zero pixels")
|
|
return torch.from_numpy(combined_mask).unsqueeze(0)
|
|
|
|
except model_management.InterruptProcessingException:
|
|
raise
|
|
except Exception as e:
|
|
print(f"Error creating unified mask: {e}")
|
|
image_shape = image.shape
|
|
return torch.zeros((1, image_shape[1], image_shape[2]), dtype=torch.float32)
|
|
|
|
|
|
def process_single_face_sampling(self, cropped_face, face_mask, model, vae, positive, negative,
|
|
steps, cfg_scale, sampler, scheduler, denoise_strength, seed,
|
|
sampling_mask_blur_size, sampling_mask_blur_strength):
|
|
try:
|
|
face_latent = self.encode_image_to_latent(cropped_face, vae)
|
|
latent_samples = face_latent["samples"]
|
|
|
|
B, C, H, W = latent_samples.shape
|
|
resized_mask = self.process_inpaint_mask(face_mask, H, W, latent_samples.device,
|
|
sampling_mask_blur_size, sampling_mask_blur_strength)
|
|
|
|
conditioned_positive, conditioned_negative = self.prepare_inpaint_conditioning(
|
|
positive, negative, latent_samples, resized_mask.unsqueeze(1)
|
|
)
|
|
|
|
sampled_latent = self.run_ksampler(
|
|
model, conditioned_positive, conditioned_negative, face_latent,
|
|
steps, cfg_scale, sampler, scheduler, denoise_strength, seed,
|
|
denoise_mask=resized_mask
|
|
)
|
|
|
|
return sampled_latent
|
|
except model_management.InterruptProcessingException:
|
|
raise
|
|
except Exception as e:
|
|
print(f"Error in single face sampling: {e}")
|
|
raise e
|
|
|
|
|
|
def create_safe_fallback_outputs(self, image, processing_resolution):
|
|
"""Create safe fallback outputs that match expected formats"""
|
|
try:
|
|
final_output = image.clone() if image is not None else torch.zeros((1, 512, 512, 3), dtype=torch.float32)
|
|
|
|
empty_processed = torch.zeros((1, processing_resolution, processing_resolution, 3), dtype=torch.float32)
|
|
empty_mask = torch.zeros((1, 512, 512), dtype=torch.float32)
|
|
|
|
if image is not None:
|
|
original_np = self.tensor_to_numpy_comparison(image)
|
|
h, w = original_np.shape[:2]
|
|
black_side = np.zeros_like(original_np)
|
|
comparison_np = np.concatenate([original_np, black_side], axis=1)
|
|
empty_comparison = self.numpy_to_tensor_comparison(comparison_np)
|
|
empty_mask = torch.zeros((1, h, w), dtype=torch.float32)
|
|
else:
|
|
empty_comparison = torch.zeros((1, 512, 1024, 3), dtype=torch.float32)
|
|
|
|
return (final_output, empty_processed, empty_comparison, empty_mask)
|
|
except model_management.InterruptProcessingException:
|
|
raise
|
|
except:
|
|
fallback_img = torch.zeros((1, 512, 512, 3), dtype=torch.float32)
|
|
fallback_mask = torch.zeros((1, 512, 512), dtype=torch.float32)
|
|
return (fallback_img, fallback_img, fallback_img, fallback_mask)
|
|
|
|
|
|
def enhance_conditioning_with_face_prompts(self, positive, negative, face_positive_prompt, face_negative_prompt, replace_positive_prompt, replace_negative_prompt, clip):
|
|
try:
|
|
has_positive_prompt = face_positive_prompt and face_positive_prompt.strip() != ""
|
|
has_negative_prompt = face_negative_prompt and face_negative_prompt.strip() != ""
|
|
|
|
if not has_positive_prompt and not has_negative_prompt:
|
|
return positive, negative
|
|
|
|
if clip is None:
|
|
print("Face prompts provided but no CLIP model connected - ignoring face prompts")
|
|
print("Connect a CLIP model to use face_positive_prompt and face_negative_prompt")
|
|
return positive, negative
|
|
|
|
enhanced_positive = self.modify_conditioning_with_prompt(
|
|
positive, face_positive_prompt, "positive", clip, replace_positive_prompt
|
|
) if has_positive_prompt else positive
|
|
|
|
enhanced_negative = self.modify_conditioning_with_prompt(
|
|
negative, face_negative_prompt, "negative", clip, replace_negative_prompt
|
|
) if has_negative_prompt else negative
|
|
|
|
return enhanced_positive, enhanced_negative
|
|
|
|
except model_management.InterruptProcessingException:
|
|
raise
|
|
except Exception as e:
|
|
print(f"Error enhancing conditioning with face prompts: {e}")
|
|
print("Falling back to original conditioning")
|
|
return positive, negative
|
|
|
|
def modify_conditioning_with_prompt(self, conditioning, face_prompt, prompt_type, clip, overwrite_mode):
|
|
try:
|
|
if not face_prompt or not face_prompt.strip():
|
|
return conditioning
|
|
|
|
tokens = clip.tokenize(face_prompt.strip())
|
|
new_cond, new_pooled = clip.encode_from_tokens(tokens, return_pooled=True)
|
|
|
|
if overwrite_mode:
|
|
enhanced_conditioning = []
|
|
for original_item in conditioning:
|
|
new_opts = original_item[1].copy()
|
|
new_opts["pooled_output"] = new_pooled
|
|
enhanced_conditioning.append([new_cond, new_opts])
|
|
return enhanced_conditioning
|
|
else:
|
|
if not conditioning:
|
|
return [[new_cond, {"pooled_output": new_pooled}]]
|
|
|
|
new_opts = conditioning[0][1].copy()
|
|
new_opts["pooled_output"] = new_pooled
|
|
new_item = [new_cond, new_opts]
|
|
|
|
return [new_item] + conditioning
|
|
|
|
except model_management.InterruptProcessingException:
|
|
raise
|
|
except Exception as e:
|
|
print(f"Error modifying {prompt_type} conditioning: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
return conditioning
|
|
|
|
def tensor_to_numpy_comparison(self, tensor):
|
|
return self.fast_tensor_to_numpy(tensor, target_range=(0, 255))
|
|
|
|
def numpy_to_tensor_comparison(self, np_array):
|
|
return self.fast_numpy_to_tensor(np_array, input_range=(0, 255))
|
|
|
|
def _perform_color_correction(self, processed_face_np, original_crop_np, correction_strength):
|
|
"""
|
|
An efficient helper that performs color correction on NumPy arrays.
|
|
It expects the processed face and the corresponding crop from the original image.
|
|
"""
|
|
try:
|
|
original_face_resized = cv2.resize(original_crop_np,
|
|
(processed_face_np.shape[1], processed_face_np.shape[0]),
|
|
interpolation=cv2.INTER_CUBIC)
|
|
|
|
corrected_face = self.lab_color_transfer(processed_face_np, original_face_resized)
|
|
|
|
if correction_strength < 1.0:
|
|
corrected_face = processed_face_np * (1.0 - correction_strength) + corrected_face * correction_strength
|
|
|
|
return np.clip(corrected_face, 0.0, 1.0)
|
|
|
|
except model_management.InterruptProcessingException:
|
|
raise
|
|
except Exception as e:
|
|
print(f"Error in _perform_color_correction helper: {e}")
|
|
return processed_face_np
|
|
|
|
def lab_color_transfer(self, source_image, target_image):
|
|
try:
|
|
if source_image.shape != target_image.shape:
|
|
target_image = cv2.resize(target_image, (source_image.shape[1], source_image.shape[0]))
|
|
|
|
source_uint8 = (np.clip(source_image, 0, 1) * 255).astype(np.uint8)
|
|
target_uint8 = (np.clip(target_image, 0, 1) * 255).astype(np.uint8)
|
|
|
|
source_lab = cv2.cvtColor(source_uint8, cv2.COLOR_RGB2LAB).astype(np.float32)
|
|
target_lab = cv2.cvtColor(target_uint8, cv2.COLOR_RGB2LAB).astype(np.float32)
|
|
|
|
source_mean = np.mean(source_lab.reshape(-1, 3), axis=0)
|
|
source_std = np.std(source_lab.reshape(-1, 3), axis=0)
|
|
|
|
target_mean = np.mean(target_lab.reshape(-1, 3), axis=0)
|
|
target_std = np.std(target_lab.reshape(-1, 3), axis=0)
|
|
|
|
result_lab = source_lab.copy()
|
|
|
|
for i in range(3):
|
|
if source_std[i] > 1e-6:
|
|
result_lab[:, :, i] = (result_lab[:, :, i] - source_mean[i]) * (target_std[i] / source_std[i]) + target_mean[i]
|
|
|
|
result_lab = np.clip(result_lab, 0, 255)
|
|
result_rgb = cv2.cvtColor(result_lab.astype(np.uint8), cv2.COLOR_LAB2RGB)
|
|
|
|
return result_rgb.astype(np.float32) / 255.0
|
|
|
|
except model_management.InterruptProcessingException:
|
|
raise
|
|
except Exception as e:
|
|
print(f"Error in LAB color transfer: {e}")
|
|
return source_image
|
|
|
|
def is_empty_detection(self, cropped_face, restore_info):
|
|
try:
|
|
if isinstance(restore_info, list):
|
|
return len(restore_info) == 0 or all(info.get("original_image_size", (0, 0))[0] == 0 for info in restore_info)
|
|
else:
|
|
return restore_info.get("original_image_size", (0, 0))[0] == 0
|
|
except model_management.InterruptProcessingException:
|
|
raise
|
|
except:
|
|
return True
|
|
|
|
def run_inpaint_sampling(self, cropped_face, face_mask, model, vae, positive, negative,
|
|
steps, cfg_scale, sampler, scheduler, denoise_strength, seed,
|
|
sampling_mask_blur_size, sampling_mask_blur_strength):
|
|
try:
|
|
result = self.process_single_face_sampling(
|
|
cropped_face, face_mask, model, vae, positive, negative,
|
|
steps, cfg_scale, sampler, scheduler, denoise_strength, seed,
|
|
sampling_mask_blur_size, sampling_mask_blur_strength
|
|
)
|
|
return result
|
|
except model_management.InterruptProcessingException:
|
|
raise
|
|
except Exception as e:
|
|
print(f"Error in inpaint sampling: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
return {"samples": torch.zeros((1, 4, 64, 64), dtype=torch.float32)}
|
|
|
|
def process_inpaint_mask(self, face_mask, latent_height, latent_width, device, blur_size, blur_strength):
|
|
try:
|
|
mask_tensor = face_mask.reshape(1, 1, face_mask.shape[-2], face_mask.shape[-1])
|
|
resized_mask = torch.nn.functional.interpolate(mask_tensor, size=(latent_height, latent_width), mode='bilinear', align_corners=False).squeeze(1)
|
|
|
|
resized_mask_np = resized_mask.cpu().numpy().squeeze()
|
|
|
|
if blur_size > 1 and blur_strength > 0:
|
|
base_sigma = (blur_size - 1) / 8.0
|
|
actual_sigma = base_sigma * (1 + np.tanh(blur_strength - 1) * 2)
|
|
|
|
if blur_size % 2 == 0:
|
|
blur_size += 1
|
|
|
|
blurred_mask = cv2.GaussianBlur(resized_mask_np, (blur_size, blur_size), actual_sigma)
|
|
resized_mask = torch.from_numpy(blurred_mask).unsqueeze(0).to(device)
|
|
|
|
return resized_mask
|
|
except model_management.InterruptProcessingException:
|
|
raise
|
|
except Exception as e:
|
|
print(f"Error processing inpaint mask: {e}")
|
|
return resized_mask
|
|
|
|
|
|
def encode_image_to_latent(self, image, vae):
|
|
try:
|
|
with torch.no_grad():
|
|
encoded = vae.encode(image)
|
|
|
|
return {"samples": encoded}
|
|
|
|
except model_management.InterruptProcessingException:
|
|
raise
|
|
except Exception as e:
|
|
print(f"Error encoding image to latent: {e}")
|
|
raise e
|
|
|
|
def prepare_inpaint_conditioning(self, positive, negative, latent_samples, resized_mask):
|
|
try:
|
|
conditioned_positive = []
|
|
for p in positive:
|
|
new_cond_dict = p[1].copy()
|
|
new_cond_dict["concat_latent_image"] = latent_samples
|
|
new_cond_dict["concat_mask"] = resized_mask
|
|
conditioned_positive.append([p[0], new_cond_dict])
|
|
|
|
conditioned_negative = []
|
|
for n in negative:
|
|
new_cond_dict = n[1].copy()
|
|
new_cond_dict["concat_latent_image"] = latent_samples
|
|
new_cond_dict["concat_mask"] = resized_mask
|
|
conditioned_negative.append([n[0], new_cond_dict])
|
|
|
|
return conditioned_positive, conditioned_negative
|
|
|
|
except model_management.InterruptProcessingException:
|
|
raise
|
|
except Exception as e:
|
|
print(f"Error preparing inpaint conditioning: {e}")
|
|
return positive, negative
|
|
|
|
def run_ksampler(self, model, positive, negative, latent, steps, cfg, sampler_name, scheduler, denoise, seed, denoise_mask=None):
|
|
try:
|
|
check_for_interruption()
|
|
device = model_management.get_torch_device()
|
|
latent_image = latent["samples"]
|
|
|
|
if denoise_mask is not None:
|
|
denoise_mask = denoise_mask.to(device)
|
|
|
|
try:
|
|
noise = comfy.sample.prepare_noise(latent_image, seed, device=device)
|
|
except TypeError:
|
|
noise = comfy.sample.prepare_noise(latent_image, seed)
|
|
noise = noise.to(device)
|
|
positive_cond = self.prepare_conditioning_for_sampling(positive, device)
|
|
negative_cond = self.prepare_conditioning_for_sampling(negative, device)
|
|
|
|
try:
|
|
previewer = latent_preview.get_previewer(device, model.model.latent_format)
|
|
except:
|
|
previewer = None
|
|
|
|
pbar = comfy.utils.ProgressBar(steps)
|
|
|
|
def preview_callback(step, x0, x, total_steps):
|
|
preview_bytes = None
|
|
if previewer:
|
|
try:
|
|
preview_bytes = previewer.decode_latent_to_preview_image("JPEG", x0)
|
|
except:
|
|
pass
|
|
pbar.update_absolute(step + 1, total_steps, preview_bytes)
|
|
|
|
|
|
sampler = comfy.samplers.KSampler(
|
|
model, steps=steps, device=device, sampler=sampler_name,
|
|
scheduler=scheduler, denoise=denoise, model_options=model.model_options,
|
|
)
|
|
|
|
samples = sampler.sample(
|
|
noise, positive_cond, negative_cond, cfg=cfg,
|
|
latent_image=latent_image, denoise_mask=denoise_mask,
|
|
start_step=0, last_step=steps, force_full_denoise=False,
|
|
callback=preview_callback,
|
|
)
|
|
|
|
return {"samples": samples}
|
|
|
|
except model_management.InterruptProcessingException:
|
|
raise
|
|
except Exception as e:
|
|
print(f"Error in KSampler: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
latent_shape = latent.get("samples", torch.zeros((1, 4, 64, 64))).shape
|
|
return {"samples": torch.zeros(latent_shape, device=model_management.get_torch_device())}
|
|
|
|
def prepare_conditioning_for_sampling(self, conditioning, device):
|
|
try:
|
|
if not conditioning:
|
|
return conditioning
|
|
|
|
prepared_conditioning = []
|
|
|
|
for cond_item in conditioning:
|
|
if cond_item[0].device == device:
|
|
cond_tensor = cond_item[0]
|
|
else:
|
|
cond_tensor = cond_item[0].to(device)
|
|
|
|
cond_dict = {}
|
|
for key, value in cond_item[1].items():
|
|
if torch.is_tensor(value) and value.device != device:
|
|
cond_dict[key] = value.to(device)
|
|
else:
|
|
cond_dict[key] = value
|
|
|
|
prepared_conditioning.append([cond_tensor, cond_dict])
|
|
|
|
return prepared_conditioning
|
|
|
|
except model_management.InterruptProcessingException:
|
|
raise
|
|
except Exception as e:
|
|
print(f"Error preparing conditioning: {e}")
|
|
return conditioning
|
|
|
|
def restore_faces_unified(self, processed_face_batch, restore_info, blend_softness):
|
|
try:
|
|
if isinstance(restore_info, list):
|
|
restore_info = restore_info[0]
|
|
|
|
original_image_np = restore_info["original_image"].squeeze(0)
|
|
if original_image_np.max() <= 1.0:
|
|
final_result_np = original_image_np.copy()
|
|
else:
|
|
final_result_np = original_image_np.astype(np.float32) / 255.0
|
|
|
|
processed_float = processed_face_batch.squeeze(0).cpu().numpy()
|
|
|
|
crop_coords = restore_info["crop_coords"]
|
|
crop_x1, crop_y1, crop_x2, crop_y2 = crop_coords
|
|
|
|
actual_crop_w, actual_crop_h = crop_x2 - crop_x1, crop_y2 - crop_y1
|
|
|
|
if actual_crop_w <= 0 or actual_crop_h <= 0:
|
|
print(f"Skipping restoration due to invalid crop dimensions: {actual_crop_w}x{actual_crop_h}")
|
|
return torch.from_numpy(restore_info["original_image"])
|
|
|
|
resized_processed_face = self.high_quality_resize(
|
|
processed_float, (actual_crop_w, actual_crop_h)
|
|
)
|
|
|
|
original_crop = final_result_np[crop_y1:crop_y2, crop_x1:crop_x2]
|
|
|
|
full_blend_mask = restore_info.get("blend_mask")
|
|
if full_blend_mask is None:
|
|
print(f"Warning: No blend_mask found, using fallback")
|
|
paste_mask = np.ones((actual_crop_h, actual_crop_w), dtype=np.float32)
|
|
else:
|
|
paste_mask = full_blend_mask[crop_y1:crop_y2, crop_x1:crop_x2]
|
|
if paste_mask.shape != (actual_crop_h, actual_crop_w):
|
|
paste_mask = self.high_quality_resize_mask(paste_mask, (actual_crop_w, actual_crop_h))
|
|
|
|
compositing_mask = self.create_compositing_blend_mask(paste_mask, blend_softness)
|
|
restore_info["compositing_mask_used"] = compositing_mask
|
|
compositing_mask = np.clip(compositing_mask, 0.0, 1.0)
|
|
compositing_mask_3d = np.stack([compositing_mask] * 3, axis=-1)
|
|
|
|
pasted_crop = resized_processed_face * compositing_mask_3d + original_crop * (1.0 - compositing_mask_3d)
|
|
|
|
final_result_np[crop_y1:crop_y2, crop_x1:crop_x2] = pasted_crop
|
|
|
|
final_tensor = torch.from_numpy(np.clip(final_result_np, 0.0, 1.0)).unsqueeze(0)
|
|
|
|
return final_tensor
|
|
|
|
except model_management.InterruptProcessingException:
|
|
raise
|
|
except Exception as e:
|
|
print(f"Error in unified restoration: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
return torch.from_numpy(restore_info["original_image"])
|
|
|
|
def _unsharp_mask(self, image, amount=0.5, sigma=0.8, threshold=0.02):
|
|
try:
|
|
if image.size == 0: return image
|
|
blurred = cv2.GaussianBlur(image, (0, 0), sigma)
|
|
detail_mask = image - blurred
|
|
if threshold > 0:
|
|
contrast_mask = np.abs(detail_mask) >= threshold
|
|
detail_mask = np.where(contrast_mask, detail_mask, 0)
|
|
sharpened = image + detail_mask * amount
|
|
return np.clip(sharpened, 0.0, 1.0)
|
|
except model_management.InterruptProcessingException:
|
|
raise
|
|
except Exception as e:
|
|
print(f"Error in unsharp mask: {e}")
|
|
return image
|
|
|
|
def high_quality_resize(self, image_float, target_size):
|
|
try:
|
|
target_w, target_h = target_size
|
|
current_h, current_w = image_float.shape[:2]
|
|
|
|
if target_w <= 0 or target_h <= 0:
|
|
return image_float
|
|
if current_w == target_w and current_h == target_h:
|
|
return image_float
|
|
|
|
if image_float.dtype != np.float32:
|
|
image_float = image_float.astype(np.float32)
|
|
|
|
scale_factor = min(target_w / current_w, target_h / current_h)
|
|
|
|
if scale_factor > 1.0:
|
|
return cv2.resize(image_float, (target_w, target_h), interpolation=cv2.INTER_LANCZOS4)
|
|
elif scale_factor > 0.75:
|
|
return cv2.resize(image_float, (target_w, target_h), interpolation=cv2.INTER_CUBIC)
|
|
else:
|
|
resized = cv2.resize(image_float, (target_w, target_h), interpolation=cv2.INTER_AREA)
|
|
if scale_factor < 0.5:
|
|
amount = 0.3
|
|
sigma = 0.6
|
|
blurred = cv2.GaussianBlur(resized, (0, 0), sigma)
|
|
detail = resized - blurred
|
|
sharpened = resized + detail * amount
|
|
return np.clip(sharpened, 0.0, 1.0)
|
|
return resized
|
|
|
|
except model_management.InterruptProcessingException:
|
|
raise
|
|
except Exception as e:
|
|
print(f"Error in high quality resize: {e}")
|
|
return cv2.resize(image_float, target_size, interpolation=cv2.INTER_AREA)
|
|
|
|
def high_quality_resize_mask(self, mask, target_size):
|
|
try:
|
|
target_w, target_h = target_size
|
|
current_h, current_w = mask.shape[:2]
|
|
scale_factor = min(target_w / current_w, target_h / current_h)
|
|
|
|
interpolation = cv2.INTER_LANCZOS4 if scale_factor > 1.0 else cv2.INTER_LINEAR
|
|
if mask.dtype != np.float32: mask = mask.astype(np.float32)
|
|
resized_mask = cv2.resize(mask, (target_w, target_h), interpolation=interpolation)
|
|
if scale_factor != 1.0: resized_mask = self.sharpen_mask_edges(resized_mask)
|
|
return np.clip(resized_mask, 0.0, 1.0)
|
|
except Exception as e:
|
|
print(f"Error in mask resize: {e}")
|
|
return cv2.resize(mask, target_size, interpolation=cv2.INTER_LINEAR)
|
|
|
|
def sharpen_mask_edges(self, mask, strength=0.3):
|
|
try:
|
|
kernel = np.array([[-1, -1, -1], [-1, 9, -1], [-1, -1, -1]]) * strength
|
|
kernel[1,1] = 1 + strength * 8
|
|
sharpened = cv2.filter2D(mask, -1, kernel)
|
|
return np.clip(sharpened, 0.0, 1.0)
|
|
except Exception as e:
|
|
print(f"Error in mask sharpening: {e}")
|
|
return mask
|
|
|
|
def create_compositing_blend_mask(self, clean_sam_mask, blend_softness):
|
|
"""
|
|
Creates a superior hybrid compositing mask. It preserves the core details of the
|
|
new generation while also expanding the blend outward for seamless, full coverage.
|
|
"""
|
|
try:
|
|
if blend_softness <= 0:
|
|
return (clean_sam_mask > 0.5).astype(np.float32)
|
|
h, w = clean_sam_mask.shape[:2]
|
|
core_mask = (clean_sam_mask > 0.5).astype(np.uint8) * 255
|
|
contours, _ = cv2.findContours(core_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
if not contours:
|
|
return np.zeros_like(clean_sam_mask, dtype=np.float32)
|
|
|
|
largest_contour = max(contours, key=cv2.contourArea)
|
|
_, _, mask_w, mask_h = cv2.boundingRect(largest_contour)
|
|
mask_size = max(mask_w, mask_h, 1)
|
|
blend_ratio = blend_softness / 400.0
|
|
expand_amount = int(mask_size * blend_ratio * 0.8)
|
|
if expand_amount > 0:
|
|
expand_kernel_size = expand_amount * 2 + 1
|
|
expand_kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (expand_kernel_size, expand_kernel_size))
|
|
expanded_for_blur = cv2.dilate(core_mask, expand_kernel, iterations=1)
|
|
else:
|
|
expanded_for_blur = core_mask.copy()
|
|
blur_amount = int(mask_size * blend_ratio * 1.2)
|
|
blur_kernel_size = max(3, blur_amount * 2 + 1)
|
|
|
|
feathering_zone = cv2.GaussianBlur(expanded_for_blur, (blur_kernel_size, blur_kernel_size), 0)
|
|
final_mask_uint8 = np.maximum(core_mask, feathering_zone)
|
|
final_mask_float = final_mask_uint8.astype(np.float32) / 255.0
|
|
fade_pixels = int(min(h, w) * 0.02)
|
|
if fade_pixels > 0:
|
|
for i in range(fade_pixels):
|
|
fade_factor = i / float(fade_pixels)
|
|
final_mask_float[i, :] *= fade_factor
|
|
final_mask_float[h-1-i, :] *= fade_factor
|
|
final_mask_float[:, i] *= fade_factor
|
|
final_mask_float[:, w-1-i] *= fade_factor
|
|
|
|
final_mask_float[0, :], final_mask_float[-1, :], final_mask_float[:, 0], final_mask_float[:, -1] = 0, 0, 0, 0
|
|
return np.clip(final_mask_float, 0.0, 1.0)
|
|
except model_management.InterruptProcessingException:
|
|
raise
|
|
except Exception as e:
|
|
print(f"Error creating hybrid blend mask: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
return (clean_sam_mask > 0.5).astype(np.float32)
|
|
|
|
def fast_tensor_to_numpy(self, tensor, target_range=(0, 255)):
|
|
try:
|
|
if tensor.device != torch.device('cpu'):
|
|
tensor = tensor.cpu()
|
|
|
|
np_array = tensor.squeeze().numpy()
|
|
|
|
if target_range == (0, 255):
|
|
if np_array.max() <= 1.0:
|
|
return (np_array * 255.0).astype(np.uint8)
|
|
else:
|
|
return np.clip(np_array, 0, 255).astype(np.uint8)
|
|
else:
|
|
if np_array.max() > 1.0:
|
|
return np_array / 255.0
|
|
else:
|
|
return np_array
|
|
|
|
except model_management.InterruptProcessingException:
|
|
raise
|
|
except Exception as e:
|
|
print(f"Error in fast tensor conversion: {e}")
|
|
return safe_tensor_to_numpy(tensor, target_range)
|
|
|
|
def fast_numpy_to_tensor(self, np_array, input_range=(0, 255)):
|
|
try:
|
|
if input_range == (0, 255):
|
|
if np_array.dtype == np.uint8:
|
|
tensor_data = np_array.astype(np.float32) / 255.0
|
|
else:
|
|
tensor_data = np.clip(np_array / 255.0, 0.0, 1.0)
|
|
else:
|
|
tensor_data = np.clip(np_array, 0.0, 1.0).astype(np.float32)
|
|
|
|
return torch.from_numpy(tensor_data).unsqueeze(0)
|
|
|
|
except model_management.InterruptProcessingException:
|
|
raise
|
|
except Exception as e:
|
|
print(f"Error in fast numpy conversion: {e}")
|
|
return safe_numpy_to_tensor(np_array, input_range) |