nodes/FuouM-ComfyUI-FirstOrderMM
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor

Browse Source
This commit is contained in:
FuouM
2024-08-10 19:36:06 +07:00
parent 4a2149feb3
commit 8d5c84c0ec
12 changed files with 328 additions and 331 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+6
View File
@@ -7,6 +7,7 @@ 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)
4. [Learning Motion Refinement for Unsupervised Face Animation](https://github.com/JialeTao/MRFA/)
More will come soon
@@ -69,6 +70,9 @@ Doesn't need any
* `avd`: similar to `relative`, may yield better but more "jittery/jumpy" result
* `find_best_frame`: Same as FOMM
### MRFA
## Installation
1. Clone the repo to `ComfyUI/custom_nodes/`
@@ -129,3 +133,5 @@ resnet18-5c106cde.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`.
For **MRFA**, download the model from [Pretrained models](https://github.com/JialeTao/MRFA/?tab=readme-ov-file#pretrained-models) section and place it here: `mrfa_module/models/`
fomm_loader.py
-35
View File
@@ -1,35 +0,0 @@
import yaml
import torch
from .modules.generator import OcclusionAwareGenerator
from .modules.keypoint_detector import KPDetector
def load_checkpoint(
config_path: str, checkpoint_path: str, cpu=False
) -> tuple[OcclusionAwareGenerator, KPDetector]:
with open(config_path) as f:
config = yaml.full_load(f)
generator = OcclusionAwareGenerator(
**config["model_params"]["generator_params"],
**config["model_params"]["common_params"],
)
kp_detector = KPDetector(
**config["model_params"]["kp_detector_params"],
**config["model_params"]["common_params"],
)
if not cpu:
generator.cuda()
kp_detector.cuda()
checkpoint = torch.load(checkpoint_path)
else:
checkpoint = torch.load(checkpoint_path, map_location=torch.device("cpu"))
generator.load_state_dict(checkpoint["generator"])
kp_detector.load_state_dict(checkpoint["kp_detector"])
generator.eval()
kp_detector.eval()
return generator, kp_detector
articulate_inference.py → inference_articulate.py
View File
fomm_inference.py → inference_fomm.py
+40 -5
View File
@@ -1,11 +1,12 @@
import torch
import numpy as np
import torch
import yaml
from comfy.utils import ProgressBar
from tqdm import tqdm
from scipy.spatial import ConvexHull
from tqdm import tqdm
from .modules.generator import OcclusionAwareGenerator
from .modules.keypoint_detector import KPDetector
from .utils import tensor_to_numpy
try:
@@ -24,6 +25,7 @@ def normalize_kp_best_frame(kp):
kp[:, :2] = kp[:, :2] / area
return kp
def normalize_kp(
kp_source,
kp_driving,
@@ -55,11 +57,12 @@ def normalize_kp(
return kp_new
def find_best_frame(source_image, driving_video: list, cpu=False) -> int:
if not face_alignment_available:
print("Face Alignment module were not found. Best frame set to 0.")
return 0
print("Initializing Face Alignment...")
fa = face_alignment.FaceAlignment(
face_alignment.LandmarksType.TWO_D,
@@ -93,6 +96,7 @@ def find_best_frame(source_image, driving_video: list, cpu=False) -> int:
return frame_num
def inference(
source_image,
driving_video: list,
@@ -182,4 +186,35 @@ def inference_best_frame(
predictions_first = predictions_first[::-1] # Reverse the first half back
predictions.extend(predictions_first)
predictions.extend(predictions_second)
return predictions
return predictions
def load_checkpoint(
config_path: str, checkpoint_path: str, cpu=False
) -> tuple[OcclusionAwareGenerator, KPDetector]:
with open(config_path) as f:
config = yaml.full_load(f)
generator = OcclusionAwareGenerator(
**config["model_params"]["generator_params"],
**config["model_params"]["common_params"],
)
kp_detector = KPDetector(
**config["model_params"]["kp_detector_params"],
**config["model_params"]["common_params"],
)
if not cpu:
generator.cuda()
kp_detector.cuda()
checkpoint = torch.load(checkpoint_path)
else:
checkpoint = torch.load(checkpoint_path, map_location=torch.device("cpu"))
generator.load_state_dict(checkpoint["generator"])
kp_detector.load_state_dict(checkpoint["kp_detector"])
generator.eval()
kp_detector.eval()
return generator, kp_detector
inference_partswap.py
+274
View File
@@ -0,0 +1,274 @@
import numpy as np
import torch
import torch.nn.functional as F
import yaml
from comfy.utils import ProgressBar
from tqdm import tqdm
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 face_parse_seg(
source_image, face_parser_model, in_size=(512, 512), out_size=(64, 64)
):
seg = F.interpolate(source_image, size=in_size)
seg = (seg - face_parser_model.mean) / face_parser_model.std
seg = torch.softmax(face_parser_model(seg)[0], dim=1)
seg = F.interpolate(seg, size=out_size)
return seg
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(
block_expansion=64,
num_blocks=5,
max_features=1024,
num_kp=10,
num_channels=3,
estimate_occlusion_map=True,
scale_factor=0.25,
)
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.forward_partswap(
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,
face_parser_model=None,
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)
# print(f"{seg_source['segmentation'].shape}")
num_frames = target_video.shape[2]
pbar = ProgressBar(num_frames)
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 face_parser_model is not None:
blend_mask = face_parse_seg(
source_image if use_source_seg else target_frame,
face_parser_model,
(512, 512),
(64, 64),
)
else:
blend_mask = (
seg_source["segmentation"]
if use_source_seg
else 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())
# print(f"{blend_mask.shape=}")
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]
)
pbar.update_absolute(frame_idx, num_frames)
return seg_source, seg_targets, predictions
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 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
spline_inference.py → inference_spline.py
+1 -1
View File
@@ -5,7 +5,7 @@ import yaml
import tqdm
from comfy.utils import ProgressBar
from .fomm_inference import find_best_frame
from .inference_fomm import find_best_frame
from .spline_module.avd_network import AVDNetwork
from .spline_module.dense_motion import DenseMotionNetwork
partswap_inference.py
-84
View File
@@ -1,84 +0,0 @@
import numpy as np
import torch
import torch.nn.functional as F
from tqdm import tqdm
from comfy.utils import ProgressBar
from .modules.segmentation_module import SegmentationModule
from .partswap_model import PartSwapGenerator
def face_parse_seg(
source_image, face_parser_model, in_size=(512, 512), out_size=(64, 64)
):
seg = F.interpolate(source_image, size=in_size)
seg = (seg - face_parser_model.mean) / face_parser_model.std
seg = torch.softmax(face_parser_model(seg)[0], dim=1)
seg = F.interpolate(seg, size=out_size)
return seg
def partswap_inference(
swap_indices: list[int],
source_image,
target_video,
reconstruction_module: PartSwapGenerator,
segmentation_module: SegmentationModule,
use_source_seg: bool,
face_parser_model=None,
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)
# print(f"{seg_source['segmentation'].shape}")
num_frames = target_video.shape[2]
pbar = ProgressBar(num_frames)
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 face_parser_model is not None:
blend_mask = face_parse_seg(
source_image if use_source_seg else target_frame,
face_parser_model,
(512, 512),
(64, 64),
)
else:
blend_mask = (
seg_source["segmentation"]
if use_source_seg
else 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())
print(f"{blend_mask.shape=}")
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]
)
pbar.update_absolute(frame_idx, num_frames)
return seg_source, seg_targets, predictions
partswap_loader.py
-92
View File
@@ -1,92 +0,0 @@
import torch
import yaml
from .modules.segmentation_module import SegmentationModule
from .partswap_model import PartSwapGenerator
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 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
partswap_model.py
-105
View File
@@ -1,105 +0,0 @@
import torch.nn.functional as F
from .modules.dense_motion import DenseMotionNetwork
from .modules.reconstruction_module import ReconstructionModule
from .modules.util import AntiAliasInterpolation2d
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(
block_expansion=64,
num_blocks=5,
max_features=1024,
num_kp=10,
num_channels=3,
estimate_occlusion_map=True,
scale_factor=0.25,
)
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.forward_partswap(
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
pyproject.toml
+1 -1
View File
@@ -1,6 +1,6 @@
[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)"
description = "ComfyUI-native nodes to run First Order Motion Model for Image Animation and its non-diffusion-based successors. [a/https://github.com/AliaksandrSiarohin/first-order-model](https://github.com/AliaksandrSiarohin/first-order-model)"
version = "1.0.3"
license = { file = "LICENSE" }
dependencies = ["numpy", "torch", "scipy", "yaml", "matplotlib"]
requirements.txt
+1 -1
View File
@@ -1,5 +1,5 @@
numpy
torch
scipy
yaml
pyyaml
matplotlib
run.py
+5 -7
View File
@@ -2,14 +2,13 @@
@author: Fuou Marinas
@title: ComfyUI-FirstOrderMM
@nickname: FOMM
@description: Run First Order Motion Model for Image Animation in ComfyUI.
@description: ComfyUI-native nodes to run First Order Motion Model for Image Animation and its non-diffusion-based successors.
"""
from pathlib import Path
import torch
from .articulate_inference import articulate_inference
from .constants import (
ARTICULATE_CFG_PATH,
ARTICULATE_MODEL_PATH,
@@ -28,12 +27,11 @@ from .constants import (
supported_model_names,
)
from .face_parsing.face_parsing_loader import load_face_parser_model
from .fomm_inference import inference, inference_best_frame
from .fomm_loader import load_checkpoint
from .partswap_inference import partswap_inference
from .partswap_loader import load_partswap_checkpoint
from .inference_articulate import articulate_inference
from .inference_fomm import inference, inference_best_frame, load_checkpoint
from .inference_partswap import load_partswap_checkpoint, partswap_inference
from .inference_spline import spline_inference
from .seg_viz import visualize_frame
from .spline_inference import spline_inference
from .utils import (
build_seg_arguments,
get_config_path,