refactor
This commit is contained in:
FuouM
+26
-5
@@ -9,23 +9,44 @@ supported_model_names = [
|
||||
# "bair",
|
||||
]
|
||||
|
||||
non_cpk_model_names = [
|
||||
"fashion"
|
||||
]
|
||||
non_cpk_model_names = ["fashion"]
|
||||
|
||||
default_model_name = "vox-adv"
|
||||
|
||||
checkpoint_folder_link = "https://github.com/graphemecluster/first-order-model-demo/releases/download/checkpoints"
|
||||
|
||||
checkpoint_suffix = ".pth.tar"
|
||||
config_suffix = "256.yaml"
|
||||
|
||||
config_folder = "config"
|
||||
checkpoint_folder = "checkpoints"
|
||||
|
||||
partswap_model_config_dict = {
|
||||
"vox-5segments" : "vox-256-sem-5segments.yaml",
|
||||
"vox-5segments": "vox-256-sem-5segments.yaml",
|
||||
"vox-10segments": "vox-256-sem-10segments.yaml",
|
||||
"vox-15segments": "vox-256-sem-15segments.yaml",
|
||||
}
|
||||
|
||||
default_partswap_model_name = "vox-10segments"
|
||||
partswap_model_names = list(partswap_model_config_dict.keys())
|
||||
|
||||
default_partswap_model_name = "vox-10segments"
|
||||
|
||||
seg10_segment_names = [
|
||||
"m_0_background",
|
||||
"m_1_lf_ear",
|
||||
"m_2_mouth",
|
||||
"m_3_lf_fr_head",
|
||||
"m_4_rt_ear",
|
||||
"m_5_frnt_face",
|
||||
"m_6_lt_cheek",
|
||||
"m_7_eyes",
|
||||
"m_8_rt_fr_head",
|
||||
"m_9_shoulder",
|
||||
"m_10_hair_top",
|
||||
]
|
||||
|
||||
seg_model_dict = {
|
||||
"vox-5segments": [],
|
||||
"vox-10segments": seg10_segment_names,
|
||||
"vox-15segments": [],
|
||||
}
|
||||
|
||||
@@ -0,0 +1,185 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
from comfy.utils import ProgressBar
|
||||
from tqdm import tqdm
|
||||
from scipy.spatial import ConvexHull
|
||||
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
|
||||
@@ -0,0 +1,35 @@
|
||||
import yaml
|
||||
import torch
|
||||
|
||||
from .modules.generator import OcclusionAwareGenerator
|
||||
from .modules.keypoint_detector import KPDetector
|
||||
|
||||
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
|
||||
-321
@@ -1,321 +0,0 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import yaml
|
||||
from tqdm import tqdm
|
||||
import numpy as np
|
||||
|
||||
from .constants import (
|
||||
checkpoint_folder,
|
||||
checkpoint_suffix,
|
||||
config_folder,
|
||||
partswap_model_config_dict,
|
||||
)
|
||||
from .modules.dense_motion import DenseMotionNetwork
|
||||
from .modules.reconstruction_module import ReconstructionModule
|
||||
from .modules.segmentation_module import SegmentationModule
|
||||
from .modules.util import AntiAliasInterpolation2d
|
||||
|
||||
from .sync_batchnorm.replicate import DataParallelWithCallback
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
def partial_state_dict_load(module, state_dict):
|
||||
own_state = module.state_dict()
|
||||
for name, param in state_dict.items():
|
||||
if name not in own_state:
|
||||
continue
|
||||
|
||||
if isinstance(param, torch.nn.Parameter):
|
||||
# backwards compatibility for serialized parameters
|
||||
param = param.data
|
||||
own_state[name].copy_(param)
|
||||
|
||||
|
||||
def load_reconstruction_module(module, checkpoint):
|
||||
if "generator" in checkpoint:
|
||||
partial_state_dict_load(module, checkpoint["generator"])
|
||||
else:
|
||||
module.load_state_dict(checkpoint["reconstruction_module"])
|
||||
|
||||
|
||||
def load_segmentation_module(module, checkpoint):
|
||||
if "kp_detector" in checkpoint:
|
||||
partial_state_dict_load(module, checkpoint["kp_detector"])
|
||||
module.state_dict()["affine.weight"].copy_(
|
||||
checkpoint["kp_detector"]["jacobian.weight"]
|
||||
)
|
||||
module.state_dict()["affine.bias"].copy_(
|
||||
checkpoint["kp_detector"]["jacobian.bias"]
|
||||
)
|
||||
module.state_dict()["shift.weight"].copy_(
|
||||
checkpoint["kp_detector"]["kp.weight"]
|
||||
)
|
||||
module.state_dict()["shift.bias"].copy_(checkpoint["kp_detector"]["kp.bias"])
|
||||
if "semantic_seg.weight" in checkpoint["kp_detector"]:
|
||||
module.state_dict()["segmentation.weight"].copy_(
|
||||
checkpoint["kp_detector"]["semantic_seg.weight"]
|
||||
)
|
||||
module.state_dict()["segmentation.bias"].copy_(
|
||||
checkpoint["kp_detector"]["semantic_seg.bias"]
|
||||
)
|
||||
else:
|
||||
print("Segmentation part initialized at random.")
|
||||
else:
|
||||
module.load_state_dict(checkpoint["segmentation_module"])
|
||||
|
||||
|
||||
def get_partswap_model_file_name(model_name: str):
|
||||
return f"{checkpoint_folder}/{model_name}{checkpoint_suffix}"
|
||||
|
||||
|
||||
def load_partswap_checkpoint(
|
||||
config_path: str,
|
||||
checkpoint_path: str,
|
||||
blend_scale: float,
|
||||
use_fomm: bool,
|
||||
cpu=False,
|
||||
):
|
||||
with open(config_path) as f:
|
||||
config = yaml.full_load(f)
|
||||
|
||||
reconstruction_module = PartSwapGenerator(
|
||||
blend_scale=blend_scale,
|
||||
first_order_motion_model=use_fomm,
|
||||
**config["model_params"]["reconstruction_module_params"],
|
||||
**config["model_params"]["common_params"],
|
||||
)
|
||||
|
||||
segmentation_module = SegmentationModule(
|
||||
**config["model_params"]["segmentation_module_params"],
|
||||
**config["model_params"]["common_params"],
|
||||
)
|
||||
|
||||
if cpu:
|
||||
checkpoint = torch.load(checkpoint_path, map_location=torch.device("cpu"))
|
||||
else:
|
||||
segmentation_module.cuda()
|
||||
checkpoint = torch.load(checkpoint_path)
|
||||
|
||||
load_reconstruction_module(reconstruction_module, checkpoint)
|
||||
load_segmentation_module(segmentation_module, checkpoint)
|
||||
|
||||
if not cpu:
|
||||
reconstruction_module = DataParallelWithCallback(reconstruction_module)
|
||||
segmentation_module = DataParallelWithCallback(segmentation_module)
|
||||
|
||||
reconstruction_module.eval()
|
||||
segmentation_module.eval()
|
||||
|
||||
return reconstruction_module, segmentation_module
|
||||
|
||||
|
||||
class PartSwapGenerator(ReconstructionModule):
|
||||
def __init__(self, blend_scale=1, first_order_motion_model=False, **kwargs):
|
||||
super(PartSwapGenerator, self).__init__(**kwargs)
|
||||
if blend_scale == 1:
|
||||
self.blend_downsample = lambda x: x
|
||||
else:
|
||||
self.blend_downsample = AntiAliasInterpolation2d(1, blend_scale)
|
||||
|
||||
if first_order_motion_model:
|
||||
self.dense_motion_network = DenseMotionNetwork()
|
||||
else:
|
||||
self.dense_motion_network = None
|
||||
|
||||
def forward(
|
||||
self,
|
||||
source_image,
|
||||
target_image,
|
||||
seg_target,
|
||||
seg_source,
|
||||
blend_mask,
|
||||
use_source_segmentation=False,
|
||||
):
|
||||
# Encoding of source image
|
||||
enc_source = self.first(source_image)
|
||||
for i in range(len(self.down_blocks)):
|
||||
enc_source = self.down_blocks[i](enc_source)
|
||||
|
||||
# Encoding of target image
|
||||
enc_target = self.first(target_image)
|
||||
for i in range(len(self.down_blocks)):
|
||||
enc_target = self.down_blocks[i](enc_target)
|
||||
|
||||
output_dict = {}
|
||||
# Compute flow field for source image
|
||||
if self.dense_motion_network is None:
|
||||
segment_motions = self.segment_motion(seg_target, seg_source)
|
||||
segment_motions = segment_motions.permute(0, 1, 4, 2, 3)
|
||||
mask = seg_target["segmentation"].unsqueeze(2)
|
||||
deformation = (segment_motions * mask).sum(dim=1)
|
||||
deformation = deformation.permute(0, 2, 3, 1)
|
||||
else:
|
||||
motion = self.dense_motion_network(
|
||||
source_image=source_image, seg_target=seg_target, seg_source=seg_source
|
||||
)
|
||||
deformation = motion["deformation"]
|
||||
|
||||
# Deform source encoding according to the motion
|
||||
enc_source = self.deform_input(enc_source, deformation)
|
||||
|
||||
if self.estimate_visibility:
|
||||
if self.dense_motion_network is None:
|
||||
visibility = seg_source["segmentation"][:, 1:].sum(
|
||||
dim=1, keepdim=True
|
||||
) * (
|
||||
1
|
||||
- seg_target["segmentation"][:, 1:]
|
||||
.sum(dim=1, keepdim=True)
|
||||
.detach()
|
||||
)
|
||||
visibility = 1 - visibility
|
||||
else:
|
||||
visibility = motion["visibility"]
|
||||
|
||||
if (
|
||||
enc_source.shape[2] != visibility.shape[2]
|
||||
or enc_source.shape[3] != visibility.shape[3]
|
||||
):
|
||||
visibility = F.interpolate(
|
||||
visibility, size=enc_source.shape[2:], mode="bilinear"
|
||||
)
|
||||
enc_source = enc_source * visibility
|
||||
|
||||
blend_mask = self.blend_downsample(blend_mask)
|
||||
# If source segmentation is provided use it should be deformed before blending
|
||||
if use_source_segmentation:
|
||||
blend_mask = self.deform_input(blend_mask, deformation)
|
||||
|
||||
out = enc_target * (1 - blend_mask) + enc_source * blend_mask
|
||||
|
||||
out = self.bottleneck(out)
|
||||
|
||||
for i in range(len(self.up_blocks)):
|
||||
out = self.up_blocks[i](out)
|
||||
|
||||
out = self.final(out)
|
||||
out = F.sigmoid(out)
|
||||
|
||||
output_dict["prediction"] = out
|
||||
|
||||
return output_dict
|
||||
|
||||
def partswap_inference(
|
||||
swap_indices: list[int],
|
||||
source_image,
|
||||
target_video,
|
||||
reconstruction_module: PartSwapGenerator,
|
||||
segmentation_module: SegmentationModule,
|
||||
use_source_seg: bool,
|
||||
hard_edges=False,
|
||||
cpu=False,
|
||||
) -> list:
|
||||
with torch.no_grad():
|
||||
predictions = []
|
||||
seg_targets = []
|
||||
if not cpu:
|
||||
source_image = source_image.cuda()
|
||||
target_video = target_video.cuda()
|
||||
print("Using GPU")
|
||||
|
||||
seg_source = segmentation_module(source_image)
|
||||
|
||||
num_frames = target_video.shape[2]
|
||||
for frame_idx in tqdm(range(num_frames)):
|
||||
target_frame = target_video[:, :, frame_idx]
|
||||
seg_target = segmentation_module(target_frame)
|
||||
seg_targets.append(seg_target)
|
||||
if use_source_seg:
|
||||
blend_mask = seg_source["segmentation"]
|
||||
else:
|
||||
blend_mask = seg_target["segmentation"]
|
||||
|
||||
blend_mask = blend_mask[:, swap_indices].sum(dim=1, keepdim=True)
|
||||
if hard_edges:
|
||||
blend_mask = (blend_mask > 0.5).type(blend_mask.type())
|
||||
|
||||
out = reconstruction_module(
|
||||
source_image,
|
||||
target_frame,
|
||||
seg_source=seg_source,
|
||||
seg_target=seg_target,
|
||||
blend_mask=blend_mask,
|
||||
use_source_segmentation=use_source_seg,
|
||||
)
|
||||
|
||||
predictions.append(np.transpose(out['prediction'].data.cpu().numpy(), [0, 2, 3, 1])[0])
|
||||
|
||||
return seg_source, seg_targets, predictions
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import matplotlib.patches as mpatches
|
||||
from matplotlib.figure import Figure
|
||||
from matplotlib.backends.backend_agg import FigureCanvasAgg
|
||||
from PIL import Image
|
||||
|
||||
def visualize_frame(image, seg_source, hard_edges=True, colormap='gist_rainbow', alpha=0.6):
|
||||
# Ensure image is the right shape and on CPU
|
||||
image = image.squeeze(0).cpu() # Remove batch dimension
|
||||
|
||||
# Extract segmentation from seg_source dictionary
|
||||
mask = seg_source["segmentation"]
|
||||
|
||||
# Interpolate mask to match image size if necessary
|
||||
if mask.shape[2:] != image.shape[1:]:
|
||||
mask = F.interpolate(mask, size=image.shape[1:], mode='bilinear')
|
||||
|
||||
if hard_edges:
|
||||
mask = (torch.max(mask, dim=1, keepdim=True)[0] == mask).float()
|
||||
|
||||
colormap = plt.get_cmap(colormap)
|
||||
num_segments = mask.shape[1]
|
||||
mask = mask.squeeze(0).permute(1, 2, 0).cpu().numpy()
|
||||
color_mask = np.zeros((256, 256, 3), dtype=np.float32)
|
||||
patches = []
|
||||
for i in range(num_segments):
|
||||
if i != 0:
|
||||
color = np.array(colormap((i - 1) / (num_segments - 1)))[:3]
|
||||
else:
|
||||
color = np.array((0, 0, 0))
|
||||
patches.append(mpatches.Patch(color=color, label=str(i)))
|
||||
color_mask += mask[..., i:(i+1)] * color.reshape(1, 1, 3)
|
||||
|
||||
# Convert image to numpy and normalize to [0, 1] range if necessary
|
||||
image_np = image.permute(1, 2, 0).numpy()
|
||||
if image_np.max() > 1:
|
||||
image_np = image_np / 255.0
|
||||
|
||||
# Overlay the color mask on the original image
|
||||
overlaid_image = image_np * (1 - alpha) + color_mask * alpha
|
||||
|
||||
# Ensure all values are in [0, 1] range
|
||||
overlaid_image = np.clip(overlaid_image, 0, 1)
|
||||
|
||||
# Create a figure and axis
|
||||
fig = Figure(figsize=(10, 10))
|
||||
canvas = FigureCanvasAgg(fig)
|
||||
ax = fig.add_subplot(111)
|
||||
|
||||
# Display the image
|
||||
ax.imshow(overlaid_image)
|
||||
ax.axis('off')
|
||||
|
||||
# Add legend
|
||||
ax.legend(handles=patches, loc='upper right', bbox_to_anchor=(1, 1), fontsize='small')
|
||||
|
||||
# Render the figure to numpy array
|
||||
fig.tight_layout(pad=0)
|
||||
canvas.draw()
|
||||
result_image = np.frombuffer(canvas.tostring_rgb(), dtype='uint8')
|
||||
result_image = result_image.reshape(canvas.get_width_height()[::-1] + (3,))
|
||||
|
||||
# # Resize to 256x256 if necessary
|
||||
# if result_image.shape[:2] != (256, 256):
|
||||
# result_image = np.array(Image.fromarray(result_image).resize((256, 256)))
|
||||
|
||||
# Convert to torch tensor and add batch dimension to get [1, 256, 256, 3]
|
||||
result_tensor = torch.from_numpy(result_image).unsqueeze(0).float() / 255.0
|
||||
|
||||
return result_tensor
|
||||
@@ -0,0 +1,60 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
|
||||
from comfy.utils import ProgressBar
|
||||
from .modules.segmentation_module import SegmentationModule
|
||||
from .partswap_model import PartSwapGenerator
|
||||
|
||||
|
||||
def partswap_inference(
|
||||
swap_indices: list[int],
|
||||
source_image,
|
||||
target_video,
|
||||
reconstruction_module: PartSwapGenerator,
|
||||
segmentation_module: SegmentationModule,
|
||||
use_source_seg: bool,
|
||||
hard_edges=False,
|
||||
cpu=False,
|
||||
) -> list:
|
||||
with torch.no_grad():
|
||||
predictions = []
|
||||
seg_targets = []
|
||||
if not cpu:
|
||||
source_image = source_image.cuda()
|
||||
target_video = target_video.cuda()
|
||||
print("Using GPU")
|
||||
|
||||
seg_source = segmentation_module(source_image)
|
||||
|
||||
num_frames = target_video.shape[2]
|
||||
pbar = ProgressBar(num_frames)
|
||||
for frame_idx in tqdm(range(num_frames)):
|
||||
target_frame = target_video[:, :, frame_idx]
|
||||
seg_target = segmentation_module(target_frame)
|
||||
seg_targets.append(seg_target)
|
||||
if use_source_seg:
|
||||
blend_mask = seg_source["segmentation"]
|
||||
else:
|
||||
blend_mask = seg_target["segmentation"]
|
||||
|
||||
blend_mask = blend_mask[:, swap_indices].sum(dim=1, keepdim=True)
|
||||
if hard_edges:
|
||||
blend_mask = (blend_mask > 0.5).type(blend_mask.type())
|
||||
|
||||
out = reconstruction_module(
|
||||
source_image,
|
||||
target_frame,
|
||||
seg_source=seg_source,
|
||||
seg_target=seg_target,
|
||||
blend_mask=blend_mask,
|
||||
use_source_segmentation=use_source_seg,
|
||||
)
|
||||
|
||||
predictions.append(
|
||||
np.transpose(out["prediction"].data.cpu().numpy(), [0, 2, 3, 1])[0]
|
||||
)
|
||||
|
||||
pbar.update_absolute(frame_idx, num_frames)
|
||||
|
||||
return seg_source, seg_targets, predictions
|
||||
@@ -0,0 +1,92 @@
|
||||
import torch
|
||||
import yaml
|
||||
|
||||
from .modules.segmentation_module import SegmentationModule
|
||||
from .partswap_model import PartSwapGenerator
|
||||
from .sync_batchnorm.replicate import DataParallelWithCallback
|
||||
|
||||
|
||||
def partial_state_dict_load(module, state_dict):
|
||||
own_state = module.state_dict()
|
||||
for name, param in state_dict.items():
|
||||
if name not in own_state:
|
||||
continue
|
||||
|
||||
if isinstance(param, torch.nn.Parameter):
|
||||
# backwards compatibility for serialized parameters
|
||||
param = param.data
|
||||
own_state[name].copy_(param)
|
||||
|
||||
|
||||
def load_reconstruction_module(module, checkpoint):
|
||||
if "generator" in checkpoint:
|
||||
partial_state_dict_load(module, checkpoint["generator"])
|
||||
else:
|
||||
module.load_state_dict(checkpoint["reconstruction_module"])
|
||||
|
||||
|
||||
def load_segmentation_module(module, checkpoint):
|
||||
if "kp_detector" in checkpoint:
|
||||
partial_state_dict_load(module, checkpoint["kp_detector"])
|
||||
module.state_dict()["affine.weight"].copy_(
|
||||
checkpoint["kp_detector"]["jacobian.weight"]
|
||||
)
|
||||
module.state_dict()["affine.bias"].copy_(
|
||||
checkpoint["kp_detector"]["jacobian.bias"]
|
||||
)
|
||||
module.state_dict()["shift.weight"].copy_(
|
||||
checkpoint["kp_detector"]["kp.weight"]
|
||||
)
|
||||
module.state_dict()["shift.bias"].copy_(checkpoint["kp_detector"]["kp.bias"])
|
||||
if "semantic_seg.weight" in checkpoint["kp_detector"]:
|
||||
module.state_dict()["segmentation.weight"].copy_(
|
||||
checkpoint["kp_detector"]["semantic_seg.weight"]
|
||||
)
|
||||
module.state_dict()["segmentation.bias"].copy_(
|
||||
checkpoint["kp_detector"]["semantic_seg.bias"]
|
||||
)
|
||||
else:
|
||||
print("Segmentation part initialized at random.")
|
||||
else:
|
||||
module.load_state_dict(checkpoint["segmentation_module"])
|
||||
|
||||
|
||||
def load_partswap_checkpoint(
|
||||
config_path: str,
|
||||
checkpoint_path: str,
|
||||
blend_scale: float,
|
||||
use_fomm: bool,
|
||||
cpu=False,
|
||||
):
|
||||
with open(config_path) as f:
|
||||
config = yaml.full_load(f)
|
||||
|
||||
reconstruction_module = PartSwapGenerator(
|
||||
blend_scale=blend_scale,
|
||||
first_order_motion_model=use_fomm,
|
||||
**config["model_params"]["reconstruction_module_params"],
|
||||
**config["model_params"]["common_params"],
|
||||
)
|
||||
|
||||
segmentation_module = SegmentationModule(
|
||||
**config["model_params"]["segmentation_module_params"],
|
||||
**config["model_params"]["common_params"],
|
||||
)
|
||||
|
||||
if cpu:
|
||||
checkpoint = torch.load(checkpoint_path, map_location=torch.device("cpu"))
|
||||
else:
|
||||
segmentation_module.cuda()
|
||||
checkpoint = torch.load(checkpoint_path)
|
||||
|
||||
load_reconstruction_module(reconstruction_module, checkpoint)
|
||||
load_segmentation_module(segmentation_module, checkpoint)
|
||||
|
||||
if not cpu:
|
||||
reconstruction_module = DataParallelWithCallback(reconstruction_module)
|
||||
segmentation_module = DataParallelWithCallback(segmentation_module)
|
||||
|
||||
reconstruction_module.eval()
|
||||
segmentation_module.eval()
|
||||
|
||||
return reconstruction_module, segmentation_module
|
||||
@@ -0,0 +1,97 @@
|
||||
import torch.nn.functional as F
|
||||
|
||||
from .modules.dense_motion import DenseMotionNetwork
|
||||
from .modules.reconstruction_module import ReconstructionModule
|
||||
from .modules.util import AntiAliasInterpolation2d
|
||||
|
||||
|
||||
class PartSwapGenerator(ReconstructionModule):
|
||||
def __init__(self, blend_scale=1, first_order_motion_model=False, **kwargs):
|
||||
super(PartSwapGenerator, self).__init__(**kwargs)
|
||||
if blend_scale == 1:
|
||||
self.blend_downsample = lambda x: x
|
||||
else:
|
||||
self.blend_downsample = AntiAliasInterpolation2d(1, blend_scale)
|
||||
|
||||
if first_order_motion_model:
|
||||
self.dense_motion_network = DenseMotionNetwork()
|
||||
else:
|
||||
self.dense_motion_network = None
|
||||
|
||||
def forward(
|
||||
self,
|
||||
source_image,
|
||||
target_image,
|
||||
seg_target,
|
||||
seg_source,
|
||||
blend_mask,
|
||||
use_source_segmentation=False,
|
||||
):
|
||||
# Encoding of source image
|
||||
enc_source = self.first(source_image)
|
||||
for i in range(len(self.down_blocks)):
|
||||
enc_source = self.down_blocks[i](enc_source)
|
||||
|
||||
# Encoding of target image
|
||||
enc_target = self.first(target_image)
|
||||
for i in range(len(self.down_blocks)):
|
||||
enc_target = self.down_blocks[i](enc_target)
|
||||
|
||||
output_dict = {}
|
||||
# Compute flow field for source image
|
||||
if self.dense_motion_network is None:
|
||||
segment_motions = self.segment_motion(seg_target, seg_source)
|
||||
segment_motions = segment_motions.permute(0, 1, 4, 2, 3)
|
||||
mask = seg_target["segmentation"].unsqueeze(2)
|
||||
deformation = (segment_motions * mask).sum(dim=1)
|
||||
deformation = deformation.permute(0, 2, 3, 1)
|
||||
else:
|
||||
motion = self.dense_motion_network(
|
||||
source_image=source_image, seg_target=seg_target, seg_source=seg_source
|
||||
)
|
||||
deformation = motion["deformation"]
|
||||
|
||||
# Deform source encoding according to the motion
|
||||
enc_source = self.deform_input(enc_source, deformation)
|
||||
|
||||
if self.estimate_visibility:
|
||||
if self.dense_motion_network is None:
|
||||
visibility = seg_source["segmentation"][:, 1:].sum(
|
||||
dim=1, keepdim=True
|
||||
) * (
|
||||
1
|
||||
- seg_target["segmentation"][:, 1:]
|
||||
.sum(dim=1, keepdim=True)
|
||||
.detach()
|
||||
)
|
||||
visibility = 1 - visibility
|
||||
else:
|
||||
visibility = motion["visibility"]
|
||||
|
||||
if (
|
||||
enc_source.shape[2] != visibility.shape[2]
|
||||
or enc_source.shape[3] != visibility.shape[3]
|
||||
):
|
||||
visibility = F.interpolate(
|
||||
visibility, size=enc_source.shape[2:], mode="bilinear"
|
||||
)
|
||||
enc_source = enc_source * visibility
|
||||
|
||||
blend_mask = self.blend_downsample(blend_mask)
|
||||
# If source segmentation is provided use it should be deformed before blending
|
||||
if use_source_segmentation:
|
||||
blend_mask = self.deform_input(blend_mask, deformation)
|
||||
|
||||
out = enc_target * (1 - blend_mask) + enc_source * blend_mask
|
||||
|
||||
out = self.bottleneck(out)
|
||||
|
||||
for i in range(len(self.up_blocks)):
|
||||
out = self.up_blocks[i](out)
|
||||
|
||||
out = self.final(out)
|
||||
out = F.sigmoid(out)
|
||||
|
||||
output_dict["prediction"] = out
|
||||
|
||||
return output_dict
|
||||
@@ -2,3 +2,4 @@ numpy
|
||||
torch
|
||||
scipy
|
||||
yaml
|
||||
matplotlib
|
||||
@@ -6,16 +6,27 @@
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
from .constants import default_model_name, supported_model_names, config_folder
|
||||
from .fomm_inference import inference, inference_best_frame
|
||||
from .fomm_loader import load_checkpoint
|
||||
from .partswap_inference import partswap_inference
|
||||
from .partswap_loader import load_partswap_checkpoint
|
||||
|
||||
from .constants import (
|
||||
config_folder,
|
||||
default_model_name,
|
||||
default_partswap_model_name,
|
||||
partswap_model_config_dict,
|
||||
partswap_model_names,
|
||||
supported_model_names,
|
||||
)
|
||||
from .seg_viz import visualize_frame
|
||||
from .utils import (
|
||||
build_seg_arguments,
|
||||
get_config_path,
|
||||
get_model_file_name,
|
||||
inference,
|
||||
inference_best_frame,
|
||||
load_checkpoint,
|
||||
get_partswap_model_file_name,
|
||||
out_video,
|
||||
reshape_image,
|
||||
)
|
||||
@@ -67,8 +78,8 @@ class FOMM_Runner:
|
||||
|
||||
def todo(
|
||||
self,
|
||||
source_image: torch.Tensor,
|
||||
driving_video_input: torch.Tensor,
|
||||
source_image,
|
||||
driving_video_input,
|
||||
model_name: str,
|
||||
frame_rate: float,
|
||||
relative_movement: bool,
|
||||
@@ -119,74 +130,18 @@ class FOMM_Runner:
|
||||
)
|
||||
|
||||
|
||||
from .constants import partswap_model_config_dict, default_partswap_model_name
|
||||
from .part_swap import (
|
||||
get_partswap_model_file_name,
|
||||
load_partswap_checkpoint,
|
||||
partswap_inference,
|
||||
visualize_frame,
|
||||
)
|
||||
|
||||
|
||||
def serialize_integers(int_list):
|
||||
return "_".join(map(str, int_list))
|
||||
|
||||
|
||||
def deserialize_integers(int_string):
|
||||
return list(map(int, int_string.split('_')))
|
||||
return list(map(int, int_string.split("_")))
|
||||
|
||||
|
||||
class FOMM_Seg10Chooser:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"m_0_background": (
|
||||
"BOOLEAN",
|
||||
{"default": False},
|
||||
),
|
||||
"m_1_lf_ear": (
|
||||
"BOOLEAN",
|
||||
{"default": True},
|
||||
),
|
||||
"m_2_mouth": (
|
||||
"BOOLEAN",
|
||||
{"default": True},
|
||||
),
|
||||
"m_3_lf_fr_head": (
|
||||
"BOOLEAN",
|
||||
{"default": True},
|
||||
),
|
||||
"m_4_rt_ear": (
|
||||
"BOOLEAN",
|
||||
{"default": True},
|
||||
),
|
||||
"m_5_frnt_face": (
|
||||
"BOOLEAN",
|
||||
{"default": True},
|
||||
),
|
||||
"m_6_lt_cheek": (
|
||||
"BOOLEAN",
|
||||
{"default": True},
|
||||
),
|
||||
"m_7_eyes": (
|
||||
"BOOLEAN",
|
||||
{"default": True},
|
||||
),
|
||||
"m_8_rt_fr_head": (
|
||||
"BOOLEAN",
|
||||
{"default": True},
|
||||
),
|
||||
"m_9_shoulder": (
|
||||
"BOOLEAN",
|
||||
{"default": True},
|
||||
),
|
||||
"m_10_hair_top": (
|
||||
"BOOLEAN",
|
||||
{"default": True},
|
||||
),
|
||||
},
|
||||
}
|
||||
return build_seg_arguments("vox-10segments")
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
|
||||
@@ -195,41 +150,12 @@ class FOMM_Seg10Chooser:
|
||||
FUNCTION = "todo"
|
||||
CATEGORY = "FirstOrderMM"
|
||||
|
||||
def todo(
|
||||
self,
|
||||
m_0_background: bool,
|
||||
m_1_lf_ear: bool,
|
||||
m_2_mouth: bool,
|
||||
m_3_lf_fr_head: bool,
|
||||
m_4_rt_ear: bool,
|
||||
m_5_frnt_face: bool,
|
||||
m_6_lt_cheek: bool,
|
||||
m_7_eyes: bool,
|
||||
m_8_rt_fr_head: bool,
|
||||
m_9_shoulder: bool,
|
||||
m_10_hair_top: bool,
|
||||
):
|
||||
seg_list = []
|
||||
segments = [
|
||||
m_0_background,
|
||||
m_1_lf_ear,
|
||||
m_2_mouth,
|
||||
m_3_lf_fr_head,
|
||||
m_4_rt_ear,
|
||||
m_5_frnt_face,
|
||||
m_6_lt_cheek,
|
||||
m_7_eyes,
|
||||
m_8_rt_fr_head,
|
||||
m_9_shoulder,
|
||||
m_10_hair_top,
|
||||
]
|
||||
|
||||
for i, seg in enumerate(segments):
|
||||
print(f"Segment {i}: {seg}, Type: {type(seg)}, Truthy: {bool(seg)}")
|
||||
if seg:
|
||||
seg_list.append(i)
|
||||
def todo(self, **args):
|
||||
print(args)
|
||||
assert len(args) == 11, f"Expected 11 arguments, got {len(args)}"
|
||||
segments = list(args.values())
|
||||
seg_list = [i for i, seg in enumerate(segments) if seg]
|
||||
seg_list = serialize_integers(seg_list)
|
||||
print("Final seg_list:", seg_list)
|
||||
return (seg_list,)
|
||||
|
||||
|
||||
@@ -241,7 +167,7 @@ class FOMM_Partswap:
|
||||
"source_image": ("IMAGE",),
|
||||
"driving_video_input": ("IMAGE",),
|
||||
"model_name": (
|
||||
list(partswap_model_config_dict.keys()),
|
||||
partswap_model_names,
|
||||
{"default": default_partswap_model_name},
|
||||
),
|
||||
"frame_rate": ("FLOAT", {"default": 30.0}),
|
||||
@@ -340,14 +266,14 @@ class FOMM_Partswap:
|
||||
)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"FOMM_Runner": FOMM_Runner,
|
||||
"FOMM_Partswap": FOMM_Partswap,
|
||||
"FOMM_Seg10Chooser": FOMM_Seg10Chooser,
|
||||
}
|
||||
# NODE_CLASS_MAPPINGS = {
|
||||
# "FOMM_Runner": FOMM_Runner,
|
||||
# "FOMM_Partswap": FOMM_Partswap,
|
||||
# "FOMM_Seg10Chooser": FOMM_Seg10Chooser,
|
||||
# }
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"FOMM_Runner": "FOMM Runner",
|
||||
"FOMM_Partswap": "FOMM Partswap",
|
||||
"FOMM_Seg10Chooser": "FOMM_Seg10Chooser",
|
||||
}
|
||||
# NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
# "FOMM_Runner": "FOMM Runner",
|
||||
# "FOMM_Partswap": "FOMM Partswap",
|
||||
# "FOMM_Seg10Chooser": "FOMM_Seg10Chooser",
|
||||
# }
|
||||
|
||||
+134
@@ -0,0 +1,134 @@
|
||||
import matplotlib.patches as mpatches
|
||||
import matplotlib.pyplot as plt
|
||||
from matplotlib.backends.backend_agg import FigureCanvasAgg
|
||||
from matplotlib.figure import Figure
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
def preprocess_image_and_mask(image, seg_source):
|
||||
image = image.squeeze(0).cpu() # Remove batch dimension
|
||||
mask = seg_source["segmentation"]
|
||||
|
||||
if mask.shape[2:] != image.shape[1:]:
|
||||
mask = F.interpolate(mask, size=image.shape[1:], mode='bilinear')
|
||||
|
||||
return image, mask
|
||||
|
||||
def apply_hard_edges(mask):
|
||||
return (torch.max(mask, dim=1, keepdim=True)[0] == mask).float()
|
||||
|
||||
def create_color_mask(mask, colormap='gist_rainbow'):
|
||||
colormap = plt.get_cmap(colormap)
|
||||
num_segments = mask.shape[1]
|
||||
mask = mask.squeeze(0).permute(1, 2, 0).cpu().numpy()
|
||||
color_mask = np.zeros((256, 256, 3), dtype=np.float32)
|
||||
patches = []
|
||||
|
||||
for i in range(num_segments):
|
||||
color = np.array(colormap((i - 1) / (num_segments - 1)))[:3] if i != 0 else np.array((0, 0, 0))
|
||||
patches.append(mpatches.Patch(color=color, label=str(i)))
|
||||
color_mask += mask[..., i:(i+1)] * color.reshape(1, 1, 3)
|
||||
|
||||
return color_mask, patches
|
||||
|
||||
def overlay_mask(image, color_mask, alpha=0.6):
|
||||
image_np = image.permute(1, 2, 0).numpy()
|
||||
if image_np.max() > 1:
|
||||
image_np = image_np / 255.0
|
||||
|
||||
overlaid_image = image_np * (1 - alpha) + color_mask * alpha
|
||||
return np.clip(overlaid_image, 0, 1)
|
||||
|
||||
def create_figure(overlaid_image, patches):
|
||||
fig = Figure(figsize=(10, 10))
|
||||
canvas = FigureCanvasAgg(fig)
|
||||
ax = fig.add_subplot(111)
|
||||
|
||||
ax.imshow(overlaid_image)
|
||||
ax.axis('off')
|
||||
ax.legend(handles=patches, loc='upper right', bbox_to_anchor=(1, 1), fontsize='small')
|
||||
|
||||
fig.tight_layout(pad=0)
|
||||
return fig, canvas
|
||||
|
||||
def render_figure_to_tensor(canvas):
|
||||
canvas.draw()
|
||||
result_image = np.frombuffer(canvas.tostring_rgb(), dtype='uint8')
|
||||
result_image = result_image.reshape(canvas.get_width_height()[::-1] + (3,))
|
||||
result_tensor = torch.from_numpy(result_image).unsqueeze(0).float() / 255.0
|
||||
return result_tensor
|
||||
|
||||
def visualize_frame(image, seg_source, hard_edges=True, colormap='gist_rainbow', alpha=0.6):
|
||||
image, mask = preprocess_image_and_mask(image, seg_source)
|
||||
|
||||
if hard_edges:
|
||||
mask = apply_hard_edges(mask)
|
||||
|
||||
color_mask, patches = create_color_mask(mask, colormap)
|
||||
overlaid_image = overlay_mask(image, color_mask, alpha)
|
||||
fig, canvas = create_figure(overlaid_image, patches)
|
||||
result_tensor = render_figure_to_tensor(canvas)
|
||||
|
||||
return result_tensor
|
||||
|
||||
# def visualize_frame(image, seg_source, hard_edges=True, colormap='gist_rainbow', alpha=0.6):
|
||||
# # Ensure image is the right shape and on CPU
|
||||
# image = image.squeeze(0).cpu() # Remove batch dimension
|
||||
|
||||
# # Extract segmentation from seg_source dictionary
|
||||
# mask = seg_source["segmentation"]
|
||||
|
||||
# # Interpolate mask to match image size if necessary
|
||||
# if mask.shape[2:] != image.shape[1:]:
|
||||
# mask = F.interpolate(mask, size=image.shape[1:], mode='bilinear')
|
||||
|
||||
# if hard_edges:
|
||||
# mask = (torch.max(mask, dim=1, keepdim=True)[0] == mask).float()
|
||||
|
||||
# colormap = plt.get_cmap(colormap)
|
||||
# num_segments = mask.shape[1]
|
||||
# mask = mask.squeeze(0).permute(1, 2, 0).cpu().numpy()
|
||||
# color_mask = np.zeros((256, 256, 3), dtype=np.float32)
|
||||
# patches = []
|
||||
# for i in range(num_segments):
|
||||
# if i != 0:
|
||||
# color = np.array(colormap((i - 1) / (num_segments - 1)))[:3]
|
||||
# else:
|
||||
# color = np.array((0, 0, 0))
|
||||
# patches.append(mpatches.Patch(color=color, label=str(i)))
|
||||
# color_mask += mask[..., i:(i+1)] * color.reshape(1, 1, 3)
|
||||
|
||||
# # Convert image to numpy and normalize to [0, 1] range if necessary
|
||||
# image_np = image.permute(1, 2, 0).numpy()
|
||||
# if image_np.max() > 1:
|
||||
# image_np = image_np / 255.0
|
||||
|
||||
# # Overlay the color mask on the original image
|
||||
# overlaid_image = image_np * (1 - alpha) + color_mask * alpha
|
||||
|
||||
# # Ensure all values are in [0, 1] range
|
||||
# overlaid_image = np.clip(overlaid_image, 0, 1)
|
||||
|
||||
# # Create a figure and axis
|
||||
# fig = Figure(figsize=(10, 10))
|
||||
# canvas = FigureCanvasAgg(fig)
|
||||
# ax = fig.add_subplot(111)
|
||||
|
||||
# # Display the image
|
||||
# ax.imshow(overlaid_image)
|
||||
# ax.axis('off')
|
||||
|
||||
# # Add legend
|
||||
# ax.legend(handles=patches, loc='upper right', bbox_to_anchor=(1, 1), fontsize='small')
|
||||
|
||||
# # Render the figure to numpy array
|
||||
# fig.tight_layout(pad=0)
|
||||
# canvas.draw()
|
||||
# result_image = np.frombuffer(canvas.tostring_rgb(), dtype='uint8')
|
||||
# result_image = result_image.reshape(canvas.get_width_height()[::-1] + (3,))
|
||||
|
||||
# # Convert to torch tensor and add batch dimension to get [1, 256, 256, 3]
|
||||
# result_tensor = torch.from_numpy(result_image).unsqueeze(0).float() / 255.0
|
||||
|
||||
# return result_tensor
|
||||
@@ -1,35 +1,28 @@
|
||||
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. ")
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import yaml
|
||||
from comfy.utils import ProgressBar
|
||||
from scipy.spatial import ConvexHull
|
||||
from tqdm import tqdm
|
||||
|
||||
from .constants import (
|
||||
checkpoint_folder,
|
||||
checkpoint_folder_link,
|
||||
checkpoint_suffix,
|
||||
config_folder,
|
||||
config_suffix,
|
||||
non_cpk_model_names,
|
||||
partswap_model_names,
|
||||
seg_model_dict,
|
||||
supported_model_names,
|
||||
)
|
||||
from .modules.generator import OcclusionAwareGenerator
|
||||
from .modules.keypoint_detector import KPDetector
|
||||
|
||||
|
||||
def check_supported(entity_name: str, entity_support_list: list[str]):
|
||||
assert (
|
||||
entity_name in entity_support_list
|
||||
), f"Unknown {entity_name}, expected {entity_support_list}"
|
||||
|
||||
|
||||
def get_model_file_name(model_name: str) -> str:
|
||||
assert (
|
||||
model_name in supported_model_names
|
||||
), f"Unknown model. Got {model_name}, expected {supported_model_names}"
|
||||
check_supported(model_name, supported_model_names)
|
||||
|
||||
if model_name in non_cpk_model_names:
|
||||
model_file_name = f"{checkpoint_folder}/{model_name}{checkpoint_suffix}"
|
||||
@@ -46,51 +39,23 @@ def get_model_link(model_name: str) -> str:
|
||||
|
||||
|
||||
def get_config_path(model_name: str) -> str:
|
||||
config_path = f"{config_folder}/{model_name}-256.yaml"
|
||||
config_path = f"{config_folder}/{model_name}-{config_suffix}"
|
||||
return config_path
|
||||
|
||||
|
||||
def get_partswap_model_file_name(model_name: str):
|
||||
return f"{checkpoint_folder}/{model_name}{checkpoint_suffix}"
|
||||
|
||||
|
||||
def out_video(predictions: list[np.ndarray]):
|
||||
out_tensor_list = []
|
||||
for i in predictions:
|
||||
# out_img = i.astype(np.float32) / 255.0
|
||||
out_img = torch.from_numpy(i)
|
||||
out_tensor_list.append(out_img)
|
||||
images = torch.stack(out_tensor_list, dim=0)
|
||||
return images
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
def reshape_image(image, target_size=(256, 256)):
|
||||
# Check if image is 4D (batch, height, width, channels)
|
||||
if image.dim() != 4:
|
||||
@@ -107,177 +72,21 @@ def reshape_image(image, target_size=(256, 256)):
|
||||
return resized_image
|
||||
|
||||
|
||||
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 tensor_to_numpy(tensor):
|
||||
# Convert [1, 3, H, W] tensor to [H, W, 3] numpy array
|
||||
return tensor.squeeze(0).permute(1, 2, 0).cpu().numpy()
|
||||
|
||||
|
||||
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.")
|
||||
def build_seg_arguments(model_name: str):
|
||||
check_supported(model_name, partswap_model_names)
|
||||
segment_names = seg_model_dict[model_name]
|
||||
required = dict()
|
||||
for segment_name in segment_names:
|
||||
default = True
|
||||
if segment_name.endswith("background"):
|
||||
default = False
|
||||
required[segment_name] = ("BOOLEAN", {"default": default})
|
||||
|
||||
source_np = tensor_to_numpy(source_image)
|
||||
kp_source = fa.get_landmarks(255 * source_np)[0]
|
||||
kp_source = normalize_kp_best_frame(kp_source)
|
||||
arguments = {"required": required}
|
||||
|
||||
norm = float("inf")
|
||||
frame_num = 0
|
||||
|
||||
num_frames = driving_video.shape[1]
|
||||
|
||||
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 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
|
||||
# print(f"{adapt_movement_scale=}")
|
||||
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 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
|
||||
return arguments
|
||||
|
||||
Reference in New Issue
Block a user