diff --git a/constants.py b/constants.py index c53c7de..1e43992 100644 --- a/constants.py +++ b/constants.py @@ -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" \ No newline at end of file +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": [], +} diff --git a/fomm_inference.py b/fomm_inference.py new file mode 100644 index 0000000..2665374 --- /dev/null +++ b/fomm_inference.py @@ -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 \ No newline at end of file diff --git a/fomm_loader.py b/fomm_loader.py new file mode 100644 index 0000000..b100543 --- /dev/null +++ b/fomm_loader.py @@ -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 \ No newline at end of file diff --git a/part_swap.py b/part_swap.py deleted file mode 100644 index 5a14e25..0000000 --- a/part_swap.py +++ /dev/null @@ -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 \ No newline at end of file diff --git a/partswap_inference.py b/partswap_inference.py new file mode 100644 index 0000000..79a9859 --- /dev/null +++ b/partswap_inference.py @@ -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 diff --git a/partswap_loader.py b/partswap_loader.py new file mode 100644 index 0000000..55c599c --- /dev/null +++ b/partswap_loader.py @@ -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 diff --git a/partswap_model.py b/partswap_model.py new file mode 100644 index 0000000..49085c0 --- /dev/null +++ b/partswap_model.py @@ -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 diff --git a/requirements.txt b/requirements.txt index 054a554..fe18382 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,3 +2,4 @@ numpy torch scipy yaml +matplotlib \ No newline at end of file diff --git a/run.py b/run.py index 399fa4e..0120909 100644 --- a/run.py +++ b/run.py @@ -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", +# } diff --git a/seg_viz.py b/seg_viz.py new file mode 100644 index 0000000..d78663c --- /dev/null +++ b/seg_viz.py @@ -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 \ No newline at end of file diff --git a/utils.py b/utils.py index 12894cd..213abb2 100644 --- a/utils.py +++ b/utils.py @@ -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