Upload model.py

since this file is out of the root path, here is another file updated needed for the previous PR.

models/model.py

@@ -1,6 +1,6 @@
 import os
+from abc import abstractmethod
 from typing import Dict
-from typing import Optional
 from typing import Tuple
 
 import kornia
@@ -10,28 +10,30 @@ import torch.nn as nn
 import torch.nn.functional as F
 from loguru import logger
 
-from arcface_torch.backbones.iresnet import iresnet100
-from configs.train_config import TrainConfig
 from Deep3DFaceRecon_pytorch.models.bfm import ParametricFaceModel
 from Deep3DFaceRecon_pytorch.models.networks import ReconNetWrapper
 from HRNet.hrnet import HighResolutionNet
+from arcface_torch.backbones.iresnet import iresnet100
 from models.discriminator import Discriminator
 from models.gan_loss import GANLoss
 from models.generator import Generator
 from models.init_weight import init_net
 
-
 class HifiFace:
     def __init__(
         self,
         identity_extractor_config,
-        is_training=True,
-        device="cpu",
-        load_checkpoint: Optional[Tuple[str, int]] = None,
+        generator_path,
+        is_training=False,
+        device="cpu"
     ):
         super(HifiFace, self).__init__()
+        self.d_optimizer = None
+        self.g_optimizer = None
         self.generator = Generator(identity_extractor_config)
         self.is_training = is_training
+        self.device = device
+        self.generator_path = generator_path
 
         if self.is_training:
             self.lr = TrainConfig().lr
@@ -80,10 +82,9 @@ class HifiFace:
 
             self.dilation_kernel = torch.ones(5, 5)
 
-        if load_checkpoint is not None:
-            self.load(load_checkpoint[0], load_checkpoint[1])
+        self.load_checkpoint()
 
-        self.setup(device)
+        self.setup(self.device)
 
     def save(self, path, idx=None):
         os.makedirs(path, exist_ok=True)
@@ -100,18 +101,9 @@ class HifiFace:
             torch.save(self.generator.state_dict(), g_path)
             torch.save(self.discriminator.state_dict(), d_path)
 
-    def load(self, path, idx=None):
-        if idx is None:
-            g_path = os.path.join(path, "generator.pth")
-            d_path = os.path.join(path, "discriminator.pth")
-        else:
-            g_path = os.path.join(path, f"generator_{idx}.pth")
-            d_path = os.path.join(path, f"discriminator_{idx}.pth")
-        logger.info(f"Loading generator from {g_path}")
-        self.generator.load_state_dict(torch.load(g_path, map_location="cpu"))
-        if self.is_training:
-            logger.info(f"Loading discriminator from {d_path}")
-            self.discriminator.load_state_dict(torch.load(d_path, map_location="cpu"))
+    @abstractmethod
+    def load_checkpoint(self):
+        pass
 
     def setup(self, device):
         self.generator.to(device)
@@ -399,37 +391,18 @@ class HifiFace:
         }
 
 
-if __name__ == "__main__":
-    import torch
-    import cv2
-    from configs.train_config import TrainConfig
-
-    identity_extractor_config = TrainConfig().identity_extractor_config
-
-    model = HifiFace(identity_extractor_config, is_training=True)
-
-    # src = cv2.imread("/home/xuehongyang/data/test1.jpg")
-    # tgt = cv2.imread("/home/xuehongyang/data/test2.jpg")
-    # src = cv2.cvtColor(src, cv2.COLOR_BGR2RGB)
-    # tgt = cv2.cvtColor(tgt, cv2.COLOR_BGR2RGB)
-    # src = cv2.resize(src, (256, 256))
-    # tgt = cv2.resize(tgt, (256, 256))
-    # src = src.transpose(2, 0, 1)[None, ...]
-    # tgt = tgt.transpose(2, 0, 1)[None, ...]
-    # source_img = torch.from_numpy(src).float() / 255.0
-    # target_img = torch.from_numpy(tgt).float() / 255.0
-    # same_id_mask = torch.Tensor([1]).unsqueeze(0)
-    # tgt_mask = target_img[:, 0, :, :].unsqueeze(1)
-    # if torch.cuda.is_available():
-    #     model.to("cuda:3")
-    #     source_img = source_img.to("cuda:3")
-    #     target_img = target_img.to("cuda:3")
-    #     tgt_mask = tgt_mask.to("cuda:3")
-    #     same_id_mask = same_id_mask.to("cuda:3")
-    #     source_img = source_img.repeat(16, 1, 1, 1)
-    #     target_img = target_img.repeat(16, 1, 1, 1)
-    #     tgt_mask = tgt_mask.repeat(16, 1, 1, 1)
-    #     same_id_mask = same_id_mask.repeat(16, 1)
-    # while True:
-    #     x = model.optimize(source_img, target_img, tgt_mask, same_id_mask)
-    #     print(x[0]["loss_generator"])
+class HifiFaceST(HifiFace):
+    def __init__(self, identity_extractor_config, device, generator_path):
+        super().__init__(identity_extractor_config, device=device, generator_path=generator_path)
+
+    def load_checkpoint(self):
+        self.generator.load_state_dict(torch.load(self.generator_path, map_location=self.device))
+        logger.info(f"Loading generator from {self.generator_path}")
+
+class HifiFaceWGM(HifiFace):
+    def __init__(self, identity_extractor_config, device, generator_path):
+        super().__init__(identity_extractor_config, device=device, generator_path=generator_path)
+
+    def load_checkpoint(self):
+        self.generator.load_state_dict(torch.load(self.generator_path, map_location=self.device))
+        logger.info(f"Loading generator from {self.generator_path}")