comply with optimized generator

alievk · alievk · commit 295b366f8a75 · 2020-05-10T12:29:59.000+03:00
diff --git a/afy/predictor_local.py b/afy/predictor_local.py
@@ -2,7 +2,7 @@
 import torch
 import yaml
 from modules.keypoint_detector import KPDetector
-from modules.generator import OcclusionAwareGenerator
+from modules.generator_optim import OcclusionAwareGenerator
 from sync_batchnorm import DataParallelWithCallback
 import numpy as np
 import face_alignment
@@ -45,9 +45,6 @@ def load_checkpoints(self):
         generator.load_state_dict(checkpoint['generator'])
         kp_detector.load_state_dict(checkpoint['kp_detector'])
     
-        generator = DataParallelWithCallback(generator)
-        kp_detector = DataParallelWithCallback(kp_detector)
-    
         generator.eval()
         kp_detector.eval()
         
@@ -60,6 +57,14 @@ def set_source_image(self, source_image):
         self.source = to_tensor(source_image).to(self.device)
         self.kp_source = self.kp_detector(self.source)
 
+        if self.enc_downscale > 1:
+            h, w = int(self.source.shape[2] / self.enc_downscale), int(self.source.shape[3] / self.enc_downscale)
+            source_enc = torch.nn.functional.interpolate(self.source, size=(h, w), mode='bilinear')
+        else:
+            source_enc = self.source
+
+        self.generator.encode_source(source_enc)
+
     def predict(self, driving_frame):
         with torch.no_grad():
             driving = to_tensor(driving_frame).to(self.device)
@@ -74,17 +79,7 @@ def predict(self, driving_frame):
                                 kp_driving_initial=self.kp_driving_initial, use_relative_movement=self.relative,
                                 use_relative_jacobian=self.relative, adapt_movement_scale=self.adapt_movement_scale)
 
-            if self.enc_downscale > 1:
-                h, w = int(source.shape[2] / self.enc_downscale), int(source.shape[3] / self.enc_downscale)
-                source_enc = torch.nn.functional.interpolate(source, size=(h, w), mode='bilinear')
-            else:
-                source_enc = None
-
-            try:
-                out = self.generator(self.source, kp_source=self.kp_source, kp_driving=kp_norm, source_image_enc=source_enc, optim_ret=True)
-            except TypeError:
-                Once('\n*** Please update FOMM:\ncd fomm\ngit pull\n')
-                out = self.generator(self.source, kp_source=self.kp_source, kp_driving=kp_norm)
+            out = self.generator(self.source, kp_source=self.kp_source, kp_driving=kp_norm)
 
             out = np.transpose(out['prediction'].data.cpu().numpy(), [0, 2, 3, 1])[0]
             out = (np.clip(out, 0, 1) * 255).astype(np.uint8)
