28 lines
615 B
Python
28 lines
615 B
Python
import torch
|
|
|
|
def load_a_ckpt(ckpt_path: str):
|
|
state_dict = torch.load(ckpt_path)
|
|
if hasattr(state_dict, "_metadata"):
|
|
del state_dict._metadata
|
|
return state_dict
|
|
|
|
|
|
def inference_ucast(
|
|
src_img: torch.Tensor,
|
|
style_img: torch.Tensor,
|
|
netAE,
|
|
netDec_B,
|
|
device,
|
|
):
|
|
real_A = src_img.permute(0, 3, 1, 2).to(device)
|
|
real_B = style_img.permute(0, 3, 1, 2).to(device)
|
|
|
|
with torch.no_grad():
|
|
real_A_feat = netAE.forward(real_A, real_B)
|
|
fake_B = netDec_B.forward(real_A_feat)
|
|
|
|
if device.type != "cpu":
|
|
torch.cuda.empty_cache()
|
|
|
|
return fake_B
|