Flax pipeline: width and height are ignored

[pcuenca]

Describe the bug

width and height are passed to the generate function, and we do check that they are multiples of 8, but then we always create the latents using the standard size.

Reproduction

images = pipeline(prompt_ids, p_params, rng, width=448, height=832, jit=True)[0]
images.shape
# output: (8, 1, 512, 512, 3)

Logs

No response

System Info

Tested on TPU v3-8 with diffusers @ e48ca0f

[NightMachinery]

Does the model work with a custom resolution? Isn't it only trained on 512x512?

How high can we go on Colab's TPUs?

[patrickvonplaten]

Haven't tested the max resolution that's possible yet - would also be interested in such numbers :-)

[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.