275 lines
9.1 KiB
Python
275 lines
9.1 KiB
Python
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
|