FluxInpaintPipeline error when im using DataParallel to use multiple gpus

[Himasnhu-AT]

Describe the bug

line 175, in forward
for t in chain(self.module.parameters(), self.module.buffers()):
^^^^^^^^^^^^^^^^^^^^^^
File "/opt/tensorflow/lib/python3.12/site-packages/diffusers/configuration_utils.py", line 143, in getattr
raise AttributeError(f"'{type(self).name}' object has no attribute '{name}'")
AttributeError: 'FluxFillPipeline' object has no attribute 'parameters'

Reproduction

if torch.cuda.device_count() > 1:
pipe = torch.nn.DataParallel(pipe)
pipe = pipe.to("cuda")

Logs

System Info

Ubuntu AWS deep learning AMI

g5.12xlarge
4A10 GPUS

Who can help?

No response

[a-r-r-o-w]

DP needs to be applied on torch modules. The pipelines are not torch modules, so it is expected to not work. Instead, you can initialise the modules separately, wrap with DP, then create a per-device pipeline by initialising with the modules.

[Himasnhu-AT]

okay, is there a way to make it part of the pipeline itself as it'll be easier for other to use it too

[Himasnhu-AT]

here's im doing it:

pipe = FluxFillPipeline.from_pretrained("black-forest-labs/FLUX.1-Fill-dev", torch_dtype=torch.bfloat16, 
device_map="balanced", max_memory={0:"23GB", 1: "23GB", 2: "23GB", 3: "23GB"})


    image = pipe(
        # mask_image_latent=vae.encode(controlImage),
        prompt=prompt,
        prompt_2=prompt,
        image=image,
        mask_image=mask,
        height=height,
        width=width,
        guidance_scale=guidance_scale,
        # strength=strength,
        num_inference_steps=num_inference_steps,
        # generator=torch.Generator(device='cuda').manual_seed(seed),
        # generator=torch.Generator().manual_seed(seed),
        # lora_scale=0.75 // not supported in this version
    ).images[0]

Error:

Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! 
(when checking argument for argument mat2 in method wrapper_CUDA_mm)

[a-r-r-o-w]

okay, is there a way to make it part of the pipeline itself as it'll be easier for other to use it too

Not really. Data parallel should be applied per torch module, not on the container holding the modules. Pipeline objects should be created per process to avoid this error. The simple way of doing distributed inference with diffusers would look something like:

import torch
import torch.distributed as dist
from diffusers import FluxPipeline

dist.init_process_group("nccl")
world_size = dist.get_world_size()
rank = dist.get_rank()

torch.cuda.set_device(rank)

pipe = FluxPipeline.from_pretrained(...).to("cuda")

with open("prompts.txt", "r") as file:
    prompts = [line.strip() for line in file.readlines()]

num_chunks = len(prompts) // world_size
prompts = prompts[rank * num_chunks : (rank + 1) * num_chunks]

for i, prompt in enumerate(prompts):
    output = pipe(prompt, ...).images[0]  # or, do batched inference per rank
    output.save(f"{rank}-{i}-{escape_special_characters(prompt)}.png")

dist.destroy_process_group()

Then run the code using:

torchrun --nodes 1 --nproc_per_node <NUM_GPUS> test.py

It's best to do it similar to this way because Diffusers schedulers are not thread safe anyway, so you need to create an instance per rank.

[joangava]

def quick_sort(arr):
if len(arr) <= 1:
return arr
else:
pivot = arr[len(arr) // 2]
left = [x for x in arr if x < pivot]
middle = [x for x in arr if x == pivot]
right = [x for x in arr if x > pivot]
return quick_sort(left) + middle + quick_sort(right)