More inference tests

libs/spandrel/spandrel/architectures/KBNet/__init__.py

@@ -91,7 +91,7 @@ def _load_l(self, state_dict: StateDict):
             state_dict,
             architecture=self,
             purpose="Restoration",
-            tags=["L"],
+            tags=["L", f"{dim}dim"],
             supports_half=False,
             supports_bfloat16=True,
             scale=1,
@@ -117,12 +117,12 @@ def _load_s(self, state_dict: StateDict):
         enc_count = get_seq_len(state_dict, "encoders")
         enc_blk_nums = [1] * enc_count
         for i in range(enc_count):
-            enc_blk_nums[i] = get_seq_len(state_dict, "encoders." + str(i))
+            enc_blk_nums[i] = get_seq_len(state_dict, f"encoders.{i}")
 
         dec_count = get_seq_len(state_dict, "decoders")
         dec_blk_nums = [1] * dec_count
         for i in range(dec_count):
-            dec_blk_nums[i] = get_seq_len(state_dict, "decoders." + str(i))
+            dec_blk_nums[i] = get_seq_len(state_dict, f"decoders.{i}")
 
         # in code: ffn_ch = int(c * ffn_scale)
         temp_c = state_dict["middle_blks.0.conv4.weight"].shape[1]
@@ -148,7 +148,11 @@ def _load_s(self, state_dict: StateDict):
             state_dict,
             architecture=self,
             purpose="Restoration",
-            tags=["S"],
+            tags=[
+                "S",
+                f"{width}w",
+                *(["lightweight"] if lightweight else []),
+            ],
             supports_half=False,
             supports_bfloat16=True,
             scale=1,

tests/__snapshots__/test_KBNet.ambr

@@ -0,0 +1,42 @@
+# serializer version: 1
+# name: test_derain
+  ImageModelDescriptor(
+    architecture=KBNetArch(
+      id='KBNet',
+      name='KBNet',
+    ),
+    input_channels=3,
+    output_channels=3,
+    purpose='Restoration',
+    scale=1,
+    size_requirements=SizeRequirements(minimum=0, multiple_of=16, square=False),
+    supports_bfloat16=True,
+    supports_half=False,
+    tags=list([
+      'L',
+      '48dim',
+    ]),
+    tiling=<ModelTiling.SUPPORTED: 1>,
+  )
+# ---
+# name: test_sidd
+  ImageModelDescriptor(
+    architecture=KBNetArch(
+      id='KBNet',
+      name='KBNet',
+    ),
+    input_channels=3,
+    output_channels=3,
+    purpose='Restoration',
+    scale=1,
+    size_requirements=SizeRequirements(minimum=0, multiple_of=1, square=False),
+    supports_bfloat16=True,
+    supports_half=False,
+    tags=list([
+      'S',
+      '64w',
+      'lightweight',
+    ]),
+    tiling=<ModelTiling.SUPPORTED: 1>,
+  )
+# ---

tests/test_FFTformer.py

@@ -2,6 +2,8 @@
 
 from .util import (
     ModelFile,
+    TestImage,
+    assert_image_inference,
     assert_loads_correctly,
     assert_size_requirements,
     disallowed_props,
@@ -38,3 +40,4 @@ def test_fftformer_GoPro(snapshot):
     model = file.load_model()
     assert model == snapshot(exclude=disallowed_props)
     assert isinstance(model.model, FFTformer)
+    assert_image_inference(file, model, [TestImage.SR_32])

tests/test_KBNet.py

@@ -1,6 +1,14 @@
 from spandrel.architectures.KBNet import KBNet_l, KBNet_s, KBNetArch
 
-from .util import assert_loads_correctly, skip_if_unchanged
+from .util import (
+    ModelFile,
+    TestImage,
+    assert_image_inference,
+    assert_loads_correctly,
+    assert_size_requirements,
+    disallowed_props,
+    skip_if_unchanged,
+)
 
 skip_if_unchanged(__file__)
 
@@ -27,3 +35,35 @@ def test_load():
         lambda: KBNet_s(width=32, ffn_scale=3),
         lambda: KBNet_s(width=32, lightweight=True),
     )
+
+
+def test_size_requirements():
+    file = ModelFile.from_url(
+        "https://github.com/OpenModelDB/model-hub/releases/download/kbnet/1x-KBNet_derain.pth",
+    )
+    assert_size_requirements(file.load_model())
+
+    file = ModelFile.from_url(
+        "https://github.com/OpenModelDB/model-hub/releases/download/kbnet/1x-KBNet_sidd.pth",
+    )
+    assert_size_requirements(file.load_model())
+
+
+def test_derain(snapshot):
+    file = ModelFile.from_url(
+        "https://github.com/OpenModelDB/model-hub/releases/download/kbnet/1x-KBNet_derain.pth",
+    )
+    model = file.load_model()
+    assert model == snapshot(exclude=disallowed_props)
+    assert isinstance(model.model, KBNet_l)
+    assert_image_inference(file, model, [TestImage.SR_32])
+
+
+def test_sidd(snapshot):
+    file = ModelFile.from_url(
+        "https://github.com/OpenModelDB/model-hub/releases/download/kbnet/1x-KBNet_sidd.pth",
+    )
+    model = file.load_model()
+    assert model == snapshot(exclude=disallowed_props)
+    assert isinstance(model.model, KBNet_s)
+    assert_image_inference(file, model, [TestImage.JPEG_15])

tests/test_Restormer.py

@@ -2,6 +2,8 @@
 
 from .util import (
     ModelFile,
+    TestImage,
+    assert_image_inference,
     assert_loads_correctly,
     assert_size_requirements,
     disallowed_props,
@@ -82,6 +84,7 @@ def test_gaussian_color_denoising_blind(snapshot):
     model = file.load_model()
     assert model == snapshot(exclude=disallowed_props)
     assert isinstance(model.model, Restormer)
+    assert_image_inference(file, model, [TestImage.JPEG_15])
 
 
 def test_gaussian_gray_denoising_sigma25(snapshot):

tests/test_SCUNet.py

@@ -2,6 +2,8 @@
 
 from .util import (
     ModelFile,
+    TestImage,
+    assert_image_inference,
     assert_loads_correctly,
     assert_size_requirements,
     disallowed_props,
@@ -43,6 +45,7 @@ def test_SCUNet_color_GAN(snapshot):
     model = file.load_model()
     assert model == snapshot(exclude=disallowed_props)
     assert isinstance(model.model, SCUNet)
+    assert_image_inference(file, model, [TestImage.JPEG_15])
 
 
 def test_SCUNet_color_RSNR(snapshot):