diff --git a/constants.py b/constants.py index a8e7a34..9373269 100644 --- a/constants.py +++ b/constants.py @@ -25,12 +25,18 @@ partswap_model_config_dict = { "vox-5segments": "vox-256-sem-5segments.yaml", "vox-10segments": "vox-256-sem-10segments.yaml", "vox-15segments": "vox-256-sem-15segments.yaml", + "vox-cpk": "vox-256-sem-10segments.yaml", + "vox-adv-cpk": "vox-256-sem-10segments.yaml", } +partswap_fomm_model_names = ["vox-cpk", "vox-adv-cpk"] + partswap_model_length_dict = { - "vox-5segments": 6, + "vox-5segments": 6, "vox-10segments": 11, - "vox-15segments": 15 + "vox-15segments": 15, + "vox-cpk": 19, + "vox-adv-cpk": 19, } partswap_model_names = list(partswap_model_config_dict.keys()) @@ -57,7 +63,7 @@ seg5_segment_names = [ "m_2_frnt_face", "m_3_shoulder", "m_4_hair", - "m_5_fr_head" + "m_5_fr_head", ] seg_15_segment_names = [ @@ -76,11 +82,13 @@ seg_15_segment_names = [ "m_12_lt_ear", "m_13_lt_hair", "m_14_rt_hair", - "m_15_neck" + "m_15_neck", ] seg_model_dict = { "vox-5segments": seg5_segment_names, "vox-10segments": seg10_segment_names, "vox-15segments": seg_15_segment_names, + "vox": seg10_segment_names, + "vox-adv": seg10_segment_names, } diff --git a/modules/dense_motion.py b/modules/dense_motion.py index 56d5ac5..1767bec 100644 --- a/modules/dense_motion.py +++ b/modules/dense_motion.py @@ -29,13 +29,13 @@ class DenseMotionNetwork(nn.Module): if self.scale_factor != 1: self.down = AntiAliasInterpolation2d(num_channels, self.scale_factor) - def create_heatmap_representations(self, source_image, kp_driving, kp_source): + def create_heatmap_representations(self, source_image, kp_driving, kp_source, kp_key='value'): """ Eq 6. in the paper H_k(z) """ spatial_size = source_image.shape[2:] - gaussian_driving = kp2gaussian(kp_driving, spatial_size=spatial_size, kp_variance=self.kp_variance) - gaussian_source = kp2gaussian(kp_source, spatial_size=spatial_size, kp_variance=self.kp_variance) + gaussian_driving = kp2gaussian(kp_driving, spatial_size=spatial_size, kp_variance=self.kp_variance, kp_key=kp_key) + gaussian_source = kp2gaussian(kp_source, spatial_size=spatial_size, kp_variance=self.kp_variance, kp_key=kp_key) heatmap = gaussian_driving - gaussian_source #adding background feature @@ -44,22 +44,22 @@ class DenseMotionNetwork(nn.Module): heatmap = heatmap.unsqueeze(2) return heatmap - def create_sparse_motions(self, source_image, kp_driving, kp_source): + def create_sparse_motions(self, source_image, kp_driving, kp_source, kp_key='value', jacob_key='jacobian'): """ Eq 4. in the paper T_{s<-d}(z) """ bs, _, h, w = source_image.shape - identity_grid = make_coordinate_grid((h, w), type=kp_source['value'].type()) + identity_grid = make_coordinate_grid((h, w), type=kp_source[kp_key].type()) identity_grid = identity_grid.view(1, 1, h, w, 2) - coordinate_grid = identity_grid - kp_driving['value'].view(bs, self.num_kp, 1, 1, 2) - if 'jacobian' in kp_driving: - jacobian = torch.matmul(kp_source['jacobian'], torch.inverse(kp_driving['jacobian'])) + coordinate_grid = identity_grid - kp_driving[kp_key].view(bs, self.num_kp, 1, 1, 2) + if jacob_key in kp_driving: + jacobian = torch.matmul(kp_source[jacob_key], torch.inverse(kp_driving[jacob_key])) jacobian = jacobian.unsqueeze(-3).unsqueeze(-3) jacobian = jacobian.repeat(1, 1, h, w, 1, 1) coordinate_grid = torch.matmul(jacobian, coordinate_grid.unsqueeze(-1)) coordinate_grid = coordinate_grid.squeeze(-1) - driving_to_source = coordinate_grid + kp_source['value'].view(bs, self.num_kp, 1, 1, 2) + driving_to_source = coordinate_grid + kp_source[kp_key].view(bs, self.num_kp, 1, 1, 2) #adding background feature identity_grid = identity_grid.repeat(bs, 1, 1, 1, 1) @@ -111,3 +111,36 @@ class DenseMotionNetwork(nn.Module): out_dict['occlusion_map'] = occlusion_map return out_dict + + def forward_partswap(self, source_image, seg_target, seg_source): + if self.scale_factor != 1: + source_image = self.down(source_image) + + bs, _, h, w = source_image.shape + + out_dict = dict() + heatmap_representation = self.create_heatmap_representations(source_image, seg_target, seg_source, kp_key='shift') + sparse_motion = self.create_sparse_motions(source_image, seg_target, seg_source, kp_key='shift', jacob_key='affine') + deformed_source = self.create_deformed_source_image(source_image, sparse_motion) + out_dict['sparse_deformed'] = deformed_source + + input = torch.cat([heatmap_representation, deformed_source], dim=2) + input = input.view(bs, -1, h, w) + + prediction = self.hourglass(input) + + mask = self.mask(prediction) + mask = F.softmax(mask, dim=1) + out_dict['mask'] = mask + mask = mask.unsqueeze(2) + sparse_motion = sparse_motion.permute(0, 1, 4, 2, 3) + deformation = (sparse_motion * mask).sum(dim=1) + deformation = deformation.permute(0, 2, 3, 1) + + out_dict['deformation'] = deformation + + if self.occlusion: + visibility = torch.sigmoid(self.occlusion(prediction)) + out_dict['visibility'] = visibility + + return out_dict diff --git a/modules/util.py b/modules/util.py index f9e5eaf..3b2e5d5 100644 --- a/modules/util.py +++ b/modules/util.py @@ -6,11 +6,11 @@ import torch from ..sync_batchnorm.batchnorm import SynchronizedBatchNorm2d as BatchNorm2d -def kp2gaussian(kp, spatial_size, kp_variance): +def kp2gaussian(kp, spatial_size, kp_variance, kp_key='value'): """ Transform a keypoint into gaussian like representation """ - mean = kp['value'] + mean = kp[kp_key] coordinate_grid = make_coordinate_grid(spatial_size, mean.type()) number_of_leading_dimensions = len(mean.shape) - 1 diff --git a/partswap_model.py b/partswap_model.py index 49085c0..ee55dea 100644 --- a/partswap_model.py +++ b/partswap_model.py @@ -14,7 +14,15 @@ class PartSwapGenerator(ReconstructionModule): self.blend_downsample = AntiAliasInterpolation2d(1, blend_scale) if first_order_motion_model: - self.dense_motion_network = DenseMotionNetwork() + self.dense_motion_network = DenseMotionNetwork( + block_expansion=64, + num_blocks=5, + max_features=1024, + num_kp=10, + num_channels=3, + estimate_occlusion_map=True, + scale_factor=0.25, + ) else: self.dense_motion_network = None @@ -46,7 +54,7 @@ class PartSwapGenerator(ReconstructionModule): deformation = (segment_motions * mask).sum(dim=1) deformation = deformation.permute(0, 2, 3, 1) else: - motion = self.dense_motion_network( + motion = self.dense_motion_network.forward_partswap( source_image=source_image, seg_target=seg_target, seg_source=seg_source ) deformation = motion["deformation"] diff --git a/run.py b/run.py index 2d563cc..951090c 100644 --- a/run.py +++ b/run.py @@ -6,22 +6,23 @@ """ from pathlib import Path -import torch -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 +import torch from .constants import ( config_folder, default_model_name, default_partswap_model_name, + partswap_fomm_model_names, partswap_model_config_dict, + partswap_model_length_dict, partswap_model_names, supported_model_names, - partswap_model_length_dict, ) +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 .seg_viz import visualize_frame from .utils import ( build_seg_arguments, @@ -273,8 +274,12 @@ class FOMM_Partswap: driving_video_org = reshape_image(driving_video_input, (256, 256)) driving_video = driving_video_org.unsqueeze(0).permute(0, 2, 1, 3, 4) + use_fomm = False + if model_name in partswap_fomm_model_names: + use_fomm = True + reconstruction_module, segmentation_module = load_partswap_checkpoint( - config_path, checkpoint_path, blend_scale, use_fomm=False + config_path, checkpoint_path, blend_scale, use_fomm=use_fomm ) print(f"Chosen {chosen_seg_indices}") swap_indices = deserialize_integers(chosen_seg_indices) @@ -294,7 +299,9 @@ class FOMM_Partswap: seg_source, seg_targets, predictions = partswap_inference(**params) seg_src_viz = visualize_frame(source_image, seg_source, alpha=viz_alpha) - seg_tgt_viz = visualize_frame(driving_video_org[0], seg_targets[0], alpha=viz_alpha) + seg_tgt_viz = visualize_frame( + driving_video_org[0], seg_targets[0], alpha=viz_alpha + ) output_images = out_video(predictions)