Upload fs_model.py

SimSwap/models/fs_model.py

@@ -0,0 +1,17 @@
+
+import torch
+import os
+from .base_model import BaseModel
+from . import networks
+
+class FSModel(BaseModel):
+    def initialize(self, opt):
+        self.opt = opt
+        self.device = torch.device("cpu")
+        self.save_dir = os.path.join(opt.checkpoints_dir, opt.name)
+
+        self.netG = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG, self.device)
+        self.netArc = networks.define_arcface(opt, self.device)
+
+        self.load_network(self.netG, 'G', opt.which_epoch)
+        self.load_network(self.netArc, 'arcface', opt.which_epoch)