[DC-AE] support tiling for DC-AE

[chenjy2003]

What does this PR do?

This PR will add tiling support for DC-AE (Deep Compression Autoencoder for Efficient High-Resolution Diffusion Models) into the diffusers lib, in order to reduce memory consumption for encoding and decoding super-high-resolution images like 4096x4096.

Could you please have a look? Thanks! @sayakpaul @yiyixuxu
Cc: @lawrence-cj

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[elismasilva]

After this PR merged, It will be necessary to implement StableDiffusionMixIn inheritance in SANA pipelines in order to enable tiling and slicing.

@chenjy2003 I think you can implement this: See this signature class with StableDiffusionMixin

class SanaPipeline(DiffusionPipeline, StableDiffusionMixin, SanaLoraLoaderMixin):
    r"""
    Pipeline for text-to-image generation using [Sana](https://huggingface.co/papers/2410.10629).
    """

Ive tested locally with your changes and working.

pipe.enable_vae_slicing()
pipe.enable_vae_tiling()

[a-r-r-o-w]

Thanks @lawrence-cj! The changes look good.

As mentioned by the @elismasilva, we need to add some methods to the pipeline. Let's not derive from StableDiffusionMixin however, since it contains enable_freeu and disable_freeu.

Would just copy this into the pipeline:

    def enable_vae_slicing(self):
        r"""
        Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
        compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
        """
        self.vae.enable_slicing()

    def disable_vae_slicing(self):
        r"""
        Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to
        computing decoding in one step.
        """
        self.vae.disable_slicing()

    def enable_vae_tiling(self):
        r"""
        Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
        compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
        processing larger images.
        """
        self.vae.enable_tiling()

    def disable_vae_tiling(self):
        r"""
        Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to
        computing decoding in one step.
        """
        self.vae.disable_tiling()

[FurkanGozukara]

Dear @lawrence-cj please add this feature to your official SANA pipeline

[chenjy2003]

@elismasilva @a-r-r-o-w I've added the four methods enable_vae_slicing, disable_vae_slicing, enable_vae_tiling, and disable_vae_tiling to SanaPipeline.

[elismasilva]

@elismasilva @a-r-r-o-w I've added the four methods enable_vae_slicing, disable_vae_slicing, enable_vae_tiling, and disable_vae_tiling to SanaPipeline.

ok i think you forgot to do the same for this pipeline
https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/pag/pipeline_pag_sana.py

[chenjy2003]

@elismasilva @a-r-r-o-w I've added the four methods enable_vae_slicing, disable_vae_slicing, enable_vae_tiling, and disable_vae_tiling to SanaPipeline.

ok i think you forgot to do the same for this pipeline https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/pag/pipeline_pag_sana.py

Thanks for pointing out! @lawrence-cj Could you please double-check whether no other SANA pipelines need to be modified?

[lawrence-cj]

We have two pipelines in diffusers. I'm ok with it! Thanks @chenjy2003 for hard work!

[nitinmukesh]

Subscribing for it to be merged
NVlabs/Sana#138

[FurkanGozukara]

what is difference between pipeline_pag_sana.py and pipeline_sana.py

[chenjy2003]

what is difference between pipeline_pag_sana.py and pipeline_sana.py

@lawrence-cj Here is a question for you.

[nitinmukesh]

@FurkanGozukara @chenjy2003

It was answered here
NVlabs/Sana#107 (comment)

[FurkanGozukara]

@FurkanGozukara @chenjy2003

It was answered here NVlabs/Sana#107 (comment)

which one is recommended it doesnt tell much

[nitinmukesh]

@FurkanGozukara @chenjy2003

which one is recommended it doesnt tell much

Based on my test results PAG quality is 60-70% lower. The question remains why create PAG pipeline or in which specific scenario it will be useful.

[FurkanGozukara]

@FurkanGozukara @chenjy2003

which one is recommended it doesnt tell much

Based on my test results PAG quality is 60-70% lower. The question remains why create PAG pipeline or in which specific scenario it will be useful.

thanks a lot

[lawrence-cj]

Based on my test results PAG quality is 60-70% lower. The question remains why create PAG pipeline or in which specific scenario it will be useful.

PAG is a function. Anyone can test and try to improve it. We just did the basic work for the community.

[nitinmukesh]

PAG is a function. Anyone can test and try to improve it. We just did the basic work for the community.

Makes sense. So it's primarily for developers to improve or enhance it.

[bghira]

this guy Furkan should just be blocked to be honest, creates too much noise and demands a lot from others without thanks

[lawrence-cj]

Makes sense. So it's primarily for developers to improve or enhance it.

Correct.