make examples in docs device agnostic

[kaixuanliu]

Following #12161, which aims to fix #14575

[kaixuanliu]

@stevhliu Hi, could you help review this as well? Thx!

[yao-matrix]

@IlyasMoutawwakil , pls help review, thx.

[stevhliu]

hi, thanks for your PR! i think here for Diffusers, we'd like to keep it as "cuda" for now but we're also thinking more about how we can make it clear that Diffusers also supports other devices. let me know if you have any ideas!

[IlyasMoutawwakil]

yeh this complicates snippets that are meant to be minimal, with transformers "auto" is already a common and well known value for devie_map, but here this more like pure pytorch so users should just replace "cuda" with "xpu" not use a get_device util.

[yao-matrix]

Possible to pass device_map="auto" to from_pretrained(), then we can avoid to() which is not best practice and will lead to spike CPU memory usage as staging buffer. @IlyasMoutawwakil @kaixuanliu @stevhliu

[IlyasMoutawwakil]

i'm not sure how optimized is device_map api in diffusers (zero-copy or not), cc @sayakpaul

[sayakpaul]

device_map="cuda" follows the same strategy as transformers. Also, what do you mean by zero-copy? Something like the following?

with torch.device("meta"):
    model = init_model()
model.load_state_dict(..., assign=True)

If so, yes.

[sayakpaul]

"auto", within the context of pipelines changes meaning, though. A DiffusionPipeline is NOT a single model, but a collection of different models connected in a directed graph. So, we shouldn't use "auto", IMO.

[kaixuanliu]

For pipelines, we use "balanced" or get_device: L124, here in diffusers we already use get_device to do device-agnostic processing, so I think in the doc it is ok to use get_device utils, WDYT? @IlyasMoutawwakil @sayakpaul @yao-matrix , or we can use device = torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else "cpu" instead. device_map=balanced is another option, but it is not suitable for all the cases.

[sayakpaul]

IMO we're still complicating the things that are supposed to be very simpler. I am not sure if the approaches adapted in Transformers directly transfer because of the complexities around the pipeline system. Perhaps it's better to document that in a central manner.

[kaixuanliu]

Then can we add comment like device = "cuda" # or "mps", "xpu", "cpu" ? This will not bring actual change to all the existing examples.

[yao-matrix]

@kaixuanliu , could you change to use device_map="auto" rather than to()?

[github-actions]

Hi @kaixuanliu, thanks for the PR! It does not appear to link an issue it fixes. If this PR addresses an existing issue, please add a closing keyword (e.g. Fixes #1234) to the PR description so the issue is linked. See the contribution guide for more details. If this PR intentionally does not fix a tracked issue, a maintainer can add the no-issue-needed label to silence this reminder.

Please note that PRs without a linked issue are likely to be automatically closed 10 days after this notice.

Once the PR links an issue (or gets the no-issue-needed label), you can ignore this message — it stays here as a comment, but it no longer applies.

[kaixuanliu]

@sayakpaul @stevhliu @IlyasMoutawwakil Hi, I have updated the code by adding comment。 This is the least intrusive method I can think of... Is it OK now?

[stevhliu]

thanks, i think this is better and we can also add a sentence or two in the Installation doc that says something like "Diffusers runs on any accelerator, swap in ..."

[kaixuanliu]

@stevhliu Thx for your advice. done for it.

[stevhliu]

thanks! 🤗

[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.