Fix ESRGAN loading + add more community model tests (#27)

* Fix ESRGAN loading + add more community model tests * Add image tests for newly added model tests
chaiNNer-org · Nov 20, 2023 · 809aff6 · 809aff6
1 parent 58d4233
commit 809aff6
Show file tree

Hide file tree

Showing 15 changed files with 87 additions and 2 deletions.
diff --git a/src/spandrel/__helpers/main_registry.py b/src/spandrel/__helpers/main_registry.py
@@ -177,8 +177,6 @@ def _detect(state_dict: StateDict) -> bool:
             _has_keys(
                 "model.0.weight",
                 "model.1.sub.0.RDB1.conv1.0.weight",
-                "model.2.weight",
-                "model.4.weight",
             )(state)
             or _has_keys(
                 "conv_first.weight",

diff --git a/tests/__snapshots__/test_ESRGAN.ambr b/tests/__snapshots__/test_ESRGAN.ambr
@@ -44,6 +44,51 @@
     ]),
   )
 # ---
+# name: test_ESRGAN_community_2x
+  SRModelDescriptor(
+    architecture='ESRGAN',
+    input_channels=3,
+    output_channels=3,
+    scale=2,
+    size_requirements=SizeRequirements(minimum=None, multiple_of=None, square=False),
+    supports_bfloat16=True,
+    supports_half=True,
+    tags=list([
+      '64nf',
+      '23nb',
+    ]),
+  )
+# ---
+# name: test_ESRGAN_community_4x
+  SRModelDescriptor(
+    architecture='ESRGAN',
+    input_channels=3,
+    output_channels=3,
+    scale=4,
+    size_requirements=SizeRequirements(minimum=None, multiple_of=None, square=False),
+    supports_bfloat16=True,
+    supports_half=True,
+    tags=list([
+      '32nf',
+      '12nb',
+    ]),
+  )
+# ---
+# name: test_ESRGAN_community_8x
+  SRModelDescriptor(
+    architecture='ESRGAN',
+    input_channels=3,
+    output_channels=3,
+    scale=8,
+    size_requirements=SizeRequirements(minimum=None, multiple_of=None, square=False),
+    supports_bfloat16=True,
+    supports_half=True,
+    tags=list([
+      '64nf',
+      '23nb',
+    ]),
+  )
+# ---
 # name: test_RealESRGAN_x2plus
   SRModelDescriptor(
     architecture='ESRGAN',

diff --git a/tests/images/outputs/16x16/2x-BIGOLDIES.png b/tests/images/outputs/16x16/2x-BIGOLDIES.png
diff --git a/tests/images/outputs/16x16/4x-NMKD-YandereNeo.png b/tests/images/outputs/16x16/4x-NMKD-YandereNeo.png
diff --git a/tests/images/outputs/16x16/8x-ESRGAN.png b/tests/images/outputs/16x16/8x-ESRGAN.png
diff --git a/tests/images/outputs/32x32/2x-BIGOLDIES.png b/tests/images/outputs/32x32/2x-BIGOLDIES.png
diff --git a/tests/images/outputs/32x32/4x-NMKD-YandereNeo.png b/tests/images/outputs/32x32/4x-NMKD-YandereNeo.png
diff --git a/tests/images/outputs/32x32/8x-ESRGAN.png b/tests/images/outputs/32x32/8x-ESRGAN.png
diff --git a/tests/images/outputs/64x64/2x-BIGOLDIES.png b/tests/images/outputs/64x64/2x-BIGOLDIES.png
diff --git a/tests/images/outputs/64x64/4x-NMKD-YandereNeo.png b/tests/images/outputs/64x64/4x-NMKD-YandereNeo.png
diff --git a/tests/images/outputs/64x64/8x-ESRGAN.png b/tests/images/outputs/64x64/8x-ESRGAN.png
diff --git a/tests/images/outputs/8x8/2x-BIGOLDIES.png b/tests/images/outputs/8x8/2x-BIGOLDIES.png
diff --git a/tests/images/outputs/8x8/4x-NMKD-YandereNeo.png b/tests/images/outputs/8x8/4x-NMKD-YandereNeo.png
diff --git a/tests/images/outputs/8x8/8x-ESRGAN.png b/tests/images/outputs/8x8/8x-ESRGAN.png
diff --git a/tests/test_ESRGAN.py b/tests/test_ESRGAN.py
@@ -18,6 +18,48 @@ def test_ESRGAN_community(snapshot):
     )
 
 
+def test_ESRGAN_community_2x(snapshot):
+    file = ModelFile.from_url(
+        "https://objectstorage.us-phoenix-1.oraclecloud.com/n/ax6ygfvpvzka/b/open-modeldb-files/o/2x-BIGOLDIES.pth"
+    )
+    model = ModelLoader().load_from_file(file.path)
+    assert model == snapshot(exclude=disallowed_props)
+    assert isinstance(model.model, RRDBNet)
+    assert_image_inference(
+        file,
+        model,
+        [TestImage.SR_8, TestImage.SR_16, TestImage.SR_32, TestImage.SR_64],
+    )
+
+
+def test_ESRGAN_community_4x(snapshot):
+    file = ModelFile.from_url(
+        "https://objectstorage.us-phoenix-1.oraclecloud.com/n/ax6ygfvpvzka/b/open-modeldb-files/o/4x-NMKD-YandereNeo.pth"
+    )
+    model = ModelLoader().load_from_file(file.path)
+    assert model == snapshot(exclude=disallowed_props)
+    assert isinstance(model.model, RRDBNet)
+    assert_image_inference(
+        file,
+        model,
+        [TestImage.SR_8, TestImage.SR_16, TestImage.SR_32, TestImage.SR_64],
+    )
+
+
+def test_ESRGAN_community_8x(snapshot):
+    file = ModelFile.from_url(
+        "https://objectstorage.us-phoenix-1.oraclecloud.com/n/ax6ygfvpvzka/b/open-modeldb-files/o/8x-ESRGAN.pth"
+    )
+    model = ModelLoader().load_from_file(file.path)
+    assert model == snapshot(exclude=disallowed_props)
+    assert isinstance(model.model, RRDBNet)
+    assert_image_inference(
+        file,
+        model,
+        [TestImage.SR_8, TestImage.SR_16, TestImage.SR_32, TestImage.SR_64],
+    )
+
+
 def test_BSRGAN(snapshot):
     file = ModelFile.from_url(
         "https://github.com/cszn/KAIR/releases/download/v1.0/BSRGAN.pth"