Add articulated animation (#5)

.gitignore

@@ -1,5 +1,7 @@
 # Stuff
 
+articulate_module/models/*.pth
+
 checkpoints/*
 !checkpoints/place_checkpoints_here.txt
 

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`

__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",
 }
 
 

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

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

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)

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

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)

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]

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]

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]

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]

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

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

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

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]

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

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"

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"]
 

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))
 

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
+}