[PixArt-Alpha] Introduce resolution binning

[sayakpaul]

What does this PR do?

Introduces resolution binning as originally done by @lawrence-cj in #5716.

This is particularly useful for generating non-square images.

Quality comparison

1024x2048

No Binning	Binning

1536x720

No Binning	Binning

Code

import torch

from diffusers import PixArtAlphaPipeline

prompts = [
    "A small cactus with a happy face in the Sahara desert.",
    "Pirate ship trapped in a cosmic maelstrom nebula, rendered in cosmic beach whirlpool engine, volumetric lighting, spectacular, ambient lights, light pollution, cinematic atmosphere, art nouveau style, illustration art artwork by SenseiJaye, intricate detail.",
    "stars, water, brilliantly, gorgeous large scale scene, a little girl, in the style of dreamy realism, light gold and amber, blue and pink, brilliantly illuminated in the background.",
    "nature vs human nature, surreal, UHD, 8k, hyper details, rich colors, photograph.",
]

pipe = PixArtAlphaPipeline.from_pretrained(
    "PixArt-alpha/PixArt-XL-2-1024-MS", torch_dtype=torch.float16
)
pipe.enable_model_cpu_offload()

no_resolution_binning = []
resolution_binning = []

for i, prompt in enumerate(prompts):
    generator = torch.manual_seed(2344) 
    image = pipe(prompt, height=1536, width=720, generator=generator, use_resolution_binning=False).images[0]
    image.save(f"no_resolution_binning_{i}_1536x720.png")
    no_resolution_binning.append(image)
    
    generator = torch.manual_seed(2344)
    image = pipe(prompt, height=1536, width=720, generator=generator).images[0]
    image.save(f"resolution_binning_{i}_1536x720.png")
    resolution_binning.append(image)

[sayakpaul]

@DN6 we're introducing a new dependency here - torchvision. Could you please help reviewing if I added all the necessary guardrails to not break anything?

[lawrence-cj]

No problem on my side.

[patrickvonplaten]

Very nice addition! Let's add a test as well and please no torchvision dependency 🙏

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[sayakpaul]

Eliminated torchvision as a dependency and the results look fine:

1024x2048

No Binning	Binning

1536x720

No Binning	Binning

[sayakpaul]

@patrickvonplaten I have added a pure PyTorch implementation.

The reason why we cannot do a fast test for resolution binning is that even for very resolutions, classify_height_width_bin() would return extremely high resolutions to run on a CPU. For example, for the (32, 48) resolution, the resolution will be mapped to (832, 1216) which is quite high for a CPU.

This is why I choose to add a slow test for it and include that in #5752 (as it requires changing the assertion values anyway). Is that okay?

[patrickvonplaten]

Clean!