Add official tiny model to tests

chaiNNer-org · May 11, 2024 · 0632677 · 0632677
1 parent efa4585
commit 0632677
Show file tree

Hide file tree

Showing 5 changed files with 37 additions and 0 deletions.
diff --git a/tests/__snapshots__/test_PLKSR.ambr b/tests/__snapshots__/test_PLKSR.ambr
@@ -1,4 +1,26 @@
 # serializer version: 1
+# name: test_PLKSR_official_tiny_x4
+  ImageModelDescriptor(
+    architecture=PLKSRArch(
+      id='PLKSR',
+      name='PLKSR',
+    ),
+    input_channels=3,
+    output_channels=3,
+    purpose='SR',
+    scale=4,
+    size_requirements=SizeRequirements(minimum=1, multiple_of=1, square=False),
+    supports_bfloat16=True,
+    supports_half=False,
+    tags=list([
+      '64dim',
+      '28nb',
+      '13ks',
+      'DCCM',
+    ]),
+    tiling=<ModelTiling.SUPPORTED: 1>,
+  )
+# ---
 # name: test_PLKSR_official_x2
   ImageModelDescriptor(
     architecture=PLKSRArch(

diff --git a/tests/images/outputs/16x16/PLKSR_tiny_X4_DIV2K.png b/tests/images/outputs/16x16/PLKSR_tiny_X4_DIV2K.png
diff --git a/tests/images/outputs/32x32/PLKSR_tiny_X4_DIV2K.png b/tests/images/outputs/32x32/PLKSR_tiny_X4_DIV2K.png
diff --git a/tests/images/outputs/64x64/PLKSR_tiny_X4_DIV2K.png b/tests/images/outputs/64x64/PLKSR_tiny_X4_DIV2K.png
diff --git a/tests/test_PLKSR.py b/tests/test_PLKSR.py
@@ -75,6 +75,21 @@ def test_PLKSR_official_x4(snapshot):
     )
 
 
+def test_PLKSR_official_tiny_x4(snapshot):
+    file = ModelFile.from_url(
+        "https://drive.google.com/file/d/1d8d_6TF0SrEMiX1jrnLqKnPdKjDJahOK/view",
+        name="PLKSR_tiny_X4_DIV2K.pth",
+    )
+    model = file.load_model()
+    assert model == snapshot(exclude=disallowed_props)
+    assert isinstance(model.model, PLKSR)
+    assert_image_inference(
+        file,
+        model,
+        [TestImage.SR_16, TestImage.SR_32, TestImage.SR_64],
+    )
+
+
 def test_PLKSR_official_x3(snapshot):
     file = ModelFile.from_url(
         "https://drive.google.com/file/d/1W8phbKFTOYL-AnlMJnjx2NWDHZY8jVWW/view",