221 lines
6.4 KiB
Python
221 lines
6.4 KiB
Python
import numpy as np
|
|
import torch
|
|
import yaml
|
|
from comfy.utils import ProgressBar
|
|
from scipy.spatial import ConvexHull
|
|
from tqdm import tqdm
|
|
|
|
from .modules.generator import OcclusionAwareGenerator
|
|
from .modules.keypoint_detector import KPDetector
|
|
from .utils import tensor_to_numpy
|
|
|
|
try:
|
|
import face_alignment
|
|
|
|
face_alignment_available = True
|
|
except ImportError:
|
|
face_alignment_available = False
|
|
print("Warning: face_alignment module not found. fine_best_frame will not work. ")
|
|
|
|
|
|
def normalize_kp_best_frame(kp):
|
|
kp = kp - kp.mean(axis=0, keepdims=True)
|
|
area = ConvexHull(kp[:, :2]).volume
|
|
area = np.sqrt(area)
|
|
kp[:, :2] = kp[:, :2] / area
|
|
return kp
|
|
|
|
|
|
def normalize_kp(
|
|
kp_source,
|
|
kp_driving,
|
|
kp_driving_initial,
|
|
adapt_movement_scale: bool,
|
|
use_relative_movement: bool,
|
|
use_relative_jacobian: bool,
|
|
):
|
|
if adapt_movement_scale:
|
|
source_area = ConvexHull(kp_source["value"][0].data.cpu().numpy()).volume
|
|
driving_area = ConvexHull(
|
|
kp_driving_initial["value"][0].data.cpu().numpy()
|
|
).volume
|
|
adapt_movement_scale = np.sqrt(source_area) / np.sqrt(driving_area)
|
|
else:
|
|
adapt_movement_scale = 1
|
|
kp_new = {k: v for k, v in kp_driving.items()}
|
|
|
|
if use_relative_movement:
|
|
kp_value_diff = kp_driving["value"] - kp_driving_initial["value"]
|
|
kp_value_diff *= adapt_movement_scale
|
|
kp_new["value"] = kp_value_diff + kp_source["value"]
|
|
|
|
if use_relative_jacobian:
|
|
jacobian_diff = torch.matmul(
|
|
kp_driving["jacobian"], torch.inverse(kp_driving_initial["jacobian"])
|
|
)
|
|
kp_new["jacobian"] = torch.matmul(jacobian_diff, kp_source["jacobian"])
|
|
|
|
return kp_new
|
|
|
|
|
|
def find_best_frame(source_image, driving_video: list, cpu=False) -> int:
|
|
if not face_alignment_available:
|
|
print("Face Alignment module were not found. Best frame set to 0.")
|
|
return 0
|
|
|
|
print("Initializing Face Alignment...")
|
|
fa = face_alignment.FaceAlignment(
|
|
face_alignment.LandmarksType.TWO_D,
|
|
flip_input=True,
|
|
device="cpu" if cpu else "cuda",
|
|
)
|
|
print("Face alignment model loaded.")
|
|
|
|
source_np = tensor_to_numpy(source_image)
|
|
kp_source = fa.get_landmarks(255 * source_np)[0]
|
|
kp_source = normalize_kp_best_frame(kp_source)
|
|
|
|
norm = float("inf")
|
|
frame_num = 0
|
|
|
|
num_frames = driving_video.shape[2]
|
|
|
|
for i in tqdm(range(num_frames)):
|
|
# Extract single frame from driving video
|
|
frame = driving_video[0, :, i] # Shape: [3, 256, 256]
|
|
frame_np = frame.permute(1, 2, 0).cpu().numpy()
|
|
|
|
kp_driving = fa.get_landmarks(255 * frame_np)[0]
|
|
kp_driving = normalize_kp_best_frame(kp_driving)
|
|
new_norm = (np.abs(kp_source - kp_driving) ** 2).sum()
|
|
if new_norm < norm:
|
|
norm = new_norm
|
|
frame_num = i
|
|
|
|
print(f"Best frame: {frame_num}")
|
|
|
|
return frame_num
|
|
|
|
|
|
def inference(
|
|
source_image,
|
|
driving_video: list,
|
|
generator: OcclusionAwareGenerator,
|
|
kp_detector: KPDetector,
|
|
relative_movement: bool,
|
|
relative_jacobian: bool,
|
|
adapt_movement_scale: bool,
|
|
cpu=False,
|
|
) -> list:
|
|
with torch.no_grad():
|
|
source = source_image
|
|
driving = driving_video
|
|
predictions = []
|
|
|
|
if not cpu:
|
|
source = source.cuda()
|
|
driving = driving.cuda()
|
|
print("Using GPU")
|
|
|
|
kp_source = kp_detector(source)
|
|
kp_driving_initial = kp_detector(driving[:, :, 0])
|
|
kp_norm = None
|
|
|
|
num_frames = driving.shape[2]
|
|
pbar = ProgressBar(num_frames)
|
|
|
|
for frame_idx in tqdm(range(num_frames)):
|
|
driving_frame = driving[:, :, frame_idx]
|
|
kp_driving = kp_detector(driving_frame)
|
|
kp_norm = normalize_kp(
|
|
kp_source=kp_source,
|
|
kp_driving=kp_driving,
|
|
kp_driving_initial=kp_driving_initial,
|
|
use_relative_movement=relative_movement,
|
|
use_relative_jacobian=relative_jacobian,
|
|
adapt_movement_scale=adapt_movement_scale,
|
|
)
|
|
out = generator(source, kp_source=kp_source, kp_driving=kp_norm)
|
|
out_np = np.transpose(out["prediction"].data.cpu().numpy(), [0, 2, 3, 1])[0]
|
|
predictions.append(out_np)
|
|
pbar.update_absolute(frame_idx, num_frames)
|
|
|
|
return predictions
|
|
|
|
|
|
def inference_best_frame(
|
|
source_image,
|
|
driving_video,
|
|
generator: OcclusionAwareGenerator,
|
|
kp_detector: KPDetector,
|
|
relative_movement: bool,
|
|
relative_jacobian: bool,
|
|
adapt_movement_scale: bool,
|
|
cpu=False,
|
|
) -> list:
|
|
best_frame_idx = find_best_frame(source_image, driving_video, cpu)
|
|
|
|
first_half = driving_video[
|
|
:, :, : best_frame_idx + 1
|
|
] # Include the best frame in the first half
|
|
first_half = torch.flip(first_half, dims=[2]) # Reverse the first half
|
|
second_half = driving_video[:, :, best_frame_idx + 1 :]
|
|
|
|
predictions_first = inference(
|
|
source_image,
|
|
first_half,
|
|
generator,
|
|
kp_detector,
|
|
relative_movement,
|
|
relative_jacobian,
|
|
adapt_movement_scale,
|
|
cpu,
|
|
)
|
|
predictions_second = inference(
|
|
source_image,
|
|
second_half,
|
|
generator,
|
|
kp_detector,
|
|
relative_movement,
|
|
relative_jacobian,
|
|
adapt_movement_scale,
|
|
cpu,
|
|
)
|
|
|
|
predictions = []
|
|
predictions_first = predictions_first[::-1] # Reverse the first half back
|
|
predictions.extend(predictions_first)
|
|
predictions.extend(predictions_second)
|
|
return predictions
|
|
|
|
|
|
def load_checkpoint(
|
|
config_path: str, checkpoint_path: str, cpu=False
|
|
) -> tuple[OcclusionAwareGenerator, KPDetector]:
|
|
with open(config_path) as f:
|
|
config = yaml.full_load(f)
|
|
|
|
generator = OcclusionAwareGenerator(
|
|
**config["model_params"]["generator_params"],
|
|
**config["model_params"]["common_params"],
|
|
)
|
|
kp_detector = KPDetector(
|
|
**config["model_params"]["kp_detector_params"],
|
|
**config["model_params"]["common_params"],
|
|
)
|
|
|
|
if not cpu:
|
|
generator.cuda()
|
|
kp_detector.cuda()
|
|
checkpoint = torch.load(checkpoint_path)
|
|
else:
|
|
checkpoint = torch.load(checkpoint_path, map_location=torch.device("cpu"))
|
|
|
|
generator.load_state_dict(checkpoint["generator"])
|
|
kp_detector.load_state_dict(checkpoint["kp_detector"])
|
|
|
|
generator.eval()
|
|
kp_detector.eval()
|
|
|
|
return generator, kp_detector
|