Fix ESRGAN loading + add more community model tests

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",

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',

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"