[WIP] Make FlaxLMSDiscreteScheduler jittable

[Dan8991]

PR for issue #2180 trying to make FlaxLMSDiscreteScheduler jittable

[pcuenca]

Hi @Dan8991, let us know when you are ready for a review :)

[Dan8991]

Yes @pcuenca thank you very much. Right now I am just doing some testing to better understand how the modules work so the commits contain useless files, I will squash everything when I have finished making the class jittable and I will ask you to review the code :).

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.

[Dan8991]

@pcuenca I think I am almost done with the pull request, however I have some questions in particular:

1. Since, for example, in pipeline_flax_stable_diffusion.py the loop_body function is called using jax.lax.fori_loop its input and outputs must have the same shape. However this was not the case in the previous implementation of the code since the "latents" variable contains an array ("derivatives") whose length increases over time up to size "order". In the previous code version the derivatives array was initialized with length zero however this raises an error when its size increases since the inputs and outputs of the function don't have matching shapes. A.t.m I have hardcoded the shape to the correct value and I was wondering if it would be possible to add the order argument to the set_timesteps function to be able to avoid this problem.
2. I believe that there were some bugs in the original implementation of FlaxLMSDiscreteSchedulerState in particular:

• in the step function in the append and replace functions the axis was not specified which according to the jax documentation leads to the arrays being flattened (which is unwanted in this case) so I added the axis values.
• in the step function "sigma" and "order" were computed using values from the timestep array, however in the non flax implementation the two variables depend on the index of the considered timestep so I changed it to match the non-flax implementation (which actually made the function work).
  I am not sure if I should open an issue for these bugs instead of solving them in this PR, however they are technically necessary to make FlaxLMSDiscreteScheduler jittable so probably it makes sense to solve them here, no?

3. I have noticed that whenever the generated image contains NANs the script tells the user that the image was obscured because it might have presented NSFW content, I find this error message to be non informative whenever it is caused by NANs, should I open an issue for this?
   However I am almost ready for a review, I just need to know what would be your preferred way to fix the non-matching input and output sizes that I mentioned in point 1.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[Dan8991]

@pcuenca I think I am ready for revision, however please check out my last message before starting the review and let me know if my assumptions were correct.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[patrickvonplaten]

@pcuenca do you think it makes sense to continue this PR?

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.