Flax: avoid recompilation when params change (#1096)
* Do not recompile when guidance_scale changes.
* Remove debug for simplicity.
* make style
* Make guidance_scale an array.
* Make DEBUG a constant to avoid passing it down.
* Add comments for clarification.
