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)