49 lines
1.2 KiB
Python
49 lines
1.2 KiB
Python
import torch
|
|
|
|
from . import net
|
|
|
|
|
|
class MODEL_CAST:
|
|
def __init__(
|
|
self, vgg_path: str, net_ae_path: str, net_dec_b_path: str, device
|
|
) -> None:
|
|
vgg = net.vgg
|
|
vgg.load_state_dict(torch.load(vgg_path))
|
|
vgg = torch.nn.Sequential(*list(vgg.children())[:31])
|
|
|
|
self.netAE = net.ADAIN_Encoder(vgg)
|
|
self.netDec_B = net.Decoder()
|
|
|
|
self.netAE.load_state_dict(load_a_ckpt(net_ae_path))
|
|
self.netDec_B.load_state_dict(load_a_ckpt(net_dec_b_path))
|
|
|
|
self.netAE.to(device).eval()
|
|
self.netDec_B.to(device).eval()
|
|
|
|
|
|
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
|
|
|
|
|
|
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
|