partswap seg10 working

README.md

@@ -3,6 +3,10 @@ A ComfyUI-native node to run First Order Motion Model for Image Animation
 
 https://github.com/AliaksandrSiarohin/first-order-model
 
+Now supports Face Swapping using Motion Supervised co-part Segmentation
+
+https://github.com/AliaksandrSiarohin/motion-cosegmentation
+
 ## Workflow:
 
 [FOMM.json](FOMM.json)

__init__.py

@@ -1,11 +1,16 @@
-from .run import FOMM_Runner
+from .run import FOMM_Runner, FOMM_Partswap, FOMM_Seg10Chooser
 
 NODE_CLASS_MAPPINGS = {
     "FOMM_Runner": FOMM_Runner,
+    "FOMM_Partswap": FOMM_Partswap,
+    "FOMM_Seg10Chooser": FOMM_Seg10Chooser,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
     "FOMM_Runner": "FOMM Runner",
+    "FOMM_Partswap": "FOMM Partswap",
+    "FOMM_Seg10Chooser": "FOMM Seg10 Chooser",
 }
 
+
 __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']

config/vox-256-sem-10segments.yaml

@@ -0,0 +1,92 @@
+# Dataset parameters
+dataset_params:
+  # Path to data, data can be stored in several formats: .mp4 or .gif videos, stacked .png images or folders with frames.
+  root_dir: data/vox-png
+  # Image shape, needed for staked .png format.
+  image_shape: [256, 256, 3]
+  # In case of VoxCeleb or TaiChi single video can be splitted in many chunks, or the maybe several videos for single person.
+  # In this case epoch can be a pass over different videos (if id_sampling=True) or over different chunks (if id_sampling=False)
+  # If the name of the video '12335#adsbf.mp4' the id is assumed to be 12335
+  id_sampling: True
+  # Augmentation parameters see augmentation.py for all posible augmentations
+  augmentation_params:
+    flip_param:
+      horizontal_flip: True
+      time_flip: True
+    jitter_param:
+      brightness: 0.1
+      contrast: 0.1
+      saturation: 0.1
+      hue: 0.1
+
+# Defines model architecture
+model_params:
+  common_params:
+    # Number of segments
+    num_segments: 10
+    # Number of channels per image
+    num_channels: 3
+    # Use only shift and no-affine part
+    estimate_affine_part: True
+  segmentation_module_params:
+     # Softmax temperature for shift heatmaps
+     temperature: 0.1
+     # Number of features mutliplier
+     block_expansion: 32
+     # Maximum allowed number of features
+     max_features: 1024
+     # Number of block in Unet. Can be increased or decreased depending or resolution.
+     num_blocks: 5
+     # Segmentations is predicted on smaller images for better performance,
+     # scale_factor=0.25 means that 256x256 image will be resized to 64x64
+     scale_factor: 0.25
+  reconstruction_module_params:
+    # Number of features mutliplier
+    block_expansion: 64
+    # Maximum allowed number of features
+    max_features: 512
+    # Number of downsampling blocks in Jonson architecture.
+    # Can be increased or decreased depending or resolution.
+    num_down_blocks: 2
+    # Number of ResBlocks  in Jonson architecture.
+    num_bottleneck_blocks: 6
+    # Use visibility map or not
+    estimate_visibility: True
+
+# Parameters of training
+train_params:
+  # Number of training epochs
+  num_epochs: 5
+  # For better i/o performance when number of videos is small number of epochs can be multiplied by this number.
+  # Thus effectively with num_repeats=100 each epoch is 100 times larger.
+  num_repeats: 50
+  # Learning rates
+  lr_segmentation: 2.0e-4
+  lr_reconstruction: 2.0e-4
+  batch_size: 20
+  # Scales for perceptual pyramide loss. If scales = [1, 0.5, 0.25, 0.125] and image resolution is 256x256,
+  # than the loss will be computer on resolutions 256x256, 128x128, 64x64, 32x32.
+  scales: [1, 0.5, 0.25, 0.125]
+  # Save checkpoint this frequently. If checkpoint_freq=50, checkpoint will be saved every 50 epochs.
+  checkpoint_freq: 50
+  # Parameters of transform for equivariance loss
+  transform_params:
+    # Sigma for affine part
+    sigma_affine: 0.05
+    # Sigma for deformation part
+    sigma_tps: 0.005
+    # Number of point in the deformation grid
+    points_tps: 5
+  loss_weights:
+    equivariance: 10
+    perceptual: [10, 10, 10, 10, 10]
+
+
+# Visualization parameters
+visualizer_params:
+  # Draw keypoints (shifts in affine transformations) of this size, increase or decrease depending on resolution
+  kp_size: 5
+  # Draw white border around images
+  draw_border: True
+  # Color map for keypoints
+  colormap: 'gist_rainbow'

config/vox-256-sem-15segments.yaml

@@ -0,0 +1,92 @@
+# Dataset parameters
+dataset_params:
+  # Path to data, data can be stored in several formats: .mp4 or .gif videos, stacked .png images or folders with frames.
+  root_dir: data/vox-png
+  # Image shape, needed for staked .png format.
+  image_shape: [256, 256, 3]
+  # In case of VoxCeleb or TaiChi single video can be splitted in many chunks, or the maybe several videos for single person.
+  # In this case epoch can be a pass over different videos (if id_sampling=True) or over different chunks (if id_sampling=False)
+  # If the name of the video '12335#adsbf.mp4' the id is assumed to be 12335
+  id_sampling: True
+  # Augmentation parameters see augmentation.py for all posible augmentations
+  augmentation_params:
+    flip_param:
+      horizontal_flip: True
+      time_flip: True
+    jitter_param:
+      brightness: 0.1
+      contrast: 0.1
+      saturation: 0.1
+      hue: 0.1
+
+# Defines model architecture
+model_params:
+  common_params:
+    # Number of segments
+    num_segments: 15
+    # Number of channels per image
+    num_channels: 3
+    # Use only shift and no-affine part
+    estimate_affine_part: True
+  segmentation_module_params:
+     # Softmax temperature for shift heatmaps
+     temperature: 0.1
+     # Number of features mutliplier
+     block_expansion: 32
+     # Maximum allowed number of features
+     max_features: 1024
+     # Number of block in Unet. Can be increased or decreased depending or resolution.
+     num_blocks: 5
+     # Segmentations is predicted on smaller images for better performance,
+     # scale_factor=0.25 means that 256x256 image will be resized to 64x64
+     scale_factor: 0.25
+  reconstruction_module_params:
+    # Number of features mutliplier
+    block_expansion: 64
+    # Maximum allowed number of features
+    max_features: 512
+    # Number of downsampling blocks in Jonson architecture.
+    # Can be increased or decreased depending or resolution.
+    num_down_blocks: 2
+    # Number of ResBlocks  in Jonson architecture.
+    num_bottleneck_blocks: 6
+    # Use visibility map or not
+    estimate_visibility: True
+
+# Parameters of training
+train_params:
+  # Number of training epochs
+  num_epochs: 5
+  # For better i/o performance when number of videos is small number of epochs can be multiplied by this number.
+  # Thus effectively with num_repeats=100 each epoch is 100 times larger.
+  num_repeats: 50
+  # Learning rates
+  lr_segmentation: 2.0e-4
+  lr_reconstruction: 2.0e-4
+  batch_size: 20
+  # Scales for perceptual pyramide loss. If scales = [1, 0.5, 0.25, 0.125] and image resolution is 256x256,
+  # than the loss will be computer on resolutions 256x256, 128x128, 64x64, 32x32.
+  scales: [1, 0.5, 0.25, 0.125]
+  # Save checkpoint this frequently. If checkpoint_freq=50, checkpoint will be saved every 50 epochs.
+  checkpoint_freq: 50
+  # Parameters of transform for equivariance loss
+  transform_params:
+    # Sigma for affine part
+    sigma_affine: 0.05
+    # Sigma for deformation part
+    sigma_tps: 0.005
+    # Number of point in the deformation grid
+    points_tps: 5
+  loss_weights:
+    equivariance: 10
+    perceptual: [10, 10, 10, 10, 10]
+
+
+# Visualization parameters
+visualizer_params:
+  # Draw keypoints (shifts in affine transformations) of this size, increase or decrease depending on resolution
+  kp_size: 5
+  # Draw white border around images
+  draw_border: True
+  # Color map for keypoints
+  colormap: 'gist_rainbow'

config/vox-256-sem-5segments.yaml

@@ -0,0 +1,92 @@
+# Dataset parameters
+dataset_params:
+  # Path to data, data can be stored in several formats: .mp4 or .gif videos, stacked .png images or folders with frames.
+  root_dir: data/vox-png
+  # Image shape, needed for staked .png format.
+  image_shape: [256, 256, 3]
+  # In case of VoxCeleb or TaiChi single video can be splitted in many chunks, or the maybe several videos for single person.
+  # In this case epoch can be a pass over different videos (if id_sampling=True) or over different chunks (if id_sampling=False)
+  # If the name of the video '12335#adsbf.mp4' the id is assumed to be 12335
+  id_sampling: True
+  # Augmentation parameters see augmentation.py for all posible augmentations
+  augmentation_params:
+    flip_param:
+      horizontal_flip: True
+      time_flip: True
+    jitter_param:
+      brightness: 0.1
+      contrast: 0.1
+      saturation: 0.1
+      hue: 0.1
+
+# Defines model architecture
+model_params:
+  common_params:
+    # Number of segments
+    num_segments: 5
+    # Number of channels per image
+    num_channels: 3
+    # Use only shift and no-affine part
+    estimate_affine_part: True
+  segmentation_module_params:
+     # Softmax temperature for shift heatmaps
+     temperature: 0.1
+     # Number of features mutliplier
+     block_expansion: 32
+     # Maximum allowed number of features
+     max_features: 1024
+     # Number of block in Unet. Can be increased or decreased depending or resolution.
+     num_blocks: 5
+     # Segmentations is predicted on smaller images for better performance,
+     # scale_factor=0.25 means that 256x256 image will be resized to 64x64
+     scale_factor: 0.25
+  reconstruction_module_params:
+    # Number of features mutliplier
+    block_expansion: 64
+    # Maximum allowed number of features
+    max_features: 512
+    # Number of downsampling blocks in Jonson architecture.
+    # Can be increased or decreased depending or resolution.
+    num_down_blocks: 2
+    # Number of ResBlocks  in Jonson architecture.
+    num_bottleneck_blocks: 6
+    # Use visibility map or not
+    estimate_visibility: True
+
+# Parameters of training
+train_params:
+  # Number of training epochs
+  num_epochs: 5
+  # For better i/o performance when number of videos is small number of epochs can be multiplied by this number.
+  # Thus effectively with num_repeats=100 each epoch is 100 times larger.
+  num_repeats: 50
+  # Learning rates
+  lr_segmentation: 2.0e-4
+  lr_reconstruction: 2.0e-4
+  batch_size: 20
+  # Scales for perceptual pyramide loss. If scales = [1, 0.5, 0.25, 0.125] and image resolution is 256x256,
+  # than the loss will be computer on resolutions 256x256, 128x128, 64x64, 32x32.
+  scales: [1, 0.5, 0.25, 0.125]
+  # Save checkpoint this frequently. If checkpoint_freq=50, checkpoint will be saved every 50 epochs.
+  checkpoint_freq: 50
+  # Parameters of transform for equivariance loss
+  transform_params:
+    # Sigma for affine part
+    sigma_affine: 0.05
+    # Sigma for deformation part
+    sigma_tps: 0.005
+    # Number of point in the deformation grid
+    points_tps: 5
+  loss_weights:
+    equivariance: 10
+    perceptual: [10, 10, 10, 10, 10]
+
+
+# Visualization parameters
+visualizer_params:
+  # Draw keypoints (shifts in affine transformations) of this size, increase or decrease depending on resolution
+  kp_size: 5
+  # Draw white border around images
+  draw_border: True
+  # Color map for keypoints
+  colormap: 'gist_rainbow'

constants.py

@@ -20,4 +20,12 @@ checkpoint_folder_link = "https://github.com/graphemecluster/first-order-model-d
 checkpoint_suffix = ".pth.tar"
 
 config_folder = "config"
-checkpoint_folder = "checkpoints"
+checkpoint_folder = "checkpoints"
+
+partswap_model_config_dict = {
+    "vox-5segments" : "vox-256-sem-5segments.yaml",
+    "vox-10segments": "vox-256-sem-10segments.yaml",
+    "vox-15segments": "vox-256-sem-15segments.yaml",
+}
+
+default_partswap_model_name = "vox-10segments"

modules/reconstruction_module.py

@@ -0,0 +1,109 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+from .util import ResBlock2d, SameBlock2d, UpBlock2d, DownBlock2d
+from .util import make_coordinate_grid
+
+
+class ReconstructionModule(nn.Module):
+    """
+    Reconstruct target from source and segmentation of target and part transformations
+    """
+
+    def __init__(self, num_channels, block_expansion, max_features, num_down_blocks,
+                 num_bottleneck_blocks, num_segments, estimate_visibility=False, **kwargs):
+        super(ReconstructionModule, self).__init__()
+
+        self.first = SameBlock2d(num_channels, block_expansion, kernel_size=(7, 7), padding=(3, 3))
+
+        down_blocks = []
+        for i in range(num_down_blocks):
+            in_features = min(max_features, block_expansion * (2 ** i))
+            out_features = min(max_features, block_expansion * (2 ** (i + 1)))
+            down_blocks.append(DownBlock2d(in_features, out_features, kernel_size=(3, 3), padding=(1, 1)))
+        self.down_blocks = nn.ModuleList(down_blocks)
+
+        up_blocks = []
+        for i in range(num_down_blocks):
+            in_features = min(max_features, block_expansion * (2 ** (num_down_blocks - i)))
+            out_features = min(max_features, block_expansion * (2 ** (num_down_blocks - i - 1)))
+            up_blocks.append(UpBlock2d(in_features, out_features, kernel_size=(3, 3), padding=(1, 1)))
+        self.up_blocks = nn.ModuleList(up_blocks)
+
+        self.bottleneck = torch.nn.Sequential()
+        in_features = min(max_features, block_expansion * (2 ** num_down_blocks))
+        for i in range(num_bottleneck_blocks):
+            self.bottleneck.add_module('r' + str(i), ResBlock2d(in_features, kernel_size=(3, 3), padding=(1, 1)))
+
+        self.final = nn.Conv2d(block_expansion, num_channels, kernel_size=(7, 7), padding=(3, 3))
+        self.estimate_visibility = estimate_visibility
+        self.num_channels = num_channels
+        self.num_segments = num_segments
+
+    def segment_motion(self, seg_target, seg_source):
+        bs, _, h, w = seg_target['segmentation'].shape
+        identity_grid = make_coordinate_grid((h, w), type=seg_source['shift'].type())
+        identity_grid = identity_grid.view(1, 1, h, w, 2)
+        coordinate_grid = identity_grid - seg_target['shift'].view(bs, self.num_segments, 1, 1, 2)
+        if 'affine' in seg_target:
+            affine = torch.matmul(seg_source['affine'], torch.inverse(seg_target['affine']))
+            affine = affine.unsqueeze(-3).unsqueeze(-3)
+            affine = affine.repeat(1, 1, h, w, 1, 1)
+            coordinate_grid = torch.matmul(affine, coordinate_grid.unsqueeze(-1))
+            coordinate_grid = coordinate_grid.squeeze(-1)
+
+        target_to_source = coordinate_grid + seg_source['shift'].view(bs, self.num_segments, 1, 1, 2)
+
+        # adding background feature
+        identity_grid = identity_grid.repeat(bs, 1, 1, 1, 1)
+        return torch.cat([identity_grid, target_to_source], dim=1)
+
+    def deform_input(self, inp, deformation):
+        _, h_old, w_old, _ = deformation.shape
+        _, _, h, w = inp.shape
+        if h_old != h or w_old != w:
+            deformation = deformation.permute(0, 3, 1, 2)
+            deformation = F.interpolate(deformation, size=(h, w), mode='bilinear')
+            deformation = deformation.permute(0, 2, 3, 1)
+        return F.grid_sample(inp, deformation)
+
+    def forward(self, source_image, seg_target, seg_source):
+        # Encoding (downsampling) part
+        out = self.first(source_image)
+        for i in range(len(self.down_blocks)):
+            out = self.down_blocks[i](out)
+
+        output_dict = {}
+
+        # Computing segment motion
+        segment_motions = self.segment_motion(seg_target, seg_source)
+        segment_motions = segment_motions.permute(0, 1, 4, 2, 3)
+        mask = seg_target['segmentation'].unsqueeze(2)
+        deformation = (segment_motions * mask).sum(dim=1)
+        deformation = deformation.permute(0, 2, 3, 1)
+
+        # Deform according to motion
+        out = self.deform_input(out, deformation)
+        output_dict["deformed"] = self.deform_input(source_image, deformation)
+
+        if self.estimate_visibility:
+            visibility = seg_source['segmentation'][:, 1:].sum(dim=1, keepdim=True) * \
+                         (1 - seg_target['segmentation'][:, 1:].sum(dim=1, keepdim=True).detach())
+            visibility = 1 - visibility
+
+            if out.shape[2] != visibility.shape[2] or out.shape[3] != visibility.shape[3]:
+                visibility = F.interpolate(visibility, size=out.shape[2:], mode='bilinear')
+            out = out * visibility
+            output_dict['visibility_maps'] = visibility
+
+        out = self.bottleneck(out)
+
+        for i in range(len(self.up_blocks)):
+            out = self.up_blocks[i](out)
+
+        out = self.final(out)
+        out = F.sigmoid(out)
+
+        output_dict["prediction"] = out
+
+        return output_dict

modules/segmentation_module.py

@@ -0,0 +1,78 @@
+from torch import nn
+import torch
+import torch.nn.functional as F
+from .util import Hourglass, make_coordinate_grid, AntiAliasInterpolation2d
+
+
+class SegmentationModule(nn.Module):
+    """
+    Computing a segmentation map and affine transformations.
+    """
+
+    def __init__(self, block_expansion, num_segments, num_channels, max_features,
+                 num_blocks, temperature, estimate_affine_part=False, scale_factor=1):
+        super(SegmentationModule, self).__init__()
+
+        self.predictor = Hourglass(block_expansion, in_features=num_channels,
+                                   max_features=max_features, num_blocks=num_blocks)
+        self.num_segments = num_segments
+        self.shift = nn.Conv2d(in_channels=self.predictor.out_filters, out_channels=num_segments, kernel_size=(7, 7),
+                               padding=(3, 3))
+
+        if estimate_affine_part:
+            self.affine = nn.Conv2d(in_channels=self.predictor.out_filters,
+                                    out_channels=4 * num_segments, kernel_size=(7, 7), padding=(3, 3))
+            self.affine.weight.data.zero_()
+            self.affine.bias.data.copy_(torch.tensor([1, 0, 0, 1] * num_segments, dtype=torch.float))
+        else:
+            self.affine = None
+
+        self.segmentation = nn.Conv2d(in_channels=self.predictor.out_filters,
+                                      out_channels=(1 + num_segments), kernel_size=(7, 7), padding=(3, 3))
+
+        self.temperature = temperature
+        self.scale_factor = scale_factor
+        if self.scale_factor != 1:
+            self.down = AntiAliasInterpolation2d(num_channels, self.scale_factor)
+
+    def gaussian2kp(self, heatmap):
+        """
+        Extract the mean and the variance from a heatmap
+        """
+        shape = heatmap.shape
+        heatmap = heatmap.unsqueeze(-1)
+        grid = make_coordinate_grid(shape[2:], heatmap.type()).unsqueeze_(0).unsqueeze_(0)
+        result = (heatmap * grid).sum(dim=(2, 3))
+
+        return result
+
+    def forward(self, x):
+        if self.scale_factor != 1:
+            x = self.down(x)
+
+        out = {}
+        feature_map = self.predictor(x)
+        out['segmentation'] = F.softmax(self.segmentation(feature_map), dim=1)
+
+        prediction = self.shift(feature_map)
+
+        final_shape = prediction.shape
+        heatmap = prediction.view(final_shape[0], final_shape[1], -1)
+        heatmap = F.softmax(heatmap / self.temperature, dim=2)
+        heatmap = heatmap.view(*final_shape)
+
+        out['shift'] = self.gaussian2kp(heatmap)
+
+        if self.affine is not None:
+            affine_map = self.affine(feature_map)
+            affine_map = affine_map.reshape(final_shape[0], self.num_segments, 4, final_shape[2],
+                                                final_shape[3])
+            heatmap = heatmap.unsqueeze(2)
+
+            affine = heatmap * affine_map
+            affine = affine.view(final_shape[0], final_shape[1], 4, -1)
+            affine = affine.sum(dim=-1)
+            affine = affine.view(affine.shape[0], affine.shape[1], 2, 2)
+            out['affine'] = affine
+
+        return out

part_swap.py

@@ -0,0 +1,321 @@
+import torch
+import torch.nn.functional as F
+import yaml
+from tqdm import tqdm
+import numpy as np
+
+from .constants import (
+    checkpoint_folder,
+    checkpoint_suffix,
+    config_folder,
+    partswap_model_config_dict,
+)
+from .modules.dense_motion import DenseMotionNetwork
+from .modules.reconstruction_module import ReconstructionModule
+from .modules.segmentation_module import SegmentationModule
+from .modules.util import AntiAliasInterpolation2d
+
+from .sync_batchnorm.replicate import DataParallelWithCallback
+
+
+
+
+
+def partial_state_dict_load(module, state_dict):
+    own_state = module.state_dict()
+    for name, param in state_dict.items():
+        if name not in own_state:
+            continue
+
+        if isinstance(param, torch.nn.Parameter):
+            # backwards compatibility for serialized parameters
+            param = param.data
+        own_state[name].copy_(param)
+
+
+def load_reconstruction_module(module, checkpoint):
+    if "generator" in checkpoint:
+        partial_state_dict_load(module, checkpoint["generator"])
+    else:
+        module.load_state_dict(checkpoint["reconstruction_module"])
+
+
+def load_segmentation_module(module, checkpoint):
+    if "kp_detector" in checkpoint:
+        partial_state_dict_load(module, checkpoint["kp_detector"])
+        module.state_dict()["affine.weight"].copy_(
+            checkpoint["kp_detector"]["jacobian.weight"]
+        )
+        module.state_dict()["affine.bias"].copy_(
+            checkpoint["kp_detector"]["jacobian.bias"]
+        )
+        module.state_dict()["shift.weight"].copy_(
+            checkpoint["kp_detector"]["kp.weight"]
+        )
+        module.state_dict()["shift.bias"].copy_(checkpoint["kp_detector"]["kp.bias"])
+        if "semantic_seg.weight" in checkpoint["kp_detector"]:
+            module.state_dict()["segmentation.weight"].copy_(
+                checkpoint["kp_detector"]["semantic_seg.weight"]
+            )
+            module.state_dict()["segmentation.bias"].copy_(
+                checkpoint["kp_detector"]["semantic_seg.bias"]
+            )
+        else:
+            print("Segmentation part initialized at random.")
+    else:
+        module.load_state_dict(checkpoint["segmentation_module"])
+
+
+def get_partswap_model_file_name(model_name: str):
+    return f"{checkpoint_folder}/{model_name}{checkpoint_suffix}"
+
+
+def load_partswap_checkpoint(
+    config_path: str,
+    checkpoint_path: str,
+    blend_scale: float,
+    use_fomm: bool,
+    cpu=False,
+):
+    with open(config_path) as f:
+        config = yaml.full_load(f)
+
+    reconstruction_module = PartSwapGenerator(
+        blend_scale=blend_scale,
+        first_order_motion_model=use_fomm,
+        **config["model_params"]["reconstruction_module_params"],
+        **config["model_params"]["common_params"],
+    )
+
+    segmentation_module = SegmentationModule(
+        **config["model_params"]["segmentation_module_params"],
+        **config["model_params"]["common_params"],
+    )
+
+    if cpu:
+        checkpoint = torch.load(checkpoint_path, map_location=torch.device("cpu"))
+    else:
+        segmentation_module.cuda()
+        checkpoint = torch.load(checkpoint_path)
+
+    load_reconstruction_module(reconstruction_module, checkpoint)
+    load_segmentation_module(segmentation_module, checkpoint)
+
+    if not cpu:
+        reconstruction_module = DataParallelWithCallback(reconstruction_module)
+        segmentation_module = DataParallelWithCallback(segmentation_module)
+
+    reconstruction_module.eval()
+    segmentation_module.eval()
+
+    return reconstruction_module, segmentation_module
+
+
+class PartSwapGenerator(ReconstructionModule):
+    def __init__(self, blend_scale=1, first_order_motion_model=False, **kwargs):
+        super(PartSwapGenerator, self).__init__(**kwargs)
+        if blend_scale == 1:
+            self.blend_downsample = lambda x: x
+        else:
+            self.blend_downsample = AntiAliasInterpolation2d(1, blend_scale)
+
+        if first_order_motion_model:
+            self.dense_motion_network = DenseMotionNetwork()
+        else:
+            self.dense_motion_network = None
+
+    def forward(
+        self,
+        source_image,
+        target_image,
+        seg_target,
+        seg_source,
+        blend_mask,
+        use_source_segmentation=False,
+    ):
+        # Encoding of source image
+        enc_source = self.first(source_image)
+        for i in range(len(self.down_blocks)):
+            enc_source = self.down_blocks[i](enc_source)
+
+        # Encoding of target image
+        enc_target = self.first(target_image)
+        for i in range(len(self.down_blocks)):
+            enc_target = self.down_blocks[i](enc_target)
+
+        output_dict = {}
+        # Compute flow field for source image
+        if self.dense_motion_network is None:
+            segment_motions = self.segment_motion(seg_target, seg_source)
+            segment_motions = segment_motions.permute(0, 1, 4, 2, 3)
+            mask = seg_target["segmentation"].unsqueeze(2)
+            deformation = (segment_motions * mask).sum(dim=1)
+            deformation = deformation.permute(0, 2, 3, 1)
+        else:
+            motion = self.dense_motion_network(
+                source_image=source_image, seg_target=seg_target, seg_source=seg_source
+            )
+            deformation = motion["deformation"]
+
+        # Deform source encoding according to the motion
+        enc_source = self.deform_input(enc_source, deformation)
+
+        if self.estimate_visibility:
+            if self.dense_motion_network is None:
+                visibility = seg_source["segmentation"][:, 1:].sum(
+                    dim=1, keepdim=True
+                ) * (
+                    1
+                    - seg_target["segmentation"][:, 1:]
+                    .sum(dim=1, keepdim=True)
+                    .detach()
+                )
+                visibility = 1 - visibility
+            else:
+                visibility = motion["visibility"]
+
+            if (
+                enc_source.shape[2] != visibility.shape[2]
+                or enc_source.shape[3] != visibility.shape[3]
+            ):
+                visibility = F.interpolate(
+                    visibility, size=enc_source.shape[2:], mode="bilinear"
+                )
+            enc_source = enc_source * visibility
+
+        blend_mask = self.blend_downsample(blend_mask)
+        # If source segmentation is provided use it should be deformed before blending
+        if use_source_segmentation:
+            blend_mask = self.deform_input(blend_mask, deformation)
+
+        out = enc_target * (1 - blend_mask) + enc_source * blend_mask
+
+        out = self.bottleneck(out)
+
+        for i in range(len(self.up_blocks)):
+            out = self.up_blocks[i](out)
+
+        out = self.final(out)
+        out = F.sigmoid(out)
+
+        output_dict["prediction"] = out
+
+        return output_dict
+    
+def partswap_inference(
+    swap_indices: list[int],
+    source_image,
+    target_video,
+    reconstruction_module: PartSwapGenerator,
+    segmentation_module: SegmentationModule,
+    use_source_seg: bool,
+    hard_edges=False,
+    cpu=False,
+) -> list:
+    with torch.no_grad():
+        predictions = []
+        seg_targets = []
+        if not cpu:
+            source_image = source_image.cuda()
+            target_video = target_video.cuda()
+            print("Using GPU")
+
+        seg_source = segmentation_module(source_image)
+
+        num_frames = target_video.shape[2]
+        for frame_idx in tqdm(range(num_frames)):
+            target_frame = target_video[:, :, frame_idx]
+            seg_target = segmentation_module(target_frame)
+            seg_targets.append(seg_target)
+            if use_source_seg:
+                blend_mask = seg_source["segmentation"]
+            else:
+                blend_mask = seg_target["segmentation"]
+
+            blend_mask = blend_mask[:, swap_indices].sum(dim=1, keepdim=True)
+            if hard_edges:
+                blend_mask = (blend_mask > 0.5).type(blend_mask.type())
+
+            out = reconstruction_module(
+                source_image,
+                target_frame,
+                seg_source=seg_source,
+                seg_target=seg_target,
+                blend_mask=blend_mask,
+                use_source_segmentation=use_source_seg,
+            )
+            
+            predictions.append(np.transpose(out['prediction'].data.cpu().numpy(), [0, 2, 3, 1])[0])
+
+        return seg_source, seg_targets, predictions
+    
+import matplotlib.pyplot as plt
+import matplotlib.patches as mpatches
+from matplotlib.figure import Figure
+from matplotlib.backends.backend_agg import FigureCanvasAgg
+from PIL import Image 
+
+def visualize_frame(image, seg_source, hard_edges=True, colormap='gist_rainbow', alpha=0.6):
+    # Ensure image is the right shape and on CPU
+    image = image.squeeze(0).cpu()  # Remove batch dimension
+    
+    # Extract segmentation from seg_source dictionary
+    mask = seg_source["segmentation"]
+    
+    # Interpolate mask to match image size if necessary
+    if mask.shape[2:] != image.shape[1:]:
+        mask = F.interpolate(mask, size=image.shape[1:], mode='bilinear')
+    
+    if hard_edges:
+        mask = (torch.max(mask, dim=1, keepdim=True)[0] == mask).float()
+    
+    colormap = plt.get_cmap(colormap)
+    num_segments = mask.shape[1]
+    mask = mask.squeeze(0).permute(1, 2, 0).cpu().numpy()
+    color_mask = np.zeros((256, 256, 3), dtype=np.float32)
+    patches = []
+    for i in range(num_segments):
+        if i != 0:
+            color = np.array(colormap((i - 1) / (num_segments - 1)))[:3]
+        else:
+            color = np.array((0, 0, 0))
+        patches.append(mpatches.Patch(color=color, label=str(i)))
+        color_mask += mask[..., i:(i+1)] * color.reshape(1, 1, 3)
+    
+    # Convert image to numpy and normalize to [0, 1] range if necessary
+    image_np = image.permute(1, 2, 0).numpy()
+    if image_np.max() > 1:
+        image_np = image_np / 255.0
+    
+    # Overlay the color mask on the original image
+    overlaid_image = image_np * (1 - alpha) + color_mask * alpha
+    
+    # Ensure all values are in [0, 1] range
+    overlaid_image = np.clip(overlaid_image, 0, 1)
+    
+    # Create a figure and axis
+    fig = Figure(figsize=(10, 10))
+    canvas = FigureCanvasAgg(fig)
+    ax = fig.add_subplot(111)
+    
+    # Display the image
+    ax.imshow(overlaid_image)
+    ax.axis('off')
+    
+    # Add legend
+    ax.legend(handles=patches, loc='upper right', bbox_to_anchor=(1, 1), fontsize='small')
+    
+    # Render the figure to numpy array
+    fig.tight_layout(pad=0)
+    canvas.draw()
+    result_image = np.frombuffer(canvas.tostring_rgb(), dtype='uint8')
+    result_image = result_image.reshape(canvas.get_width_height()[::-1] + (3,))
+    
+    # # Resize to 256x256 if necessary
+    # if result_image.shape[:2] != (256, 256):
+    #     result_image = np.array(Image.fromarray(result_image).resize((256, 256)))
+    
+    # Convert to torch tensor and add batch dimension to get [1, 256, 256, 3]
+    result_tensor = torch.from_numpy(result_image).unsqueeze(0).float() / 255.0
+    
+    return result_tensor

run.py

@@ -9,7 +9,7 @@ from pathlib import Path
 
 import torch
 
-from .constants import default_model_name, supported_model_names
+from .constants import default_model_name, supported_model_names, config_folder
 from .utils import (
     get_config_path,
     get_model_file_name,
@@ -119,10 +119,235 @@ class FOMM_Runner:
         )
 
 
+from .constants import partswap_model_config_dict, default_partswap_model_name
+from .part_swap import (
+    get_partswap_model_file_name,
+    load_partswap_checkpoint,
+    partswap_inference,
+    visualize_frame,
+)
+
+
+def serialize_integers(int_list):
+    return "_".join(map(str, int_list))
+
+
+def deserialize_integers(int_string):
+    return list(map(int, int_string.split('_')))
+
+
+class FOMM_Seg10Chooser:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "m_0_background": (
+                    "BOOLEAN",
+                    {"default": False},
+                ),
+                "m_1_lf_ear": (
+                    "BOOLEAN",
+                    {"default": True},
+                ),
+                "m_2_mouth": (
+                    "BOOLEAN",
+                    {"default": True},
+                ),
+                "m_3_lf_fr_head": (
+                    "BOOLEAN",
+                    {"default": True},
+                ),
+                "m_4_rt_ear": (
+                    "BOOLEAN",
+                    {"default": True},
+                ),
+                "m_5_frnt_face": (
+                    "BOOLEAN",
+                    {"default": True},
+                ),
+                "m_6_lt_cheek": (
+                    "BOOLEAN",
+                    {"default": True},
+                ),
+                "m_7_eyes": (
+                    "BOOLEAN",
+                    {"default": True},
+                ),
+                "m_8_rt_fr_head": (
+                    "BOOLEAN",
+                    {"default": True},
+                ),
+                "m_9_shoulder": (
+                    "BOOLEAN",
+                    {"default": True},
+                ),
+                "m_10_hair_top": (
+                    "BOOLEAN",
+                    {"default": True},
+                ),
+            },
+        }
+
+    RETURN_TYPES = ("STRING",)
+
+    RETURN_NAMES = ("chosen_seg_indices",)
+
+    FUNCTION = "todo"
+    CATEGORY = "FirstOrderMM"
+
+    def todo(
+        self,
+        m_0_background: bool,
+        m_1_lf_ear: bool,
+        m_2_mouth: bool,
+        m_3_lf_fr_head: bool,
+        m_4_rt_ear: bool,
+        m_5_frnt_face: bool,
+        m_6_lt_cheek: bool,
+        m_7_eyes: bool,
+        m_8_rt_fr_head: bool,
+        m_9_shoulder: bool,
+        m_10_hair_top: bool,
+    ):
+        seg_list = []
+        segments = [
+            m_0_background,
+            m_1_lf_ear,
+            m_2_mouth,
+            m_3_lf_fr_head,
+            m_4_rt_ear,
+            m_5_frnt_face,
+            m_6_lt_cheek,
+            m_7_eyes,
+            m_8_rt_fr_head,
+            m_9_shoulder,
+            m_10_hair_top,
+        ]
+
+        for i, seg in enumerate(segments):
+            print(f"Segment {i}: {seg}, Type: {type(seg)}, Truthy: {bool(seg)}")
+            if seg:
+                seg_list.append(i)
+        seg_list = serialize_integers(seg_list)
+        print("Final seg_list:", seg_list)
+        return (seg_list,)
+
+
+class FOMM_Partswap:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "source_image": ("IMAGE",),
+                "driving_video_input": ("IMAGE",),
+                "model_name": (
+                    list(partswap_model_config_dict.keys()),
+                    {"default": default_partswap_model_name},
+                ),
+                "frame_rate": ("FLOAT", {"default": 30.0}),
+                "blend_scale": (
+                    "FLOAT",
+                    {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.1},
+                ),
+                "use_source_seg": (
+                    "BOOLEAN",
+                    {"default": True},
+                ),
+                "hard_edges": (
+                    "BOOLEAN",
+                    {"default": False},
+                ),
+                "chosen_seg_indices": ("STRING", {}),
+            },
+            "optional": {"audio": ("VHS_AUDIO",)},
+        }
+
+    RETURN_TYPES = (
+        "IMAGE",
+        "IMAGE",
+        "IMAGE",
+        "VHS_AUDIO",
+        "FLOAT",
+    )
+    RETURN_NAMES = (
+        "seg_src_viz",
+        "seg_tgt_viz",
+        "images",
+        "audio",
+        "frame_rate",
+    )
+    FUNCTION = "todo"
+    CATEGORY = "FirstOrderMM"
+
+    def todo(
+        self,
+        source_image: torch.Tensor,
+        driving_video_input: torch.Tensor,
+        model_name: str,
+        frame_rate: float,
+        blend_scale: float,
+        use_source_seg: bool,
+        hard_edges: bool,
+        chosen_seg_indices: str,
+        audio=None,
+    ):
+        print(f"{type(source_image)=}")
+        print(f"{type(driving_video_input)=}")
+        print(f"{source_image.shape=}")
+        print(f"{driving_video_input.shape=}")
+        print(f"{type(audio)=}")
+        print(base_dir)
+
+        config_path = (
+            f"{base_dir}/{config_folder}/{partswap_model_config_dict[model_name]}"
+        )
+        checkpoint_path = f"{base_dir}/{get_partswap_model_file_name(model_name)}"
+
+        source_image = reshape_image(source_image, (256, 256))
+        driving_video_org = reshape_image(driving_video_input, (256, 256))
+        driving_video = driving_video_org.unsqueeze(0).permute(0, 2, 1, 3, 4)
+
+        reconstruction_module, segmentation_module = load_partswap_checkpoint(
+            config_path, checkpoint_path, blend_scale, use_fomm=False
+        )
+        print(f"Chosen {chosen_seg_indices}")
+        swap_indices = deserialize_integers(chosen_seg_indices)
+        print(f"Swapping {swap_indices}")
+
+        params = {
+            "swap_indices": swap_indices,
+            "source_image": source_image,
+            "target_video": driving_video,
+            "reconstruction_module": reconstruction_module,
+            "segmentation_module": segmentation_module,
+            "use_source_seg": use_source_seg,
+            "hard_edges": hard_edges,
+        }
+
+        seg_source, seg_targets, predictions = partswap_inference(**params)
+
+        seg_src_viz = visualize_frame(source_image, seg_source)
+        seg_tgt_viz = visualize_frame(driving_video_org[0], seg_targets[0])
+
+        output_images = out_video(predictions)
+
+        return (
+            seg_src_viz,
+            seg_tgt_viz,
+            output_images,
+            audio,
+            frame_rate,
+        )
+
+
 NODE_CLASS_MAPPINGS = {
     "FOMM_Runner": FOMM_Runner,
+    "FOMM_Partswap": FOMM_Partswap,
+    "FOMM_Seg10Chooser": FOMM_Seg10Chooser,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
     "FOMM_Runner": "FOMM Runner",
+    "FOMM_Partswap": "FOMM Partswap",
+    "FOMM_Seg10Chooser": "FOMM_Seg10Chooser",
 }