Update videoretalking/models/__init__.py

videoretalking/models/__init__.py

@@ -1,37 +1,37 @@
-import torch
-from models.DNet import DNet
-from models.LNet import LNet
-from models.ENet import ENet
-
-
-def _load(checkpoint_path):
-    map_location=None if torch.cuda.is_available() else torch.device('cpu')
-    checkpoint = torch.load(checkpoint_path, map_location=map_location)
-    return checkpoint
-
-def load_checkpoint(path, model):
-    print("Load checkpoint from: {}".format(path))
-    checkpoint = _load(path)
-    s = checkpoint["state_dict"] if 'arcface' not in path else checkpoint
-    new_s = {}
-    for k, v in s.items():
-        if 'low_res' in k:
-            continue
-        else:
-            new_s[k.replace('module.', '')] = v
-    model.load_state_dict(new_s, strict=False)
-    return model
-
-def load_network(LNet_path,ENet_path):
-    L_net = LNet()
-    L_net = load_checkpoint(LNet_path, L_net)
-    E_net = ENet(lnet=L_net)
-    model = load_checkpoint(ENet_path, E_net)
-    return model.eval()
-
-def load_DNet(DNet_path):
-    D_Net = DNet()
-    print("Load checkpoint from: {}".format(DNet_path))
-    checkpoint =  torch.load(DNet_path, map_location=lambda storage, loc: storage)
-    D_Net.load_state_dict(checkpoint['net_G_ema'], strict=False)
+import torch
+from videoretalking.models.DNet import DNet
+from videoretalking.models.LNet import LNet
+from videoretalking.models.ENet import ENet
+
+
+def _load(checkpoint_path):
+    map_location=None if torch.cuda.is_available() else torch.device('cpu')
+    checkpoint = torch.load(checkpoint_path, map_location=map_location)
+    return checkpoint
+
+def load_checkpoint(path, model):
+    print("Load checkpoint from: {}".format(path))
+    checkpoint = _load(path)
+    s = checkpoint["state_dict"] if 'arcface' not in path else checkpoint
+    new_s = {}
+    for k, v in s.items():
+        if 'low_res' in k:
+            continue
+        else:
+            new_s[k.replace('module.', '')] = v
+    model.load_state_dict(new_s, strict=False)
+    return model
+
+def load_network(LNet_path,ENet_path):
+    L_net = LNet()
+    L_net = load_checkpoint(LNet_path, L_net)
+    E_net = ENet(lnet=L_net)
+    model = load_checkpoint(ENet_path, E_net)
+    return model.eval()
+
+def load_DNet(DNet_path):
+    D_Net = DNet()
+    print("Load checkpoint from: {}".format(DNet_path))
+    checkpoint =  torch.load(DNet_path, map_location=lambda storage, loc: storage)
+    D_Net.load_state_dict(checkpoint['net_G_ema'], strict=False)
     return D_Net.eval()