diff --git a/README.md b/README.md index 1992cb2..4b2648f 100644 --- a/README.md +++ b/README.md @@ -3,6 +3,10 @@ A ComfyUI-native node to run First Order Motion Model for Image Animation https://github.com/AliaksandrSiarohin/first-order-model +Now supports Face Swapping using Motion Supervised co-part Segmentation + +https://github.com/AliaksandrSiarohin/motion-cosegmentation + ## Workflow: [FOMM.json](FOMM.json) diff --git a/__init__.py b/__init__.py index 6c98dfa..cc2303e 100644 --- a/__init__.py +++ b/__init__.py @@ -1,11 +1,16 @@ -from .run import FOMM_Runner +from .run import FOMM_Runner, FOMM_Partswap, 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 Seg10 Chooser", } + __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] \ No newline at end of file diff --git a/config/vox-256-sem-10segments.yaml b/config/vox-256-sem-10segments.yaml new file mode 100644 index 0000000..5871c97 --- /dev/null +++ b/config/vox-256-sem-10segments.yaml @@ -0,0 +1,92 @@ +# Dataset parameters +dataset_params: + # Path to data, data can be stored in several formats: .mp4 or .gif videos, stacked .png images or folders with frames. + root_dir: data/vox-png + # Image shape, needed for staked .png format. + image_shape: [256, 256, 3] + # In case of VoxCeleb or TaiChi single video can be splitted in many chunks, or the maybe several videos for single person. + # In this case epoch can be a pass over different videos (if id_sampling=True) or over different chunks (if id_sampling=False) + # If the name of the video '12335#adsbf.mp4' the id is assumed to be 12335 + id_sampling: True + # Augmentation parameters see augmentation.py for all posible augmentations + augmentation_params: + flip_param: + horizontal_flip: True + time_flip: True + jitter_param: + brightness: 0.1 + contrast: 0.1 + saturation: 0.1 + hue: 0.1 + +# Defines model architecture +model_params: + common_params: + # Number of segments + num_segments: 10 + # Number of channels per image + num_channels: 3 + # Use only shift and no-affine part + estimate_affine_part: True + segmentation_module_params: + # Softmax temperature for shift heatmaps + temperature: 0.1 + # Number of features mutliplier + block_expansion: 32 + # Maximum allowed number of features + max_features: 1024 + # Number of block in Unet. Can be increased or decreased depending or resolution. + num_blocks: 5 + # Segmentations is predicted on smaller images for better performance, + # scale_factor=0.25 means that 256x256 image will be resized to 64x64 + scale_factor: 0.25 + reconstruction_module_params: + # Number of features mutliplier + block_expansion: 64 + # Maximum allowed number of features + max_features: 512 + # Number of downsampling blocks in Jonson architecture. + # Can be increased or decreased depending or resolution. + num_down_blocks: 2 + # Number of ResBlocks in Jonson architecture. + num_bottleneck_blocks: 6 + # Use visibility map or not + estimate_visibility: True + +# Parameters of training +train_params: + # Number of training epochs + num_epochs: 5 + # For better i/o performance when number of videos is small number of epochs can be multiplied by this number. + # Thus effectively with num_repeats=100 each epoch is 100 times larger. + num_repeats: 50 + # Learning rates + lr_segmentation: 2.0e-4 + lr_reconstruction: 2.0e-4 + batch_size: 20 + # Scales for perceptual pyramide loss. If scales = [1, 0.5, 0.25, 0.125] and image resolution is 256x256, + # than the loss will be computer on resolutions 256x256, 128x128, 64x64, 32x32. + scales: [1, 0.5, 0.25, 0.125] + # Save checkpoint this frequently. If checkpoint_freq=50, checkpoint will be saved every 50 epochs. + checkpoint_freq: 50 + # Parameters of transform for equivariance loss + transform_params: + # Sigma for affine part + sigma_affine: 0.05 + # Sigma for deformation part + sigma_tps: 0.005 + # Number of point in the deformation grid + points_tps: 5 + loss_weights: + equivariance: 10 + perceptual: [10, 10, 10, 10, 10] + + +# Visualization parameters +visualizer_params: + # Draw keypoints (shifts in affine transformations) of this size, increase or decrease depending on resolution + kp_size: 5 + # Draw white border around images + draw_border: True + # Color map for keypoints + colormap: 'gist_rainbow' diff --git a/config/vox-256-sem-15segments.yaml b/config/vox-256-sem-15segments.yaml new file mode 100644 index 0000000..ade4f4d --- /dev/null +++ b/config/vox-256-sem-15segments.yaml @@ -0,0 +1,92 @@ +# Dataset parameters +dataset_params: + # Path to data, data can be stored in several formats: .mp4 or .gif videos, stacked .png images or folders with frames. + root_dir: data/vox-png + # Image shape, needed for staked .png format. + image_shape: [256, 256, 3] + # In case of VoxCeleb or TaiChi single video can be splitted in many chunks, or the maybe several videos for single person. + # In this case epoch can be a pass over different videos (if id_sampling=True) or over different chunks (if id_sampling=False) + # If the name of the video '12335#adsbf.mp4' the id is assumed to be 12335 + id_sampling: True + # Augmentation parameters see augmentation.py for all posible augmentations + augmentation_params: + flip_param: + horizontal_flip: True + time_flip: True + jitter_param: + brightness: 0.1 + contrast: 0.1 + saturation: 0.1 + hue: 0.1 + +# Defines model architecture +model_params: + common_params: + # Number of segments + num_segments: 15 + # Number of channels per image + num_channels: 3 + # Use only shift and no-affine part + estimate_affine_part: True + segmentation_module_params: + # Softmax temperature for shift heatmaps + temperature: 0.1 + # Number of features mutliplier + block_expansion: 32 + # Maximum allowed number of features + max_features: 1024 + # Number of block in Unet. Can be increased or decreased depending or resolution. + num_blocks: 5 + # Segmentations is predicted on smaller images for better performance, + # scale_factor=0.25 means that 256x256 image will be resized to 64x64 + scale_factor: 0.25 + reconstruction_module_params: + # Number of features mutliplier + block_expansion: 64 + # Maximum allowed number of features + max_features: 512 + # Number of downsampling blocks in Jonson architecture. + # Can be increased or decreased depending or resolution. + num_down_blocks: 2 + # Number of ResBlocks in Jonson architecture. + num_bottleneck_blocks: 6 + # Use visibility map or not + estimate_visibility: True + +# Parameters of training +train_params: + # Number of training epochs + num_epochs: 5 + # For better i/o performance when number of videos is small number of epochs can be multiplied by this number. + # Thus effectively with num_repeats=100 each epoch is 100 times larger. + num_repeats: 50 + # Learning rates + lr_segmentation: 2.0e-4 + lr_reconstruction: 2.0e-4 + batch_size: 20 + # Scales for perceptual pyramide loss. If scales = [1, 0.5, 0.25, 0.125] and image resolution is 256x256, + # than the loss will be computer on resolutions 256x256, 128x128, 64x64, 32x32. + scales: [1, 0.5, 0.25, 0.125] + # Save checkpoint this frequently. If checkpoint_freq=50, checkpoint will be saved every 50 epochs. + checkpoint_freq: 50 + # Parameters of transform for equivariance loss + transform_params: + # Sigma for affine part + sigma_affine: 0.05 + # Sigma for deformation part + sigma_tps: 0.005 + # Number of point in the deformation grid + points_tps: 5 + loss_weights: + equivariance: 10 + perceptual: [10, 10, 10, 10, 10] + + +# Visualization parameters +visualizer_params: + # Draw keypoints (shifts in affine transformations) of this size, increase or decrease depending on resolution + kp_size: 5 + # Draw white border around images + draw_border: True + # Color map for keypoints + colormap: 'gist_rainbow' diff --git a/config/vox-256-sem-5segments.yaml b/config/vox-256-sem-5segments.yaml new file mode 100644 index 0000000..974628f --- /dev/null +++ b/config/vox-256-sem-5segments.yaml @@ -0,0 +1,92 @@ +# Dataset parameters +dataset_params: + # Path to data, data can be stored in several formats: .mp4 or .gif videos, stacked .png images or folders with frames. + root_dir: data/vox-png + # Image shape, needed for staked .png format. + image_shape: [256, 256, 3] + # In case of VoxCeleb or TaiChi single video can be splitted in many chunks, or the maybe several videos for single person. + # In this case epoch can be a pass over different videos (if id_sampling=True) or over different chunks (if id_sampling=False) + # If the name of the video '12335#adsbf.mp4' the id is assumed to be 12335 + id_sampling: True + # Augmentation parameters see augmentation.py for all posible augmentations + augmentation_params: + flip_param: + horizontal_flip: True + time_flip: True + jitter_param: + brightness: 0.1 + contrast: 0.1 + saturation: 0.1 + hue: 0.1 + +# Defines model architecture +model_params: + common_params: + # Number of segments + num_segments: 5 + # Number of channels per image + num_channels: 3 + # Use only shift and no-affine part + estimate_affine_part: True + segmentation_module_params: + # Softmax temperature for shift heatmaps + temperature: 0.1 + # Number of features mutliplier + block_expansion: 32 + # Maximum allowed number of features + max_features: 1024 + # Number of block in Unet. Can be increased or decreased depending or resolution. + num_blocks: 5 + # Segmentations is predicted on smaller images for better performance, + # scale_factor=0.25 means that 256x256 image will be resized to 64x64 + scale_factor: 0.25 + reconstruction_module_params: + # Number of features mutliplier + block_expansion: 64 + # Maximum allowed number of features + max_features: 512 + # Number of downsampling blocks in Jonson architecture. + # Can be increased or decreased depending or resolution. + num_down_blocks: 2 + # Number of ResBlocks in Jonson architecture. + num_bottleneck_blocks: 6 + # Use visibility map or not + estimate_visibility: True + +# Parameters of training +train_params: + # Number of training epochs + num_epochs: 5 + # For better i/o performance when number of videos is small number of epochs can be multiplied by this number. + # Thus effectively with num_repeats=100 each epoch is 100 times larger. + num_repeats: 50 + # Learning rates + lr_segmentation: 2.0e-4 + lr_reconstruction: 2.0e-4 + batch_size: 20 + # Scales for perceptual pyramide loss. If scales = [1, 0.5, 0.25, 0.125] and image resolution is 256x256, + # than the loss will be computer on resolutions 256x256, 128x128, 64x64, 32x32. + scales: [1, 0.5, 0.25, 0.125] + # Save checkpoint this frequently. If checkpoint_freq=50, checkpoint will be saved every 50 epochs. + checkpoint_freq: 50 + # Parameters of transform for equivariance loss + transform_params: + # Sigma for affine part + sigma_affine: 0.05 + # Sigma for deformation part + sigma_tps: 0.005 + # Number of point in the deformation grid + points_tps: 5 + loss_weights: + equivariance: 10 + perceptual: [10, 10, 10, 10, 10] + + +# Visualization parameters +visualizer_params: + # Draw keypoints (shifts in affine transformations) of this size, increase or decrease depending on resolution + kp_size: 5 + # Draw white border around images + draw_border: True + # Color map for keypoints + colormap: 'gist_rainbow' diff --git a/constants.py b/constants.py index d69b199..c53c7de 100644 --- a/constants.py +++ b/constants.py @@ -20,4 +20,12 @@ checkpoint_folder_link = "https://github.com/graphemecluster/first-order-model-d checkpoint_suffix = ".pth.tar" config_folder = "config" -checkpoint_folder = "checkpoints" \ No newline at end of file +checkpoint_folder = "checkpoints" + +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", +} + +default_partswap_model_name = "vox-10segments" \ No newline at end of file diff --git a/modules/reconstruction_module.py b/modules/reconstruction_module.py new file mode 100644 index 0000000..bd7cf7d --- /dev/null +++ b/modules/reconstruction_module.py @@ -0,0 +1,109 @@ +import torch +from torch import nn +import torch.nn.functional as F +from .util import ResBlock2d, SameBlock2d, UpBlock2d, DownBlock2d +from .util import make_coordinate_grid + + +class ReconstructionModule(nn.Module): + """ + Reconstruct target from source and segmentation of target and part transformations + """ + + def __init__(self, num_channels, block_expansion, max_features, num_down_blocks, + num_bottleneck_blocks, num_segments, estimate_visibility=False, **kwargs): + super(ReconstructionModule, self).__init__() + + self.first = SameBlock2d(num_channels, block_expansion, kernel_size=(7, 7), padding=(3, 3)) + + down_blocks = [] + for i in range(num_down_blocks): + in_features = min(max_features, block_expansion * (2 ** i)) + out_features = min(max_features, block_expansion * (2 ** (i + 1))) + down_blocks.append(DownBlock2d(in_features, out_features, kernel_size=(3, 3), padding=(1, 1))) + self.down_blocks = nn.ModuleList(down_blocks) + + up_blocks = [] + for i in range(num_down_blocks): + in_features = min(max_features, block_expansion * (2 ** (num_down_blocks - i))) + out_features = min(max_features, block_expansion * (2 ** (num_down_blocks - i - 1))) + up_blocks.append(UpBlock2d(in_features, out_features, kernel_size=(3, 3), padding=(1, 1))) + self.up_blocks = nn.ModuleList(up_blocks) + + self.bottleneck = torch.nn.Sequential() + in_features = min(max_features, block_expansion * (2 ** num_down_blocks)) + for i in range(num_bottleneck_blocks): + self.bottleneck.add_module('r' + str(i), ResBlock2d(in_features, kernel_size=(3, 3), padding=(1, 1))) + + self.final = nn.Conv2d(block_expansion, num_channels, kernel_size=(7, 7), padding=(3, 3)) + self.estimate_visibility = estimate_visibility + self.num_channels = num_channels + self.num_segments = num_segments + + def segment_motion(self, seg_target, seg_source): + bs, _, h, w = seg_target['segmentation'].shape + identity_grid = make_coordinate_grid((h, w), type=seg_source['shift'].type()) + identity_grid = identity_grid.view(1, 1, h, w, 2) + coordinate_grid = identity_grid - seg_target['shift'].view(bs, self.num_segments, 1, 1, 2) + if 'affine' in seg_target: + affine = torch.matmul(seg_source['affine'], torch.inverse(seg_target['affine'])) + affine = affine.unsqueeze(-3).unsqueeze(-3) + affine = affine.repeat(1, 1, h, w, 1, 1) + coordinate_grid = torch.matmul(affine, coordinate_grid.unsqueeze(-1)) + coordinate_grid = coordinate_grid.squeeze(-1) + + target_to_source = coordinate_grid + seg_source['shift'].view(bs, self.num_segments, 1, 1, 2) + + # adding background feature + identity_grid = identity_grid.repeat(bs, 1, 1, 1, 1) + return torch.cat([identity_grid, target_to_source], dim=1) + + def deform_input(self, inp, deformation): + _, h_old, w_old, _ = deformation.shape + _, _, h, w = inp.shape + if h_old != h or w_old != w: + deformation = deformation.permute(0, 3, 1, 2) + deformation = F.interpolate(deformation, size=(h, w), mode='bilinear') + deformation = deformation.permute(0, 2, 3, 1) + return F.grid_sample(inp, deformation) + + def forward(self, source_image, seg_target, seg_source): + # Encoding (downsampling) part + out = self.first(source_image) + for i in range(len(self.down_blocks)): + out = self.down_blocks[i](out) + + output_dict = {} + + # Computing segment motion + 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) + + # Deform according to motion + out = self.deform_input(out, deformation) + output_dict["deformed"] = self.deform_input(source_image, deformation) + + if self.estimate_visibility: + visibility = seg_source['segmentation'][:, 1:].sum(dim=1, keepdim=True) * \ + (1 - seg_target['segmentation'][:, 1:].sum(dim=1, keepdim=True).detach()) + visibility = 1 - visibility + + if out.shape[2] != visibility.shape[2] or out.shape[3] != visibility.shape[3]: + visibility = F.interpolate(visibility, size=out.shape[2:], mode='bilinear') + out = out * visibility + output_dict['visibility_maps'] = visibility + + 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/modules/segmentation_module.py b/modules/segmentation_module.py new file mode 100644 index 0000000..e6d432b --- /dev/null +++ b/modules/segmentation_module.py @@ -0,0 +1,78 @@ +from torch import nn +import torch +import torch.nn.functional as F +from .util import Hourglass, make_coordinate_grid, AntiAliasInterpolation2d + + +class SegmentationModule(nn.Module): + """ + Computing a segmentation map and affine transformations. + """ + + def __init__(self, block_expansion, num_segments, num_channels, max_features, + num_blocks, temperature, estimate_affine_part=False, scale_factor=1): + super(SegmentationModule, self).__init__() + + self.predictor = Hourglass(block_expansion, in_features=num_channels, + max_features=max_features, num_blocks=num_blocks) + self.num_segments = num_segments + self.shift = nn.Conv2d(in_channels=self.predictor.out_filters, out_channels=num_segments, kernel_size=(7, 7), + padding=(3, 3)) + + if estimate_affine_part: + self.affine = nn.Conv2d(in_channels=self.predictor.out_filters, + out_channels=4 * num_segments, kernel_size=(7, 7), padding=(3, 3)) + self.affine.weight.data.zero_() + self.affine.bias.data.copy_(torch.tensor([1, 0, 0, 1] * num_segments, dtype=torch.float)) + else: + self.affine = None + + self.segmentation = nn.Conv2d(in_channels=self.predictor.out_filters, + out_channels=(1 + num_segments), kernel_size=(7, 7), padding=(3, 3)) + + self.temperature = temperature + self.scale_factor = scale_factor + if self.scale_factor != 1: + self.down = AntiAliasInterpolation2d(num_channels, self.scale_factor) + + def gaussian2kp(self, heatmap): + """ + Extract the mean and the variance from a heatmap + """ + shape = heatmap.shape + heatmap = heatmap.unsqueeze(-1) + grid = make_coordinate_grid(shape[2:], heatmap.type()).unsqueeze_(0).unsqueeze_(0) + result = (heatmap * grid).sum(dim=(2, 3)) + + return result + + def forward(self, x): + if self.scale_factor != 1: + x = self.down(x) + + out = {} + feature_map = self.predictor(x) + out['segmentation'] = F.softmax(self.segmentation(feature_map), dim=1) + + prediction = self.shift(feature_map) + + final_shape = prediction.shape + heatmap = prediction.view(final_shape[0], final_shape[1], -1) + heatmap = F.softmax(heatmap / self.temperature, dim=2) + heatmap = heatmap.view(*final_shape) + + out['shift'] = self.gaussian2kp(heatmap) + + if self.affine is not None: + affine_map = self.affine(feature_map) + affine_map = affine_map.reshape(final_shape[0], self.num_segments, 4, final_shape[2], + final_shape[3]) + heatmap = heatmap.unsqueeze(2) + + affine = heatmap * affine_map + affine = affine.view(final_shape[0], final_shape[1], 4, -1) + affine = affine.sum(dim=-1) + affine = affine.view(affine.shape[0], affine.shape[1], 2, 2) + out['affine'] = affine + + return out diff --git a/part_swap.py b/part_swap.py new file mode 100644 index 0000000..5a14e25 --- /dev/null +++ b/part_swap.py @@ -0,0 +1,321 @@ +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/run.py b/run.py index b23fa25..399fa4e 100644 --- a/run.py +++ b/run.py @@ -9,7 +9,7 @@ from pathlib import Path import torch -from .constants import default_model_name, supported_model_names +from .constants import default_model_name, supported_model_names, config_folder from .utils import ( get_config_path, get_model_file_name, @@ -119,10 +119,235 @@ 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('_'))) + + +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_TYPES = ("STRING",) + + RETURN_NAMES = ("chosen_seg_indices",) + + 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) + seg_list = serialize_integers(seg_list) + print("Final seg_list:", seg_list) + return (seg_list,) + + +class FOMM_Partswap: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "source_image": ("IMAGE",), + "driving_video_input": ("IMAGE",), + "model_name": ( + list(partswap_model_config_dict.keys()), + {"default": default_partswap_model_name}, + ), + "frame_rate": ("FLOAT", {"default": 30.0}), + "blend_scale": ( + "FLOAT", + {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.1}, + ), + "use_source_seg": ( + "BOOLEAN", + {"default": True}, + ), + "hard_edges": ( + "BOOLEAN", + {"default": False}, + ), + "chosen_seg_indices": ("STRING", {}), + }, + "optional": {"audio": ("VHS_AUDIO",)}, + } + + RETURN_TYPES = ( + "IMAGE", + "IMAGE", + "IMAGE", + "VHS_AUDIO", + "FLOAT", + ) + RETURN_NAMES = ( + "seg_src_viz", + "seg_tgt_viz", + "images", + "audio", + "frame_rate", + ) + FUNCTION = "todo" + CATEGORY = "FirstOrderMM" + + def todo( + self, + source_image: torch.Tensor, + driving_video_input: torch.Tensor, + model_name: str, + frame_rate: float, + blend_scale: float, + use_source_seg: bool, + hard_edges: bool, + chosen_seg_indices: str, + audio=None, + ): + print(f"{type(source_image)=}") + print(f"{type(driving_video_input)=}") + print(f"{source_image.shape=}") + print(f"{driving_video_input.shape=}") + print(f"{type(audio)=}") + print(base_dir) + + config_path = ( + f"{base_dir}/{config_folder}/{partswap_model_config_dict[model_name]}" + ) + checkpoint_path = f"{base_dir}/{get_partswap_model_file_name(model_name)}" + + source_image = reshape_image(source_image, (256, 256)) + driving_video_org = reshape_image(driving_video_input, (256, 256)) + driving_video = driving_video_org.unsqueeze(0).permute(0, 2, 1, 3, 4) + + reconstruction_module, segmentation_module = load_partswap_checkpoint( + config_path, checkpoint_path, blend_scale, use_fomm=False + ) + print(f"Chosen {chosen_seg_indices}") + swap_indices = deserialize_integers(chosen_seg_indices) + print(f"Swapping {swap_indices}") + + params = { + "swap_indices": swap_indices, + "source_image": source_image, + "target_video": driving_video, + "reconstruction_module": reconstruction_module, + "segmentation_module": segmentation_module, + "use_source_seg": use_source_seg, + "hard_edges": hard_edges, + } + + seg_source, seg_targets, predictions = partswap_inference(**params) + + seg_src_viz = visualize_frame(source_image, seg_source) + seg_tgt_viz = visualize_frame(driving_video_org[0], seg_targets[0]) + + output_images = out_video(predictions) + + return ( + seg_src_viz, + seg_tgt_viz, + output_images, + audio, + frame_rate, + ) + + 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", }