feat: allow sharding for auraflow.

[sayakpaul]

What does this PR do?

AuraFlow is one of the largest open model right now. A 5.6B parameter model maybe very hard to run on consumer GPUs having 15GBs if we are not careful.

This PR adds support to run AuraFlow Transformer in 15GBs by allowing offloading at the model level.

Example code:

from diffusers import AuraFlowTransformer2DModel
import tempfile
import torch
import os

def get_inputs():
    sample = torch.randn(1, 4, 128, 128)
    timestep = torch.randint(0, 1000, size=(1, ))
    encoder_hidden_states = torch.randn(1, 256, 2048)

    return sample, timestep, encoder_hidden_states

with torch.no_grad():
    max_memory = {0: "15GB"} # reasonable estimate for a consumer-gpu.
    with tempfile.TemporaryDirectory() as tmp_dir:
        new_model = AuraFlowTransformer2DModel.from_pretrained(
            "fal/AuraFlow", 
            subfolder="transformer",
            device_map="auto",
            max_memory=max_memory, 
            offload_folder=os.path.join("/raid/.cache/huggingface", tmp_dir)
        )

        sample, timestep, encoder_hidden_states = get_inputs()
        out = new_model(
            hidden_states=sample,
            encoder_hidden_states=encoder_hidden_states,
            timestep=timestep,
        ).sample
        print(out.shape)

I have run the offloading tests and they pass too:

pytest tests/models/transformers/test_models_pixart_transformer2d.py -k "offload"

@SunMarc FYI.

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

[SunMarc]

Nice !