Add Thin-Plate Spline Motion Model for Image Animation (#6)

.gitignore

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

README.md

@@ -3,20 +3,15 @@ A ComfyUI-native node to run First Order Motion Model for Image Animation
 
 https://github.com/AliaksandrSiarohin/first-order-model
 
-Now supports Face Swapping using Motion Supervised co-part Segmentation
+Now supports:
+1. Face Swapping using [Motion Supervised co-part Segmentation](https://github.com/AliaksandrSiarohin/motion-cosegmentation): 
+2. [Motion Representations for Articulated Animation](https://github.com/snap-research/articulated-animation)
+3. [Thin-Plate Spline Motion Model for Image Animation](https://github.com/yoyo-nb/thin-plate-spline-motion-model)
 
-https://github.com/AliaksandrSiarohin/motion-cosegmentation
-
-Now also supports Motion Representations for Articulated Animation
-
-https://github.com/snap-research/articulated-animation
+https://github.com/user-attachments/assets/b2948efb-3b44-440b-bff2-dde7b95a9946
 
 ## Workflow:
 
-
-https://github.com/FuouM/ComfyUI-FirstOrderMM/assets/57849916/e8adca97-b53c-48ce-a211-f336dbf7d7bc
-
-
 ### FOMM
 
 [FOMM.json](workflows/FOMM.json)
@@ -35,6 +30,14 @@ https://github.com/FuouM/ComfyUI-FirstOrderMM/assets/57849916/e8adca97-b53c-48ce
 
 ![Workflow Articulate](workflows/workflow_articulate.png)
 
+
+### Spline
+
+[SPLINE.json](workflows/SPLINE.json)
+
+![Workflow Spline](workflows/workflow_spline.png)
+
+
 ## Arguments
 
 ### FOMM
@@ -56,6 +59,14 @@ https://github.com/FuouM/ComfyUI-FirstOrderMM/assets/57849916/e8adca97-b53c-48ce
 
 Doesn't need any
 
+### Spline
+
+* `predict_mode`: Can be
+  * `relative`: Similar to FOMM's `relative_movement` and `adapt_movement_scale` set to True
+  * `standard`: Similar to FOMM's `adapt_movement_scale` set to False
+  * `avd`: similar to `relative`, may yield better but more "jittery/jumpy" result
+* `find_best_frame`: Same as FOMM
+
 ## Installation
 
 1. Clone the repo to `ComfyUI/custom_nodes/`
@@ -113,4 +124,6 @@ resnet18-5c106cde.pth
 79999_iter.pth
 ```
 
-For **Articulate**, download the model from [Pre-trained checkpoints](https://github.com/snap-research/articulated-animation?tab=readme-ov-file#pre-trained-checkpoints) section and place it here: `articulate_module\models\vox256.pth`
+For **Articulate**, download the model from [Pre-trained checkpoints](https://github.com/snap-research/articulated-animation?tab=readme-ov-file#pre-trained-checkpoints) section and place it here: `articulate_module/models/vox256.pth`
+
+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`.

__init__.py

@@ -5,6 +5,7 @@ from .run import (
     FOMM_Seg10Chooser,
     FOMM_Seg15Chooser,
     Articulate_Runner,
+    Spline_Runner
 )
 
 NODE_CLASS_MAPPINGS = {
@@ -14,6 +15,7 @@ NODE_CLASS_MAPPINGS = {
     "FOMM_Seg10Chooser": FOMM_Seg10Chooser,
     "FOMM_Seg15Chooser": FOMM_Seg15Chooser,
     "Articulate_Runner": Articulate_Runner,
+    "Spline_Runner": Spline_Runner,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -23,6 +25,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "FOMM_Seg10Chooser": "FOMM Seg10 Chooser",
     "FOMM_Seg15Chooser": "FOMM Seg15 Chooser",
     "Articulate_Runner": "Articulate Runner",
+    "Spline_Runner": "Spline Runner",
 }
 
 

constants.py

@@ -100,4 +100,10 @@ seg_model_dict = {
 }
 
 ARTICULATE_MODEL_PATH = "articulate_module/models/vox256.pth"
-ARTICULATE_CFG_PATH = "articulate_module/config/vox256.yaml"
+ARTICULATE_CFG_PATH = "articulate_module/config/vox256.yaml"
+
+SPLINE_MODES = ['relative', 'standard', 'avd']
+SPLINE_DEFAULT = 'relative'
+
+SPLINE_MODEL_PATH = "spline_module/models/vox.pth.tar"
+SPLINE_CFG_PATH = "spline_module/config/vox-256.yaml"

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "comfyui-firstordermm"
 description = "A ComfyUI-native node to run First Order Motion Model for Image Animation. [a/https://github.com/AliaksandrSiarohin/first-order-model](https://github.com/AliaksandrSiarohin/first-order-model)"
-version = "1.0.2"
+version = "1.0.3"
 license = { file = "LICENSE" }
 dependencies = ["numpy", "torch", "scipy", "yaml", "matplotlib"]
 

run.py

@@ -13,6 +13,10 @@ from .articulate_inference import articulate_inference
 from .constants import (
     ARTICULATE_CFG_PATH,
     ARTICULATE_MODEL_PATH,
+    SPLINE_CFG_PATH,
+    SPLINE_DEFAULT,
+    SPLINE_MODEL_PATH,
+    SPLINE_MODES,
     config_folder,
     default_model_name,
     default_partswap_model_name,
@@ -29,6 +33,7 @@ from .fomm_loader import load_checkpoint
 from .partswap_inference import partswap_inference
 from .partswap_loader import load_partswap_checkpoint
 from .seg_viz import visualize_frame
+from .spline_inference import spline_inference
 from .utils import (
     build_seg_arguments,
     get_config_path,
@@ -407,6 +412,82 @@ class Articulate_Runner:
         )
 
 
+class Spline_Runner:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "source_image": ("IMAGE",),
+                "driving_video_input": ("IMAGE",),
+                "frame_rate": ("FLOAT", {"default": 30.0}),
+                "predict_mode": (SPLINE_MODES, {"default": SPLINE_DEFAULT}),
+                "find_best_frame": (
+                    "BOOLEAN",
+                    {"default": False},
+                ),  # Generate from the frame that is the most alligned with source
+            },
+            "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,
+        frame_rate: float,
+        predict_mode: str,
+        find_best_frame: bool,
+        audio=None,
+    ):
+        print(f"{type(source_image)=}")
+        print(f"{type(driving_video_input)=}")
+        print(f"{source_image.shape=}")
+        print(f"{driving_video_input.shape=}")
+        print(f"{type(audio)=}")
+        print(base_dir)
+
+        config_path = f"{base_dir}/{SPLINE_CFG_PATH}"
+        checkpoint_path = f"{base_dir}/{SPLINE_MODEL_PATH}"
+
+        source_image = reshape_image(source_image, (256, 256))
+        driving_video = reshape_image(driving_video_input, (256, 256)).unsqueeze(0)
+        driving_video = driving_video.permute(0, 2, 1, 3, 4)
+
+        print("After reshaping")
+        print(f"{source_image.shape=}")
+        print(f"{driving_video.shape=}")
+        params = {
+            "source_image": source_image,
+            "driving_video": driving_video,
+            "config_path": config_path,
+            "checkpoint_path": checkpoint_path,
+            "predict_mode": predict_mode,
+            "find_best_frame": find_best_frame,
+        }
+
+        predictions = spline_inference(**params)
+
+        output_images = out_video(predictions)
+
+        return (
+            output_images,
+            audio,
+            frame_rate,
+        )
+
+
 def serialize_integers(int_list):
     return "_".join(map(str, int_list))
 

spline_inference.py

@@ -0,0 +1,203 @@
+import torch
+import numpy as np
+from scipy.spatial import ConvexHull
+import yaml
+import tqdm
+from comfy.utils import ProgressBar
+
+from .fomm_inference import find_best_frame
+
+from .spline_module.avd_network import AVDNetwork
+from .spline_module.dense_motion import DenseMotionNetwork
+from .spline_module.keypoint_detector import KPDetector
+from .spline_module.inpainting_network import InpaintingNetwork
+
+
+def spline_inference(
+    source_image,
+    driving_video: list,
+    config_path: str,
+    checkpoint_path: str,
+    predict_mode="relative",
+    find_best_frame=False,
+    cpu=False,
+):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    inpainting, kp_detector, dense_motion_network, avd_network = load_checkpoints(
+        config_path, checkpoint_path, device
+    )
+    
+    source_image = source_image.to(device)
+    driving_video = driving_video.to(device)
+
+    if find_best_frame:
+        predictions = spline_inference_best_frame(
+            source_image,
+            driving_video,
+            inpainting,
+            kp_detector,
+            dense_motion_network,
+            avd_network,
+            predict_mode,
+            cpu,
+        )
+    else:
+        predictions = inference(
+            source_image,
+            driving_video,
+            inpainting,
+            kp_detector,
+            dense_motion_network,
+            avd_network,
+            predict_mode,
+        )
+
+    return predictions
+
+
+def spline_inference_best_frame(
+    source_image,
+    driving_video: list,
+    inpainting,
+    kp_detector,
+    dense_motion_network,
+    avd_network,
+    predict_mode,
+    cpu=False,
+):
+    best_frame_idx = find_best_frame(source_image, driving_video, cpu)
+
+    first_half = driving_video[
+        :, :, : best_frame_idx + 1
+    ]  # Include the best frame in the first half
+    first_half = torch.flip(first_half, dims=[2])  # Reverse the first half
+    second_half = driving_video[:, :, best_frame_idx + 1 :]
+
+    predictions_first = inference(
+        source_image,
+        first_half,
+        inpainting,
+        kp_detector,
+        dense_motion_network,
+        avd_network,
+        predict_mode,
+    )
+
+    predictions_second = inference(
+        source_image,
+        second_half,
+        inpainting,
+        kp_detector,
+        dense_motion_network,
+        avd_network,
+        predict_mode,
+    )
+
+    predictions = []
+    predictions_first = predictions_first[::-1]  # Reverse the first half back
+    predictions.extend(predictions_first)
+    predictions.extend(predictions_second)
+    return predictions
+
+
+def inference(
+    source_image,
+    driving_video: list,
+    inpainting,
+    kp_detector,
+    dense_motion_network,
+    avd_network,
+    predict_mode,
+):
+    with torch.no_grad():
+        source = source_image
+        driving = driving_video
+        predictions = []
+
+        kp_source = kp_detector(source)
+        kp_driving_initial = kp_detector(driving[:, :, 0])
+
+        num_frames = driving.shape[2]
+        pbar = ProgressBar(num_frames)
+
+        for frame_idx in tqdm.tqdm(range(num_frames)):
+            driving_frame = driving[:, :, frame_idx]
+            kp_driving = kp_detector(driving_frame)
+            if predict_mode == "standard":
+                kp_norm = kp_driving
+            elif predict_mode == "relative":
+                kp_norm = relative_kp(
+                    kp_source=kp_source,
+                    kp_driving=kp_driving,
+                    kp_driving_initial=kp_driving_initial,
+                )
+            elif predict_mode == "avd":
+                kp_norm = avd_network(kp_source, kp_driving)
+
+            dense_motion = dense_motion_network(
+                source_image=source,
+                kp_driving=kp_norm,
+                kp_source=kp_source,
+                bg_param=None,
+                dropout_flag=False,
+            )
+
+            out = inpainting(source, dense_motion)
+            out_np = np.transpose(out["prediction"].data.cpu().numpy(), [0, 2, 3, 1])[0]
+            predictions.append(out_np)
+            pbar.update_absolute(frame_idx, num_frames)
+
+    return predictions
+
+
+def relative_kp(kp_source, kp_driving, kp_driving_initial):
+    source_area = ConvexHull(kp_source["fg_kp"][0].data.cpu().numpy()).volume
+    driving_area = ConvexHull(kp_driving_initial["fg_kp"][0].data.cpu().numpy()).volume
+    adapt_movement_scale = np.sqrt(source_area) / np.sqrt(driving_area)
+
+    kp_new = {k: v for k, v in kp_driving.items()}
+
+    kp_value_diff = kp_driving["fg_kp"] - kp_driving_initial["fg_kp"]
+    kp_value_diff *= adapt_movement_scale
+    kp_new["fg_kp"] = kp_value_diff + kp_source["fg_kp"]
+
+    return kp_new
+
+
+def load_checkpoints(config_path, checkpoint_path, device):
+    with open(config_path) as f:
+        config = yaml.full_load(f)
+
+    inpainting = InpaintingNetwork(
+        **config["model_params"]["generator_params"],
+        **config["model_params"]["common_params"],
+    )
+    kp_detector = KPDetector(**config["model_params"]["common_params"])
+    dense_motion_network = DenseMotionNetwork(
+        **config["model_params"]["common_params"],
+        **config["model_params"]["dense_motion_params"],
+    )
+    avd_network = AVDNetwork(
+        num_tps=config["model_params"]["common_params"]["num_tps"],
+        **config["model_params"]["avd_network_params"],
+    )
+    kp_detector.to(device)
+    dense_motion_network.to(device)
+    inpainting.to(device)
+    avd_network.to(device)
+
+    checkpoint = torch.load(checkpoint_path, map_location=device)
+
+    inpainting.load_state_dict(checkpoint["inpainting_network"])
+    kp_detector.load_state_dict(checkpoint["kp_detector"])
+    dense_motion_network.load_state_dict(checkpoint["dense_motion_network"])
+    if "avd_network" in checkpoint:
+        avd_network.load_state_dict(checkpoint["avd_network"])
+
+    inpainting.eval()
+    kp_detector.eval()
+    dense_motion_network.eval()
+    avd_network.eval()
+
+    return inpainting, kp_detector, dense_motion_network, avd_network

spline_module/avd_network.py

@@ -0,0 +1,63 @@
+import torch
+from torch import nn
+
+
+class AVDNetwork(nn.Module):
+    """
+    Animation via Disentanglement network
+    """
+
+    def __init__(self, num_tps, id_bottle_size=64, pose_bottle_size=64):
+        super(AVDNetwork, self).__init__()
+        input_size = 5 * 2 * num_tps
+        self.num_tps = num_tps
+
+        self.id_encoder = nn.Sequential(
+            nn.Linear(input_size, 256),
+            nn.BatchNorm1d(256),
+            nn.ReLU(inplace=True),
+            nn.Linear(256, 512),
+            nn.BatchNorm1d(512),
+            nn.ReLU(inplace=True),
+            nn.Linear(512, 1024),
+            nn.BatchNorm1d(1024),
+            nn.ReLU(inplace=True),
+            nn.Linear(1024, id_bottle_size),
+        )
+
+        self.pose_encoder = nn.Sequential(
+            nn.Linear(input_size, 256),
+            nn.BatchNorm1d(256),
+            nn.ReLU(inplace=True),
+            nn.Linear(256, 512),
+            nn.BatchNorm1d(512),
+            nn.ReLU(inplace=True),
+            nn.Linear(512, 1024),
+            nn.BatchNorm1d(1024),
+            nn.ReLU(inplace=True),
+            nn.Linear(1024, pose_bottle_size),
+        )
+
+        self.decoder = nn.Sequential(
+            nn.Linear(pose_bottle_size + id_bottle_size, 1024),
+            nn.BatchNorm1d(1024),
+            nn.ReLU(),
+            nn.Linear(1024, 512),
+            nn.BatchNorm1d(512),
+            nn.ReLU(),
+            nn.Linear(512, 256),
+            nn.BatchNorm1d(256),
+            nn.ReLU(),
+            nn.Linear(256, input_size),
+        )
+
+    def forward(self, kp_source, kp_random):
+        bs = kp_source["fg_kp"].shape[0]
+
+        pose_emb = self.pose_encoder(kp_random["fg_kp"].view(bs, -1))
+        id_emb = self.id_encoder(kp_source["fg_kp"].view(bs, -1))
+
+        rec = self.decoder(torch.cat([pose_emb, id_emb], dim=1))
+
+        rec = {"fg_kp": rec.view(bs, self.num_tps * 5, -1)}
+        return rec

spline_module/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

spline_module/config/mgif-256.yaml

@@ -0,0 +1,75 @@
+dataset_params:
+  root_dir: ../moving-gif
+  frame_shape: null
+  id_sampling: False
+  augmentation_params:
+    flip_param:
+      horizontal_flip: True
+      time_flip: True
+    crop_param:
+      size: [256, 256]
+    resize_param:
+      ratio: [0.9, 1.1]
+    jitter_param:
+      hue: 0.5
+
+model_params:
+  common_params:
+    num_tps: 10
+    num_channels: 3
+    bg: False
+    multi_mask: True
+  generator_params:
+    block_expansion: 64
+    max_features: 512
+    num_down_blocks: 3
+  dense_motion_params:
+    block_expansion: 64
+    max_features: 1024
+    num_blocks: 5
+    scale_factor: 0.25
+  avd_network_params:
+    id_bottle_size: 128
+    pose_bottle_size: 128
+
+
+train_params:
+  num_epochs: 100
+  num_repeats: 50
+  epoch_milestones: [70, 90]
+  lr_generator: 2.0e-4
+  batch_size: 28
+  scales: [1, 0.5, 0.25, 0.125]
+  dataloader_workers: 12
+  checkpoint_freq: 50
+  dropout_epoch: 35
+  dropout_maxp: 0.5
+  dropout_startp: 0.2
+  dropout_inc_epoch: 10
+  bg_start: 0
+  transform_params:
+    sigma_affine: 0.05
+    sigma_tps: 0.005
+    points_tps: 5
+  loss_weights:
+    perceptual: [10, 10, 10, 10, 10]
+    equivariance_value: 10
+    warp_loss: 10
+    bg: 10
+
+train_avd_params:
+  num_epochs: 100
+  num_repeats: 50
+  batch_size: 256
+  dataloader_workers: 24
+  checkpoint_freq: 10
+  epoch_milestones: [70, 90]
+  lr: 1.0e-3
+  lambda_shift: 1
+  lambda_affine: 1
+  random_scale: 0.25
+
+visualizer_params:
+  kp_size: 5
+  draw_border: True
+  colormap: 'gist_rainbow'

spline_module/config/taichi-256.yaml

@@ -0,0 +1,134 @@
+# Dataset parameters
+# Each dataset should contain 2 folders train and test
+# Each video can be represented as:
+#   - an image of concatenated frames
+#   - '.mp4' or '.gif'
+#   - folder with all frames from a specific video
+# In case of Taichi. Same (youtube) video can be splitted in many parts (chunks). Each part has a following
+# format (id)#other#info.mp4. For example '12335#adsbf.mp4' has an id 12335. In case of TaiChi id stands for youtube
+# video id.
+dataset_params:
+  # Path to data, data can be stored in several formats: .mp4 or .gif videos, stacked .png images or folders with frames.
+  root_dir: ../taichi
+  # Image shape, needed for staked .png format.
+  frame_shape: null
+  # In case of TaiChi single video can be splitted in many chunks, or the maybe several videos for single person.
+  # In this case epoch can be a pass over different videos (if id_sampling=True) or over different chunks (if id_sampling=False)
+  # If the name of the video '12335#adsbf.mp4' the id is assumed to be 12335
+  id_sampling: True
+  # Augmentation parameters see augmentation.py for all posible augmentations
+  augmentation_params:
+    flip_param:
+      horizontal_flip: True
+      time_flip: True
+    jitter_param:
+      brightness: 0.1
+      contrast: 0.1
+      saturation: 0.1
+      hue: 0.1
+
+# Defines model architecture
+model_params:
+  common_params:
+    # Number of TPS transformation
+    num_tps: 10
+    # Number of channels per image
+    num_channels: 3
+    # Whether to estimate affine background transformation
+    bg: True
+    # Whether to estimate the multi-resolution occlusion masks
+    multi_mask: True
+  generator_params:
+    # Number of features mutliplier
+    block_expansion: 64
+    # Maximum allowed number of features
+    max_features: 512
+    # Number of downsampling blocks and Upsampling blocks.
+    num_down_blocks: 3
+  dense_motion_params:
+    # Number of features mutliplier
+    block_expansion: 64
+    # Maximum allowed number of features
+    max_features: 1024
+    # Number of block in Unet. 
+    num_blocks: 5
+    # Optical flow is predicted on smaller images for better performance,
+    # scale_factor=0.25 means that 256x256 image will be resized to 64x64  
+    scale_factor: 0.25
+  avd_network_params:
+    # Bottleneck for identity branch
+    id_bottle_size: 128
+    # Bottleneck for pose branch
+    pose_bottle_size: 128
+
+# Parameters of training
+train_params:
+  # Number of training epochs 
+  num_epochs: 100
+  # For better i/o performance when number of videos is small number of epochs can be multiplied by this number.
+  # Thus effectivlly with num_repeats=100 each epoch is 100 times larger. 
+  num_repeats: 150
+  # Drop learning rate by 10 times after this epochs 
+  epoch_milestones: [70, 90]
+  # Initial learing rate for all modules
+  lr_generator: 2.0e-4
+  batch_size: 28
+  # Scales for perceptual pyramide loss. If scales = [1, 0.5, 0.25, 0.125] and image resolution is 256x256,
+  # than the loss will be computer on resolutions 256x256, 128x128, 64x64, 32x32.
+  scales: [1, 0.5, 0.25, 0.125]
+  # Dataset preprocessing cpu workers
+  dataloader_workers: 12
+  # Save checkpoint this frequently. If checkpoint_freq=50, checkpoint will be saved every 50 epochs.
+  checkpoint_freq: 50
+  # Parameters of dropout
+  # The first dropout_epoch training uses dropout operation 
+  dropout_epoch: 35
+  # The probability P will linearly increase from dropout_startp to dropout_maxp in dropout_inc_epoch epochs
+  dropout_maxp: 0.7
+  dropout_startp: 0.0
+  dropout_inc_epoch: 10
+  # Estimate affine background transformation from the bg_start epoch.
+  bg_start: 0
+  # Parameters of random TPS transformation for equivariance loss
+  transform_params:
+    # Sigma for affine part
+    sigma_affine: 0.05
+    # Sigma for deformation part
+    sigma_tps: 0.005
+    # Number of point in the deformation grid
+    points_tps: 5
+  loss_weights:
+    # Weights for perceptual loss.
+    perceptual: [10, 10, 10, 10, 10]
+    # Weights for value equivariance.
+    equivariance_value: 10
+    # Weights for warp loss.
+    warp_loss: 10
+    # Weights for bg loss.
+    bg: 10
+
+# Parameters of training (animation-via-disentanglement)
+train_avd_params:
+  # Number of training epochs, visualization is produced after each epoch.
+  num_epochs: 100
+  # For better i/o performance when number of videos is small number of epochs can be multiplied by this number.
+  # Thus effectively with num_repeats=100 each epoch is 100 times larger.
+  num_repeats: 150
+  # Batch size.
+  batch_size: 256
+  # Save checkpoint this frequently. If checkpoint_freq=50, checkpoint will be saved every 50 epochs.
+  checkpoint_freq: 10
+  # Dataset preprocessing cpu workers
+  dataloader_workers: 24
+  # Drop learning rate 10 times after this epochs
+  epoch_milestones: [70, 90]
+  # Initial learning rate
+  lr: 1.0e-3
+  # Weights for equivariance loss.
+  lambda_shift: 1
+  random_scale: 0.25
+
+visualizer_params:
+  kp_size: 5
+  draw_border: True
+  colormap: 'gist_rainbow'

spline_module/config/ted-384.yaml

@@ -0,0 +1,73 @@
+dataset_params:
+  root_dir: ../TED384-v2
+  frame_shape: null
+  id_sampling: True
+  augmentation_params:
+    flip_param:
+      horizontal_flip: True
+      time_flip: True
+    jitter_param:
+      brightness: 0.1
+      contrast: 0.1
+      saturation: 0.1
+      hue: 0.1
+
+model_params:
+  common_params:
+    num_tps: 10
+    num_channels: 3
+    bg: True
+    multi_mask: True
+  generator_params:
+    block_expansion: 64
+    max_features: 512
+    num_down_blocks: 3
+  dense_motion_params:
+    block_expansion: 64
+    max_features: 1024
+    num_blocks: 5
+    scale_factor: 0.25
+  avd_network_params:
+    id_bottle_size: 128
+    pose_bottle_size: 128
+
+
+train_params:
+  num_epochs: 100
+  num_repeats: 150
+  epoch_milestones: [70, 90]
+  lr_generator: 2.0e-4
+  batch_size: 12
+  scales: [1, 0.5, 0.25, 0.125]
+  dataloader_workers: 6
+  checkpoint_freq: 50
+  dropout_epoch: 35
+  dropout_maxp: 0.5
+  dropout_startp: 0.0
+  dropout_inc_epoch: 10
+  bg_start: 0
+  transform_params:
+    sigma_affine: 0.05
+    sigma_tps: 0.005
+    points_tps: 5
+  loss_weights:
+    perceptual: [10, 10, 10, 10, 10]
+    equivariance_value: 10
+    warp_loss: 10
+    bg: 10
+
+train_avd_params:
+  num_epochs: 30
+  num_repeats: 500
+  batch_size: 256
+  dataloader_workers: 24
+  checkpoint_freq: 10
+  epoch_milestones: [20, 25]
+  lr: 1.0e-3
+  lambda_shift: 1
+  random_scale: 0.25
+
+visualizer_params:
+  kp_size: 5
+  draw_border: True
+  colormap: 'gist_rainbow'

spline_module/config/vox-256.yaml

@@ -0,0 +1,74 @@
+dataset_params:
+  root_dir: ../vox
+  frame_shape: null
+  id_sampling: True
+  augmentation_params:
+    flip_param:
+      horizontal_flip: True
+      time_flip: True
+    jitter_param:
+      brightness: 0.1
+      contrast: 0.1
+      saturation: 0.1
+      hue: 0.1
+
+
+model_params:
+  common_params:
+    num_tps: 10
+    num_channels: 3
+    bg: True
+    multi_mask: True
+  generator_params:
+    block_expansion: 64
+    max_features: 512
+    num_down_blocks: 3
+  dense_motion_params:
+    block_expansion: 64
+    max_features: 1024
+    num_blocks: 5
+    scale_factor: 0.25
+  avd_network_params:
+    id_bottle_size: 128
+    pose_bottle_size: 128
+
+
+train_params:
+  num_epochs: 100
+  num_repeats: 75
+  epoch_milestones: [70, 90]
+  lr_generator: 2.0e-4
+  batch_size: 28
+  scales: [1, 0.5, 0.25, 0.125]
+  dataloader_workers: 12
+  checkpoint_freq: 50
+  dropout_epoch: 35
+  dropout_maxp: 0.3
+  dropout_startp: 0.1
+  dropout_inc_epoch: 10
+  bg_start: 10
+  transform_params:
+    sigma_affine: 0.05
+    sigma_tps: 0.005
+    points_tps: 5
+  loss_weights:
+    perceptual: [10, 10, 10, 10, 10]
+    equivariance_value: 10
+    warp_loss: 10
+    bg: 10
+
+train_avd_params:
+  num_epochs: 200
+  num_repeats: 300
+  batch_size: 256
+  dataloader_workers: 24
+  checkpoint_freq: 50
+  epoch_milestones: [140, 180]
+  lr: 1.0e-3
+  lambda_shift: 1
+  random_scale: 0.25
+
+visualizer_params:
+  kp_size: 5
+  draw_border: True
+  colormap: 'gist_rainbow'

spline_module/dense_motion.py

@@ -0,0 +1,241 @@
+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 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=True,
+        kp_variance=0.01,
+    ):
+        super(DenseMotionNetwork, 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_channels
+        self.maps = nn.Conv2d(
+            hourglass_output_size[-1], 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)
+                )
+            ]
+            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["fg_kp"], spatial_size=spatial_size, kp_variance=self.kp_variance
+        )
+        gaussian_source = kp2gaussian(
+            kp_source["fg_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["fg_kp"]
+        kp_2 = kp_source["fg_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)
+
+        contribution_maps = self.maps(prediction[-1])
+        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
+
+        # 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])))
+
+        out_dict["occlusion_map"] = occlusion_map  # Multi-resolution Occlusion Masks
+        return out_dict

spline_module/inpainting_network.py

@@ -0,0 +1,164 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from .util import DownBlock2d, ResBlock2d, SameBlock2d, UpBlock2d
+
+
+class InpaintingNetwork(nn.Module):
+    """
+    Inpaint the missing regions and reconstruct the Driving image.
+    """
+
+    def __init__(
+        self,
+        num_channels,
+        block_expansion,
+        max_features,
+        num_down_blocks,
+        multi_mask=True,
+        **kwargs,
+    ):
+        super(InpaintingNetwork, self).__init__()
+
+        self.num_down_blocks = num_down_blocks
+        self.multi_mask = multi_mask
+        self.first = SameBlock2d(
+            num_channels, block_expansion, kernel_size=(7, 7), padding=(3, 3)
+        )
+
+        down_blocks = []
+        up_blocks = []
+        resblock = []
+        for i in range(num_down_blocks):
+            in_features = min(max_features, block_expansion * (2**i))
+            out_features = min(max_features, block_expansion * (2 ** (i + 1)))
+            down_blocks.append(
+                DownBlock2d(
+                    in_features, out_features, kernel_size=(3, 3), padding=(1, 1)
+                )
+            )
+            decoder_in_feature = out_features * 2
+            if i == num_down_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))
+            )
+            resblock.append(
+                ResBlock2d(decoder_in_feature, 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.final = nn.Conv2d(
+            block_expansion, num_channels, kernel_size=(7, 7), padding=(3, 3)
+        )
+        self.num_channels = num_channels
+
+    def deform_input(self, inp, deformation):
+        _, h_old, w_old, _ = deformation.shape
+        _, _, h, w = inp.shape
+        if h_old != h or w_old != w:
+            deformation = deformation.permute(0, 3, 1, 2)
+            deformation = F.interpolate(
+                deformation, size=(h, w), mode="bilinear", align_corners=True
+            )
+            deformation = deformation.permute(0, 2, 3, 1)
+        return F.grid_sample(inp, deformation, align_corners=True)
+
+    def occlude_input(self, inp, occlusion_map):
+        if not self.multi_mask:
+            if (
+                inp.shape[2] != occlusion_map.shape[2]
+                or inp.shape[3] != occlusion_map.shape[3]
+            ):
+                occlusion_map = F.interpolate(
+                    occlusion_map,
+                    size=inp.shape[2:],
+                    mode="bilinear",
+                    align_corners=True,
+                )
+        out = inp * occlusion_map
+        return out
+
+    def forward(self, source_image, dense_motion):
+        out = self.first(source_image)
+        encoder_map = [out]
+        for i in range(len(self.down_blocks)):
+            out = self.down_blocks[i](out)
+            encoder_map.append(out)
+
+        output_dict = {}
+        output_dict["contribution_maps"] = dense_motion["contribution_maps"]
+        output_dict["deformed_source"] = dense_motion["deformed_source"]
+
+        occlusion_map = dense_motion["occlusion_map"]
+        output_dict["occlusion_map"] = occlusion_map
+
+        deformation = dense_motion["deformation"]
+        out_ij = self.deform_input(out.detach(), deformation)
+        out = self.deform_input(out, deformation)
+
+        out_ij = self.occlude_input(out_ij, occlusion_map[0].detach())
+        out = self.occlude_input(out, occlusion_map[0])
+
+        warped_encoder_maps = []
+        warped_encoder_maps.append(out_ij)
+
+        for i in range(self.num_down_blocks):
+            out = self.resblock[2 * i](out)
+            out = self.resblock[2 * i + 1](out)
+            out = self.up_blocks[i](out)
+
+            encode_i = encoder_map[-(i + 2)]
+            encode_ij = self.deform_input(encode_i.detach(), deformation)
+            encode_i = self.deform_input(encode_i, deformation)
+
+            occlusion_ind = 0
+            if self.multi_mask:
+                occlusion_ind = i + 1
+            encode_ij = self.occlude_input(
+                encode_ij, occlusion_map[occlusion_ind].detach()
+            )
+            encode_i = self.occlude_input(encode_i, occlusion_map[occlusion_ind])
+            warped_encoder_maps.append(encode_ij)
+
+            if i == self.num_down_blocks - 1:
+                break
+
+            out = torch.cat([out, encode_i], 1)
+
+        deformed_source = self.deform_input(source_image, deformation)
+        output_dict["deformed"] = deformed_source
+        output_dict["warped_encoder_maps"] = warped_encoder_maps
+
+        occlusion_last = occlusion_map[-1]
+        if not self.multi_mask:
+            occlusion_last = F.interpolate(
+                occlusion_last, size=out.shape[2:], mode="bilinear", align_corners=True
+            )
+
+        out = out * (1 - occlusion_last) + encode_i
+        out = self.final(out)
+        out = torch.sigmoid(out)
+        out = out * (1 - occlusion_last) + deformed_source * occlusion_last
+        output_dict["prediction"] = out
+
+        return output_dict
+
+    def get_encode(self, driver_image, occlusion_map):
+        out = self.first(driver_image)
+        encoder_map = []
+        encoder_map.append(self.occlude_input(out.detach(), occlusion_map[-1].detach()))
+        for i in range(len(self.down_blocks)):
+            out = self.down_blocks[i](out.detach())
+            out_mask = self.occlude_input(out.detach(), occlusion_map[2 - i].detach())
+            encoder_map.append(out_mask.detach())
+
+        return encoder_map

spline_module/keypoint_detector.py

@@ -0,0 +1,29 @@
+import torch
+from torch import nn
+from torchvision import models
+
+
+class KPDetector(nn.Module):
+    """
+    Predict K*5 keypoints.
+    """
+
+    def __init__(self, num_tps, **kwargs):
+        super(KPDetector, 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 = {"fg_kp": fg_kp.view(bs, self.num_tps * 5, -1)}
+
+        return out

spline_module/model.py

@@ -0,0 +1,222 @@
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torchvision import models
+
+from .util import TPS, AntiAliasInterpolation2d
+
+
+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 - 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):
+        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
+
+
+def detach_kp(kp):
+    return {key: value.detach() for key, value in kp.items()}
+
+
+class GeneratorFullModel(torch.nn.Module):
+    """
+    Merge all generator related updates into single model for better multi-gpu usage
+    """
+
+    def __init__(
+        self,
+        kp_extractor,
+        bg_predictor,
+        dense_motion_network,
+        inpainting_network,
+        train_params,
+        *kwargs,
+    ):
+        super(GeneratorFullModel, self).__init__()
+        self.kp_extractor = kp_extractor
+        self.inpainting_network = inpainting_network
+        self.dense_motion_network = dense_motion_network
+
+        self.bg_predictor = None
+        if bg_predictor:
+            self.bg_predictor = bg_predictor
+            self.bg_start = train_params["bg_start"]
+
+        self.train_params = train_params
+        self.scales = train_params["scales"]
+
+        self.pyramid = ImagePyramide(self.scales, inpainting_network.num_channels)
+        if torch.cuda.is_available():
+            self.pyramid = self.pyramid.cuda()
+
+        self.loss_weights = train_params["loss_weights"]
+        self.dropout_epoch = train_params["dropout_epoch"]
+        self.dropout_maxp = train_params["dropout_maxp"]
+        self.dropout_inc_epoch = train_params["dropout_inc_epoch"]
+        self.dropout_startp = train_params["dropout_startp"]
+
+        if sum(self.loss_weights["perceptual"]) != 0:
+            self.vgg = Vgg19()
+            if torch.cuda.is_available():
+                self.vgg = self.vgg.cuda()
+
+    def forward(self, x, epoch):
+        kp_source = self.kp_extractor(x["source"])
+        kp_driving = self.kp_extractor(x["driving"])
+        bg_param = None
+        if self.bg_predictor:
+            if epoch >= self.bg_start:
+                bg_param = self.bg_predictor(x["source"], x["driving"])
+
+        if epoch >= self.dropout_epoch:
+            dropout_flag = False
+            dropout_p = 0
+        else:
+            # 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,
+            )
+
+        dense_motion = self.dense_motion_network(
+            source_image=x["source"],
+            kp_driving=kp_driving,
+            kp_source=kp_source,
+            bg_param=bg_param,
+            dropout_flag=dropout_flag,
+            dropout_p=dropout_p,
+        )
+        generated = self.inpainting_network(x["source"], dense_motion)
+        generated.update({"kp_source": kp_source, "kp_driving": kp_driving})
+
+        loss_values = {}
+
+        pyramide_real = self.pyramid(x["driving"])
+        pyramide_generated = self.pyramid(generated["prediction"])
+
+        # reconstruction loss
+        if sum(self.loss_weights["perceptual"]) != 0:
+            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
+
+        # equivariance loss
+        if self.loss_weights["equivariance_value"] != 0:
+            transform_random = TPS(
+                mode="random",
+                bs=x["driving"].shape[0],
+                **self.train_params["transform_params"],
+            )
+            transform_grid = transform_random.transform_frame(x["driving"])
+            transformed_frame = F.grid_sample(
+                x["driving"],
+                transform_grid,
+                padding_mode="reflection",
+                align_corners=True,
+            )
+            transformed_kp = self.kp_extractor(transformed_frame)
+
+            generated["transformed_frame"] = transformed_frame
+            generated["transformed_kp"] = transformed_kp
+
+            warped = transform_random.warp_coordinates(transformed_kp["fg_kp"])
+            kp_d = kp_driving["fg_kp"]
+            value = torch.abs(kp_d - warped).mean()
+            loss_values["equivariance_value"] = (
+                self.loss_weights["equivariance_value"] * value
+            )
+
+        # warp loss
+        if self.loss_weights["warp_loss"] != 0:
+            occlusion_map = generated["occlusion_map"]
+            encode_map = self.inpainting_network.get_encode(x["driving"], occlusion_map)
+            decode_map = generated["warped_encoder_maps"]
+            value = 0
+            for i in range(len(encode_map)):
+                value += torch.abs(encode_map[i] - decode_map[-i - 1]).mean()
+
+            loss_values["warp_loss"] = self.loss_weights["warp_loss"] * value
+
+        # bg loss
+        if (
+            self.bg_predictor
+            and epoch >= self.bg_start
+            and self.loss_weights["bg"] != 0
+        ):
+            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"] = self.loss_weights["bg"] * value
+
+        return loss_values, generated

spline_module/util.py

@@ -0,0 +1,412 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+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
+
+
+def kp2gaussian(kp, spatial_size, kp_variance):
+    """
+    Transform a keypoint into gaussian like representation
+    """
+
+    coordinate_grid = make_coordinate_grid(spatial_size, kp.type()).to(kp.device)
+    number_of_leading_dimensions = len(kp.shape) - 1
+    shape = (1,) * number_of_leading_dimensions + coordinate_grid.shape
+    coordinate_grid = coordinate_grid.view(*shape)
+    repeats = kp.shape[:number_of_leading_dimensions] + (1, 1, 1)
+    coordinate_grid = coordinate_grid.repeat(*repeats)
+
+    # Preprocess kp shape
+    shape = kp.shape[:number_of_leading_dimensions] + (1, 1, 2)
+    kp = kp.view(*shape)
+
+    mean_sub = coordinate_grid - kp
+
+    out = torch.exp(-0.5 * (mean_sub**2).sum(-1) / kp_variance)
+
+    return out
+
+
+def make_coordinate_grid(spatial_size, type):
+    """
+    Create a meshgrid [-1,1] x [-1,1] of given spatial_size.
+    """
+    h, w = spatial_size
+    x = torch.arange(w).type(type)
+    y = torch.arange(h).type(type)
+
+    x = 2 * (x / (w - 1)) - 1
+    y = 2 * (y / (h - 1)) - 1
+
+    yy = y.view(-1, 1).repeat(1, w)
+    xx = x.view(1, -1).repeat(h, 1)
+
+    meshed = torch.cat([xx.unsqueeze_(2), yy.unsqueeze_(2)], 2)
+
+    return meshed
+
+
+class ResBlock2d(nn.Module):
+    """
+    Res block, preserve spatial resolution.
+    """
+
+    def __init__(self, in_features, kernel_size, padding):
+        super(ResBlock2d, self).__init__()
+        self.conv1 = nn.Conv2d(
+            in_channels=in_features,
+            out_channels=in_features,
+            kernel_size=kernel_size,
+            padding=padding,
+        )
+        self.conv2 = nn.Conv2d(
+            in_channels=in_features,
+            out_channels=in_features,
+            kernel_size=kernel_size,
+            padding=padding,
+        )
+        self.norm1 = nn.InstanceNorm2d(in_features, affine=True)
+        self.norm2 = nn.InstanceNorm2d(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 = nn.InstanceNorm2d(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 = nn.InstanceNorm2d(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 = nn.InstanceNorm2d(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]
+        # print('encoder:' ,outs[-1].shape)
+        for down_block in self.down_blocks:
+            outs.append(down_block(outs[-1]))
+            # print('encoder:' ,outs[-1].shape)
+        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 = []
+        self.out_channels = []
+        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))
+            )
+            self.out_channels.append(in_filters)
+            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_channels.append(block_expansion + in_features)
+        # self.out_filters = block_expansion + in_features
+
+    def forward(self, x, mode=0):
+        out = x.pop()
+        outs = []
+        for up_block in self.up_blocks:
+            out = up_block(out)
+            skip = x.pop()
+            out = torch.cat([out, skip], dim=1)
+            outs.append(out)
+        if mode == 0:
+            return out
+        else:
+            return outs
+
+
+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_channels = self.decoder.out_channels
+        # self.out_filters = self.decoder.out_filters
+
+    def forward(self, x, mode=0):
+        return self.decoder(self.encoder(x), mode)
+
+
+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]

workflows/SPLINE.json

@@ -0,0 +1,427 @@
+{
+  "last_node_id": 6,
+  "last_link_id": 11,
+  "nodes": [
+    {
+      "id": 6,
+      "type": "VHS_VideoCombine",
+      "pos": [
+        1236,
+        163
+      ],
+      "size": [
+        315,
+        615
+      ],
+      "flags": {},
+      "order": 4,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 11
+        },
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "link": 10
+        },
+        {
+          "name": "meta_batch",
+          "type": "VHS_BatchManager",
+          "link": null
+        },
+        {
+          "name": "vae",
+          "type": "VAE",
+          "link": null
+        },
+        {
+          "name": "frame_rate",
+          "type": "FLOAT",
+          "link": 9,
+          "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": "Spline",
+        "format": "video/h264-mp4",
+        "pix_fmt": "yuv420p",
+        "crf": 0,
+        "save_metadata": true,
+        "pingpong": false,
+        "save_output": true,
+        "videopreview": {
+          "hidden": false,
+          "paused": false,
+          "params": {
+            "filename": "Spline_00005-audio.mp4",
+            "subfolder": "",
+            "type": "output",
+            "format": "video/h264-mp4",
+            "frame_rate": 29.970029970029966
+          }
+        }
+      }
+    },
+    {
+      "id": 5,
+      "type": "VHS_VideoInfoLoaded",
+      "pos": [
+        783,
+        373
+      ],
+      "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": [
+            8
+          ],
+          "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": 2,
+      "type": "LoadImage",
+      "pos": [
+        145,
+        158
+      ],
+      "size": [
+        315,
+        314
+      ],
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            4
+          ],
+          "shape": 3
+        },
+        {
+          "name": "MASK",
+          "type": "MASK",
+          "links": null,
+          "shape": 3
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "LoadImage"
+      },
+      "widgets_values": [
+        "Rick1.png",
+        "image"
+      ]
+    },
+    {
+      "id": 3,
+      "type": "VHS_LoadVideo",
+      "pos": [
+        485,
+        232
+      ],
+      "size": [
+        235.1999969482422,
+        491.1999969482422
+      ],
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "meta_batch",
+          "type": "VHS_BatchManager",
+          "link": null
+        },
+        {
+          "name": "vae",
+          "type": "VAE",
+          "link": null
+        }
+      ],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            5
+          ],
+          "shape": 3,
+          "slot_index": 0
+        },
+        {
+          "name": "frame_count",
+          "type": "INT",
+          "links": null,
+          "shape": 3
+        },
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "links": [
+            6
+          ],
+          "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": 4,
+      "type": "Spline_Runner",
+      "pos": [
+        777,
+        161
+      ],
+      "size": [
+        380.4000244140625,
+        146
+      ],
+      "flags": {},
+      "order": 3,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "source_image",
+          "type": "IMAGE",
+          "link": 4
+        },
+        {
+          "name": "driving_video_input",
+          "type": "IMAGE",
+          "link": 5
+        },
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "link": 6
+        },
+        {
+          "name": "frame_rate",
+          "type": "FLOAT",
+          "link": 8,
+          "widget": {
+            "name": "frame_rate"
+          }
+        }
+      ],
+      "outputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "links": [
+            11
+          ],
+          "shape": 3,
+          "slot_index": 0
+        },
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "links": [
+            10
+          ],
+          "shape": 3,
+          "slot_index": 1
+        },
+        {
+          "name": "frame_rate",
+          "type": "FLOAT",
+          "links": [
+            9
+          ],
+          "shape": 3,
+          "slot_index": 2
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "Spline_Runner"
+      },
+      "widgets_values": [
+        30,
+        "relative",
+        false
+      ]
+    }
+  ],
+  "links": [
+    [
+      4,
+      2,
+      0,
+      4,
+      0,
+      "IMAGE"
+    ],
+    [
+      5,
+      3,
+      0,
+      4,
+      1,
+      "IMAGE"
+    ],
+    [
+      6,
+      3,
+      2,
+      4,
+      2,
+      "AUDIO"
+    ],
+    [
+      7,
+      3,
+      3,
+      5,
+      0,
+      "VHS_VIDEOINFO"
+    ],
+    [
+      8,
+      5,
+      0,
+      4,
+      3,
+      "FLOAT"
+    ],
+    [
+      9,
+      4,
+      2,
+      6,
+      4,
+      "FLOAT"
+    ],
+    [
+      10,
+      4,
+      1,
+      6,
+      1,
+      "AUDIO"
+    ],
+    [
+      11,
+      4,
+      0,
+      6,
+      0,
+      "IMAGE"
+    ]
+  ],
+  "groups": [],
+  "config": {},
+  "extra": {
+    "ds": {
+      "scale": 1,
+      "offset": [
+        -23,
+        68
+      ]
+    }
+  },
+  "version": 0.4
+}