diff --git a/src/spandrel/__helpers/main_registry.py b/src/spandrel/__helpers/main_registry.py index 7755389c..8a970992 100644 --- 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 index dc5b8b86..e4038bab 100644 --- 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 new file mode 100644 index 00000000..9b5ac4a8 Binary files /dev/null and b/tests/images/outputs/16x16/2x-BIGOLDIES.png differ diff --git a/tests/images/outputs/16x16/4x-NMKD-YandereNeo.png b/tests/images/outputs/16x16/4x-NMKD-YandereNeo.png new file mode 100644 index 00000000..39077070 Binary files /dev/null and b/tests/images/outputs/16x16/4x-NMKD-YandereNeo.png differ diff --git a/tests/images/outputs/16x16/8x-ESRGAN.png b/tests/images/outputs/16x16/8x-ESRGAN.png new file mode 100644 index 00000000..f2d8e4d0 Binary files /dev/null and b/tests/images/outputs/16x16/8x-ESRGAN.png differ diff --git a/tests/images/outputs/32x32/2x-BIGOLDIES.png b/tests/images/outputs/32x32/2x-BIGOLDIES.png new file mode 100644 index 00000000..1c7251aa Binary files /dev/null and b/tests/images/outputs/32x32/2x-BIGOLDIES.png differ diff --git a/tests/images/outputs/32x32/4x-NMKD-YandereNeo.png b/tests/images/outputs/32x32/4x-NMKD-YandereNeo.png new file mode 100644 index 00000000..4994bc22 Binary files /dev/null and b/tests/images/outputs/32x32/4x-NMKD-YandereNeo.png differ diff --git a/tests/images/outputs/32x32/8x-ESRGAN.png b/tests/images/outputs/32x32/8x-ESRGAN.png new file mode 100644 index 00000000..8f58b0a4 Binary files /dev/null and b/tests/images/outputs/32x32/8x-ESRGAN.png differ diff --git a/tests/images/outputs/64x64/2x-BIGOLDIES.png b/tests/images/outputs/64x64/2x-BIGOLDIES.png new file mode 100644 index 00000000..281c208d Binary files /dev/null and b/tests/images/outputs/64x64/2x-BIGOLDIES.png differ diff --git a/tests/images/outputs/64x64/4x-NMKD-YandereNeo.png b/tests/images/outputs/64x64/4x-NMKD-YandereNeo.png new file mode 100644 index 00000000..dcdc350e Binary files /dev/null and b/tests/images/outputs/64x64/4x-NMKD-YandereNeo.png differ diff --git a/tests/images/outputs/64x64/8x-ESRGAN.png b/tests/images/outputs/64x64/8x-ESRGAN.png new file mode 100644 index 00000000..46074506 Binary files /dev/null and b/tests/images/outputs/64x64/8x-ESRGAN.png differ diff --git a/tests/images/outputs/8x8/2x-BIGOLDIES.png b/tests/images/outputs/8x8/2x-BIGOLDIES.png new file mode 100644 index 00000000..f7788546 Binary files /dev/null and b/tests/images/outputs/8x8/2x-BIGOLDIES.png differ diff --git a/tests/images/outputs/8x8/4x-NMKD-YandereNeo.png b/tests/images/outputs/8x8/4x-NMKD-YandereNeo.png new file mode 100644 index 00000000..1926320b Binary files /dev/null and b/tests/images/outputs/8x8/4x-NMKD-YandereNeo.png differ diff --git a/tests/images/outputs/8x8/8x-ESRGAN.png b/tests/images/outputs/8x8/8x-ESRGAN.png new file mode 100644 index 00000000..9835bd47 Binary files /dev/null and b/tests/images/outputs/8x8/8x-ESRGAN.png differ diff --git a/tests/test_ESRGAN.py b/tests/test_ESRGAN.py index 0e87ce5f..08cab5a2 100644 --- 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"