# Distribua os componentes do modelo entre GPUs com device_map

O argumento `device_map` posiciona os componentes do modelo em GPUs separadas. Use-o quando um modelo não couber em uma única GPU e suas GPUs não possuírem um interconector rápido. Consulte [Inferência distribuída](../training/distributed_inference) para compará-lo com as outras estratégias e [Posicionamento de dispositivo](./memory#device-placement) para mais opções de `device_map`.

O exemplo abaixo carrega cada componente apenas para o estágio que precisa dele e usa `device_map` para distribuir os maiores entre duas GPUs.

## Codifique o prompt

Defina `device_map="balanced"` para posicionar os codificadores de texto em GPUs separadas. Use `max_memory` para limitar a quantidade de memória que cada GPU utiliza. Defina `transformer` e `vae` como `None` para que apenas os codificadores de texto sejam carregados.

```py
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
    )
```

Depois de obter os embeddings de texto, remova os codificadores de texto da GPU para liberar espaço para o transformer.

```py
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()
```

## Desruído com o transformer fragmentado

Defina `device_map="auto"` para permitir que o Accelerate posicione o transformer. Ele preenche a primeira GPU antes de passar para a próxima e usa a CPU ou o disco apenas quando as GPUs ficam sem memória. Em duas GPUs menores, o transformer acaba dividido entre ambas. Pesos na CPU ou no disco deixam a inferência mais lenta.

```py
from diffusers import AutoModel
import torch

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

> [!TIP]
> Verifique `transformer.hf_device_map` para ver como o transformer é distribuído entre os dispositivos. Um pipeline carregado com `device_map` também possui um atributo `hf_device_map` que mostra onde cada componente é colocado.

Carregue um pipeline com apenas o transformer e defina `output_type="latent"` para que ele retorne latentes em vez de imagens.

```py
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
```

## Decodificar os latentes

Remova o pipeline e o transformer da memória e carregue um VAE para decodificar os latentes. O VAE é pequeno o suficiente para caber em uma única GPU.

```py
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")
```

Ao carregar e descarregar seletivamente os modelos necessários em cada etapa e dividir os maiores modelos entre várias GPUs, é possível executar inferência com modelos grandes em GPUs de consumo.