Issues with patch sizes and MAXIM models

[thekevinscott]

I'm working on integrating new MAXIM models in #842, but running into an issue that I'm not sure is solvable. Pinging @danwexler because I think you're working on the same thing and am curious if you have any feedback.

These MAXIM models have fixed input sizes (more on why below), which requires images be split into patches. For some models, this works fine. For others, it does not.

The models that seem stable when splitting into patches

Deblurring

Dehazing

Deraining

Dehazing Indoor

The models that exhibit artifacting when splitting into patches

Dehazing Outdoor

Low Light Enhancement

Retouching

I'm converting both the original JAX implementation along with the TF versions. (For some reason, the converted TF versions run much faster but perform worse than their JAX counterparts; I don't know why.)

Why the models require fixed input sizes

Both the JAX and TF versions of the models require fixed input sizes when converting to Tensorflow.js. I believe the reason is because both rely on einops functions, which do not support dynamic shapes.

In the case of the TF version, there's open PRs (here and here) to support arbitrary image sizes, but they're not merged yet.

On the JAX side, similar work would have to be done to deprecate the usage of einops.

Open questions

Is there any way to adjust these models to work across images subdivided into patches? I don't think there is.

So the models could be rewritten to support arbitrary image sizes, but that I think will preclude them from being usable in the browser. Which is maybe the best that can be done for now.

[thekevinscott]

This will be resolved in #842.

I was able to convert the Jax model to accept dynamic input shapes.

If and when the TF port's PR for supporting dynamic shapes is merged, I'll try and work on porting those to Tensorflow.js as well.

[danwexler]

I'm still waiting for full arbitrary resolution support in maxim, which seems to be progressing recently.

…

On Wed, Jun 28, 2023, 6:59 AM Kevin Scott ***@***.***> wrote: This will be resolved in #842 <#842>. I was able to convert the Jax model to accept dynamic input shapes. If and when the TF port's PR for supporting dynamic shapes is merged <sayakpaul/maxim-tf#24>, I'll try and work on porting those to Tensorflow.js as well. — Reply to this email directly, view it on GitHub <#913 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAEQE2PGRNOIM7MCKZAVBIDXNQ2FPANCNFSM6AAAAAAVPW6YUA> . You are receiving this because you were mentioned.Message ID: ***@***.***>