Flax: Fix img2img and align with other pipeline (#1824)
* Flax: Add components function
* Flax: Fix img2img and align with other pipeline
* Flax: Fix PRNGKey type
* Refactor strength to start_timestep
* Fix preprocess images
* Fix processed_images dimen
* latents.shape -> latents_shape
* Fix typo
* Remove "static" comment
* Remove unnecessary optional types in _generate
* Apply doc-builder code style.
Co-authored-by: Pedro Cuenca <pedro@huggingface.co>
