TSSAT and Clean up (#10)

module_efdm/efdm_model.py

@@ -2,6 +2,7 @@ import torch
 import torch.nn.functional as F
 from torchvision import transforms
 
+from . import net
 from .function import (
     adaptive_instance_normalization,
     adaptive_mean_normalization,
@@ -12,6 +13,20 @@ from .function import (
 )
 
 
+class MODEL_EFDM:
+    def __init__(self, dec_path: str, vgg_path: str, device) -> None:
+        self.decoder = net.decoder
+        vgg = net.vgg
+
+        self.decoder.load_state_dict(torch.load(dec_path))
+        vgg.load_state_dict(torch.load(vgg_path))
+
+        self.vgg = torch.nn.Sequential(*list(vgg.children())[:31])
+
+        self.vgg.eval().to(device)
+        self.decoder.eval().to(device)
+
+
 def inference_efdm(
     src_img: torch.Tensor,
     style_img: torch.Tensor,