Image tests

.github/workflows/lint.yml

@@ -53,6 +53,7 @@ jobs:
                   python-version: '3.9'
                   cache: 'pip'
             - run: pip install .[typecheck]
+            - run: pip install .[test]
             - run: pyright src/
             - run: pyright tests/
             - run: pyright scripts/

pyproject.toml

@@ -37,7 +37,7 @@ dynamic = ["version"]
 build = ["setuptools>=46.4.0", "wheel", "build", "twine"]
 lint = ["ruff==0.1.4"]
 typecheck = ["pyright==1.1.335"]
-test = ["pytest==7.4.0", "syrupy==4.6.0"]
+test = ["pytest==7.4.0", "syrupy==4.6.0", "opencv-python==4.8.1.78"]
 
 # [tool.setuptools.dynamic]
 # version = { attr = "spandrel.VERSION" }
@@ -73,4 +73,4 @@ ignore = [
 "**/tests/**/*" = ["N802"]
 
 [tool.pytest.ini_options]
-filterwarnings = ["ignore::DeprecationWarning"]
+filterwarnings = ["ignore::DeprecationWarning", "ignore::UserWarning"]

src/spandrel/architectures/ESRGAN/arch/RRDB.py

@@ -45,6 +45,7 @@ def __init__(
         super().__init__()
 
         self.shuffle_factor = shuffle_factor
+        self.scale = scale
 
         upsample_block = {
             "upconv": B.upconv_block,

src/spandrel/architectures/SwinIR/__init__.py

@@ -178,6 +178,7 @@ def load(state_dict: StateDict) -> SRModelDescriptor[SwinIR]:
         num_feat=num_feat,
         num_in_ch=num_in_ch,
         num_out_ch=num_out_ch,
+        start_unshuffle=start_unshuffle,
     )
 
     head_length = len(depths)  # type: ignore

src/spandrel/architectures/SwinIR/arch/SwinIR.py

@@ -835,9 +835,11 @@ def __init__(
         img_range=1.0,
         upsampler="",
         resi_connection="1conv",
+        start_unshuffle=1,
         **kwargs,
     ):
         super().__init__()
+        self.start_unshuffle = start_unshuffle
         num_in_ch = in_chans
         num_out_ch = in_chans
         num_feat = 64

tests/test_Compact.py

@@ -1,7 +1,7 @@
 from spandrel import ModelLoader
 from spandrel.architectures.Compact import SRVGGNetCompact
 
-from .util import ModelFile, disallowed_props
+from .util import ModelFile, TestImage, assert_image_inference, disallowed_props
 
 
 def test_Compact_realesr_general_x4v3(snapshot):
@@ -11,6 +11,11 @@ def test_Compact_realesr_general_x4v3(snapshot):
     model = ModelLoader().load_from_file(file.path)
     assert model == snapshot(exclude=disallowed_props)
     assert isinstance(model.model, SRVGGNetCompact)
+    assert_image_inference(
+        file,
+        model,
+        [TestImage.SR_16, TestImage.SR_32, TestImage.SR_64],
+    )
 
 
 def test_Compact_community(snapshot):
@@ -20,3 +25,8 @@ def test_Compact_community(snapshot):
     model = ModelLoader().load_from_file(file.path)
     assert model == snapshot(exclude=disallowed_props)
     assert isinstance(model.model, SRVGGNetCompact)
+    assert_image_inference(
+        file,
+        model,
+        [TestImage.SR_16, TestImage.SR_32, TestImage.SR_64],
+    )

tests/test_ESRGAN.py

@@ -1,7 +1,7 @@
 from spandrel import ModelLoader
 from spandrel.architectures.ESRGAN import RRDBNet
 
-from .util import ModelFile, disallowed_props
+from .util import ModelFile, TestImage, assert_image_inference, disallowed_props
 
 
 def test_ESRGAN_community(snapshot):
@@ -11,6 +11,11 @@ def test_ESRGAN_community(snapshot):
     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_16, TestImage.SR_32, TestImage.SR_64],
+    )
 
 
 def test_BSRGAN(snapshot):
@@ -20,6 +25,11 @@ def test_BSRGAN(snapshot):
     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_16, TestImage.SR_32, TestImage.SR_64],
+    )
 
 
 def test_BSRGAN_2x(snapshot):
@@ -29,6 +39,11 @@ def test_BSRGAN_2x(snapshot):
     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_16, TestImage.SR_32, TestImage.SR_64],
+    )
 
 
 def test_RealSR_DPED(snapshot):
@@ -38,6 +53,11 @@ def test_RealSR_DPED(snapshot):
     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_16, TestImage.SR_32, TestImage.SR_64],
+    )
 
 
 def test_RealSR_JPEG(snapshot):
@@ -47,6 +67,11 @@ def test_RealSR_JPEG(snapshot):
     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_16, TestImage.SR_32, TestImage.SR_64],
+    )
 
 
 def test_RealESRGAN_x4plus(snapshot):
@@ -56,6 +81,11 @@ def test_RealESRGAN_x4plus(snapshot):
     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_16, TestImage.SR_32, TestImage.SR_64],
+    )
 
 
 def test_RealESRGAN_x2plus(snapshot):
@@ -65,6 +95,11 @@ def test_RealESRGAN_x2plus(snapshot):
     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_16, TestImage.SR_32, TestImage.SR_64],
+    )
 
 
 def test_RealESRGAN_x4plus_anime_6B(snapshot):
@@ -74,6 +109,11 @@ def test_RealESRGAN_x4plus_anime_6B(snapshot):
     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_16, TestImage.SR_32, TestImage.SR_64],
+    )
 
 
 def test_RealESRNet_x4plus(snapshot):
@@ -83,3 +123,8 @@ def test_RealESRNet_x4plus(snapshot):
     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_16, TestImage.SR_32, TestImage.SR_64],
+    )

tests/test_HAT.py

@@ -1,7 +1,7 @@
 from spandrel import ModelLoader
 from spandrel.architectures.HAT import HAT
 
-from .util import ModelFile, disallowed_props
+from .util import ModelFile, TestImage, assert_image_inference, disallowed_props
 
 
 def test_HAT_community1(snapshot):
@@ -11,6 +11,11 @@ def test_HAT_community1(snapshot):
     model = ModelLoader().load_from_file(file.path)
     assert model == snapshot(exclude=disallowed_props)
     assert isinstance(model.model, HAT)
+    assert_image_inference(
+        file,
+        model,
+        [TestImage.SR_16, TestImage.SR_32, TestImage.SR_64],
+    )
 
 
 # TODO: We don't support HAT-S models yet

tests/test_OmniSR.py

@@ -1,7 +1,7 @@
 from spandrel import ModelLoader
 from spandrel.architectures.OmniSR import OmniSR
 
-from .util import ModelFile, disallowed_props
+from .util import ModelFile, TestImage, assert_image_inference, disallowed_props
 
 
 def test_OmniSR_community1(snapshot):
@@ -11,6 +11,11 @@ def test_OmniSR_community1(snapshot):
     model = ModelLoader().load_from_file(file.path)
     assert model == snapshot(exclude=disallowed_props)
     assert isinstance(model.model, OmniSR)
+    assert_image_inference(
+        file,
+        model,
+        [TestImage.SR_16, TestImage.SR_32, TestImage.SR_64],
+    )
 
 
 def test_OmniSR_community2(snapshot):
@@ -20,3 +25,8 @@ def test_OmniSR_community2(snapshot):
     model = ModelLoader().load_from_file(file.path)
     assert model == snapshot(exclude=disallowed_props)
     assert isinstance(model.model, OmniSR)
+    assert_image_inference(
+        file,
+        model,
+        [TestImage.SR_16, TestImage.SR_32, TestImage.SR_64],
+    )

tests/test_SwiftSRGAN.py

@@ -1,7 +1,7 @@
 from spandrel import ModelLoader
 from spandrel.architectures.SwiftSRGAN import SwiftSRGAN
 
-from .util import ModelFile, disallowed_props
+from .util import ModelFile, TestImage, assert_image_inference, disallowed_props
 
 
 def test_SwiftSRGan_2x(snapshot):
@@ -11,6 +11,11 @@ def test_SwiftSRGan_2x(snapshot):
     model = ModelLoader().load_from_file(file.path)
     assert model == snapshot(exclude=disallowed_props)
     assert isinstance(model.model, SwiftSRGAN)
+    assert_image_inference(
+        file,
+        model,
+        [TestImage.SR_16, TestImage.SR_32, TestImage.SR_64],
+    )
 
 
 def test_SwiftSRGan_4x(snapshot):
@@ -20,3 +25,8 @@ def test_SwiftSRGan_4x(snapshot):
     model = ModelLoader().load_from_file(file.path)
     assert model == snapshot(exclude=disallowed_props)
     assert isinstance(model.model, SwiftSRGAN)
+    assert_image_inference(
+        file,
+        model,
+        [TestImage.SR_16, TestImage.SR_32, TestImage.SR_64],
+    )

tests/test_Swin2SR.py

@@ -1,7 +1,7 @@
 from spandrel import ModelLoader
 from spandrel.architectures.Swin2SR import Swin2SR
 
-from .util import ModelFile, disallowed_props
+from .util import ModelFile, TestImage, assert_image_inference, disallowed_props
 
 
 def test_Swin2SR_4x(snapshot):
@@ -11,3 +11,8 @@ def test_Swin2SR_4x(snapshot):
     model = ModelLoader().load_from_file(file.path)
     assert model == snapshot(exclude=disallowed_props)
     assert isinstance(model.model, Swin2SR)
+    assert_image_inference(
+        file,
+        model,
+        [TestImage.SR_16, TestImage.SR_32, TestImage.SR_64],
+    )

tests/test_SwinIR.py

@@ -1,7 +1,7 @@
 from spandrel import ModelLoader
 from spandrel.architectures.SwinIR import SwinIR
 
-from .util import ModelFile, disallowed_props
+from .util import ModelFile, TestImage, assert_image_inference, disallowed_props
 
 
 def test_SwinIR_M_s64w8_2x(snapshot):
@@ -11,6 +11,11 @@ def test_SwinIR_M_s64w8_2x(snapshot):
     model = ModelLoader().load_from_file(file.path)
     assert model == snapshot(exclude=disallowed_props)
     assert isinstance(model.model, SwinIR)
+    assert_image_inference(
+        file,
+        model,
+        [TestImage.SR_16, TestImage.SR_32, TestImage.SR_64],
+    )
 
 
 def test_SwinIR_M_s48w8_4x(snapshot):
@@ -20,6 +25,11 @@ def test_SwinIR_M_s48w8_4x(snapshot):
     model = ModelLoader().load_from_file(file.path)
     assert model == snapshot(exclude=disallowed_props)
     assert isinstance(model.model, SwinIR)
+    assert_image_inference(
+        file,
+        model,
+        [TestImage.SR_16, TestImage.SR_32, TestImage.SR_64],
+    )
 
 
 def test_SwinIR_S_2x(snapshot):
@@ -29,6 +39,11 @@ def test_SwinIR_S_2x(snapshot):
     model = ModelLoader().load_from_file(file.path)
     assert model == snapshot(exclude=disallowed_props)
     assert isinstance(model.model, SwinIR)
+    assert_image_inference(
+        file,
+        model,
+        [TestImage.SR_16, TestImage.SR_32, TestImage.SR_64],
+    )
 
 
 def test_SwinIR_L_4x(snapshot):
@@ -38,3 +53,8 @@ def test_SwinIR_L_4x(snapshot):
     model = ModelLoader().load_from_file(file.path)
     assert model == snapshot(exclude=disallowed_props)
     assert isinstance(model.model, SwinIR)
+    assert_image_inference(
+        file,
+        model,
+        [TestImage.SR_16, TestImage.SR_32, TestImage.SR_64],
+    )