google-research / maxim

[CVPR 2022 Oral] Official repository for "MAXIM: Multi-Axis MLP for Image Processing". SOTA for denoising, deblurring, deraining, dehazing, and enhancement.
https://arxiv.org/abs/2201.02973
Apache License 2.0
977 stars 105 forks source link

Support jax to tfjs conversion with dynamic models #41

Open thekevinscott opened 1 year ago

thekevinscott commented 1 year ago

I'm converting the MAXIM models into Tensorflow.js models for use in Javascript.

When using fixed sizes, the model converts fine. However, using dynamic sizes (polymorphism) fails. For a discussion of why, see here.

This PR adjusts the UpSampleRatio block to avoid the casting to int, which allows dynamically sized inputs (although inputs must still be a multiple of 64, which must be enforced outside of the model call).