import torch import cv2 import numpy as np from scipy.ndimage import gaussian_filter1d from torchvision.transforms.functional import normalize import os from tqdm import tqdm from .utils import comfy_image_to_cv2, cv2_to_comfy_image from comfy.utils import ProgressBar # Imports from vendored basicsr and facelib try: from basicsr.utils.img_util import img2tensor, tensor2img from basicsr.utils.video_util import VideoReader, VideoWriter # <-- NEW from facelib.utils.misc import is_gray except ImportError as e: print(f"ImportError in keep_processor: {e}. Check vendored 'basicsr' and 'facelib'.") # Dummy just in case, dunno if needed def img2tensor(img, bgr2rgb, float32): return None def tensor2img(tensor, rgb2bgr, min_max): return None def is_gray(img, threshold): return False class VideoReader: def __init__(self, path): self.path = path def get_frame(self): return None def get_fps(self): return 25 def close(self): pass def __len__(self): return 0 class VideoWriter: def __init__(self, path, h, w, fps): pass def write_frame(self, frame): pass def close(self): pass def interpolate_sequence(sequence): interpolated_sequence = np.copy(sequence) missing_indices = np.isnan(sequence) if np.any(missing_indices): valid_indices = ~missing_indices x = np.arange(len(sequence)) interpolated_sequence[missing_indices] = np.interp( x[missing_indices], x[valid_indices], sequence[valid_indices]) return interpolated_sequence class KEEPFaceProcessor: def __init__(self, model_pack): self.keep_net = model_pack.keep_net self.face_helper = model_pack.face_helper self.bg_upsampler = model_pack.bg_upsampler self.face_upsampler = model_pack.face_upsampler self.device = model_pack.device self.model_type_str = model_pack.model_type_str @torch.no_grad() def process_image(self, cv2_image_orig: np.ndarray, has_aligned: bool, only_center_face: bool, draw_box: bool, upscale_factor: int = 1, max_length_override: int = 20): # max_length mostly for consistency self.face_helper.upscale_factor = upscale_factor input_img_list = [cv2_image_orig.copy()] avg_landmarks_list = None if not has_aligned: # print('Detecting keypoints and preparing alignment...') # Use a list to accumulate landmarks from potentially multiple faces per image (though we only use one) raw_landmarks_list_np = [] # For single image, this loop runs once img_for_detection = input_img_list[0] self.face_helper.clean_all() self.face_helper.read_image(img_for_detection) num_det_faces = self.face_helper.get_face_landmarks_5( only_center_face=only_center_face, resize=640, eye_dist_threshold=5, only_keep_largest=True) if num_det_faces >= 1: # Take the first detected face's landmarks raw_landmarks_list_np.append(self.face_helper.all_landmarks_5[0].reshape((10,))) else: print("Warning: No face detected in the input image. Returning original.") return cv2_image_orig # Should not happen if num_det_faces >=1 if not raw_landmarks_list_np: print("Warning: Landmark detection failed. Returning original.") return cv2_image_orig # Shape (num_detected_faces_in_image, 10) raw_landmarks_np_array = np.array(raw_landmarks_list_np) # Iterate over 10 landmark coords for i_lm in range(raw_landmarks_np_array.shape[1]): raw_landmarks_np_array[:, i_lm] = interpolate_sequence(raw_landmarks_np_array[:, i_lm]) # Gaussian filter on a single item sequence doesn't change it much unless sigma is very small or padding is involved # For a single image, we can effectively skip the smoothing by using a small sigma or just reshaping # avg_landmarks_np_array = gaussian_filter1d(raw_landmarks_np_array, sigma=1, axis=0).reshape(raw_landmarks_np_array.shape[0], 5, 2) avg_landmarks_np_array = raw_landmarks_np_array.reshape(raw_landmarks_np_array.shape[0], 5, 2) avg_landmarks_list = [lm for lm in avg_landmarks_np_array] # Prepare Cropped Face Tensors cropped_face_t_list = [] # This outer loop runs once for the single input image img_to_crop = input_img_list[0] if not has_aligned: if avg_landmarks_list is None or not avg_landmarks_list: print("Error: No landmarks available for alignment. Returning original.") return img_to_crop self.face_helper.clean_all() self.face_helper.read_image(img_to_crop) self.face_helper.all_landmarks_5 = [avg_landmarks_list[0]] self.face_helper.align_warp_face() else: # has_aligned = True img_resized = cv2.resize(img_to_crop, (512, 512), interpolation=cv2.INTER_LINEAR) self.face_helper.is_gray = is_gray(img_resized, threshold=10) if self.face_helper.is_gray: print('Grayscale input: True') self.face_helper.cropped_faces = [img_resized] if not self.face_helper.cropped_faces: print("Warning: No cropped face obtained after alignment. Returning original.") return cv2_image_orig cropped_face_cv2 = self.face_helper.cropped_faces[0] cropped_face_t = img2tensor(cropped_face_cv2 / 255., bgr2rgb=True, float32=True) normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True) cropped_face_t_list.append(cropped_face_t) if not cropped_face_t_list: print("No faces to process after tensor conversion. Returning original.") return cv2_image_orig cropped_faces_tensor = torch.stack(cropped_face_t_list, dim=0).unsqueeze(0).to(self.device) # (1, 1, C, H, W) # Restore Faces with KEEP Network # print('Restoring face(s)...') # SINGLE IMAGE ADAPTATION: Duplicate the frame to make a 2-frame "sequence" # Input to net is (batch_size, num_frames, C, H, W) # For a single image, num_frames=1. We duplicate it to make it 2. duplicated_input_frames = torch.cat([cropped_faces_tensor, cropped_faces_tensor], dim=1) # (1, 2, C, H, W) # The KEEP model's forward pass in keep_arch.py takes `x` and `need_upscale` # It handles flows internally. net_output = self.keep_net(duplicated_input_frames, need_upscale=False) # (1, 2, C, H, W) # Take the first frame of the output (as both input frames were identical) final_output_tensor = net_output[:, 0:1, ...].squeeze(0) # (1, C, H, W) -> (C,H,W) if batch size was 1 # Actually (num_cropped_faces, C, H, W) # Squeeze(0) makes it (num_frames_in_output, C, H, W) # Which is (1, C, H, W) restored_faces_cv2 = [tensor2img(x, rgb2bgr=True, min_max=(-1, 1)) for x in final_output_tensor] torch.cuda.empty_cache() # Paste Faces Back # print('Pasting face(s) back...') final_restored_image = None # The original unprocessed input image img_for_pasting = input_img_list[0] self.face_helper.clean_all() if has_aligned: # Input was already aligned, so we don't paste, output is the restored face # The self.face_helper.cropped_faces might be stale here, use the input has_aligned image. img_resized = cv2.resize(img_for_pasting, (512, 512), interpolation=cv2.INTER_LINEAR) self.face_helper.is_gray = is_gray(img_resized, threshold=10) # self.face_helper.cropped_faces = [img_resized] # Storing the *original* cropped face for reference else: # Need to re-align the original image to get affine matrices for pasting back if avg_landmarks_list is None or not avg_landmarks_list: return img_for_pasting self.face_helper.read_image(img_for_pasting) self.face_helper.all_landmarks_5 = [avg_landmarks_list[0]] self.face_helper.align_warp_face() # This sets up affine matrices if not restored_faces_cv2 or not restored_faces_cv2[0].any(): # print(f"Warning: Restored face is missing or invalid. Returning original.") return cv2_image_orig self.face_helper.add_restored_face(restored_faces_cv2[0].astype('uint8')) if not has_aligned: bg_img_upscaled = None if self.bg_upsampler is not None: # print("Upsampling background...") bg_img_upscaled = self.bg_upsampler.enhance(img_for_pasting, outscale=upscale_factor)[0] # Use stored affine_matrices self.face_helper.get_inverse_affine(None) current_face_upsampler = self.face_upsampler if self.face_upsampler is not None else None final_restored_image = self.face_helper.paste_faces_to_input_image( upsample_img=bg_img_upscaled, draw_box=draw_box, face_upsampler=current_face_upsampler ) else: # has_aligned input, output is the restored face itself final_restored_image = self.face_helper.restored_faces[0] if self.face_upsampler and upscale_factor > 1: # RealESRGANer is typically x2, so outscale must match if used # The upscale_factor arg is for background. Face upsampler uses its internal scale. # print("Upsampling aligned face output...") final_restored_image = self.face_upsampler.enhance(final_restored_image, outscale=self.face_upsampler.scale)[0] if final_restored_image is None: return cv2_image_orig return final_restored_image @torch.no_grad() def process_video(self, video_path: str, output_video_path: str, has_aligned: bool, only_center_face: bool, draw_box: bool, upscale_factor: int = 1, max_clip_length: int = 20, target_fps: float = None): vidreader = VideoReader(video_path) original_fps = vidreader.get_fps() video_width = vidreader.width video_height = vidreader.height input_frames_bgr = [] for _ in tqdm(range(len(vidreader))): img = vidreader.get_frame() if img is None: break input_frames_bgr.append(img) vidreader.close() if not input_frames_bgr: # print("No frames read from video.") return None num_input_frames = len(input_frames_bgr) # print(f"Total frames to process: {num_input_frames}") # This is done once for the whole video for temporal consistency avg_landmarks_sequence = None if not has_aligned: # print('Detecting keypoints and smoothing alignment for the whole video...') raw_landmarks_list_video = [] for i in tqdm(range(num_input_frames), desc="Landmark Detection"): img = input_frames_bgr[i] self.face_helper.clean_all() self.face_helper.read_image(img) num_det_faces = self.face_helper.get_face_landmarks_5( only_center_face=only_center_face, resize=640, eye_dist_threshold=5, only_keep_largest=True) if num_det_faces >= 1: raw_landmarks_list_video.append(self.face_helper.all_landmarks_5[0].reshape((10,))) else: # Placeholder for missing faces raw_landmarks_list_video.append(np.array([np.nan] * 10)) # (num_frames, 10) raw_landmarks_np_video = np.array(raw_landmarks_list_video) if np.all(np.isnan(raw_landmarks_np_video)): # print("No faces detected in any frame. Cannot process video.") return None # Iterate over 10 landmark coordinates for i_lm_coord in range(raw_landmarks_np_video.shape[1]): raw_landmarks_np_video[:, i_lm_coord] = interpolate_sequence(raw_landmarks_np_video[:, i_lm_coord]) # Smooth landmarks over time avg_landmarks_sequence = gaussian_filter1d(raw_landmarks_np_video, sigma=5, axis=0).reshape(num_input_frames, 5, 2) print("Landmark smoothing complete.") # Process frames in clips for KEEP network all_restored_faces_cv2 = [] # List to store restored face patches (cv2 format) # print("Cropping faces and preparing for KEEP network...") all_cropped_face_tensors = [] for i in tqdm(range(num_input_frames), desc="Face Cropping"): img_bgr = input_frames_bgr[i] # Clean for each frame self.face_helper.clean_all() self.face_helper.read_image(img_bgr) if not has_aligned: if avg_landmarks_sequence is None or i >= len(avg_landmarks_sequence) or np.isnan(avg_landmarks_sequence[i]).any(): # print(f"Warning: Valid landmarks not available for frame {i}. Using previous frame's alignment or placeholder.") # Fallback: use a placeholder or try to detect individually (less ideal for consistency) # For simplicity here, if landmarks are bad, the cropped face might be bad. if avg_landmarks_sequence is not None and i < len(avg_landmarks_sequence): self.face_helper.all_landmarks_5 = [avg_landmarks_sequence[i]] else: # No landmarks, create a dummy tensor all_cropped_face_tensors.append(torch.zeros(3, 512, 512, dtype=torch.float32, device=self.device) * 0.5) # Normalized to 0 for (0.5,0.5,0.5) continue else: self.face_helper.all_landmarks_5 = [avg_landmarks_sequence[i]] self.face_helper.align_warp_face() else: # has_aligned video input img_resized = cv2.resize(img_bgr, (512, 512), interpolation=cv2.INTER_LINEAR) self.face_helper.is_gray = is_gray(img_resized, threshold=10) # if self.face_helper.is_gray: print(f'Frame {i} is grayscale.') # Can be noisy self.face_helper.cropped_faces = [img_resized] if not self.face_helper.cropped_faces: # Alignment failed or no face # print(f"Warning: No cropped face for frame {i}. Using placeholder.") all_cropped_face_tensors.append(torch.zeros(3, 512, 512, dtype=torch.float32, device=self.device) * 0.5) continue cropped_face_cv2 = self.face_helper.cropped_faces[0] cropped_face_t = img2tensor(cropped_face_cv2 / 255., bgr2rgb=True, float32=True) normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True) all_cropped_face_tensors.append(cropped_face_t) # Stack all cropped faces: (1, num_input_frames, C, H, W) if not all_cropped_face_tensors: # print("No faces could be cropped from the video.") return None batched_cropped_faces = torch.stack(all_cropped_face_tensors, dim=0).unsqueeze(0).to(self.device) print("Restoring faces with KEEP network (in clips)...") temp_restored_face_tensors = [] # Tensors from network for start_idx in tqdm(range(0, num_input_frames, max_clip_length), desc="KEEP Processing Clips"): end_idx = min(start_idx + max_clip_length, num_input_frames) current_clip_faces = batched_cropped_faces[:, start_idx:end_idx, ...] # Single frame in clip, duplicate for KEEP if current_clip_faces.shape[1] == 1: current_clip_faces_for_net = torch.cat([current_clip_faces, current_clip_faces], dim=1) clip_output_tensor = self.keep_net(current_clip_faces_for_net, need_upscale=False) # Take first frame temp_restored_face_tensors.append(clip_output_tensor[:, 0:1, ...]) elif current_clip_faces.shape[1] > 1 : clip_output_tensor = self.keep_net(current_clip_faces, need_upscale=False) temp_restored_face_tensors.append(clip_output_tensor) # If current_clip_faces.shape[1] is 0, skip (should not happen with good indexing) if not temp_restored_face_tensors: # print("KEEP network processing yielded no results.") return None # (num_frames, C, H, W) all_restored_face_tensors = torch.cat(temp_restored_face_tensors, dim=1).squeeze(0) all_restored_faces_cv2 = [tensor2img(t, rgb2bgr=True, min_max=(-1, 1)) for t in all_restored_face_tensors] del batched_cropped_faces, all_restored_face_tensors, temp_restored_face_tensors torch.cuda.empty_cache() # print("KEEP network processing complete.") # Paste faces back and write video final_fps = target_fps if target_fps is not None else original_fps # Use dimensions of the first frame after potential BG upscaling for VideoWriter # Perform one BG upscale to determine output size if bg_upsampler is used temp_first_frame_for_size = input_frames_bgr[0] if self.bg_upsampler is not None: temp_first_frame_for_size = self.bg_upsampler.enhance(temp_first_frame_for_size, outscale=upscale_factor)[0] output_video_height, output_video_width = temp_first_frame_for_size.shape[0:2] del temp_first_frame_for_size vidwriter = VideoWriter(output_video_path, output_video_height, output_video_width, final_fps) # print("Pasting faces back and writing video...") for i in tqdm(range(num_input_frames), desc="Pasting & Writing Video"): original_frame_bgr = input_frames_bgr[i] # Should not happen if lists are aligned if i >= len(all_restored_faces_cv2): # print(f"Warning: Missing restored face for frame {i}. Using original frame.") # Or a black frame etc. final_frame_to_write = original_frame_bgr # Still upscale background if requested if self.bg_upsampler: final_frame_to_write = self.bg_upsampler.enhance(final_frame_to_write, outscale=upscale_factor)[0] if final_frame_to_write.shape[0] != output_video_height or final_frame_to_write.shape[1] != output_video_width: final_frame_to_write = cv2.resize(final_frame_to_write, (output_video_width, output_video_height)) vidwriter.write_frame(final_frame_to_write) continue restored_face_patch_cv2 = all_restored_faces_cv2[i] self.face_helper.clean_all() self.face_helper.read_image(original_frame_bgr) # Set current frame in helper # Apply face_upsampler to the restored patch if needed current_restored_face_for_paste = restored_face_patch_cv2.astype('uint8') if self.face_upsampler: # If main background isn't upscaled (upscale_factor == 1), resize face back down after upsampling # to maintain geometric consistency, but with enhanced quality. if upscale_factor == 1 and getattr(self.face_upsampler, 'scale', 1) > 1: effective_face_upsampler_scale = getattr(self.face_upsampler, 'scale', 2) upscaled_face_content = self.face_upsampler.enhance(current_restored_face_for_paste, outscale=effective_face_upsampler_scale)[0] target_h_for_warp, target_w_for_warp = self.face_size[1], self.face_size[0] # H, W current_restored_face_for_paste = cv2.resize( upscaled_face_content, (target_w_for_warp, target_h_for_warp), interpolation=cv2.INTER_LANCZOS4 ) else: # Background is also upscaled or no specific resize-down needed effective_face_upsampler_scale = getattr(self.face_upsampler, 'scale', 2) # Default to 2 current_restored_face_for_paste = self.face_upsampler.enhance(current_restored_face_for_paste, outscale=effective_face_upsampler_scale)[0] self.face_helper.add_restored_face(current_restored_face_for_paste) if not has_aligned: if avg_landmarks_sequence is None or i >= len(avg_landmarks_sequence) or np.isnan(avg_landmarks_sequence[i]).any(): # print(f"Warning: Valid landmarks not available for pasting frame {i}. Using original frame content.") final_frame_to_write = original_frame_bgr if self.bg_upsampler: final_frame_to_write = self.bg_upsampler.enhance(final_frame_to_write, outscale=upscale_factor)[0] if final_frame_to_write.shape[0] != output_video_height or final_frame_to_write.shape[1] != output_video_width: final_frame_to_write = cv2.resize(final_frame_to_write, (output_video_width, output_video_height)) vidwriter.write_frame(final_frame_to_write) continue self.face_helper.all_landmarks_5 = [avg_landmarks_sequence[i]] # Set current frame's landmarks self.face_helper.align_warp_face() # This is only to re-init affine matrix for this frame bg_img_upscaled_final = None if self.bg_upsampler is not None: bg_img_upscaled_final = self.bg_upsampler.enhance(original_frame_bgr, outscale=upscale_factor)[0] self.face_helper.get_inverse_affine(None) pasted_frame = self.face_helper.paste_faces_to_input_image( upsample_img=bg_img_upscaled_final, # This is the target canvas draw_box=draw_box, face_upsampler=None # Already handled current_restored_face_for_paste ) else: # Video has_aligned (input frames are 512x512 aligned faces) # The output is just the restored (and potentially face-upsampled) face itself # Which is current_restored_face_for_paste pasted_frame = self.face_helper.restored_faces[0] # Ensure it's BGR if len(pasted_frame.shape) == 2: pasted_frame = cv2.cvtColor(pasted_frame, cv2.COLOR_GRAY2BGR) # Ensure frame is correct size for video writer if pasted_frame.shape[0] != output_video_height or pasted_frame.shape[1] != output_video_width: pasted_frame = cv2.resize(pasted_frame, (output_video_width, output_video_height), interpolation=cv2.INTER_LANCZOS4) vidwriter.write_frame(pasted_frame) vidwriter.close() print(f"Processed video saved to {output_video_path}") return output_video_path @torch.no_grad() def process_image_sequence(self, image_sequence_tensor: torch.Tensor, # Input: B, H, W, C (ComfyUI IMAGE) has_aligned_frames: bool, only_center_face: bool, draw_box: bool, upscale_factor: int = 1, # For background max_clip_length: int = 20): num_input_frames = image_sequence_tensor.shape[0] if num_input_frames == 0: return image_sequence_tensor # Return empty if input is empty # Convert ComfyUI IMAGE tensor to a list of cv2 BGR images input_frames_bgr = [] pbar = ProgressBar(num_input_frames) for i in range(num_input_frames): # comfy_image_to_cv2 expects (1,H,W,C) or (H,W,C) frame_tensor = image_sequence_tensor[i].unsqueeze(0) input_frames_bgr.append(comfy_image_to_cv2(frame_tensor)) pbar.update(1) # print(f"Processing image sequence: {num_input_frames} frames.") # Landmark Detection and Smoothing (if not has_aligned) avg_landmarks_sequence = None if not has_aligned_frames: # print('Detecting keypoints and smoothing alignment for the sequence...') raw_landmarks_list_video = [] for i in tqdm(range(num_input_frames), desc="Landmark Detection (Sequence)"): img = input_frames_bgr[i] self.face_helper.clean_all() self.face_helper.read_image(img) num_det_faces = self.face_helper.get_face_landmarks_5( only_center_face=only_center_face, resize=640, eye_dist_threshold=5, only_keep_largest=True) if num_det_faces >= 1: raw_landmarks_list_video.append(self.face_helper.all_landmarks_5[0].reshape((10,))) else: raw_landmarks_list_video.append(np.array([np.nan] * 10)) raw_landmarks_np_video = np.array(raw_landmarks_list_video) if np.all(np.isnan(raw_landmarks_np_video)): print("No faces detected in any frame. Returning original sequence.") return image_sequence_tensor for i_lm_coord in range(raw_landmarks_np_video.shape[1]): raw_landmarks_np_video[:, i_lm_coord] = interpolate_sequence(raw_landmarks_np_video[:, i_lm_coord]) avg_landmarks_sequence = gaussian_filter1d(raw_landmarks_np_video, sigma=5, axis=0).reshape(num_input_frames, 5, 2) # print("Landmark smoothing complete.") # Crop faces and prepare for KEEP network all_cropped_face_tensors = [] # print("Cropping faces for KEEP network...") pbar = ProgressBar(num_input_frames) for i in tqdm(range(num_input_frames), desc="Face Cropping (Sequence)"): img_bgr = input_frames_bgr[i] self.face_helper.clean_all() self.face_helper.read_image(img_bgr) if not has_aligned_frames: if avg_landmarks_sequence is None or i >= len(avg_landmarks_sequence) or np.isnan(avg_landmarks_sequence[i]).any(): # Fallback for missing landmarks in a frame print(f"Warning: Valid landmarks not available for frame {i} in sequence. Using placeholder crop.") all_cropped_face_tensors.append(torch.zeros(3, 512, 512, dtype=torch.float32, device=self.device) * 0.5) continue self.face_helper.all_landmarks_5 = [avg_landmarks_sequence[i]] self.face_helper.align_warp_face() else: # has_aligned_frames = True img_resized = cv2.resize(img_bgr, (512, 512), interpolation=cv2.INTER_LINEAR) self.face_helper.is_gray = is_gray(img_resized, threshold=10) self.face_helper.cropped_faces = [img_resized] if not self.face_helper.cropped_faces: print(f"Warning: No cropped face for frame {i}. Using placeholder crop.") all_cropped_face_tensors.append(torch.zeros(3, 512, 512, dtype=torch.float32, device=self.device) * 0.5) continue cropped_face_cv2 = self.face_helper.cropped_faces[0] cropped_face_t = img2tensor(cropped_face_cv2 / 255., bgr2rgb=True, float32=True) normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True) all_cropped_face_tensors.append(cropped_face_t) pbar.update(1) if not all_cropped_face_tensors: print("No faces could be cropped from the image sequence.") return image_sequence_tensor # Return original # Batch for KEEP network: (1, num_input_frames, C, H, W) batched_cropped_faces = torch.stack(all_cropped_face_tensors, dim=0).unsqueeze(0).to(self.device) # KEEP Network Processing (in clips) # print("Restoring faces with KEEP network (in clips)...") temp_restored_face_tensors = [] pbar = ProgressBar(max_clip_length) for start_idx in tqdm(range(0, num_input_frames, max_clip_length), desc="KEEP Processing Clips (Sequence)"): end_idx = min(start_idx + max_clip_length, num_input_frames) current_clip_faces = batched_cropped_faces[:, start_idx:end_idx, ...] if current_clip_faces.shape[1] == 0: continue # Skip empty clips if current_clip_faces.shape[1] == 1: current_clip_faces_for_net = torch.cat([current_clip_faces, current_clip_faces], dim=1) clip_output_tensor = self.keep_net(current_clip_faces_for_net, need_upscale=False) temp_restored_face_tensors.append(clip_output_tensor[:, 0:1, ...]) else: clip_output_tensor = self.keep_net(current_clip_faces, need_upscale=False) temp_restored_face_tensors.append(clip_output_tensor) pbar.update(1) if not temp_restored_face_tensors: print("KEEP network processing yielded no results for the sequence.") return image_sequence_tensor all_restored_face_tensors = torch.cat(temp_restored_face_tensors, dim=1).squeeze(0) all_restored_faces_cv2 = [tensor2img(t, rgb2bgr=True, min_max=(-1, 1)) for t in all_restored_face_tensors] del batched_cropped_faces, all_restored_face_tensors, temp_restored_face_tensors torch.cuda.empty_cache() print("KEEP network processing complete for sequence.") # Paste faces back onto each frame output_frames_cv2 = [] # print("Pasting faces back onto frames...") pbar = ProgressBar(num_input_frames) for i in tqdm(range(num_input_frames), desc="Pasting Faces (Sequence)"): original_frame_bgr = input_frames_bgr[i] if i >= len(all_restored_faces_cv2): print(f"Warning: Missing restored face for frame {i}. Using original frame content.") final_pasted_frame = original_frame_bgr if self.bg_upsampler: final_pasted_frame = self.bg_upsampler.enhance(final_pasted_frame, outscale=upscale_factor)[0] output_frames_cv2.append(final_pasted_frame) continue restored_face_patch_cv2 = all_restored_faces_cv2[i] self.face_helper.clean_all() self.face_helper.read_image(original_frame_bgr) current_restored_face_for_paste = restored_face_patch_cv2.astype('uint8') if self.face_upsampler: if upscale_factor == 1 and getattr(self.face_upsampler, 'scale', 1) > 1: effective_face_upsampler_scale = getattr(self.face_upsampler, 'scale', 2) upscaled_face_content = self.face_upsampler.enhance(current_restored_face_for_paste, outscale=effective_face_upsampler_scale)[0] # ------ some fix here # self.face_size is actually self.face_helper.face_size # self.face_helper.face_size is (width, height) target_w_for_warp, target_h_for_warp = self.face_helper.face_size[0], self.face_helper.face_size[1] # ------ end of fix current_restored_face_for_paste = cv2.resize( upscaled_face_content, (target_w_for_warp, target_h_for_warp), interpolation=cv2.INTER_LANCZOS4) else: effective_face_upsampler_scale = getattr(self.face_upsampler, 'scale', 2) current_restored_face_for_paste = self.face_upsampler.enhance(current_restored_face_for_paste, outscale=effective_face_upsampler_scale)[0] self.face_helper.add_restored_face(current_restored_face_for_paste) if not has_aligned_frames: if avg_landmarks_sequence is None or i >= len(avg_landmarks_sequence) or np.isnan(avg_landmarks_sequence[i]).any(): print(f"Warning: Valid landmarks not available for pasting frame {i}. Using original content.") final_pasted_frame = original_frame_bgr if self.bg_upsampler: final_pasted_frame = self.bg_upsampler.enhance(final_pasted_frame, outscale=upscale_factor)[0] output_frames_cv2.append(final_pasted_frame) continue self.face_helper.all_landmarks_5 = [avg_landmarks_sequence[i]] self.face_helper.align_warp_face() bg_img_upscaled_final = None if self.bg_upsampler is not None: bg_img_upscaled_final = self.bg_upsampler.enhance(original_frame_bgr, outscale=upscale_factor)[0] self.face_helper.get_inverse_affine(None) final_pasted_frame = self.face_helper.paste_faces_to_input_image( upsample_img=bg_img_upscaled_final, draw_box=draw_box, face_upsampler=None ) else: # has_aligned_frames = True final_pasted_frame = self.face_helper.restored_faces[0] if len(final_pasted_frame.shape) == 2: final_pasted_frame = cv2.cvtColor(final_pasted_frame, cv2.COLOR_GRAY2BGR) output_frames_cv2.append(final_pasted_frame) pbar.update(1) # Convert list of cv2 BGR images back to ComfyUI IMAGE tensor (B, H, W, C), RGB, float32 [0,1] output_tensors = [cv2_to_comfy_image(frame) for frame in output_frames_cv2] if not output_tensors: return image_sequence_tensor final_batch_tensor = torch.cat(output_tensors, dim=0) # print("Video frame processing complete.") return final_batch_tensor