460 lines
24 KiB
Python
460 lines
24 KiB
Python
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 # Import tqdm for console progress bars
|
|
from .utils import comfy_image_to_cv2, cv2_to_comfy_image
|
|
from comfy.utils import ProgressBar, tiled_scale
|
|
|
|
# Imports from vendored wm_basicsr and facelib
|
|
try:
|
|
from wm_basicsr.utils.img_util import img2tensor, tensor2img
|
|
from wm_basicsr.utils.video_util import VideoReader, VideoWriter
|
|
from wm_facelib.utils.misc import is_gray
|
|
except ImportError as e:
|
|
print(f"ImportError in keep_processor: {e}. Check vendored 'wm_basicsr' and 'facelib'.")
|
|
# Dummy just in case
|
|
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 __len__(self): return 0
|
|
def get_frame(self): return None
|
|
def get_fps(self): return 25
|
|
def close(self): pass
|
|
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_upscale_model = model_pack.bg_upscale_model
|
|
self.face_upscale_model = model_pack.face_upscale_model
|
|
self.device = model_pack.device
|
|
self.model_type_str = model_pack.model_type_str
|
|
|
|
def _run_upscaler(self, model, cv2_image):
|
|
"""Helper to run a comfy-native upscaler on a cv2 image."""
|
|
if model is None:
|
|
return cv2_image
|
|
|
|
img_tensor = cv2_to_comfy_image(cv2_image).to(self.device)
|
|
img_tensor_bchw = img_tensor.movedim(-1, -3)
|
|
s = tiled_scale(img_tensor_bchw, lambda a: model.model(a), tile_x=512, tile_y=512, overlap=64, upscale_amount=model.scale)
|
|
s = torch.clamp(s.movedim(-3, -1), min=0, max=1.0)
|
|
return comfy_image_to_cv2(s)
|
|
|
|
@torch.no_grad()
|
|
def process_image(self, cv2_image_orig: np.ndarray, final_upscale_factor: float, has_aligned: bool,
|
|
only_center_face: bool, draw_box: bool):
|
|
|
|
self.face_helper.upscale_factor = final_upscale_factor
|
|
|
|
input_img_list = [cv2_image_orig.copy()]
|
|
avg_landmarks_list = None
|
|
|
|
if not has_aligned:
|
|
raw_landmarks_list_np = []
|
|
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:
|
|
raw_landmarks_list_np.append(self.face_helper.all_landmarks_5[0].reshape((10,)))
|
|
else: return cv2_image_orig
|
|
if not raw_landmarks_list_np: return cv2_image_orig
|
|
raw_landmarks_np_array = np.array(raw_landmarks_list_np)
|
|
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])
|
|
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]
|
|
|
|
cropped_face_t_list = []
|
|
img_to_crop = input_img_list[0]
|
|
if not has_aligned:
|
|
if avg_landmarks_list is None or not avg_landmarks_list: 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:
|
|
img_resized = cv2.resize(img_to_crop, (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: 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)
|
|
cropped_faces_tensor = torch.stack(cropped_face_t_list, dim=0).unsqueeze(0).to(self.device)
|
|
duplicated_input_frames = torch.cat([cropped_faces_tensor, cropped_faces_tensor], dim=1)
|
|
net_output = self.keep_net(duplicated_input_frames, need_upscale=False)
|
|
final_output_tensor = net_output[:, 0:1, ...].squeeze(0)
|
|
restored_faces_cv2 = [tensor2img(x, rgb2bgr=True, min_max=(-1, 1)) for x in final_output_tensor]
|
|
torch.cuda.empty_cache()
|
|
|
|
img_for_pasting = input_img_list[0]
|
|
self.face_helper.clean_all()
|
|
if has_aligned:
|
|
img_resized = cv2.resize(img_for_pasting, (512, 512), interpolation=cv2.INTER_LINEAR)
|
|
self.face_helper.is_gray = is_gray(img_resized, threshold=10)
|
|
else:
|
|
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()
|
|
|
|
if not restored_faces_cv2 or not restored_faces_cv2[0].any(): return cv2_image_orig
|
|
|
|
self.face_helper.add_restored_face(restored_faces_cv2[0].astype('uint8'))
|
|
|
|
if not has_aligned:
|
|
bg_img_upscaled_by_model = self._run_upscaler(self.bg_upscale_model, img_for_pasting)
|
|
|
|
h, w, _ = img_for_pasting.shape
|
|
target_h, target_w = int(h * final_upscale_factor), int(w * final_upscale_factor)
|
|
if bg_img_upscaled_by_model.shape[0] != target_h or bg_img_upscaled_by_model.shape[1] != target_w:
|
|
bg_img_final = cv2.resize(bg_img_upscaled_by_model, (target_w, target_h), interpolation=cv2.INTER_LANCZOS4)
|
|
else:
|
|
bg_img_final = bg_img_upscaled_by_model
|
|
|
|
self.face_helper.get_inverse_affine(None)
|
|
final_restored_image = self.face_helper.paste_faces_to_input_image(
|
|
upsample_img=bg_img_final,
|
|
draw_box=draw_box,
|
|
face_upsampler=self.face_upscale_model
|
|
)
|
|
else:
|
|
final_restored_image = self.face_helper.restored_faces[0]
|
|
if self.face_upscale_model:
|
|
final_restored_image = self._run_upscaler(self.face_upscale_model, final_restored_image)
|
|
|
|
h, w, _ = final_restored_image.shape
|
|
target_h, target_w = int(512 * final_upscale_factor), int(512 * final_upscale_factor)
|
|
if h != target_h or w != target_w:
|
|
final_restored_image = cv2.resize(final_restored_image, (target_w, target_h), interpolation=cv2.INTER_LANCZOS4)
|
|
|
|
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,
|
|
final_upscale_factor: float, has_aligned: bool, only_center_face: bool, draw_box: bool,
|
|
max_clip_length: int = 20, target_fps: float = None):
|
|
|
|
vidreader = VideoReader(video_path)
|
|
original_fps = vidreader.get_fps()
|
|
num_input_frames = len(vidreader)
|
|
|
|
pbar_comfy = ProgressBar(num_input_frames * 5)
|
|
|
|
input_frames_bgr = []
|
|
for _ in tqdm(range(num_input_frames), desc="Reading video frames"):
|
|
img = vidreader.get_frame()
|
|
if img is None: break
|
|
input_frames_bgr.append(img)
|
|
pbar_comfy.update(1)
|
|
vidreader.close()
|
|
|
|
if not input_frames_bgr: return None
|
|
|
|
avg_landmarks_sequence = None
|
|
if not has_aligned:
|
|
raw_landmarks_list_video = []
|
|
for i in tqdm(range(num_input_frames), desc="Detecting face landmarks"):
|
|
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))
|
|
pbar_comfy.update(1)
|
|
raw_landmarks_np_video = np.array(raw_landmarks_list_video)
|
|
if np.all(np.isnan(raw_landmarks_np_video)): return None
|
|
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)
|
|
else:
|
|
pbar_comfy.update(num_input_frames)
|
|
|
|
all_cropped_face_tensors = []
|
|
for i in tqdm(range(num_input_frames), desc="Cropping and aligning faces"):
|
|
img_bgr = input_frames_bgr[i]
|
|
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():
|
|
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:
|
|
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:
|
|
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_comfy.update(1)
|
|
|
|
if not all_cropped_face_tensors: return None
|
|
|
|
batched_cropped_faces = torch.stack(all_cropped_face_tensors, dim=0).unsqueeze(0).to(self.device)
|
|
temp_restored_face_tensors = []
|
|
for start_idx in tqdm(range(0, num_input_frames, max_clip_length), desc="Restoring faces with KEEP"):
|
|
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] == 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, ...])
|
|
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)
|
|
pbar_comfy.update(end_idx - start_idx)
|
|
|
|
if not temp_restored_face_tensors: return None
|
|
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()
|
|
|
|
final_fps = target_fps if target_fps is not None else original_fps
|
|
h_orig, w_orig, _ = input_frames_bgr[0].shape
|
|
output_video_height, output_video_width = int(h_orig * final_upscale_factor), int(w_orig * final_upscale_factor)
|
|
|
|
vidwriter = VideoWriter(output_video_path, output_video_height, output_video_width, final_fps)
|
|
|
|
for i in tqdm(range(num_input_frames), desc="Pasting faces and writing video"):
|
|
original_frame_bgr = input_frames_bgr[i]
|
|
if i >= len(all_restored_faces_cv2):
|
|
final_frame_to_write = self._run_upscaler(self.bg_upscale_model, original_frame_bgr)
|
|
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)
|
|
self.face_helper.add_restored_face(restored_face_patch_cv2.astype('uint8'))
|
|
|
|
if not has_aligned:
|
|
if avg_landmarks_sequence is None or i >= len(avg_landmarks_sequence) or np.isnan(avg_landmarks_sequence[i]).any():
|
|
final_frame_to_write = self._run_upscaler(self.bg_upscale_model, original_frame_bgr)
|
|
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.upscale_factor = final_upscale_factor
|
|
self.face_helper.all_landmarks_5 = [avg_landmarks_sequence[i]]
|
|
self.face_helper.align_warp_face()
|
|
|
|
bg_img_upscaled_by_model = self._run_upscaler(self.bg_upscale_model, original_frame_bgr)
|
|
if bg_img_upscaled_by_model.shape[0] != output_video_height or bg_img_upscaled_by_model.shape[1] != output_video_width:
|
|
bg_img_final = cv2.resize(bg_img_upscaled_by_model, (output_video_width, output_video_height), interpolation=cv2.INTER_LANCZOS4)
|
|
else:
|
|
bg_img_final = bg_img_upscaled_by_model
|
|
|
|
self.face_helper.get_inverse_affine(None)
|
|
pasted_frame = self.face_helper.paste_faces_to_input_image(
|
|
upsample_img=bg_img_final, draw_box=draw_box, face_upsampler=self.face_upscale_model)
|
|
else:
|
|
pasted_frame = self.face_helper.restored_faces[0]
|
|
if self.face_upscale_model:
|
|
pasted_frame = self._run_upscaler(self.face_upscale_model, pasted_frame)
|
|
if len(pasted_frame.shape) == 2: pasted_frame = cv2.cvtColor(pasted_frame, cv2.COLOR_GRAY2BGR)
|
|
|
|
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)
|
|
pbar_comfy.update(1)
|
|
|
|
vidwriter.close()
|
|
return output_video_path
|
|
|
|
@torch.no_grad()
|
|
def process_image_sequence(self,
|
|
image_sequence_tensor: torch.Tensor,
|
|
final_upscale_factor: float,
|
|
has_aligned_frames: bool,
|
|
only_center_face: bool,
|
|
draw_box: bool,
|
|
max_clip_length: int = 20):
|
|
|
|
num_input_frames = image_sequence_tensor.shape[0]
|
|
if num_input_frames == 0:
|
|
return image_sequence_tensor
|
|
|
|
pbar_comfy = ProgressBar(num_input_frames * 4)
|
|
|
|
input_frames_bgr = []
|
|
for i in range(num_input_frames):
|
|
frame_tensor = image_sequence_tensor[i].unsqueeze(0)
|
|
input_frames_bgr.append(comfy_image_to_cv2(frame_tensor))
|
|
pbar_comfy.update(1)
|
|
|
|
avg_landmarks_sequence = None
|
|
if not has_aligned_frames:
|
|
raw_landmarks_list_video = []
|
|
for i in tqdm(range(num_input_frames), desc="Detecting face landmarks"):
|
|
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)):
|
|
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)
|
|
pbar_comfy.update(num_input_frames)
|
|
else:
|
|
pbar_comfy.update(num_input_frames)
|
|
|
|
all_cropped_face_tensors = []
|
|
for i in tqdm(range(num_input_frames), desc="Cropping and aligning faces"):
|
|
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():
|
|
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:
|
|
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:
|
|
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)
|
|
|
|
if not all_cropped_face_tensors:
|
|
return image_sequence_tensor
|
|
|
|
batched_cropped_faces = torch.stack(all_cropped_face_tensors, dim=0).unsqueeze(0).to(self.device)
|
|
temp_restored_face_tensors = []
|
|
for start_idx in tqdm(range(0, num_input_frames, max_clip_length), desc="Restoring faces with KEEP"):
|
|
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
|
|
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_comfy.update(end_idx - start_idx)
|
|
|
|
if not temp_restored_face_tensors:
|
|
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()
|
|
|
|
output_frames_cv2 = []
|
|
for i in tqdm(range(num_input_frames), desc="Pasting faces and finalizing frames"):
|
|
original_frame_bgr = input_frames_bgr[i]
|
|
if i >= len(all_restored_faces_cv2):
|
|
final_pasted_frame = self._run_upscaler(self.bg_upscale_model, original_frame_bgr)
|
|
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)
|
|
self.face_helper.add_restored_face(restored_face_patch_cv2.astype('uint8'))
|
|
|
|
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():
|
|
final_pasted_frame = self._run_upscaler(self.bg_upscale_model, original_frame_bgr)
|
|
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_by_model = self._run_upscaler(self.bg_upscale_model, original_frame_bgr)
|
|
h, w, _ = original_frame_bgr.shape
|
|
target_h, target_w = int(h * final_upscale_factor), int(w * final_upscale_factor)
|
|
if bg_img_upscaled_by_model.shape[0] != target_h or bg_img_upscaled_by_model.shape[1] != target_w:
|
|
bg_img_final = cv2.resize(bg_img_upscaled_by_model, (target_w, target_h), interpolation=cv2.INTER_LANCZOS4)
|
|
else:
|
|
bg_img_final = bg_img_upscaled_by_model
|
|
|
|
self.face_helper.upscale_factor = final_upscale_factor
|
|
self.face_helper.get_inverse_affine(None)
|
|
final_pasted_frame = self.face_helper.paste_faces_to_input_image(
|
|
upsample_img=bg_img_final,
|
|
draw_box=draw_box,
|
|
face_upsampler=self.face_upscale_model
|
|
)
|
|
else:
|
|
restored_face_with_model = self.face_helper.restored_faces[0]
|
|
if self.face_upscale_model:
|
|
restored_face_with_model = self._run_upscaler(self.face_upscale_model, restored_face_with_model)
|
|
|
|
h_restored, w_restored, _ = restored_face_with_model.shape
|
|
target_h, target_w = int(512 * final_upscale_factor), int(512 * final_upscale_factor)
|
|
if h_restored != target_h or w_restored != target_w:
|
|
final_pasted_frame = cv2.resize(restored_face_with_model, (target_w, target_h), interpolation=cv2.INTER_LANCZOS4)
|
|
else:
|
|
final_pasted_frame = restored_face_with_model
|
|
|
|
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_comfy.update(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)
|
|
return final_batch_tensor |