feat: cuda device_map for pipelines.

[sayakpaul]

What does this PR do?

TL;DR: This PR adds a device_map option at the pipeline-level to speed up end-to-end pipeline loading on a target device.

To benefit from #11904, users have to follow this pattern:

from diffusers import AutoModel, DiffusionPipeline
import torch 

# first initialize the model on device_map to benefit from fast loading
model = AutoModel.from_pretrained(..., device_map="cuda", torch_dtype=...)

# initialize the pipeline
pipe = DiffusionPipeline.from_pretrained(..., transformer=model, torch_dtype=...)

# place other stuff on cuda
for name, component in pipe.components.items():
     if name != "model_already_loaded_above":
         component.to("cuda")

# run inference
...

We could improve the UX a bit by letting the users pass a device_map="cuda" (or whatever valid value) WHILE initializing the pipe. This PR tackles that =>

from diffusers import DiffusionPipeline
import torch

pipe = DiffusionPipeline.from_pretrained(..., device_map="cuda", torch_dtype=...)
...

For pipelines like Flux, passing device_map for loading the text encoders might be as same as doing to(). However, for pipelines like Qwen-Image, that use a mid-range model like Qwen25VL-7B, passing device_map="cuda" while initializing the pipeline should be beneficial (of course, the target device should have enough VRAM to support this). Below are the results I got for Qwen-Image (with cold cache):

time: 8.494s (no device_map)
time: 6.678s (device_map)

Code

import time
t_ini = time.time()

import torch
from diffusers import DiffusionPipeline
print(f"import time: {time.time() - t_ini:.3f}s")

model_id = "Qwen/Qwen-Image"

t0 = time.time()
torch.cuda.synchronize()
print(f"CUDA sync time: {time.time() - t0:.3f}s")

print("starting pipe load")
t1 = time.time()
pipe = DiffusionPipeline.from_pretrained(
    model_id, torch_dtype=torch.bfloat16, device_map="cuda"
)
torch.cuda.synchronize()
t2 = time.time()

diff = t2 - t1
print(f"time: {diff:.3f}s")

print(getattr(pipe, "hf_device_map", None))
_ = pipe("dog", num_inference_steps=2)

Any objections?

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[DN6]

LGTM. Could we add a simple fast GPU test

[sayakpaul]

@DN6 done! I have added a test, too.

[DN6]

I would prefer not to import from testing_utils for non-test modules (in fact we should move this module out of src). It would be better to redefine the relevant torch_device functionality in torch_utils and import from there.

[sayakpaul]

Yeah my thoughts too. We should actually move torch_device to utils.