diff --git a/.gitignore b/.gitignore index 7e54636..f42bb28 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,7 @@ # Stuff +articulate_module/models/*.pth + checkpoints/* !checkpoints/place_checkpoints_here.txt diff --git a/README.md b/README.md index 2f8f066..7791ae2 100644 --- a/README.md +++ b/README.md @@ -7,6 +7,10 @@ Now supports Face Swapping using Motion Supervised co-part Segmentation https://github.com/AliaksandrSiarohin/motion-cosegmentation +Now also supports Motion Representations for Articulated Animation + +https://github.com/snap-research/articulated-animation + ## Workflow: @@ -15,15 +19,21 @@ https://github.com/FuouM/ComfyUI-FirstOrderMM/assets/57849916/e8adca97-b53c-48ce ### FOMM -[FOMM.json](FOMM.json) +[FOMM.json](workflows/FOMM.json) -![FOMM Workflow](workflow.png) +![FOMM Workflow](workflows/workflow.png) ### Part Swap -[FOMM_PARTSWAP.json](FOMM_PARTSWAP.json) +[FOMM_PARTSWAP.json](workflows/FOMM_PARTSWAP.json) -![Partswap Workflow](workflow_fomm_partswap.png) +![Partswap Workflow](workflows/workflow_fomm_partswap.png) + +### Articulate + +[ARTICULATE.json](workflows/ARTICULATE.json) + +![Workflow Articulate](workflows/workflow_articulate.png) ## Arguments @@ -42,6 +52,10 @@ https://github.com/FuouM/ComfyUI-FirstOrderMM/assets/57849916/e8adca97-b53c-48ce * `use_face_parser`: For Seg-based models, may help with cleaning up residual background (should only use `15seg` with this). TODO: Additional cleanup face_parser masks. Should definitely be used for FOMM models * `viz_alpha`: Opacity of the segments in the visualization +### Articulate + +Doesn't need any + ## Installation 1. Clone the repo to `ComfyUI/custom_nodes/` @@ -98,3 +112,5 @@ face_parsing_model.py resnet18-5c106cde.pth 79999_iter.pth ``` + +For **Articulate**, download the model from [Pre-trained checkpoints](https://github.com/snap-research/articulated-animation?tab=readme-ov-file#pre-trained-checkpoints) section and place it here: `articulate_module\models\vox256.pth` \ No newline at end of file diff --git a/__init__.py b/__init__.py index ccb00b0..d85937d 100644 --- a/__init__.py +++ b/__init__.py @@ -4,6 +4,7 @@ from .run import ( FOMM_Seg5Chooser, FOMM_Seg10Chooser, FOMM_Seg15Chooser, + Articulate_Runner, ) NODE_CLASS_MAPPINGS = { @@ -12,6 +13,7 @@ NODE_CLASS_MAPPINGS = { "FOMM_Seg5Chooser": FOMM_Seg5Chooser, "FOMM_Seg10Chooser": FOMM_Seg10Chooser, "FOMM_Seg15Chooser": FOMM_Seg15Chooser, + "Articulate_Runner": Articulate_Runner, } NODE_DISPLAY_NAME_MAPPINGS = { @@ -20,6 +22,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "FOMM_Seg5Chooser": "FOMM Seg5 Chooser", "FOMM_Seg10Chooser": "FOMM Seg10 Chooser", "FOMM_Seg15Chooser": "FOMM Seg15 Chooser", + "Articulate_Runner": "Articulate Runner", } diff --git a/articulate_inference.py b/articulate_inference.py new file mode 100644 index 0000000..0d2c1a4 --- /dev/null +++ b/articulate_inference.py @@ -0,0 +1,200 @@ +import numpy as np +import torch +import tqdm +import yaml +from comfy.utils import ProgressBar +from scipy.spatial import ConvexHull + +from .articulate_module.avd_network import AVDNetwork +from .articulate_module.bg_motion_predictor import BGMotionPredictor +from .articulate_module.generator import Generator +from .articulate_module.region_predictor import RegionPredictor +from .sync_batchnorm.replicate import DataParallelWithCallback + + +def articulate_inference( + source_image, + driving_video: list, + config_path: str, + checkpoint_path: str, + estimate_affine=True, + pca_based=True, +): + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + with open(config_path) as f: + config = yaml.full_load(f) + + generator, region_predictor, bg_predictor, avd_network = init_models( + config, estimate_affine, pca_based + ) + generator = generator.to(device) + region_predictor = region_predictor.to(device) + bg_predictor = bg_predictor.to(device) + avd_network = avd_network.to(device) + + animate_params = config["animate_params"] + + load_cpk(checkpoint_path, generator, region_predictor, bg_predictor, avd_network) + + if torch.cuda.is_available(): + generator = DataParallelWithCallback(generator) + region_predictor = DataParallelWithCallback(region_predictor) + avd_network = DataParallelWithCallback(avd_network) + + generator.eval() + region_predictor.eval() + avd_network.eval() + + source_frame = source_image + driving = driving_video + predictions = [] + + num_frames = driving.shape[2] + pbar = ProgressBar(num_frames) + + with torch.no_grad(): + source_region_params = region_predictor(source_frame) + driving_region_params_initial = region_predictor(driving_video[:, :, 0]) + + for frame_idx in tqdm.tqdm(range(num_frames)): + driving_frame = driving[:, :, frame_idx] + driving_region_params = region_predictor(driving_frame) + new_region_params = get_animation_region_params( + source_region_params, + driving_region_params, + driving_region_params_initial, + mode=animate_params["mode"], + avd_network=avd_network, + ) + out = generator( + source_frame, + source_region_params=source_region_params, + driving_region_params=new_region_params, + ) + + out["driving_region_params"] = driving_region_params + out["source_region_params"] = source_region_params + out["new_region_params"] = new_region_params + + # visualization = Visualizer(**config["visualizer_params"]).visualize( + # source=source_frame, driving=driving_frame, out=out + # ) / 255.0 + + prediction = out["prediction"].data.cpu().numpy() + prediction = np.transpose(prediction, [0, 2, 3, 1]).squeeze(0) + + # visualizations.append(visualization) + predictions.append(prediction) + + pbar.update_absolute(frame_idx, num_frames) + + # print(f"{predictions[0].shape=}") + # print(f"{visualizations[0].shape=}") + + # return predictions, visualizations + return predictions + + +def init_models(config, estimate_affine, pca_based): + generator = Generator( + num_regions=config["model_params"]["num_regions"], + num_channels=config["model_params"]["num_channels"], + revert_axis_swap=config["model_params"]["revert_axis_swap"], + **config["model_params"]["generator_params"], + ) + + config["model_params"]["region_predictor_params"]["pca_based"] = pca_based + + region_predictor = RegionPredictor( + num_regions=config["model_params"]["num_regions"], + num_channels=config["model_params"]["num_channels"], + estimate_affine=estimate_affine, + **config["model_params"]["region_predictor_params"], + ) + + bg_predictor = BGMotionPredictor( + num_channels=config["model_params"]["num_channels"], + **config["model_params"]["bg_predictor_params"], + ) + + avd_network = AVDNetwork( + num_regions=config["model_params"]["num_regions"], + **config["model_params"]["avd_network_params"], + ) + + return generator, region_predictor, bg_predictor, avd_network + + +def load_cpk( + checkpoint_path, + generator=None, + region_predictor=None, + bg_predictor=None, + avd_network=None, + optimizer_reconstruction=None, + optimizer_avd=None, +): + checkpoint = torch.load(checkpoint_path) + # print(checkpoint.keys()) + if generator is not None: + generator.load_state_dict(checkpoint["generator"]) + if region_predictor is not None: + region_predictor.load_state_dict(checkpoint["region_predictor"]) + if bg_predictor is not None: + bg_predictor.load_state_dict(checkpoint["bg_predictor"]) + if avd_network is not None: + if "avd_network" in checkpoint: + avd_network.load_state_dict(checkpoint["avd_network"]) + + if optimizer_reconstruction is not None: + optimizer_reconstruction.load_state_dict(checkpoint["optimizer_reconstruction"]) + return checkpoint["epoch_reconstruction"] + + if optimizer_avd is not None: + if "optimizer_avd" in checkpoint: + optimizer_avd.load_state_dict(checkpoint["optimizer_avd"]) + return checkpoint["epoch_avd"] + return 0 + + return 0 + + +def get_animation_region_params( + source_region_params, + driving_region_params, + driving_region_params_initial, + mode="standard", + avd_network=None, + adapt_movement_scale=True, +): + assert mode in ["standard", "relative", "avd"] + new_region_params = {k: v for k, v in driving_region_params.items()} + if mode == "standard": + return new_region_params + elif mode == "relative": + source_area = ConvexHull( + source_region_params["shift"][0].data.cpu().numpy() + ).volume + driving_area = ConvexHull( + driving_region_params_initial["shift"][0].data.cpu().numpy() + ).volume + movement_scale = np.sqrt(source_area) / np.sqrt(driving_area) + + shift_diff = ( + driving_region_params["shift"] - driving_region_params_initial["shift"] + ) + shift_diff *= movement_scale + new_region_params["shift"] = shift_diff + source_region_params["shift"] + + affine_diff = torch.matmul( + driving_region_params["affine"], + torch.inverse(driving_region_params_initial["affine"]), + ) + new_region_params["affine"] = torch.matmul( + affine_diff, source_region_params["affine"] + ) + return new_region_params + elif mode == "avd": + new_region_params = avd_network(source_region_params, driving_region_params) + return new_region_params diff --git a/articulate_module/avd_network.py b/articulate_module/avd_network.py new file mode 100644 index 0000000..f1b12cb --- /dev/null +++ b/articulate_module/avd_network.py @@ -0,0 +1,90 @@ +""" +Copyright Snap Inc. 2021. This sample code is made available by Snap Inc. for informational purposes only. +No license, whether implied or otherwise, is granted in or to such code (including any rights to copy, modify, +publish, distribute and/or commercialize such code), unless you have entered into a separate agreement for such rights. +Such code is provided as-is, without warranty of any kind, express or implied, including any warranties of merchantability, +title, fitness for a particular purpose, non-infringement, or that such code is free of defects, errors or viruses. +In no event will Snap Inc. be liable for any damages or losses of any kind arising from the sample code or your use thereof. +""" +import torch +from torch import nn + + +class AVDNetwork(nn.Module): + """ + Animation via Disentanglement network + """ + + def __init__(self, num_regions, id_bottle_size=64, pose_bottle_size=64, revert_axis_swap=True): + super(AVDNetwork, self).__init__() + input_size = (2 + 4) * num_regions + self.num_regions = num_regions + self.revert_axis_swap = revert_axis_swap + + self.id_encoder = nn.Sequential( + nn.Linear(input_size, 256), + nn.BatchNorm1d(256), + nn.ReLU(inplace=True), + nn.Linear(256, 512), + nn.BatchNorm1d(512), + nn.ReLU(inplace=True), + nn.Linear(512, 1024), + nn.BatchNorm1d(1024), + nn.ReLU(inplace=True), + nn.Linear(1024, id_bottle_size) + ) + + self.pose_encoder = nn.Sequential( + nn.Linear(input_size, 256), + nn.BatchNorm1d(256), + nn.ReLU(inplace=True), + nn.Linear(256, 512), + nn.BatchNorm1d(512), + nn.ReLU(inplace=True), + nn.Linear(512, 1024), + nn.BatchNorm1d(1024), + nn.ReLU(inplace=True), + nn.Linear(1024, pose_bottle_size) + ) + + self.decoder = nn.Sequential( + nn.Linear(pose_bottle_size + id_bottle_size, 1024), + nn.BatchNorm1d(1024), + nn.ReLU(), + nn.Linear(1024, 512), + nn.BatchNorm1d(512), + nn.ReLU(), + nn.Linear(512, 256), + nn.BatchNorm1d(256), + nn.ReLU(), + nn.Linear(256, input_size) + ) + + @staticmethod + def region_params_to_emb(x): + mean = x['shift'] + jac = x['affine'] + emb = torch.cat([mean, jac.view(jac.shape[0], jac.shape[1], -1)], dim=-1) + emb = emb.view(emb.shape[0], -1) + return emb + + def emb_to_region_params(self, emb): + emb = emb.view(emb.shape[0], self.num_regions, 6) + mean = emb[:, :, :2] + jac = emb[:, :, 2:].view(emb.shape[0], emb.shape[1], 2, 2) + return {'shift': mean, 'affine': jac} + + def forward(self, x_id, x_pose, alpha=0.2): + if self.revert_axis_swap: + affine = torch.matmul(x_id['affine'], torch.inverse(x_pose['affine'])) + sign = torch.sign(affine[:, :, 0:1, 0:1]) + x_id = {'affine': x_id['affine'] * sign, 'shift': x_id['shift']} + + pose_emb = self.pose_encoder(self.region_params_to_emb(x_pose)) + id_emb = self.id_encoder(self.region_params_to_emb(x_id)) + + rec = self.decoder(torch.cat([pose_emb, id_emb], dim=1)) + + rec = self.emb_to_region_params(rec) + rec['covar'] = torch.matmul(rec['affine'], rec['affine'].permute(0, 1, 3, 2)) + return rec diff --git a/articulate_module/batchnorm.py b/articulate_module/batchnorm.py new file mode 100644 index 0000000..a5f8ff0 --- /dev/null +++ b/articulate_module/batchnorm.py @@ -0,0 +1,315 @@ +# -*- coding: utf-8 -*- +# File : batchnorm.py +# Author : Jiayuan Mao +# Email : maojiayuan@gmail.com +# Date : 27/01/2018 +# +# This file is part of Synchronized-BatchNorm-PyTorch. +# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch +# Distributed under MIT License. + +import collections + +# import torch +import torch.nn.functional as F + +from torch.nn.modules.batchnorm import _BatchNorm +from torch.nn.parallel._functions import ReduceAddCoalesced, Broadcast + +from .comm import SyncMaster + +__all__ = ['SynchronizedBatchNorm1d', 'SynchronizedBatchNorm2d', 'SynchronizedBatchNorm3d'] + + +def _sum_ft(tensor): + """sum over the first and last dimention""" + return tensor.sum(dim=0).sum(dim=-1) + + +def _unsqueeze_ft(tensor): + """add new dementions at the front and the tail""" + return tensor.unsqueeze(0).unsqueeze(-1) + + +_ChildMessage = collections.namedtuple('_ChildMessage', ['sum', 'ssum', 'sum_size']) +_MasterMessage = collections.namedtuple('_MasterMessage', ['sum', 'inv_std']) + + +class _SynchronizedBatchNorm(_BatchNorm): + def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True): + super(_SynchronizedBatchNorm, self).__init__(num_features, eps=eps, momentum=momentum, affine=affine) + + self._sync_master = SyncMaster(self._data_parallel_master) + + self._is_parallel = False + self._parallel_id = None + self._slave_pipe = None + + def forward(self, input): + # If it is not parallel computation or is in evaluation mode, use PyTorch's implementation. + if not (self._is_parallel and self.training): + return F.batch_norm( + input, self.running_mean, self.running_var, self.weight, self.bias, + self.training, self.momentum, self.eps) + + # Resize the input to (B, C, -1). + input_shape = input.size() + input = input.view(input.size(0), self.num_features, -1) + + # Compute the sum and square-sum. + sum_size = input.size(0) * input.size(2) + input_sum = _sum_ft(input) + input_ssum = _sum_ft(input ** 2) + + # Reduce-and-broadcast the statistics. + if self._parallel_id == 0: + mean, inv_std = self._sync_master.run_master(_ChildMessage(input_sum, input_ssum, sum_size)) + else: + mean, inv_std = self._slave_pipe.run_slave(_ChildMessage(input_sum, input_ssum, sum_size)) + + # Compute the output. + if self.affine: + # MJY:: Fuse the multiplication for speed. + output = (input - _unsqueeze_ft(mean)) * _unsqueeze_ft(inv_std * self.weight) + _unsqueeze_ft(self.bias) + else: + output = (input - _unsqueeze_ft(mean)) * _unsqueeze_ft(inv_std) + + # Reshape it. + return output.view(input_shape) + + def __data_parallel_replicate__(self, ctx, copy_id): + self._is_parallel = True + self._parallel_id = copy_id + + # parallel_id == 0 means master device. + if self._parallel_id == 0: + ctx.sync_master = self._sync_master + else: + self._slave_pipe = ctx.sync_master.register_slave(copy_id) + + def _data_parallel_master(self, intermediates): + """Reduce the sum and square-sum, compute the statistics, and broadcast it.""" + + # Always using same "device order" makes the ReduceAdd operation faster. + # Thanks to:: Tete Xiao (http://tetexiao.com/) + intermediates = sorted(intermediates, key=lambda i: i[1].sum.get_device()) + + to_reduce = [i[1][:2] for i in intermediates] + to_reduce = [j for i in to_reduce for j in i] # flatten + target_gpus = [i[1].sum.get_device() for i in intermediates] + + sum_size = sum([i[1].sum_size for i in intermediates]) + sum_, ssum = ReduceAddCoalesced.apply(target_gpus[0], 2, *to_reduce) + mean, inv_std = self._compute_mean_std(sum_, ssum, sum_size) + + broadcasted = Broadcast.apply(target_gpus, mean, inv_std) + + outputs = [] + for i, rec in enumerate(intermediates): + outputs.append((rec[0], _MasterMessage(*broadcasted[i*2:i*2+2]))) + + return outputs + + def _compute_mean_std(self, sum_, ssum, size): + """Compute the mean and standard-deviation with sum and square-sum. This method + also maintains the moving average on the master device.""" + assert size > 1, 'BatchNorm computes unbiased standard-deviation, which requires size > 1.' + mean = sum_ / size + sumvar = ssum - sum_ * mean + unbias_var = sumvar / (size - 1) + bias_var = sumvar / size + + self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * mean.data + self.running_var = (1 - self.momentum) * self.running_var + self.momentum * unbias_var.data + + return mean, bias_var.clamp(self.eps) ** -0.5 + + +class SynchronizedBatchNorm1d(_SynchronizedBatchNorm): + r"""Applies Synchronized Batch Normalization over a 2d or 3d input that is seen as a + mini-batch. + + .. math:: + + y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta + + This module differs from the built-in PyTorch BatchNorm1d as the mean and + standard-deviation are reduced across all devices during training. + + For example, when one uses `nn.DataParallel` to wrap the network during + training, PyTorch's implementation normalize the tensor on each device using + the statistics only on that device, which accelerated the computation and + is also easy to implement, but the statistics might be inaccurate. + Instead, in this synchronized version, the statistics will be computed + over all training samples distributed on multiple devices. + + Note that, for one-GPU or CPU-only case, this module behaves exactly same + as the built-in PyTorch implementation. + + The mean and standard-deviation are calculated per-dimension over + the mini-batches and gamma and beta are learnable parameter vectors + of size C (where C is the input size). + + During training, this layer keeps a running estimate of its computed mean + and variance. The running sum is kept with a default momentum of 0.1. + + During evaluation, this running mean/variance is used for normalization. + + Because the BatchNorm is done over the `C` dimension, computing statistics + on `(N, L)` slices, it's common terminology to call this Temporal BatchNorm + + Args: + num_features: num_features from an expected input of size + `batch_size x num_features [x width]` + eps: a value added to the denominator for numerical stability. + Default: 1e-5 + momentum: the value used for the running_mean and running_var + computation. Default: 0.1 + affine: a boolean value that when set to ``True``, gives the layer learnable + affine parameters. Default: ``True`` + + Shape: + - Input: :math:`(N, C)` or :math:`(N, C, L)` + - Output: :math:`(N, C)` or :math:`(N, C, L)` (same shape as input) + + Examples: + >>> # With Learnable Parameters + >>> m = SynchronizedBatchNorm1d(100) + >>> # Without Learnable Parameters + >>> m = SynchronizedBatchNorm1d(100, affine=False) + >>> input = torch.autograd.Variable(torch.randn(20, 100)) + >>> output = m(input) + """ + + def _check_input_dim(self, input): + if input.dim() != 2 and input.dim() != 3: + raise ValueError('expected 2D or 3D input (got {}D input)' + .format(input.dim())) + super(SynchronizedBatchNorm1d, self)._check_input_dim(input) + + +class SynchronizedBatchNorm2d(_SynchronizedBatchNorm): + r"""Applies Batch Normalization over a 4d input that is seen as a mini-batch + of 3d inputs + + .. math:: + + y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta + + This module differs from the built-in PyTorch BatchNorm2d as the mean and + standard-deviation are reduced across all devices during training. + + For example, when one uses `nn.DataParallel` to wrap the network during + training, PyTorch's implementation normalize the tensor on each device using + the statistics only on that device, which accelerated the computation and + is also easy to implement, but the statistics might be inaccurate. + Instead, in this synchronized version, the statistics will be computed + over all training samples distributed on multiple devices. + + Note that, for one-GPU or CPU-only case, this module behaves exactly same + as the built-in PyTorch implementation. + + The mean and standard-deviation are calculated per-dimension over + the mini-batches and gamma and beta are learnable parameter vectors + of size C (where C is the input size). + + During training, this layer keeps a running estimate of its computed mean + and variance. The running sum is kept with a default momentum of 0.1. + + During evaluation, this running mean/variance is used for normalization. + + Because the BatchNorm is done over the `C` dimension, computing statistics + on `(N, H, W)` slices, it's common terminology to call this Spatial BatchNorm + + Args: + num_features: num_features from an expected input of + size batch_size x num_features x height x width + eps: a value added to the denominator for numerical stability. + Default: 1e-5 + momentum: the value used for the running_mean and running_var + computation. Default: 0.1 + affine: a boolean value that when set to ``True``, gives the layer learnable + affine parameters. Default: ``True`` + + Shape: + - Input: :math:`(N, C, H, W)` + - Output: :math:`(N, C, H, W)` (same shape as input) + + Examples: + >>> # With Learnable Parameters + >>> m = SynchronizedBatchNorm2d(100) + >>> # Without Learnable Parameters + >>> m = SynchronizedBatchNorm2d(100, affine=False) + >>> input = torch.autograd.Variable(torch.randn(20, 100, 35, 45)) + >>> output = m(input) + """ + + def _check_input_dim(self, input): + if input.dim() != 4: + raise ValueError('expected 4D input (got {}D input)' + .format(input.dim())) + super(SynchronizedBatchNorm2d, self)._check_input_dim(input) + + +class SynchronizedBatchNorm3d(_SynchronizedBatchNorm): + r"""Applies Batch Normalization over a 5d input that is seen as a mini-batch + of 4d inputs + + .. math:: + + y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta + + This module differs from the built-in PyTorch BatchNorm3d as the mean and + standard-deviation are reduced across all devices during training. + + For example, when one uses `nn.DataParallel` to wrap the network during + training, PyTorch's implementation normalize the tensor on each device using + the statistics only on that device, which accelerated the computation and + is also easy to implement, but the statistics might be inaccurate. + Instead, in this synchronized version, the statistics will be computed + over all training samples distributed on multiple devices. + + Note that, for one-GPU or CPU-only case, this module behaves exactly same + as the built-in PyTorch implementation. + + The mean and standard-deviation are calculated per-dimension over + the mini-batches and gamma and beta are learnable parameter vectors + of size C (where C is the input size). + + During training, this layer keeps a running estimate of its computed mean + and variance. The running sum is kept with a default momentum of 0.1. + + During evaluation, this running mean/variance is used for normalization. + + Because the BatchNorm is done over the `C` dimension, computing statistics + on `(N, D, H, W)` slices, it's common terminology to call this Volumetric BatchNorm + or Spatio-temporal BatchNorm + + Args: + num_features: num_features from an expected input of + size batch_size x num_features x depth x height x width + eps: a value added to the denominator for numerical stability. + Default: 1e-5 + momentum: the value used for the running_mean and running_var + computation. Default: 0.1 + affine: a boolean value that when set to ``True``, gives the layer learnable + affine parameters. Default: ``True`` + + Shape: + - Input: :math:`(N, C, D, H, W)` + - Output: :math:`(N, C, D, H, W)` (same shape as input) + + Examples: + >>> # With Learnable Parameters + >>> m = SynchronizedBatchNorm3d(100) + >>> # Without Learnable Parameters + >>> m = SynchronizedBatchNorm3d(100, affine=False) + >>> input = torch.autograd.Variable(torch.randn(20, 100, 35, 45, 10)) + >>> output = m(input) + """ + + def _check_input_dim(self, input): + if input.dim() != 5: + raise ValueError('expected 5D input (got {}D input)' + .format(input.dim())) + super(SynchronizedBatchNorm3d, self)._check_input_dim(input) diff --git a/articulate_module/bg_motion_predictor.py b/articulate_module/bg_motion_predictor.py new file mode 100644 index 0000000..af9de2d --- /dev/null +++ b/articulate_module/bg_motion_predictor.py @@ -0,0 +1,57 @@ +""" +Copyright Snap Inc. 2021. This sample code is made available by Snap Inc. for informational purposes only. +No license, whether implied or otherwise, is granted in or to such code (including any rights to copy, modify, +publish, distribute and/or commercialize such code), unless you have entered into a separate agreement for such rights. +Such code is provided as-is, without warranty of any kind, express or implied, including any warranties of merchantability, +title, fitness for a particular purpose, non-infringement, or that such code is free of defects, errors or viruses. +In no event will Snap Inc. be liable for any damages or losses of any kind arising from the sample code or your use thereof. +""" + +from torch import nn +import torch +from .util import Encoder + + +class BGMotionPredictor(nn.Module): + """ + Module for background estimation, return single transformation, parametrized as 3x3 matrix. + """ + + def __init__(self, block_expansion, num_channels, max_features, num_blocks, bg_type='zero'): + super(BGMotionPredictor, self).__init__() + assert bg_type in ['zero', 'shift', 'affine', 'perspective'] + + self.bg_type = bg_type + if self.bg_type != 'zero': + self.encoder = Encoder(block_expansion, in_features=num_channels * 2, max_features=max_features, + num_blocks=num_blocks) + in_features = min(max_features, block_expansion * (2 ** num_blocks)) + if self.bg_type == 'perspective': + self.fc = nn.Linear(in_features, 8) + self.fc.weight.data.zero_() + self.fc.bias.data.copy_(torch.tensor([1, 0, 0, 0, 1, 0, 0, 0], dtype=torch.float)) + elif self.bg_type == 'affine': + self.fc = nn.Linear(in_features, 6) + self.fc.weight.data.zero_() + self.fc.bias.data.copy_(torch.tensor([1, 0, 0, 0, 1, 0], dtype=torch.float)) + elif self.bg_type == 'shift': + self.fc = nn.Linear(in_features, 2) + self.fc.weight.data.zero_() + self.fc.bias.data.copy_(torch.tensor([0, 0], dtype=torch.float)) + + def forward(self, source_image, driving_image): + bs = source_image.shape[0] + out = torch.eye(3).unsqueeze(0).repeat(bs, 1, 1).type(source_image.type()) + if self.bg_type != 'zero': + prediction = self.encoder(torch.cat([source_image, driving_image], dim=1)) + prediction = prediction[-1].mean(dim=(2, 3)) + prediction = self.fc(prediction) + if self.bg_type == 'shift': + out[:, :2, 2] = prediction + elif self.bg_type == 'affine': + out[:, :2, :] = prediction.view(bs, 2, 3) + elif self.bg_type == 'perspective': + out[:, :2, :] = prediction[:, :6].view(bs, 2, 3) + out[:, 2, :2] = prediction[:, 6:].view(bs, 2) + + return out diff --git a/articulate_module/comm.py b/articulate_module/comm.py new file mode 100644 index 0000000..922f8c4 --- /dev/null +++ b/articulate_module/comm.py @@ -0,0 +1,137 @@ +# -*- coding: utf-8 -*- +# File : comm.py +# Author : Jiayuan Mao +# Email : maojiayuan@gmail.com +# Date : 27/01/2018 +# +# This file is part of Synchronized-BatchNorm-PyTorch. +# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch +# Distributed under MIT License. + +import queue +import collections +import threading + +__all__ = ['FutureResult', 'SlavePipe', 'SyncMaster'] + + +class FutureResult(object): + """A thread-safe future implementation. Used only as one-to-one pipe.""" + + def __init__(self): + self._result = None + self._lock = threading.Lock() + self._cond = threading.Condition(self._lock) + + def put(self, result): + with self._lock: + assert self._result is None, 'Previous result has\'t been fetched.' + self._result = result + self._cond.notify() + + def get(self): + with self._lock: + if self._result is None: + self._cond.wait() + + res = self._result + self._result = None + return res + + +_MasterRegistry = collections.namedtuple('MasterRegistry', ['result']) +_SlavePipeBase = collections.namedtuple('_SlavePipeBase', ['identifier', 'queue', 'result']) + + +class SlavePipe(_SlavePipeBase): + """Pipe for master-slave communication.""" + + def run_slave(self, msg): + self.queue.put((self.identifier, msg)) + ret = self.result.get() + self.queue.put(True) + return ret + + +class SyncMaster(object): + """An abstract `SyncMaster` object. + + - During the replication, as the data parallel will trigger an callback of each module, all slave devices should + call `register(id)` and obtain an `SlavePipe` to communicate with the master. + - During the forward pass, master device invokes `run_master`, all messages from slave devices will be collected, + and passed to a registered callback. + - After receiving the messages, the master device should gather the information and determine to message passed + back to each slave devices. + """ + + def __init__(self, master_callback): + """ + + Args: + master_callback: a callback to be invoked after having collected messages from slave devices. + """ + self._master_callback = master_callback + self._queue = queue.Queue() + self._registry = collections.OrderedDict() + self._activated = False + + def __getstate__(self): + return {'master_callback': self._master_callback} + + def __setstate__(self, state): + self.__init__(state['master_callback']) + + def register_slave(self, identifier): + """ + Register an slave device. + + Args: + identifier: an identifier, usually is the device id. + + Returns: a `SlavePipe` object which can be used to communicate with the master device. + + """ + if self._activated: + assert self._queue.empty(), 'Queue is not clean before next initialization.' + self._activated = False + self._registry.clear() + future = FutureResult() + self._registry[identifier] = _MasterRegistry(future) + return SlavePipe(identifier, self._queue, future) + + def run_master(self, master_msg): + """ + Main entry for the master device in each forward pass. + The messages were first collected from each devices (including the master device), and then + an callback will be invoked to compute the message to be sent back to each devices + (including the master device). + + Args: + master_msg: the message that the master want to send to itself. This will be placed as the first + message when calling `master_callback`. For detailed usage, see `_SynchronizedBatchNorm` for an example. + + Returns: the message to be sent back to the master device. + + """ + self._activated = True + + intermediates = [(0, master_msg)] + for i in range(self.nr_slaves): + intermediates.append(self._queue.get()) + + results = self._master_callback(intermediates) + assert results[0][0] == 0, 'The first result should belongs to the master.' + + for i, res in results: + if i == 0: + continue + self._registry[i].result.put(res) + + for i in range(self.nr_slaves): + assert self._queue.get() is True + + return results[0][1] + + @property + def nr_slaves(self): + return len(self._registry) diff --git a/articulate_module/config/taichi256.yaml b/articulate_module/config/taichi256.yaml new file mode 100644 index 0000000..4dac9b1 --- /dev/null +++ b/articulate_module/config/taichi256.yaml @@ -0,0 +1,104 @@ +#Copyright Snap Inc. 2021. This sample code is made available by Snap Inc. for informational purposes only. +#No license, whether implied or otherwise, is granted in or to such code (including any rights to copy, modify, +#publish, distribute and/or commercialize such code), unless you have entered into a separate agreement for such rights. +#Such code is provided as-is, without warranty of any kind, express or implied, including any warranties of merchantability, +#title, fitness for a particular purpose, non-infringement, or that such code is free of defects, errors or viruses. +#In no event will Snap Inc. be liable for any damages or losses of any kind arising from the sample code or your use thereof. + + +dataset_params: + root_dir: data/taichi-png + frame_shape: null + id_sampling: True + pairs_list: null + augmentation_params: + flip_param: + horizontal_flip: True + time_flip: True + jitter_param: + brightness: 0.1 + contrast: 0.1 + saturation: 0.1 + hue: 0.1 + +model_params: + num_regions: 10 + num_channels: 3 + estimate_affine: True + revert_axis_swap: True + bg_predictor_params: + block_expansion: 32 + max_features: 1024 + num_blocks: 5 + bg_type: 'affine' + region_predictor_params: + temperature: 0.1 + block_expansion: 32 + max_features: 1024 + scale_factor: 0.25 + num_blocks: 5 + pca_based: True + pad: 0 + fast_svd: False + generator_params: + block_expansion: 64 + max_features: 512 + num_down_blocks: 2 + num_bottleneck_blocks: 6 + skips: True + pixelwise_flow_predictor_params: + block_expansion: 64 + max_features: 1024 + num_blocks: 5 + scale_factor: 0.25 + use_deformed_source: True + use_covar_heatmap: True + estimate_occlusion_map: True + avd_network_params: + id_bottle_size: 64 + pose_bottle_size: 64 + +train_params: + num_epochs: 100 + num_repeats: 150 + epoch_milestones: [60, 90] + lr: 2.0e-4 + batch_size: 14 + dataloader_workers: 6 + checkpoint_freq: 50 + scales: [1, 0.5, 0.25, 0.125] + transform_params: + sigma_affine: 0.05 + sigma_tps: 0.005 + points_tps: 5 + loss_weights: + perceptual: [10, 10, 10, 10, 10] + equivariance_shift: 10 + equivariance_affine: 10 + +train_avd_params: + num_epochs: 30 + num_repeats: 150 + batch_size: 256 + dataloader_workers: 12 + checkpoint_freq: 50 + epoch_milestones: [20, 25] + lr: 1.0e-3 + lambda_shift: 1 + lambda_affine: 1 + random_scale: 0.25 + +reconstruction_params: + num_videos: 1000 + format: '.mp4' + +animate_params: + num_pairs: 50 + format: '.mp4' + mode: 'avd' + +visualizer_params: + kp_size: 5 + draw_border: True + colormap: 'gist_rainbow' + region_bg_color: [1, 1, 1] diff --git a/articulate_module/config/ted-youtube384.yaml b/articulate_module/config/ted-youtube384.yaml new file mode 100644 index 0000000..478fc4a --- /dev/null +++ b/articulate_module/config/ted-youtube384.yaml @@ -0,0 +1,105 @@ +#Copyright Snap Inc. 2021. This sample code is made available by Snap Inc. for informational purposes only. +#No license, whether implied or otherwise, is granted in or to such code (including any rights to copy, modify, +#publish, distribute and/or commercialize such code), unless you have entered into a separate agreement for such rights. +#Such code is provided as-is, without warranty of any kind, express or implied, including any warranties of merchantability, +#title, fitness for a particular purpose, non-infringement, or that such code is free of defects, errors or viruses. +#In no event will Snap Inc. be liable for any damages or losses of any kind arising from the sample code or your use thereof. + + +dataset_params: + root_dir: data/ted-youtube-384 + frame_shape: null + id_sampling: True + pairs_list: null + augmentation_params: + flip_param: + horizontal_flip: True + time_flip: True + jitter_param: + brightness: 0.1 + contrast: 0.1 + saturation: 0.1 + hue: 0.1 + +model_params: + num_regions: 20 + num_channels: 3 + estimate_affine: True + revert_axis_swap: True + + bg_predictor_params: + block_expansion: 32 + max_features: 1024 + num_blocks: 5 + bg_type: 'affine' + region_predictor_params: + temperature: 0.1 + block_expansion: 32 + max_features: 1024 + scale_factor: 0.25 + num_blocks: 5 + pca_based: True + pad: 0 + generator_params: + block_expansion: 64 + max_features: 512 + num_down_blocks: 2 + num_bottleneck_blocks: 6 + skips: True + pixelwise_flow_predictor_params: + block_expansion: 64 + max_features: 1024 + num_blocks: 5 + scale_factor: 0.25 + use_deformed_source: True + use_covar_heatmap: True + estimate_occlusion_map: True + avd_network_params: + id_bottle_size: 64 + pose_bottle_size: 64 + +train_params: + num_epochs: 100 + num_repeats: 150 + epoch_milestones: [60, 90] + lr: 2.0e-4 + batch_size: 20 + dataloader_workers: 6 + checkpoint_freq: 50 + use_sync_bn: False + scales: [1, 0.5, 0.25, 0.125] + transform_params: + sigma_affine: 0.05 + sigma_tps: 0.005 + points_tps: 5 + loss_weights: + perceptual: [10, 10, 10, 10, 10] + equivariance_shift: 10 + equivariance_affine: 10 + +train_avd_params: + num_epochs: 30 + num_repeats: 500 + batch_size: 256 + dataloader_workers: 12 + checkpoint_freq: 50 + epoch_milestones: [20, 25] + lr: 1.0e-3 + lambda_shift: 1 + lambda_affine: 1 + random_scale: 0.25 + +reconstruction_params: + num_videos: 1000 + format: '.mp4' + +animate_params: + num_pairs: 50 + format: '.mp4' + mode: 'avd' + +visualizer_params: + kp_size: 5 + draw_border: True + colormap: 'gist_rainbow' + region_bg_color: [1, 1, 1] diff --git a/articulate_module/config/ted384.yaml b/articulate_module/config/ted384.yaml new file mode 100644 index 0000000..d076ee1 --- /dev/null +++ b/articulate_module/config/ted384.yaml @@ -0,0 +1,105 @@ +#Copyright Snap Inc. 2021. This sample code is made available by Snap Inc. for informational purposes only. +#No license, whether implied or otherwise, is granted in or to such code (including any rights to copy, modify, +#publish, distribute and/or commercialize such code), unless you have entered into a separate agreement for such rights. +#Such code is provided as-is, without warranty of any kind, express or implied, including any warranties of merchantability, +#title, fitness for a particular purpose, non-infringement, or that such code is free of defects, errors or viruses. +#In no event will Snap Inc. be liable for any damages or losses of any kind arising from the sample code or your use thereof. + + +dataset_params: + root_dir: data/TED384-v2 + frame_shape: null + id_sampling: True + pairs_list: null + augmentation_params: + flip_param: + horizontal_flip: True + time_flip: True + jitter_param: + brightness: 0.1 + contrast: 0.1 + saturation: 0.1 + hue: 0.1 + +model_params: + num_regions: 10 + num_channels: 3 + estimate_affine: True + revert_axis_swap: True + + bg_predictor_params: + block_expansion: 32 + max_features: 1024 + num_blocks: 5 + bg_type: 'affine' + region_predictor_params: + temperature: 0.1 + block_expansion: 32 + max_features: 1024 + scale_factor: 0.25 + num_blocks: 5 + pca_based: True + pad: 0 + fast_svd: False + generator_params: + block_expansion: 64 + max_features: 512 + num_down_blocks: 2 + num_bottleneck_blocks: 6 + skips: True + pixelwise_flow_predictor_params: + block_expansion: 64 + max_features: 1024 + num_blocks: 5 + scale_factor: 0.25 + use_deformed_source: True + use_covar_heatmap: True + estimate_occlusion_map: True + avd_network_params: + id_bottle_size: 64 + pose_bottle_size: 64 + +train_params: + num_epochs: 100 + num_repeats: 150 + epoch_milestones: [60, 90] + lr: 2.0e-4 + batch_size: 6 + dataloader_workers: 6 + checkpoint_freq: 50 + scales: [1, 0.5, 0.25, 0.125] + transform_params: + sigma_affine: 0.05 + sigma_tps: 0.005 + points_tps: 5 + loss_weights: + perceptual: [10, 10, 10, 10, 10] + equivariance_shift: 10 + equivariance_affine: 10 + +train_avd_params: + num_epochs: 30 + num_repeats: 150 + batch_size: 256 + dataloader_workers: 12 + checkpoint_freq: 50 + epoch_milestones: [20, 25] + lr: 1.0e-3 + lambda_shift: 1 + lambda_affine: 1 + random_scale: 0.25 + +reconstruction_params: + num_videos: 1000 + format: '.mp4' + +animate_params: + num_pairs: 50 + format: '.mp4' + mode: 'avd' + +visualizer_params: + kp_size: 5 + draw_border: True + colormap: 'gist_rainbow' + region_bg_color: [1, 1, 1] diff --git a/articulate_module/config/vox256.yaml b/articulate_module/config/vox256.yaml new file mode 100644 index 0000000..9a49924 --- /dev/null +++ b/articulate_module/config/vox256.yaml @@ -0,0 +1,203 @@ +#Copyright Snap Inc. 2021. This sample code is made available by Snap Inc. for informational purposes only. +#No license, whether implied or otherwise, is granted in or to such code (including any rights to copy, modify, +#publish, distribute and/or commercialize such code), unless you have entered into a separate agreement for such rights. +#Such code is provided as-is, without warranty of any kind, express or implied, including any warranties of merchantability, +#title, fitness for a particular purpose, non-infringement, or that such code is free of defects, errors or viruses. +#In no event will Snap Inc. be liable for any damages or losses of any kind arising from the sample code or your use thereof. + + +# Dataset parameters +# Each dataset should contain 2 folders train and test +# Each video can be represented as: +# - an image of concatenated frames +# - '.mp4' or '.gif' +# - folder with all frames from a specific video +dataset_params: + # Path to data, data can be stored in several formats: .mp4 or .gif videos, stacked .png images or folders with frames. + # Folder with frames is preferred format for training, since it is the fastest. + root_dir: data/vox-png + # Shape to resize all frames to, specify null if resizing is not needed + frame_shape: null + # In case of Vox 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 identities (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 + # List with pairs for animation, null for random pairs + pairs_list: null + # Augmentation parameters see augmentation.py for all possible 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 architecture of the model +model_params: + # Number of regions + num_regions: 10 + # Number of channels, for RGB image it is always 3 + num_channels: 3 + # Enable estimation of affine parameters for each region, + # set to False if only region centers (keypoints) need to be estimated + estimate_affine: True + # Svd can perform random axis swap between source and driving if singular values are close to each other + # Set to True to avoid axis swap between source and driving + revert_axis_swap: True + + # Parameters of background prediction network based on simple Unet-like encoder. + bg_predictor_params: + # Number of features multiplier + block_expansion: 32 + # Maximum allowed number of features + max_features: 1024 + # Number of block in the Encoder. + num_blocks: 5 + # Type of background movement model, select one from ['zero', 'shift', 'affine', 'perspective'] + bg_type: 'affine' + + # Parameters of the region prediction network based on Unet + region_predictor_params: + # Softmax temperature for heatmaps + temperature: 0.1 + # Number of features multiplier + block_expansion: 32 + # Maximum allowed number of features + max_features: 1024 + # Regions 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 + # Number of block in Unet. Can be increased or decreased depending or resolution. + num_blocks: 5 + # Either to use pca_based estimation of affine parameters of regression based + pca_based: True + # Either to use fast_svd (https://github.com/KinglittleQ/torch-batch-svd) or standard pytorch svd + # Fast svd may produce not meaningful regions if used along with revert_axis_swap + fast_svd: False + + # Parameters of Generator, based on Jonson architecture + generator_params: + # Number of features multiplier + block_expansion: 64 + # Maximum allowed number of features + max_features: 512 + # Number of down-sampling 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 + # To use skip connections or no. + skips: True + # Parameters of pixelwise flow predictor based on Unet + pixelwise_flow_predictor_params: + # Number of features multiplier + block_expansion: 64 + # Maximum allowed number of features + max_features: 1024 + # Number of block in Unet. Can be increased or decreased depending or resolution. + num_blocks: 5 + # Flow predictor operates on the smaller images for better performance, + # scale_factor=0.25 means that 256x256 image will be resized to 64x64 + scale_factor: 0.25 + # Set to True in order to use deformed source images using sparse flow + use_deformed_source: True + # Set to False in order to render region heatmaps with fixed covariance + # True for covariance estimate using region_predictor + use_covar_heatmap: True + # Set to False to disable occlusion mask estimation + estimate_occlusion_map: True + + # Parameter for animation-via-disentanglement (avd) network + avd_network_params: + # Bottleneck for identity branch + id_bottle_size: 64 + # Bottleneck for pose branch + pose_bottle_size: 64 + +# Parameters of training (reconstruction) +train_params: + # Number of training epochs, visualization is produced after each epoch. + num_epochs: 100 + # 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: 75 + # Drop learning rate 10 times after this epochs + epoch_milestones: [60, 90] + # Initial learning rate + lr: 2.0e-4 + # Batch size. (14 is batch size for one V100 gpu). + batch_size: 14 + # Either to use sync_bn or not, enabling sync_bn will significantly slow the training time + use_sync_bn: False + # Dataset preprocessing cpu workers + dataloader_workers: 6 + # Save checkpoint this frequently. If checkpoint_freq=50, checkpoint will be saved every 50 epochs. + checkpoint_freq: 50 + # 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] + # Parameters of transform for equivariance loss + transform_params: + sigma_affine: 0.05 + sigma_tps: 0.005 + points_tps: 5 + loss_weights: + # Weights for perceptual pyramide loss. Note that here you can only specify weight across the layer, and + # weights across the resolution will be the same. + perceptual: [10, 10, 10, 10, 10] + # Weights for equivariance loss. + equivariance_shift: 10 + equivariance_affine: 10 + +# Parameters of training (animation-via-disentanglement) +train_avd_params: + # Number of training epochs, visualization is produced after each epoch. + num_epochs: 30 + # 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: 2000 + # Batch size. + batch_size: 256 + # Dataset preprocessing cpu workers + dataloader_workers: 20 + # Save checkpoint this frequently. If checkpoint_freq=50, checkpoint will be saved every 50 epochs. + checkpoint_freq: 50 + # Drop learning rate 10 times after this epochs + epoch_milestones: [20, 25] + # Initial learning rate + lr: 1.0e-3 + # Maximum value for scaling augmentation, that emulate objects with different shape + random_scale: 0.25 + # Weights for equivariance loss. + lambda_shift: 1 + lambda_affine: 1 + +# Parameters of reconstruction mode +reconstruction_params: + # Maximum number of videos for reconstruction + num_videos: 1000 + # Format for visualization, note that results will be also stored in staked .png for evaluation. + format: '.mp4' + +# Parameters of animation +animate_params: + # Maximum number of pairs for animation, the pairs will be either taken from pairs_list or random. + num_pairs: 50 + # Format for visualization + format: '.mp4' + # Animation mode, select one from [avd, standard, relative] + mode: 'avd' + +# Parameters of visualization +visualizer_params: + # Size of keypoints + kp_size: 5 + # Draw border between images or not + draw_border: True + # Colormap for regions and keypoints visualization + colormap: 'gist_rainbow' + # Background color for region visualization + region_bg_color: [1, 1, 1] diff --git a/articulate_module/generator.py b/articulate_module/generator.py new file mode 100644 index 0000000..5b9dafe --- /dev/null +++ b/articulate_module/generator.py @@ -0,0 +1,126 @@ +""" +Copyright Snap Inc. 2021. This sample code is made available by Snap Inc. for informational purposes only. +No license, whether implied or otherwise, is granted in or to such code (including any rights to copy, modify, +publish, distribute and/or commercialize such code), unless you have entered into a separate agreement for such rights. +Such code is provided as-is, without warranty of any kind, express or implied, including any warranties of merchantability, +title, fitness for a particular purpose, non-infringement, or that such code is free of defects, errors or viruses. +In no event will Snap Inc. be liable for any damages or losses of any kind arising from the sample code or your use thereof. +""" + +import torch +from torch import nn +import torch.nn.functional as F +from .util import ResBlock2d, SameBlock2d, UpBlock2d, DownBlock2d +from .pixelwise_flow_predictor import PixelwiseFlowPredictor + + +class Generator(nn.Module): + """ + Generator that given source image and region parameters try to transform image according to movement trajectories + induced by region parameters. Generator follows Johnson architecture. + """ + + def __init__(self, num_channels, num_regions, block_expansion, max_features, num_down_blocks, + num_bottleneck_blocks, pixelwise_flow_predictor_params=None, skips=False, revert_axis_swap=True): + super(Generator, self).__init__() + + if pixelwise_flow_predictor_params is not None: + self.pixelwise_flow_predictor = PixelwiseFlowPredictor(num_regions=num_regions, num_channels=num_channels, + revert_axis_swap=revert_axis_swap, + **pixelwise_flow_predictor_params) + else: + self.pixelwise_flow_predictor = None + + 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.num_channels = num_channels + self.skips = skips + + @staticmethod + def deform_input(inp, optical_flow): + _, h_old, w_old, _ = optical_flow.shape + _, _, h, w = inp.shape + if h_old != h or w_old != w: + optical_flow = optical_flow.permute(0, 3, 1, 2) + optical_flow = F.interpolate(optical_flow, size=(h, w), mode='bilinear') + optical_flow = optical_flow.permute(0, 2, 3, 1) + return F.grid_sample(inp, optical_flow) + + def apply_optical(self, input_previous=None, input_skip=None, motion_params=None): + if motion_params is not None: + if 'occlusion_map' in motion_params: + occlusion_map = motion_params['occlusion_map'] + else: + occlusion_map = None + deformation = motion_params['optical_flow'] + input_skip = self.deform_input(input_skip, deformation) + + if occlusion_map is not None: + if input_skip.shape[2] != occlusion_map.shape[2] or input_skip.shape[3] != occlusion_map.shape[3]: + occlusion_map = F.interpolate(occlusion_map, size=input_skip.shape[2:], mode='bilinear') + if input_previous is not None: + input_skip = input_skip * occlusion_map + input_previous * (1 - occlusion_map) + else: + input_skip = input_skip * occlusion_map + out = input_skip + else: + out = input_previous if input_previous is not None else input_skip + return out + + def forward(self, source_image, driving_region_params, source_region_params, bg_params=None): + out = self.first(source_image) + skips = [out] + for i in range(len(self.down_blocks)): + out = self.down_blocks[i](out) + skips.append(out) + + output_dict = {} + if self.pixelwise_flow_predictor is not None: + motion_params = self.pixelwise_flow_predictor(source_image=source_image, + driving_region_params=driving_region_params, + source_region_params=source_region_params, + bg_params=bg_params) + output_dict["deformed"] = self.deform_input(source_image, motion_params['optical_flow']) + if 'occlusion_map' in motion_params: + output_dict['occlusion_map'] = motion_params['occlusion_map'] + else: + motion_params = None + + out = self.apply_optical(input_previous=None, input_skip=out, motion_params=motion_params) + + out = self.bottleneck(out) + for i in range(len(self.up_blocks)): + if self.skips: + out = self.apply_optical(input_skip=skips[-(i + 1)], input_previous=out, motion_params=motion_params) + out = self.up_blocks[i](out) + if self.skips: + out = self.apply_optical(input_skip=skips[0], input_previous=out, motion_params=motion_params) + out = self.final(out) + out = F.sigmoid(out) + + if self.skips: + out = self.apply_optical(input_skip=source_image, input_previous=out, motion_params=motion_params) + + output_dict["prediction"] = out + + return output_dict diff --git a/articulate_module/models/.gitkeep b/articulate_module/models/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/articulate_module/pixelwise_flow_predictor.py b/articulate_module/pixelwise_flow_predictor.py new file mode 100644 index 0000000..6c83cc7 --- /dev/null +++ b/articulate_module/pixelwise_flow_predictor.py @@ -0,0 +1,137 @@ +""" +Copyright Snap Inc. 2021. This sample code is made available by Snap Inc. for informational purposes only. +No license, whether implied or otherwise, is granted in or to such code (including any rights to copy, modify, +publish, distribute and/or commercialize such code), unless you have entered into a separate agreement for such rights. +Such code is provided as-is, without warranty of any kind, express or implied, including any warranties of merchantability, +title, fitness for a particular purpose, non-infringement, or that such code is free of defects, errors or viruses. +In no event will Snap Inc. be liable for any damages or losses of any kind arising from the sample code or your use thereof. +""" + +from torch import nn +import torch.nn.functional as F +import torch +from .util import Hourglass, AntiAliasInterpolation2d, make_coordinate_grid, region2gaussian +from .util import to_homogeneous, from_homogeneous + + +class PixelwiseFlowPredictor(nn.Module): + """ + Module that predicts a pixelwise flow from sparse motion representation given by + source_region_params and driving_region_params + """ + + def __init__(self, block_expansion, num_blocks, max_features, num_regions, num_channels, + estimate_occlusion_map=False, scale_factor=1, region_var=0.01, + use_covar_heatmap=False, use_deformed_source=True, revert_axis_swap=False): + super(PixelwiseFlowPredictor, self).__init__() + self.hourglass = Hourglass(block_expansion=block_expansion, + in_features=(num_regions + 1) * (num_channels * use_deformed_source + 1), + max_features=max_features, num_blocks=num_blocks) + + self.mask = nn.Conv2d(self.hourglass.out_filters, num_regions + 1, kernel_size=(7, 7), padding=(3, 3)) + + if estimate_occlusion_map: + self.occlusion = nn.Conv2d(self.hourglass.out_filters, 1, kernel_size=(7, 7), padding=(3, 3)) + else: + self.occlusion = None + + self.num_regions = num_regions + self.scale_factor = scale_factor + self.region_var = region_var + self.use_covar_heatmap = use_covar_heatmap + self.use_deformed_source = use_deformed_source + self.revert_axis_swap = revert_axis_swap + + if self.scale_factor != 1: + self.down = AntiAliasInterpolation2d(num_channels, self.scale_factor) + + def create_heatmap_representations(self, source_image, driving_region_params, source_region_params): + """ + Eq 6. in the paper H_k(z) + """ + spatial_size = source_image.shape[2:] + covar = self.region_var if not self.use_covar_heatmap else driving_region_params['covar'] + gaussian_driving = region2gaussian(driving_region_params['shift'], covar=covar, spatial_size=spatial_size) + covar = self.region_var if not self.use_covar_heatmap else source_region_params['covar'] + gaussian_source = region2gaussian(source_region_params['shift'], covar=covar, spatial_size=spatial_size) + + heatmap = gaussian_driving - gaussian_source + + # adding background feature + zeros = torch.zeros(heatmap.shape[0], 1, spatial_size[0], spatial_size[1]) + heatmap = torch.cat([zeros.type(heatmap.type()), heatmap], dim=1) + heatmap = heatmap.unsqueeze(2) + return heatmap + + def create_sparse_motions(self, source_image, driving_region_params, source_region_params, bg_params=None): + bs, _, h, w = source_image.shape + identity_grid = make_coordinate_grid((h, w), type=source_region_params['shift'].type()) + identity_grid = identity_grid.view(1, 1, h, w, 2) + coordinate_grid = identity_grid - driving_region_params['shift'].view(bs, self.num_regions, 1, 1, 2) + if 'affine' in driving_region_params: + affine = torch.matmul(source_region_params['affine'], torch.inverse(driving_region_params['affine'])) + if self.revert_axis_swap: + affine = affine * torch.sign(affine[:, :, 0:1, 0:1]) + 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) + + driving_to_source = coordinate_grid + source_region_params['shift'].view(bs, self.num_regions, 1, 1, 2) + + # adding background feature + if bg_params is None: + bg_grid = identity_grid.repeat(bs, 1, 1, 1, 1) + else: + bg_grid = identity_grid.repeat(bs, 1, 1, 1, 1) + bg_grid = to_homogeneous(bg_grid) + bg_grid = torch.matmul(bg_params.view(bs, 1, 1, 1, 3, 3), bg_grid.unsqueeze(-1)).squeeze(-1) + bg_grid = from_homogeneous(bg_grid) + + sparse_motions = torch.cat([bg_grid, driving_to_source], dim=1) + + return sparse_motions + + def create_deformed_source_image(self, source_image, sparse_motions): + bs, _, h, w = source_image.shape + source_repeat = source_image.unsqueeze(1).unsqueeze(1).repeat(1, self.num_regions + 1, 1, 1, 1, 1) + source_repeat = source_repeat.view(bs * (self.num_regions + 1), -1, h, w) + sparse_motions = sparse_motions.view((bs * (self.num_regions + 1), h, w, -1)) + sparse_deformed = F.grid_sample(source_repeat, sparse_motions) + sparse_deformed = sparse_deformed.view((bs, self.num_regions + 1, -1, h, w)) + return sparse_deformed + + def forward(self, source_image, driving_region_params, source_region_params, bg_params=None): + 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, driving_region_params, + source_region_params) + sparse_motion = self.create_sparse_motions(source_image, driving_region_params, + source_region_params, bg_params=bg_params) + deformed_source = self.create_deformed_source_image(source_image, sparse_motion) + if self.use_deformed_source: + predictor_input = torch.cat([heatmap_representation, deformed_source], dim=2) + else: + predictor_input = heatmap_representation + predictor_input = predictor_input.view(bs, -1, h, w) + + prediction = self.hourglass(predictor_input) + + mask = self.mask(prediction) + mask = F.softmax(mask, dim=1) + 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['optical_flow'] = deformation + + if self.occlusion: + occlusion_map = torch.sigmoid(self.occlusion(prediction)) + out_dict['occlusion_map'] = occlusion_map + + return out_dict diff --git a/articulate_module/region_predictor.py b/articulate_module/region_predictor.py new file mode 100644 index 0000000..c98bb58 --- /dev/null +++ b/articulate_module/region_predictor.py @@ -0,0 +1,121 @@ +""" +Copyright Snap Inc. 2021. This sample code is made available by Snap Inc. for informational purposes only. +No license, whether implied or otherwise, is granted in or to such code (including any rights to copy, modify, +publish, distribute and/or commercialize such code), unless you have entered into a separate agreement for such rights. +Such code is provided as-is, without warranty of any kind, express or implied, including any warranties of merchantability, +title, fitness for a particular purpose, non-infringement, or that such code is free of defects, errors or viruses. +In no event will Snap Inc. be liable for any damages or losses of any kind arising from the sample code or your use thereof. +""" + +from torch import nn +import torch +import torch.nn.functional as F +from .util import Hourglass, make_coordinate_grid, AntiAliasInterpolation2d +from torch.linalg import svd as svd_torch + +def svd(covar, fast=False): + # if fast: + # from torch_batch_svd import svd as fast_svd + # return fast_svd(covar) + # else: + + # u, s, v = torch.svd(covar.cpu()) + # u, s, vh = svd_torch(covar.cpu()) + u, s, vh = svd_torch(covar) + v = vh.mH + s = s.to(covar.device) + u = u.to(covar.device) + v = v.to(covar.device) + return u, s, v + + +class RegionPredictor(nn.Module): + """ + Region estimating. Estimate affine parameters of the region. + """ + + def __init__(self, block_expansion, num_regions, num_channels, max_features, + num_blocks, temperature, estimate_affine=False, scale_factor=1, + pca_based=False, fast_svd=False, pad=3): + super(RegionPredictor, self).__init__() + self.predictor = Hourglass(block_expansion, in_features=num_channels, + max_features=max_features, num_blocks=num_blocks) + + self.regions = nn.Conv2d(in_channels=self.predictor.out_filters, out_channels=num_regions, kernel_size=(7, 7), + padding=pad) + + # FOMM-like regression based representation + if estimate_affine and not pca_based: + self.jacobian = nn.Conv2d(in_channels=self.predictor.out_filters, + out_channels=4, kernel_size=(7, 7), padding=pad) + self.jacobian.weight.data.zero_() + self.jacobian.bias.data.copy_(torch.tensor([1, 0, 0, 1], dtype=torch.float)) + else: + self.jacobian = None + + self.temperature = temperature + self.scale_factor = scale_factor + self.pca_based = pca_based + self.fast_svd = fast_svd + + if self.scale_factor != 1: + self.down = AntiAliasInterpolation2d(num_channels, self.scale_factor) + + def region2affine(self, region): + shape = region.shape + region = region.unsqueeze(-1) + grid = make_coordinate_grid(shape[2:], region.type()).unsqueeze_(0).unsqueeze_(0) + mean = (region * grid).sum(dim=(2, 3)) + + region_params = {'shift': mean} + + if self.pca_based: + mean_sub = grid - mean.unsqueeze(-2).unsqueeze(-2) + covar = torch.matmul(mean_sub.unsqueeze(-1), mean_sub.unsqueeze(-2)) + covar = covar * region.unsqueeze(-1) + covar = covar.sum(dim=(2, 3)) + region_params['covar'] = covar + + return region_params + + def forward(self, x): + if self.scale_factor != 1: + x = self.down(x) + + feature_map = self.predictor(x) + prediction = self.regions(feature_map) + + final_shape = prediction.shape + region = prediction.view(final_shape[0], final_shape[1], -1) + region = F.softmax(region / self.temperature, dim=2) + region = region.view(*final_shape) + + region_params = self.region2affine(region) + region_params['heatmap'] = region + + # Regression-based estimation + if self.jacobian is not None: + jacobian_map = self.jacobian(feature_map) + jacobian_map = jacobian_map.reshape(final_shape[0], 1, 4, final_shape[2], + final_shape[3]) + region = region.unsqueeze(2) + + jacobian = region * jacobian_map + jacobian = jacobian.view(final_shape[0], final_shape[1], 4, -1) + jacobian = jacobian.sum(dim=-1) + jacobian = jacobian.view(jacobian.shape[0], jacobian.shape[1], 2, 2) + region_params['affine'] = jacobian + region_params['covar'] = torch.matmul(jacobian, jacobian.permute(0, 1, 3, 2)) + elif self.pca_based: + covar = region_params['covar'] + shape = covar.shape + covar = covar.view(-1, 2, 2) + u, s, v = svd(covar, self.fast_svd) + d = torch.diag_embed(s ** 0.5) + sqrt = torch.matmul(u, d) + sqrt = sqrt.view(*shape) + region_params['affine'] = sqrt + region_params['u'] = u + region_params['d'] = d + + return region_params diff --git a/articulate_module/util.py b/articulate_module/util.py new file mode 100644 index 0000000..4bc1fb7 --- /dev/null +++ b/articulate_module/util.py @@ -0,0 +1,271 @@ +""" +Copyright Snap Inc. 2021. This sample code is made available by Snap Inc. for informational purposes only. +No license, whether implied or otherwise, is granted in or to such code (including any rights to copy, modify, +publish, distribute and/or commercialize such code), unless you have entered into a separate agreement for such rights. +Such code is provided as-is, without warranty of any kind, express or implied, including any warranties of merchantability, +title, fitness for a particular purpose, non-infringement, or that such code is free of defects, errors or viruses. +In no event will Snap Inc. be liable for any damages or losses of any kind arising from the sample code or your use thereof. +""" + +from torch import nn + +import torch.nn.functional as F +import torch +from .batchnorm import SynchronizedBatchNorm2d as BatchNorm2d + + +def region2gaussian(center, covar, spatial_size): + """ + Transform a region parameters into gaussian like heatmap + """ + mean = center + + coordinate_grid = make_coordinate_grid(spatial_size, mean.type()) + number_of_leading_dimensions = len(mean.shape) - 1 + shape = (1,) * number_of_leading_dimensions + coordinate_grid.shape + coordinate_grid = coordinate_grid.view(*shape) + repeats = mean.shape[:number_of_leading_dimensions] + (1, 1, 1) + coordinate_grid = coordinate_grid.repeat(*repeats) + + # Preprocess kp shape + shape = mean.shape[:number_of_leading_dimensions] + (1, 1, 2) + mean = mean.view(*shape) + + mean_sub = (coordinate_grid - mean) + if type(covar) == float: + out = torch.exp(-0.5 * (mean_sub ** 2).sum(-1) / covar) + else: + shape = mean.shape[:number_of_leading_dimensions] + (1, 1, 2, 2) + covar_inverse = torch.inverse(covar).view(*shape) + under_exp = torch.matmul(torch.matmul(mean_sub.unsqueeze(-2), covar_inverse), mean_sub.unsqueeze(-1)) + out = torch.exp(-0.5 * under_exp.sum(dim=(-1, -2))) + + return out + + +def make_coordinate_grid(spatial_size, type): + """ + Create a meshgrid [-1,1] x [-1,1] of given spatial_size. + """ + h, w = spatial_size + x = torch.arange(w).type(type) + y = torch.arange(h).type(type) + + x = (2 * (x / (w - 1)) - 1) + y = (2 * (y / (h - 1)) - 1) + + yy = y.view(-1, 1).repeat(1, w) + xx = x.view(1, -1).repeat(h, 1) + + meshed = torch.cat([xx.unsqueeze_(2), yy.unsqueeze_(2)], 2) + + return meshed + + +class ResBlock2d(nn.Module): + """ + Res block, preserve spatial resolution. + """ + + def __init__(self, in_features, kernel_size, padding): + super(ResBlock2d, self).__init__() + self.conv1 = nn.Conv2d(in_channels=in_features, out_channels=in_features, kernel_size=kernel_size, + padding=padding) + self.conv2 = nn.Conv2d(in_channels=in_features, out_channels=in_features, kernel_size=kernel_size, + padding=padding) + self.norm1 = BatchNorm2d(in_features, affine=True) + self.norm2 = BatchNorm2d(in_features, affine=True) + + def forward(self, x): + out = self.norm1(x) + out = F.relu(out) + out = self.conv1(out) + out = self.norm2(out) + out = F.relu(out) + out = self.conv2(out) + out += x + return out + + +class UpBlock2d(nn.Module): + """ + Upsampling block for use in decoder. + """ + + def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1): + super(UpBlock2d, self).__init__() + + self.conv = nn.Conv2d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size, + padding=padding, groups=groups) + self.norm = BatchNorm2d(out_features, affine=True) + + def forward(self, x): + out = F.interpolate(x, scale_factor=2) + out = self.conv(out) + out = self.norm(out) + out = F.relu(out) + return out + + +class DownBlock2d(nn.Module): + """ + Downsampling block for use in encoder. + """ + + def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1): + super(DownBlock2d, self).__init__() + self.conv = nn.Conv2d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size, + padding=padding, groups=groups) + self.norm = BatchNorm2d(out_features, affine=True) + self.pool = nn.AvgPool2d(kernel_size=(2, 2)) + + def forward(self, x): + out = self.conv(x) + out = self.norm(out) + out = F.relu(out) + out = self.pool(out) + return out + + +class SameBlock2d(nn.Module): + """ + Simple block, preserve spatial resolution. + """ + + def __init__(self, in_features, out_features, groups=1, kernel_size=3, padding=1): + super(SameBlock2d, self).__init__() + self.conv = nn.Conv2d(in_channels=in_features, out_channels=out_features, + kernel_size=kernel_size, padding=padding, groups=groups) + self.norm = BatchNorm2d(out_features, affine=True) + + def forward(self, x): + out = self.conv(x) + out = self.norm(out) + out = F.relu(out) + return out + + +class Encoder(nn.Module): + """ + Hourglass Encoder + """ + + def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256): + super(Encoder, self).__init__() + + down_blocks = [] + for i in range(num_blocks): + down_blocks.append(DownBlock2d(in_features if i == 0 else min(max_features, block_expansion * (2 ** i)), + min(max_features, block_expansion * (2 ** (i + 1))), + kernel_size=3, padding=1)) + self.down_blocks = nn.ModuleList(down_blocks) + + def forward(self, x): + outs = [x] + for down_block in self.down_blocks: + outs.append(down_block(outs[-1])) + return outs + + +class Decoder(nn.Module): + """ + Hourglass Decoder + """ + + def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256): + super(Decoder, self).__init__() + + up_blocks = [] + + for i in range(num_blocks)[::-1]: + in_filters = (1 if i == num_blocks - 1 else 2) * min(max_features, block_expansion * (2 ** (i + 1))) + out_filters = min(max_features, block_expansion * (2 ** i)) + up_blocks.append(UpBlock2d(in_filters, out_filters, kernel_size=3, padding=1)) + + self.up_blocks = nn.ModuleList(up_blocks) + self.out_filters = block_expansion + in_features + + def forward(self, x): + out = x.pop() + for up_block in self.up_blocks: + out = up_block(out) + skip = x.pop() + out = torch.cat([out, skip], dim=1) + return out + + +class Hourglass(nn.Module): + """ + Hourglass architecture. + """ + + def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256): + super(Hourglass, self).__init__() + self.encoder = Encoder(block_expansion, in_features, num_blocks, max_features) + self.decoder = Decoder(block_expansion, in_features, num_blocks, max_features) + self.out_filters = self.decoder.out_filters + + def forward(self, x): + return self.decoder(self.encoder(x)) + + +class AntiAliasInterpolation2d(nn.Module): + """ + Band-limited downsampling, for better preservation of the input signal. + """ + + def __init__(self, channels, scale): + super(AntiAliasInterpolation2d, self).__init__() + sigma = (1 / scale - 1) / 2 + kernel_size = 2 * round(sigma * 4) + 1 + self.ka = kernel_size // 2 + self.kb = self.ka - 1 if kernel_size % 2 == 0 else self.ka + + kernel_size = [kernel_size, kernel_size] + sigma = [sigma, sigma] + # The gaussian kernel is the product of the + # gaussian function of each dimension. + kernel = 1 + meshgrids = torch.meshgrid( + [ + torch.arange(size, dtype=torch.float32) + for size in kernel_size + ] + ) + for size, std, mgrid in zip(kernel_size, sigma, meshgrids): + mean = (size - 1) / 2 + kernel *= torch.exp(-(mgrid - mean) ** 2 / (2 * std ** 2)) + + # Make sure sum of values in gaussian kernel equals 1. + kernel = kernel / torch.sum(kernel) + # Reshape to depthwise convolutional weight + kernel = kernel.view(1, 1, *kernel.size()) + kernel = kernel.repeat(channels, *[1] * (kernel.dim() - 1)) + + self.register_buffer('weight', kernel) + self.groups = channels + self.scale = scale + inv_scale = 1 / scale + self.int_inv_scale = int(inv_scale) + + def forward(self, input): + if self.scale == 1.0: + return input + + out = F.pad(input, (self.ka, self.kb, self.ka, self.kb)) + out = F.conv2d(out, weight=self.weight, groups=self.groups) + out = out[:, :, ::self.int_inv_scale, ::self.int_inv_scale] + + return out + + +def to_homogeneous(coordinates): + ones_shape = list(coordinates.shape) + ones_shape[-1] = 1 + ones = torch.ones(ones_shape).type(coordinates.type()) + + return torch.cat([coordinates, ones], dim=-1) + + +def from_homogeneous(coordinates): + return coordinates[..., :2] / coordinates[..., 2:3] diff --git a/articulate_module/visualizer.py b/articulate_module/visualizer.py new file mode 100644 index 0000000..7406229 --- /dev/null +++ b/articulate_module/visualizer.py @@ -0,0 +1,145 @@ +import matplotlib.pyplot as plt +import numpy as np +import torch.nn.functional as F +from skimage.draw import circle_perimeter as circle + + +class Visualizer: + def __init__( + self, + kp_size=5, + draw_border=False, + colormap="gist_rainbow", + region_bg_color=(0, 0, 0), + ): + self.kp_size = kp_size + self.draw_border = draw_border + self.colormap = plt.get_cmap(colormap) + self.region_bg_color = np.array(region_bg_color) + + def draw_image_with_kp(self, image, kp_array): + image = np.copy(image) + spatial_size = np.array(image.shape[:2][::-1])[np.newaxis] + kp_array = spatial_size * (kp_array + 1) / 2 + num_regions = kp_array.shape[0] + for kp_ind, kp in enumerate(kp_array): + rr, cc = circle(int(kp[1]), int(kp[0]), self.kp_size, shape=image.shape[:2]) + image[rr, cc] = np.array(self.colormap(kp_ind / num_regions))[:3] + return image + + def create_image_column_with_kp(self, images, kp): + image_array = np.array( + [self.draw_image_with_kp(v, k) for v, k in zip(images, kp)] + ) + return self.create_image_column(image_array) + + def create_image_column(self, images): + if self.draw_border: + images = np.copy(images) + images[:, :, [0, -1]] = (1, 1, 1) + images[:, :, [0, -1]] = (1, 1, 1) + return np.concatenate(list(images), axis=0) + + def create_image_grid(self, *args): + out = [] + for arg in args: + if type(arg) == tuple: + out.append(self.create_image_column_with_kp(arg[0], arg[1])) + else: + out.append(self.create_image_column(arg)) + return np.concatenate(out, axis=1) + + def visualize(self, driving, source, out): + images = [] + + # Source image with region centers + source = source.data.cpu() + source_region_params = out["source_region_params"]["shift"].data.cpu().numpy() + source = np.transpose(source, [0, 2, 3, 1]) + images.append((source, source_region_params)) + + # Equivariance visualization + if "transformed_frame" in out: + transformed = out["transformed_frame"].data.cpu().numpy() + transformed = np.transpose(transformed, [0, 2, 3, 1]) + transformed_kp = ( + out["transformed_region_params"]["shift"].data.cpu().numpy() + ) + images.append((transformed, transformed_kp)) + + # Driving image with region centers + driving_region_params = out["driving_region_params"]["shift"].data.cpu().numpy() + driving = driving.data.cpu().numpy() + driving = np.transpose(driving, [0, 2, 3, 1]) + images.append((driving, driving_region_params)) + + # Deformed image + if "deformed" in out: + deformed = out["deformed"].data.cpu().numpy() + deformed = np.transpose(deformed, [0, 2, 3, 1]) + images.append(deformed) + + # Result + prediction = out["prediction"].data.cpu().numpy() + prediction = np.transpose(prediction, [0, 2, 3, 1]) + images.append(prediction) + + # Occlusion map + if "occlusion_map" in out: + occlusion_map = out["occlusion_map"].data.cpu().repeat(1, 3, 1, 1) + occlusion_map = F.interpolate(occlusion_map, size=source.shape[1:3]).numpy() + occlusion_map = np.transpose(occlusion_map, [0, 2, 3, 1]) + images.append(occlusion_map) + + # Heatmaps visualizations + if "heatmap" in out["driving_region_params"]: + driving_heatmap = F.interpolate( + out["driving_region_params"]["heatmap"], size=source.shape[1:3] + ) + driving_heatmap = np.transpose( + driving_heatmap.data.cpu().numpy(), [0, 2, 3, 1] + ) + images.append( + draw_colored_heatmap( + driving_heatmap, self.colormap, self.region_bg_color + ) + ) + + if "heatmap" in out["source_region_params"]: + source_heatmap = F.interpolate( + out["source_region_params"]["heatmap"], size=source.shape[1:3] + ) + source_heatmap = np.transpose( + source_heatmap.data.cpu().numpy(), [0, 2, 3, 1] + ) + images.append( + draw_colored_heatmap( + source_heatmap, self.colormap, self.region_bg_color + ) + ) + + image = self.create_image_grid(*images) + image = (255 * image).astype(np.uint8) + return image + + +def draw_colored_heatmap(heatmap, colormap, bg_color): + parts = [] + weights = [] + bg_color = np.array(bg_color).reshape((1, 1, 1, 3)) + num_regions = heatmap.shape[-1] + for i in range(num_regions): + color = np.array(colormap(i / num_regions))[:3] + color = color.reshape((1, 1, 1, 3)) + part = heatmap[:, :, :, i : (i + 1)] + part = part / np.max(part, axis=(1, 2), keepdims=True) + weights.append(part) + + color_part = part * color + parts.append(color_part) + + weight = sum(weights) + bg_weight = 1 - np.minimum(1, weight) + weight = np.maximum(1, weight) + result = sum(parts) / weight + bg_weight * bg_color + return result diff --git a/constants.py b/constants.py index 2eade3a..2b0691a 100644 --- a/constants.py +++ b/constants.py @@ -98,3 +98,6 @@ seg_model_dict = { "vox": seg10_segment_names, "vox-adv": seg10_segment_names, } + +ARTICULATE_MODEL_PATH = "articulate_module/models/vox256.pth" +ARTICULATE_CFG_PATH = "articulate_module/config/vox256.yaml" \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 1a85784..ffe24f1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui-firstordermm" description = "A ComfyUI-native node to run First Order Motion Model for Image Animation. [a/https://github.com/AliaksandrSiarohin/first-order-model](https://github.com/AliaksandrSiarohin/first-order-model)" -version = "1.0.1" +version = "1.0.2" license = { file = "LICENSE" } dependencies = ["numpy", "torch", "scipy", "yaml", "matplotlib"] diff --git a/run.py b/run.py index 013f8f0..2564109 100644 --- a/run.py +++ b/run.py @@ -9,7 +9,10 @@ from pathlib import Path import torch +from .articulate_inference import articulate_inference from .constants import ( + ARTICULATE_CFG_PATH, + ARTICULATE_MODEL_PATH, config_folder, default_model_name, default_partswap_model_name, @@ -338,6 +341,72 @@ class FOMM_Partswap: ) +class Articulate_Runner: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "source_image": ("IMAGE",), + "driving_video_input": ("IMAGE",), + "frame_rate": ("FLOAT", {"default": 30.0}), + }, + "optional": {"audio": ("AUDIO",)}, + } + + RETURN_TYPES = ( + "IMAGE", + "AUDIO", + "FLOAT", + ) + RETURN_NAMES = ( + "images", + "audio", + "frame_rate", + ) + FUNCTION = "todo" + CATEGORY = "FirstOrderMM" + + def todo( + self, + source_image: torch.Tensor, + driving_video_input: torch.Tensor, + frame_rate: float, + audio=None, + ): + print(f"{source_image.shape=}") + print(f"{driving_video_input.shape=}") + print(f"{type(audio)=}") + print(base_dir) + + config_path = f"{base_dir}/{ARTICULATE_CFG_PATH}" + checkpoint_path = f"{base_dir}/{ARTICULATE_MODEL_PATH}" + + source_image = reshape_image(source_image, (256, 256)) + driving_video = reshape_image(driving_video_input, (256, 256)).unsqueeze(0) + driving_video = driving_video.permute(0, 2, 1, 3, 4) + + print("After reshaping") + print(f"{source_image.shape=}") + print(f"{driving_video.shape=}") + + params = { + "source_image": source_image, + "driving_video": driving_video, + "config_path": config_path, + "checkpoint_path": checkpoint_path, + } + + predictions = articulate_inference(**params) + + output_images = out_video(predictions) + + return ( + output_images, + audio, + frame_rate, + ) + + def serialize_integers(int_list): return "_".join(map(str, int_list)) diff --git a/workflows/ARTICULATE.json b/workflows/ARTICULATE.json new file mode 100644 index 0000000..a5f1e76 --- /dev/null +++ b/workflows/ARTICULATE.json @@ -0,0 +1,423 @@ +{ + "last_node_id": 17, + "last_link_id": 58, + "nodes": [ + { + "id": 2, + "type": "LoadImage", + "pos": [ + 242, + 94 + ], + "size": { + "0": 315, + "1": 314 + }, + "flags": {}, + "order": 0, + "mode": 0, + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 52 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "MASK", + "type": "MASK", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "Rick1.png", + "image" + ] + }, + { + "id": 3, + "type": "VHS_LoadVideo", + "pos": [ + 323, + 465 + ], + "size": [ + 235.1999969482422, + 491.1999969482422 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [ + { + "name": "meta_batch", + "type": "VHS_BatchManager", + "link": null + }, + { + "name": "vae", + "type": "VAE", + "link": null + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 53 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "frame_count", + "type": "INT", + "links": null, + "shape": 3 + }, + { + "name": "audio", + "type": "AUDIO", + "links": [ + 54 + ], + "shape": 3, + "slot_index": 2 + }, + { + "name": "video_info", + "type": "VHS_VIDEOINFO", + "links": [ + 3 + ], + "shape": 3, + "slot_index": 3 + } + ], + "properties": { + "Node name for S&R": "VHS_LoadVideo" + }, + "widgets_values": { + "video": "damedane.mp4", + "force_rate": 0, + "force_size": "Disabled", + "custom_width": 512, + "custom_height": 512, + "frame_load_cap": 0, + "skip_first_frames": 0, + "select_every_nth": 1, + "choose video to upload": "image", + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "frame_load_cap": 0, + "skip_first_frames": 0, + "force_rate": 0, + "filename": "damedane.mp4", + "type": "input", + "format": "video/mp4", + "select_every_nth": 1 + } + } + } + }, + { + "id": 4, + "type": "VHS_VideoInfoLoaded", + "pos": [ + 647, + 346 + ], + "size": { + "0": 304.79998779296875, + "1": 106 + }, + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + { + "name": "video_info", + "type": "VHS_VIDEOINFO", + "link": 3 + } + ], + "outputs": [ + { + "name": "fps🟦", + "type": "FLOAT", + "links": [ + 55 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "frame_count🟦", + "type": "INT", + "links": null, + "shape": 3 + }, + { + "name": "duration🟦", + "type": "FLOAT", + "links": null, + "shape": 3 + }, + { + "name": "width🟦", + "type": "INT", + "links": null, + "shape": 3 + }, + { + "name": "height🟦", + "type": "INT", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "VHS_VideoInfoLoaded" + }, + "widgets_values": {} + }, + { + "id": 15, + "type": "VHS_VideoCombine", + "pos": [ + 1065, + 135 + ], + "size": [ + 315, + 615 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 56 + }, + { + "name": "audio", + "type": "AUDIO", + "link": 57 + }, + { + "name": "meta_batch", + "type": "VHS_BatchManager", + "link": null + }, + { + "name": "vae", + "type": "VAE", + "link": null + }, + { + "name": "frame_rate", + "type": "FLOAT", + "link": 58, + "widget": { + "name": "frame_rate" + } + } + ], + "outputs": [ + { + "name": "Filenames", + "type": "VHS_FILENAMES", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "VHS_VideoCombine" + }, + "widgets_values": { + "frame_rate": 8, + "loop_count": 0, + "filename_prefix": "Articulate", + "format": "video/h264-mp4", + "pix_fmt": "yuv420p", + "crf": 0, + "save_metadata": true, + "pingpong": false, + "save_output": false, + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "filename": "Articulate_00001-audio.mp4", + "subfolder": "", + "type": "temp", + "format": "video/h264-mp4", + "frame_rate": 29.970029970029966 + } + } + } + }, + { + "id": 17, + "type": "Articulate_Runner", + "pos": [ + 646, + 145 + ], + "size": { + "0": 380.4000244140625, + "1": 146 + }, + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "source_image", + "type": "IMAGE", + "link": 52 + }, + { + "name": "driving_video_input", + "type": "IMAGE", + "link": 53 + }, + { + "name": "audio", + "type": "AUDIO", + "link": 54 + }, + { + "name": "frame_rate", + "type": "FLOAT", + "link": 55, + "widget": { + "name": "frame_rate" + } + } + ], + "outputs": [ + { + "name": "images", + "type": "IMAGE", + "links": [ + 56 + ], + "shape": 3 + }, + { + "name": "audio", + "type": "AUDIO", + "links": [ + 57 + ], + "shape": 3 + }, + { + "name": "frame_rate", + "type": "FLOAT", + "links": [ + 58 + ], + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "Articulate_Runner" + }, + "widgets_values": [ + 30 + ] + } + ], + "links": [ + [ + 3, + 3, + 3, + 4, + 0, + "VHS_VIDEOINFO" + ], + [ + 52, + 2, + 0, + 17, + 0, + "IMAGE" + ], + [ + 53, + 3, + 0, + 17, + 1, + "IMAGE" + ], + [ + 54, + 3, + 2, + 17, + 2, + "AUDIO" + ], + [ + 55, + 4, + 0, + 17, + 3, + "FLOAT" + ], + [ + 56, + 17, + 0, + 15, + 0, + "IMAGE" + ], + [ + 57, + 17, + 1, + 15, + 1, + "AUDIO" + ], + [ + 58, + 17, + 2, + 15, + 4, + "FLOAT" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 0.9090909090909091, + "offset": [ + 1.9174305033795593, + 14.883996994740842 + ] + } + }, + "version": 0.4 +} \ No newline at end of file diff --git a/FOMM.json b/workflows/FOMM.json similarity index 100% rename from FOMM.json rename to workflows/FOMM.json diff --git a/FOMM_PARTSWAP.json b/workflows/FOMM_PARTSWAP.json similarity index 100% rename from FOMM_PARTSWAP.json rename to workflows/FOMM_PARTSWAP.json diff --git a/workflow.png b/workflows/workflow.png similarity index 100% rename from workflow.png rename to workflows/workflow.png diff --git a/workflows/workflow_articulate.png b/workflows/workflow_articulate.png new file mode 100644 index 0000000..da37ce7 Binary files /dev/null and b/workflows/workflow_articulate.png differ diff --git a/workflow_fomm_partswap.png b/workflows/workflow_fomm_partswap.png similarity index 100% rename from workflow_fomm_partswap.png rename to workflows/workflow_fomm_partswap.png