Add new Maxim models

dev/browser/package.json

@@ -12,6 +12,14 @@
     "@upscalerjs/esrgan-slim": "workspace:*",
     "@upscalerjs/esrgan-thick": "workspace:*",
     "@upscalerjs/pixel-upsampler": "workspace:*",
+    "@upscalerjs/maxim-experiments": "workspace:*",
+    "@upscalerjs/maxim-deblurring": "workspace:*",
+    "@upscalerjs/maxim-deraining": "workspace:*",
+    "@upscalerjs/maxim-denoising": "workspace:*",
+    "@upscalerjs/maxim-dehazing-indoor": "workspace:*",
+    "@upscalerjs/maxim-dehazing-outdoor": "workspace:*",
+    "@upscalerjs/maxim-enhancement": "workspace:*",
+    "@upscalerjs/maxim-retouching": "workspace:*",
     "seedrandom": "3.0.5",
     "upscaler": "workspace:*",
     "vite": "^4.4.9"

dev/browser/specific-model/index.ts

@@ -54,23 +54,26 @@ const loadModel = async (packageName: string, modelName: string) => {
 
   const img = await makeImg(fixture, `Original: ${packageName}/${modelName}`, 1);
   const modelPath = getModelPath(packageName, modelJSON.path);
+  console.log(modelJSON)
   const upscaledImg = await upscaleImage({
     ...modelJSON,
     path: modelPath,
-  }, img);
+  }, img, 64, 2);
   await makeImg(upscaledImg, `Upscaled: ${packageName}/${modelName}`);
 };
 
 
-const upscaleImage = async (model: ModelDefinition, img: HTMLImageElement | HTMLCanvasElement, patchSize: undefined | number = 64) => {
+const upscaleImage = async (model: ModelDefinition, img: HTMLImageElement | HTMLCanvasElement, patchSize?: number, padding?: number) => {
   status.innerHTML = 'Starting';
   const upscaler = new Upscaler({
     model,
   });
   status.innerHTML = 'Upscaling...';
   const start = performance.now();
+  console.log({ patchSize, padding })
   const upscaledImg = await upscaler.upscale(img, {
     patchSize,
+    padding,
     progress: console.log,
   });
   console.log(`Duration: ${((performance.now() - start) / 1000).toFixed(2)}s`);

docs/blog/2023-12-01-maxim-models.md

@@ -0,0 +1,152 @@
+---
+title: MAXIM models
+description: New models based on the MAXIM architecture released
+slug: maxim-models
+authors: kscott
+tags: [announcement]
+hide_table_of_contents: false
+---
+
+I've released Tensorflow.js ports of the [MAXIM family of models](https://arxiv.org/abs/2201.02973), supporting the ability to deblur, denoise, derain, dehaze, retouch, and low-light enhance images.
+
+<!--truncate-->
+
+## Motivation
+
+UpscalerJS was originally created in 2020 with a primary goal of upscaling images. So far, all the models have been exclusively focused on super resolution.
+
+Today I'm releasing a new family of MAXIM models for UpscalerJS that enable a variety of new image enhancement techniques, including:
+
+- [Deblurring](https://upscalerjs.com/models/available/maxim-deblurring)
+- [Denoising](https://upscalerjs.com/models/available/maxim-denoising)
+- [Deraining](https://upscalerjs.com/models/available/maxim-deraining)
+- Dehazing (both [indoor](https://upscalerjs.com/models/available/maxim-dehazing-indoor) and [outdoor](https://upscalerjs.com/models/available/maxim-dehazing-outdoor))
+- [Low Light Enhancement](https://upscalerjs.com/models/available/maxim-enhancement)
+- [Retouching](https://upscalerjs.com/models/available/maxim-retouching)
+
+These models are available in Javascript via UpscalerJS, and run in both the browser and Node.js.
+
+## MAXIM
+
+MAXIM is the architecture at the heart of these new models.
+
+The MAXIM paper ([MAXIM: Multi-Axis MLP for Image Processing](https://arxiv.org/abs/2201.02973)) was published in 2022, and was nominated as one of the best papers at CVPR 2022. The MAXIM architecture is capable, via training, of supporting a variety of image enhancement tasks. It's also an efficient architecture, making it particularly appealing for JavaScript applications.
+
+> We present a multi-axis MLP based architecture called MAXIM, that can serve as an efficient and flexible general-purpose vision backbone for image processing tasks ... Our extensive experimental results show that the proposed MAXIM model achieves state-of-the-art performance on more than ten benchmarks across a range of image processing tasks, including denoising, deblurring, deraining, dehazing, and enhancement while requiring fewer or comparable numbers of parameters and FLOPs than competitive models. 
+
+Google Research [published an implementation in Jax](https://github.com/google-research/maxim), and additional ports were made available in [Tensorflow](https://github.com/sayakpaul/maxim-tf) and [Pytorch](https://github.com/vztu/maxim-pytorch/tree/main/maxim_pytorch).
+
+## Getting Started
+
+You can run MAXIM models in the browser or Node.js. To get started, install your desired model:
+
+```bash
+npm install @upscalerjs/maxim-deblurring
+```
+
+And provide the model as an argument to UpscalerJS:
+
+```javascript
+import model from '@upscalerjs/maxim-deblurring'
+const upscaler = new Upscaler({
+  model,
+})
+```
+
+For model-specific instructions, [check out the specific model page](/models) you're interested in.
+
+## Samples
+
+Below are some samples of each image enhancement task:
+
+### Deblurring
+
+*Before*
+
+![Deblurring Before](/assets/sample-images/maxim-deblurring/fixture.png)
+
+*After*
+
+![Deblurring After](/assets/sample-images/maxim-deblurring/result.png)
+
+### Denoising
+
+*Before*
+
+![Denoising Before](/assets/sample-images/maxim-denoising/fixture.png)
+
+*After*
+
+![Denoising After](/assets/sample-images/maxim-denoising/result.png)
+
+### Deraining
+
+*Before*
+
+![Deraining Before](/assets/sample-images/maxim-deraining/fixture.png)
+
+*After*
+
+![Deraining After](/assets/sample-images/maxim-deraining/result.png)
+
+### Low Light Enhancement
+
+*Before*
+
+![Low Light Enhancement Before](/assets/sample-images/maxim-enhancement/fixture.png)
+
+*After*
+
+![Low Light Enhancement After](/assets/sample-images/maxim-enhancement/result.png)
+
+### Retouching
+
+*Before*
+
+![Retouching Before](/assets/sample-images/maxim-retouching/fixture.png)
+
+*After*
+
+![Retouching After](/assets/sample-images/maxim-retouching/result.png)
+
+### Dehazing Indoor
+
+*Before*
+
+![Dehazing Indoor Before](/assets/sample-images/maxim-dehazing-indoor/fixture.png)
+
+*After*
+
+![Dehazing Indoor After](/assets/sample-images/maxim-dehazing-indoor/result.png)
+
+### Dehazing Outdoor
+
+*Before*
+
+![Dehazing Outdoor Before](/assets/sample-images/maxim-dehazing-outdoor/fixture.png)
+
+*After*
+
+![Dehazing Outdoor After](/assets/sample-images/maxim-dehazing-outdoor/result.png)
+
+## Technical Information
+
+My original attempts at getting this working leveraged both the [Jax repo](https://github.com/google-research/maxim/) as well as the [Tensorflow port](https://github.com/sayakpaul/maxim-tf/tree/main). The ported Jax model exhibited close-to-identical fidelity with its Python implementation, but the Tensorflow port was far more performant in the browser, albiet with noticeably inferior fidelity.
+
+Both implementations originally required fixed size inputs in order to port to Tensorflow.js. Fixed size inputs require chunking images, which can break models that rely on a holisitic view of the image, specifically the Dehazing models, Enhancement model, and Retouching model. (For samples of the artifacting this produces and a longer discussion, [see this Github issue](https://github.com/thekevinscott/UpscalerJS/issues/913).)
+
+I modified the MAXIM Jax code to support dynamic image input sizes, [and opened a PR here](https://github.com/google-research/maxim/pull/41). [This PR is also integrated in my fork of the MAXIM code](https://github.com/thekevinscott/maxim).
+
+(There is also an [open PR against the Tensorflow repo exploring dynamic sizes](https://github.com/sayakpaul/maxim-tf/pull/24); when it gets merged, I'll explore porting it to Tensorflow.js as well.) 
+
+If you'd like to check out the Tensorflow implementation port (noticeably faster, noticeably worse inference, and a fixed size input) these models are available under the `maxim-experiments` repo in the UpscalerJS repo. Clone the repo, pull the models (`dvc pull`) and you'll see fixed-input models of `64` and `256` pixel sizes respectively for each task.
+
+## Converting the files yourself
+
+If you'd like to convert the original Jax or Tensorflow model files yourself, instructions are [in these Jupyter notebooks](https://github.com/upscalerjs/maxim). Feel free to open an issue on Github if you run into issues or questions.
+
+----
+
+The past few years have seen an explosion of innovation in the image enhancement space, and I hope to continue bringing the latest innovations to Javascript. MAXIM is a first step towards enabling Javascript-based image enhancement tasks beyond super resolution in UpscalerJS.
+
+If you have particular models you'd like to see available via UpscalerJS, [let me know in Github](https://github.com/thekevinscott/UpscalerJS/discussions/new?category=ideas). If you use MAXIM in your work, [I'd love to hear about it](https://github.com/thekevinscott/UpscalerJS/discussions/new?category=show-and-tell)! And if you run into questions or find bugs, [please open a bug report](https://github.com/thekevinscott/UpscalerJS/issues/new?assignees=thekevinscott&labels=&projects=&template=bug_report.md&title=).

docs/docs/documentation/troubleshooting.md

@@ -65,27 +65,25 @@ This likely means one of two things:
 - You are using `upscaler`, instead of `upscaler/node`; [check out the guide on Node.js here](/documentation/guides/node/nodejs).
 - You are using `import` syntax instead of `require` syntax; if so, try switching to `require('upscaler')`. For more information on this, [see this Github issue](https://github.com/thekevinscott/UpscalerJS/issues/554#issuecomment-1344108954).
 
-## Missing Model Scale
+## Missing Model Path
 
 If you see an error like:
 
 ```
-Error: You must provide a "scale" when providing a model definition
+Error: You must provide a "path" when providing a model definition
 ```
 
-You've passed a `null` or `undefined` scale argument in the `model` argument to UpscalerJS:
+You've passed a `null` or `undefined` path argument in the `model` argument to UpscalerJS:
 
 ```javascript
 const upscaler = new Upscaler({
   model: {
-    scale: null,
+    path: null,
   },
 })
 ```
 
-Every model must have an explicit `scale` defined.
-
-Ensure you pass a valid `scale` argument in the `model` payload. [See the guide on models for more information](/documentation/guides/browser/models).
+Ensure you pass a valid `path` argument in the `model` payload. [See the guide on models for more information](/documentation/guides/browser/models).
 
 ## Invalid Warmup Value
 

docs/docusaurus.config.js

@@ -179,6 +179,30 @@ const config = {
                 label: 'Super Resolution',
                 to: '/models#upscaling',
               },
+              {
+                label: 'Deblurring',
+                to: '/models#deblurring',
+              },
+              {
+                label: 'Denoising',
+                to: '/models#denoising',
+              },
+              {
+                label: 'Deraining',
+                to: '/models#deraining',
+              },
+              {
+                label: 'Dehazing',
+                to: '/models#dehazing',
+              },
+              {
+                label: 'Low Light Enhancement',
+                to: '/models#low-light-enhancement',
+              },
+              {
+                label: 'Retouching',
+                to: '/models#retouching',
+              },
             ],
           },
         ],

docs/src/components/homepage/header/demo-video/images.ts

@@ -1,5 +1,26 @@
 import { IExampleImage } from "./example-images/example-images";
 
+const makeMAXIM = (model: string, action: string): IExampleImage => ({
+  sizes: {
+    original: 256,
+    enhanced: 256,
+  },
+  original: {
+    src: `/assets/homepage-demo/originals/${model}.png`,
+    labels: {
+      short: 'Original',
+      long: 'Original',
+    },
+  },
+  enhanced: {
+    src: `/assets/homepage-demo/enhanced/${model}.png`,
+    labels: {
+      short: `@upscalerjs/maxim-${model}`,
+      long: `${action} using [\`@upscalerjs/maxim-${model}\`](/models/available/maxim-${model}) model`,
+    },
+  },
+});
+
 const makeESRGAN = (image: string): IExampleImage => ({
   sizes: {
     original: 128,
@@ -22,6 +43,9 @@ const makeESRGAN = (image: string): IExampleImage => ({
 });
 
 const images: Record<string, IExampleImage> = {
+  maximDenoising: makeMAXIM('denoising', 'Denoised'),
+  maximEnhancement: makeMAXIM('enhancement', 'Enhanced'),
+  maximDeblurring: makeMAXIM('deblurring', 'Deblurred'),
   esrganFlower: makeESRGAN('flower'),
   esrganFace3: makeESRGAN('face3'),
   esrganFace2: makeESRGAN('face2'),
@@ -31,6 +55,9 @@ const images: Record<string, IExampleImage> = {
 export const IMAGES: IExampleImage[] = [
   images.esrganFlower,
   images.esrganFace2,
+  images.maximEnhancement,
   images.esrganFace3,
+  images.maximDeblurring,
   images.esrganFace1,
+  images.maximDenoising,
 ].filter(Boolean);

docs/src/components/modelCards/modelCard/modelCard.module.scss

@@ -52,6 +52,8 @@
   img {
     flex: 1;
     width: 50%;
+    aspect-ratio: 1 / 1;
+    object-fit: cover;
   }
 }
 

docs/src/theme/DocItem/Layout/index.js

@@ -65,7 +65,7 @@ export default function DocItemLayout({children}) {
                 {docTOC.mobile}
                 <DocItemContent>
                   {children}
-                  </DocItemContent>
+                </DocItemContent>
                 <DocItemFooter />
               </article>
               <DocItemPaginator />

internals/common/src/models.ts

@@ -137,28 +137,38 @@ interface AvailableModel {
   value: string | PackageJSONExport;
 }
 
-export const getSupportedPlatforms = async (packageName: string, modelName: string): Promise<Environment[]> => {
+export const getSupportedPlatforms = async (packageName: string, modelName: string, key = 'supportedPlatforms'): Promise<Environment[]> => {
   if (!packageName) {
     throw new Error('Missing package name')
   }
   const packageJSONPath = path.resolve(MODELS_DIR, packageName);
   const packageJSON = await getPackageJSON(packageJSONPath);
-  return packageJSON['@upscalerjs']?.models?.[modelName]?.supportedPlatforms || ['clientside', 'serverside'];
+  const supportedPlatforms: undefined | string[] = packageJSON['@upscalerjs']?.models?.[modelName]?.[key];
+  if (supportedPlatforms === undefined) {
+    return ['clientside', 'serverside'];
+  }
+  return Array.from(new Set(supportedPlatforms.map(platform => {
+    if (['node', 'node-gpu'].includes(platform)) {
+      return 'serverside';
+    }
+    return 'clientside';
+  })));
 };
 
-const getPackagesAndModelsMatchingEnvironment = async (environment: Environment, packagesAndModels: ModelInformation[]) => {
+const getPackagesAndModelsMatchingEnvironment = async (environment: Environment, packagesAndModels: ModelInformation[], CI = false) => {
   const filteredPackagesAndModels: ModelInformation[] = [];
   await Promise.all(packagesAndModels.map(async (modelInformation) => {
     const supportedPlatforms = await getSupportedPlatforms(modelInformation.packageDirectoryName, modelInformation.modelName);
-    if (supportedPlatforms.includes(environment)) {
+    const supportedCIPlatforms = await getSupportedPlatforms(modelInformation.packageDirectoryName, modelInformation.modelName, 'CI');
+    if (supportedPlatforms.includes(environment) && (CI === false || supportedCIPlatforms.includes(environment))) {
       filteredPackagesAndModels.push(modelInformation);
     }
   }));
   return filteredPackagesAndModels;
 };
 
 
-export const getPackagesAndModelsForEnvironment = async (environment: Environment): Promise<ModelInformation[]> => {
+export const getPackagesAndModelsForEnvironment = async (environment: Environment, CI = false): Promise<ModelInformation[]> => {
   const packagesAndModels = await ALL_MODELS;
-  return getPackagesAndModelsMatchingEnvironment(environment, packagesAndModels);
+  return getPackagesAndModelsMatchingEnvironment(environment, packagesAndModels, CI);
 };

models/esrgan-medium/package.json

@@ -96,7 +96,8 @@
     "models": {
       "./8x": {
         "supportedPlatforms": [
-          "node"
+          "node",
+          "node-gpu"
         ]
       }
     },

models/esrgan-slim/package.json

@@ -96,7 +96,8 @@
     "models": {
       "./8x": {
         "supportedPlatforms": [
-          "node"
+          "node",
+          "node-gpu"
         ]
       }
     },