Solved minor bugs

TMSolution · Apr 21, 2022 · 52abca6 · 52abca6
1 parent 113d9d1
commit 52abca6
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 25 deletions.
diff --git a/jupyter_notebooks/novel_pose_and_expression_synthesis.ipynb b/jupyter_notebooks/novel_pose_and_expression_synthesis.ipynb
diff --git a/nha/models/nha_optimizer.py b/nha/models/nha_optimizer.py
@@ -216,7 +216,11 @@ def __init__(self, max_frame_id, w_lap, w_silh, w_semantic_hair, body_part_weigh
         # loss definitions
         self._leaky_hinge = LeakyHingeLoss(0.0, 1.0, 0.3)
         self._masked_L1 = MaskedCriterion(torch.nn.L1Loss(reduction="none"))
-        self._perceptual_loss = NoSubmoduleWrapper(ResNetLOSS())  # don't store perc_loss weights as model weights
+
+        if Path("assets/InsightFace/backbone.pth").exists():
+            self._perceptual_loss = NoSubmoduleWrapper(ResNetLOSS())  # don't store perc_loss weights as model weights
+        else:
+            self._perceptual_loss = None
 
         # training stage
         self.fit_residuals = False
@@ -257,7 +261,11 @@ def on_train_start(self) -> None:
             shutil.copy(self.hparams["body_part_weights"], body_part_weights_path, follow_symlinks=True)
 
         # moves perceptual loss to own device
-        self._perceptual_loss.to(self.device)
+        try:
+            self._perceptual_loss.to(self.device)
+        except AttributeError:
+            raise AttributeError("You have to download the backbone weights for the perceptual loss. Please refer"
+                                 "to the 'Installation' section of the README")
 
         # hard setting lr
         flame_optim, offset_optim, tex_optim, joint_flame_optim, off_resid_optim, all_resid_optim = self.optimizers()

diff --git a/nha/optimization/perceptual_loss.py b/nha/optimization/perceptual_loss.py
@@ -6,8 +6,9 @@
 from torchvision.transforms import Normalize
 from collections import namedtuple
 import sys
+from pathlib import Path
 
-sys.path.append("deps")
+sys.path.append(str((Path(__file__).parents[2]/"deps")))
 from InsightFace.recognition.arcface_torch.backbones import get_model
 
 

diff --git a/setup.py b/setup.py
@@ -7,6 +7,7 @@
     python_requires=">=3.9",
     install_requires=[
         "ConfigArgParse",
+        "numpy==1.21"
         "torch@https://download.pytorch.org/whl/cu113/torch-1.10.2%2Bcu113-cp39-cp39-linux_x86_64.whl",
         "torchvision@https://download.pytorch.org/whl/cu113/torchvision-0.11.3%2Bcu113-cp39-cp39-linux_x86_64.whl",
         "pytorch3d@https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py39_cu113_pyt1100/pytorch3d-0.6.1-cp39-cp39-linux_x86_64.whl",