Using Flux.1-dev on EC2 instance

[Devank-Garg]

Hi, I am trying to using flux.1-dev on an EC2 cluster (g5.12xlarge). Which has 4 A10s.

import torch
from diffusers import FluxPipeline
from accelerate import PartialState

model_path = "black-forest-labs/FLUX.1-dev"

pipe = FluxPipeline.from_pretrained(model_path, torch_dtype=torch.bfloat16, device_map="balanced", max_memory={0:"24GB", 1:"24GB", 2:"24GB", 3:"24GB"})
#pipe.enable_model_cpu_offload() #save some VRAM by offloading the model to CPU. Remove this if you have enough GPU power
#pipe.enable_sequential_cpu_offload()

print(pipe.hf_device_map)

prompt = " cat holding a sign that says hello world lovely, 8k"
image = pipe(
prompt,
height=1024,
width=1024,
guidance_scale=3.5,
output_type="pil",
num_inference_steps=20,
max_sequence_length=512,

).images[0]
image.save("flux-dev.png")

I am trying to use this snippet but I am getting
{'transformer': 'cpu', 'text_encoder_2': 0, 'text_encoder': 1, 'vae': 2}
from print(pipe.hf_device_map)

Can anyone tell me what else can I do to increase inference speed?