From 6d80ba25371cff2cdd0e295b0049709f4ce12ba4 Mon Sep 17 00:00:00 2001
From: Joshua Lochner <admin@xenova.com>
Date: Wed, 8 Nov 2023 06:14:43 +0200
Subject: [PATCH] Add image-to-image unit tests

---
 tests/pipelines.test.js | 38 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/tests/pipelines.test.js b/tests/pipelines.test.js
index 26a905779..a3dcc0301 100644
--- a/tests/pipelines.test.js
+++ b/tests/pipelines.test.js
@@ -1326,6 +1326,44 @@ describe('Pipelines', () => {
         }, MAX_TEST_EXECUTION_TIME);
     });
 
+    describe('Image-to-image', () => {
+
+        // List all models which will be tested
+        const models = [
+            'caidas/swin2SR-classical-sr-x2-64',
+        ];
+
+        it(models[0], async () => {
+            let upscaler = await pipeline('image-to-image', m(models[0]));
+
+            // Input is 3x3 => padded to 8x8 => upscaled to 16x16
+            let url = 'https://huggingface.co/datasets/Xenova/transformers.js-docs/resolve/main/pattern_3x3.png';
+
+            // single
+            {
+                let outputs = await upscaler(url);
+                expect(outputs.width).toEqual(16);
+                expect(outputs.height).toEqual(16);
+                expect(outputs.channels).toEqual(3);
+                expect(outputs.data).toHaveLength(768);
+            }
+
+            // batched
+            {
+                let outputs = await upscaler([url, url]);
+                expect(outputs).toHaveLength(2);
+                for (let output of outputs) {
+                    expect(output.width).toEqual(16);
+                    expect(output.height).toEqual(16);
+                    expect(output.channels).toEqual(3);
+                    expect(output.data).toHaveLength(768);
+                }
+            }
+
+            await upscaler.dispose();
+        }, MAX_TEST_EXECUTION_TIME);
+    });
+
     describe('Document question answering', () => {
 
         // List all models which will be tested