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

add fomm models to partswap

Browse Source
This commit is contained in:
FuouM
2024-07-01 11:50:31 +07:00
parent 0194960d6d
commit dac0fb62e1
5 changed files with 81 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
constants.py
+12 -4
View File
@@ -25,12 +25,18 @@ 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",
"vox-cpk": "vox-256-sem-10segments.yaml",
"vox-adv-cpk": "vox-256-sem-10segments.yaml",
}
partswap_fomm_model_names = ["vox-cpk", "vox-adv-cpk"]
partswap_model_length_dict = {
"vox-5segments": 6,
"vox-5segments": 6,
"vox-10segments": 11,
"vox-15segments": 15
"vox-15segments": 15,
"vox-cpk": 19,
"vox-adv-cpk": 19,
}
partswap_model_names = list(partswap_model_config_dict.keys())
@@ -57,7 +63,7 @@ seg5_segment_names = [
"m_2_frnt_face",
"m_3_shoulder",
"m_4_hair",
"m_5_fr_head"
"m_5_fr_head",
]
seg_15_segment_names = [
@@ -76,11 +82,13 @@ seg_15_segment_names = [
"m_12_lt_ear",
"m_13_lt_hair",
"m_14_rt_hair",
"m_15_neck"
"m_15_neck",
]
seg_model_dict = {
"vox-5segments": seg5_segment_names,
"vox-10segments": seg10_segment_names,
"vox-15segments": seg_15_segment_names,
"vox": seg10_segment_names,
"vox-adv": seg10_segment_names,
}
modules/dense_motion.py
+42 -9
View File
@@ -29,13 +29,13 @@ class DenseMotionNetwork(nn.Module):
if self.scale_factor != 1:
self.down = AntiAliasInterpolation2d(num_channels, self.scale_factor)
def create_heatmap_representations(self, source_image, kp_driving, kp_source):
def create_heatmap_representations(self, source_image, kp_driving, kp_source, kp_key='value'):
"""
Eq 6. in the paper H_k(z)
"""
spatial_size = source_image.shape[2:]
gaussian_driving = kp2gaussian(kp_driving, spatial_size=spatial_size, kp_variance=self.kp_variance)
gaussian_source = kp2gaussian(kp_source, spatial_size=spatial_size, kp_variance=self.kp_variance)
gaussian_driving = kp2gaussian(kp_driving, spatial_size=spatial_size, kp_variance=self.kp_variance, kp_key=kp_key)
gaussian_source = kp2gaussian(kp_source, spatial_size=spatial_size, kp_variance=self.kp_variance, kp_key=kp_key)
heatmap = gaussian_driving - gaussian_source
#adding background feature
@@ -44,22 +44,22 @@ class DenseMotionNetwork(nn.Module):
heatmap = heatmap.unsqueeze(2)
return heatmap
def create_sparse_motions(self, source_image, kp_driving, kp_source):
def create_sparse_motions(self, source_image, kp_driving, kp_source, kp_key='value', jacob_key='jacobian'):
"""
Eq 4. in the paper T_{s<-d}(z)
"""
bs, _, h, w = source_image.shape
identity_grid = make_coordinate_grid((h, w), type=kp_source['value'].type())
identity_grid = make_coordinate_grid((h, w), type=kp_source[kp_key].type())
identity_grid = identity_grid.view(1, 1, h, w, 2)
coordinate_grid = identity_grid - kp_driving['value'].view(bs, self.num_kp, 1, 1, 2)
if 'jacobian' in kp_driving:
jacobian = torch.matmul(kp_source['jacobian'], torch.inverse(kp_driving['jacobian']))
coordinate_grid = identity_grid - kp_driving[kp_key].view(bs, self.num_kp, 1, 1, 2)
if jacob_key in kp_driving:
jacobian = torch.matmul(kp_source[jacob_key], torch.inverse(kp_driving[jacob_key]))
jacobian = jacobian.unsqueeze(-3).unsqueeze(-3)
jacobian = jacobian.repeat(1, 1, h, w, 1, 1)
coordinate_grid = torch.matmul(jacobian, coordinate_grid.unsqueeze(-1))
coordinate_grid = coordinate_grid.squeeze(-1)
driving_to_source = coordinate_grid + kp_source['value'].view(bs, self.num_kp, 1, 1, 2)
driving_to_source = coordinate_grid + kp_source[kp_key].view(bs, self.num_kp, 1, 1, 2)
#adding background feature
identity_grid = identity_grid.repeat(bs, 1, 1, 1, 1)
@@ -111,3 +111,36 @@ class DenseMotionNetwork(nn.Module):
out_dict['occlusion_map'] = occlusion_map
return out_dict
def forward_partswap(self, source_image, seg_target, seg_source):
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, seg_target, seg_source, kp_key='shift')
sparse_motion = self.create_sparse_motions(source_image, seg_target, seg_source, kp_key='shift', jacob_key='affine')
deformed_source = self.create_deformed_source_image(source_image, sparse_motion)
out_dict['sparse_deformed'] = deformed_source
input = torch.cat([heatmap_representation, deformed_source], dim=2)
input = input.view(bs, -1, h, w)
prediction = self.hourglass(input)
mask = self.mask(prediction)
mask = F.softmax(mask, dim=1)
out_dict['mask'] = mask
mask = mask.unsqueeze(2)
sparse_motion = sparse_motion.permute(0, 1, 4, 2, 3)
deformation = (sparse_motion * mask).sum(dim=1)
deformation = deformation.permute(0, 2, 3, 1)
out_dict['deformation'] = deformation
if self.occlusion:
visibility = torch.sigmoid(self.occlusion(prediction))
out_dict['visibility'] = visibility
return out_dict
modules/util.py
+2 -2
View File
@@ -6,11 +6,11 @@ import torch
from ..sync_batchnorm.batchnorm import SynchronizedBatchNorm2d as BatchNorm2d
def kp2gaussian(kp, spatial_size, kp_variance):
def kp2gaussian(kp, spatial_size, kp_variance, kp_key='value'):
"""
Transform a keypoint into gaussian like representation
"""
mean = kp['value']
mean = kp[kp_key]
coordinate_grid = make_coordinate_grid(spatial_size, mean.type())
number_of_leading_dimensions = len(mean.shape) - 1
partswap_model.py
+10 -2
View File
@@ -14,7 +14,15 @@ class PartSwapGenerator(ReconstructionModule):
self.blend_downsample = AntiAliasInterpolation2d(1, blend_scale)
if first_order_motion_model:
self.dense_motion_network = DenseMotionNetwork()
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
@@ -46,7 +54,7 @@ class PartSwapGenerator(ReconstructionModule):
deformation = (segment_motions * mask).sum(dim=1)
deformation = deformation.permute(0, 2, 3, 1)
else:
motion = self.dense_motion_network(
motion = self.dense_motion_network.forward_partswap(
source_image=source_image, seg_target=seg_target, seg_source=seg_source
)
deformation = motion["deformation"]
run.py
+15 -8
View File
@@ -6,22 +6,23 @@
"""
from pathlib import Path
import torch
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
import torch
from .constants import (
config_folder,
default_model_name,
default_partswap_model_name,
partswap_fomm_model_names,
partswap_model_config_dict,
partswap_model_length_dict,
partswap_model_names,
supported_model_names,
partswap_model_length_dict,
)
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 .seg_viz import visualize_frame
from .utils import (
build_seg_arguments,
@@ -273,8 +274,12 @@ class FOMM_Partswap:
driving_video_org = reshape_image(driving_video_input, (256, 256))
driving_video = driving_video_org.unsqueeze(0).permute(0, 2, 1, 3, 4)
use_fomm = False
if model_name in partswap_fomm_model_names:
use_fomm = True
reconstruction_module, segmentation_module = load_partswap_checkpoint(
config_path, checkpoint_path, blend_scale, use_fomm=False
config_path, checkpoint_path, blend_scale, use_fomm=use_fomm
)
print(f"Chosen {chosen_seg_indices}")
swap_indices = deserialize_integers(chosen_seg_indices)
@@ -294,7 +299,9 @@ class FOMM_Partswap:
seg_source, seg_targets, predictions = partswap_inference(**params)
seg_src_viz = visualize_frame(source_image, seg_source, alpha=viz_alpha)
seg_tgt_viz = visualize_frame(driving_video_org[0], seg_targets[0], alpha=viz_alpha)
seg_tgt_viz = visualize_frame(
driving_video_org[0], seg_targets[0], alpha=viz_alpha
)
output_images = out_video(predictions)