add mrfa node

.gitignore

@@ -2,6 +2,7 @@
 
 module_articulate/models/*.pth
 module_spline/models/*.tar
+module_mrfa/models/*.pth
 checkpoints/*
 !checkpoints/place_checkpoints_here.txt
 

README.md

@@ -40,6 +40,11 @@ https://github.com/user-attachments/assets/b2948efb-3b44-440b-bff2-dde7b95a9946
 
 ![Workflow Spline](workflows/workflow_spline.png)
 
+### MRFA
+
+[MRFA.json](workflows/MRFA.json)
+
+![Workflow MRFA](workflows/workflow_mrfa.png)
 
 ## Arguments
 
@@ -73,6 +78,10 @@ Doesn't need any
 
 ### MRFA
 
+* `model_name`: `celebvhq` or `vox`, which is trained on different datasets
+* `use_relative`: Whether to use relative mode or not (absolute mode). Absolute mode is similar to FOMM's `adapt_movement_scale` set to False
+* `relative_movement`, `relative_jacobian`, `adapt_movement_scale`: Same as FOMM
+
 ## Installation
 
 1. Clone the repo to `ComfyUI/custom_nodes/`
@@ -134,4 +143,9 @@ For **Articulate**, download the model from [Pre-trained checkpoints](https://gi
 
 For **Spline**, download the model from [Pre-trained models](https://github.com/yoyo-nb/thin-plate-spline-motion-model?tab=readme-ov-file#pre-trained-models) section and place it here: `spline_module/models/vox.pth.tar`. To use `find_best_frame`, install `face-alignment`.
 
-For **MRFA**, download the model from [Pretrained models](https://github.com/JialeTao/MRFA/?tab=readme-ov-file#pretrained-models) section and place it here: `mrfa_module/models/`
+For **MRFA**, download `MRFACheckpoints.tar.gz` from [Pretrained models](https://github.com/JialeTao/MRFA/?tab=readme-ov-file#pretrained-models) section and decompress it. Inside there should be two files. Place them here: 
+
+```
+mrfa_module/models/celebvhq.pth
+mrfa_module/models/vox.pth
+```

__init__.py

@@ -5,7 +5,8 @@ from .run import (
     FOMM_Seg10Chooser,
     FOMM_Seg15Chooser,
     Articulate_Runner,
-    Spline_Runner
+    Spline_Runner,
+    MRFA_Runner
 )
 
 NODE_CLASS_MAPPINGS = {
@@ -16,6 +17,7 @@ NODE_CLASS_MAPPINGS = {
     "FOMM_Seg15Chooser": FOMM_Seg15Chooser,
     "Articulate_Runner": Articulate_Runner,
     "Spline_Runner": Spline_Runner,
+    "MRFA_Runner": MRFA_Runner,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -26,6 +28,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "FOMM_Seg15Chooser": "FOMM Seg15 Chooser",
     "Articulate_Runner": "Articulate Runner",
     "Spline_Runner": "Spline Runner",
+    "MRFA_Runner": "MRFA Runner",
 }
 
 

constants.py

@@ -106,4 +106,16 @@ SPLINE_MODES = ['relative', 'standard', 'avd']
 SPLINE_DEFAULT = 'relative'
 
 SPLINE_MODEL_PATH = "module_spline/models/vox.pth.tar"
-SPLINE_CFG_PATH = "module_spline/config/vox-256.yaml"
+SPLINE_CFG_PATH = "module_spline/config/vox-256.yaml"
+
+MRFA_MODEL_NAMES = ["celebvhq", "vox"]
+MRFA_MODEL_PATHS = {
+    "celebvhq": "module_mrfa/models/celebvhq.pth",
+    "vox": "module_mrfa/models/vox.pth",
+}
+MRFA_DEFAULT_MODEL = "celebvhq"
+
+MRFA_CFG_PATHS = {
+    "celebvhq": "module_mrfa/configs/celebvhq.yaml",
+    "vox": "module_mrfa/configs/vox1.yaml",
+}

inference_mrfa.py

@@ -0,0 +1,131 @@
+import numpy as np
+import torch
+import tqdm
+import yaml
+from comfy.utils import ProgressBar
+from scipy.spatial import ConvexHull
+
+from .module_mrfa.model import MRFA
+from .module_mrfa.util import AntiAliasInterpolation2d
+
+
+def mrfa_inference(
+    source_image: torch.Tensor,
+    driving_video: torch.Tensor,
+    config_path: str,
+    checkpoint_path: str,
+    use_relative=True,
+    relative_movement=True,
+    relative_jacobian=True,
+    adapt_movement_scale=False,
+):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    with open(config_path) as f:
+        config = yaml.full_load(f)
+
+    model = MRFA(config)
+    model = torch.nn.DataParallel(model)
+
+    checkpoint = torch.load(checkpoint_path, map_location=device)
+    model.load_state_dict(checkpoint["model"], strict=False)
+    model = model.module
+    model.eval().to(device)
+
+    kp_detector = model.encoder
+    dense_motion_network = model.dense_motion
+    decoder = model.decoder
+    down = AntiAliasInterpolation2d(3, 0.25).to(device)
+
+    source = source_image.to(device)  # 4D
+    driving = driving_video.to(device)  #
+
+    num_frames = driving.shape[2]
+    pbar = ProgressBar(num_frames)
+
+    predictions = []
+    with torch.no_grad():
+        # visualizations = []
+
+        kp_source = kp_detector.forward(source)
+        kp_driving_initial = kp_detector.forward(driving[:, :, 0])
+
+        for frame_idx in tqdm.tqdm(
+            range(num_frames),
+            desc=f"Generating {'Relative' if use_relative else 'Absolute'} mode",
+        ):
+            driving_frame = driving[:, :, frame_idx]
+            if use_relative:
+                kp_driving = kp_detector.forward(driving_frame)
+                kp_norm = normalize_kp(
+                    kp_source=kp_source,
+                    kp_driving=kp_driving,
+                    kp_driving_initial=kp_driving_initial,
+                    use_relative_movement=relative_movement,
+                    use_relative_jacobian=relative_jacobian,
+                    adapt_movement_scale=adapt_movement_scale,
+                )
+                dense_motion = dense_motion_network.forward(
+                    source, kp_norm, kp_source, bg_param=None
+                )
+
+                if config["train_params"]["prior_model"] == "tpsm":
+                    kp_s_value = kp_source["kp"].view(source.shape[0], -1, 5, 2).mean(2)
+                    kp_d_value = kp_norm["kp"].view(driving.shape[0], -1, 5, 2).mean(2)
+                else:
+                    kp_s_value = kp_source["kp"]
+                    kp_d_value = kp_norm["kp"]
+
+                out, _, _ = decoder.forward(
+                    kp_s_value,
+                    kp_d_value,
+                    dense_motion,
+                    img=down(source),
+                    img_full=source,
+                )
+
+                # visualization = Visualizer(**config['visualizer_params']).visualize(source=source, driving=driving_frame, out=out)
+                # visualizations.append(visualization)
+
+            else:
+                abs_inp = {"source": source, "driving": driving_frame}
+                out, _, _, _, _ = model.forward(abs_inp, is_train=False)
+
+            prediction = out.data.cpu().numpy()
+            prediction = np.transpose(prediction, [0, 2, 3, 1]).squeeze(0)
+            predictions.append(prediction)
+
+            pbar.update_absolute(frame_idx, num_frames)
+
+    return predictions
+
+
+def normalize_kp(
+    kp_source,
+    kp_driving,
+    kp_driving_initial,
+    adapt_movement_scale=False,
+    use_relative_movement=False,
+    use_relative_jacobian=False,
+):
+    if adapt_movement_scale:
+        source_area = ConvexHull(kp_source["kp"][0].data.cpu().numpy()).volume
+        driving_area = ConvexHull(kp_driving_initial["kp"][0].data.cpu().numpy()).volume
+        adapt_movement_scale = np.sqrt(source_area) / np.sqrt(driving_area)
+    else:
+        adapt_movement_scale = 1
+
+    kp_new = {k: v for k, v in kp_driving.items()}
+
+    if use_relative_movement:
+        kp_value_diff = kp_driving["kp"] - kp_driving_initial["kp"]
+        kp_value_diff *= adapt_movement_scale
+        kp_new["kp"] = kp_value_diff + kp_source["kp"]
+
+        if use_relative_jacobian:
+            jacobian_diff = torch.matmul(
+                kp_driving["jacobian"], torch.inverse(kp_driving_initial["jacobian"])
+            )
+            kp_new["jacobian"] = torch.matmul(jacobian_diff, kp_source["jacobian"])
+
+    return kp_new

module_mrfa/bg_motion_predictor.py

@@ -0,0 +1,29 @@
+import torch
+from torch import nn
+from torchvision import models
+
+
+class BGMotionPredictor(nn.Module):
+    """
+    Module for background estimation, return single transformation, parametrized as 3x3 matrix. The third row is [0 0 1]
+    """
+
+    def __init__(self):
+        super(BGMotionPredictor, self).__init__()
+        self.bg_encoder = models.resnet18(pretrained=False)
+        self.bg_encoder.conv1 = nn.Conv2d(
+            6, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False
+        )
+        num_features = self.bg_encoder.fc.in_features
+        self.bg_encoder.fc = nn.Linear(num_features, 6)
+        self.bg_encoder.fc.weight.data.zero_()
+        self.bg_encoder.fc.bias.data.copy_(
+            torch.tensor([1, 0, 0, 0, 1, 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())
+        prediction = self.bg_encoder(torch.cat([source_image, driving_image], dim=1))
+        out[:, :2, :] = prediction.view(bs, 2, 3)
+        return out

module_mrfa/configs/celebvhq.yaml

@@ -0,0 +1,186 @@
+dataset_params:
+  root_dir: ./data/celebvhq256-png/
+  split_len: 11
+  frame_shape: [256, 256, 3]
+  id_sampling: True
+  pairs_list: 
+  augmentation_params:
+    flip_param:
+      horizontal_flip: True 
+      time_flip: True
+    jitter_param:
+      brightness: 0.1
+      contrast: 0.1
+      saturation: 0.1
+      hue: 0.1
+
+dense_motion:
+  block_expansion: 64
+  max_features: 1024
+  num_blocks: 5
+  scale_factor: 0.25
+  num_kp: 10
+  num_channels: 3
+  estimate_occlusion_map: True
+fomm_kp_detector:
+  block_expansion: 32
+  num_kp: 10
+  num_channels: 3
+  max_features: 1024
+  num_blocks: 5
+  temperature: 0.1
+  scale_factor: 0.25
+  estimate_jacobian: True
+  estimate_occlusion: False
+tpsm_kp_detector:
+  num_tps: 10
+tpsm_dense_motion:
+  block_expansion: 64
+  max_features: 1024
+  num_blocks: 5
+  scale_factor: 0.25
+  num_tps: 10
+  num_channels: 3
+
+raft_flow:
+  prior_only: False
+  num_kp: 10
+  dim: 256
+  size: 256
+  generator:
+    num_channels: 3
+    block_expansion: 64 
+    max_features: 512
+    num_up_blocks: 5
+  driving_encoder:
+    in_features: 10
+    block_expansion: 32 
+    max_features: 512
+    num_blocks: 5
+  source_encoder:
+    in_features: 13
+    block_expansion: 32 
+    max_features: 512
+    num_blocks: 5
+
+train_params:
+  ## [mtia, tpsm, fomm]
+  prior_model: mtia
+  num_epochs: 100
+  num_repeats: 5
+  epoch_milestones: [60, 90]
+  # ### enable this for tpsm
+  # epoch_milestones: [70, 90]
+  # ### enable this for tpsm
+
+  lr: 2.0e-4
+  batch_size: 80 
+  scales: [1, 0.5, 0.25, 0.125]
+  clip_grad: True
+  clip: 10
+  bg_start: 0
+  checkpoint_freq: 100
+  ### for tpsm
+  dropout_epoch: 35 
+  dropout_maxp: 0.3
+  dropout_startp: 0.1
+  dropout_inc_epoch: 10
+  ### for tpsm
+  
+  transform_params:
+    sigma_affine: 0.05
+    sigma_tps: 0.005
+    points_tps: 5
+  loss_weights:
+    perceptual: [10, 10, 10, 10, 10]
+    equivariance: 10
+    equivariance_jacobian: 10
+    # ### enable this for tpsm
+    # equivariance_jacobian: 0
+    # ### enable this for tpsm
+
+reconstruction_params:
+  num_videos: 1000
+  format: '.mp4'
+
+animate_params:
+  num_pairs: 50
+  format: '.mp4'
+  normalization_params:
+    adapt_movement_scale: False
+    use_relative_movement: True
+    use_relative_jacobian: True
+
+visualizer_params:
+  draw_border: True
+  colormap: 'gist_rainbow'
+
+mtia_kp_detector:
+  MODEL:
+    ESTIMATE_JACOBIAN: True
+    DATA_PREPROCESS: False
+    FIX_IMG2MOTION_ATTENTION: False
+    TRANSFORMER_DEPTH: 12
+    TRANSFORMER_HEADS: 8
+    DIM: 192 # 4*4*3
+    INIT_WEIGHTS: False
+    NAME: pose_tokenpose_b
+    NUM_JOINTS: 10
+    PRETRAINED: ''
+    PATCH_SIZE: [4, 4]
+    IMAGE_SIZE: [256, 256]
+    HEATMAP_SIZE: [64, 64]
+
+    ## default
+    TAG_PER_JOINT: True
+    HIDDEN_HEATMAP_DIM: -1
+    MULTI_TRANSFORMER_DEPTH: [12, 12]
+    MULTI_TRANSFORMER_HEADS: [16, 16]
+    MULTI_DIM: [48, 48]
+    NUM_BRANCHES: 1
+    BASE_CHANNEL: 32
+    TRANSFORMER_MLP_RATIO: 3
+    POS_EMBEDDING_TYPE: 'sine-full'
+    TEMPERATURE: 0.1
+    TARGET_TYPE: gaussian
+    INIT: True
+
+    SIGMA: 2
+    EXTRA:
+      PRETRAINED_LAYERS:
+      - 'conv1'
+      - 'bn1'
+      - 'conv2'
+      - 'bn2'
+      - 'layer1'
+      - 'transition1'
+      - 'stage2'
+      - 'transition2'
+      - 'stage3'
+      FINAL_CONV_KERNEL: 1
+      STAGE2:
+        NUM_MODULES: 1
+        NUM_BRANCHES: 2
+        BLOCK: BASIC
+        NUM_BLOCKS:
+        - 4
+        - 4
+        NUM_CHANNELS:
+        - 32
+        - 64
+        FUSE_METHOD: SUM
+      STAGE3:
+        NUM_MODULES: 4
+        NUM_BRANCHES: 3
+        BLOCK: BASIC
+        NUM_BLOCKS:
+        - 4
+        - 4
+        - 4
+        NUM_CHANNELS:
+        - 32
+        - 64
+        - 128
+        FUSE_METHOD: SUM
+      ## default
+  

module_mrfa/configs/vox1.yaml

@@ -0,0 +1,186 @@
+dataset_params:
+  root_dir: ./data/voxceleb1-png/
+  split_len: 
+  frame_shape: [256, 256, 3]
+  id_sampling: True
+  pairs_list: 
+  augmentation_params:
+    flip_param:
+      horizontal_flip: True 
+      time_flip: True
+    jitter_param:
+      brightness: 0.1
+      contrast: 0.1
+      saturation: 0.1
+      hue: 0.1
+
+dense_motion:
+  block_expansion: 64
+  max_features: 1024
+  num_blocks: 5
+  scale_factor: 0.25
+  num_kp: 10
+  num_channels: 3
+  estimate_occlusion_map: True
+fomm_kp_detector:
+  block_expansion: 32
+  num_kp: 10
+  num_channels: 3
+  max_features: 1024
+  num_blocks: 5
+  temperature: 0.1
+  scale_factor: 0.25
+  estimate_jacobian: True
+  estimate_occlusion: False
+tpsm_kp_detector:
+  num_tps: 10
+tpsm_dense_motion:
+  block_expansion: 64
+  max_features: 1024
+  num_blocks: 5
+  scale_factor: 0.25
+  num_tps: 10
+  num_channels: 3
+
+raft_flow:
+  prior_only: False
+  num_kp: 10
+  dim: 256
+  size: 256
+  generator:
+    num_channels: 3
+    block_expansion: 64 
+    max_features: 512
+    num_up_blocks: 5
+  driving_encoder:
+    in_features: 10
+    block_expansion: 32 
+    max_features: 512
+    num_blocks: 5
+  source_encoder:
+    in_features: 13
+    block_expansion: 32 
+    max_features: 512
+    num_blocks: 5
+
+train_params:
+  ## [mtia, tpsm, fomm]
+  prior_model: mtia
+  num_epochs: 100
+  num_repeats: 150
+  epoch_milestones: [60, 90]
+  # ### enable this for tpsm
+  # epoch_milestones: [70, 90]
+  # ### enable this for tpsm
+
+  lr: 2.0e-4
+  batch_size: 80 
+  scales: [1, 0.5, 0.25, 0.125]
+  clip_grad: True
+  clip: 10
+  bg_start: 1000
+  checkpoint_freq: 100
+  ### for tpsm
+  dropout_epoch: 35 
+  dropout_maxp: 0.3
+  dropout_startp: 0.1
+  dropout_inc_epoch: 10
+  ### for tpsm
+  
+  transform_params:
+    sigma_affine: 0.05
+    sigma_tps: 0.005
+    points_tps: 5
+  loss_weights:
+    perceptual: [10, 10, 10, 10, 10]
+    equivariance: 10
+    equivariance_jacobian: 10
+    # ### enable this for tpsm
+    # equivariance_jacobian: 0
+    # ### enable this for tpsm
+
+reconstruction_params:
+  num_videos: 1000
+  format: '.mp4'
+
+animate_params:
+  num_pairs: 50
+  format: '.mp4'
+  normalization_params:
+    adapt_movement_scale: False
+    use_relative_movement: True
+    use_relative_jacobian: True
+
+visualizer_params:
+  draw_border: True
+  colormap: 'gist_rainbow'
+
+mtia_kp_detector:
+  MODEL:
+    ESTIMATE_JACOBIAN: True
+    DATA_PREPROCESS: False
+    FIX_IMG2MOTION_ATTENTION: False
+    TRANSFORMER_DEPTH: 12
+    TRANSFORMER_HEADS: 8
+    DIM: 192 # 4*4*3
+    INIT_WEIGHTS: False
+    NAME: pose_tokenpose_b
+    NUM_JOINTS: 10
+    PRETRAINED: ''
+    PATCH_SIZE: [4, 4]
+    IMAGE_SIZE: [256, 256]
+    HEATMAP_SIZE: [64, 64]
+
+    ## default
+    TAG_PER_JOINT: True
+    HIDDEN_HEATMAP_DIM: -1
+    MULTI_TRANSFORMER_DEPTH: [12, 12]
+    MULTI_TRANSFORMER_HEADS: [16, 16]
+    MULTI_DIM: [48, 48]
+    NUM_BRANCHES: 1
+    BASE_CHANNEL: 32
+    TRANSFORMER_MLP_RATIO: 3
+    POS_EMBEDDING_TYPE: 'sine-full'
+    TEMPERATURE: 0.1
+    TARGET_TYPE: gaussian
+    INIT: True
+
+    SIGMA: 2
+    EXTRA:
+      PRETRAINED_LAYERS:
+      - 'conv1'
+      - 'bn1'
+      - 'conv2'
+      - 'bn2'
+      - 'layer1'
+      - 'transition1'
+      - 'stage2'
+      - 'transition2'
+      - 'stage3'
+      FINAL_CONV_KERNEL: 1
+      STAGE2:
+        NUM_MODULES: 1
+        NUM_BRANCHES: 2
+        BLOCK: BASIC
+        NUM_BLOCKS:
+        - 4
+        - 4
+        NUM_CHANNELS:
+        - 32
+        - 64
+        FUSE_METHOD: SUM
+      STAGE3:
+        NUM_MODULES: 4
+        NUM_BRANCHES: 3
+        BLOCK: BASIC
+        NUM_BLOCKS:
+        - 4
+        - 4
+        - 4
+        NUM_CHANNELS:
+        - 32
+        - 64
+        - 128
+        FUSE_METHOD: SUM
+      ## default
+  

module_mrfa/dense_motion.py

@@ -0,0 +1,435 @@
+import math
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from .util import (
+    TPS,
+    AntiAliasInterpolation2d,
+    Hourglass,
+    UpBlock2d,
+    from_homogeneous,
+    kp2gaussian,
+    make_coordinate_grid,
+    to_homogeneous,
+)
+
+
+class DenseMotionNetwork(nn.Module):
+    """
+    Module that predicting a dense motion from sparse motion representation given by kp_source and kp_driving
+    """
+
+    def __init__(
+        self,
+        block_expansion,
+        num_blocks,
+        max_features,
+        num_kp,
+        num_channels,
+        estimate_occlusion_map=True,
+        scale_factor=1,
+        kp_variance=0.01,
+    ):
+        super(DenseMotionNetwork, self).__init__()
+
+        infeatures = num_kp + 1
+        self.infeatures = infeatures
+
+        self.hourglass = Hourglass(
+            block_expansion=block_expansion,
+            in_features=infeatures * (num_channels + 1),
+            max_features=max_features,
+            num_blocks=num_blocks,
+        )
+        self.mask = nn.Conv2d(
+            self.hourglass.out_filters, infeatures, 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_kp = num_kp
+        self.scale_factor = scale_factor
+        self.kp_variance = kp_variance
+
+        if self.scale_factor != 1:
+            self.down = AntiAliasInterpolation2d(num_channels, self.scale_factor)
+
+    def create_heatmap_representations(self, source_image, kp_driving, kp_source):
+        spatial_size = source_image.shape[2:]
+        gaussian_driving = kp2gaussian(
+            kp_driving["kp"], spatial_size=spatial_size, kp_variance=self.kp_variance
+        )
+        gaussian_source = kp2gaussian(
+            kp_source["kp"], spatial_size=spatial_size, kp_variance=self.kp_variance
+        )
+        heatmap = gaussian_driving - gaussian_source
+
+        # adding background feature
+        zeros = torch.zeros(heatmap.shape[0], 1, spatial_size[0], spatial_size[1]).type(
+            heatmap.type()
+        )
+        heatmap = torch.cat([zeros, heatmap], dim=1)
+        heatmap = heatmap.unsqueeze(2)
+        return heatmap
+
+    def create_sparse_motions(self, source_image, kp_driving, kp_source, bg_param=None):
+        bs, _, h, w = source_image.shape
+        identity_grid = make_coordinate_grid((h, w), type=kp_source["kp"].type())
+        identity_grid = identity_grid.view(1, 1, h, w, 2)
+        coordinate_grid = identity_grid - kp_driving["kp"].view(
+            bs, self.num_kp, 1, 1, 2
+        )
+        if "jacobian" in kp_driving:
+            jacobian = torch.matmul(
+                kp_source["jacobian"], torch.inverse(kp_driving["jacobian"])
+            )
+            jacobian = jacobian.unsqueeze(-3).unsqueeze(-3)
+            jacobian = jacobian.repeat(1, 1, h, w, 1, 1)
+            coordinate_grid = torch.matmul(jacobian, coordinate_grid.unsqueeze(-1))
+            coordinate_grid = coordinate_grid.squeeze(-1)
+            # n * 10 * h * w * 2
+
+        driving_to_source = coordinate_grid + kp_source["kp"].view(
+            bs, self.num_kp, 1, 1, 2
+        )
+
+        # #adding background feature
+        # identity_grid = identity_grid.repeat(bs, 1, 1, 1, 1)
+
+        # adding background feature
+        if bg_param 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_param.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_kp + 1, 1, 1, 1, 1)
+        )
+        source_repeat = source_repeat.view(bs * (self.num_kp + 1), -1, h, w)
+        sparse_motions = sparse_motions.view((bs * (self.num_kp + 1), h, w, -1))
+        sparse_deformed = F.grid_sample(source_repeat, sparse_motions)
+        sparse_deformed = sparse_deformed.view((bs, self.num_kp + 1, -1, h, w))
+        return sparse_deformed
+
+    def dropout_softmax(self, X, P):
+        """
+        Dropout for TPS transformations. Eq(7) and Eq(8) in the paper.
+        """
+        drop = (
+            (torch.rand(X.shape[0], X.shape[1]) < (1 - P)).type(X.type()).to(X.device)
+        )
+        drop[..., 0] = 1
+        drop = drop.repeat(X.shape[2], X.shape[3], 1, 1).permute(2, 3, 0, 1)
+
+        maxx = X.max(1).values.unsqueeze_(1)
+        X = X - maxx
+        X_exp = X.exp()
+        X[:, 1:, ...] /= 1 - P
+        mask_bool = drop == 0
+        X_exp = X_exp.masked_fill(mask_bool, 0)
+        partition = X_exp.sum(dim=1, keepdim=True) + 1e-6
+        return X_exp / partition
+
+    def forward(
+        self,
+        source_image,
+        kp_driving,
+        kp_source,
+        bg_param=None,
+        dropout_flag=False,
+        dropout_p=0,
+    ):
+        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, kp_driving, kp_source
+        )
+        sparse_motion = self.create_sparse_motions(
+            source_image, kp_driving, kp_source, bg_param=bg_param
+        )
+        # n * 11 * h * w * 2
+        deformed_source = self.create_deformed_source_image(source_image, sparse_motion)
+        # n * 11 * 3 * h * w
+        out_dict["sparse_deformed"] = deformed_source
+
+        input = torch.cat([heatmap_representation, deformed_source], dim=2)
+        # n * 11(101,22) * 4 * h * w
+        input = input.view(bs, -1, h, w)
+
+        prediction = self.hourglass(input)
+
+        mask = self.mask(prediction)
+        out_dict["logit_mask"] = mask
+        if dropout_flag:
+            mask = self.dropout_softmax(mask, dropout_p)
+        else:
+            mask = F.softmax(mask, dim=1)
+        # mask = F.softmax(mask, dim=1)
+        out_dict["mask"] = mask
+        mask = mask.unsqueeze(2)
+        sparse_motion = sparse_motion.permute(0, 1, 4, 2, 3)
+        # n * 11 * 2 * h * w
+        deformation = (sparse_motion * mask).sum(dim=1)
+        deformation = deformation.permute(0, 2, 3, 1)
+        # n * h * w * 2
+
+        out_dict["deformation"] = deformation
+
+        if self.occlusion is not None:
+            # occlusion_map = torch.sigmoid(self.occlusion(prediction))
+            occlusion_map = self.occlusion(prediction)
+            out_dict["occlusion"] = occlusion_map
+
+        return out_dict
+
+
+class TPSDenseMotionNetwork(nn.Module):
+    """
+    Module that estimating an optical flow and multi-resolution occlusion masks
+                        from K TPS transformations and an affine transformation.
+    """
+
+    def __init__(
+        self,
+        block_expansion,
+        num_blocks,
+        max_features,
+        num_tps,
+        num_channels,
+        scale_factor=0.25,
+        bg=False,
+        multi_mask=False,
+        kp_variance=0.01,
+    ):
+        super(TPSDenseMotionNetwork, self).__init__()
+
+        if scale_factor != 1:
+            self.down = AntiAliasInterpolation2d(num_channels, scale_factor)
+        self.scale_factor = scale_factor
+        self.multi_mask = multi_mask
+
+        self.hourglass = Hourglass(
+            block_expansion=block_expansion,
+            in_features=(num_channels * (num_tps + 1) + num_tps * 5 + 1),
+            max_features=max_features,
+            num_blocks=num_blocks,
+        )
+
+        hourglass_output_size = self.hourglass.out_filters
+        # self.maps = nn.Conv2d(hourglass_output_size[-1], num_tps + 1, kernel_size=(7, 7), padding=(3, 3))
+        self.maps = nn.Conv2d(
+            hourglass_output_size, num_tps + 1, kernel_size=(7, 7), padding=(3, 3)
+        )
+
+        if multi_mask:
+            up = []
+            self.up_nums = int(math.log(1 / scale_factor, 2))
+            self.occlusion_num = 4
+
+            channel = [hourglass_output_size[-1] // (2**i) for i in range(self.up_nums)]
+            for i in range(self.up_nums):
+                up.append(
+                    UpBlock2d(channel[i], channel[i] // 2, kernel_size=3, padding=1)
+                )
+            self.up = nn.ModuleList(up)
+
+            channel = [
+                hourglass_output_size[-i - 1]
+                for i in range(self.occlusion_num - self.up_nums)[::-1]
+            ]
+            for i in range(self.up_nums):
+                channel.append(hourglass_output_size[-1] // (2 ** (i + 1)))
+            occlusion = []
+
+            for i in range(self.occlusion_num):
+                occlusion.append(
+                    nn.Conv2d(channel[i], 1, kernel_size=(7, 7), padding=(3, 3))
+                )
+            self.occlusion = nn.ModuleList(occlusion)
+        else:
+            # occlusion = [nn.Conv2d(hourglass_output_size[-1], 1, kernel_size=(7, 7), padding=(3, 3))]
+            occlusion = [
+                nn.Conv2d(hourglass_output_size, 1, kernel_size=(7, 7), padding=(3, 3))
+            ]
+            self.occlusion = nn.ModuleList(occlusion)
+
+        self.num_tps = num_tps
+        self.bg = bg
+        self.kp_variance = kp_variance
+
+    def create_heatmap_representations(self, source_image, kp_driving, kp_source):
+        spatial_size = source_image.shape[2:]
+        gaussian_driving = kp2gaussian(
+            kp_driving["kp"], spatial_size=spatial_size, kp_variance=self.kp_variance
+        )
+        gaussian_source = kp2gaussian(
+            kp_source["kp"], spatial_size=spatial_size, kp_variance=self.kp_variance
+        )
+        heatmap = gaussian_driving - gaussian_source
+
+        zeros = (
+            torch.zeros(heatmap.shape[0], 1, spatial_size[0], spatial_size[1])
+            .type(heatmap.type())
+            .to(heatmap.device)
+        )
+        heatmap = torch.cat([zeros, heatmap], dim=1)
+
+        return heatmap
+
+    def create_transformations(self, source_image, kp_driving, kp_source, bg_param):
+        # K TPS transformaions
+        bs, _, h, w = source_image.shape
+        kp_1 = kp_driving["kp"]
+        kp_2 = kp_source["kp"]
+        kp_1 = kp_1.view(bs, -1, 5, 2)
+        kp_2 = kp_2.view(bs, -1, 5, 2)
+        trans = TPS(mode="kp", bs=bs, kp_1=kp_1, kp_2=kp_2)
+        driving_to_source = trans.transform_frame(source_image)
+
+        identity_grid = make_coordinate_grid((h, w), type=kp_1.type()).to(kp_1.device)
+        identity_grid = identity_grid.view(1, 1, h, w, 2)
+        identity_grid = identity_grid.repeat(bs, 1, 1, 1, 1)
+
+        # affine background transformation
+        if bg_param is not None:
+            identity_grid = to_homogeneous(identity_grid)
+            identity_grid = torch.matmul(
+                bg_param.view(bs, 1, 1, 1, 3, 3), identity_grid.unsqueeze(-1)
+            ).squeeze(-1)
+            identity_grid = from_homogeneous(identity_grid)
+
+        transformations = torch.cat([identity_grid, driving_to_source], dim=1)
+        return transformations
+
+    def create_deformed_source_image(self, source_image, transformations):
+        bs, _, h, w = source_image.shape
+        source_repeat = (
+            source_image.unsqueeze(1)
+            .unsqueeze(1)
+            .repeat(1, self.num_tps + 1, 1, 1, 1, 1)
+        )
+        source_repeat = source_repeat.view(bs * (self.num_tps + 1), -1, h, w)
+        transformations = transformations.view((bs * (self.num_tps + 1), h, w, -1))
+        deformed = F.grid_sample(source_repeat, transformations, align_corners=True)
+        deformed = deformed.view((bs, self.num_tps + 1, -1, h, w))
+        return deformed
+
+    def dropout_softmax(self, X, P):
+        """
+        Dropout for TPS transformations. Eq(7) and Eq(8) in the paper.
+        """
+        drop = (
+            (torch.rand(X.shape[0], X.shape[1]) < (1 - P)).type(X.type()).to(X.device)
+        )
+        drop[..., 0] = 1
+        drop = drop.repeat(X.shape[2], X.shape[3], 1, 1).permute(2, 3, 0, 1)
+
+        maxx = X.max(1).values.unsqueeze_(1)
+        X = X - maxx
+        X_exp = X.exp()
+        X[:, 1:, ...] /= 1 - P
+        mask_bool = drop == 0
+        X_exp = X_exp.masked_fill(mask_bool, 0)
+        partition = X_exp.sum(dim=1, keepdim=True) + 1e-6
+        return X_exp / partition
+
+    def forward(
+        self,
+        source_image,
+        kp_driving,
+        kp_source,
+        bg_param=None,
+        dropout_flag=False,
+        dropout_p=0,
+    ):
+        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, kp_driving, kp_source
+        )
+        transformations = self.create_transformations(
+            source_image, kp_driving, kp_source, bg_param
+        )
+        deformed_source = self.create_deformed_source_image(
+            source_image, transformations
+        )
+        out_dict["deformed_source"] = deformed_source
+        # out_dict['transformations'] = transformations
+        deformed_source = deformed_source.view(bs, -1, h, w)
+        input = torch.cat([heatmap_representation, deformed_source], dim=1)
+        input = input.view(bs, -1, h, w)
+
+        # prediction = self.hourglass(input, mode = 1)
+        prediction = self.hourglass(input)
+
+        # contribution_maps = self.maps(prediction[-1])
+        contribution_maps = self.maps(prediction)
+        if dropout_flag:
+            contribution_maps = self.dropout_softmax(contribution_maps, dropout_p)
+        else:
+            contribution_maps = F.softmax(contribution_maps, dim=1)
+        out_dict["contribution_maps"] = contribution_maps
+        out_dict["mask"] = contribution_maps
+
+        # Combine the K+1 transformations
+        # Eq(6) in the paper
+        contribution_maps = contribution_maps.unsqueeze(2)
+        transformations = transformations.permute(0, 1, 4, 2, 3)
+        deformation = (transformations * contribution_maps).sum(dim=1)
+        deformation = deformation.permute(0, 2, 3, 1)
+
+        out_dict["deformation"] = deformation  # Optical Flow
+
+        occlusion_map = []
+        if self.multi_mask:
+            for i in range(self.occlusion_num - self.up_nums):
+                occlusion_map.append(
+                    torch.sigmoid(
+                        self.occlusion[i](
+                            prediction[self.up_nums - self.occlusion_num + i]
+                        )
+                    )
+                )
+            prediction = prediction[-1]
+            for i in range(self.up_nums):
+                prediction = self.up[i](prediction)
+                occlusion_map.append(
+                    torch.sigmoid(
+                        self.occlusion[i + self.occlusion_num - self.up_nums](
+                            prediction
+                        )
+                    )
+                )
+        else:
+            # occlusion_map.append(torch.sigmoid(self.occlusion[0](prediction[-1])))
+            occlusion_map = self.occlusion[0](prediction)
+
+        out_dict["occlusion"] = occlusion_map  # Multi-resolution Occlusion Masks
+        return out_dict

module_mrfa/generator.py

@@ -0,0 +1,88 @@
+# import math
+
+import torch
+# import torch.nn.functional as F
+from torch import nn
+
+from .util import ChannelBlock2d, DownBlock2d, ResBlock2d, SameBlock2d, UpBlock2d
+
+
+class OcclusionAwareGenerator(nn.Module):
+    def __init__(self, num_channels, block_expansion, max_features, num_up_blocks):
+        super(OcclusionAwareGenerator, self).__init__()
+
+        self.num_up_blocks = num_up_blocks
+        self.first = SameBlock2d(
+            num_channels, block_expansion, kernel_size=(7, 7), padding=(3, 3)
+        )
+        down_blocks = []
+        up_blocks = []
+        resblock = []
+        channel_block = []
+        for i in range(num_up_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)
+                )
+            )
+            decoder_in_feature = out_features
+            if i == num_up_blocks - 1:
+                decoder_in_feature = out_features
+            up_blocks.append(
+                UpBlock2d(
+                    decoder_in_feature, in_features, kernel_size=(3, 3), padding=(1, 1)
+                )
+            )
+            resblock.append(
+                ResBlock2d(decoder_in_feature, kernel_size=(3, 3), padding=(1, 1))
+            )
+            channel_block.append(
+                ChannelBlock2d(
+                    decoder_in_feature * 2, kernel_size=(3, 3), padding=(1, 1)
+                )
+            )
+        self.down_blocks = nn.ModuleList(down_blocks)
+        self.up_blocks = nn.ModuleList(up_blocks[::-1])
+        self.resblock = nn.ModuleList(resblock[::-1])
+        self.channel_block = nn.ModuleList(channel_block[::-1])
+        self.final = nn.Conv2d(
+            block_expansion, num_channels, kernel_size=(7, 7), padding=(3, 3)
+        )
+
+    def encode(self, x):
+        f = []
+        out = self.first(x)
+        f.append(out)
+        for i in range(len(self.down_blocks)):
+            out = self.down_blocks[i](out)
+            f.append(out)
+        # 256 128 64 32,...,16,8
+        return f[::-1]
+
+    def decode(self, warp_f, warp_img, occlusion, warp_f_c=None, occlusion_c=None):
+        # warp_f 32 64 128 256
+        out = warp_f[0] * occlusion[0]
+        # out = warp_f[0]
+        if warp_f_c is not None:
+            out_c = warp_f_c[0]
+            out = torch.cat([out, out_c], dim=1)
+        for i in range(self.num_up_blocks):
+            if warp_f_c is not None:
+                out = self.channel_block[i](out)
+            out = self.resblock[i](out)
+            out = self.up_blocks[i](out)
+            out = warp_f[i + 1] * occlusion[i + 1] + out * (1 - occlusion[i + 1])
+            if warp_f_c is not None and i != self.num_up_blocks - 1:
+                out_c = warp_f_c[i + 1]
+                out = torch.cat([out, out_c], dim=1)
+        out = self.final(out)
+        out = torch.sigmoid(out)
+        out = out * (1 - occlusion[-1]) + warp_img * occlusion[-1]
+        return out
+
+    def forward(self, x):
+        x = self.encode(x)
+        x = self.decode(x)
+        return x

module_mrfa/kp_detector.py

@@ -0,0 +1,213 @@
+# from collections import OrderedDict
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+# from torch.jit.annotations import Dict
+# from torch.nn import BatchNorm2d
+from torchvision import models
+
+# from torchvision.models._utils import IntermediateLayerGetter
+from .util import AntiAliasInterpolation2d, DownBlock2d, Hourglass, make_coordinate_grid
+
+
+class KPDetector(nn.Module):
+    """
+    Detecting a keypoints. Return keypoint position and jacobian near each keypoint.
+    """
+
+    def __init__(
+        self,
+        block_expansion=32,
+        num_kp=15,
+        num_channels=3,
+        max_features=1024,
+        num_blocks=5,
+        temperature=0.1,
+        scale_factor=0.25,
+        estimate_jacobian=False,
+        estimate_occlusion=False,
+    ):
+        super(KPDetector, self).__init__()
+
+        self.predictor = Hourglass(
+            block_expansion,
+            in_features=num_channels,
+            max_features=max_features,
+            num_blocks=num_blocks,
+        )
+
+        self.kp = nn.Conv2d(
+            in_channels=self.predictor.out_filters,
+            out_channels=num_kp,
+            kernel_size=(7, 7),
+            padding=0,
+        )
+        self.num_kp = num_kp
+
+        self.estimate_jacobian = estimate_jacobian
+        if estimate_jacobian:
+            self.num_jacobian_maps = 1
+            self.jacobian = nn.Conv2d(
+                in_channels=self.predictor.out_filters,
+                out_channels=4 * self.num_jacobian_maps,
+                kernel_size=(7, 7),
+                padding=0,
+            )
+            self.jacobian.weight.data.zero_()
+            self.jacobian.bias.data.copy_(
+                torch.tensor([1, 0, 0, 1] * self.num_jacobian_maps, dtype=torch.float)
+            )
+
+        self.estimate_occlusion = estimate_occlusion
+        if estimate_occlusion:
+            kp_occlusion = [
+                DownBlock2d(
+                    self.predictor.out_filters,
+                    block_expansion,
+                    kernel_size=3,
+                    padding=1,
+                )
+            ]
+            kp_occlusion.append(
+                DownBlock2d(
+                    block_expansion, block_expansion * 2, kernel_size=3, padding=1
+                )
+            )
+            kp_occlusion.append(
+                DownBlock2d(
+                    block_expansion * 2, block_expansion * 3, kernel_size=3, padding=1
+                )
+            )
+            kp_occlusion.append(
+                DownBlock2d(
+                    block_expansion * 3, block_expansion * 4, kernel_size=3, padding=1
+                )
+            )
+            kp_occlusion.append(
+                nn.Conv2d(
+                    block_expansion * 4, num_kp, kernel_size=(4, 4), padding=0, stride=4
+                )
+            )
+            self.kp_occlusion = nn.Sequential(*kp_occlusion)
+
+        self.temperature = temperature
+        self.scale_factor = scale_factor
+        if self.scale_factor != 1:
+            self.down = AntiAliasInterpolation2d(num_channels, self.scale_factor)
+
+    def create_part_flows(self, source_image, kp_driving, kp_source, bg_params=None):
+        """
+        Eq 4. in the paper T_{s<-d}(z)
+        """
+        bs, _, h, w = source_image.shape
+        num_kp = kp_driving["kp"].shape[1]
+        identity_grid = make_coordinate_grid((h, w), type=kp_source["kp"].type())
+        identity_grid = identity_grid.view(1, 1, h, w, 2)
+        coordinate_grid = identity_grid - kp_driving["kp"].view(bs, num_kp, 1, 1, 2)
+        if "jacobian" in kp_driving:
+            jacobian = torch.matmul(
+                kp_source["jacobian"], torch.inverse(kp_driving["jacobian"])
+            )
+            jacobian = jacobian.unsqueeze(-3).unsqueeze(-3)
+            jacobian = jacobian.repeat(1, 1, h, w, 1, 1)
+            coordinate_grid = torch.matmul(jacobian, coordinate_grid.unsqueeze(-1))
+            coordinate_grid = coordinate_grid.squeeze(-1)
+            # n * 10 * h * w * 2
+        driving_to_source = coordinate_grid + kp_source["kp"].view(bs, num_kp, 1, 1, 2)
+        flow = driving_to_source - identity_grid
+        flow = (h - 1) * flow / 2.0
+        return flow.view(bs, num_kp * 2, h, w)
+
+        # # #adding background feature
+        # # identity_grid = identity_grid.repeat(bs, 1, 1, 1, 1)
+
+        # # 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 gaussian2kp(self, heatmap):
+        """
+        Extract the mean and from a heatmap
+        """
+        shape = heatmap.shape
+        heatmap = heatmap.unsqueeze(-1)
+        grid = (
+            make_coordinate_grid(shape[2:], heatmap.type()).unsqueeze_(0).unsqueeze_(0)
+        )
+        kp = (heatmap * grid).sum(dim=(2, 3))  # N * 10 * 2
+        # kp = {'value': value}
+
+        return kp
+
+    def forward(self, x):
+        if self.scale_factor != 1:
+            x = self.down(x)
+        feature_map = self.predictor(x)
+        prediction = self.kp(feature_map)
+        final_shape = prediction.shape
+        heatmap = prediction.view(final_shape[0], final_shape[1], -1)
+        heatmap = F.softmax(heatmap / self.temperature, dim=2)
+        heatmap = heatmap.view(*final_shape)
+        out = self.gaussian2kp(heatmap)
+        if self.estimate_jacobian:
+            jacobian_map = self.jacobian(feature_map)
+            jacobian_map = jacobian_map.reshape(
+                final_shape[0],
+                self.num_jacobian_maps,
+                4,
+                final_shape[2],
+                final_shape[3],
+            )
+            heatmap = heatmap.unsqueeze(2)
+            jacobian = heatmap * 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)
+            # N * 10 *2 *2
+            out = {"kp": out, "jacobian": jacobian}
+        else:
+            out = {"kp": out}
+
+        if self.estimate_occlusion:
+            # N*10*1*1
+            kp_occlusion = self.kp_occlusion(feature_map)
+            kp_occlusion = torch.sigmoid(kp_occlusion)
+            out["kp_occlusion"] = kp_occlusion
+        # out['heatmap'] = heatmap
+        # out['feature_map'] = feature_map
+        return out
+
+
+class TPSKPDetector(nn.Module):
+    """
+    Predict K*5 keypoints.
+    """
+
+    def __init__(self, num_tps, **kwargs):
+        super(TPSKPDetector, self).__init__()
+        self.num_tps = num_tps
+
+        self.fg_encoder = models.resnet18(pretrained=False)
+        num_features = self.fg_encoder.fc.in_features
+        self.fg_encoder.fc = nn.Linear(num_features, num_tps * 5 * 2)
+
+    def forward(self, image):
+        fg_kp = self.fg_encoder(image)
+        (
+            bs,
+            _,
+        ) = fg_kp.shape
+        fg_kp = torch.sigmoid(fg_kp)
+        fg_kp = fg_kp * 2 - 1
+        out = {"kp": fg_kp.view(bs, self.num_tps * 5, -1)}
+
+        return out

module_mrfa/model.py

@@ -0,0 +1,294 @@
+# from __future__ import absolute_import, division, print_function
+
+# import logging
+
+# import timm
+# import math
+# import os
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import grad
+from torchvision import models
+
+from .bg_motion_predictor import BGMotionPredictor
+from .dense_motion import DenseMotionNetwork, TPSDenseMotionNetwork
+
+# from .generator import OcclusionAwareGenerator
+from .kp_detector import KPDetector, TPSKPDetector
+from .raft import RaftFlow
+from .transformer.pose_tokenpose_b import get_pose_net
+from .util import AntiAliasInterpolation2d, make_coordinate_grid
+
+
+class Transform:
+    """
+    Random tps transformation for equivariance constraints. See Sec 3.3
+    """
+
+    def __init__(self, bs, **kwargs):
+        noise = torch.normal(
+            mean=0, std=kwargs["sigma_affine"] * torch.ones([bs, 2, 3])
+        )
+        self.theta = noise + torch.eye(2, 3).view(1, 2, 3)
+        self.bs = bs
+
+        if ("sigma_tps" in kwargs) and ("points_tps" in kwargs):
+            self.tps = True
+            self.control_points = make_coordinate_grid(
+                (kwargs["points_tps"], kwargs["points_tps"]), type=noise.type()
+            )
+            self.control_points = self.control_points.unsqueeze(0)
+            self.control_params = torch.normal(
+                mean=0,
+                std=kwargs["sigma_tps"]
+                * torch.ones([bs, 1, kwargs["points_tps"] ** 2]),
+            )
+        else:
+            self.tps = False
+
+    def transform_frame(self, frame):
+        grid = make_coordinate_grid(frame.shape[2:], type=frame.type()).unsqueeze(0)
+        grid = grid.view(1, frame.shape[2] * frame.shape[3], 2)
+        grid = self.warp_coordinates(grid).view(
+            self.bs, frame.shape[2], frame.shape[3], 2
+        )
+        return F.grid_sample(frame, grid, padding_mode="reflection")
+
+    def warp_coordinates(self, coordinates):
+        theta = self.theta.type(coordinates.type())
+        theta = theta.unsqueeze(1)
+        # N * (HW) * 2 * 1
+        transformed = (
+            torch.matmul(theta[:, :, :, :2], coordinates.unsqueeze(-1))
+            + theta[:, :, :, 2:]
+        )
+        transformed = transformed.squeeze(-1)
+
+        if self.tps:
+            control_points = self.control_points.type(coordinates.type())
+            control_params = self.control_params.type(coordinates.type())
+            # 1 * (HW) * 1 * 2 - 1 * 1 * 25 *2
+            distances = coordinates.view(
+                coordinates.shape[0], -1, 1, 2
+            ) - control_points.view(1, 1, -1, 2)
+            distances = torch.abs(distances).sum(-1)
+
+            result = distances**2
+            result = result * torch.log(distances + 1e-6)
+            result = result * control_params
+            result = result.sum(dim=2).view(self.bs, coordinates.shape[1], 1)
+            transformed = transformed + result
+
+        return transformed
+
+    def jacobian(self, coordinates):
+        new_coordinates = self.warp_coordinates(coordinates)
+        grad_x = grad(new_coordinates[..., 0].sum(), coordinates, create_graph=True)
+        grad_y = grad(new_coordinates[..., 1].sum(), coordinates, create_graph=True)
+        jacobian = torch.cat([grad_x[0].unsqueeze(-2), grad_y[0].unsqueeze(-2)], dim=-2)
+        return jacobian
+
+
+class Vgg19(torch.nn.Module):
+    """
+    Vgg19 network for perceptual loss. See Sec 3.3.
+    """
+
+    def __init__(self, requires_grad=False):
+        super(Vgg19, self).__init__()
+        vgg_pretrained_features = models.vgg19(pretrained=True).features
+        self.slice1 = torch.nn.Sequential()
+        self.slice2 = torch.nn.Sequential()
+        self.slice3 = torch.nn.Sequential()
+        self.slice4 = torch.nn.Sequential()
+        self.slice5 = torch.nn.Sequential()
+        for x in range(2):
+            self.slice1.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(2, 7):
+            self.slice2.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(7, 12):
+            self.slice3.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(12, 21):
+            self.slice4.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(21, 30):
+            self.slice5.add_module(str(x), vgg_pretrained_features[x])
+
+        self.mean = torch.nn.Parameter(
+            data=torch.Tensor(np.array([0.485, 0.456, 0.406]).reshape((1, 3, 1, 1))),
+            requires_grad=False,
+        )
+        self.std = torch.nn.Parameter(
+            data=torch.Tensor(np.array([0.229, 0.224, 0.225]).reshape((1, 3, 1, 1))),
+            requires_grad=False,
+        )
+
+        if not requires_grad:
+            for param in self.parameters():
+                param.requires_grad = False
+
+    def forward(self, X):
+        # X = X.clamp(-1, 1)
+        # X = X / 2 + 0.5
+        X = (X - self.mean) / self.std
+        h_relu1 = self.slice1(X)
+        h_relu2 = self.slice2(h_relu1)
+        h_relu3 = self.slice3(h_relu2)
+        h_relu4 = self.slice4(h_relu3)
+        h_relu5 = self.slice5(h_relu4)
+        out = [h_relu1, h_relu2, h_relu3, h_relu4, h_relu5]
+        return out
+
+
+class ImagePyramide(torch.nn.Module):
+    """
+    Create image pyramide for computing pyramide perceptual loss. See Sec 3.3
+    """
+
+    def __init__(self, scales, num_channels=3):
+        super(ImagePyramide, self).__init__()
+        downs = {}
+        for scale in scales:
+            downs[str(scale).replace(".", "-")] = AntiAliasInterpolation2d(
+                num_channels, scale
+            )
+        self.downs = nn.ModuleDict(downs)
+
+    def forward(self, x):
+        out_dict = {}
+        for scale, down_module in self.downs.items():
+            out_dict["prediction_" + str(scale).replace("-", ".")] = down_module(x)
+        return out_dict
+
+
+class MRFA(nn.Module):
+    def __init__(self, cfg, **kwargs):
+        super(MRFA, self).__init__()
+
+        self.cfg = cfg
+        train_params = cfg["train_params"]
+        self.train_params = train_params
+        self.scales = train_params["scales"]
+        self.loss_weights = train_params["loss_weights"]
+        self.pyramid = ImagePyramide(self.scales).cuda()
+        if sum(self.loss_weights["perceptual"]) != 0:
+            self.vgg = Vgg19().cuda()
+
+        ##################################################
+        if train_params["prior_model"] == "fomm":
+            self.encoder = KPDetector(**cfg["fomm_kp_detector"])
+            self.dense_motion = DenseMotionNetwork(**cfg["dense_motion"])
+        elif train_params["prior_model"] == "tpsm":
+            self.encoder = TPSKPDetector(**cfg["tpsm_kp_detector"])
+            self.dense_motion = TPSDenseMotionNetwork(**cfg["tpsm_dense_motion"])
+            self.dropout_epoch = train_params["dropout_epoch"]
+            self.dropout_maxp = train_params["dropout_maxp"]
+            self.dropout_startp = train_params["dropout_startp"]
+            self.dropout_inc_epoch = train_params["dropout_inc_epoch"]
+        elif train_params["prior_model"] == "mtia":
+            self.encoder = get_pose_net(cfg["mtia_kp_detector"], is_train=True)
+            self.dense_motion = DenseMotionNetwork(**cfg["dense_motion"])
+        else:
+            print("please specify the prior motion model")
+
+        if train_params["bg_start"] < train_params["num_epochs"]:
+            self.bg_predictor = BGMotionPredictor()
+        self.bg_start = train_params["bg_start"]
+
+        self.decoder = RaftFlow(**cfg["raft_flow"])
+        self.down = AntiAliasInterpolation2d(3, 0.25)
+
+    def forward(self, x, epoch=100, is_train=True):
+        kp_s = self.encoder(x["source"])
+        kp_d = self.encoder(x["driving"])
+
+        img_down = self.down(x["source"])
+        if epoch >= self.bg_start:
+            bg_param = self.bg_predictor(x["source"], x["driving"])
+        else:
+            bg_param = None
+
+        if self.train_params["prior_model"] == "tpsm" and epoch < self.dropout_epoch:
+            # dropout_p will linearly increase from dropout_startp to dropout_maxp
+            dropout_flag = True
+            dropout_p = min(
+                epoch / self.dropout_inc_epoch * self.dropout_maxp
+                + self.dropout_startp,
+                self.dropout_maxp,
+            )
+        else:
+            dropout_flag = False
+            dropout_p = 0
+        dense_motion = self.dense_motion(
+            x["source"],
+            kp_d,
+            kp_s,
+            bg_param=bg_param,
+            dropout_flag=dropout_flag,
+            dropout_p=dropout_p,
+        )
+
+        if self.train_params["prior_model"] == "tpsm":
+            kp_s_value = kp_s["kp"].view(x["source"].shape[0], -1, 5, 2).mean(2)
+            kp_d_value = kp_d["kp"].view(x["driving"].shape[0], -1, 5, 2).mean(2)
+        else:
+            kp_s_value = kp_s["kp"]
+            kp_d_value = kp_d["kp"]
+
+        gen, warp_img, occlusion = self.decoder(
+            kp_s_value, kp_d_value, dense_motion, img=img_down, img_full=x["source"]
+        )
+        warp_img = torch.cat([warp_img, occlusion.repeat(1, 3, 1, 1)], dim=3)
+
+        loss_values = {}
+
+        if not is_train:
+            return gen, warp_img, loss_values, kp_s["kp"], kp_d["kp"]
+
+        pyramide_real = self.pyramid(x["driving"])
+        pyramide_generated = self.pyramid(gen)
+        value_total = 0
+        for scale in self.scales:
+            x_vgg = self.vgg(pyramide_generated["prediction_" + str(scale)])
+            y_vgg = self.vgg(pyramide_real["prediction_" + str(scale)])
+
+            for i, weight in enumerate(self.loss_weights["perceptual"]):
+                value = torch.abs(x_vgg[i] - y_vgg[i].detach()).mean()
+                value_total += self.loss_weights["perceptual"][i] * value
+            loss_values["perceptual"] = value_total
+
+        if self.loss_weights["equivariance"] != 0:
+            transform = Transform(
+                x["driving"].shape[0], **self.train_params["transform_params"]
+            )
+            transformed_frame = transform.transform_frame(x["driving"])
+            transformed_kp = self.encoder(transformed_frame)
+            value = torch.abs(
+                kp_d["kp"] - transform.warp_coordinates(transformed_kp["kp"])
+            )
+            value = value.mean()
+            loss_values["equivariance"] = self.loss_weights["equivariance"] * value
+
+            if self.loss_weights["equivariance_jacobian"] != 0:
+                jacobian_transformed = torch.matmul(
+                    transform.jacobian(transformed_kp["kp"]), transformed_kp["jacobian"]
+                )
+                normed_driving = torch.inverse(kp_d["jacobian"])
+                normed_transformed = jacobian_transformed
+                value = torch.matmul(normed_driving, normed_transformed)
+                eye = torch.eye(2).view(1, 1, 2, 2).type(value.type())
+                value = torch.abs(eye - value)
+                loss_values["equivariance_jacobian"] = (
+                    self.loss_weights["equivariance_jacobian"] * value
+                )
+
+        if epoch >= self.bg_start:
+            bg_param_reverse = self.bg_predictor(x["driving"], x["source"])
+            value = torch.matmul(bg_param, bg_param_reverse)
+            eye = torch.eye(3).view(1, 1, 3, 3).type(value.type())
+            value = torch.abs(eye - value).mean()
+            loss_values["bg"] = 10 * value
+
+        return gen, warp_img, loss_values, kp_s["kp"], kp_d["kp"]

module_mrfa/raft.py

@@ -0,0 +1,456 @@
+import math
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+from timm.models.layers.weight_init import trunc_normal_
+from torch import nn
+
+from .generator import OcclusionAwareGenerator
+from .util import (
+    # AntiAliasInterpolation2d,
+    Hourglass,
+    batch_bilinear_sampler,
+    bilinear_sampler,
+    coords_grid,
+    kp2gaussian,
+    # make_coordinate_grid,
+)
+
+
+class CorrBlock:
+    def __init__(self, corr, num_levels=2, radius=3):
+        self.num_levels = num_levels
+        self.radius = radius
+        self.corr_pyramid = []
+
+        self.corr_pyramid.append(corr)
+        for i in range(self.num_levels - 1):
+            corr = F.avg_pool2d(corr, 2, stride=2)
+            self.corr_pyramid.append(corr)
+
+    def __call__(self, coords):
+        r = self.radius
+        coords = coords.permute(0, 2, 3, 1)
+        batch, h1, w1, _ = coords.shape
+
+        out_pyramid = []
+        for i in range(self.num_levels):
+            corr = self.corr_pyramid[i]
+            dx = torch.linspace(-r, r, 2 * r + 1, device=coords.device)
+            dy = torch.linspace(-r, r, 2 * r + 1, device=coords.device)
+            delta = torch.stack(torch.meshgrid(dy, dx), axis=-1)
+
+            centroid_lvl = coords.reshape(batch * h1 * w1, 1, 1, 2) / 2**i
+            delta_lvl = delta.view(1, 2 * r + 1, 2 * r + 1, 2)
+            coords_lvl = centroid_lvl + delta_lvl
+
+            if batch > 1 and h1 >= 128:
+                corr = batch_bilinear_sampler(
+                    corr, coords_lvl, h=h1, w=w1, mini_batch=1
+                )
+            else:
+                corr = bilinear_sampler(corr, coords_lvl)
+            # corr = bilinear_sampler(corr, coords_lvl)
+            corr = corr.view(batch, h1, w1, -1)
+            out_pyramid.append(corr)
+
+        out = torch.cat(out_pyramid, dim=-1)
+        return out.permute(0, 3, 1, 2).contiguous().float()
+
+
+class BasicMotionEncoder(nn.Module):
+    def __init__(self, num_levels=2, radius=3):
+        super(BasicMotionEncoder, self).__init__()
+        cor_planes = num_levels * (2 * radius + 1) ** 2
+        self.convc1 = nn.Conv2d(cor_planes, 128, 1, padding=0)
+        self.convc2 = nn.Conv2d(128, 96, 3, padding=1)
+        self.convf1 = nn.Conv2d(2, 128, 7, padding=3)
+        self.convf2 = nn.Conv2d(128, 64, 3, padding=1)
+        self.conv = nn.Conv2d(64 + 96, 128 - 2, 3, padding=1)
+
+    def forward(self, delta_flow, corr):
+        cor = F.relu(self.convc1(corr))
+        cor = F.relu(self.convc2(cor))
+        flo = F.relu(self.convf1(delta_flow))
+        flo = F.relu(self.convf2(flo))
+
+        cor_flo = torch.cat([cor, flo], dim=1)
+        out = F.relu(self.conv(cor_flo))
+        return torch.cat([out, delta_flow], dim=1)
+
+
+class RefineFlow(nn.Module):
+    def __init__(self):
+        super(RefineFlow, self).__init__()
+        self.convc1 = nn.Conv2d(192, 128, 3, padding=1)
+        # self.convc2 = nn.Conv2d(128, 64, 3, padding=1)
+        self.conv1 = nn.Conv2d(256, 128, 3, padding=1)
+        self.conv2 = nn.Conv2d(128, 2, 3, padding=1)
+        self.convo1 = nn.Conv2d(256, 128, 3, padding=1)
+        self.convo2 = nn.Conv2d(128, 1, 3, padding=1)
+
+    def forward(self, m_f, warp_f):
+        c = F.relu(self.convc1(warp_f))
+        ## c = F.relu(self.convc2(c))
+        inp = torch.cat([m_f, c], dim=1)
+        flow = self.conv2(F.relu(self.conv1(inp)))
+        occ = self.convo2(F.relu(self.convo1(inp)))
+        out = torch.cat([flow, occ], dim=1)
+        return out, inp
+        # return flow, inp
+
+
+class RaftFlow(nn.Module):
+    def __init__(
+        self,
+        prior_only=False,
+        num_kp=10,
+        dim=256,
+        size=256,
+        generator=None,
+        driving_encoder=None,
+        source_encoder=None,
+    ):
+        super(RaftFlow, self).__init__()
+
+        self.scale = dim**-0.5
+        self.size = size
+        ## basic flow resolution of the prior motion model
+        self.h = size // 4
+        self.w = size // 4
+        self.prior_only = prior_only
+        self.generator = OcclusionAwareGenerator(**generator)
+        # self.down =  AntiAliasInterpolation2d(3, 0.25)
+
+        ## start iteration from resolution of size//2**5
+        channels = {
+            size // 32: 512,
+            size // 16: 512,
+            size // 8: 512,
+            size // 4: 256,
+            size // 2: 128,
+            size: 64,
+            # 512: 32,
+        }
+        self.total_iter = int(math.log(2**5, 2)) + 1
+        self.num_iter = int(math.log(2**5, 2)) + 1
+        self.basic_res_index = int(math.log(self.h // (size // 32), 2))
+
+        if not self.prior_only:
+            ### driving and source structure encoder
+            self.kp = Hourglass(**driving_encoder)
+            self.kp_img = Hourglass(**source_encoder)
+            self.kp_head = nn.Conv2d(
+                in_channels=self.kp.out_filters,
+                out_channels=dim,
+                kernel_size=1,
+                padding=0,
+            )
+            self.kp_img_head = nn.Conv2d(
+                in_channels=self.kp_img.out_filters,
+                out_channels=dim,
+                kernel_size=1,
+                padding=0,
+            )
+            # self._make_position_embedding(64, 64, num_kp, 'sine-full')
+            # self.pos_embedding = rearrange(self.pos_embedding, 'b (h w) c -> b c h w', h=self.h, w=self.w)
+            self.pos_embedding = nn.Parameter(torch.zeros(1, num_kp, self.h, self.w))
+            trunc_normal_(self.pos_embedding, std=0.02)
+            ### driving and source structure encoder
+
+            ### flow updater
+            self.corr_enc = BasicMotionEncoder()
+            self.refine = RefineFlow()
+            ## project different-scale warped features to the same dimension
+            self.to_context = nn.ModuleList()
+            for i in range(self.num_iter):
+                f_channel = channels[(size // 32) * (2**i)]
+                self.to_context.append(nn.Conv2d(f_channel, 192, 1, padding=0))
+            ### flow updater
+
+    def forward(self, kp_s, kp_d, dense_motion, img, img_full):
+        feature = self.generator.encode(img_full)
+        if img is None:
+            img = self.down(img_full)
+        b, _, h, w = img.shape
+
+        out_warp_f = []
+        out_occlusion = []
+        out_flow = []
+
+        out_warp_f_c = []
+        out_occlusion_c = []
+
+        ### prior only
+        if self.prior_only:
+            flow = dense_motion["deformation"]
+            occlusion = dense_motion["occlusion"]
+            for i in range(self.total_iter):
+                if flow.shape[2] != feature[i].shape[2]:
+                    flow_res = F.interpolate(
+                        flow.permute(0, 3, 1, 2),
+                        size=feature[i].shape[2:],
+                        mode="bilinear",
+                        align_corners=True,
+                    )
+                    occlusion_res = F.interpolate(
+                        occlusion,
+                        size=feature[i].shape[2:],
+                        mode="bilinear",
+                        align_corners=True,
+                    )
+                else:
+                    flow_res = flow.permute(0, 3, 1, 2)
+                    occlusion_res = occlusion
+                out_warp_f.append(
+                    F.grid_sample(feature[i], flow_res.permute(0, 2, 3, 1))
+                )
+                out_occlusion.append(F.sigmoid(occlusion_res))
+            warp_img = F.grid_sample(img_full, flow_res.permute(0, 2, 3, 1))
+            out = self.generator.decode(out_warp_f, warp_img, out_occlusion)
+            for i in range(len(out_occlusion)):
+                out_occlusion[i] = F.interpolate(
+                    out_occlusion[i],
+                    size=self.size,
+                    mode="bilinear",
+                    align_corners=True,
+                )
+            occlusion = torch.cat(out_occlusion, dim=3)
+            return out, warp_img, occlusion
+        ### prior only
+
+        ### compute correlation volume at basic resolution
+        kp_s = kp2gaussian(kp_s, (h, w), 0.1) + self.pos_embedding
+        kp_d = kp2gaussian(kp_d, (h, w), 0.1) + self.pos_embedding
+        fe_s = self.kp_img(torch.cat([kp_s, img], dim=1))
+        fe_d = self.kp(kp_d)
+        k_s = self.kp_img_head(fe_s)
+        q_d = self.kp_head(fe_d)
+        f_s = rearrange(k_s, "b c h w -> b (h w) c", h=self.h, w=self.w)
+        f_d = rearrange(q_d, "b c h w -> b (h w) c", h=self.h, w=self.w)
+        corr_volume = torch.einsum("bic,bjc->bij", f_d, f_s) * self.scale
+        ### compute correlation volume at basic resolution
+
+        ### prior motion initialization
+        id_grid = coords_grid(b, self.h, self.w, corr_volume.device)
+        init_flow = (self.h - 1) * (
+            dense_motion["deformation"].permute(0, 3, 1, 2) + 1
+        ) / 2.0 - id_grid
+        init_occlusion = dense_motion["occlusion"]
+        ### prior motion initialization
+
+        # ### NPMR only initialization
+        # id_grid = coords_grid(b, self.h, self.w, corr_volume.device)
+        # id_grid = rearrange(id_grid, 'b c h w -> b (h w) c', h=self.h, w=self.w)
+        # init_flow = torch.einsum('bij,bjc->bic', corr_volume.softmax(-1), id_grid)
+        # init_flow = rearrange(init_flow, 'b (h w) c-> b c h w', h=self.h, w=self.w)
+        # id_grid = rearrange(id_grid, 'b (h w) c -> b c h w', h=self.h, w=self.w)
+        # init_flow = init_flow - id_grid
+        # init_occlusion = None
+        # ### NPMR only initialization
+
+        ## start iter resolution: size//32 = (size//4) // 8, where size//4 defines the basic flow resolution
+        flow = (
+            F.interpolate(
+                init_flow, scale_factor=1.0 / 8.0, mode="bilinear", align_corners=True
+            )
+            / 8.0
+        )
+        occlusion = (
+            F.interpolate(
+                init_occlusion,
+                scale_factor=1.0 / 8.0,
+                mode="bilinear",
+                align_corners=True,
+            )
+            if init_occlusion is not None
+            else init_occlusion
+        )
+        ## reshape for pooling the driving image dimension, fitting different resolutions
+        corr_volume = rearrange(
+            corr_volume, "b (h w) n -> (b n) h w", h=self.h, w=self.w
+        ).unsqueeze(1)
+
+        for i in range(self.total_iter):
+            ## identity grid with current-iter resolution
+            id_grid = coords_grid(
+                b,
+                self.size // 32 * (2**i),
+                self.size // 32 * (2**i),
+                corr_volume.device,
+            )
+            flow_sample = flow
+            id_grid_sample = id_grid
+
+            ### computing upsampled or downsampled structure correlation volume
+            if i < self.basic_res_index:
+                corr_volume_res = F.avg_pool2d(
+                    corr_volume,
+                    2 ** (self.basic_res_index - i),
+                    stride=2 ** (self.basic_res_index - i),
+                )
+                scale = 2 ** (self.basic_res_index - i)
+            elif i == self.basic_res_index:
+                corr_volume_res = corr_volume
+                scale = 1
+            elif i > self.basic_res_index:
+                ## for resolution larger than basic flow resolution, sampling with the basic flow resolution
+                corr_volume_res = corr_volume
+                scale = 0.5 ** (i - self.basic_res_index)
+                flow_sample = (
+                    F.interpolate(
+                        flow, size=self.h, mode="bilinear", align_corners=True
+                    )
+                    * scale
+                )
+                id_grid_sample = coords_grid(b, self.h, self.w, corr_volume.device)
+                scale = 1
+            ### computing upsampled or downsampled structure correlation volume
+
+            if i < self.num_iter:
+                ## reshape for pooling the source image dimension, thus to obtain correlation pyramids
+                corr_volume_res = rearrange(
+                    corr_volume_res, "(b n) c h w -> (b h w) c n", n=self.h * self.w
+                )
+                corr_volume_res = rearrange(
+                    corr_volume_res, "b c (p q) -> b c p q", p=self.h, q=self.w
+                )
+                ## building correlation pyramids
+                corr_fn = CorrBlock(corr_volume_res)
+                ## sampling correlation features from the pyramids
+                corr = corr_fn((flow_sample + id_grid_sample) * scale)
+                if i > self.basic_res_index:
+                    ## reshape from basic resolution to current resolution
+                    corr = F.interpolate(
+                        corr, size=flow.shape[2], mode="bilinear", align_corners=True
+                    )
+                ## non-prior-based motion feature
+                m_f = self.corr_enc(flow, corr)
+
+                warp_f = bilinear_sampler(
+                    feature[i], (flow + id_grid).permute(0, 2, 3, 1)
+                )
+                warp_f = F.relu(self.to_context[i](warp_f))
+
+                d_flow, _ = self.refine(m_f, warp_f)
+                flow_w = flow + d_flow[:, 0:2, :, :]
+                d_occ = d_flow[:, 2:, :, :]
+                if occlusion is not None:
+                    occlusion = occlusion + d_occ
+                else:
+                    occlusion = d_occ
+            else:
+                flow_w = flow
+
+            out = bilinear_sampler(feature[i], (flow_w + id_grid).permute(0, 2, 3, 1))
+            out_occlusion.append(F.sigmoid(occlusion))
+            out_warp_f.append(out)
+
+            ## coarse warping (prior motion flow)
+            if i != self.basic_res_index:
+                flow_res = F.interpolate(
+                    dense_motion["deformation"].permute(0, 3, 1, 2),
+                    size=feature[i].shape[2:],
+                    mode="bilinear",
+                    align_corners=True,
+                )
+                occlusion_res = F.interpolate(
+                    dense_motion["occlusion"],
+                    size=feature[i].shape[2:],
+                    mode="bilinear",
+                    align_corners=True,
+                )
+            else:
+                flow_res = dense_motion["deformation"].permute(0, 3, 1, 2)
+                occlusion_res = dense_motion["occlusion"]
+            out_warp_f_c.append(F.grid_sample(feature[i], flow_res.permute(0, 2, 3, 1)))
+            out_occlusion_c.append(F.sigmoid(occlusion_res))
+            ## coarse warping (prior motion flow)
+
+            ### updating flow by directly scaling initial flow
+            if i < self.num_iter - 1:
+                ## scale between basic resolution and updated resolution
+                scale = 2 ** (self.basic_res_index - i) / 2.0
+                d_f = (
+                    F.interpolate(
+                        d_flow[:, 0:2, :, :],
+                        scale_factor=2,
+                        mode="bilinear",
+                        align_corners=True,
+                    )
+                    * 2
+                )
+                ## flow refinement of current iteration
+                flow = (
+                    d_f
+                    + F.interpolate(
+                        init_flow,
+                        size=self.size // 32 * (2 ** (i + 1)),
+                        mode="bilinear",
+                        align_corners=True,
+                    )
+                    / scale
+                )
+                ## flow refinement of previous iterations
+                if i == 0:
+                    d_f_pre = d_f
+                else:
+                    flow = (
+                        flow
+                        + F.interpolate(
+                            d_f_pre, scale_factor=2, mode="bilinear", align_corners=True
+                        )
+                        * 2
+                    )
+                    d_f_pre = (
+                        d_f
+                        + F.interpolate(
+                            d_f_pre, scale_factor=2, mode="bilinear", align_corners=True
+                        )
+                        * 2
+                    )
+                ## occlusion refinement
+                d_occ = F.interpolate(
+                    d_occ, scale_factor=2, mode="bilinear", align_corners=True
+                )
+                occlusion = d_occ + F.interpolate(
+                    init_occlusion,
+                    size=self.size // 32 * (2 ** (i + 1)),
+                    mode="bilinear",
+                    align_corners=True,
+                )
+                if i == 0:
+                    d_occ_pre = d_occ
+                else:
+                    occlusion = occlusion + F.interpolate(
+                        d_occ_pre, scale_factor=2, mode="bilinear", align_corners=True
+                    )
+                    d_occ_pre = d_occ + F.interpolate(
+                        d_occ_pre, scale_factor=2, mode="bilinear", align_corners=True
+                    )
+            ## for num_iter < total_iter
+            # elif i < self.total_iter-1:
+            #     ## without refinement
+            #     flow = F.interpolate(flow, scale_factor=2, mode='bilinear',align_corners=True) * 2
+            #     occlusion = F.interpolate(occlusion, scale_factor=2, mode='bilinear',align_corners=True)
+
+        warp_img = bilinear_sampler(img_full, (flow + id_grid).permute(0, 2, 3, 1))
+        out = self.generator.decode(
+            out_warp_f, warp_img, out_occlusion, out_warp_f_c, out_occlusion_c
+        )
+
+        out_occlusion.append(F.sigmoid(init_occlusion))
+        out_occlusion_vis = []
+        for i in range(len(out_occlusion)):
+            out_occlusion_vis.append(
+                F.interpolate(
+                    out_occlusion[i],
+                    size=self.size,
+                    mode="bilinear",
+                    align_corners=True,
+                )
+            )
+        occlusion = torch.cat(out_occlusion_vis, dim=3)
+
+        return out, warp_img, occlusion

module_mrfa/transformer/hr_base.py

@@ -0,0 +1,727 @@
+# ------------------------------------------------------------------------------
+# Copyright (c) Microsoft
+# Licensed under the MIT License.
+# Written by Bin Xiao (Bin.Xiao@microsoft.com)
+# Modified by Sen Yang (yangsenius@seu.edu.cn)
+# ------------------------------------------------------------------------------
+
+import copy
+import logging
+import os
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+BN_MOMENTUM = 0.1
+logger = logging.getLogger(__name__)
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(
+        in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False
+    )
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(Bottleneck, self).__init__()
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
+        self.conv2 = nn.Conv2d(
+            planes, planes, kernel_size=3, stride=stride, padding=1, bias=False
+        )
+        self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
+        self.conv3 = nn.Conv2d(
+            planes, planes * self.expansion, kernel_size=1, bias=False
+        )
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion, momentum=BN_MOMENTUM)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+def _get_clones(module, N):
+    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])
+
+
+def _get_activation_fn(activation):
+    """Return an activation function given a string"""
+    if activation == "relu":
+        return F.relu
+    if activation == "gelu":
+        return F.gelu
+    if activation == "glu":
+        return F.glu
+    raise RuntimeError(f"activation should be relu/gelu, not {activation}.")
+
+
+class HighResolutionModule(nn.Module):
+    def __init__(
+        self,
+        num_branches,
+        blocks,
+        num_blocks,
+        num_inchannels,
+        num_channels,
+        fuse_method,
+        multi_scale_output=True,
+    ):
+        super(HighResolutionModule, self).__init__()
+        self._check_branches(
+            num_branches, blocks, num_blocks, num_inchannels, num_channels
+        )
+
+        self.num_inchannels = num_inchannels
+        self.fuse_method = fuse_method
+        self.num_branches = num_branches
+
+        self.multi_scale_output = multi_scale_output
+
+        self.branches = self._make_branches(
+            num_branches, blocks, num_blocks, num_channels
+        )
+        self.fuse_layers = self._make_fuse_layers()
+        self.relu = nn.ReLU(True)
+
+    def _check_branches(
+        self, num_branches, blocks, num_blocks, num_inchannels, num_channels
+    ):
+        if num_branches != len(num_blocks):
+            error_msg = "NUM_BRANCHES({}) <> NUM_BLOCKS({})".format(
+                num_branches, len(num_blocks)
+            )
+            logger.error(error_msg)
+            raise ValueError(error_msg)
+
+        if num_branches != len(num_channels):
+            error_msg = "NUM_BRANCHES({}) <> NUM_CHANNELS({})".format(
+                num_branches, len(num_channels)
+            )
+            logger.error(error_msg)
+            raise ValueError(error_msg)
+
+        if num_branches != len(num_inchannels):
+            error_msg = "NUM_BRANCHES({}) <> NUM_INCHANNELS({})".format(
+                num_branches, len(num_inchannels)
+            )
+            logger.error(error_msg)
+            raise ValueError(error_msg)
+
+    def _make_one_branch(self, branch_index, block, num_blocks, num_channels, stride=1):
+        downsample = None
+        if (
+            stride != 1
+            or self.num_inchannels[branch_index]
+            != num_channels[branch_index] * block.expansion
+        ):
+            downsample = nn.Sequential(
+                nn.Conv2d(
+                    self.num_inchannels[branch_index],
+                    num_channels[branch_index] * block.expansion,
+                    kernel_size=1,
+                    stride=stride,
+                    bias=False,
+                ),
+                nn.BatchNorm2d(
+                    num_channels[branch_index] * block.expansion, momentum=BN_MOMENTUM
+                ),
+            )
+
+        layers = []
+        layers.append(
+            block(
+                self.num_inchannels[branch_index],
+                num_channels[branch_index],
+                stride,
+                downsample,
+            )
+        )
+        self.num_inchannels[branch_index] = num_channels[branch_index] * block.expansion
+        for i in range(1, num_blocks[branch_index]):
+            layers.append(
+                block(self.num_inchannels[branch_index], num_channels[branch_index])
+            )
+
+        return nn.Sequential(*layers)
+
+    def _make_branches(self, num_branches, block, num_blocks, num_channels):
+        branches = []
+
+        for i in range(num_branches):
+            branches.append(self._make_one_branch(i, block, num_blocks, num_channels))
+
+        return nn.ModuleList(branches)
+
+    def _make_fuse_layers(self):
+        if self.num_branches == 1:
+            return None
+
+        num_branches = self.num_branches
+        num_inchannels = self.num_inchannels
+        fuse_layers = []
+        for i in range(num_branches if self.multi_scale_output else 1):
+            fuse_layer = []
+            for j in range(num_branches):
+                if j > i:
+                    fuse_layer.append(
+                        nn.Sequential(
+                            nn.Conv2d(
+                                num_inchannels[j],
+                                num_inchannels[i],
+                                1,
+                                1,
+                                0,
+                                bias=False,
+                            ),
+                            nn.BatchNorm2d(num_inchannels[i]),
+                            nn.Upsample(scale_factor=2 ** (j - i), mode="nearest"),
+                        )
+                    )
+                elif j == i:
+                    fuse_layer.append(None)
+                else:
+                    conv3x3s = []
+                    for k in range(i - j):
+                        if k == i - j - 1:
+                            num_outchannels_conv3x3 = num_inchannels[i]
+                            conv3x3s.append(
+                                nn.Sequential(
+                                    nn.Conv2d(
+                                        num_inchannels[j],
+                                        num_outchannels_conv3x3,
+                                        3,
+                                        2,
+                                        1,
+                                        bias=False,
+                                    ),
+                                    nn.BatchNorm2d(num_outchannels_conv3x3),
+                                )
+                            )
+                        else:
+                            num_outchannels_conv3x3 = num_inchannels[j]
+                            conv3x3s.append(
+                                nn.Sequential(
+                                    nn.Conv2d(
+                                        num_inchannels[j],
+                                        num_outchannels_conv3x3,
+                                        3,
+                                        2,
+                                        1,
+                                        bias=False,
+                                    ),
+                                    nn.BatchNorm2d(num_outchannels_conv3x3),
+                                    nn.ReLU(True),
+                                )
+                            )
+                    fuse_layer.append(nn.Sequential(*conv3x3s))
+            fuse_layers.append(nn.ModuleList(fuse_layer))
+
+        return nn.ModuleList(fuse_layers)
+
+    def get_num_inchannels(self):
+        return self.num_inchannels
+
+    def forward(self, x):
+        if self.num_branches == 1:
+            return [self.branches[0](x[0])]
+
+        for i in range(self.num_branches):
+            x[i] = self.branches[i](x[i])
+
+        x_fuse = []
+
+        for i in range(len(self.fuse_layers)):
+            y = x[0] if i == 0 else self.fuse_layers[i][0](x[0])
+            for j in range(1, self.num_branches):
+                if i == j:
+                    y = y + x[j]
+                else:
+                    y = y + self.fuse_layers[i][j](x[j])
+            x_fuse.append(self.relu(y))
+
+        return x_fuse
+
+
+blocks_dict = {"BASIC": BasicBlock, "BOTTLENECK": Bottleneck}
+
+
+class HRNET_base(nn.Module):
+    def __init__(self, cfg, **kwargs):
+        self.inplanes = 64
+        extra = cfg["MODEL"]["EXTRA"]
+        super(HRNET_base, self).__init__()
+
+        # stem net
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM)
+        self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM)
+        self.relu = nn.ReLU(inplace=True)
+        self.layer1 = self._make_layer(Bottleneck, 64, 4)
+
+        self.stage2_cfg = extra["STAGE2"]
+        num_channels = self.stage2_cfg["NUM_CHANNELS"]
+        block = blocks_dict[self.stage2_cfg["BLOCK"]]
+        num_channels = [
+            num_channels[i] * block.expansion for i in range(len(num_channels))
+        ]
+        self.transition1 = self._make_transition_layer([256], num_channels)
+        self.stage2, pre_stage_channels = self._make_stage(
+            self.stage2_cfg, num_channels
+        )
+
+        self.stage3_cfg = extra["STAGE3"]
+        num_channels = self.stage3_cfg["NUM_CHANNELS"]
+        block = blocks_dict[self.stage3_cfg["BLOCK"]]
+        num_channels = [
+            num_channels[i] * block.expansion for i in range(len(num_channels))
+        ]
+        self.transition2 = self._make_transition_layer(pre_stage_channels, num_channels)
+        self.stage3, pre_stage_channels = self._make_stage(
+            self.stage3_cfg, num_channels, multi_scale_output=False
+        )
+
+        self.pretrained_layers = extra["PRETRAINED_LAYERS"]
+
+    def _make_transition_layer(self, num_channels_pre_layer, num_channels_cur_layer):
+        num_branches_cur = len(num_channels_cur_layer)
+        num_branches_pre = len(num_channels_pre_layer)
+
+        transition_layers = []
+        for i in range(num_branches_cur):
+            if i < num_branches_pre:
+                if num_channels_cur_layer[i] != num_channels_pre_layer[i]:
+                    transition_layers.append(
+                        nn.Sequential(
+                            nn.Conv2d(
+                                num_channels_pre_layer[i],
+                                num_channels_cur_layer[i],
+                                3,
+                                1,
+                                1,
+                                bias=False,
+                            ),
+                            nn.BatchNorm2d(num_channels_cur_layer[i]),
+                            nn.ReLU(inplace=True),
+                        )
+                    )
+                else:
+                    transition_layers.append(None)
+            else:
+                conv3x3s = []
+                for j in range(i + 1 - num_branches_pre):
+                    inchannels = num_channels_pre_layer[-1]
+                    outchannels = (
+                        num_channels_cur_layer[i]
+                        if j == i - num_branches_pre
+                        else inchannels
+                    )
+                    conv3x3s.append(
+                        nn.Sequential(
+                            nn.Conv2d(inchannels, outchannels, 3, 2, 1, bias=False),
+                            nn.BatchNorm2d(outchannels),
+                            nn.ReLU(inplace=True),
+                        )
+                    )
+                transition_layers.append(nn.Sequential(*conv3x3s))
+
+        return nn.ModuleList(transition_layers)
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(
+                    self.inplanes,
+                    planes * block.expansion,
+                    kernel_size=1,
+                    stride=stride,
+                    bias=False,
+                ),
+                nn.BatchNorm2d(planes * block.expansion, momentum=BN_MOMENTUM),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def _make_stage(self, layer_config, num_inchannels, multi_scale_output=True):
+        num_modules = layer_config["NUM_MODULES"]
+        num_branches = layer_config["NUM_BRANCHES"]
+        num_blocks = layer_config["NUM_BLOCKS"]
+        num_channels = layer_config["NUM_CHANNELS"]
+        block = blocks_dict[layer_config["BLOCK"]]
+        fuse_method = layer_config["FUSE_METHOD"]
+
+        modules = []
+        for i in range(num_modules):
+            # multi_scale_output is only used last module
+            if not multi_scale_output and i == num_modules - 1:
+                reset_multi_scale_output = False
+            else:
+                reset_multi_scale_output = True
+
+            modules.append(
+                HighResolutionModule(
+                    num_branches,
+                    block,
+                    num_blocks,
+                    num_inchannels,
+                    num_channels,
+                    fuse_method,
+                    reset_multi_scale_output,
+                )
+            )
+            num_inchannels = modules[-1].get_num_inchannels()
+
+        return nn.Sequential(*modules), num_inchannels
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.conv2(x)
+        x = self.bn2(x)
+        x = self.relu(x)
+        x = self.layer1(x)
+
+        x_list = []
+        for i in range(self.stage2_cfg["NUM_BRANCHES"]):
+            if self.transition1[i] is not None:
+                x_list.append(self.transition1[i](x))
+            else:
+                x_list.append(x)
+        y_list = self.stage2(x_list)
+
+        x_list = []
+        for i in range(self.stage3_cfg["NUM_BRANCHES"]):
+            if self.transition2[i] is not None:
+                x_list.append(self.transition2[i](y_list[-1]))
+            else:
+                x_list.append(y_list[i])
+        y_list = self.stage3(x_list)
+        return y_list[0]
+
+    def init_weights(self, pretrained="", print_load_info=False):
+        logger.info("=> init weights from normal distribution")
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.normal_(m.weight, std=0.001)
+                for name, _ in m.named_parameters():
+                    if name in ["bias"]:
+                        nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.ConvTranspose2d):
+                nn.init.normal_(m.weight, std=0.001)
+                for name, _ in m.named_parameters():
+                    if name in ["bias"]:
+                        nn.init.constant_(m.bias, 0)
+
+        if os.path.isfile(pretrained):
+            pretrained_state_dict = torch.load(
+                pretrained, map_location=torch.device("cpu")
+            )
+            logger.info("=> loading pretrained model {}".format(pretrained))
+
+            existing_state_dict = {}
+            for name, m in pretrained_state_dict.items():
+                if (
+                    name.split(".")[0] in self.pretrained_layers
+                    and name in self.state_dict()
+                    or self.pretrained_layers[0] is "*"
+                ):
+                    existing_state_dict[name] = m
+                    if print_load_info:
+                        print(":: {} is loaded from {}".format(name, pretrained))
+            self.load_state_dict(existing_state_dict, strict=False)
+        elif pretrained:
+            logger.error("=> please download pre-trained models first!")
+            raise ValueError("{} is not exist!".format(pretrained))
+
+
+class HRNET4_base(nn.Module):
+    def __init__(self, cfg, **kwargs):
+        self.inplanes = 64
+        extra = cfg["MODEL"]["EXTRA"]
+        super(HRNET4_base, self).__init__()
+
+        # stem net
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM)
+        self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM)
+        self.relu = nn.ReLU(inplace=True)
+        self.layer1 = self._make_layer(Bottleneck, 64, 4)
+
+        # stage 2
+        self.stage2_cfg = extra["STAGE2"]
+        num_channels = self.stage2_cfg["NUM_CHANNELS"]
+        block = blocks_dict[self.stage2_cfg["BLOCK"]]
+        num_channels = [
+            num_channels[i] * block.expansion for i in range(len(num_channels))
+        ]
+        self.transition1 = self._make_transition_layer([256], num_channels)
+        self.stage2, pre_stage_channels = self._make_stage(
+            self.stage2_cfg, num_channels
+        )
+
+        # stage 3
+        self.stage3_cfg = extra["STAGE3"]
+        num_channels = self.stage3_cfg["NUM_CHANNELS"]
+        block = blocks_dict[self.stage3_cfg["BLOCK"]]
+        num_channels = [
+            num_channels[i] * block.expansion for i in range(len(num_channels))
+        ]
+        self.transition2 = self._make_transition_layer(pre_stage_channels, num_channels)
+        self.stage3, pre_stage_channels = self._make_stage(
+            self.stage3_cfg, num_channels
+        )
+
+        # stage 4
+        self.stage4_cfg = extra["STAGE4"]
+        num_channels = self.stage4_cfg["NUM_CHANNELS"]
+        block = blocks_dict[self.stage4_cfg["BLOCK"]]
+        num_channels = [
+            num_channels[i] * block.expansion for i in range(len(num_channels))
+        ]
+        self.transition3 = self._make_transition_layer(pre_stage_channels, num_channels)
+        self.stage4, pre_stage_channels = self._make_stage(
+            self.stage4_cfg, num_channels, multi_scale_output=False
+        )
+
+        self.pretrained_layers = extra["PRETRAINED_LAYERS"]
+
+    def _make_transition_layer(self, num_channels_pre_layer, num_channels_cur_layer):
+        num_branches_cur = len(num_channels_cur_layer)
+        num_branches_pre = len(num_channels_pre_layer)
+
+        transition_layers = []
+        for i in range(num_branches_cur):
+            if i < num_branches_pre:
+                if num_channels_cur_layer[i] != num_channels_pre_layer[i]:
+                    transition_layers.append(
+                        nn.Sequential(
+                            nn.Conv2d(
+                                num_channels_pre_layer[i],
+                                num_channels_cur_layer[i],
+                                3,
+                                1,
+                                1,
+                                bias=False,
+                            ),
+                            nn.BatchNorm2d(num_channels_cur_layer[i]),
+                            nn.ReLU(inplace=True),
+                        )
+                    )
+                else:
+                    transition_layers.append(None)
+            else:
+                conv3x3s = []
+                for j in range(i + 1 - num_branches_pre):
+                    inchannels = num_channels_pre_layer[-1]
+                    outchannels = (
+                        num_channels_cur_layer[i]
+                        if j == i - num_branches_pre
+                        else inchannels
+                    )
+                    conv3x3s.append(
+                        nn.Sequential(
+                            nn.Conv2d(inchannels, outchannels, 3, 2, 1, bias=False),
+                            nn.BatchNorm2d(outchannels),
+                            nn.ReLU(inplace=True),
+                        )
+                    )
+                transition_layers.append(nn.Sequential(*conv3x3s))
+
+        return nn.ModuleList(transition_layers)
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(
+                    self.inplanes,
+                    planes * block.expansion,
+                    kernel_size=1,
+                    stride=stride,
+                    bias=False,
+                ),
+                nn.BatchNorm2d(planes * block.expansion, momentum=BN_MOMENTUM),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def _make_stage(self, layer_config, num_inchannels, multi_scale_output=True):
+        num_modules = layer_config["NUM_MODULES"]
+        num_branches = layer_config["NUM_BRANCHES"]
+        num_blocks = layer_config["NUM_BLOCKS"]
+        num_channels = layer_config["NUM_CHANNELS"]
+        block = blocks_dict[layer_config["BLOCK"]]
+        fuse_method = layer_config["FUSE_METHOD"]
+
+        modules = []
+        for i in range(num_modules):
+            # multi_scale_output is only used last module
+            if not multi_scale_output and i == num_modules - 1:
+                reset_multi_scale_output = False
+            else:
+                reset_multi_scale_output = True
+
+            modules.append(
+                HighResolutionModule(
+                    num_branches,
+                    block,
+                    num_blocks,
+                    num_inchannels,
+                    num_channels,
+                    fuse_method,
+                    reset_multi_scale_output,
+                )
+            )
+            num_inchannels = modules[-1].get_num_inchannels()
+
+        return nn.Sequential(*modules), num_inchannels
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.conv2(x)
+        x = self.bn2(x)
+        x = self.relu(x)
+        x = self.layer1(x)
+
+        # stage 2
+        x_list = []
+        for i in range(self.stage2_cfg["NUM_BRANCHES"]):
+            if self.transition1[i] is not None:
+                x_list.append(self.transition1[i](x))
+            else:
+                x_list.append(x)
+        y_list = self.stage2(x_list)
+
+        # stage 3
+        x_list = []
+        for i in range(self.stage3_cfg["NUM_BRANCHES"]):
+            if self.transition2[i] is not None:
+                x_list.append(self.transition2[i](y_list[-1]))
+            else:
+                x_list.append(y_list[i])
+        y_list = self.stage3(x_list)
+
+        # stage 4
+        x_list = []
+        for i in range(self.stage4_cfg["NUM_BRANCHES"]):
+            if self.transition3[i] is not None:
+                x_list.append(self.transition3[i](y_list[-1]))
+            else:
+                x_list.append(y_list[i])
+        y_list = self.stage4(x_list)
+
+        return y_list[0]
+
+    def init_weights(self, pretrained="", print_load_info=False):
+        logger.info("=> init weights from normal distribution")
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.normal_(m.weight, std=0.001)
+                for name, _ in m.named_parameters():
+                    if name in ["bias"]:
+                        nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.ConvTranspose2d):
+                nn.init.normal_(m.weight, std=0.001)
+                for name, _ in m.named_parameters():
+                    if name in ["bias"]:
+                        nn.init.constant_(m.bias, 0)
+
+        if os.path.isfile(pretrained):
+            pretrained_state_dict = torch.load(pretrained)
+            logger.info("=> loading pretrained model {}".format(pretrained))
+
+            existing_state_dict = {}
+            for name, m in pretrained_state_dict.items():
+                if (
+                    name.split(".")[0] in self.pretrained_layers
+                    and name in self.state_dict()
+                    or self.pretrained_layers[0] is "*"
+                ):
+                    existing_state_dict[name] = m
+                    if print_load_info:
+                        print(":: {} is loaded from {}".format(name, pretrained))
+            self.load_state_dict(existing_state_dict, strict=False)
+        elif pretrained:
+            logger.error("=> please download pre-trained models first!")
+            raise ValueError("{} is not exist!".format(pretrained))

module_mrfa/transformer/pose_tokenpose_b.py

@@ -0,0 +1,61 @@
+import torch
+import torch.nn as nn
+
+from .hr_base import HRNET_base
+from .tokenpose_base import TokenPose_TB_base
+
+
+class TokenPose_B(nn.Module):
+    def __init__(self, cfg, **kwargs):
+        super(TokenPose_B, self).__init__()
+
+        extra = cfg["MODEL"]["EXTRA"]
+        self.cfg = cfg
+        print(cfg["MODEL"])
+        ##################################################
+        self.pre_feature = HRNET_base(cfg, **kwargs)
+        self.transformer = TokenPose_TB_base(
+            feature_size=[cfg["MODEL"]["IMAGE_SIZE"][1] // 4, cfg["MODEL"]["IMAGE_SIZE"][0] // 4],
+            patch_size=[cfg["MODEL"]["PATCH_SIZE"][1], cfg["MODEL"]["PATCH_SIZE"][0]],
+            num_keypoints=cfg["MODEL"]["NUM_JOINTS"],
+            dim=cfg["MODEL"]["DIM"],
+            channels=cfg["MODEL"]["BASE_CHANNEL"],
+            depth=cfg["MODEL"]["TRANSFORMER_DEPTH"],
+            heads=cfg["MODEL"]["TRANSFORMER_HEADS"],
+            mlp_dim=cfg["MODEL"]["DIM"] * cfg["MODEL"]["TRANSFORMER_MLP_RATIO"],
+            apply_init=cfg["MODEL"]["INIT"],
+            hidden_heatmap_dim=cfg["MODEL"]["HEATMAP_SIZE"][1]
+            * cfg["MODEL"]["HEATMAP_SIZE"][0]
+            // 8,
+            heatmap_dim=cfg["MODEL"]["HEATMAP_SIZE"][1] * cfg["MODEL"]["HEATMAP_SIZE"][0],
+            heatmap_size=[cfg["MODEL"]["HEATMAP_SIZE"][1], cfg["MODEL"]["HEATMAP_SIZE"][0]],
+            pos_embedding_type=cfg["MODEL"]["POS_EMBEDDING_TYPE"],
+            estimate_jacobian=cfg["MODEL"]["ESTIMATE_JACOBIAN"],
+            temperature=cfg["MODEL"]["TEMPERATURE"],
+            fix_img2motion_attention=cfg["MODEL"]["FIX_IMG2MOTION_ATTENTION"],
+        )
+        ###################################################
+
+    def forward(self, x):
+        if self.cfg["MODEL"]["DATA_PREPROCESS"]:
+            mean = torch.tensor([0.485, 0.456, 0.406]).to(x.device)
+            std = torch.tensor([0.229, 0.224, 0.225]).to(x.device)
+            mean = mean.view(1, 3, 1, 1)
+            std = std.view(1, 3, 1, 1)
+            x = x - mean
+            x = x / std
+        x = self.pre_feature(x)
+        x = self.transformer(x)
+        return x
+
+    def init_weights(self, pretrained=""):
+        self.pre_feature.init_weights(pretrained)
+
+
+def get_pose_net(cfg, is_train, **kwargs):
+    # cfg = convert_dict_to_attrit_dict(cfg)
+    model = TokenPose_B(cfg, **kwargs)
+    if is_train and cfg["MODEL"]["INIT_WEIGHTS"]:
+        model.init_weights(cfg["MODEL"]["PRETRAINED"])
+
+    return model

module_mrfa/util.py

@@ -0,0 +1,503 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.nn import BatchNorm2d
+
+
+class AttributeDict(dict):
+    def __getattr__(self, attr):
+        try:
+            return self[attr]
+        except AttributeError as e:
+            raise AttributeError(f"{e}")
+
+    def __setattr__(self, attr, value):
+        self[attr] = value
+
+
+def convert_dict_to_attrit_dict(dict_instance):
+    att_dict = AttributeDict()
+    for key, value in dict_instance.items():
+        if isinstance(value, dict):
+            value = convert_dict_to_attrit_dict(value)
+        att_dict.__setattr__(key, value)
+    return att_dict
+
+
+def bilinear_sampler(img, coords, mode="bilinear", mask=False):
+    """Wrapper for grid_sample, uses pixel coordinates"""
+    H, W = img.shape[-2:]
+    xgrid, ygrid = coords.split([1, 1], dim=-1)
+    xgrid = 2 * xgrid / (W - 1) - 1
+    ygrid = 2 * ygrid / (H - 1) - 1
+
+    grid = torch.cat([xgrid, ygrid], dim=-1)
+    img = F.grid_sample(img, grid, align_corners=True)
+    if mask:
+        mask = (xgrid > -1) & (ygrid > -1) & (xgrid < 1) & (ygrid < 1)
+        return img, mask.float()
+    return img
+
+
+def batch_bilinear_sampler(
+    img, coords, mode="bilinear", mask=False, h=256, w=256, mini_batch=4
+):
+    """Wrapper for grid_sample, uses pixel coordinates"""
+    batch = img.shape[0] // (h * w)
+    # img = img.view(batch,h,w,img.shape[1],img.shape[2],img.shape[3])
+    # coords = coords.view(batch,h,w,coords.shape[1],coords.shape[2],coords.shape[3])
+    num_iter = batch // mini_batch
+    sampled = []
+    for i in range(num_iter):
+        img_i = img[i * mini_batch * h * w : (i + 1) * mini_batch * h * w, :, :, :]
+        coords_i = coords[
+            i * mini_batch * h * w : (i + 1) * mini_batch * h * w, :, :, :
+        ]
+        sampled.append(bilinear_sampler(img_i, coords_i))
+    return torch.cat(sampled, dim=0)
+
+
+def coords_grid(batch, ht, wd, device):
+    coords = torch.meshgrid(
+        torch.arange(ht, device=device), torch.arange(wd, device=device)
+    )
+    coords = torch.stack(coords[::-1], dim=0).float()
+    return coords[None].repeat(batch, 1, 1, 1)
+
+
+def kp2gaussian(kp, spatial_size, kp_variance):
+    """
+    Transform a keypoint into gaussian like representation
+    """
+    # mean = kp['value']
+    mean = kp
+    # N * 10 * 2
+
+    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
+    # 1,1,h,w,2
+    coordinate_grid = coordinate_grid.view(*shape)
+    repeats = mean.shape[:number_of_leading_dimensions] + (1, 1, 1)
+    # N * 10 * 1 *1 *1
+    coordinate_grid = coordinate_grid.repeat(*repeats)
+    # N,10,h,w,2
+
+    # Preprocess kp shape
+    shape = mean.shape[:number_of_leading_dimensions] + (1, 1, 2)
+    # N * 10 * 1 * 1 * 2
+    mean = mean.view(*shape)
+
+    mean_sub = coordinate_grid - mean
+
+    out = torch.exp(-0.5 * (mean_sub**2).sum(-1) / kp_variance)
+    # N * 10 * h * w
+
+    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)
+    # xx [[-1,-1/3,1/3,1],[-1,-1/3,1/3,1],[-1,-1/3,1/3,1],[-1,-1/3,1/3,1]]
+    # yy [[-1,-1,-1,-1],[-1/3,-1/3,-1/3,-1/3],[1/3,1/3,1/3,1/3],[1,1,1,1]]
+
+    meshed = torch.cat([xx.unsqueeze_(2), yy.unsqueeze_(2)], 2)
+
+    return meshed
+
+
+class ChannelBlock2d(nn.Module):
+    """
+    Res block, preserve spatial resolution.
+    """
+
+    def __init__(self, in_features, kernel_size, padding):
+        super(ChannelBlock2d, self).__init__()
+        self.conv1 = nn.Conv2d(
+            in_channels=in_features,
+            out_channels=in_features // 2,
+            kernel_size=kernel_size,
+            padding=padding,
+        )
+        # self.conv2 = nn.Conv2d(in_channels=in_features//2, out_channels=in_features//2, 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 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
+
+    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 = F.interpolate(out, scale_factor=(self.scale, self.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]
+
+
+class TPS:
+    """
+    TPS transformation, mode 'kp' for Eq(2) in the paper, mode 'random' for equivariance loss.
+    """
+
+    def __init__(self, mode, bs, **kwargs):
+        self.bs = bs
+        self.mode = mode
+        if mode == "random":
+            noise = torch.normal(
+                mean=0, std=kwargs["sigma_affine"] * torch.ones([bs, 2, 3])
+            )
+            self.theta = noise + torch.eye(2, 3).view(1, 2, 3)
+            self.control_points = make_coordinate_grid(
+                (kwargs["points_tps"], kwargs["points_tps"]), type=noise.type()
+            )
+            self.control_points = self.control_points.unsqueeze(0)
+            self.control_params = torch.normal(
+                mean=0,
+                std=kwargs["sigma_tps"]
+                * torch.ones([bs, 1, kwargs["points_tps"] ** 2]),
+            )
+        elif mode == "kp":
+            kp_1 = kwargs["kp_1"]
+            kp_2 = kwargs["kp_2"]
+            device = kp_1.device
+            kp_type = kp_1.type()
+            self.gs = kp_1.shape[1]
+            n = kp_1.shape[2]
+            K = torch.norm(kp_1[:, :, :, None] - kp_1[:, :, None, :], dim=4, p=2)
+            K = K**2
+            K = K * torch.log(K + 1e-9)
+
+            one1 = (
+                torch.ones(self.bs, kp_1.shape[1], kp_1.shape[2], 1)
+                .to(device)
+                .type(kp_type)
+            )
+            kp_1p = torch.cat([kp_1, one1], 3)
+
+            zero = torch.zeros(self.bs, kp_1.shape[1], 3, 3).to(device).type(kp_type)
+            P = torch.cat([kp_1p, zero], 2)
+            L = torch.cat([K, kp_1p.permute(0, 1, 3, 2)], 2)
+            L = torch.cat([L, P], 3)
+
+            zero = torch.zeros(self.bs, kp_1.shape[1], 3, 2).to(device).type(kp_type)
+            Y = torch.cat([kp_2, zero], 2)
+            one = torch.eye(L.shape[2]).expand(L.shape).to(device).type(kp_type) * 0.01
+            L = L + one
+
+            param = torch.matmul(torch.inverse(L), Y)
+            self.theta = param[:, :, n:, :].permute(0, 1, 3, 2)
+
+            self.control_points = kp_1
+            self.control_params = param[:, :, :n, :]
+        else:
+            raise Exception("Error TPS mode")
+
+    def transform_frame(self, frame):
+        grid = (
+            make_coordinate_grid(frame.shape[2:], type=frame.type())
+            .unsqueeze(0)
+            .to(frame.device)
+        )
+        grid = grid.view(1, frame.shape[2] * frame.shape[3], 2)
+        shape = [self.bs, frame.shape[2], frame.shape[3], 2]
+        if self.mode == "kp":
+            shape.insert(1, self.gs)
+        grid = self.warp_coordinates(grid).view(*shape)
+        return grid
+
+    def warp_coordinates(self, coordinates):
+        theta = self.theta.type(coordinates.type()).to(coordinates.device)
+        control_points = self.control_points.type(coordinates.type()).to(
+            coordinates.device
+        )
+        control_params = self.control_params.type(coordinates.type()).to(
+            coordinates.device
+        )
+
+        if self.mode == "kp":
+            transformed = (
+                torch.matmul(theta[:, :, :, :2], coordinates.permute(0, 2, 1))
+                + theta[:, :, :, 2:]
+            )
+
+            distances = coordinates.view(
+                coordinates.shape[0], 1, 1, -1, 2
+            ) - control_points.view(self.bs, control_points.shape[1], -1, 1, 2)
+
+            distances = distances**2
+            result = distances.sum(-1)
+            result = result * torch.log(result + 1e-9)
+            result = torch.matmul(result.permute(0, 1, 3, 2), control_params)
+            transformed = transformed.permute(0, 1, 3, 2) + result
+
+        elif self.mode == "random":
+            theta = theta.unsqueeze(1)
+            transformed = (
+                torch.matmul(theta[:, :, :, :2], coordinates.unsqueeze(-1))
+                + theta[:, :, :, 2:]
+            )
+            transformed = transformed.squeeze(-1)
+            ances = coordinates.view(
+                coordinates.shape[0], -1, 1, 2
+            ) - control_points.view(1, 1, -1, 2)
+            distances = ances**2
+
+            result = distances.sum(-1)
+            result = result * torch.log(result + 1e-9)
+            result = result * control_params
+            result = result.sum(dim=2).view(self.bs, coordinates.shape[1], 1)
+            transformed = transformed + result
+        else:
+            raise Exception("Error TPS mode")
+
+        return transformed

pyproject.toml

@@ -1,9 +1,9 @@
 [project]
 name = "comfyui-firstordermm"
 description = "ComfyUI-native nodes to run First Order Motion Model for Image Animation and its non-diffusion-based successors. [a/https://github.com/AliaksandrSiarohin/first-order-model](https://github.com/AliaksandrSiarohin/first-order-model)"
-version = "1.0.3"
+version = "1.0.4"
 license = { file = "LICENSE" }
-dependencies = ["numpy", "torch", "scipy", "yaml", "matplotlib"]
+dependencies = ["numpy", "torch", "scipy", "pyyaml", "matplotlib", "einops", "timm"]
 
 [project.urls]
 Repository = "https://github.com/FuouM/ComfyUI-FirstOrderMM"

requirements.txt

@@ -2,4 +2,6 @@ numpy
 torch
 scipy
 pyyaml
-matplotlib
+matplotlib
+einops
+timm

run.py

@@ -12,6 +12,10 @@ import torch
 from .constants import (
     ARTICULATE_CFG_PATH,
     ARTICULATE_MODEL_PATH,
+    MRFA_CFG_PATHS,
+    MRFA_DEFAULT_MODEL,
+    MRFA_MODEL_NAMES,
+    MRFA_MODEL_PATHS,
     SPLINE_CFG_PATH,
     SPLINE_DEFAULT,
     SPLINE_MODEL_PATH,
@@ -29,6 +33,7 @@ from .constants import (
 from .face_parsing.face_parsing_loader import load_face_parser_model
 from .inference_articulate import articulate_inference
 from .inference_fomm import inference, inference_best_frame, load_checkpoint
+from .inference_mrfa import mrfa_inference
 from .inference_partswap import load_partswap_checkpoint, partswap_inference
 from .inference_spline import spline_inference
 from .seg_viz import visualize_frame
@@ -449,7 +454,7 @@ class Spline_Runner:
         find_best_frame: bool,
         audio=None,
     ):
-        print(f"{type(source_image)=}")
+        print(f"{type(source_image)=}")  # [B, H, W, C]
         print(f"{type(driving_video_input)=}")
         print(f"{source_image.shape=}")
         print(f"{driving_video_input.shape=}")
@@ -486,6 +491,99 @@ class Spline_Runner:
         )
 
 
+class MRFA_Runner:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "source_image": ("IMAGE",),
+                "driving_video_input": ("IMAGE",),
+                "model_name": (MRFA_MODEL_NAMES, {"default": MRFA_DEFAULT_MODEL}),
+                "frame_rate": ("FLOAT", {"default": 30.0}),
+                "use_relative": (
+                    "BOOLEAN",
+                    {"default": True},
+                ),
+                "relative_movement": (
+                    "BOOLEAN",
+                    {"default": True},
+                ),
+                "relative_jacobian": (
+                    "BOOLEAN",
+                    {"default": True},
+                ),
+                "adapt_movement_scale": (
+                    "BOOLEAN",
+                    {"default": False},
+                ),
+            },
+            "optional": {"audio": ("AUDIO",)},
+        }
+
+    RETURN_TYPES = (
+        "IMAGE",
+        "AUDIO",
+        "FLOAT",
+    )
+    RETURN_NAMES = (
+        "images",
+        "audio",
+        "frame_rate",
+    )
+    FUNCTION = "todo"
+    CATEGORY = "FirstOrderMM"
+
+    def todo(
+        self,
+        source_image,
+        driving_video_input,
+        model_name: str,
+        frame_rate: float,
+        use_relative: bool,
+        relative_movement: bool,
+        relative_jacobian: bool,
+        adapt_movement_scale: bool,
+        audio=None,
+    ):
+        print(f"{type(source_image)=}")  # [B, H, W, C]
+        print(f"{type(driving_video_input)=}")
+        print(f"{source_image.shape=}")
+        print(f"{driving_video_input.shape=}")
+        print(f"{type(audio)=}")
+        print(base_dir)
+
+        config_path = f"{base_dir}/{MRFA_CFG_PATHS[model_name]}"
+        checkpoint_path = f"{base_dir}/{MRFA_MODEL_PATHS[model_name]}"
+
+        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,
+            "use_relative": use_relative,
+            "relative_movement": relative_movement,
+            "relative_jacobian": relative_jacobian,
+            "adapt_movement_scale": adapt_movement_scale,
+        }
+
+        predictions = mrfa_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/MRFA.json

@@ -0,0 +1,423 @@
+{
+  "last_node_id": 7,
+  "last_link_id": 18,
+  "nodes": [
+    {
+      "id": 2,
+      "type": "LoadImage",
+      "pos": [
+        83,
+        105
+      ],
+      "size": {
+        "0": 315,
+        "1": 314
+      },
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            13
+          ],
+          "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": 5,
+      "type": "VHS_VideoInfoLoaded",
+      "pos": [
+        753,
+        383
+      ],
+      "size": {
+        "0": 304.79998779296875,
+        "1": 106
+      },
+      "flags": {},
+      "order": 2,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "video_info",
+          "type": "VHS_VIDEOINFO",
+          "link": 7
+        }
+      ],
+      "outputs": [
+        {
+          "name": "fps🟦",
+          "type": "FLOAT",
+          "links": [
+            18
+          ],
+          "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": 6,
+      "type": "VHS_VideoCombine",
+      "pos": [
+        1236,
+        163
+      ],
+      "size": [
+        315,
+        615
+      ],
+      "flags": {},
+      "order": 4,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 15
+        },
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "link": 16
+        },
+        {
+          "name": "meta_batch",
+          "type": "VHS_BatchManager",
+          "link": null
+        },
+        {
+          "name": "vae",
+          "type": "VAE",
+          "link": null
+        },
+        {
+          "name": "frame_rate",
+          "type": "FLOAT",
+          "link": 17,
+          "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": "MRFA",
+        "format": "video/h264-mp4",
+        "pix_fmt": "yuv420p",
+        "crf": 0,
+        "save_metadata": true,
+        "pingpong": false,
+        "save_output": true,
+        "videopreview": {
+          "hidden": false,
+          "paused": false,
+          "params": {
+            "filename": "MRFA_00003.mp4",
+            "subfolder": "",
+            "type": "output",
+            "format": "video/h264-mp4",
+            "frame_rate": 29.970029970029966
+          }
+        }
+      }
+    },
+    {
+      "id": 3,
+      "type": "VHS_LoadVideo",
+      "pos": [
+        430,
+        196
+      ],
+      "size": [
+        235.1999969482422,
+        491.1999969482422
+      ],
+      "flags": {
+        "collapsed": false
+      },
+      "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": [
+            14
+          ],
+          "shape": 3,
+          "slot_index": 0
+        },
+        {
+          "name": "frame_count",
+          "type": "INT",
+          "links": null,
+          "shape": 3
+        },
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "links": [],
+          "shape": 3,
+          "slot_index": 2
+        },
+        {
+          "name": "video_info",
+          "type": "VHS_VIDEOINFO",
+          "links": [
+            7
+          ],
+          "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": 7,
+      "type": "MRFA_Runner",
+      "pos": [
+        779,
+        107
+      ],
+      "size": {
+        "0": 380.4000244140625,
+        "1": 218
+      },
+      "flags": {},
+      "order": 3,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "source_image",
+          "type": "IMAGE",
+          "link": 13
+        },
+        {
+          "name": "driving_video_input",
+          "type": "IMAGE",
+          "link": 14
+        },
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "link": null
+        },
+        {
+          "name": "frame_rate",
+          "type": "FLOAT",
+          "link": 18,
+          "widget": {
+            "name": "frame_rate"
+          }
+        }
+      ],
+      "outputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "links": [
+            15
+          ],
+          "shape": 3,
+          "slot_index": 0
+        },
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "links": [
+            16
+          ],
+          "shape": 3,
+          "slot_index": 1
+        },
+        {
+          "name": "frame_rate",
+          "type": "FLOAT",
+          "links": [
+            17
+          ],
+          "shape": 3,
+          "slot_index": 2
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "MRFA_Runner"
+      },
+      "widgets_values": [
+        "celebvhq",
+        30,
+        true,
+        true,
+        true,
+        true
+      ]
+    }
+  ],
+  "links": [
+    [
+      7,
+      3,
+      3,
+      5,
+      0,
+      "VHS_VIDEOINFO"
+    ],
+    [
+      13,
+      2,
+      0,
+      7,
+      0,
+      "IMAGE"
+    ],
+    [
+      14,
+      3,
+      0,
+      7,
+      1,
+      "IMAGE"
+    ],
+    [
+      15,
+      7,
+      0,
+      6,
+      0,
+      "IMAGE"
+    ],
+    [
+      16,
+      7,
+      1,
+      6,
+      1,
+      "AUDIO"
+    ],
+    [
+      17,
+      7,
+      2,
+      6,
+      4,
+      "FLOAT"
+    ],
+    [
+      18,
+      5,
+      0,
+      7,
+      3,
+      "FLOAT"
+    ]
+  ],
+  "groups": [],
+  "config": {},
+  "extra": {
+    "ds": {
+      "scale": 1,
+      "offset": [
+        -17,
+        60
+      ]
+    }
+  },
+  "version": 0.4
+}