Gradient checkpointing not applied to UNet mid_block

[laksjdjf]

The mid block of the SDXL is huge, so fixing it significantly reduces VRAM usage.

I have tested the following changes and have seen great results.
main...laksjdjf:diffusers:mid_block_gradient_checkpointing

However, there seem to be several variations of midblock, including UNetMidBlock2DSimpleCrossAttn and UNetMidBlock3DCrossAttn, and I am not sure what to do with them.

By the way, Kohya's trainer applies gradient checkpointing to all blocks.
https://github.com/kohya-ss/sd-scripts/blob/4072f723c12822e2fa1b2e076cc1f90b8f4e30c9/library/sdxl_original_unet.py#L1035-L1041

[sayakpaul]

Cc: @patrickvonplaten

[patrickvonplaten]

Great catch - yes for SDXL we should indeed apply gradient checkpointing to the midblock as well :-)

Would you like to open a PR for it? This would be a great addition for the community I believe :-)