Split model components across GPUs with device_map

The device_map argument places model components on separate GPUs. Use it when a model doesn’t fit on one GPU and your GPUs don’t have a fast interconnect. Refer to Distributed inference to compare it with the other strategies, and to Device placement for more device_map options.

The example below loads each component only for the stage that needs it, and uses device_map to spread the large ones across two GPUs.

Encode the prompt

Set device_map="balanced" to place the text encoders on separate GPUs. Use max_memory to cap how much memory each GPU uses. Set transformer and vae to None so only the text encoders load.

from diffusers import FluxPipeline
import torch

prompt = """
cinematic film still of a cat sipping a margarita in a pool in Palm Springs, California
highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain
"""

pipeline = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    transformer=None,
    vae=None,
    device_map="balanced",
    max_memory={0: "16GB", 1: "16GB"},
    dtype=torch.bfloat16
)
with torch.no_grad():
    print("Encoding prompts.")
    prompt_embeds, pooled_prompt_embeds, text_ids = pipeline.encode_prompt(
        prompt=prompt, prompt_2=None, max_sequence_length=512
    )

Once you have the text embeddings, remove the text encoders from the GPU to make room for the transformer.

import gc

def flush():
    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.reset_max_memory_allocated()
    torch.cuda.reset_peak_memory_stats()

del pipeline.text_encoder
del pipeline.text_encoder_2
del pipeline.tokenizer
del pipeline.tokenizer_2
del pipeline

flush()

Denoise with the sharded transformer

Set device_map="auto" to let Accelerate place the transformer. It fills the first GPU before moving to the next, and uses the CPU or disk only when the GPUs run out of memory. On two smaller GPUs, the transformer ends up split across both. Weights on the CPU or disk slow down inference.

from diffusers import AutoModel
import torch

transformer = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    device_map="auto",
    dtype=torch.bfloat16
)

Check transformer.hf_device_map to see how the transformer is distributed across devices. A pipeline loaded with device_map also has a hf_device_map attribute that shows where each component is placed.

Load a pipeline with only the transformer, and set output_type="latent" so it returns latents instead of images.

pipeline = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    text_encoder=None,
    text_encoder_2=None,
    tokenizer=None,
    tokenizer_2=None,
    vae=None,
    transformer=transformer,
    dtype=torch.bfloat16
)

print("Running denoising.")
height, width = 768, 1360
latents = pipeline(
    prompt_embeds=prompt_embeds,
    pooled_prompt_embeds=pooled_prompt_embeds,
    num_inference_steps=50,
    guidance_scale=3.5,
    height=height,
    width=width,
    output_type="latent",
).images

Decode the latents

Remove the pipeline and transformer from memory and load a VAE to decode the latents. The VAE is small enough to fit on one GPU.

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

del pipeline.transformer
del pipeline
del transformer

flush()

vae = AutoencoderKL.from_pretrained(
    "black-forest-labs/FLUX.1-dev", subfolder="vae", dtype=torch.bfloat16
).to("cuda")  # or "mps", "xpu", "cpu"
vae_scale_factor = 2 ** (len(vae.config.block_out_channels) - 1)
image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)

with torch.no_grad():
    print("Running decoding.")
    latents = FluxPipeline._unpack_latents(latents, height, width, vae_scale_factor)
    latents = (latents / vae.config.scaling_factor) + vae.config.shift_factor

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

By selectively loading and unloading the models you need at a given stage and sharding the largest models across multiple GPUs, you can run inference with large models on consumer GPUs.

Update on GitHub

Diffusers

Split model components across GPUs with device_map

Encode the prompt

Denoise with the sharded transformer

Decode the latents