v1.0.6-post1 - Auto detect old/new(current) architecture for the load…

…ed model (whoopsie)
rlaphoenix · Oct 22, 2019 · a339d98 · a339d98
1 parent 67916fd
commit a339d98
Show file tree

Hide file tree

Showing 9 changed files with 57 additions and 3 deletions.
diff --git a/.idea/.gitignore b/.idea/.gitignore
diff --git a/.idea/VSGAN.iml b/.idea/VSGAN.iml
diff --git a/.idea/dictionaries/pragma.xml b/.idea/dictionaries/pragma.xml
diff --git a/.idea/inspectionProfiles/profiles_settings.xml b/.idea/inspectionProfiles/profiles_settings.xml
diff --git a/.idea/misc.xml b/.idea/misc.xml
diff --git a/.idea/modules.xml b/.idea/modules.xml
diff --git a/.idea/vcs.xml b/.idea/vcs.xml
diff --git a/setup.py b/setup.py
@@ -5,7 +5,7 @@
 
 setup(
     name="vsgan",
-    version="1.0.6",
+    version="1.0.6-post1",
     author="PRAGMA",
     author_email="[email protected]",
     description="VapourSynth GAN Implementation using RRDBNet, based on ESRGAN's implementation",

diff --git a/vsgan/__init__.py b/vsgan/__init__.py
@@ -26,12 +26,13 @@ def __init__(self, device="cuda"):
     def load_model(self, model, scale):
         self.model_file = model
         self.model_scale = scale
-        # attempt to use Old Arch, and if that fails, attempt to use Current Arch
+        # attempt to use New Arch, and if that fails, attempt to use Old Arch
         # if both fail to be loaded, it will raise it's original exception
-        for arch in range(1):
+        for arch in range(2):
             self.rrdb_net_model = self.get_rrdb_net_arch(arch)
             try:
                 self.rrdb_net_model.load_state_dict(torch.load(self.model_file), strict=True)
+                break
             except RuntimeError:
                 if arch == 1:
                     raise