93 lines
3.0 KiB
Python
93 lines
3.0 KiB
Python
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
|