diffusers: device_map="auto" in AutoPipelineForText2Image raises error

Describe the bug

Thanks for ur contributions! I’m using AutoPipelineForText2Image and set device_map="auto" when initializing, but the program raises Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cuda:0! when inferring. Even I manually put the essential components on target devices(text_encoder and unet on cuda:0, vae on cuda:1) and built a pipeline through StableDiffusionPipeline this keeps happening. Should I rewrite the pipeline process in my code?

Reproduction

This is a simple script for reproduction:

from diffusers import AutoPipelineForText2Image

pipeline = AutoPipelineForText2Image.from_pretrained(
	"stabilityai/stable-diffusion-2-1", device_map="auto"
)

generator = torch.Generator().manual_seed(31)
image = pipeline("Astronaut in a jungle, cold color palette, muted colors, detailed, 8k", generator=generator).images[0]

Logs

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

System Info

I’m using python=3.11, diffusers=0.24.0, accelerate=0.25.0, torch=2.1.2 with CUDA=11.8

Who can help?

@yiyixuxu @DN6