feat: allow flux transformer to be sharded during inference

[sayakpaul]

What does this PR do?

Adds support to shard the Flux transformer across multiple devices.

Here's how to run it:

generate_embeddings.py

from diffusers import FluxPipeline
import torch

pipeline = FluxPipeline.from_pretrained(
    "black-forest-labs/flux.1-dev", transformer=None, vae=None, torch_dtype=torch.bfloat16
)
pipeline.enable_model_cpu_offload()

prompt = "a cute fish holding a sign saying 'hello world!'"

with torch.no_grad():
    prompt_embeds, pooled_prompt_embeds, text_ids = pipeline.encode_prompt(
        prompt, prompt_2=None, max_sequence_length=512, num_images_per_prompt=4
    )

torch.save(prompt_embeds, "prompt_embeds.pt")
torch.save(pooled_prompt_embeds, "pooled_prompt_embeds.pt")
torch.save(text_ids, "prompt_attention_mask.pt")

This will generate and serialize the embeddings to the disk.

run_denoising_loop.py

from diffusers import FluxTransformer2DModel, FluxPipeline
import torch 

max_memory = {0: "16GB", 1: "16GB"}
ckpt_id = "black-forest-labs/flux.1-dev"
model = FluxTransformer2DModel.from_pretrained(
    ckpt_id, 
    subfolder="transformer",
    device_map="auto",
    max_memory=max_memory, 
    torch_dtype=torch.bfloat16,
)
pipeline = FluxPipeline.from_pretrained(
    ckpt_id,
    transformer=model,
    text_encoder=None,
    text_encoder_2=None,
    tokenizer=None,
    tokenizer_2=None,
    vae=None,
    torch_dtype=torch.bfloat16,
)

height, width = 768, 1360
latents = pipeline(
    prompt_embeds=torch.load("prompt_embeds.pt"),
    pooled_prompt_embeds=torch.load("pooled_prompt_embeds.pt"),
    num_inference_steps=50,
    guidance_scale=3.5,
    height=height,
    width=width,
    output_type="latent",
).images
print(f"{latents.shape=}")
torch.save(latents, "latents.pt")

This does the following:

• Shows how to distribute the Flux transformer into two GPUs (mimicking 16GBs for each).
• Run the denoising loop with the serialized text embeddings.
• Serializes the final latents for decoding.

decode_latents.py

from diffusers import FluxPipeline, AutoencoderKL
from diffusers.image_processor import VaeImageProcessor
import torch

vae = AutoencoderKL.from_pretrained("black-forest-labs/flux.1-dev", subfolder="vae", torch_dtype=torch.bfloat16).to(
    "cuda"
)

latents = torch.load("latents.pt")
height, width = 768, 1360
vae_scale_factor = 2 ** (len(vae.config.block_out_channels))
image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)

latents = FluxPipeline._unpack_latents(latents, height, width, vae_scale_factor)
latents = (latents / vae.config.scaling_factor) + vae.config.shift_factor

with torch.no_grad():
    image = vae.decode(latents, return_dict=False)[0]
image = image_processor.postprocess(image, output_type="pil")
image[0].save("image.png")

We get:

("a cute fish holding a sign saying 'hello world!'")

The tests were run with the following command:

 CUDA_VISIBLE_DEVICES=0,1 pytest tests/models/transformers/test_models_transformer_flux.py -k "offload"

[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 ! You mean CUDA_VISIBLE_DEVICES=0,1 pytest tests/models/transformers/test_models_transformer_flux.py -k "offload" instead no ? Otherwise, it is only using one gpu

[sayakpaul]

@SunMarc that's right!

[notdanilo]

*** Looking for the merge button ***

Damn! I don't have write access.

[DN6]

LGTM 👍🏽

[sayakpaul]

Failing test is unrelated.